radix_trie/
subtrie.rs

1use crate::keys::*;
2use crate::TrieNode;
3use crate::{SubTrie, SubTrieMut, SubTrieResult};
4use std::borrow::Borrow;
5
6use nibble_vec::Nibblet;
7
8impl<'a, K, V> SubTrie<'a, K, V>
9where
10    K: TrieKey,
11{
12    /// Look up the value for the given key, which should be an extension of this subtrie's key.
13    ///
14    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
15    /// form *must* match those for the key type
16    pub fn get<Q: ?Sized>(&self, key: &Q) -> SubTrieResult<&V>
17    where
18        K: Borrow<Q>,
19        Q: TrieKey,
20    {
21        subtrie_get(&self.prefix, self.node, key)
22    }
23}
24
25fn subtrie_get<'a, K, Q: ?Sized, V>(
26    prefix: &Nibblet,
27    node: &'a TrieNode<K, V>,
28    key: &Q,
29) -> SubTrieResult<&'a V>
30where
31    K: TrieKey,
32    K: Borrow<Q>,
33    Q: TrieKey,
34{
35    let key_enc = key.encode();
36    match match_keys(0, prefix, &key_enc) {
37        KeyMatch::Full => Ok(node.value()),
38        KeyMatch::FirstPrefix => Ok(node
39            .get(&stripped(key_enc, prefix))
40            .and_then(TrieNode::value)),
41        _ => Err(()),
42    }
43}
44
45impl<'a, K, V> SubTrieMut<'a, K, V>
46where
47    K: TrieKey,
48{
49    /// Mutable reference to the node's value.
50    pub fn value_mut(&mut self) -> Option<&mut V> {
51        self.node.value_mut()
52    }
53
54    /// Look up the value for the given key, which should be an extension of this subtrie's key.
55    ///
56    /// The key may be any borrowed form of the trie's key type, but TrieKey on the borrowed
57    /// form *must* match those for the key type
58    pub fn get<Q: ?Sized>(&self, key: &Q) -> SubTrieResult<&V>
59    where
60        K: Borrow<Q>,
61        Q: TrieKey,
62    {
63        subtrie_get(&self.prefix, &*self.node, key)
64    }
65
66    /// Insert a value in this subtrie. The key should be an extension of this subtrie's key.
67    pub fn insert(&mut self, key: K, value: V) -> SubTrieResult<V> {
68        let key_enc = key.encode();
69        let previous = match match_keys(0, &self.prefix, &key_enc) {
70            KeyMatch::Full => self.node.replace_value(key, value),
71            KeyMatch::FirstPrefix => self
72                .node
73                .insert(key, value, stripped(key_enc, &self.prefix)),
74            _ => {
75                return Err(());
76            }
77        };
78
79        if previous.is_none() {
80            *self.length += 1;
81        }
82
83        Ok(previous)
84    }
85
86    /// Remove a value from this subtrie. The key should be an extension of this subtrie's 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    pub fn remove<Q: ?Sized>(&mut self, key: &Q) -> SubTrieResult<V>
91    where
92        K: Borrow<Q>,
93        Q: TrieKey,
94    {
95        let key_enc = key.encode();
96        let removed = match match_keys(0, &self.prefix, &key_enc) {
97            KeyMatch::Full => self.node.take_value(key),
98            KeyMatch::FirstPrefix => self.node.remove(key),
99            _ => {
100                return Err(());
101            }
102        };
103
104        if removed.is_some() {
105            *self.length -= 1;
106        }
107
108        Ok(removed)
109    }
110}
111
112fn stripped(mut key: Nibblet, prefix: &Nibblet) -> Nibblet {
113    key.split(prefix.len())
114}
radix_trie/subtrie.rs

radix_trie/
subtrie.rs
