Skip to main content

fluentbase_types/
bincode.rs

1use crate::Bytes;
2pub use ::bincode::{
3    config::{Config, Configuration, Fixint, LittleEndian},
4    de::{read::Reader, Decoder},
5    error::DecodeError,
6    *,
7};
8use core::ops::{Deref, DerefMut};
9
10#[derive(Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
11pub struct ZeroCopyBytes(pub Bytes);
12
13impl ZeroCopyBytes {
14    pub fn new() -> Self {
15        Self(Bytes::new())
16    }
17}
18
19impl From<Bytes> for ZeroCopyBytes {
20    fn from(value: Bytes) -> Self {
21        ZeroCopyBytes(value)
22    }
23}
24impl From<ZeroCopyBytes> for Bytes {
25    fn from(value: ZeroCopyBytes) -> Self {
26        value.0
27    }
28}
29
30impl Deref for ZeroCopyBytes {
31    type Target = Bytes;
32    #[inline]
33    fn deref(&self) -> &Self::Target {
34        &self.0
35    }
36}
37
38impl DerefMut for ZeroCopyBytes {
39    #[inline]
40    fn deref_mut(&mut self) -> &mut Self::Target {
41        &mut self.0
42    }
43}
44
45impl AsRef<[u8]> for ZeroCopyBytes {
46    #[inline]
47    fn as_ref(&self) -> &[u8] {
48        self.0.as_ref()
49    }
50}
51
52pub struct BytesReader {
53    bytes: Bytes,
54}
55
56impl BytesReader {
57    pub fn new(bytes: Bytes) -> Self {
58        Self { bytes }
59    }
60}
61
62/// Extension methods for a Reader that is backed by `Bytes`.
63pub trait BytesReaderExt: Reader {
64    /// Take the next `n` bytes as a zero-copy `Bytes` slice and advance the reader.
65    fn take_bytes(&mut self, n: usize) -> Result<Bytes, DecodeError>;
66}
67
68impl BytesReaderExt for BytesReader {
69    #[inline]
70    fn take_bytes(&mut self, n: usize) -> Result<Bytes, DecodeError> {
71        if n > self.bytes.len() {
72            return Err(DecodeError::UnexpectedEnd {
73                additional: n - self.bytes.len(),
74            });
75        }
76        let out = self.bytes.slice(..n);
77        self.bytes = self.bytes.slice(n..);
78        Ok(out)
79    }
80}
81
82impl Reader for BytesReader {
83    #[inline(always)]
84    fn read(&mut self, bytes: &mut [u8]) -> Result<(), DecodeError> {
85        if bytes.len() > self.bytes.len() {
86            return Err(DecodeError::UnexpectedEnd {
87                additional: bytes.len() - self.bytes.len(),
88            });
89        }
90        let (read_slice, _remaining) = self.bytes.split_at(bytes.len());
91        bytes.copy_from_slice(read_slice);
92        self.bytes = self.bytes.slice(bytes.len()..);
93        Ok(())
94    }
95
96    #[inline]
97    fn peek_read(&mut self, n: usize) -> Option<&[u8]> {
98        self.bytes.get(..n)
99    }
100
101    #[inline]
102    fn consume(&mut self, n: usize) {
103        self.bytes = self.bytes.slice(n..);
104    }
105}
106
107pub trait DecodeBytes<Context>: Sized {
108    /// Attempt to decode this type with the given [bincode::Decode].
109    fn decode_bytes<D: Decoder<Context = Context, R = BytesReader>>(
110        decoder: &mut D,
111    ) -> Result<Self, DecodeError>;
112}
113
114impl<Context> DecodeBytes<Context> for ZeroCopyBytes {
115    fn decode_bytes<D>(decoder: &mut D) -> Result<Self, DecodeError>
116    where
117        D: Decoder<Context = Context, R = BytesReader>,
118    {
119        let len_u64 = u64::decode(decoder)?;
120        let len: usize = len_u64
121            .try_into()
122            .map_err(|_| DecodeError::OutsideUsizeRange(len_u64))?;
123
124        decoder.claim_container_read::<u8>(len)?;
125
126        // Now we *know* the reader is BytesReader, so we can zero-copy:
127        let bytes = decoder.reader().take_bytes(len)?;
128        Ok(ZeroCopyBytes(bytes))
129    }
130}
131impl<Context> DecodeBytes<Context> for Option<ZeroCopyBytes> {
132    fn decode_bytes<D>(decoder: &mut D) -> Result<Self, DecodeError>
133    where
134        D: Decoder<Context = Context, R = BytesReader>,
135    {
136        let variant = match u8::decode(decoder)? {
137            0 => Ok(None),
138            1 => Ok(Some(ZeroCopyBytes::new())),
139            x => Err(DecodeError::UnexpectedVariant {
140                found: x as u32,
141                allowed: &error::AllowedEnumVariants::Range { max: 1, min: 0 },
142                type_name: core::any::type_name::<Option<ZeroCopyBytes>>(),
143            }),
144        }?;
145        match variant {
146            Some(_) => {
147                let val = ZeroCopyBytes::decode_bytes(decoder)?;
148                Ok(Some(val))
149            }
150            None => Ok(None),
151        }
152    }
153}
154
155pub fn decode_from_bytes<D: DecodeBytes<()>, C: Config>(
156    src: Bytes,
157    config: C,
158) -> Result<(D, usize), DecodeError> {
159    let original_len_bytes = src.len();
160    let reader = BytesReader::new(src);
161    let mut decoder = de::DecoderImpl::<_, C, ()>::new(reader, config, ());
162    let result = D::decode_bytes(&mut decoder)?;
163    let bytes_read = original_len_bytes - decoder.reader().bytes.len();
164    Ok((result, bytes_read))
165}