radix_trie/
iter.rs

Help
1//! Iterators over key-value pairs, keys, values and child subtries.
2
3use std::iter::{FilterMap, FromIterator, Map};
4use std::slice;
5
6use crate::TrieNode;
7use crate::{SubTrie, Trie, TrieKey};
8
9use nibble_vec::Nibblet;
10
11// MY EYES.
12type Child<K, V> = Box<TrieNode<K, V>>;
13type RawChildIter<'a, K, V> = slice::Iter<'a, Option<Child<K, V>>>;
14type ChildMapFn<'a, K, V> = fn(&'a Option<Child<K, V>>) -> Option<&'a Child<K, V>>;
15type ChildIter<'a, K, V> = FilterMap<RawChildIter<'a, K, V>, ChildMapFn<'a, K, V>>;
16
17/// Iterator over the keys and values of a Trie.
18pub struct Iter<'a, K: 'a, V: 'a> {
19    root: &'a TrieNode<K, V>,
20    root_visited: bool,
21    stack: Vec<ChildIter<'a, K, V>>,
22}
23
24impl<'a, K, V> Iter<'a, K, V> {
25    // TODO: make this private somehow (and same for the other iterators).
26    pub fn new(root: &'a TrieNode<K, V>) -> Iter<'a, K, V> {
27        Iter {
28            root: root,
29            root_visited: false,
30            stack: vec![],
31        }
32    }
33}
34
35/// Iterator over the keys of a Trie.
36pub struct Keys<'a, K: 'a, V: 'a> {
37    inner: Map<Iter<'a, K, V>, KeyMapFn<'a, K, V>>,
38}
39
40type KeyMapFn<'a, K, V> = fn((&'a K, &'a V)) -> &'a K;
41
42impl<'a, K, V> Keys<'a, K, V> {
43    pub fn new(iter: Iter<'a, K, V>) -> Keys<'a, K, V> {
44        fn first<'b, K, V>((k, _): (&'b K, &'b V)) -> &'b K {
45            k
46        }
47        Keys {
48            inner: iter.map(first),
49        }
50    }
51}
52
53impl<'a, K, V> Iterator for Keys<'a, K, V> {
54    type Item = &'a K;
55
56    fn next(&mut self) -> Option<&'a K> {
57        self.inner.next()
58    }
59}
60
61/// Iterator over the values of a Trie.
62pub struct Values<'a, K: 'a, V: 'a> {
63    inner: Map<Iter<'a, K, V>, ValueMapFn<'a, K, V>>,
64}
65
66type ValueMapFn<'a, K, V> = fn((&'a K, &'a V)) -> &'a V;
67
68impl<'a, K, V> Values<'a, K, V> {
69    pub fn new(iter: Iter<'a, K, V>) -> Values<'a, K, V> {
70        fn second<'b, K, V>((_, v): (&'b K, &'b V)) -> &'b V {
71            v
72        }
73        Values {
74            inner: iter.map(second),
75        }
76    }
77}
78
79impl<'a, K, V> Iterator for Values<'a, K, V> {
80    type Item = &'a V;
81
82    fn next(&mut self) -> Option<&'a V> {
83        self.inner.next()
84    }
85}
86
87/// Iterator over the child subtries of a trie.
88pub struct Children<'a, K: 'a, V: 'a> {
89    prefix: Nibblet,
90    inner: ChildIter<'a, K, V>,
91}
92
93impl<'a, K, V> Children<'a, K, V> {
94    pub fn new(key: Nibblet, node: &'a TrieNode<K, V>) -> Self {
95        Children {
96            prefix: key,
97            inner: node.child_iter(),
98        }
99    }
100}
101
102impl<'a, K, V> Iterator for Children<'a, K, V> {
103    type Item = SubTrie<'a, K, V>;
104
105    fn next(&mut self) -> Option<SubTrie<'a, K, V>> {
106        self.inner.next().map(|node| SubTrie {
107            prefix: self.prefix.clone().join(&node.key),
108            node: node,
109        })
110    }
111}
112
113impl<K, V> TrieNode<K, V> {
114    /// Helper function to get all the non-empty children of a node.
115    fn child_iter(&self) -> ChildIter<K, V> {
116        fn id<K, V>(x: &Option<Child<K, V>>) -> Option<&Child<K, V>> {
117            x.as_ref()
118        }
119
120        self.children.iter().filter_map(id)
121    }
122
123    /// Get the key and value of a node as a pair.
124    fn kv_as_pair(&self) -> Option<(&K, &V)> {
125        self.key_value.as_ref().map(|kv| (&kv.key, &kv.value))
126    }
127}
128
129enum IterAction<'a, K: 'a, V: 'a> {
130    Push(&'a TrieNode<K, V>),
131    Pop,
132}
133
134impl<'a, K, V> Iterator for Iter<'a, K, V> {
135    type Item = (&'a K, &'a V);
136
137    fn next(&mut self) -> Option<Self::Item> {
138        use self::IterAction::*;
139
140        // Visit each node as it is reached from its parent (with special root handling).
141        if !self.root_visited {
142            self.root_visited = true;
143            self.stack.push(self.root.child_iter());
144            if let Some(kv) = self.root.kv_as_pair() {
145                return Some(kv);
146            }
147        }
148
149        loop {
150            let action = match self.stack.last_mut() {
151                Some(stack_top) => match stack_top.next() {
152                    Some(child) => Push(child),
153                    None => Pop,
154                },
155                None => return None,
156            };
157
158            match action {
159                Push(trie) => {
160                    self.stack.push(trie.child_iter());
161                    if let Some(kv) = trie.kv_as_pair() {
162                        return Some(kv);
163                    }
164                }
165                Pop => {
166                    self.stack.pop();
167                }
168            }
169        }
170    }
171}
172
173impl<K, V> FromIterator<(K, V)> for Trie<K, V>
174where
175    K: TrieKey,
176{
177    fn from_iter<T>(iter: T) -> Trie<K, V>
178    where
179        T: IntoIterator<Item = (K, V)>,
180    {
181        let mut trie = Trie::new();
182        for (k, v) in iter {
183            trie.insert(k, v);
184        }
185        trie
186    }
187}
radix_trie/iter.rs

radix_trie/
iter.rs
