s2n_codec/encoder/
scatter.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::{Encoder, EncoderBuffer};
5
6pub struct Buffer<'a> {
7    inner: EncoderBuffer<'a>,
8    #[cfg(feature = "bytes")]
9    extra: Option<bytes::Bytes>,
10}
11
12impl<'a> Buffer<'a> {
13    /// Ensures an extra bytes are written into the main EncoderBuffer.
14    #[inline]
15    pub fn flatten(&mut self) -> &mut EncoderBuffer<'a> {
16        self.flush();
17        &mut self.inner
18    }
19}
20
21/// Implement a version with `bytes` enabled
22#[cfg(feature = "bytes")]
23impl<'a> Buffer<'a> {
24    /// Initializes a buffer without any extra bytes at the end
25    #[inline]
26    pub fn new(inner: EncoderBuffer<'a>) -> Self {
27        Self { inner, extra: None }
28    }
29
30    /// Initializes the buffer with extra bytes
31    ///
32    /// NOTE the EncoderBuffer position should not include the extra bytes. This ensures the bytes
33    /// can be "flushed" into the EncoderBuffer if another write happens or `flatten` is called.
34    #[inline]
35    pub fn new_with_extra(inner: EncoderBuffer<'a>, extra: Option<bytes::Bytes>) -> Self {
36        Self { inner, extra }
37    }
38
39    /// Converts the buffer into its inner parts
40    ///
41    /// NOTE: the EncoderBuffer position will not include the extra bytes. The caller will need to
42    /// account for this data.
43    #[inline]
44    pub fn into_inner(self) -> (EncoderBuffer<'a>, Option<bytes::Bytes>) {
45        (self.inner, self.extra)
46    }
47
48    /// Returns the inner EncoderBuffer, along with the current extra bytes
49    #[inline]
50    pub fn inner_mut(&mut self) -> (&mut EncoderBuffer<'a>, &Option<bytes::Bytes>) {
51        (&mut self.inner, &self.extra)
52    }
53
54    /// Resets the encoder to its initial state while also dropping any extra bytes at the end
55    #[inline]
56    pub fn clear(&mut self) {
57        self.inner.set_position(0);
58        self.extra = None;
59    }
60
61    #[inline]
62    fn flush(&mut self) {
63        // move the extra bytes into the main buffer
64        if let Some(extra) = self.extra.take() {
65            self.inner.write_slice(&extra);
66        }
67    }
68}
69
70/// Include a functional implementation when `bytes` is not available
71#[cfg(not(feature = "bytes"))]
72impl<'a> Buffer<'a> {
73    #[inline]
74    pub fn new(inner: EncoderBuffer<'a>) -> Self {
75        Self { inner }
76    }
77
78    #[inline]
79    pub fn new_with_extra(inner: EncoderBuffer<'a>, extra: Option<&'static [u8]>) -> Self {
80        debug_assert!(extra.is_none());
81        Self { inner }
82    }
83
84    #[inline]
85    pub fn into_inner(self) -> (EncoderBuffer<'a>, Option<&'static [u8]>) {
86        (self.inner, None)
87    }
88
89    #[inline]
90    pub fn inner_mut(&mut self) -> (&mut EncoderBuffer<'a>, &Option<&'static [u8]>) {
91        (&mut self.inner, &None)
92    }
93
94    #[inline]
95    pub fn clear(&mut self) {
96        self.inner.set_position(0);
97    }
98
99    #[inline(always)]
100    fn flush(&mut self) {}
101}
102
103impl Encoder for Buffer<'_> {
104    /// We have special handling for writes that include `Bytes` so signal that
105    #[cfg(feature = "bytes")]
106    const SPECIALIZES_BYTES: bool = true;
107
108    #[inline]
109    fn write_sized<F: FnOnce(&mut [u8])>(&mut self, len: usize, write: F) {
110        // we need to flush the extra bytes if we have them
111        self.flush();
112        self.inner.write_sized(len, write)
113    }
114
115    #[inline]
116    fn write_slice(&mut self, slice: &[u8]) {
117        // we need to flush the extra bytes if we have them
118        self.flush();
119        self.inner.write_slice(slice);
120    }
121
122    #[inline]
123    #[cfg(feature = "bytes")]
124    fn write_bytes(&mut self, bytes: bytes::Bytes) {
125        // we need to flush the extra bytes if we have them and replace them with the new one
126        self.flush();
127        // ensure the underlying buffer is big enough for the write
128        self.inner.assert_capacity(bytes.len());
129        // store the extra bytes and defer the copy
130        self.extra = Some(bytes);
131    }
132
133    #[inline]
134    fn write_zerocopy<
135        T: zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Unaligned,
136        F: FnOnce(&mut T),
137    >(
138        &mut self,
139        write: F,
140    ) {
141        // we need to flush the extra bytes if we have them
142        self.flush();
143        self.inner.write_zerocopy(write);
144    }
145
146    #[inline]
147    fn write_repeated(&mut self, count: usize, value: u8) {
148        // we need to flush the extra bytes if we have them
149        self.flush();
150        self.inner.write_repeated(count, value)
151    }
152
153    #[inline]
154    fn capacity(&self) -> usize {
155        self.inner.capacity()
156    }
157
158    #[inline]
159    #[cfg(feature = "bytes")]
160    fn len(&self) -> usize {
161        let mut len = self.inner.len();
162
163        // make sure our len includes the extra bytes if we have them
164        if let Some(extra) = self.extra.as_ref() {
165            len += extra.len();
166        }
167
168        len
169    }
170
171    #[inline]
172    #[cfg(not(feature = "bytes"))]
173    fn len(&self) -> usize {
174        self.inner.len()
175    }
176}
s2n_codec/encoder/scatter.rs

s2n_codec/encoder/
scatter.rs
