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.get());
50            }
51            if let Some(left) = node.left {
52                self.nodes.push(left.get());
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.get());
106            }
107            if let Some(left) = node.left {
108                self.nodes.push(left.get());
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.get());
220            }
221            if let Some(left) = node.left {
222                self.nodes.push(left.get());
223            }
224            if let Some(v) = &mut node.value {
225                return Some((&node.prefix, v));
226            }
227        }
228        None
229    }
230}
231
232/// A mutable iterator over values of [`PrefixMap`]. This iterator yields elements in lexicographic
233/// order.
234#[derive(Default)]
235pub struct ValuesMut<'a, P, T> {
236    // # Safety
237    // You must ensure that there only ever exists one such iterator for each tree. You may create
238    // multiple such iterators for the same tree if you start with distinct starting nodes! This
239    // ensures that any one iteration will never yield elements of the other iterator.
240    pub(crate) inner: IterMut<'a, P, T>,
241}
242
243impl<'a, P, T> Iterator for ValuesMut<'a, P, T> {
244    type Item = &'a mut T;
245
246    fn next(&mut self) -> Option<Self::Item> {
247        self.inner.next().map(|(_, v)| v)
248    }
249}
250
251pub(super) fn lpm_children_iter_start<P: Prefix, T>(table: &Table<P, T>, prefix: &P) -> Vec<usize> {
252    let mut idx = 0;
253    let mut cur_p = &table[idx].prefix;
254
255    loop {
256        if cur_p.eq(prefix) {
257            break vec![idx];
258        }
259        let right = to_right(cur_p, prefix);
260        match table.get_child(idx, right).map(|x| x.get()) {
261            Some(c) => {
262                cur_p = &table[c].prefix;
263                if cur_p.contains(prefix) {
264                    // continue traversal
265                    idx = c;
266                } else if prefix.contains(cur_p) {
267                    break vec![c];
268                } else {
269                    break vec![];
270                }
271            }
272            None => break vec![],
273        }
274    }
275}
276
277impl<P, T> FromIterator<(P, T)> for PrefixMap<P, T>
278where
279    P: Prefix,
280{
281    fn from_iter<I: IntoIterator<Item = (P, T)>>(iter: I) -> Self {
282        let mut map = Self::new();
283        iter.into_iter().for_each(|(p, v)| {
284            map.insert(p, v);
285        });
286        map
287    }
288}
289
290/// An iterator that yields all items in a `PrefixMap` that covers a given prefix (including the
291/// prefix itself if preseint). See [`PrefixMap::cover`] for how to create this iterator.
292pub struct Cover<'a, 'p, P, T> {
293    pub(super) table: &'a Table<P, T>,
294    pub(super) idx: Option<usize>,
295    pub(super) prefix: &'p P,
296}
297
298impl<'a, P, T> Iterator for Cover<'a, '_, P, T>
299where
300    P: Prefix,
301{
302    type Item = (&'a P, &'a T);
303
304    fn next(&mut self) -> Option<Self::Item> {
305        // check if self.idx is None. If so, then check if the first branch is present in the map
306        if self.idx.is_none() {
307            self.idx = Some(0);
308            let entry = &self.table[0usize];
309            if let Some(value) = entry.value.as_ref() {
310                return Some((&entry.prefix, value));
311            }
312        }
313
314        // if we reach here, then self.idx is not None!
315
316        loop {
317            let map::Direction::Enter { next, .. } =
318                self.table.get_direction(self.idx.unwrap(), self.prefix)
319            else {
320                return None;
321            };
322            self.idx = Some(next.get());
323            let entry = &self.table[next.get()];
324            if let Some(value) = entry.value.as_ref() {
325                return Some((&entry.prefix, value));
326            }
327        }
328    }
329}
330
331/// An iterator that yields all keys (prefixes) in a `PrefixMap` that covers a given prefix
332/// (including the prefix itself if preseint). See [`PrefixMap::cover_keys`] for how to create this
333/// iterator.
334pub struct CoverKeys<'a, 'p, P, T>(pub(super) Cover<'a, 'p, P, T>);
335
336impl<'a, P, T> Iterator for CoverKeys<'a, '_, P, T>
337where
338    P: Prefix,
339{
340    type Item = &'a P;
341
342    fn next(&mut self) -> Option<Self::Item> {
343        self.0.next().map(|(p, _)| p)
344    }
345}
346
347/// An iterator that yields all values in a `PrefixMap` that covers a given prefix (including the
348/// prefix itself if preseint). See [`PrefixMap::cover_values`] for how to create this iterator.
349pub struct CoverValues<'a, 'p, P, T>(pub(super) Cover<'a, 'p, P, T>);
350
351impl<'a, P, T> Iterator for CoverValues<'a, '_, P, T>
352where
353    P: Prefix,
354{
355    type Item = &'a T;
356
357    fn next(&mut self) -> Option<Self::Item> {
358        self.0.next().map(|(_, t)| t)
359    }
360}
prefix_trie/map/iter.rs

prefix_trie/map/
iter.rs
