prefix_trie/map/
iter.rs

1//! Module that contains the implementation for the iterators
2
3use map::Table;
4
5use crate::*;
6
7use super::Node;
8
9/// An iterator over all entries of a [`PrefixMap`] in lexicographic order.
10pub struct Iter<'a, P, T> {
11    pub(super) table: Option<&'a Table<P, T>>,
12    pub(super) nodes: Vec<usize>,
13}
14
15impl<P, T> Clone for Iter<'_, P, T> {
16    fn clone(&self) -> Self {
17        Self {
18            table: self.table,
19            nodes: self.nodes.clone(),
20        }
21    }
22}
23
24impl<P, T> Default for Iter<'_, P, T> {
25    fn default() -> Self {
26        Self {
27            table: None,
28            nodes: Vec::new(),
29        }
30    }
31}
32
33impl<'a, P, T> Iter<'a, P, T> {
34    pub(crate) fn new(table: &'a Table<P, T>, nodes: Vec<usize>) -> Self {
35        Self {
36            table: Some(table),
37            nodes,
38        }
39    }
40}
41
42impl<'a, P, T> Iterator for Iter<'a, P, T> {
43    type Item = (&'a P, &'a T);
44
45    fn next(&mut self) -> Option<(&'a P, &'a T)> {
46        while let Some(cur) = self.nodes.pop() {
47            let node = &self.table.as_ref()?[cur];
48            if let Some(right) = node.right {
49                self.nodes.push(right);
50            }
51            if let Some(left) = node.left {
52                self.nodes.push(left);
53            }
54            if let Some(v) = &node.value {
55                return Some((&node.prefix, v));
56            }
57        }
58        None
59    }
60}
61
62/// An iterator over all prefixes of a [`PrefixMap`] in lexicographic order.
63#[derive(Clone, Default)]
64pub struct Keys<'a, P, T> {
65    pub(crate) inner: Iter<'a, P, T>,
66}
67
68impl<'a, P, T> Iterator for Keys<'a, P, T> {
69    type Item = &'a P;
70
71    fn next(&mut self) -> Option<&'a P> {
72        self.inner.next().map(|(k, _)| k)
73    }
74}
75
76/// An iterator over all values of a [`PrefixMap`] in lexicographic order of their associated
77/// prefixes.
78#[derive(Clone, Default)]
79pub struct Values<'a, P, T> {
80    pub(crate) inner: Iter<'a, P, T>,
81}
82
83impl<'a, P, T> Iterator for Values<'a, P, T> {
84    type Item = &'a T;
85
86    fn next(&mut self) -> Option<&'a T> {
87        self.inner.next().map(|(_, v)| v)
88    }
89}
90
91/// An iterator over all owned entries of a [`PrefixMap`] in lexicographic order.
92#[derive(Clone)]
93pub struct IntoIter<P, T> {
94    pub(super) table: Vec<Node<P, T>>,
95    pub(super) nodes: Vec<usize>,
96}
97
98impl<P: Prefix, T> Iterator for IntoIter<P, T> {
99    type Item = (P, T);
100
101    fn next(&mut self) -> Option<(P, T)> {
102        while let Some(cur) = self.nodes.pop() {
103            let node = &mut self.table[cur];
104            if let Some(right) = node.right {
105                self.nodes.push(right);
106            }
107            if let Some(left) = node.left {
108                self.nodes.push(left);
109            }
110            if let Some(v) = node.value.take() {
111                return Some((std::mem::replace(&mut node.prefix, P::zero()), v));
112            }
113        }
114        None
115    }
116}
117
118/// An iterator over all prefixes of a [`PrefixMap`] in lexicographic order.
119#[derive(Clone)]
120pub struct IntoKeys<P, T> {
121    pub(super) inner: IntoIter<P, T>,
122}
123
124impl<P: Prefix, T> Iterator for IntoKeys<P, T> {
125    type Item = P;
126
127    fn next(&mut self) -> Option<P> {
128        self.inner.next().map(|(k, _)| k)
129    }
130}
131
132/// An iterator over all values of a [`PrefixMap`] in lexicographic order of their associated
133/// prefix.
134#[derive(Clone)]
135pub struct IntoValues<P, T> {
136    pub(super) inner: IntoIter<P, T>,
137}
138
139impl<P: Prefix, T> Iterator for IntoValues<P, T> {
140    type Item = T;
141
142    fn next(&mut self) -> Option<T> {
143        self.inner.next().map(|(_, v)| v)
144    }
145}
146
147impl<P: Prefix, T> IntoIterator for PrefixMap<P, T> {
148    type Item = (P, T);
149
150    type IntoIter = IntoIter<P, T>;
151
152    fn into_iter(self) -> Self::IntoIter {
153        IntoIter {
154            table: self.table.into_inner(),
155            nodes: vec![0],
156        }
157    }
158}
159
160impl<'a, P, T> IntoIterator for &'a PrefixMap<P, T> {
161    type Item = (&'a P, &'a T);
162
163    type IntoIter = Iter<'a, P, T>;
164
165    fn into_iter(self) -> Self::IntoIter {
166        // Safety: we own an immutable reference, and `Iter` will only ever read the table.
167        Iter::new(&self.table, vec![0])
168    }
169}
170
171/// A mutable iterator over a [`PrefixMap`]. This iterator yields elements in lexicographic order of
172/// their associated prefix.
173pub struct IterMut<'a, P, T> {
174    pub(super) table: Option<&'a Table<P, T>>,
175    pub(super) nodes: Vec<usize>,
176}
177
178impl<P, T> Default for IterMut<'_, P, T> {
179    fn default() -> Self {
180        Self {
181            table: None,
182            nodes: Vec::new(),
183        }
184    }
185}
186
187impl<'a, P, T> IterMut<'a, P, T> {
188    /// # Safety
189    /// - First, you must ensure that 'a is tied to a mutable reference of the original table.
190    /// - Second, you are allowed to create mutiple such `IterMut`s, as long as none of the root
191    ///   nodes is the parent of another root node (of any of the iterators). This also applies if
192    ///   you only create a single iterator with multiple root nodes.
193    ///
194    /// The iterator will only ever access its roots or its children.
195    pub(crate) unsafe fn new(table: &'a Table<P, T>, nodes: Vec<usize>) -> Self {
196        Self {
197            table: Some(table),
198            nodes,
199        }
200    }
201}
202
203impl<'a, P, T> Iterator for IterMut<'a, P, T> {
204    type Item = (&'a P, &'a mut T);
205
206    fn next(&mut self) -> Option<Self::Item> {
207        while let Some(cur) = self.nodes.pop() {
208            // Safety:
209            // In the following, we assume that there are no two iterators that may reach the same
210            // sub-tree (see the safety comment above).
211            //
212            // The iterator borrows from &'a mut PrefixMap, see `PrefixMap::iter_mut` where 'a is
213            // linked to a mutable reference. Then, we must ensure that we only ever construct a
214            // mutable reference to each element exactly once. We ensure this by the fact that we
215            // iterate over a tree. Thus, each node is visited exactly once.
216            let node: &'a mut Node<P, T> = unsafe { self.table.as_ref()?.get_mut(cur) };
217
218            if let Some(right) = node.right {
219                self.nodes.push(right);
220            }
221            if let Some(left) = node.left {
222                self.nodes.push(left);
223            }
224            if node.value.is_some() {
225                let v = node.value.as_mut().unwrap();
226                return Some((&node.prefix, v));
227            }
228        }
229        None
230    }
231}
232
233/// A mutable iterator over values of [`PrefixMap`]. This iterator yields elements in lexicographic
234/// order.
235#[derive(Default)]
236pub struct ValuesMut<'a, P, T> {
237    // # Safety
238    // You must ensure that there only ever exists one such iterator for each tree. You may create
239    // multiple such iterators for the same tree if you start with distinct starting nodes! This
240    // ensures that any one iteration will never yield elements of the other iterator.
241    pub(crate) inner: IterMut<'a, P, T>,
242}
243
244impl<'a, P, T> Iterator for ValuesMut<'a, P, T> {
245    type Item = &'a mut T;
246
247    fn next(&mut self) -> Option<Self::Item> {
248        self.inner.next().map(|(_, v)| v)
249    }
250}
251
252pub(super) fn lpm_children_iter_start<P: Prefix, T>(table: &Table<P, T>, prefix: &P) -> Vec<usize> {
253    let mut idx = 0;
254    let mut cur_p = &table[idx].prefix;
255
256    loop {
257        if cur_p.eq(prefix) {
258            break vec![idx];
259        }
260        let right = to_right(cur_p, prefix);
261        match table.get_child(idx, right) {
262            Some(c) => {
263                cur_p = &table[c].prefix;
264                if cur_p.contains(prefix) {
265                    // continue traversal
266                    idx = c;
267                } else if prefix.contains(cur_p) {
268                    break vec![c];
269                } else {
270                    break vec![];
271                }
272            }
273            None => break vec![],
274        }
275    }
276}
277
278impl<P, T> FromIterator<(P, T)> for PrefixMap<P, T>
279where
280    P: Prefix,
281{
282    fn from_iter<I: IntoIterator<Item = (P, T)>>(iter: I) -> Self {
283        let mut map = Self::new();
284        iter.into_iter().for_each(|(p, v)| {
285            map.insert(p, v);
286        });
287        map
288    }
289}
290
291/// An iterator that yields all items in a `PrefixMap` that covers a given prefix (including the
292/// prefix itself if preseint). See [`PrefixMap::cover`] for how to create this iterator.
293pub struct Cover<'a, 'p, P, T> {
294    pub(super) table: &'a Table<P, T>,
295    pub(super) idx: Option<usize>,
296    pub(super) prefix: &'p P,
297}
298
299impl<'a, P, T> Iterator for Cover<'a, '_, P, T>
300where
301    P: Prefix,
302{
303    type Item = (&'a P, &'a T);
304
305    fn next(&mut self) -> Option<Self::Item> {
306        // check if self.idx is None. If so, then check if the first branch is present in the map
307        if self.idx.is_none() {
308            self.idx = Some(0);
309            let entry = &self.table[0];
310            if let Some(value) = entry.value.as_ref() {
311                return Some((&entry.prefix, value));
312            }
313        }
314
315        // if we reach here, then self.idx is not None!
316
317        loop {
318            let map::Direction::Enter { next, .. } =
319                self.table.get_direction(self.idx.unwrap(), self.prefix)
320            else {
321                return None;
322            };
323            self.idx = Some(next);
324            let entry = &self.table[next];
325            if let Some(value) = entry.value.as_ref() {
326                return Some((&entry.prefix, value));
327            }
328        }
329    }
330}
331
332/// An iterator that yields all keys (prefixes) in a `PrefixMap` that covers a given prefix
333/// (including the prefix itself if preseint). See [`PrefixMap::cover_keys`] for how to create this
334/// iterator.
335pub struct CoverKeys<'a, 'p, P, T>(pub(super) Cover<'a, 'p, P, T>);
336
337impl<'a, P, T> Iterator for CoverKeys<'a, '_, P, T>
338where
339    P: Prefix,
340{
341    type Item = &'a P;
342
343    fn next(&mut self) -> Option<Self::Item> {
344        self.0.next().map(|(p, _)| p)
345    }
346}
347
348/// An iterator that yields all values in a `PrefixMap` that covers a given prefix (including the
349/// prefix itself if preseint). See [`PrefixMap::cover_values`] for how to create this iterator.
350pub struct CoverValues<'a, 'p, P, T>(pub(super) Cover<'a, 'p, P, T>);
351
352impl<'a, P, T> Iterator for CoverValues<'a, '_, P, T>
353where
354    P: Prefix,
355{
356    type Item = &'a T;
357
358    fn next(&mut self) -> Option<Self::Item> {
359        self.0.next().map(|(_, t)| t)
360    }
361}
prefix_trie/map/iter.rs

prefix_trie/map/
iter.rs
