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

spacetimedb_sats/
buffer.rs
