prefix_trie/map/
iter.rs

1//! Module that contains the implementation for the iterators
2
3use num_traits::Zero;
4
5use crate::{
6    Prefix,
7    {
8        allocator::{Loc, RawPtr},
9        node::{child_bit, LexIterElem, MaskedLexIter},
10        table::{DataIdx, Table, K},
11        PrefixMap,
12    },
13};
14
15/// An iterator over all entries of a [`PrefixMap`] in lexicographic order.
16pub struct Iter<'a, P: Prefix, T> {
17    table: Option<&'a Table<T>>,
18    stack: Vec<MaskedLexIter<P::R>>,
19}
20
21impl<P: Prefix, T> Clone for Iter<'_, P, T> {
22    fn clone(&self) -> Self {
23        Self {
24            table: self.table,
25            stack: self.stack.clone(),
26        }
27    }
28}
29
30impl<P: Prefix, T> Default for Iter<'_, P, T> {
31    fn default() -> Self {
32        Self {
33            table: None,
34            stack: Vec::new(),
35        }
36    }
37}
38
39impl<'a, P: Prefix, T> Iter<'a, P, T> {
40    /// Create a new iterator that iterates over all
41    pub(crate) fn new(table: &'a Table<T>) -> Self {
42        // SAFETY: `Loc::root()` is always valid. The `&'a Table<T>` borrow prevents
43        // structural mutations for the lifetime of the iterator, keeping all Locs valid.
44        let lex = unsafe { table.lex_iter(Loc::root(), 0, <P::R as Zero>::zero()) };
45        Self::at_node(table, lex)
46    }
47
48    pub(crate) fn at_node(table: &'a Table<T>, lex: MaskedLexIter<P::R>) -> Self {
49        let stack = vec![lex];
50        Self {
51            table: Some(table),
52            stack,
53        }
54    }
55
56    pub(crate) fn from_stack(table: &'a Table<T>, stack: Vec<MaskedLexIter<P::R>>) -> Self {
57        Self {
58            table: Some(table),
59            stack,
60        }
61    }
62}
63
64impl<'a, P: Prefix, T> Iterator for Iter<'a, P, T> {
65    type Item = (P, &'a T);
66
67    fn next(&mut self) -> Option<(P, &'a T)> {
68        let table = self.table?;
69        while let Some(lex_iter) = self.stack.last_mut() {
70            let key = *lex_iter.key();
71            let Some(next) = lex_iter.next() else {
72                self.stack.pop();
73                continue;
74            };
75
76            match next {
77                LexIterElem::Data(idx) => {
78                    // SAFETY: `Iter` holds `&'a Table`; no structural changes occur during
79                    // immutable iteration, so `idx.node` remains valid.
80                    let Some(r) = (unsafe { idx.resolve(table) }) else {
81                        continue;
82                    };
83                    let p = r.prefix(key);
84                    return Some((p, r.get()));
85                }
86                LexIterElem::Child(next_loc, depth, next_key) => {
87                    // SAFETY: `next_loc` came from the previous `lex_iter` step; the borrow
88                    // of `table` prevents structural mutations, so `next_loc` remains valid.
89                    self.stack
90                        .push(unsafe { table.lex_iter(next_loc, depth, next_key) })
91                }
92            }
93        }
94        None
95    }
96}
97
98/// An iterator over all prefixes of a [`PrefixMap`] in lexicographic order.
99#[derive(Clone, Default)]
100pub struct Keys<'a, P: Prefix, T>(pub(crate) Iter<'a, P, T>);
101
102impl<'a, P: Prefix, T> Iterator for Keys<'a, P, T> {
103    type Item = P;
104
105    fn next(&mut self) -> Option<P> {
106        self.0.next().map(|(k, _)| k)
107    }
108}
109
110/// An iterator over all values of a [`PrefixMap`] in lexicographic order of their associated
111/// prefixes.
112#[derive(Clone, Default)]
113pub struct Values<'a, P: Prefix, T>(pub(crate) Iter<'a, P, T>);
114
115impl<'a, P: Prefix, T> Iterator for Values<'a, P, T> {
116    type Item = &'a T;
117
118    fn next(&mut self) -> Option<&'a T> {
119        self.0.next().map(|(_, v)| v)
120    }
121}
122
123/// A mutable iterator over a [`PrefixMap`]. This iterator yields elements in lexicographic order of
124/// their associated prefix.
125pub struct IterMut<'a, P: Prefix, T> {
126    // SAFETY
127    // 1. This structure must be created using a mutable reference to the table (see Self::new).
128    // 2. The tree cannot have cycles (ensured by the way we create the tree.)
129    table: Option<(&'a Table<T>, RawPtr<T>)>,
130    stack: Vec<MaskedLexIter<P::R>>,
131}
132
133impl<P: Prefix, T> Default for IterMut<'_, P, T> {
134    fn default() -> Self {
135        Self {
136            table: None,
137            stack: Vec::new(),
138        }
139    }
140}
141
142impl<'a, P: Prefix, T> IterMut<'a, P, T> {
143    /// Create a new iterator that iterates over all
144    pub(crate) fn new(table: &'a mut Table<T>) -> Self {
145        // SAFETY: `Loc::root()` is always valid. The `&'a mut Table<T>` borrow prevents
146        // structural mutations for the lifetime of the iterator, keeping all Locs valid.
147        let lex = unsafe { table.lex_iter(Loc::root(), 0, <P::R as Zero>::zero()) };
148        Self::at_node(table, lex)
149    }
150
151    pub(crate) fn at_node(table: &'a mut Table<T>, lex: MaskedLexIter<P::R>) -> Self {
152        let ptr = table.raw_cells();
153        Self {
154            table: Some((table, ptr)),
155            stack: vec![lex],
156        }
157    }
158
159    pub(crate) fn from_stack(table: &'a mut Table<T>, stack: Vec<MaskedLexIter<P::R>>) -> Self {
160        let ptr = table.raw_cells();
161        Self {
162            table: Some((table, ptr)),
163            stack,
164        }
165    }
166}
167
168impl<'a, P: Prefix, T> Iterator for IterMut<'a, P, T> {
169    type Item = (P, &'a mut T);
170
171    fn next(&mut self) -> Option<Self::Item> {
172        let (table, mut ptr) = self.table?;
173        while let Some(lex_iter) = self.stack.last_mut() {
174            let key = *lex_iter.key();
175            let Some(next) = lex_iter.next() else {
176                self.stack.pop();
177                continue;
178            };
179
180            match next {
181                LexIterElem::Data(idx) => {
182                    // SAFETY: `IterMut` holds exclusive access to the table data; no structural
183                    // changes occur, so `idx.node` remains valid.
184                    let Some(r) = (unsafe { idx.resolve(table) }) else {
185                        continue;
186                    };
187                    let p = r.prefix(key);
188                    // SAFETY: `IterMut` was created from `&'a mut Table`; ptr was obtained
189                    // from that same table. The tree is acyclic so each node is visited at most
190                    // once, guaranteeing no two live `&mut T` references to the same slot.
191                    return Some((p, unsafe { r.unsafe_get_mut(&mut ptr) }));
192                }
193                LexIterElem::Child(next_idx, depth, next_key) => {
194                    // SAFETY: `next_idx` came from the previous `lex_iter` step; the exclusive
195                    // borrow of `table` prevents structural mutations, so `next_idx` remains valid.
196                    self.stack
197                        .push(unsafe { table.lex_iter(next_idx, depth, next_key) })
198                }
199            }
200        }
201        None
202    }
203}
204
205/// A mutable iterator over values of [`PrefixMap`]. This iterator yields elements in lexicographic
206/// order.
207#[derive(Default)]
208pub struct ValuesMut<'a, P: Prefix, T>(pub(crate) IterMut<'a, P, T>);
209
210impl<'a, P: Prefix, T> Iterator for ValuesMut<'a, P, T> {
211    type Item = &'a mut T;
212
213    fn next(&mut self) -> Option<Self::Item> {
214        self.0.next().map(|(_, v)| v)
215    }
216}
217
218/// An iterator over all owned entries of a [`PrefixMap`] in lexicographic order.
219#[derive(Clone)]
220pub struct IntoIter<P: Prefix, T> {
221    table: Table<T>,
222    stack: Vec<MaskedLexIter<P::R>>,
223}
224
225impl<P: Prefix, T> IntoIter<P, T> {
226    /// Create a new iterator that iterates over all
227    pub(crate) fn new(table: Table<T>) -> Self {
228        // SAFETY: `Loc::root()` is always valid. `IntoIter` owns the table, preventing
229        // any external structural mutations for its lifetime.
230        let lex = unsafe { table.lex_iter(Loc::root(), 0, <P::R as Zero>::zero()) };
231        Self::at_node(table, lex)
232    }
233
234    pub(crate) fn at_node(table: Table<T>, lex: MaskedLexIter<P::R>) -> Self {
235        let stack = vec![lex];
236        Self { table, stack }
237    }
238}
239
240impl<P: Prefix, T> Iterator for IntoIter<P, T> {
241    type Item = (P, T);
242
243    fn next(&mut self) -> Option<(P, T)> {
244        while let Some(lex_iter) = self.stack.last_mut() {
245            let key = *lex_iter.key();
246            let Some(next) = lex_iter.next() else {
247                self.stack.pop();
248                continue;
249            };
250
251            match next {
252                LexIterElem::Data(idx) => {
253                    // SAFETY: IntoIter owns the table; no external mutations possible.
254                    // `take()` clears the bitmap bit, so the MaskedLexIter snapshot may
255                    // yield already-taken bits; `resolve_mut` re-checks and returns None.
256                    let Some(r) = (unsafe { idx.resolve_mut(&mut self.table) }) else {
257                        continue;
258                    };
259                    let p = r.prefix(key);
260                    return Some((p, r.take()));
261                }
262                LexIterElem::Child(next_loc, depth, next_key) => {
263                    // SAFETY: `IntoIter` owns the table and makes no structural mutations;
264                    // `next_loc` came from the previous `lex_iter` step and remains valid.
265                    self.stack
266                        .push(unsafe { self.table.lex_iter(next_loc, depth, next_key) })
267                }
268            }
269        }
270        None
271    }
272}
273
274/// An iterator over all prefixes of a [`PrefixMap`] in lexicographic order.
275#[derive(Clone)]
276pub struct IntoKeys<P: Prefix, T>(pub(super) IntoIter<P, T>);
277
278impl<P: Prefix, T> Iterator for IntoKeys<P, T> {
279    type Item = P;
280
281    fn next(&mut self) -> Option<P> {
282        self.0.next().map(|(k, _)| k)
283    }
284}
285
286/// An iterator over all values of a [`PrefixMap`] in lexicographic order of their associated
287/// prefix.
288#[derive(Clone)]
289pub struct IntoValues<P: Prefix, T>(pub(super) IntoIter<P, T>);
290
291impl<P: Prefix, T> Iterator for IntoValues<P, T> {
292    type Item = T;
293
294    fn next(&mut self) -> Option<T> {
295        self.0.next().map(|(_, v)| v)
296    }
297}
298
299impl<P: Prefix, T> IntoIterator for PrefixMap<P, T> {
300    type Item = (P, T);
301
302    type IntoIter = IntoIter<P, T>;
303
304    fn into_iter(self) -> Self::IntoIter {
305        IntoIter::new(self.table)
306    }
307}
308
309impl<'a, P: Prefix, T> IntoIterator for &'a PrefixMap<P, T> {
310    type Item = (P, &'a T);
311
312    type IntoIter = Iter<'a, P, T>;
313
314    fn into_iter(self) -> Self::IntoIter {
315        Iter::new(&self.table)
316    }
317}
318
319impl<P, T> FromIterator<(P, T)> for PrefixMap<P, T>
320where
321    P: Prefix,
322{
323    fn from_iter<I: IntoIterator<Item = (P, T)>>(iter: I) -> Self {
324        let mut map = Self::new();
325        iter.into_iter().for_each(|(p, v)| {
326            map.insert(p, v);
327        });
328        map
329    }
330}
331
332/// An iterator that yields all items in a `PrefixMap` that cover (are a superset of) a given
333/// prefix (including the prefix itself if present). See [`PrefixMap::cover`] for how to
334/// create this iterator.
335pub struct Cover<'a, P: Prefix, T> {
336    table: &'a Table<T>,
337    next_loc: Option<Loc>,
338    lpm_elements: Vec<DataIdx>,
339    key: P::R,
340    prefix_len: u32,
341    depth: u32,
342}
343
344impl<'a, P: Prefix, T> Cover<'a, P, T> {
345    pub(crate) fn new(table: &'a PrefixMap<P, T>, prefix: &P) -> Self {
346        let key = prefix.repr();
347        let prefix_len = prefix.prefix_len() as u32;
348        let lpm_elements = Vec::new();
349        let next_loc = Some(Loc::root());
350        Self {
351            table: &table.table,
352            next_loc,
353            lpm_elements,
354            key,
355            prefix_len,
356            depth: 0,
357        }
358    }
359}
360
361impl<'a, P: Prefix, T> Iterator for Cover<'a, P, T>
362where
363    P: Prefix,
364{
365    type Item = (P, &'a T);
366
367    fn next(&mut self) -> Option<Self::Item> {
368        loop {
369            let Some(loc) = self.lpm_elements.pop() else {
370                // fill the lpm_nodes
371                let loc = self.next_loc?; // if none, then exit.
372                                          // SAFETY: `Cover` holds `&'a Table<T>` — no structural mutations are possible
373                                          // for the lifetime `'a`. `loc` starts as `Loc::root()` and is only updated to
374                                          // the result of a prior `child()` call, so it is always a valid, live node.
375                self.lpm_elements = unsafe {
376                    self.table
377                        .data_ancestors(loc, self.depth, self.key, self.prefix_len)
378                }
379                .collect();
380                self.lpm_elements.reverse();
381                let prefix_in_this_node = self.prefix_len < self.depth + K;
382                self.next_loc = if prefix_in_this_node {
383                    None
384                } else {
385                    let child_bit = child_bit(self.depth, self.key);
386                    // SAFETY: same invariant as above.
387                    unsafe { self.table.child(loc, child_bit) }
388                };
389                self.depth += K;
390                continue;
391            };
392
393            // SAFETY: `loc` came from `data_ancestors`, which only yields bits set in the
394            // bitmap. `Cover` holds `&'a Table<T>` and never mutates node structure.
395            let r = unsafe { loc.resolve(self.table) }
396                .expect("Cover: data_ancestors yielded stale idx");
397            return Some((r.prefix(self.key), r.get()));
398        }
399    }
400}
401
402/// An iterator that yields all keys (prefixes) in a `PrefixMap` that covers a given prefix
403/// (including the prefix itself if present). See [`PrefixMap::cover_keys`] for how to create this
404/// iterator.
405pub struct CoverKeys<'a, P: Prefix, T>(pub(super) Cover<'a, P, T>);
406
407impl<'a, P: Prefix, T> Iterator for CoverKeys<'a, P, T>
408where
409    P: Prefix,
410{
411    type Item = P;
412
413    fn next(&mut self) -> Option<Self::Item> {
414        self.0.next().map(|(p, _)| p)
415    }
416}
417
418/// An iterator that yields all values in a `PrefixMap` that covers a given prefix (including the
419/// prefix itself if present). See [`PrefixMap::cover_values`] for how to create this iterator.
420pub struct CoverValues<'a, P: Prefix, T>(pub(super) Cover<'a, P, T>);
421
422impl<'a, P: Prefix, T> Iterator for CoverValues<'a, P, T>
423where
424    P: Prefix,
425{
426    type Item = &'a T;
427
428    fn next(&mut self) -> Option<Self::Item> {
429        self.0.next().map(|(_, t)| t)
430    }
431}
432
433/// Create a lex iterator that covers the given prefix.
434pub(crate) fn lpm_children_iter_start<P: Prefix, T>(
435    table: &Table<T>,
436    prefix: &P,
437) -> MaskedLexIter<P::R> {
438    let key = prefix.repr();
439    let prefix_len = prefix.prefix_len() as u32;
440    table.lex_iter_at(key, prefix_len)
441}
prefix_trie/map/iter.rs

prefix_trie/map/
iter.rs
