ego_tree/
iter.rs

1use std::iter::FusedIterator;
2use std::ops::Range;
3use std::{slice, vec};
4
5use crate::{Node, NodeId, NodeRef, Tree};
6
7/// Iterator that moves out of a tree in insert order.
8#[derive(Debug)]
9pub struct IntoIter<T>(vec::IntoIter<Node<T>>);
10impl<T> ExactSizeIterator for IntoIter<T> {}
11impl<T> FusedIterator for IntoIter<T> {}
12impl<T> Iterator for IntoIter<T> {
13    type Item = T;
14    fn next(&mut self) -> Option<Self::Item> {
15        self.0.next().map(|node| node.value)
16    }
17    fn size_hint(&self) -> (usize, Option<usize>) {
18        self.0.size_hint()
19    }
20}
21impl<T> DoubleEndedIterator for IntoIter<T> {
22    fn next_back(&mut self) -> Option<Self::Item> {
23        self.0.next_back().map(|node| node.value)
24    }
25}
26
27/// Iterator over values in insert order.
28#[derive(Debug)]
29pub struct Values<'a, T: 'a>(slice::Iter<'a, Node<T>>);
30impl<'a, T: 'a> Clone for Values<'a, T> {
31    fn clone(&self) -> Self {
32        Values(self.0.clone())
33    }
34}
35impl<'a, T: 'a> ExactSizeIterator for Values<'a, T> {}
36impl<'a, T: 'a> FusedIterator for Values<'a, T> {}
37impl<'a, T: 'a> Iterator for Values<'a, T> {
38    type Item = &'a T;
39    fn next(&mut self) -> Option<Self::Item> {
40        self.0.next().map(|node| &node.value)
41    }
42    fn size_hint(&self) -> (usize, Option<usize>) {
43        self.0.size_hint()
44    }
45}
46impl<'a, T: 'a> DoubleEndedIterator for Values<'a, T> {
47    fn next_back(&mut self) -> Option<Self::Item> {
48        self.0.next_back().map(|node| &node.value)
49    }
50}
51
52/// Mutable iterator over values in insert order.
53#[derive(Debug)]
54pub struct ValuesMut<'a, T: 'a>(slice::IterMut<'a, Node<T>>);
55impl<'a, T: 'a> ExactSizeIterator for ValuesMut<'a, T> {}
56impl<'a, T: 'a> FusedIterator for ValuesMut<'a, T> {}
57impl<'a, T: 'a> Iterator for ValuesMut<'a, T> {
58    type Item = &'a mut T;
59    fn next(&mut self) -> Option<Self::Item> {
60        self.0.next().map(|node| &mut node.value)
61    }
62    fn size_hint(&self) -> (usize, Option<usize>) {
63        self.0.size_hint()
64    }
65}
66impl<'a, T: 'a> DoubleEndedIterator for ValuesMut<'a, T> {
67    fn next_back(&mut self) -> Option<Self::Item> {
68        self.0.next_back().map(|node| &mut node.value)
69    }
70}
71
72/// Iterator over nodes in insert order.
73#[derive(Debug)]
74pub struct Nodes<'a, T: 'a> {
75    tree: &'a Tree<T>,
76    iter: Range<usize>,
77}
78impl<'a, T: 'a> Clone for Nodes<'a, T> {
79    fn clone(&self) -> Self {
80        Self {
81            tree: self.tree,
82            iter: self.iter.clone(),
83        }
84    }
85}
86impl<'a, T: 'a> ExactSizeIterator for Nodes<'a, T> {}
87impl<'a, T: 'a> FusedIterator for Nodes<'a, T> {}
88impl<'a, T: 'a> Iterator for Nodes<'a, T> {
89    type Item = NodeRef<'a, T>;
90    fn next(&mut self) -> Option<Self::Item> {
91        self.iter
92            .next()
93            .map(|i| unsafe { self.tree.get_unchecked(NodeId::from_index(i)) })
94    }
95    fn size_hint(&self) -> (usize, Option<usize>) {
96        self.iter.size_hint()
97    }
98}
99impl<'a, T: 'a> DoubleEndedIterator for Nodes<'a, T> {
100    fn next_back(&mut self) -> Option<Self::Item> {
101        self.iter
102            .next_back()
103            .map(|i| unsafe { self.tree.get_unchecked(NodeId::from_index(i)) })
104    }
105}
106
107impl<T> IntoIterator for Tree<T> {
108    type Item = T;
109    type IntoIter = IntoIter<T>;
110    fn into_iter(self) -> Self::IntoIter {
111        IntoIter(self.vec.into_iter())
112    }
113}
114
115impl<T> Tree<T> {
116    /// Returns an iterator over values in insert order.
117    pub fn values(&self) -> Values<T> {
118        Values(self.vec.iter())
119    }
120
121    /// Returns a mutable iterator over values in insert order.
122    pub fn values_mut(&mut self) -> ValuesMut<T> {
123        ValuesMut(self.vec.iter_mut())
124    }
125
126    /// Returns an iterator over nodes in insert order.
127    pub fn nodes(&self) -> Nodes<T> {
128        Nodes {
129            tree: self,
130            iter: 0..self.vec.len(),
131        }
132    }
133}
134
135macro_rules! axis_iterators {
136    ($(#[$m:meta] $i:ident($f:path);)*) => {
137        $(
138            #[$m]
139            #[derive(Debug)]
140            pub struct $i<'a, T: 'a>(Option<NodeRef<'a, T>>);
141            impl<'a, T: 'a> Clone for $i<'a, T> {
142                fn clone(&self) -> Self {
143                    $i(self.0)
144                }
145            }
146            impl<'a, T: 'a> FusedIterator for $i<'a, T> {}
147            impl<'a, T: 'a> Iterator for $i<'a, T> {
148                type Item = NodeRef<'a, T>;
149                fn next(&mut self) -> Option<Self::Item> {
150                    let node = self.0.take();
151                    self.0 = node.as_ref().and_then($f);
152                    node
153                }
154            }
155        )*
156    };
157}
158
159axis_iterators! {
160    /// Iterator over ancestors.
161    Ancestors(NodeRef::parent);
162
163    /// Iterator over previous siblings.
164    PrevSiblings(NodeRef::prev_sibling);
165
166    /// Iterator over next siblings.
167    NextSiblings(NodeRef::next_sibling);
168
169    /// Iterator over first children.
170    FirstChildren(NodeRef::first_child);
171
172    /// Iterator over last children.
173    LastChildren(NodeRef::last_child);
174}
175
176/// Iterator over children.
177#[derive(Debug)]
178pub struct Children<'a, T: 'a> {
179    front: Option<NodeRef<'a, T>>,
180    back: Option<NodeRef<'a, T>>,
181}
182impl<'a, T: 'a> Clone for Children<'a, T> {
183    fn clone(&self) -> Self {
184        Self {
185            front: self.front,
186            back: self.back,
187        }
188    }
189}
190impl<'a, T: 'a> FusedIterator for Children<'a, T> {}
191impl<'a, T: 'a> Iterator for Children<'a, T> {
192    type Item = NodeRef<'a, T>;
193    fn next(&mut self) -> Option<Self::Item> {
194        if self.front == self.back {
195            let node = self.front.take();
196            self.back = None;
197            node
198        } else {
199            let node = self.front.take();
200            self.front = node.as_ref().and_then(NodeRef::next_sibling);
201            node
202        }
203    }
204}
205impl<'a, T: 'a> DoubleEndedIterator for Children<'a, T> {
206    fn next_back(&mut self) -> Option<Self::Item> {
207        if self.back == self.front {
208            let node = self.back.take();
209            self.front = None;
210            node
211        } else {
212            let node = self.back.take();
213            self.back = node.as_ref().and_then(NodeRef::prev_sibling);
214            node
215        }
216    }
217}
218
219/// Open or close edge of a node.
220#[derive(Debug)]
221pub enum Edge<'a, T: 'a> {
222    /// Open.
223    Open(NodeRef<'a, T>),
224    /// Close.
225    Close(NodeRef<'a, T>),
226}
227impl<'a, T: 'a> Copy for Edge<'a, T> {}
228impl<'a, T: 'a> Clone for Edge<'a, T> {
229    fn clone(&self) -> Self {
230        *self
231    }
232}
233impl<'a, T: 'a> Eq for Edge<'a, T> {}
234impl<'a, T: 'a> PartialEq for Edge<'a, T> {
235    fn eq(&self, other: &Self) -> bool {
236        match (*self, *other) {
237            (Edge::Open(a), Edge::Open(b)) | (Edge::Close(a), Edge::Close(b)) => a == b,
238            _ => false,
239        }
240    }
241}
242
243/// Iterator which traverses a subtree.
244#[derive(Debug)]
245pub struct Traverse<'a, T: 'a> {
246    root: Option<NodeRef<'a, T>>,
247    edge: Option<Edge<'a, T>>,
248}
249impl<'a, T: 'a> Clone for Traverse<'a, T> {
250    fn clone(&self) -> Self {
251        Self {
252            root: self.root,
253            edge: self.edge,
254        }
255    }
256}
257impl<'a, T: 'a> FusedIterator for Traverse<'a, T> {}
258impl<'a, T: 'a> Iterator for Traverse<'a, T> {
259    type Item = Edge<'a, T>;
260    fn next(&mut self) -> Option<Self::Item> {
261        match self.edge {
262            None => {
263                if let Some(root) = self.root {
264                    self.edge = Some(Edge::Open(root));
265                }
266            }
267            Some(Edge::Open(node)) => {
268                if let Some(first_child) = node.first_child() {
269                    self.edge = Some(Edge::Open(first_child));
270                } else {
271                    self.edge = Some(Edge::Close(node));
272                }
273            }
274            Some(Edge::Close(node)) => {
275                if node == self.root.unwrap() {
276                    self.root = None;
277                    self.edge = None;
278                } else if let Some(next_sibling) = node.next_sibling() {
279                    self.edge = Some(Edge::Open(next_sibling));
280                } else {
281                    self.edge = node.parent().map(Edge::Close);
282                }
283            }
284        }
285        self.edge
286    }
287}
288
289/// Iterator over a node and its descendants.
290#[derive(Debug)]
291pub struct Descendants<'a, T: 'a>(Traverse<'a, T>);
292impl<'a, T: 'a> Clone for Descendants<'a, T> {
293    fn clone(&self) -> Self {
294        Descendants(self.0.clone())
295    }
296}
297impl<'a, T: 'a> FusedIterator for Descendants<'a, T> {}
298impl<'a, T: 'a> Iterator for Descendants<'a, T> {
299    type Item = NodeRef<'a, T>;
300    fn next(&mut self) -> Option<Self::Item> {
301        for edge in &mut self.0 {
302            if let Edge::Open(node) = edge {
303                return Some(node);
304            }
305        }
306        None
307    }
308}
309
310impl<'a, T: 'a> NodeRef<'a, T> {
311    /// Returns an iterator over ancestors.
312    pub fn ancestors(&self) -> Ancestors<'a, T> {
313        Ancestors(self.parent())
314    }
315
316    /// Returns an iterator over previous siblings.
317    pub fn prev_siblings(&self) -> PrevSiblings<'a, T> {
318        PrevSiblings(self.prev_sibling())
319    }
320
321    /// Returns an iterator over next siblings.
322    pub fn next_siblings(&self) -> NextSiblings<'a, T> {
323        NextSiblings(self.next_sibling())
324    }
325
326    /// Returns an iterator over first children.
327    pub fn first_children(&self) -> FirstChildren<'a, T> {
328        FirstChildren(self.first_child())
329    }
330
331    /// Returns an iterator over last children.
332    pub fn last_children(&self) -> LastChildren<'a, T> {
333        LastChildren(self.last_child())
334    }
335
336    /// Returns an iterator over children.
337    pub fn children(&self) -> Children<'a, T> {
338        Children {
339            front: self.first_child(),
340            back: self.last_child(),
341        }
342    }
343
344    /// Returns an iterator which traverses the subtree starting at this node.
345    pub fn traverse(&self) -> Traverse<'a, T> {
346        Traverse {
347            root: Some(*self),
348            edge: None,
349        }
350    }
351
352    /// Returns an iterator over this node and its descendants.
353    pub fn descendants(&self) -> Descendants<'a, T> {
354        Descendants(self.traverse())
355    }
356}