radix_trie/
trie.rs

1use crate::traversal::DescendantResult::*;
2use crate::TrieNode;
3use crate::{SubTrie, SubTrieMut, Trie, TrieCommon, TrieKey};
4use std::borrow::Borrow;
5
6use nibble_vec::Nibblet;
7
8impl<K, V> Trie<K, V>
9where
10    K: TrieKey,
11{
12    /// Create an empty Trie.
13    #[inline]
14    pub fn new() -> Trie<K, V> {
15        Trie {
16            length: 0,
17            node: TrieNode::new(),
18        }
19    }
20
21    /// Fetch a reference to the given key's corresponding value, if any.
22    ///
23    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
24    /// form *must* match those for the key type.
25    #[inline]
26    pub fn get<Q: ?Sized>(&self, key: &Q) -> Option<&V>
27    where
28        K: Borrow<Q>,
29        Q: TrieKey,
30    {
31        let key_fragments = key.encode();
32        self.node
33            .get(&key_fragments)
34            .and_then(|t| t.value_checked(key))
35    }
36
37    /// Fetch a mutable reference to the given key's corresponding value, if any.
38    ///
39    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
40    /// form *must* match those for the key type.
41    #[inline]
42    pub fn get_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V>
43    where
44        K: Borrow<Q>,
45        Q: TrieKey,
46    {
47        let key_fragments = key.encode();
48        self.node
49            .get_mut(&key_fragments)
50            .and_then(|t| t.value_checked_mut(key))
51    }
52
53    /// Insert the given key-value pair, returning any previous value associated with the key.
54    #[inline]
55    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
56        let key_fragments = key.encode();
57        let result = self.node.insert(key, value, key_fragments);
58        if result.is_none() {
59            self.length += 1;
60        }
61        result
62    }
63
64    /// Remove the value associated with the given key.
65    ///
66    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
67    /// form *must* match those for the key type.
68    #[inline]
69    pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
70    where
71        K: Borrow<Q>,
72        Q: TrieKey,
73    {
74        let removed = self.node.remove(key);
75        if removed.is_some() {
76            self.length -= 1;
77        }
78        removed
79    }
80
81    /// Get a mutable reference to the value stored at this node, if any.
82    pub fn value_mut(&mut self) -> Option<&mut V> {
83        self.node.value_mut()
84    }
85
86    /// Fetch a reference to the subtrie for a given key.
87    ///
88    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
89    /// form *must* match those for the key type.
90    #[inline]
91    pub fn subtrie<'a, Q: ?Sized>(&'a self, key: &Q) -> Option<SubTrie<'a, K, V>>
92    where
93        K: Borrow<Q>,
94        Q: TrieKey,
95    {
96        let key_fragments = key.encode();
97        self.node
98            .get(&key_fragments)
99            .map(|node| node.as_subtrie(key_fragments))
100    }
101
102    /// Fetch a mutable reference to the subtrie for a given key.
103    ///
104    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
105    /// form *must* match those for the key type.
106    #[inline]
107    pub fn subtrie_mut<'a, Q: ?Sized>(&'a mut self, key: &Q) -> Option<SubTrieMut<'a, K, V>>
108    where
109        K: Borrow<Q>,
110        Q: TrieKey,
111    {
112        let key_fragments = key.encode();
113        let length_ref = &mut self.length;
114        self.node
115            .get_mut(&key_fragments)
116            .map(move |node| node.as_subtrie_mut(key_fragments, length_ref))
117    }
118
119    /// Fetch a reference to the closest ancestor node of the given key.
120    ///
121    /// If `key` is encoded as byte-vector `b`, return the node `n` in the tree
122    /// such that `n`'s key's byte-vector is the longest possible prefix of `b`, and `n`
123    /// has a value.
124    ///
125    /// Invariant: `result.is_some() => result.key_value.is_some()`.
126    ///
127    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
128    /// form *must* match those for the key type.
129    #[inline]
130    pub fn get_ancestor<'a, Q: ?Sized>(&'a self, key: &Q) -> Option<SubTrie<'a, K, V>>
131    where
132        K: Borrow<Q>,
133        Q: TrieKey,
134    {
135        let mut key_fragments = key.encode();
136        self.node
137            .get_ancestor(&key_fragments)
138            .map(|(node, node_key_len)| {
139                key_fragments.split(node_key_len);
140                node.as_subtrie(key_fragments)
141            })
142    }
143
144    /// Fetch the closest ancestor *value* for a given key.
145    ///
146    /// See `get_ancestor` for precise semantics, this is just a shortcut.
147    ///
148    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
149    /// form *must* match those for the key type.
150    #[inline]
151    pub fn get_ancestor_value<Q: ?Sized>(&self, key: &Q) -> Option<&V>
152    where
153        K: Borrow<Q>,
154        Q: TrieKey,
155    {
156        self.get_ancestor(key).and_then(|t| t.node.value())
157    }
158
159    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
160    /// form *must* match those for the key type
161    #[inline]
162    pub fn get_raw_ancestor<'a, Q: ?Sized>(&'a self, key: &Q) -> SubTrie<'a, K, V>
163    where
164        K: Borrow<Q>,
165        Q: TrieKey,
166    {
167        let mut nv = key.encode();
168        let (ancestor_node, depth) = self.node.get_raw_ancestor(&nv);
169        nv.split(depth);
170        ancestor_node.as_subtrie(nv)
171    }
172
173    /// Fetch the closest descendant for a given key.
174    ///
175    /// If the key is in the trie, this is the same as `subtrie`.
176    ///
177    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
178    /// form *must* match those for the key type
179    #[inline]
180    pub fn get_raw_descendant<'a, Q: ?Sized>(&'a self, key: &Q) -> Option<SubTrie<'a, K, V>>
181    where
182        K: Borrow<Q>,
183        Q: TrieKey,
184    {
185        let mut nv = key.encode();
186        self.node.get_raw_descendant(&nv).map(|desc| {
187            let (node, prefix) = match desc {
188                NoModification(node) => (node, nv),
189                ExtendKey(node, depth, extension) => {
190                    nv.split(depth);
191                    (node, nv.join(extension))
192                }
193            };
194            node.as_subtrie(prefix)
195        })
196    }
197
198    /// Take a function `f` and apply it to the value stored at `key`.
199    ///
200    /// If no value is stored at `key`, store `default`.
201    #[inline]
202    pub fn map_with_default<F>(&mut self, key: K, f: F, default: V)
203    where
204        F: Fn(&mut V),
205    {
206        {
207            if let Some(v) = self.get_mut(&key) {
208                f(v);
209                return;
210            }
211        }
212        self.insert(key, default);
213    }
214
215    /// Check that the Trie invariants are satisfied - you shouldn't ever have to call this!
216    /// Quite slow!
217    #[doc(hidden)]
218    pub fn check_integrity(&self) -> bool {
219        let (ok, length) = self.node.check_integrity_recursive(&Nibblet::new());
220        ok && length == self.length
221    }
222}
223
224impl<K, V> PartialEq for Trie<K, V>
225where
226    K: TrieKey,
227    V: PartialEq,
228{
229    #[inline]
230    fn eq(&self, other: &Trie<K, V>) -> bool {
231        if self.len() != other.len() {
232            return false;
233        }
234
235        self.iter()
236            .all(|(key, value)| other.get(key).map_or(false, |v| *value == *v))
237    }
238}
239
240impl<K: TrieKey, V> Default for Trie<K, V> {
241    #[inline]
242    fn default() -> Self {
243        Self::new()
244    }
245}
radix_trie/trie.rs

radix_trie/
trie.rs
