spacetimedb_sats/
bsatn.rs

1use core::mem::MaybeUninit;
2
3use crate::buffer::{BufReader, BufWriter, CountWriter};
4use crate::de::{BasicSmallVecVisitor, Deserialize, DeserializeSeed, Deserializer as _};
5use crate::ser::Serialize;
6use crate::{ProductValue, Typespace, WithTypespace};
7use bytes::BytesMut;
8use ser::BsatnError;
9use smallvec::SmallVec;
10
11pub mod de;
12pub mod eq;
13pub mod ser;
14
15pub use de::Deserializer;
16pub use ser::Serializer;
17
18pub use crate::buffer::DecodeError;
19pub use ser::BsatnError as EncodeError;
20
21/// Serialize `value` into the buffered writer `w` in the BSATN format.
22#[inline]
23pub fn to_writer<W: BufWriter, T: Serialize + ?Sized>(w: &mut W, value: &T) -> Result<(), EncodeError> {
24    value.serialize_into_bsatn(Serializer::new(w))
25}
26
27/// Serialize `value` into a `Vec<u8>` in the BSATN format.
28pub fn to_vec<T: Serialize + ?Sized>(value: &T) -> Result<Vec<u8>, EncodeError> {
29    let mut v = Vec::new();
30    to_writer(&mut v, value)?;
31    Ok(v)
32}
33
34/// Computes the size of `val` when BSATN encoding without actually encoding.
35pub fn to_len<T: Serialize + ?Sized>(value: &T) -> Result<usize, EncodeError> {
36    let mut writer = CountWriter::default();
37    to_writer(&mut writer, value)?;
38    Ok(writer.finish())
39}
40
41/// Deserialize a `T` from the BSATN format in the buffered `reader`.
42pub fn from_reader<'de, T: Deserialize<'de>>(reader: &mut impl BufReader<'de>) -> Result<T, DecodeError> {
43    T::deserialize(Deserializer::new(reader))
44}
45
46/// Deserialize a `T` from the BSATN format in `bytes`.
47pub fn from_slice<'de, T: Deserialize<'de>>(bytes: &'de [u8]) -> Result<T, DecodeError> {
48    from_reader(&mut &*bytes)
49}
50
51/// Decode `bytes` to the value type of `ty: S`.
52pub fn decode<'a, 'de, S: ?Sized>(
53    ty: &'a S,
54    bytes: &mut impl BufReader<'de>,
55) -> Result<<WithTypespace<'a, S> as DeserializeSeed<'de>>::Output, DecodeError>
56where
57    WithTypespace<'a, S>: DeserializeSeed<'de>,
58{
59    crate::WithTypespace::empty(ty).deserialize(Deserializer::new(bytes))
60}
61
62macro_rules! codec_funcs {
63    ($ty:ty) => {
64        impl $ty {
65            pub fn decode<'a>(bytes: &mut impl BufReader<'a>) -> Result<Self, DecodeError> {
66                from_reader(bytes)
67            }
68
69            pub fn encode(&self, bytes: &mut impl BufWriter) {
70                to_writer(bytes, self).unwrap()
71            }
72        }
73    };
74    (val: $ty:ty) => {
75        impl $ty {
76            /// Decode a value from `bytes` typed at `ty`.
77            pub fn decode<'a>(
78                ty: &<Self as crate::Value>::Type,
79                bytes: &mut impl BufReader<'a>,
80            ) -> Result<Self, DecodeError> {
81                decode(ty, bytes)
82            }
83
84            /// Decode a vector of values from `bytes` with each value typed at `ty`.
85            pub fn decode_smallvec<'a>(
86                ty: &<Self as crate::Value>::Type,
87                bytes: &mut impl BufReader<'a>,
88            ) -> Result<SmallVec<[Self; 1]>, DecodeError> {
89                Deserializer::new(bytes).deserialize_array_seed(
90                    BasicSmallVecVisitor,
91                    crate::WithTypespace::new(&Typespace::new(Vec::new()), ty),
92                )
93            }
94
95            pub fn encode(&self, bytes: &mut impl BufWriter) {
96                to_writer(bytes, self).unwrap()
97            }
98        }
99    };
100}
101
102codec_funcs!(crate::AlgebraicType);
103codec_funcs!(crate::ProductType);
104codec_funcs!(crate::SumType);
105codec_funcs!(crate::ProductTypeElement);
106codec_funcs!(crate::SumTypeVariant);
107
108codec_funcs!(val: crate::AlgebraicValue);
109codec_funcs!(val: crate::ProductValue);
110codec_funcs!(val: crate::SumValue);
111
112/// Provides a view over a buffer that can reserve an additional `len` bytes
113/// and then provide those as an uninitialized buffer to write into.
114pub trait BufReservedFill {
115    /// Reserves space for `len` in `self` and then runs `fill` to fill it,
116    /// adding `len` to the total length of `self`.
117    ///
118    /// # Safety
119    ///
120    /// `fill` must initialize every byte in the slice.
121    unsafe fn reserve_and_fill(&mut self, len: usize, fill: impl FnOnce(&mut [MaybeUninit<u8>]));
122}
123
124impl BufReservedFill for Vec<u8> {
125    unsafe fn reserve_and_fill(&mut self, len: usize, fill: impl FnOnce(&mut [MaybeUninit<u8>])) {
126        // Get an uninitialized slice within `self` of `len` bytes.
127        let start = self.len();
128        self.reserve(len);
129        let sink = &mut self.spare_capacity_mut()[..len];
130
131        // Run the filling logic.
132        fill(sink);
133
134        // SAFETY: Caller promised that `sink` was fully initialized,
135        // which entails that we initialized `start .. start + len`,
136        // so now we have initialized up to `start + len`.
137        unsafe { self.set_len(start + len) }
138    }
139}
140
141impl BufReservedFill for BytesMut {
142    unsafe fn reserve_and_fill(&mut self, len: usize, fill: impl FnOnce(&mut [MaybeUninit<u8>])) {
143        // Get an uninitialized slice within `self` of `len` bytes.
144        let start = self.len();
145        self.reserve(len);
146        let sink = &mut self.spare_capacity_mut()[..len];
147
148        // Run the filling logic.
149        fill(sink);
150
151        // SAFETY: Caller promised that `sink` was fully initialized,
152        // which entails that we initialized `start .. start + len`,
153        // so now we have initialized up to `start + len`.
154        unsafe { self.set_len(start + len) }
155    }
156}
157
158/// Types that can be encoded to BSATN.
159///
160/// Implementations of this trait may be more efficient than directly calling [`to_vec`].
161/// In particular, for `spacetimedb_table::table::RowRef`, this method will use a `StaticLayout` if one is available,
162/// avoiding expensive runtime type dispatch.
163pub trait ToBsatn {
164    /// BSATN-encode the row referred to by `self` into a freshly-allocated `Vec<u8>`.
165    fn to_bsatn_vec(&self) -> Result<Vec<u8>, BsatnError>;
166
167    /// BSATN-encode the row referred to by `self` into `buf`,
168    /// pushing `self`'s bytes onto the end of `buf`, similar to [`Vec::extend`].
169    fn to_bsatn_extend(&self, buf: &mut (impl BufWriter + BufReservedFill)) -> Result<(), BsatnError>;
170
171    /// Returns the static size of the type of this object.
172    ///
173    /// When this returns `Some(_)` there is also a `StaticLayout`.
174    fn static_bsatn_size(&self) -> Option<u16>;
175}
176
177impl<T: ToBsatn> ToBsatn for &T {
178    fn to_bsatn_vec(&self) -> Result<Vec<u8>, BsatnError> {
179        T::to_bsatn_vec(*self)
180    }
181    fn to_bsatn_extend(&self, buf: &mut (impl BufWriter + BufReservedFill)) -> Result<(), BsatnError> {
182        T::to_bsatn_extend(*self, buf)
183    }
184    fn static_bsatn_size(&self) -> Option<u16> {
185        T::static_bsatn_size(*self)
186    }
187}
188
189impl ToBsatn for ProductValue {
190    fn to_bsatn_vec(&self) -> Result<Vec<u8>, BsatnError> {
191        to_vec(self)
192    }
193
194    fn to_bsatn_extend(&self, buf: &mut (impl BufWriter + BufReservedFill)) -> Result<(), BsatnError> {
195        to_writer(buf, self)
196    }
197
198    fn static_bsatn_size(&self) -> Option<u16> {
199        None
200    }
201}
202
203mod private_is_primitive_type {
204    pub trait Sealed {}
205}
206/// A primitive type.
207/// This is purely intended for use in `crates/bindings-macro`.
208///
209/// # Safety
210///
211/// Implementing this guarantees that the type has no padding, recursively.
212#[doc(hidden)]
213pub unsafe trait IsPrimitiveType: private_is_primitive_type::Sealed {}
214macro_rules! is_primitive_type {
215    ($($prim:ty),*) => {
216        $(
217            impl private_is_primitive_type::Sealed for $prim {}
218            // SAFETY:  the type is primitive and has no padding.
219            unsafe impl IsPrimitiveType for $prim {}
220        )*
221    };
222}
223is_primitive_type!(u8, i8, u16, i16, u32, i32, u64, i64, u128, i128, f32, f64);
224
225/// Enforces that a type is a primitive.
226/// This is purely intended for use in `crates/bindings-macro`.
227pub const fn assert_is_primitive_type<T: IsPrimitiveType>() {}
228
229#[cfg(test)]
230mod tests {
231    use super::{to_vec, DecodeError};
232    use crate::proptest::generate_typed_value;
233    use crate::{meta_type::MetaType, AlgebraicType, AlgebraicValue};
234    use proptest::prelude::*;
235    use proptest::proptest;
236
237    #[test]
238    fn type_to_binary_equivalent() {
239        check_type(&AlgebraicType::meta_type());
240    }
241
242    #[track_caller]
243    fn check_type(ty: &AlgebraicType) {
244        let mut through_value = Vec::new();
245        ty.as_value().encode(&mut through_value);
246        let mut direct = Vec::new();
247        ty.encode(&mut direct);
248        assert_eq!(direct, through_value);
249    }
250
251    proptest! {
252        #[test]
253        fn bsatn_enc_de_roundtrips((ty, val) in generate_typed_value()) {
254            let bytes = to_vec(&val).unwrap();
255            let val_decoded = AlgebraicValue::decode(&ty, &mut &bytes[..]).unwrap();
256            prop_assert_eq!(val, val_decoded);
257        }
258
259        #[test]
260        fn bsatn_non_zero_one_u8_aint_bool(val in 2u8..) {
261            let bytes = [val];
262            prop_assert_eq!(
263                AlgebraicValue::decode(&AlgebraicType::Bool, &mut &bytes[..]),
264                Err(DecodeError::InvalidBool(val))
265            );
266        }
267    }
268}