slice_rbtree/forest/iterators/
mod.rs

1//! Iterators over [`RBTree`](crate::tree::RBTree) and [`RBForest`](crate::forest::RBForest)
2use borsh::{BorshDeserialize, BorshSerialize};
3use core::cmp::Ord;
4use core::fmt;
5use core::iter::FusedIterator;
6
7use super::RBForest;
8
9/// An iterator over key-value pairs ordered by key
10pub struct PairsIterator<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize>
11where
12    K: Ord + BorshDeserialize + BorshSerialize,
13    V: BorshDeserialize + BorshSerialize,
14{
15    next_node: Option<usize>,
16    tree: &'a RBForest<'b, K, V, KSIZE, VSIZE>,
17}
18impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> PairsIterator<'a, 'b, K, V, KSIZE, VSIZE>
19where
20    K: Ord + BorshDeserialize + BorshSerialize,
21    V: BorshDeserialize + BorshSerialize,
22{
23    pub(super) fn from_raw_parts(
24        tree: &'a RBForest<'b, K, V, KSIZE, VSIZE>,
25        next_node: Option<usize>,
26    ) -> Self {
27        Self { next_node, tree }
28    }
29}
30
31impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> Iterator
32    for PairsIterator<'a, 'b, K, V, KSIZE, VSIZE>
33where
34    K: Ord + BorshDeserialize + BorshSerialize,
35    V: BorshDeserialize + BorshSerialize,
36{
37    type Item = (K, V);
38
39    fn next(&mut self) -> Option<Self::Item> {
40        self.next_node.map(|mut id| {
41            let nodes = &self.tree.nodes;
42
43            let key = K::deserialize(&mut nodes[id].key.as_slice()).expect("Key corrupted");
44            let value = V::deserialize(&mut nodes[id].value.as_slice()).expect("Value corrupted");
45
46            // find next
47            if let Some(right_id) = nodes[id].right() {
48                self.next_node = Some(self.tree.min(right_id as usize));
49            } else {
50                self.next_node = None;
51                while let Some(parent_id) = nodes[id].parent() {
52                    let parent_id = parent_id as usize;
53                    if Some(id as u32) == nodes[parent_id].left() {
54                        self.next_node = Some(parent_id);
55                        break;
56                    } else {
57                        id = parent_id;
58                    }
59                }
60            }
61
62            (key, value)
63        })
64    }
65}
66
67impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> FusedIterator
68    for PairsIterator<'a, 'b, K, V, KSIZE, VSIZE>
69where
70    K: Ord + BorshDeserialize + BorshSerialize,
71    V: BorshDeserialize + BorshSerialize,
72{
73}
74
75impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> fmt::Debug
76    for PairsIterator<'a, 'b, K, V, KSIZE, VSIZE>
77where
78    K: Ord + BorshDeserialize + BorshSerialize + fmt::Debug,
79    V: BorshDeserialize + BorshSerialize + fmt::Debug,
80{
81    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
82        let PairsIterator { next_node, tree } = self;
83        let new_iter = PairsIterator {
84            next_node: *next_node,
85            tree,
86        };
87        f.debug_map().entries(new_iter).finish()
88    }
89}
90
91/// An ordered iterator over keys
92pub struct KeysIterator<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize>
93where
94    K: Ord + BorshDeserialize + BorshSerialize,
95    V: BorshDeserialize + BorshSerialize,
96{
97    next_node: Option<usize>,
98    tree: &'a RBForest<'b, K, V, KSIZE, VSIZE>,
99}
100
101impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> KeysIterator<'a, 'b, K, V, KSIZE, VSIZE>
102where
103    K: Ord + BorshDeserialize + BorshSerialize,
104    V: BorshDeserialize + BorshSerialize,
105{
106    pub(super) fn from_raw_parts(
107        tree: &'a RBForest<'b, K, V, KSIZE, VSIZE>,
108        next_node: Option<usize>,
109    ) -> Self {
110        Self { next_node, tree }
111    }
112}
113
114impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> Iterator
115    for KeysIterator<'a, 'b, K, V, KSIZE, VSIZE>
116where
117    K: Ord + BorshDeserialize + BorshSerialize,
118    V: BorshDeserialize + BorshSerialize,
119{
120    type Item = K;
121
122    fn next(&mut self) -> Option<Self::Item> {
123        self.next_node.map(|mut id| {
124            let nodes = &self.tree.nodes;
125
126            let key = K::deserialize(&mut nodes[id].key.as_slice()).expect("Key corrupted");
127
128            // find next
129            if let Some(right_id) = nodes[id].right() {
130                self.next_node = Some(self.tree.min(right_id as usize));
131            } else {
132                self.next_node = None;
133                while let Some(parent_id) = nodes[id].parent() {
134                    let parent_id = parent_id as usize;
135                    if Some(id as u32) == nodes[parent_id].left() {
136                        self.next_node = Some(parent_id);
137                        break;
138                    } else {
139                        id = parent_id;
140                    }
141                }
142            }
143
144            key
145        })
146    }
147}
148
149impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> FusedIterator
150    for KeysIterator<'a, 'b, K, V, KSIZE, VSIZE>
151where
152    K: Ord + BorshDeserialize + BorshSerialize,
153    V: BorshDeserialize + BorshSerialize,
154{
155}
156
157impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> fmt::Debug
158    for KeysIterator<'a, 'b, K, V, KSIZE, VSIZE>
159where
160    K: Ord + BorshDeserialize + BorshSerialize + fmt::Debug,
161    V: BorshDeserialize + BorshSerialize + fmt::Debug,
162{
163    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
164        let KeysIterator { next_node, tree } = self;
165        let new_iter = KeysIterator {
166            next_node: *next_node,
167            tree,
168        };
169        f.debug_set().entries(new_iter).finish()
170    }
171}
172
173/// An iterator over values ordered by key
174pub struct ValuesIterator<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize>
175where
176    K: Ord + BorshDeserialize + BorshSerialize,
177    V: BorshDeserialize + BorshSerialize,
178{
179    next_node: Option<usize>,
180    tree: &'a RBForest<'b, K, V, KSIZE, VSIZE>,
181}
182
183impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize>
184    ValuesIterator<'a, 'b, K, V, KSIZE, VSIZE>
185where
186    K: Ord + BorshDeserialize + BorshSerialize,
187    V: BorshDeserialize + BorshSerialize,
188{
189    pub(super) fn from_raw_parts(
190        tree: &'a RBForest<'b, K, V, KSIZE, VSIZE>,
191        next_node: Option<usize>,
192    ) -> Self {
193        Self { next_node, tree }
194    }
195}
196
197impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> Iterator
198    for ValuesIterator<'a, 'b, K, V, KSIZE, VSIZE>
199where
200    K: Ord + BorshDeserialize + BorshSerialize,
201    V: BorshDeserialize + BorshSerialize,
202{
203    type Item = V;
204
205    fn next(&mut self) -> Option<Self::Item> {
206        self.next_node.map(|mut id| {
207            let nodes = &self.tree.nodes;
208
209            let value = V::deserialize(&mut nodes[id].value.as_slice()).expect("Value corrupted");
210
211            // find next
212            if let Some(right_id) = nodes[id].right() {
213                self.next_node = Some(self.tree.min(right_id as usize));
214            } else {
215                self.next_node = None;
216                while let Some(parent_id) = nodes[id].parent() {
217                    let parent_id = parent_id as usize;
218                    if Some(id as u32) == nodes[parent_id].left() {
219                        self.next_node = Some(parent_id);
220                        break;
221                    } else {
222                        id = parent_id;
223                    }
224                }
225            }
226
227            value
228        })
229    }
230}
231
232impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> FusedIterator
233    for ValuesIterator<'a, 'b, K, V, KSIZE, VSIZE>
234where
235    K: Ord + BorshDeserialize + BorshSerialize,
236    V: BorshDeserialize + BorshSerialize,
237{
238}
239
240impl<'a, 'b, K, V, const KSIZE: usize, const VSIZE: usize> fmt::Debug
241    for ValuesIterator<'a, 'b, K, V, KSIZE, VSIZE>
242where
243    K: Ord + BorshDeserialize + BorshSerialize + fmt::Debug,
244    V: BorshDeserialize + BorshSerialize + fmt::Debug,
245{
246    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
247        let ValuesIterator { next_node, tree } = self;
248        let new_iter = ValuesIterator {
249            next_node: *next_node,
250            tree,
251        };
252        f.debug_set().entries(new_iter).finish()
253    }
254}
slice_rbtree/forest/iterators/mod.rs

slice_rbtree/forest/iterators/
mod.rs
