Skip to main content

radix_trie/
keys.rs

1use endian_type::{BigEndian, LittleEndian};
2use nibble_vec::Nibblet;
3use std::path::{Path, PathBuf};
4
5/// Trait for types which can be used to key a Radix Trie.
6///
7/// Types that implement this trait should be convertible to a vector of half-bytes (nibbles)
8/// such that no two instances of the type convert to the same vector.
9/// To protect against faulty behaviour, the trie will **panic** if it finds two distinct keys
10/// of type `K` which encode to the same `Nibblet`, so be careful!
11///
12/// If you would like to implement this trait for your own type, you need to implement
13/// *either* `encode_bytes` or `encode`. You only need to implement one of the two.
14/// If you don't implement one, your code will **panic** as soon you use the trie.
15/// There is no performance penalty for implementing `encode_bytes` instead of `encode`,
16/// so it is preferred except in the case where you require half-byte precision.
17///
18/// Many standard types implement this trait already. Integer types are encoded *big-endian*
19/// by default but can be encoded little-endian using the `LittleEndian<T>` wrapper type.
20pub trait TrieKey: PartialEq + Eq {
21    /// Encode a value as a vector of bytes.
22    fn encode_bytes(&self) -> Vec<u8> {
23        panic!("implement this method or TrieKey::encode");
24    }
25
26    /// Encode a value as a NibbleVec.
27    #[inline]
28    fn encode(&self) -> Nibblet {
29        Nibblet::from_byte_vec(self.encode_bytes())
30    }
31}
32
33/// Key comparison result.
34#[derive(Debug)]
35pub enum KeyMatch {
36    /// The keys match up to the given index.
37    Partial(usize),
38    /// The first key is a prefix of the second.
39    FirstPrefix,
40    /// The second key is a prefix of the first.
41    SecondPrefix,
42    /// The keys match exactly.
43    Full,
44}
45
46/// Compare two Trie keys.
47///
48/// Compares `first[start_idx .. ]` to `second`, i.e. only looks at a slice of the first key.
49#[inline]
50pub fn match_keys(start_idx: usize, first: &Nibblet, second: &Nibblet) -> KeyMatch {
51    let first_len = first.len() - start_idx;
52    let min_length = ::std::cmp::min(first_len, second.len());
53
54    for i in 0..min_length {
55        if first.get(start_idx + i) != second.get(i) {
56            return KeyMatch::Partial(i);
57        }
58    }
59
60    match (first_len, second.len()) {
61        (x, y) if x < y => KeyMatch::FirstPrefix,
62        (x, y) if x == y => KeyMatch::Full,
63        _ => KeyMatch::SecondPrefix,
64    }
65}
66
67/// Check two keys for equality and panic if they differ.
68#[inline]
69pub fn check_keys<K: ?Sized>(key1: &K, key2: &K)
70where
71    K: TrieKey,
72{
73    if *key1 != *key2 {
74        panic!("multiple-keys with the same bit representation.");
75    }
76}
77
78// --- TrieKey Implementations for standard types --- ///
79
80// This blanket implementation goes into play when specialization is stabilized
81// impl<T> TrieKey for T where T: Into<Vec<u8>> + Clone + Eq + PartialEq {
82// fn encode_bytes(&self) -> Vec<u8> {
83// self.clone().into()
84// }
85// }
86
87impl TrieKey for Vec<u8> {
88    #[inline]
89    fn encode_bytes(&self) -> Vec<u8> {
90        self.clone()
91    }
92}
93
94impl TrieKey for [u8] {
95    #[inline]
96    fn encode_bytes(&self) -> Vec<u8> {
97        self.to_vec()
98    }
99}
100
101impl TrieKey for String {
102    #[inline]
103    fn encode_bytes(&self) -> Vec<u8> {
104        self.as_bytes().encode_bytes()
105    }
106}
107
108impl TrieKey for str {
109    #[inline]
110    fn encode_bytes(&self) -> Vec<u8> {
111        self.as_bytes().encode_bytes()
112    }
113}
114
115impl<T: ?Sized + TrieKey> TrieKey for &T {
116    #[inline]
117    fn encode_bytes(&self) -> Vec<u8> {
118        (**self).encode_bytes()
119    }
120}
121
122impl<T: ?Sized + TrieKey> TrieKey for &mut T {
123    #[inline]
124    fn encode_bytes(&self) -> Vec<u8> {
125        (**self).encode_bytes()
126    }
127}
128
129impl TrieKey for i8 {
130    #[inline]
131    fn encode_bytes(&self) -> Vec<u8> {
132        let mut v: Vec<u8> = Vec::with_capacity(1);
133        v.push(*self as u8);
134        v
135    }
136}
137
138impl TrieKey for u8 {
139    #[inline]
140    fn encode_bytes(&self) -> Vec<u8> {
141        let mut v: Vec<u8> = Vec::with_capacity(1);
142        v.push(*self);
143        v
144    }
145}
146
147#[cfg(unix)]
148impl TrieKey for PathBuf {
149    fn encode_bytes(&self) -> Vec<u8> {
150        use std::ffi::OsString;
151        use std::os::unix::ffi::OsStringExt;
152        let str: OsString = self.clone().into();
153        str.into_vec()
154    }
155}
156
157#[cfg(unix)]
158impl TrieKey for Path {
159    fn encode_bytes(&self) -> Vec<u8> {
160        use std::os::unix::ffi::OsStrExt;
161        self.as_os_str().as_bytes().encode_bytes()
162    }
163}
164
165#[cfg(windows)]
166impl TrieKey for PathBuf {
167    fn encode_bytes(&self) -> Vec<u8> {
168        self.as_os_str().as_encoded_bytes().to_vec()
169    }
170}
171
172#[cfg(windows)]
173impl TrieKey for Path {
174    fn encode_bytes(&self) -> Vec<u8> {
175        self.as_os_str().as_encoded_bytes().to_vec()
176    }
177}
178
179impl<T> TrieKey for LittleEndian<T>
180where
181    T: Eq + Copy,
182{
183    fn encode_bytes(&self) -> Vec<u8> {
184        self.as_byte_slice().encode_bytes()
185    }
186}
187
188impl<T> TrieKey for BigEndian<T>
189where
190    T: Eq + Copy,
191{
192    fn encode_bytes(&self) -> Vec<u8> {
193        self.as_byte_slice().to_vec()
194    }
195}
196
197macro_rules! int_keys {
198    ( $( $t:ty ),* ) => {
199        $(
200        impl TrieKey for $t {
201            fn encode_bytes(&self) -> Vec<u8> {
202                let be: BigEndian<$t> = From::from(*self);
203                be.encode_bytes()
204            }
205        }
206        )*
207    };
208}
209
210int_keys!(u16, u32, u64, i16, i32, i64, usize, isize);
211
212macro_rules! vec_int_keys {
213  ( $( $t:ty ),* ) => {
214      $(
215      impl TrieKey for Vec<$t> {
216          fn encode_bytes(&self) -> Vec<u8> {
217              let mut v = Vec::<u8>::with_capacity(self.len() * std::mem::size_of::<$t>());
218              for u in self {
219                  v.extend_from_slice(&u.to_be_bytes());
220              }
221              v
222          }
223      }
224      )*
225   };
226}
227
228vec_int_keys!(u16, u32, u64, i16, i32, i64, usize, isize);
229
230#[cfg(test)]
231mod test {
232    pub trait DefaultTrieKey {
233        fn encode_bytes(&self) -> Vec<u8>;
234    }
235
236    impl<T: Into<Vec<u8>> + Clone + PartialEq + Eq> DefaultTrieKey for T {
237        #[inline]
238        fn encode_bytes(&self) -> Vec<u8> {
239            self.clone().into()
240        }
241    }
242
243    macro_rules! encode_bytes {
244        ($e:expr) => {
245            (&$e).encode_bytes()
246        };
247    }
248
249    #[test]
250    fn test_autoref_specialization() {
251        let _ = encode_bytes!([0_u8]);
252        let _ = encode_bytes!("hello");
253        let _ = encode_bytes!("hello".to_string());
254    }
255}