radix_trie/
trie_node.rs

Help
1use crate::keys::*;
2use crate::{SubTrie, SubTrieMut, BRANCH_FACTOR};
3use std::borrow::Borrow;
4use std::default::Default;
5
6use nibble_vec::Nibblet;
7
8#[derive(Debug, Clone)]
9pub struct TrieNode<K, V> {
10    /// Key fragments/bits associated with this node, such that joining the keys from all
11    /// parent nodes and this node is equal to the bit-encoding of this node's key.
12    pub key: Nibblet,
13
14    /// The key and value stored at this node.
15    pub key_value: Option<Box<KeyValue<K, V>>>,
16
17    /// The number of children which are Some rather than None.
18    pub child_count: usize,
19
20    /// The children of this node stored such that the first nibble of each child key
21    /// dictates the child's bucket.
22    pub children: [Option<Box<TrieNode<K, V>>>; BRANCH_FACTOR],
23}
24
25#[derive(Debug, Clone)]
26pub struct KeyValue<K, V> {
27    pub key: K,
28    pub value: V,
29}
30
31impl<K, V> TrieNode<K, V>
32where
33    K: TrieKey,
34{
35    /// Create a value-less, child-less TrieNode.
36    #[inline]
37    pub fn new() -> TrieNode<K, V> {
38        TrieNode {
39            key: Nibblet::new(),
40            key_value: None,
41            children: Default::default(),
42            child_count: 0,
43        }
44    }
45
46    /// Create a TrieNode with no children.
47    #[inline]
48    pub fn with_key_value(key_fragments: Nibblet, key: K, value: V) -> TrieNode<K, V> {
49        TrieNode {
50            key: key_fragments,
51            key_value: Some(Box::new(KeyValue {
52                key: key,
53                value: value,
54            })),
55            children: Default::default(),
56            child_count: 0,
57        }
58    }
59
60    /// Get the key stored at this node, if any.
61    #[inline]
62    pub fn key(&self) -> Option<&K> {
63        self.key_value.as_ref().map(|kv| &kv.key)
64    }
65
66    /// Get the value stored at this node, if any.
67    #[inline]
68    pub fn value(&self) -> Option<&V> {
69        self.key_value.as_ref().map(|kv| &kv.value)
70    }
71
72    /// Get a mutable reference to the value stored at this node, if any.
73    #[inline]
74    pub fn value_mut(&mut self) -> Option<&mut V> {
75        self.key_value.as_mut().map(|kv| &mut kv.value)
76    }
77
78    /// Get the value whilst checking a key match.
79    ///
80    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
81    /// form *must* match those for the key type.
82    #[inline]
83    pub fn value_checked<Q: ?Sized>(&self, key: &Q) -> Option<&V>
84    where
85        K: Borrow<Q>,
86        Q: TrieKey,
87    {
88        self.key_value.as_ref().map(|kv| {
89            check_keys(kv.key.borrow(), key);
90            &kv.value
91        })
92    }
93
94    /// Get a mutable value whilst checking a key match.
95    ///
96    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
97    /// form *must* match those for the key type.
98    #[inline]
99    pub fn value_checked_mut<Q: ?Sized>(&mut self, key: &Q) -> Option<&mut V>
100    where
101        K: Borrow<Q>,
102        Q: TrieKey,
103    {
104        self.key_value.as_mut().map(|kv| {
105            check_keys(kv.key.borrow(), key);
106            &mut kv.value
107        })
108    }
109
110    /// Compute the number of keys and values in this node's subtrie.
111    #[inline]
112    pub fn compute_size(&self) -> usize {
113        let mut size = self.key_value.is_some() as usize;
114
115        for child in &self.children {
116            if let Some(ref child) = *child {
117                // TODO: could unroll this recursion
118                size += child.compute_size();
119            }
120        }
121
122        size
123    }
124
125    /// Add a child at the given index, given that none exists there already.
126    #[inline]
127    pub fn add_child(&mut self, idx: usize, node: Box<TrieNode<K, V>>) {
128        debug_assert!(self.children[idx].is_none());
129        self.child_count += 1;
130        self.children[idx] = Some(node);
131    }
132
133    /// Remove a child at the given index, if it exists.
134    #[inline]
135    pub fn take_child(&mut self, idx: usize) -> Option<Box<TrieNode<K, V>>> {
136        self.children[idx].take().map(|node| {
137            self.child_count -= 1;
138            node
139        })
140    }
141
142    /// Helper function for removing the single child of a node.
143    #[inline]
144    pub fn take_only_child(&mut self) -> Box<TrieNode<K, V>> {
145        debug_assert_eq!(self.child_count, 1);
146        for i in 0..BRANCH_FACTOR {
147            if let Some(child) = self.take_child(i) {
148                return child;
149            }
150        }
151        unreachable!("node with child_count 1 has no actual children");
152    }
153
154    /// Set the key and value of a node, given that it currently lacks one.
155    #[inline]
156    pub fn add_key_value(&mut self, key: K, value: V) {
157        debug_assert!(self.key_value.is_none());
158        self.key_value = Some(Box::new(KeyValue { key, value }));
159    }
160
161    /// Move the value out of a node, whilst checking that its key is as expected.
162    /// Can panic (see check_keys).
163    ///
164    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
165    /// form *must* match those for the key type
166    #[inline]
167    pub fn take_value<Q: ?Sized>(&mut self, key: &Q) -> Option<V>
168    where
169        K: Borrow<Q>,
170        Q: TrieKey,
171    {
172        self.key_value.take().map(|kv| {
173            check_keys(kv.key.borrow(), key);
174            kv.value
175        })
176    }
177
178    /// Replace a value, returning the previous value if there was one.
179    #[inline]
180    pub fn replace_value(&mut self, key: K, value: V) -> Option<V> {
181        // TODO: optimise this?
182        let previous = self.take_value(&key);
183        self.add_key_value(key, value);
184        previous
185    }
186
187    /// Get a reference to this node if it has a value.
188    #[inline]
189    pub fn as_value_node(&self) -> Option<&TrieNode<K, V>> {
190        self.key_value.as_ref().map(|_| self)
191    }
192
193    /// Split a node at a given index in its key, transforming it into a prefix node of its
194    /// previous self.
195    #[inline]
196    pub fn split(&mut self, idx: usize) {
197        // Extract all the parts of the suffix node, starting with the key.
198        let key = self.key.split(idx);
199
200        // Key-value.
201        let key_value = self.key_value.take();
202
203        // Children.
204        let mut children: [Option<Box<TrieNode<K, V>>>; BRANCH_FACTOR] = Default::default();
205
206        for (i, child) in self.children.iter_mut().enumerate() {
207            if child.is_some() {
208                children[i] = child.take();
209            }
210        }
211
212        // Child count.
213        let child_count = self.child_count;
214        self.child_count = 1;
215
216        // Insert the collected items below what is now an empty prefix node.
217        let bucket = key.get(0) as usize;
218        self.children[bucket] = Some(Box::new(TrieNode {
219            key: key,
220            key_value,
221            children,
222            child_count,
223        }));
224    }
225    #[inline]
226    pub fn as_subtrie(&self, prefix: Nibblet) -> SubTrie<K, V> {
227        SubTrie {
228            prefix: prefix,
229            node: self,
230        }
231    }
232    #[inline]
233    pub fn as_subtrie_mut<'a>(
234        &'a mut self,
235        prefix: Nibblet,
236        length: &'a mut usize,
237    ) -> SubTrieMut<'a, K, V> {
238        SubTrieMut {
239            prefix: prefix,
240            length: length,
241            node: self,
242        }
243    }
244
245    /// Check the integrity of a trie subtree (quite costly).
246    /// Return true and the size of the subtree if all checks are successful,
247    /// or false and a junk value if any test fails.
248    pub fn check_integrity_recursive(&self, prefix: &Nibblet) -> (bool, usize) {
249        let mut sub_tree_size = 0;
250        let is_root = prefix.len() == 0;
251
252        // Check that no value-less, non-root nodes have only 1 child.
253        if !is_root && self.child_count == 1 && self.key_value.is_none() {
254            println!("Value-less node with a single child.");
255            return (false, sub_tree_size);
256        }
257
258        // Check that all non-root key vector's have length > 1.
259        if !is_root && self.key.len() == 0 {
260            println!("Key length is 0 at non-root node.");
261            return (false, sub_tree_size);
262        }
263
264        // Check that the child count matches the actual number of children.
265        let child_count = self
266            .children
267            .iter()
268            .fold(0, |acc, e| acc + (e.is_some() as usize));
269
270        if child_count != self.child_count {
271            println!(
272                "Child count error, recorded: {}, actual: {}",
273                self.child_count, child_count
274            );
275            return (false, sub_tree_size);
276        }
277
278        // Compute the key fragments for this node, according to the trie.
279        let trie_key = prefix.clone().join(&self.key);
280
281        // Account for this node in the size check, and check its key.
282        if let Some(ref kv) = self.key_value {
283            sub_tree_size += 1;
284
285            let actual_key = kv.key.encode();
286
287            if trie_key != actual_key {
288                return (false, sub_tree_size);
289            }
290        }
291
292        // Recursively check children.
293        for i in 0..BRANCH_FACTOR {
294            if let Some(ref child) = self.children[i] {
295                match child.check_integrity_recursive(&trie_key) {
296                    (false, _) => return (false, sub_tree_size),
297                    (true, child_size) => sub_tree_size += child_size,
298                }
299            }
300        }
301
302        (true, sub_tree_size)
303    }
304}
305
306impl<K: TrieKey, V> Default for TrieNode<K, V> {
307    fn default() -> Self {
308        Self::new()
309    }
310}
radix_trie/trie_node.rs

radix_trie/
trie_node.rs
