ndn_store/
trie.rs

1use std::collections::HashMap;
2use std::sync::{Arc, RwLock};
3
4use ndn_packet::{Name, NameComponent};
5
6/// A concurrent name trie mapping name prefixes to values of type `V`.
7///
8/// Uses per-node `RwLock` so that many readers can descend simultaneously
9/// without holding parent locks. This is the shared structure used by both
10/// the FIB and the StrategyTable.
11pub struct NameTrie<V: Clone + Send + Sync + 'static> {
12    root: Arc<RwLock<TrieNode<V>>>,
13}
14
15struct TrieNode<V> {
16    entry: Option<V>,
17    children: HashMap<NameComponent, Arc<RwLock<TrieNode<V>>>>,
18}
19
20impl<V> TrieNode<V> {
21    fn new() -> Self {
22        Self {
23            entry: None,
24            children: HashMap::new(),
25        }
26    }
27}
28
29impl<V: Clone + Send + Sync + 'static> NameTrie<V> {
30    pub fn new() -> Self {
31        Self {
32            root: Arc::new(RwLock::new(TrieNode::new())),
33        }
34    }
35
36    /// Longest-prefix match — returns the value at the deepest matching node.
37    pub fn lpm(&self, name: &Name) -> Option<V> {
38        let root = self.root.read().unwrap();
39        let mut best = root.entry.clone();
40        let mut current: Arc<RwLock<TrieNode<V>>> = Arc::clone(&self.root);
41        drop(root);
42
43        for component in name.components() {
44            let child_arc = {
45                let node = current.read().unwrap();
46                node.children.get(component).map(Arc::clone)
47            };
48            match child_arc {
49                None => break,
50                Some(child) => {
51                    let node = child.read().unwrap();
52                    if node.entry.is_some() {
53                        best = node.entry.clone();
54                    }
55                    drop(node);
56                    current = child;
57                }
58            }
59        }
60        best
61    }
62
63    /// Exact-prefix lookup — only returns a value if `name` exactly matches
64    /// a registered prefix.
65    pub fn get(&self, name: &Name) -> Option<V> {
66        let mut current = Arc::clone(&self.root);
67        for component in name.components() {
68            let child = {
69                let node = current.read().unwrap();
70                node.children.get(component).map(Arc::clone)
71            };
72            match child {
73                None => return None,
74                Some(c) => current = c,
75            }
76        }
77        let node = current.read().unwrap();
78        node.entry.clone()
79    }
80
81    /// Insert or replace the value at `name`.
82    pub fn insert(&self, name: &Name, value: V) {
83        let mut current = Arc::clone(&self.root);
84        for component in name.components() {
85            let child = {
86                let mut node = current.write().unwrap();
87                node.children
88                    .entry(component.clone())
89                    .or_insert_with(|| Arc::new(RwLock::new(TrieNode::new())))
90                    .clone()
91            };
92            current = child;
93        }
94        let mut node = current.write().unwrap();
95        node.entry = Some(value);
96    }
97
98    /// Remove the value at exactly `name`. Does not prune empty nodes.
99    pub fn remove(&self, name: &Name) {
100        let mut current = Arc::clone(&self.root);
101        for component in name.components() {
102            let child = {
103                let node = current.read().unwrap();
104                node.children.get(component).map(Arc::clone)
105            };
106            match child {
107                None => return,
108                Some(c) => current = c,
109            }
110        }
111        let mut node = current.write().unwrap();
112        node.entry = None;
113    }
114
115    /// Walk the entire trie and return all `(Name, V)` pairs in depth-first order.
116    pub fn dump(&self) -> Vec<(Name, V)> {
117        let mut out = Vec::new();
118        dump_subtree(&self.root, &mut Vec::new(), &mut out);
119        out
120    }
121
122    /// Collect all values stored at or below `prefix` in the trie.
123    ///
124    /// Used for prefix-based eviction (e.g. `cs erase /prefix`).
125    pub fn descendants(&self, prefix: &Name) -> Vec<V> {
126        let mut current = Arc::clone(&self.root);
127        for component in prefix.components() {
128            let child = {
129                let node = current.read().unwrap();
130                node.children.get(component).map(Arc::clone)
131            };
132            match child {
133                None => return Vec::new(),
134                Some(c) => current = c,
135            }
136        }
137        let mut out = Vec::new();
138        collect_subtree(&current, &mut out);
139        out
140    }
141
142    /// Returns the first value found at or below `prefix` in the trie.
143    ///
144    /// Used for `CanBePrefix` CS lookups: walk to the Interest name position,
145    /// then return any Data stored at or below that node. The traversal order
146    /// within a level is unspecified (HashMap iteration order).
147    pub fn first_descendant(&self, prefix: &Name) -> Option<V> {
148        let mut current = Arc::clone(&self.root);
149        for component in prefix.components() {
150            let child = {
151                let node = current.read().unwrap();
152                node.children.get(component).map(Arc::clone)
153            };
154            match child {
155                None => return None,
156                Some(c) => current = c,
157            }
158        }
159        first_in_subtree(&current)
160    }
161}
162
163impl<V: Clone + Send + Sync + 'static> Default for NameTrie<V> {
164    fn default() -> Self {
165        Self::new()
166    }
167}
168
169/// Depth-first collection of all (Name, V) entries at or below `node`.
170fn dump_subtree<V: Clone + Send + Sync + 'static>(
171    node: &Arc<RwLock<TrieNode<V>>>,
172    path: &mut Vec<NameComponent>,
173    out: &mut Vec<(Name, V)>,
174) {
175    let guard = node.read().unwrap();
176    if let Some(v) = &guard.entry {
177        out.push((Name::from_components(path.iter().cloned()), v.clone()));
178    }
179    // Collect children first to release the lock before recursing.
180    let children: Vec<(NameComponent, Arc<RwLock<TrieNode<V>>>)> = guard
181        .children
182        .iter()
183        .map(|(k, v)| (k.clone(), Arc::clone(v)))
184        .collect();
185    drop(guard);
186    for (comp, child) in children {
187        path.push(comp);
188        dump_subtree(&child, path, out);
189        path.pop();
190    }
191}
192
193/// Collect all values at or below `node` into `out`.
194fn collect_subtree<V: Clone + Send + Sync + 'static>(
195    node: &Arc<RwLock<TrieNode<V>>>,
196    out: &mut Vec<V>,
197) {
198    let guard = node.read().unwrap();
199    if let Some(v) = &guard.entry {
200        out.push(v.clone());
201    }
202    let children: Vec<Arc<RwLock<TrieNode<V>>>> = guard.children.values().map(Arc::clone).collect();
203    drop(guard);
204    for child in children {
205        collect_subtree(&child, out);
206    }
207}
208
209/// Depth-first search for the first value at or below `node`.
210fn first_in_subtree<V: Clone>(node: &Arc<RwLock<TrieNode<V>>>) -> Option<V> {
211    let guard = node.read().unwrap();
212    if let Some(v) = &guard.entry {
213        return Some(v.clone());
214    }
215    for child in guard.children.values() {
216        if let Some(v) = first_in_subtree(child) {
217            return Some(v);
218        }
219    }
220    None
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use bytes::Bytes;
227    use ndn_packet::{Name, NameComponent};
228
229    fn name(components: &[&str]) -> Name {
230        Name::from_components(
231            components
232                .iter()
233                .map(|s| NameComponent::generic(Bytes::copy_from_slice(s.as_bytes()))),
234        )
235    }
236
237    // ── lpm ──────────────────────────────────────────────────────────────────
238
239    #[test]
240    fn lpm_empty_trie_returns_none() {
241        let trie: NameTrie<u32> = NameTrie::new();
242        assert!(trie.lpm(&name(&["a", "b"])).is_none());
243    }
244
245    #[test]
246    fn lpm_exact_match() {
247        let trie: NameTrie<u32> = NameTrie::new();
248        trie.insert(&name(&["a", "b"]), 42);
249        assert_eq!(trie.lpm(&name(&["a", "b"])), Some(42));
250    }
251
252    #[test]
253    fn lpm_prefix_wins() {
254        let trie: NameTrie<u32> = NameTrie::new();
255        trie.insert(&name(&["a"]), 1);
256        trie.insert(&name(&["a", "b"]), 2);
257        // Query /a/b/c — most specific match is /a/b.
258        assert_eq!(trie.lpm(&name(&["a", "b", "c"])), Some(2));
259    }
260
261    #[test]
262    fn lpm_shorter_prefix_fallback() {
263        let trie: NameTrie<u32> = NameTrie::new();
264        trie.insert(&name(&["a"]), 1);
265        // /a/b is not in the trie; fallback to /a.
266        assert_eq!(trie.lpm(&name(&["a", "b"])), Some(1));
267    }
268
269    #[test]
270    fn lpm_root_matches_everything() {
271        let trie: NameTrie<u32> = NameTrie::new();
272        trie.insert(&Name::root(), 99);
273        assert_eq!(trie.lpm(&name(&["x", "y", "z"])), Some(99));
274    }
275
276    // ── get (exact) ───────────────────────────────────────────────────────────
277
278    #[test]
279    fn get_returns_none_for_missing_prefix() {
280        let trie: NameTrie<u32> = NameTrie::new();
281        trie.insert(&name(&["a", "b"]), 5);
282        assert!(trie.get(&name(&["a"])).is_none());
283        assert!(trie.get(&name(&["a", "b", "c"])).is_none());
284    }
285
286    #[test]
287    fn get_returns_exact_entry() {
288        let trie: NameTrie<u32> = NameTrie::new();
289        trie.insert(&name(&["a", "b"]), 7);
290        assert_eq!(trie.get(&name(&["a", "b"])), Some(7));
291    }
292
293    // ── insert / remove ───────────────────────────────────────────────────────
294
295    #[test]
296    fn insert_replaces_value() {
297        let trie: NameTrie<u32> = NameTrie::new();
298        trie.insert(&name(&["a"]), 1);
299        trie.insert(&name(&["a"]), 2);
300        assert_eq!(trie.get(&name(&["a"])), Some(2));
301    }
302
303    #[test]
304    fn remove_clears_entry() {
305        let trie: NameTrie<u32> = NameTrie::new();
306        trie.insert(&name(&["a", "b"]), 10);
307        trie.remove(&name(&["a", "b"]));
308        assert!(trie.get(&name(&["a", "b"])).is_none());
309    }
310
311    #[test]
312    fn remove_nonexistent_is_noop() {
313        let trie: NameTrie<u32> = NameTrie::new();
314        trie.remove(&name(&["x"])); // should not panic
315    }
316
317    // ── first_descendant ─────────────────────────────────────────────────────
318
319    #[test]
320    fn first_descendant_exact_node_has_value() {
321        let trie: NameTrie<u32> = NameTrie::new();
322        trie.insert(&name(&["a", "b"]), 5);
323        // first_descendant of /a/b finds the value at /a/b itself.
324        assert_eq!(trie.first_descendant(&name(&["a", "b"])), Some(5));
325    }
326
327    #[test]
328    fn first_descendant_finds_child() {
329        let trie: NameTrie<u32> = NameTrie::new();
330        trie.insert(&name(&["a", "b", "c"]), 42);
331        // first_descendant of /a/b finds /a/b/c.
332        assert_eq!(trie.first_descendant(&name(&["a", "b"])), Some(42));
333    }
334
335    #[test]
336    fn first_descendant_missing_prefix_returns_none() {
337        let trie: NameTrie<u32> = NameTrie::new();
338        trie.insert(&name(&["x"]), 1);
339        assert!(trie.first_descendant(&name(&["y"])).is_none());
340    }
341
342    #[test]
343    fn first_descendant_empty_prefix_returns_any() {
344        let trie: NameTrie<u32> = NameTrie::new();
345        trie.insert(&name(&["a"]), 10);
346        // first_descendant of root (empty prefix) returns some value.
347        assert!(trie.first_descendant(&Name::root()).is_some());
348    }
349
350    #[test]
351    fn first_descendant_no_children_no_value_returns_none() {
352        let trie: NameTrie<u32> = NameTrie::new();
353        trie.insert(&name(&["a"]), 1);
354        // /a/b exists in path as intermediate node but has no value or children.
355        assert!(trie.first_descendant(&name(&["a", "b"])).is_none());
356    }
357}
ndn_store/trie.rs

ndn_store/
trie.rs
