binary_tree/
iter.rs

1//! Generic iterators.
2//!
3//! This module is not meant for the end-user.
4
5use Node;
6use NodeMut;
7use unbox::Unbox;
8
9#[derive(PartialEq)]
10enum IterAction {
11    Left,
12    Right,
13}
14
15pub struct Iter<'a, T>
16    where T: Node + 'a
17{
18    stack: Vec<(&'a T, IterAction)>,
19}
20
21impl<'a, T> Iter<'a, T>
22    where T: Node + 'a
23{
24    pub fn new(root: Option<&'a T>) -> Iter<'a, T> {
25        Iter { stack: root.map_or(vec![], |node| vec![(node, IterAction::Left)]) }
26    }
27}
28
29impl<'a, T> Iterator for Iter<'a, T>
30    where T: Node + 'a
31{
32    type Item = &'a T::Value;
33
34    fn next(&mut self) -> Option<&'a T::Value> {
35        if let Some((mut subtree, action)) = self.stack.pop() {
36            if action == IterAction::Left {
37                while let Some(st) = subtree.left() {
38                    self.stack.push((&*subtree, IterAction::Right));
39                    subtree = st;
40                }
41            }
42            if let Some(st) = subtree.right() {
43                self.stack.push((&*st, IterAction::Left));
44            }
45            Some(subtree.value())
46        } else {
47            None
48        }
49    }
50}
51
52pub struct IntoIter<T>
53    where T: NodeMut,
54          T::NodePtr: Unbox<Target=T>
55{
56    stack: Vec<(T::NodePtr, IterAction)>,
57}
58
59impl<T> IntoIter<T>
60    where T: NodeMut,
61          T::NodePtr: Unbox<Target=T>
62{
63    pub fn new(root: Option<T::NodePtr>) -> IntoIter<T> {
64        IntoIter { stack: root.map_or(vec![], |node| vec![(node, IterAction::Left)]) }
65    }
66}
67
68impl<T> Iterator for IntoIter<T>
69    where T: NodeMut,
70          T::NodePtr: Unbox<Target=T>
71{
72    type Item = T::Value;
73
74    fn next(&mut self) -> Option<T::Value> {
75        if let Some((mut subtree, action)) = self.stack.pop() {
76            if action == IterAction::Left {
77                while let Some(left) = subtree.detach_left() {
78                    self.stack.push((subtree, IterAction::Right));
79                    subtree = left;
80                }
81            }
82            let (value, _, right) = subtree.unbox().into_parts();
83            if let Some(st) = right {
84                self.stack.push((st, IterAction::Left));
85            }
86            Some(value)
87        } else {
88            None
89        }
90    }
91}
92
93impl<T> Drop for IntoIter<T>
94    where T: NodeMut,
95          T::NodePtr: Unbox<Target=T>
96{
97    fn drop(&mut self) {
98        for _ in self {}
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use NodeMut;
105    use test::TestNode;
106    use super::Iter;
107    use super::IntoIter;
108
109    #[test]
110    fn iteration() {
111        let mut ct = Box::new(TestNode::new(7));
112        let mut ct_l = Box::new(TestNode::new(8));
113        ct_l.insert_right(Some(Box::new(TestNode::new(12))));
114        ct.insert_left(Some(ct_l));
115        ct.insert_right(Some(Box::new(TestNode::new(5))));
116
117        {
118            let vals: Vec<_> = Iter::new(Some(&*ct)).collect();
119            assert_eq!(vals, [&8, &12, &7, &5]);
120        }
121
122        let node_mi: IntoIter<TestNode<_>> = IntoIter::new(Some(ct));
123        let vals: Vec<_> = node_mi.collect();
124        assert_eq!(vals, [8, 12, 7, 5]);
125    }
126}