spacetimedb_sats/
buffer.rs

1//! Minimal utility for reading & writing the types we need to internal storage,
2//! without relying on types in third party libraries like `bytes::Bytes`, etc.
3//! Meant to be kept slim and trim for use across both native and WASM.
4
5use crate::{i256, u256};
6use core::cell::Cell;
7use core::fmt;
8use core::str::Utf8Error;
9
10/// An error that occurred when decoding.
11#[derive(Debug, Clone, PartialEq, Eq)]
12pub enum DecodeError {
13    /// Not enough data was provided in the input.
14    BufferLength {
15        for_type: &'static str,
16        expected: usize,
17        given: usize,
18    },
19    /// Length did not match the statically expected length.
20    InvalidLen { expected: usize, given: usize },
21    /// The tag does not exist for the sum.
22    InvalidTag { tag: u8, sum_name: Option<String> },
23    /// Expected data to be UTF-8 but it wasn't.
24    InvalidUtf8,
25    /// Expected the byte to be 0 or 1 to be a valid bool.
26    InvalidBool(u8),
27    /// Custom error not in the other variants of `DecodeError`.
28    Other(String),
29}
30
31impl fmt::Display for DecodeError {
32    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
33        match self {
34            DecodeError::BufferLength {
35                for_type,
36                expected,
37                given,
38            } => write!(f, "data too short for {for_type}: Expected {expected}, given {given}"),
39            DecodeError::InvalidLen { expected, given } => {
40                write!(f, "unexpected data length: Expected {expected}, given {given}")
41            }
42            DecodeError::InvalidTag { tag, sum_name } => {
43                write!(
44                    f,
45                    "unknown tag {tag:#x} for sum type {}",
46                    sum_name.as_deref().unwrap_or("<unknown>")
47                )
48            }
49            DecodeError::InvalidUtf8 => f.write_str("invalid utf8"),
50            DecodeError::InvalidBool(byte) => write!(f, "byte {byte} not valid as `bool` (must be 0 or 1)"),
51            DecodeError::Other(err) => f.write_str(err),
52        }
53    }
54}
55impl From<DecodeError> for String {
56    fn from(err: DecodeError) -> Self {
57        err.to_string()
58    }
59}
60impl std::error::Error for DecodeError {}
61
62impl From<Utf8Error> for DecodeError {
63    fn from(_: Utf8Error) -> Self {
64        DecodeError::InvalidUtf8
65    }
66}
67
68/// A buffered writer of some kind.
69pub trait BufWriter {
70    /// Writes the `slice` to the buffer.
71    ///
72    /// This is the only method implementations are required to define.
73    /// All other methods are provided.
74    fn put_slice(&mut self, slice: &[u8]);
75
76    /// Writes a `u8` to the buffer in little-endian (LE) encoding.
77    fn put_u8(&mut self, val: u8) {
78        self.put_slice(&val.to_le_bytes())
79    }
80
81    /// Writes a `u16` to the buffer in little-endian (LE) encoding.
82    fn put_u16(&mut self, val: u16) {
83        self.put_slice(&val.to_le_bytes())
84    }
85
86    /// Writes a `u32` to the buffer in little-endian (LE) encoding.
87    fn put_u32(&mut self, val: u32) {
88        self.put_slice(&val.to_le_bytes())
89    }
90
91    /// Writes a `u64` to the buffer in little-endian (LE) encoding.
92    fn put_u64(&mut self, val: u64) {
93        self.put_slice(&val.to_le_bytes())
94    }
95
96    /// Writes a `u128` to the buffer in little-endian (LE) encoding.
97    fn put_u128(&mut self, val: u128) {
98        self.put_slice(&val.to_le_bytes())
99    }
100
101    /// Writes a `u256` to the buffer in little-endian (LE) encoding.
102    fn put_u256(&mut self, val: u256) {
103        self.put_slice(&val.to_le_bytes())
104    }
105
106    /// Writes an `i8` to the buffer in little-endian (LE) encoding.
107    fn put_i8(&mut self, val: i8) {
108        self.put_slice(&val.to_le_bytes())
109    }
110
111    /// Writes an `i16` to the buffer in little-endian (LE) encoding.
112    fn put_i16(&mut self, val: i16) {
113        self.put_slice(&val.to_le_bytes())
114    }
115
116    /// Writes an `i32` to the buffer in little-endian (LE) encoding.
117    fn put_i32(&mut self, val: i32) {
118        self.put_slice(&val.to_le_bytes())
119    }
120
121    /// Writes an `i64` to the buffer in little-endian (LE) encoding.
122    fn put_i64(&mut self, val: i64) {
123        self.put_slice(&val.to_le_bytes())
124    }
125
126    /// Writes an `i128` to the buffer in little-endian (LE) encoding.
127    fn put_i128(&mut self, val: i128) {
128        self.put_slice(&val.to_le_bytes())
129    }
130
131    /// Writes an `i256` to the buffer in little-endian (LE) encoding.
132    fn put_i256(&mut self, val: i256) {
133        self.put_slice(&val.to_le_bytes())
134    }
135}
136
137macro_rules! get_int {
138    ($self:ident, $int:ident) => {
139        match $self.get_array_chunk() {
140            Some(&arr) => Ok($int::from_le_bytes(arr)),
141            None => Err(DecodeError::BufferLength {
142                for_type: stringify!($int),
143                expected: std::mem::size_of::<$int>(),
144                given: $self.remaining(),
145            }),
146        }
147    };
148}
149
150/// A buffered reader of some kind.
151///
152/// The lifetime `'de` allows the output of deserialization to borrow from the input.
153pub trait BufReader<'de> {
154    /// Reads and returns a chunk of `.len() = size` advancing the cursor iff `self.remaining() >= size`.
155    fn get_chunk(&mut self, size: usize) -> Option<&'de [u8]>;
156
157    /// Returns the number of bytes left to read in the input.
158    fn remaining(&self) -> usize;
159
160    /// Reads and returns a chunk of `.len() = N` as an array, advancing the cursor.
161    #[inline]
162    fn get_array_chunk<const N: usize>(&mut self) -> Option<&'de [u8; N]> {
163        self.get_chunk(N)?.try_into().ok()
164    }
165
166    /// Reads and returns a byte slice of `.len() = size` advancing the cursor.
167    #[inline]
168    fn get_slice(&mut self, size: usize) -> Result<&'de [u8], DecodeError> {
169        self.get_chunk(size).ok_or_else(|| DecodeError::BufferLength {
170            for_type: "[u8]",
171            expected: size,
172            given: self.remaining(),
173        })
174    }
175
176    /// Reads an array of type `[u8; N]` from the input.
177    #[inline]
178    fn get_array<const N: usize>(&mut self) -> Result<&'de [u8; N], DecodeError> {
179        self.get_array_chunk().ok_or_else(|| DecodeError::BufferLength {
180            for_type: "[u8; _]",
181            expected: N,
182            given: self.remaining(),
183        })
184    }
185
186    /// Reads a `u8` in little endian (LE) encoding from the input.
187    ///
188    /// This method is provided for convenience
189    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
190    #[inline]
191    fn get_u8(&mut self) -> Result<u8, DecodeError> {
192        get_int!(self, u8)
193    }
194
195    /// Reads a `u16` in little endian (LE) encoding from the input.
196    ///
197    /// This method is provided for convenience
198    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
199    #[inline]
200    fn get_u16(&mut self) -> Result<u16, DecodeError> {
201        get_int!(self, u16)
202    }
203
204    /// Reads a `u32` in little endian (LE) encoding from the input.
205    ///
206    /// This method is provided for convenience
207    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
208    #[inline]
209    fn get_u32(&mut self) -> Result<u32, DecodeError> {
210        get_int!(self, u32)
211    }
212
213    /// Reads a `u64` in little endian (LE) encoding from the input.
214    ///
215    /// This method is provided for convenience
216    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
217    #[inline]
218    fn get_u64(&mut self) -> Result<u64, DecodeError> {
219        get_int!(self, u64)
220    }
221
222    /// Reads a `u128` in little endian (LE) encoding from the input.
223    ///
224    /// This method is provided for convenience
225    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
226    #[inline]
227    fn get_u128(&mut self) -> Result<u128, DecodeError> {
228        get_int!(self, u128)
229    }
230
231    /// Reads a `u256` in little endian (LE) encoding from the input.
232    ///
233    /// This method is provided for convenience
234    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
235    #[inline]
236    fn get_u256(&mut self) -> Result<u256, DecodeError> {
237        get_int!(self, u256)
238    }
239
240    /// Reads an `i8` in little endian (LE) encoding from the input.
241    ///
242    /// This method is provided for convenience
243    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
244    #[inline]
245    fn get_i8(&mut self) -> Result<i8, DecodeError> {
246        get_int!(self, i8)
247    }
248
249    /// Reads an `i16` in little endian (LE) encoding from the input.
250    ///
251    /// This method is provided for convenience
252    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
253    #[inline]
254    fn get_i16(&mut self) -> Result<i16, DecodeError> {
255        get_int!(self, i16)
256    }
257
258    /// Reads an `i32` in little endian (LE) encoding from the input.
259    ///
260    /// This method is provided for convenience
261    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
262    #[inline]
263    fn get_i32(&mut self) -> Result<i32, DecodeError> {
264        get_int!(self, i32)
265    }
266
267    /// Reads an `i64` in little endian (LE) encoding from the input.
268    ///
269    /// This method is provided for convenience
270    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
271    #[inline]
272    fn get_i64(&mut self) -> Result<i64, DecodeError> {
273        get_int!(self, i64)
274    }
275
276    /// Reads an `i128` in little endian (LE) encoding from the input.
277    ///
278    /// This method is provided for convenience
279    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
280    #[inline]
281    fn get_i128(&mut self) -> Result<i128, DecodeError> {
282        get_int!(self, i128)
283    }
284
285    /// Reads an `i256` in little endian (LE) encoding from the input.
286    ///
287    /// This method is provided for convenience
288    /// and is derived from [`get_chunk`](BufReader::get_chunk)'s definition.
289    #[inline]
290    fn get_i256(&mut self) -> Result<i256, DecodeError> {
291        get_int!(self, i256)
292    }
293}
294
295impl BufWriter for Vec<u8> {
296    fn put_slice(&mut self, slice: &[u8]) {
297        self.extend_from_slice(slice);
298    }
299}
300
301impl BufWriter for &mut [u8] {
302    fn put_slice(&mut self, slice: &[u8]) {
303        if self.len() < slice.len() {
304            panic!("not enough buffer space")
305        }
306        let (buf, rest) = std::mem::take(self).split_at_mut(slice.len());
307        buf.copy_from_slice(slice);
308        *self = rest;
309    }
310}
311
312/// A [`BufWriter`] that only counts the bytes.
313#[derive(Default)]
314pub struct CountWriter {
315    /// The number of bytes counted thus far.
316    num_bytes: usize,
317}
318
319impl CountWriter {
320    /// Consumes the counter and returns the final count.
321    pub fn finish(self) -> usize {
322        self.num_bytes
323    }
324}
325
326impl BufWriter for CountWriter {
327    fn put_slice(&mut self, slice: &[u8]) {
328        self.num_bytes += slice.len();
329    }
330}
331
332/// A [`BufWriter`] that writes the bytes to two writers `W1` and `W2`.
333pub struct TeeWriter<W1, W2> {
334    pub w1: W1,
335    pub w2: W2,
336}
337
338impl<W1: BufWriter, W2: BufWriter> TeeWriter<W1, W2> {
339    pub fn new(w1: W1, w2: W2) -> Self {
340        Self { w1, w2 }
341    }
342}
343
344impl<W1: BufWriter, W2: BufWriter> BufWriter for TeeWriter<W1, W2> {
345    fn put_slice(&mut self, slice: &[u8]) {
346        self.w1.put_slice(slice);
347        self.w2.put_slice(slice);
348    }
349}
350
351impl<'de> BufReader<'de> for &'de [u8] {
352    #[inline]
353    fn get_chunk(&mut self, size: usize) -> Option<&'de [u8]> {
354        let (ret, rest) = self.split_at_checked(size)?;
355        *self = rest;
356        Some(ret)
357    }
358
359    #[inline]
360    fn get_array_chunk<const N: usize>(&mut self) -> Option<&'de [u8; N]> {
361        let (ret, rest) = self.split_first_chunk()?;
362        *self = rest;
363        Some(ret)
364    }
365
366    #[inline(always)]
367    fn remaining(&self) -> usize {
368        self.len()
369    }
370}
371
372/// A cursor based [`BufReader<'de>`] implementation.
373#[derive(Debug)]
374pub struct Cursor<I> {
375    /// The underlying input read from.
376    pub buf: I,
377    /// The position within the reader.
378    pub pos: Cell<usize>,
379}
380
381impl<I> Cursor<I> {
382    /// Returns a new cursor on the provided `buf` input.
383    ///
384    /// The cursor starts at the beginning of `buf`.
385    pub fn new(buf: I) -> Self {
386        Self { buf, pos: 0.into() }
387    }
388}
389
390impl<'de, I: AsRef<[u8]>> BufReader<'de> for &'de Cursor<I> {
391    #[inline]
392    fn get_chunk(&mut self, size: usize) -> Option<&'de [u8]> {
393        // "Read" the slice `buf[pos..size]`.
394        let buf = &self.buf.as_ref()[self.pos.get()..];
395        let ret = buf.get(..size)?;
396
397        // Advance the cursor by `size` bytes.
398        self.pos.set(self.pos.get() + size);
399
400        Some(ret)
401    }
402
403    #[inline]
404    fn get_array_chunk<const N: usize>(&mut self) -> Option<&'de [u8; N]> {
405        // "Read" the slice `buf[pos..size]`.
406        let buf = &self.buf.as_ref()[self.pos.get()..];
407        let ret = buf.first_chunk()?;
408
409        // Advance the cursor by `size` bytes.
410        self.pos.set(self.pos.get() + N);
411
412        Some(ret)
413    }
414
415    fn remaining(&self) -> usize {
416        self.buf.as_ref().len() - self.pos.get()
417    }
418}
419
420#[cfg(test)]
421mod tests {
422    use crate::buffer::{BufReader, BufWriter};
423
424    #[test]
425    fn test_simple_encode_decode() {
426        let mut writer: Vec<u8> = vec![];
427        writer.put_u8(5);
428        writer.put_u32(6);
429        writer.put_u64(7);
430
431        let arr_val = [1, 2, 3, 4];
432        writer.put_slice(&arr_val[..]);
433
434        let mut reader = writer.as_slice();
435        assert_eq!(reader.get_u8().unwrap(), 5);
436        assert_eq!(reader.get_u32().unwrap(), 6);
437        assert_eq!(reader.get_u64().unwrap(), 7);
438
439        let slice = reader.get_slice(4).unwrap();
440        assert_eq!(slice, arr_val);
441
442        // reading beyond buffer end should fail
443        assert!(reader.get_slice(1).is_err());
444        assert!(reader.get_slice(123).is_err());
445        assert!(reader.get_u64().is_err());
446        assert!(reader.get_u32().is_err());
447        assert!(reader.get_u8().is_err());
448    }
449}
spacetimedb_sats/buffer.rs

spacetimedb_sats/
buffer.rs
