Skip to main content

ts_bart/iptrie/lookup/
iter.rs

1use crate::{PrefixReadOps, iptrie, node::NodePrefixIter};
2
3/// Iterator over prefixes matching a particular address.
4///
5/// Built by [`iptrie::lookup_address_all`].
6pub struct LookupIter<'n, T> {
7    /// The first item, if there was a child entry.
8    pub(super) value: Option<&'n T>,
9
10    /// The stack of nodes in the trie that matched the address query.
11    pub(super) stack: heapless::Vec<(&'n dyn PrefixReadOps<T = T>, u8), { iptrie::MAX_DEPTH }>,
12
13    /// Intermediate partially-yielded state for prefix matches on the current stack item.
14    pub(super) prefix_iter: Option<NodePrefixIter<'n, T>>,
15}
16
17impl<T> Default for LookupIter<'_, T> {
18    fn default() -> Self {
19        Self {
20            value: None,
21            stack: Default::default(),
22            prefix_iter: None,
23        }
24    }
25}
26
27impl<'n, T> Iterator for LookupIter<'n, T>
28where
29    T: 'static,
30{
31    type Item = &'n T;
32
33    fn next(&mut self) -> Option<Self::Item> {
34        // There was a terminal leaf or fringe value: return it first.
35        if let Some(val) = self.value.take() {
36            return Some(val);
37        }
38
39        walk_stack(&mut self.stack, &mut self.prefix_iter)
40    }
41}
42
43pub(super) fn walk_stack<'n, T>(
44    stack: &mut heapless::Vec<(&'n dyn PrefixReadOps<T = T>, u8), { iptrie::MAX_DEPTH }>,
45    prefix_iter: &mut Option<NodePrefixIter<'n, T>>,
46) -> Option<&'n T>
47where
48    T: 'static,
49{
50    while let Some(&(node, octet)) = stack.last() {
51        if let Some(iter) = prefix_iter {
52            if let Some((_idx, val)) = iter.next() {
53                return Some(val);
54            }
55
56            prefix_iter.take();
57            stack.pop();
58            continue;
59        }
60
61        // Populate state for the next iteration of the loop
62        *prefix_iter = Some(NodePrefixIter::for_octet(node, octet));
63    }
64
65    None
66}