1
use std::borrow::Cow;
2
use std::fmt::Display;
3
use std::ops::{Deref, DerefMut};
4

            
5
use ordered_varint::Variable;
6
use serde::{Deserialize, Serialize};
7

            
8
use super::ByteSource;
9
use 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
156
#[derive(Debug, Default, Clone, Copy, Eq, PartialEq, Ord, PartialOrd, Hash)]
80
pub struct VarInt<T>(pub T)
81
where
82
    T: VariableInteger;
83

            
84
impl<T> Serialize for VarInt<T>
85
where
86
    T: Serialize + VariableInteger,
87
{
88
12
    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
89
12
    where
90
12
        S: serde::Serializer,
91
12
    {
92
12
        self.0.serialize(serializer)
93
12
    }
94
}
95

            
96
impl<'de, T> Deserialize<'de> for VarInt<T>
97
where
98
    T: Deserialize<'de> + VariableInteger,
99
{
100
12
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
101
12
    where
102
12
        D: serde::Deserializer<'de>,
103
12
    {
104
12
        T::deserialize(deserializer).map(Self)
105
12
    }
106
}
107

            
108
impl<T> Display for VarInt<T>
109
where
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

            
117
impl<T> From<T> for VarInt<T>
118
where
119
    T: VariableInteger,
120
{
121
12
    fn from(value: T) -> Self {
122
12
        Self(value)
123
12
    }
124
}
125

            
126
macro_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
120
            fn $method(self, rhs: T) -> Self::Output {
135
120
                Self(self.0.$method(rhs))
136
120
            }
137
        }
138
    };
139
}
140

            
141
impl_varint_op!(Add, add);
142
impl_varint_op!(Sub, sub);
143
impl_varint_op!(Mul, mul);
144
impl_varint_op!(Div, div);
145
impl_varint_op!(Rem, rem);
146
impl_varint_op!(BitAnd, bitand);
147
impl_varint_op!(BitOr, bitor);
148
impl_varint_op!(BitXor, bitxor);
149
impl_varint_op!(Shl, shl);
150
impl_varint_op!(Shr, shr);
151

            
152
impl<T> std::ops::Not for VarInt<T>
153
where
154
    T: std::ops::Not<Output = T> + VariableInteger,
155
{
156
    type Output = Self;
157

            
158
12
    fn not(self) -> Self::Output {
159
12
        Self(self.0.not())
160
12
    }
161
}
162

            
163
impl<T> Deref for VarInt<T>
164
where
165
    T: VariableInteger,
166
{
167
    type Target = T;
168

            
169
    fn deref(&self) -> &Self::Target {
170
        &self.0
171
    }
172
}
173

            
174
impl<T> DerefMut for VarInt<T>
175
where
176
    T: VariableInteger,
177
{
178
    fn deref_mut(&mut self) -> &mut Self::Target {
179
        &mut self.0
180
    }
181
}
182

            
183
impl<'k, T> Key<'k> for VarInt<T>
184
where
185
    T: VariableInteger,
186
{
187
    const CAN_OWN_BYTES: bool = false;
188

            
189
12
    fn from_ord_bytes<'e>(bytes: ByteSource<'k, 'e>) -> Result<Self, Self::Error> {
190
12
        T::decode_variable(bytes.as_ref()).map(Self)
191
12
    }
192
}
193

            
194
impl<T> KeyEncoding for VarInt<T>
195
where
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
12
    fn as_ord_bytes(&self) -> Result<Cow<'_, [u8]>, Self::Error> {
214
12
        let mut output = Vec::with_capacity(16);
215
12
        self.0.encode_variable(&mut output)?;
216
12
        Ok(Cow::Owned(output))
217
12
    }
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.
223
pub trait VariableInteger: Variable + Send + Sync + Clone + sealed::Sealed {}
224

            
225
mod sealed {
226
    pub trait Sealed {
227
        fn describe_contents<Visitor>(visitor: &mut Visitor)
228
        where
229
            Visitor: crate::key::KeyVisitor;
230
    }
231
}
232

            
233
macro_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
12
            fn from(value: VarInt<$type>) -> Self {
248
12
                value.0
249
12
            }
250
        }
251

            
252
        #[test]
253
12
        fn $test() {
254
12
            let i = VarInt::<$type>::from(0);
255
12
            let r = 0;
256
12
            let i = i + 2;
257
12
            let r = r + 2;
258
12
            assert_eq!(i, VarInt(r));
259
12
            let i = i - 1;
260
12
            let r = r - 1;
261
12
            assert_eq!(i, VarInt(r));
262
12
            let i = i * 6;
263
12
            let r = r * 6;
264
12
            assert_eq!(i, VarInt(r));
265
12
            let i = i / 3;
266
12
            let r = r / 3;
267
12
            assert_eq!(i, VarInt(r));
268
12
            let i = i % 2;
269
12
            let r = r % 2;
270
12
            assert_eq!(i, VarInt(r));
271
12
            let i = !i;
272
12
            let r = !r;
273
12
            assert_eq!(i, VarInt(r));
274
12
            let i = i >> 1;
275
12
            let r = r >> 1;
276
12
            assert_eq!(i, VarInt(r));
277
12
            let i = i << 1;
278
12
            let r = r << 1;
279
12
            assert_eq!(i, VarInt(r));
280
12
            let i = i & 0xF;
281
12
            let r = r & 0xF;
282
12
            assert_eq!(i, VarInt(r));
283
12
            let i = i | 0x70;
284
12
            let r = r | 0x70;
285
12
            assert_eq!(i, VarInt(r));
286
12
            let i = i ^ 0x18;
287
12
            let r = r ^ 0x18;
288
12
            assert_eq!(i, VarInt(r));
289
12
            assert_eq!(Into::<$type>::into(i), r);
290

            
291
12
            let encoded = i.as_ord_bytes().unwrap();
292
12
            let decoded =
293
12
                VarInt::<$type>::from_ord_bytes(crate::key::ByteSource::Borrowed(&encoded))
294
12
                    .unwrap();
295
12
            assert_eq!(i, decoded);
296

            
297
12
            let pot = transmog_pot::Pot::default();
298
12
            let pot_encoded = transmog::Format::serialize(&pot, &i).unwrap();
299
12
            let decoded =
300
12
                transmog::OwnedDeserializer::deserialize_owned(&pot, &pot_encoded).unwrap();
301
12
            assert_eq!(i, decoded);
302
12
        }
303
    };
304
}
305

            
306
1
impl_variable_integer!(u8, KeyKind::Unsigned, varint_u8_tests);
307
1
impl_variable_integer!(u16, KeyKind::Unsigned, varint_u16_tests);
308
1
impl_variable_integer!(u32, KeyKind::Unsigned, varint_u32_tests);
309
1
impl_variable_integer!(u64, KeyKind::Unsigned, varint_u64_tests);
310
1
impl_variable_integer!(u128, KeyKind::Unsigned, varint_u126_tests);
311
1
impl_variable_integer!(usize, KeyKind::Unsigned, varint_usize_tests);
312
1
impl_variable_integer!(i8, KeyKind::Signed, varint_i8_tests);
313
1
impl_variable_integer!(i16, KeyKind::Signed, varint_i16_tests);
314
1
impl_variable_integer!(i32, KeyKind::Signed, varint_i32_tests);
315
1
impl_variable_integer!(i64, KeyKind::Signed, varint_i64_tests);
316
1
impl_variable_integer!(i128, KeyKind::Signed, varint_i128_tests);
317
1
impl_variable_integer!(isize, KeyKind::Signed, varint_isize_tests);