forest_ds/
iter.rs

1/// Implements iteration over a tree
2use crate::{entry::Entry, id::NodeId, node::Node, tree::Tree};
3
4impl<T> Tree<T> {
5    #[must_use]
6    pub fn iter(&self) -> Iter<T> {
7        Iter {
8            current: self.first_node,
9            nodes: &self.nodes,
10        }
11    }
12
13    #[must_use]
14    pub fn iter_from(&self, id: &NodeId) -> Iter<T> {
15        Iter {
16            current: self.index(id),
17            nodes: &self.nodes,
18        }
19    }
20
21    #[must_use]
22    pub fn iter_mut(&mut self) -> IterMut<T> {
23        IterMut {
24            current: self.first_node,
25            nodes: self
26                .nodes
27                .iter_mut()
28                .filter_map(|entry| match entry {
29                    Entry::Free { .. } => None,
30                    Entry::Occupied(node) => Some(Node {
31                        value: Some(&mut node.value),
32                        parent: node.parent,
33                        first_child: node.first_child,
34                        last_child: node.last_child,
35                        next_sibling: node.next_sibling,
36                        prev_sibling: node.prev_sibling,
37                    }),
38                })
39                .collect(),
40        }
41    }
42
43    #[must_use]
44    pub fn into_iterator(self) -> IntoIter<T> {
45        IntoIter {
46            current: self.first_node,
47            nodes: self
48                .nodes
49                .into_iter()
50                .filter_map(|entry| match entry {
51                    Entry::Free { .. } => None,
52                    Entry::Occupied(node) => Some(Node {
53                        value: Some(node.value),
54                        parent: node.parent,
55                        first_child: node.first_child,
56                        last_child: node.last_child,
57                        next_sibling: node.next_sibling,
58                        prev_sibling: node.prev_sibling,
59                    }),
60                })
61                .collect(),
62        }
63    }
64}
65
66impl<'a, T> IntoIterator for &'a Tree<T> {
67    type Item = &'a T;
68
69    type IntoIter = Iter<'a, T>;
70
71    fn into_iter(self) -> Self::IntoIter {
72        self.iter()
73    }
74}
75
76impl<T> IntoIterator for Tree<T> {
77    type Item = T;
78
79    type IntoIter = IntoIter<T>;
80
81    fn into_iter(self) -> Self::IntoIter {
82        self.into_iterator()
83    }
84}
85
86#[derive(Debug)]
87pub struct Iter<'a, T> {
88    current: Option<usize>,
89    nodes: &'a [Entry<T>],
90}
91
92impl<'a, T> Iterator for Iter<'a, T> {
93    type Item = &'a T;
94
95    fn next(&mut self) -> Option<Self::Item> {
96        self.current.take().map(|current| {
97            let node = &self.nodes[current];
98
99            if let Some(child) = node.unwrap_ref().first_child {
100                self.current = Some(child);
101            } else if let Some(sibling) = node.unwrap_ref().next_sibling {
102                self.current = Some(sibling);
103            } else {
104                // Start from the current node
105                let mut next = node.unwrap_ref();
106                // Cycle to the parent to search for the next sibling or go up the tree
107                while let Some(parent_index) = next.parent {
108                    next = self.nodes[parent_index].unwrap_ref();
109                    if next.next_sibling.is_some() {
110                        break;
111                    }
112                }
113                // If next.sibling is Some we have the next node, otherwise both next.parent is None
114                // and next.next_sibling is None
115                self.current = next.next_sibling;
116            }
117
118            &node.unwrap_ref().value
119        })
120    }
121}
122
123#[derive(Debug)]
124pub struct IterMut<'a, T> {
125    current: Option<usize>,
126    nodes: Vec<Node<Option<&'a mut T>>>,
127}
128
129impl<'a, T> Iterator for IterMut<'a, T> {
130    type Item = &'a mut T;
131
132    fn next(&mut self) -> Option<Self::Item> {
133        self.current.take().and_then(|current| {
134            // TODO: this could done more safely
135            let node = &self.nodes[current];
136
137            if let Some(child) = node.first_child {
138                self.current = Some(child);
139            } else if let Some(sibling) = node.next_sibling {
140                self.current = Some(sibling);
141            } else {
142                // Start from the current node
143                let mut next = node;
144                // Cycle to the parent to search for the next sibling or go up the tree
145                while let Some(parent_index) = next.parent {
146                    next = &self.nodes[parent_index];
147                    if next.next_sibling.is_some() {
148                        break;
149                    }
150                }
151                // If next.sibling is Some we have the next node, otherwise both next.parent is None
152                // and next.next_sibling is None
153                self.current = next.next_sibling;
154            }
155
156            self.nodes[current].value.take()
157        })
158    }
159}
160
161#[derive(Debug)]
162pub struct IntoIter<T> {
163    current: Option<usize>,
164    nodes: Vec<Node<Option<T>>>,
165}
166
167impl<T> Iterator for IntoIter<T> {
168    type Item = T;
169
170    fn next(&mut self) -> Option<Self::Item> {
171        self.current.take().and_then(|current| {
172            let node = &self.nodes[current];
173
174            if let Some(child) = node.first_child {
175                self.current = Some(child);
176            } else if let Some(sibling) = node.next_sibling {
177                self.current = Some(sibling);
178            } else {
179                // Start from the current node
180                let mut next = node;
181                // Cycle to the parent to search for the next sibling or go up the tree
182                while let Some(parent_index) = next.parent {
183                    next = &self.nodes[parent_index];
184                    if next.next_sibling.is_some() {
185                        break;
186                    }
187                }
188                // If next.sibling is Some we have the next node, otherwise both next.parent is None
189                // and next.next_sibling is None
190                self.current = next.next_sibling;
191            }
192
193            self.nodes[current].value.take()
194        })
195    }
196}
197
198#[cfg(test)]
199mod test {
200    use crate::tree::Tree;
201    use pretty_assertions::assert_eq;
202
203    #[test]
204    fn should_return_none_on_empty() {
205        let tree = Tree::<i32>::new();
206
207        assert_eq!(None, tree.iter().next());
208    }
209
210    #[test]
211    fn should_iter_children() {
212        let mut tree = Tree::<i32>::new();
213
214        tree.append_child(1);
215        tree.append_child(2);
216
217        let mut iter = tree.iter();
218
219        assert_eq!(1, *iter.next().unwrap());
220        assert_eq!(2, *iter.next().unwrap());
221    }
222
223    #[test]
224    fn should_iter_siblings() {
225        let mut tree = Tree::<i32>::new();
226
227        tree.append_sibling(1);
228        tree.append_sibling(2);
229
230        dbg!(&tree);
231
232        let mut iter = tree.iter();
233
234        assert_eq!(1, *iter.next().unwrap());
235        assert_eq!(2, *iter.next().unwrap());
236    }
237
238    #[test]
239    fn mut_should_return_none_on_empty() {
240        let mut tree = Tree::<i32>::new();
241
242        assert_eq!(None, tree.iter_mut().next());
243    }
244
245    #[test]
246    fn mut_should_iter_children() {
247        let mut tree = Tree::<i32>::new();
248
249        tree.append_child(1);
250        tree.append_child(2);
251
252        let mut iter = tree.iter_mut();
253
254        assert_eq!(1, *iter.next().unwrap());
255        assert_eq!(2, *iter.next().unwrap());
256    }
257
258    #[test]
259    fn mut_should_iter_siblings() {
260        let mut tree = Tree::<i32>::new();
261
262        tree.append_sibling(1);
263        tree.append_sibling(2);
264
265        dbg!(&tree);
266
267        let mut iter = tree.iter_mut();
268
269        assert_eq!(1, *iter.next().unwrap());
270        assert_eq!(2, *iter.next().unwrap());
271    }
272
273    #[test]
274    fn into_should_return_none_on_empty() {
275        let tree = Tree::<i32>::new();
276
277        assert_eq!(None, tree.into_iterator().next());
278    }
279
280    #[test]
281    fn into_should_iter_children() {
282        let mut tree = Tree::<i32>::new();
283
284        tree.append_child(1);
285        tree.append_child(2);
286
287        let mut iter = tree.into_iterator();
288
289        assert_eq!(1, iter.next().unwrap());
290        assert_eq!(2, iter.next().unwrap());
291    }
292
293    #[test]
294    fn into_should_iter_siblings() {
295        let mut tree = Tree::<i32>::new();
296
297        tree.append_sibling(1);
298        tree.append_sibling(2);
299
300        dbg!(&tree);
301
302        let mut iter = tree.into_iterator();
303
304        assert_eq!(1, iter.next().unwrap());
305        assert_eq!(2, iter.next().unwrap());
306    }
307}
forest_ds/iter.rs

forest_ds/
iter.rs
