champ_trie/
iter.rs

1//! Iterator types for CHAMP maps.
2
3use safe_bump::Idx;
4
5use crate::node::{self, Entry, Node};
6use crate::store::ChampStore;
7
8/// Iterator over references to key-value pairs in a [`ChampMap`](crate::ChampMap).
9pub struct Iter<'a, K, V> {
10    entries: Vec<(&'a K, &'a V)>,
11    pos: usize,
12}
13
14impl<'a, K, V> Iter<'a, K, V> {
15    /// Creates an iterator by collecting all live entries via DFS.
16    pub fn new<S: ChampStore<K, V>>(store: &'a S, root: Option<Idx<Node<K, V>>>) -> Self {
17        let mut entries = Vec::new();
18        if let Some(idx) = root {
19            collect(store, idx, &mut entries);
20        }
21        Self { entries, pos: 0 }
22    }
23}
24
25impl<'a, K, V> Iterator for Iter<'a, K, V> {
26    type Item = (&'a K, &'a V);
27
28    fn next(&mut self) -> Option<Self::Item> {
29        if self.pos < self.entries.len() {
30            let item = self.entries[self.pos];
31            self.pos += 1;
32            Some(item)
33        } else {
34            None
35        }
36    }
37
38    fn size_hint(&self) -> (usize, Option<usize>) {
39        let remaining = self.entries.len() - self.pos;
40        (remaining, Some(remaining))
41    }
42}
43
44impl<K, V> ExactSizeIterator for Iter<'_, K, V> {}
45
46/// DFS collect all `(&K, &V)` from the subtree rooted at `node_idx`.
47fn collect<'a, K, V, S: ChampStore<K, V>>(
48    store: &'a S,
49    node_idx: Idx<Node<K, V>>,
50    out: &mut Vec<(&'a K, &'a V)>,
51) {
52    match *store.get_node(node_idx) {
53        Node::Inner {
54            data_map,
55            node_map,
56            data_start,
57            children_start,
58            ..
59        } => {
60            let data_len = data_map.count_ones() as usize;
61            let children_len = node_map.count_ones() as usize;
62
63            for i in 0..data_len {
64                let e: &'a Entry<K, V> = store.get_entry(node::offset(data_start, i));
65                out.push((&e.key, &e.value));
66            }
67
68            for i in 0..children_len {
69                let child = *store.get_child(node::offset(children_start, i));
70                collect(store, child, out);
71            }
72        }
73        Node::Collision {
74            entries_start,
75            entries_len,
76            ..
77        } => {
78            for i in 0..usize::from(entries_len) {
79                let e: &'a Entry<K, V> = store.get_entry(node::offset(entries_start, i));
80                out.push((&e.key, &e.value));
81            }
82        }
83    }
84}
champ_trie/iter.rs

champ_trie/
iter.rs
