scapegoat/tree/
iter.rs

1use core::iter::FusedIterator;
2
3use tinyvec::ArrayVec;
4
5use super::node::Node;
6use super::node_dispatch::SmallNode;
7use super::tree::{Idx, SgTree};
8
9// Immutable Reference Iterator ----------------------------------------------------------------------------------------
10
11/// Uses iterative in-order tree traversal algorithm.
12/// Maintains a small stack of arena indexes (won't contain all indexes simultaneously for a balanced tree).
13pub struct Iter<'a, K: Default, V: Default, const N: usize> {
14    bst: &'a SgTree<K, V, N>,
15    idx_stack: ArrayVec<[usize; N]>,
16    total_cnt: usize,
17    spent_cnt: usize,
18}
19
20impl<'a, K: Ord + Default, V: Default, const N: usize> Iter<'a, K, V, N> {
21    pub fn new(bst: &'a SgTree<K, V, N>) -> Self {
22        let mut ordered_iter = Iter {
23            bst,
24            idx_stack: ArrayVec::<[usize; N]>::new(),
25            total_cnt: bst.len(),
26            spent_cnt: 0,
27        };
28
29        if let Some(root_idx) = ordered_iter.bst.opt_root_idx {
30            let mut curr_idx = root_idx;
31            loop {
32                let node = &ordered_iter.bst.arena[curr_idx];
33                match node.left_idx() {
34                    Some(lt_idx) => {
35                        ordered_iter.idx_stack.push(curr_idx);
36                        curr_idx = lt_idx;
37                    }
38                    None => {
39                        ordered_iter.idx_stack.push(curr_idx);
40                        break;
41                    }
42                }
43            }
44        }
45
46        ordered_iter
47    }
48}
49
50impl<'a, K: Ord + Default, V: Default, const N: usize> Iterator for Iter<'a, K, V, N> {
51    type Item = (&'a K, &'a V);
52
53    fn next(&mut self) -> Option<Self::Item> {
54        match self.idx_stack.pop() {
55            Some(pop_idx) => {
56                let node = &self.bst.arena[pop_idx];
57                if let Some(gt_idx) = node.right_idx() {
58                    let mut curr_idx = gt_idx;
59                    loop {
60                        let node = &self.bst.arena[curr_idx];
61                        match node.left_idx() {
62                            Some(lt_idx) => {
63                                self.idx_stack.push(curr_idx);
64                                curr_idx = lt_idx;
65                            }
66                            None => {
67                                self.idx_stack.push(curr_idx);
68                                break;
69                            }
70                        }
71                    }
72                }
73
74                let node = &self.bst.arena[pop_idx];
75                self.spent_cnt += 1;
76                Some((node.key(), node.val()))
77            }
78            None => None,
79        }
80    }
81}
82
83impl<'a, K: Ord + Default, V: Default, const N: usize> ExactSizeIterator for Iter<'a, K, V, N> {
84    fn len(&self) -> usize {
85        debug_assert!(self.spent_cnt <= self.total_cnt);
86        self.total_cnt - self.spent_cnt
87    }
88}
89
90impl<'a, K: Ord + Default, V: Default, const N: usize> FusedIterator for Iter<'a, K, V, N> {}
91
92// Mutable Reference Iterator ------------------------------------------------------------------------------------------
93
94pub struct IterMut<'a, K, V, const N: usize> {
95    arena_iter_mut: core::slice::IterMut<'a, Option<Node<K, V, Idx>>>,
96}
97
98impl<'a, K: Ord + Default, V: Default, const N: usize> IterMut<'a, K, V, N> {
99    pub fn new(bst: &'a mut SgTree<K, V, N>) -> Self {
100        bst.sort_arena();
101        IterMut {
102            arena_iter_mut: bst.arena.iter_mut(),
103        }
104    }
105}
106
107impl<'a, K: Ord + Default, V: Default, const N: usize> Iterator for IterMut<'a, K, V, N> {
108    type Item = (&'a K, &'a mut V);
109
110    fn next(&mut self) -> Option<Self::Item> {
111        match self.arena_iter_mut.next() {
112            Some(Some(node)) => Some(node.get_mut()),
113            _ => None,
114        }
115    }
116}
117
118impl<'a, K: Ord + Default, V: Default, const N: usize> DoubleEndedIterator
119    for IterMut<'a, K, V, N>
120{
121    fn next_back(&mut self) -> Option<Self::Item> {
122        match self.arena_iter_mut.next_back() {
123            Some(Some(node)) => Some(node.get_mut()),
124            _ => None,
125        }
126    }
127}
128
129impl<'a, K: Ord + Default, V: Default, const N: usize> ExactSizeIterator for IterMut<'a, K, V, N> {
130    fn len(&self) -> usize {
131        self.arena_iter_mut.len()
132    }
133}
134
135impl<'a, K: Ord + Default, V: Default, const N: usize> FusedIterator for IterMut<'a, K, V, N> {}
136
137// Consuming Iterator --------------------------------------------------------------------------------------------------
138
139/// Cheats a little by using internal flattening logic to sort, instead of re-implementing proper traversal.
140/// Maintains a shrinking list of arena indexes, initialized with all of them.
141pub struct IntoIter<K: Default, V: Default, const N: usize> {
142    bst: SgTree<K, V, N>,
143    sorted_idxs: ArrayVec<[usize; N]>,
144}
145
146impl<K: Ord + Default, V: Default, const N: usize> IntoIter<K, V, N> {
147    pub fn new(bst: SgTree<K, V, N>) -> Self {
148        let mut ordered_iter = IntoIter {
149            bst,
150            sorted_idxs: ArrayVec::<[usize; N]>::new(),
151        };
152
153        if let Some(root_idx) = ordered_iter.bst.opt_root_idx {
154            ordered_iter.sorted_idxs = ordered_iter.bst.flatten_subtree_to_sorted_idxs(root_idx);
155            ordered_iter.sorted_idxs.reverse();
156        }
157
158        ordered_iter
159    }
160}
161
162impl<K: Ord + Default, V: Default, const N: usize> Iterator for IntoIter<K, V, N> {
163    type Item = (K, V);
164
165    fn next(&mut self) -> Option<Self::Item> {
166        match self.sorted_idxs.pop() {
167            Some(idx) => match self.bst.priv_remove_by_idx(idx) {
168                Some((key, val)) => Some((key, val)),
169                None => {
170                    debug_assert!(false, "Use of invalid index in consuming iterator!");
171                    None
172                }
173            },
174            None => None,
175        }
176    }
177}
178
179impl<K: Ord + Default, V: Default, const N: usize> ExactSizeIterator for IntoIter<K, V, N> {
180    fn len(&self) -> usize {
181        self.sorted_idxs.len()
182    }
183}
184
185impl<K: Ord + Default, V: Default, const N: usize> FusedIterator for IntoIter<K, V, N> {}
scapegoat/tree/iter.rs

scapegoat/tree/
iter.rs
