bonsaidb_core/key/
varint.rs

1use std::borrow::Cow;
2use std::fmt::Display;
3use std::ops::{Deref, DerefMut};
4
5use ordered_varint::Variable;
6use serde::{Deserialize, Serialize};
7
8use super::ByteSource;
9use crate::key::{Key, KeyEncoding, KeyKind};
10
11/// A wrapper type for Rust's built-in integer types that encodes with variable
12/// length and implements the [`Key`] trait.
13///
14/// This type supports all of Rust's built-in integer types:
15///
16/// - `u8`
17/// - `u16`
18/// - `u32`
19/// - `u64`
20/// - `u128`
21/// - `usize`
22/// - `i8`
23/// - `i16`
24/// - `i32`
25/// - `i64`
26/// - `i128`
27/// - `isize`
28///
29/// ```rust
30/// use bonsaidb_core::key::{Key, KeyEncoding, VarInt};
31///
32/// #[derive(Key, Default, Clone)]
33/// # #[key(core = bonsaidb_core)]
34/// struct UserId(u64);
35///
36/// // `UserId` type will always encode to 8 bytes, since u64 will encode
37/// // using `u64::to_be_bytes`.
38/// let default_key_len = UserId::default().as_ord_bytes().unwrap().len();
39/// assert_eq!(default_key_len, 8);
40/// let another_key_len = UserId(u64::MAX).as_ord_bytes().unwrap().len();
41/// assert_eq!(another_key_len, 8);
42///
43/// #[derive(Key, Default, Clone)]
44/// # #[key(core = bonsaidb_core)]
45/// struct UserIdVariable(VarInt<u64>);
46///
47/// // However, `UserIdVariable` will be able to encode in as little as 1 byte,
48/// // but can take up to 9 bytes if the entire u64 range is utilized.
49/// let default_key_len = UserIdVariable::default().as_ord_bytes().unwrap().len();
50/// assert_eq!(default_key_len, 1);
51/// let another_key_len = UserIdVariable(VarInt(u64::MAX))
52///     .as_ord_bytes()
53///     .unwrap()
54///     .len();
55/// assert_eq!(another_key_len, 9);
56/// ```
57///
58///
59/// # Why does this type exist?
60///
61/// The [`Key`] trait is implemented for all of Rust's native integer types by
62/// using `to_be_bytes()`/`from_be_bytes()`. This provides some benefits: very
63/// fast encoding and decoding, and known-width encoding is faster to decode.
64///
65/// This type uses [`ordered_varint`] to encode the types using a variable
66/// length encoding that is still compatible with the [`Key`] trait. This allows
67/// a value of 0 to encode as a single byte while still preserving the correct
68/// sort order required by `Key`.
69///
70/// Additionally, this encoding format allows for upgrading the in-memory size
71/// transparently if the value range needs increases over time. This only works
72/// between types that are signed the same.
73///
74/// # Behavior with Serde
75///
76/// This type implements [`serde::Serialize`] and [`serde::Deserialize`]
77/// transparently, as many serialization formats implement native variable
78/// integer encoding, and do not benefit from an ordered implementation.
79#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
80pub struct VarInt<T>(pub T)
81where
82    T: VariableInteger;
83
84impl<T> Serialize for VarInt<T>
85where
86    T: Serialize + VariableInteger,
87{
88    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
89    where
90        S: serde::Serializer,
91    {
92        self.0.serialize(serializer)
93    }
94}
95
96impl<'de, T> Deserialize<'de> for VarInt<T>
97where
98    T: Deserialize<'de> + VariableInteger,
99{
100    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
101    where
102        D: serde::Deserializer<'de>,
103    {
104        T::deserialize(deserializer).map(Self)
105    }
106}
107
108impl<T> Display for VarInt<T>
109where
110    T: Display + VariableInteger,
111{
112    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
113        self.0.fmt(f)
114    }
115}
116
117impl<T> From<T> for VarInt<T>
118where
119    T: VariableInteger,
120{
121    fn from(value: T) -> Self {
122        Self(value)
123    }
124}
125
126macro_rules! impl_varint_op {
127    ($trait:ident, $method:ident) => {
128        impl<T> std::ops::$trait<T> for VarInt<T>
129        where
130            T: std::ops::$trait<T, Output = T> + VariableInteger,
131        {
132            type Output = Self;
133
134            fn $method(self, rhs: T) -> Self::Output {
135                Self(self.0.$method(rhs))
136            }
137        }
138    };
139}
140
141impl_varint_op!(Add, add);
142impl_varint_op!(Sub, sub);
143impl_varint_op!(Mul, mul);
144impl_varint_op!(Div, div);
145impl_varint_op!(Rem, rem);
146impl_varint_op!(BitAnd, bitand);
147impl_varint_op!(BitOr, bitor);
148impl_varint_op!(BitXor, bitxor);
149impl_varint_op!(Shl, shl);
150impl_varint_op!(Shr, shr);
151
152impl<T> std::ops::Not for VarInt<T>
153where
154    T: std::ops::Not<Output = T> + VariableInteger,
155{
156    type Output = Self;
157
158    fn not(self) -> Self::Output {
159        Self(self.0.not())
160    }
161}
162
163impl<T> Deref for VarInt<T>
164where
165    T: VariableInteger,
166{
167    type Target = T;
168
169    fn deref(&self) -> &Self::Target {
170        &self.0
171    }
172}
173
174impl<T> DerefMut for VarInt<T>
175where
176    T: VariableInteger,
177{
178    fn deref_mut(&mut self) -> &mut Self::Target {
179        &mut self.0
180    }
181}
182
183impl<'k, T> Key<'k> for VarInt<T>
184where
185    T: VariableInteger,
186{
187    const CAN_OWN_BYTES: bool = false;
188
189    fn from_ord_bytes<'e>(bytes: ByteSource<'k, 'e>) -> Result<Self, Self::Error> {
190        T::decode_variable(bytes.as_ref()).map(Self)
191    }
192}
193
194impl<T> KeyEncoding for VarInt<T>
195where
196    T: VariableInteger,
197{
198    type Error = std::io::Error;
199
200    const LENGTH: Option<usize> = None;
201
202    fn describe<Visitor>(visitor: &mut Visitor)
203    where
204        Visitor: super::KeyVisitor,
205    {
206        visitor.visit_composite(
207            super::CompositeKind::Struct(Cow::Borrowed("bonsaidb::core::key::VarInt")),
208            1,
209        );
210        T::describe_contents(visitor);
211    }
212
213    fn as_ord_bytes(&self) -> Result<Cow<'_, [u8]>, Self::Error> {
214        let mut output = Vec::with_capacity(16);
215        self.0.encode_variable(&mut output)?;
216        Ok(Cow::Owned(output))
217    }
218}
219
220/// A type that is compatible with [`VarInt`].
221///
222/// This trait is implemented by all of Rust's built-in integer types.
223pub trait VariableInteger: Variable + Send + Sync + Clone + sealed::Sealed {}
224
225mod sealed {
226    pub trait Sealed {
227        fn describe_contents<Visitor>(visitor: &mut Visitor)
228        where
229            Visitor: crate::key::KeyVisitor;
230    }
231}
232
233macro_rules! impl_variable_integer {
234    ($type:ty, $kind:expr, $test:ident) => {
235        impl VariableInteger for $type {}
236
237        impl sealed::Sealed for $type {
238            fn describe_contents<Visitor>(visitor: &mut Visitor)
239            where
240                Visitor: super::KeyVisitor,
241            {
242                visitor.visit_type($kind);
243            }
244        }
245
246        impl From<VarInt<$type>> for $type {
247            fn from(value: VarInt<$type>) -> Self {
248                value.0
249            }
250        }
251
252        #[test]
253        fn $test() {
254            let i = VarInt::<$type>::from(0);
255            let r = 0;
256            let i = i + 2;
257            let r = r + 2;
258            assert_eq!(i, VarInt(r));
259            let i = i - 1;
260            let r = r - 1;
261            assert_eq!(i, VarInt(r));
262            let i = i * 6;
263            let r = r * 6;
264            assert_eq!(i, VarInt(r));
265            let i = i / 3;
266            let r = r / 3;
267            assert_eq!(i, VarInt(r));
268            let i = i % 2;
269            let r = r % 2;
270            assert_eq!(i, VarInt(r));
271            let i = !i;
272            let r = !r;
273            assert_eq!(i, VarInt(r));
274            let i = i >> 1;
275            let r = r >> 1;
276            assert_eq!(i, VarInt(r));
277            let i = i << 1;
278            let r = r << 1;
279            assert_eq!(i, VarInt(r));
280            let i = i & 0xF;
281            let r = r & 0xF;
282            assert_eq!(i, VarInt(r));
283            let i = i | 0x70;
284            let r = r | 0x70;
285            assert_eq!(i, VarInt(r));
286            let i = i ^ 0x18;
287            let r = r ^ 0x18;
288            assert_eq!(i, VarInt(r));
289            assert_eq!(Into::<$type>::into(i), r);
290
291            let encoded = i.as_ord_bytes().unwrap();
292            let decoded =
293                VarInt::<$type>::from_ord_bytes(crate::key::ByteSource::Borrowed(&encoded))
294                    .unwrap();
295            assert_eq!(i, decoded);
296
297            let pot = transmog_pot::Pot::default();
298            let pot_encoded = transmog::Format::serialize(&pot, &i).unwrap();
299            let decoded =
300                transmog::OwnedDeserializer::deserialize_owned(&pot, &pot_encoded).unwrap();
301            assert_eq!(i, decoded);
302        }
303    };
304}
305
306impl_variable_integer!(u8, KeyKind::Unsigned, varint_u8_tests);
307impl_variable_integer!(u16, KeyKind::Unsigned, varint_u16_tests);
308impl_variable_integer!(u32, KeyKind::Unsigned, varint_u32_tests);
309impl_variable_integer!(u64, KeyKind::Unsigned, varint_u64_tests);
310impl_variable_integer!(u128, KeyKind::Unsigned, varint_u126_tests);
311impl_variable_integer!(usize, KeyKind::Unsigned, varint_usize_tests);
312impl_variable_integer!(i8, KeyKind::Signed, varint_i8_tests);
313impl_variable_integer!(i16, KeyKind::Signed, varint_i16_tests);
314impl_variable_integer!(i32, KeyKind::Signed, varint_i32_tests);
315impl_variable_integer!(i64, KeyKind::Signed, varint_i64_tests);
316impl_variable_integer!(i128, KeyKind::Signed, varint_i128_tests);
317impl_variable_integer!(isize, KeyKind::Signed, varint_isize_tests);