Skip to main content

ax_codec_core/
buffer.rs

1use crate::{BufferReader, BufferWriter, DecodeError, EncodeError};
2
3#[cfg(feature = "alloc")]
4pub struct VecWriter {
5    inner: alloc::vec::Vec<u8>,
6}
7
8#[cfg(feature = "alloc")]
9impl VecWriter {
10    #[inline]
11    pub fn new() -> Self {
12        Self {
13            inner: alloc::vec::Vec::new(),
14        }
15    }
16
17    #[inline]
18    pub fn with_capacity(cap: usize) -> Self {
19        Self {
20            inner: alloc::vec::Vec::with_capacity(cap),
21        }
22    }
23
24    #[inline]
25    pub fn into_vec(self) -> alloc::vec::Vec<u8> {
26        self.inner
27    }
28
29    #[inline]
30    pub fn as_slice(&self) -> &[u8] {
31        &self.inner
32    }
33}
34
35#[cfg(feature = "alloc")]
36impl Default for VecWriter {
37    #[inline]
38    fn default() -> Self {
39        Self::new()
40    }
41}
42
43#[cfg(feature = "alloc")]
44impl BufferWriter for VecWriter {
45    #[inline(always)]
46    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodeError> {
47        self.inner.extend_from_slice(buf);
48        Ok(())
49    }
50}
51
52#[cfg(feature = "alloc")]
53pub struct PooledVecWriter {
54    inner: alloc::vec::Vec<u8>,
55    recycled: bool,
56}
57
58#[cfg(feature = "alloc")]
59impl PooledVecWriter {
60    #[inline]
61    pub fn new() -> Self {
62        Self {
63            inner: crate::pool::take(256),
64            recycled: false,
65        }
66    }
67
68    #[inline]
69    pub fn with_capacity(cap: usize) -> Self {
70        Self {
71            inner: crate::pool::take(cap),
72            recycled: false,
73        }
74    }
75
76    #[inline]
77    pub fn as_slice(&self) -> &[u8] {
78        &self.inner
79    }
80
81    #[inline]
82    pub fn into_vec(mut self) -> alloc::vec::Vec<u8> {
83        self.recycled = true;
84        core::mem::take(&mut self.inner)
85    }
86}
87
88#[cfg(feature = "alloc")]
89impl Default for PooledVecWriter {
90    #[inline]
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96#[cfg(feature = "alloc")]
97impl Drop for PooledVecWriter {
98    fn drop(&mut self) {
99        if !self.recycled {
100            let buf = core::mem::take(&mut self.inner);
101            crate::pool::recycle(buf);
102        }
103    }
104}
105
106#[cfg(feature = "alloc")]
107impl BufferWriter for PooledVecWriter {
108    #[inline(always)]
109    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodeError> {
110        self.inner.extend_from_slice(buf);
111        Ok(())
112    }
113}
114
115pub struct SliceReader<'a> {
116    buf: &'a [u8],
117    pos: usize,
118}
119
120impl<'a> SliceReader<'a> {
121    #[inline]
122    pub fn new(buf: &'a [u8]) -> Self {
123        Self { buf, pos: 0 }
124    }
125
126    #[inline]
127    pub fn remaining_len(&self) -> usize {
128        self.buf.len().saturating_sub(self.pos)
129    }
130
131    #[inline]
132    pub fn position(&self) -> usize {
133        self.pos
134    }
135}
136
137impl<'a> BufferReader<'a> for SliceReader<'a> {
138    #[inline(always)]
139    fn peek(&self) -> Option<u8> {
140        self.buf.get(self.pos).copied()
141    }
142
143    #[inline(always)]
144    fn next(&mut self) -> Option<u8> {
145        let byte = self.buf.get(self.pos).copied()?;
146        self.pos += 1;
147        Some(byte)
148    }
149
150    #[inline(always)]
151    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), DecodeError> {
152        let end = self
153            .pos
154            .checked_add(buf.len())
155            .filter(|&end| end <= self.buf.len())
156            .ok_or(DecodeError::UnexpectedEOF)?;
157        buf.copy_from_slice(&self.buf[self.pos..end]);
158        self.pos = end;
159        Ok(())
160    }
161
162    #[inline(always)]
163    fn remaining(&self) -> &'a [u8] {
164        &self.buf[self.pos..]
165    }
166
167    #[inline(always)]
168    fn advance(&mut self, n: usize) -> Result<(), DecodeError> {
169        let end = self
170            .pos
171            .checked_add(n)
172            .filter(|&end| end <= self.buf.len())
173            .ok_or(DecodeError::UnexpectedEOF)?;
174        self.pos = end;
175        Ok(())
176    }
177}
178
179pub struct StackWriter<const N: usize> {
180    buf: [u8; N],
181    pos: usize,
182}
183
184impl<const N: usize> Default for StackWriter<N> {
185    #[inline]
186    fn default() -> Self {
187        Self::new()
188    }
189}
190
191impl<const N: usize> StackWriter<N> {
192    #[inline]
193    pub fn new() -> Self {
194        Self {
195            buf: [0u8; N],
196            pos: 0,
197        }
198    }
199
200    #[inline]
201    pub fn as_slice(&self) -> &[u8] {
202        &self.buf[..self.pos]
203    }
204}
205
206impl<const N: usize> BufferWriter for StackWriter<N> {
207    #[inline]
208    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodeError> {
209        let end = self
210            .pos
211            .checked_add(buf.len())
212            .ok_or(EncodeError::InsufficientCapacity)?;
213        if end > N {
214            return Err(EncodeError::InsufficientCapacity);
215        }
216        self.buf[self.pos..end].copy_from_slice(buf);
217        self.pos = end;
218        Ok(())
219    }
220}