s2n_codec/encoder/
buffer.rs

1// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::encoder::Encoder;
5
6/// EncoderBuffer is a buffer for writing to a mutable slice
7#[derive(Debug, Eq, PartialEq, PartialOrd, Ord)]
8pub struct EncoderBuffer<'a> {
9    bytes: &'a mut [u8],
10    position: usize,
11}
12
13impl<'a> EncoderBuffer<'a> {
14    /// Creates a new `EncoderBuffer`
15    #[inline]
16    pub fn new(bytes: &'a mut [u8]) -> Self {
17        Self { bytes, position: 0 }
18    }
19
20    /// Sets the write cursor to a new position
21    ///
22    /// # Panics
23    /// Panics when `position > capacity`
24    #[inline]
25    pub fn set_position(&mut self, position: usize) {
26        debug_assert!(
27            position <= self.capacity(),
28            "position {position} exceeded capacity of {}",
29            self.capacity()
30        );
31        self.position = position;
32    }
33
34    /// Advances the write cursor by offset
35    ///
36    /// # Panics
37    /// Panics when `position > capacity`
38    #[inline]
39    pub fn advance_position(&mut self, offset: usize) {
40        let position = self.position + offset;
41        self.set_position(position)
42    }
43
44    /// Splits off the used buffer from the remaining bytes
45    #[inline]
46    pub fn split_off(self) -> (&'a mut [u8], &'a mut [u8]) {
47        self.bytes.split_at_mut(self.position)
48    }
49
50    /// Splits the used buffer from the remaining bytes
51    #[inline]
52    pub fn split_mut(&mut self) -> (&mut [u8], &mut [u8]) {
53        self.bytes.split_at_mut(self.position)
54    }
55
56    /// Returns the written bytes as a mutable slice
57    #[inline]
58    pub fn as_mut_slice(&mut self) -> &mut [u8] {
59        unsafe { self.bytes.get_unchecked_mut(..self.position) }
60    }
61
62    #[inline]
63    pub(crate) fn assert_capacity(&self, len: usize) {
64        debug_assert!(
65            len <= self.remaining_capacity(),
66            "not enough buffer capacity. wanted: {}, available: {}",
67            len,
68            self.remaining_capacity()
69        );
70    }
71}
72
73impl Encoder for EncoderBuffer<'_> {
74    #[inline]
75    fn write_sized<F: FnOnce(&mut [u8])>(&mut self, len: usize, write: F) {
76        self.assert_capacity(len);
77        let end = self.position + len;
78        let bytes = unsafe {
79            // Safety: bounds already checked
80            self.bytes.get_unchecked_mut(self.position..end)
81        };
82        write(bytes);
83        self.position = end;
84    }
85
86    #[inline]
87    fn write_slice(&mut self, slice: &[u8]) {
88        self.write_sized(slice.len(), |dest| dest.copy_from_slice(slice));
89    }
90
91    #[inline]
92    fn write_repeated(&mut self, count: usize, value: u8) {
93        self.write_sized(count, |dest| {
94            for byte in dest {
95                *byte = value;
96            }
97        })
98    }
99
100    #[inline]
101    fn write_zerocopy<
102        T: zerocopy::IntoBytes + zerocopy::FromBytes + zerocopy::Unaligned,
103        F: FnOnce(&mut T),
104    >(
105        &mut self,
106        write: F,
107    ) {
108        let len = core::mem::size_of::<T>();
109        self.write_sized(len, |bytes| {
110            let value = unsafe {
111                // The `zerocopy` markers ensure this is a safe operation
112                &mut *(bytes as *mut _ as *mut T)
113            };
114            write(value)
115        })
116    }
117
118    #[inline]
119    fn capacity(&self) -> usize {
120        self.bytes.len()
121    }
122
123    #[inline]
124    fn len(&self) -> usize {
125        self.position
126    }
127}
s2n_codec/encoder/buffer.rs

s2n_codec/encoder/
buffer.rs
