ax_codec_bytes/
bytes_impl.rs

1use ax_codec_core::{BufferReader, BufferWriter, DecodeError, EncodeError};
2use bytes::{Bytes, BytesMut};
3
4#[derive(Debug, Clone)]
5pub struct BytesReader<'a> {
6    remaining: &'a [u8],
7}
8
9impl<'a> BytesReader<'a> {
10    #[inline]
11    pub fn new(bytes: &'a Bytes) -> Self {
12        Self {
13            remaining: bytes.as_ref(),
14        }
15    }
16
17    #[inline]
18    pub fn remaining_len(&self) -> usize {
19        self.remaining.len()
20    }
21}
22
23impl<'a> BufferReader<'a> for BytesReader<'a> {
24    #[inline]
25    fn peek(&self) -> Option<u8> {
26        self.remaining.first().copied()
27    }
28
29    #[inline]
30    fn next(&mut self) -> Option<u8> {
31        let byte = self.remaining.first().copied()?;
32        self.remaining = &self.remaining[1..];
33        Some(byte)
34    }
35
36    #[inline]
37    fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), DecodeError> {
38        if buf.len() > self.remaining.len() {
39            return Err(DecodeError::UnexpectedEOF);
40        }
41        buf.copy_from_slice(&self.remaining[..buf.len()]);
42        self.remaining = &self.remaining[buf.len()..];
43        Ok(())
44    }
45
46    #[inline]
47    fn remaining(&self) -> &'a [u8] {
48        self.remaining
49    }
50
51    #[inline]
52    fn advance(&mut self, n: usize) -> Result<(), DecodeError> {
53        if n > self.remaining.len() {
54            return Err(DecodeError::UnexpectedEOF);
55        }
56        self.remaining = &self.remaining[n..];
57        Ok(())
58    }
59}
60
61#[derive(Debug, Clone)]
62pub struct BytesMutWriter {
63    inner: BytesMut,
64}
65
66impl Default for BytesMutWriter {
67    fn default() -> Self {
68        Self::new()
69    }
70}
71
72impl BytesMutWriter {
73    #[inline]
74    pub fn new() -> Self {
75        Self {
76            inner: BytesMut::new(),
77        }
78    }
79
80    #[inline]
81    pub fn with_capacity(cap: usize) -> Self {
82        Self {
83            inner: BytesMut::with_capacity(cap),
84        }
85    }
86
87    #[inline]
88    pub fn freeze(self) -> Bytes {
89        self.inner.freeze()
90    }
91
92    #[inline]
93    pub fn as_slice(&self) -> &[u8] {
94        &self.inner
95    }
96}
97
98impl BufferWriter for BytesMutWriter {
99    #[inline]
100    fn write_all(&mut self, buf: &[u8]) -> Result<(), EncodeError> {
101        self.inner.extend_from_slice(buf);
102        Ok(())
103    }
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109    use ax_codec_core::{Decode, Encode};
110
111    #[test]
112    fn bytes_reader_roundtrip_u32() {
113        let val: u32 = 0xDEADBEEF;
114        let mut w = BytesMutWriter::new();
115        val.encode(&mut w).unwrap();
116        let bytes = w.freeze();
117
118        let mut r = BytesReader::new(&bytes);
119        let decoded = u32::decode(&mut r).unwrap();
120        assert_eq!(val, decoded);
121    }
122
123    #[test]
124    fn bytes_reader_borrow_str() {
125        use ax_codec_core::View;
126        use ax_codec_core::varint;
127
128        let s = "hello";
129        let mut w = BytesMutWriter::new();
130        varint::encode_uvarint(s.len() as u64, &mut w).unwrap();
131        w.write_all(s.as_bytes()).unwrap();
132        let bytes = w.freeze();
133
134        let mut r = BytesReader::new(&bytes);
135        let decoded: &str = View::view(&mut r).unwrap();
136        assert_eq!(decoded, "hello");
137    }
138
139    #[test]
140    fn bytes_mut_writer_grows() {
141        let mut w = BytesMutWriter::new();
142        for _ in 0..100 {
143            w.write_all(b"abcd").unwrap();
144        }
145        assert_eq!(w.as_slice().len(), 400);
146    }
147}
ax_codec_bytes/bytes_impl.rs

ax_codec_bytes/
bytes_impl.rs
