quotick/radix_trie/
trie.rs

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

quotick/radix_trie/
trie.rs
