turbo_persistence/
key.rs

1use std::{cmp::min, hash::Hasher};
2
3/// A trait for keys that can be used for hashing.
4pub trait KeyBase {
5    /// Returns the length of the key in bytes.
6    fn len(&self) -> usize;
7    fn is_empty(&self) -> bool {
8        self.len() == 0
9    }
10    /// Hashes the key. It should not include the structure of the key, only the data. E.g. `([1,
11    /// 2], [3, 4])` should hash the same as `[1, 2, 3, 4]`.
12    fn hash<H: Hasher>(&self, state: &mut H);
13}
14
15impl KeyBase for &'_ [u8] {
16    fn len(&self) -> usize {
17        <[u8]>::len(self)
18    }
19
20    fn is_empty(&self) -> bool {
21        <[u8]>::is_empty(self)
22    }
23
24    fn hash<H: Hasher>(&self, state: &mut H) {
25        state.write(self);
26    }
27}
28
29impl<const N: usize> KeyBase for [u8; N] {
30    fn len(&self) -> usize {
31        N
32    }
33
34    fn is_empty(&self) -> bool {
35        N > 0
36    }
37
38    fn hash<H: Hasher>(&self, state: &mut H) {
39        state.write(self);
40    }
41}
42
43impl KeyBase for Vec<u8> {
44    fn len(&self) -> usize {
45        self.len()
46    }
47
48    fn is_empty(&self) -> bool {
49        self.is_empty()
50    }
51
52    fn hash<H: Hasher>(&self, state: &mut H) {
53        state.write(self);
54    }
55}
56
57impl KeyBase for Box<[u8]> {
58    fn len(&self) -> usize {
59        (**self).len()
60    }
61
62    fn is_empty(&self) -> bool {
63        (**self).is_empty()
64    }
65
66    fn hash<H: Hasher>(&self, state: &mut H) {
67        state.write(self);
68    }
69}
70
71impl KeyBase for u8 {
72    fn len(&self) -> usize {
73        1
74    }
75
76    fn is_empty(&self) -> bool {
77        false
78    }
79
80    fn hash<H: Hasher>(&self, state: &mut H) {
81        state.write_u8(*self);
82    }
83}
84
85impl<A: KeyBase, B: KeyBase> KeyBase for (A, B) {
86    fn len(&self) -> usize {
87        let (a, b) = self;
88        a.len() + b.len()
89    }
90
91    fn is_empty(&self) -> bool {
92        let (a, b) = self;
93        a.is_empty() && b.is_empty()
94    }
95
96    fn hash<H: Hasher>(&self, state: &mut H) {
97        let (a, b) = self;
98        KeyBase::hash(a, state);
99        KeyBase::hash(b, state);
100    }
101}
102
103impl<T: KeyBase> KeyBase for &'_ T {
104    fn len(&self) -> usize {
105        (*self).len()
106    }
107
108    fn is_empty(&self) -> bool {
109        (*self).is_empty()
110    }
111
112    fn hash<H: Hasher>(&self, state: &mut H) {
113        (*self).hash(state)
114    }
115}
116
117/// A trait for keys that can be used to query the database. They need to allow hashing and
118/// comparison with a byte slice (total order).
119pub trait QueryKey: KeyBase {
120    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering;
121    fn eq(&self, key: &[u8]) -> bool;
122}
123
124impl QueryKey for &'_ [u8] {
125    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
126        Ord::cmp(self, &key)
127    }
128
129    fn eq(&self, key: &[u8]) -> bool {
130        PartialEq::eq(*self, key)
131    }
132}
133
134impl<const N: usize> QueryKey for [u8; N] {
135    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
136        Ord::cmp(&self[..], key)
137    }
138    fn eq(&self, key: &[u8]) -> bool {
139        PartialEq::eq(self, key)
140    }
141}
142
143impl QueryKey for Vec<u8> {
144    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
145        Ord::cmp(&**self, key)
146    }
147    fn eq(&self, key: &[u8]) -> bool {
148        PartialEq::eq(self.as_slice(), key)
149    }
150}
151
152impl QueryKey for Box<[u8]> {
153    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
154        Ord::cmp(&**self, key)
155    }
156    fn eq(&self, key: &[u8]) -> bool {
157        PartialEq::eq(&**self, key)
158    }
159}
160
161impl QueryKey for u8 {
162    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
163        Ord::cmp(&[*self][..], key)
164    }
165    fn eq(&self, key: &[u8]) -> bool {
166        PartialEq::eq(&[*self][..], key)
167    }
168}
169
170impl<A: QueryKey, B: QueryKey> QueryKey for (A, B) {
171    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
172        let (a, b) = self;
173        let len = a.len();
174        let key_len = key.len();
175        let (key_part, value_part) = key.split_at(min(key_len, len));
176        a.cmp(key_part).then_with(|| b.cmp(value_part))
177    }
178    fn eq(&self, key: &[u8]) -> bool {
179        let (a, b) = self;
180        let len = a.len();
181        let key_len = key.len();
182        let (key_part, value_part) = &key.split_at(min(key_len, len));
183        a.eq(key_part) && b.eq(value_part)
184    }
185}
186
187impl<T: QueryKey> QueryKey for &'_ T {
188    fn cmp(&self, key: &[u8]) -> std::cmp::Ordering {
189        (*self).cmp(key)
190    }
191    fn eq(&self, key: &[u8]) -> bool {
192        (*self).eq(key)
193    }
194}
195
196/// A trait for keys that can be stored in the database. They need to allow hashing and comparison.
197pub trait StoreKey: KeyBase + Ord {
198    /// The key's bytes.
199    fn as_slice(&self) -> &[u8];
200
201    fn write_to(&self, buf: &mut Vec<u8>) {
202        buf.extend_from_slice(self.as_slice());
203    }
204}
205
206impl<const N: usize> StoreKey for [u8; N] {
207    fn as_slice(&self) -> &[u8] {
208        &self[..]
209    }
210}
211
212impl StoreKey for Vec<u8> {
213    fn as_slice(&self) -> &[u8] {
214        self
215    }
216}
217
218impl StoreKey for Box<[u8]> {
219    fn as_slice(&self) -> &[u8] {
220        self
221    }
222}
223
224impl StoreKey for &'_ [u8] {
225    fn as_slice(&self) -> &[u8] {
226        self
227    }
228}
229
230impl<T: StoreKey> StoreKey for &'_ T {
231    fn as_slice(&self) -> &[u8] {
232        (*self).as_slice()
233    }
234}
235
236/// Hashes a key with a fast, deterministic hash function.
237pub fn hash_key(key: &impl KeyBase) -> u64 {
238    let mut hasher = xxhash_rust::xxh3::Xxh3Default::new();
239    key.hash(&mut hasher);
240    hasher.finish()
241}
242
243#[cfg(test)]
244mod tests {
245    use std::cmp::Ordering;
246
247    use crate::{QueryKey, key::hash_key};
248
249    #[test]
250    fn tuple() {
251        let key = (&[1, 2], &[3, 4]);
252        assert_eq!(QueryKey::cmp(&key, &[1, 2, 3, 4]), Ordering::Equal);
253        assert_eq!(QueryKey::cmp(&key, &[1, 2, 3, 3]), Ordering::Greater);
254        assert_eq!(QueryKey::cmp(&key, &[1, 2, 3, 5]), Ordering::Less);
255        assert_eq!(QueryKey::cmp(&key, &[0, 2, 3, 4]), Ordering::Greater);
256        assert_eq!(QueryKey::cmp(&key, &[2, 2, 3, 4]), Ordering::Less);
257        assert_eq!(QueryKey::cmp(&key, &[1, 2, 3, 4, 5]), Ordering::Less);
258        assert_eq!(QueryKey::cmp(&key, &[1, 2, 3]), Ordering::Greater);
259        assert_eq!(QueryKey::cmp(&key, &[1, 2]), Ordering::Greater);
260        assert_eq!(QueryKey::cmp(&key, &[1]), Ordering::Greater);
261        assert_eq!(QueryKey::cmp(&key, &[]), Ordering::Greater);
262    }
263
264    #[test]
265    fn hash() {
266        let h1 = hash_key(&[1, 2, 3, 4]);
267        let h2 = hash_key(&(&[1, 2], &[3, 4]));
268        let h3 = hash_key(&(vec![1, 2, 3], 4u8));
269        assert_eq!(h2, h1);
270        assert_eq!(h3, h1);
271    }
272}
turbo_persistence/key.rs

turbo_persistence/
key.rs
