1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126

//! Generic iterators.
//!
//! This module is not meant for the end-user.

use Node;
use NodeMut;
use unbox::Unbox;

#[derive(PartialEq)]
enum IterAction {
    Left,
    Right,
}

pub struct Iter<'a, T>
    where T: Node + 'a
{
    stack: Vec<(&'a T, IterAction)>,
}

impl<'a, T> Iter<'a, T>
    where T: Node + 'a
{
    pub fn new(root: Option<&'a T>) -> Iter<'a, T> {
        Iter { stack: root.map_or(vec![], |node| vec![(node, IterAction::Left)]) }
    }
}

impl<'a, T> Iterator for Iter<'a, T>
    where T: Node + 'a
{
    type Item = &'a T::Value;

    fn next(&mut self) -> Option<&'a T::Value> {
        if let Some((mut subtree, action)) = self.stack.pop() {
            if action == IterAction::Left {
                while let Some(st) = subtree.left() {
                    self.stack.push((&*subtree, IterAction::Right));
                    subtree = st;
                }
            }
            if let Some(st) = subtree.right() {
                self.stack.push((&*st, IterAction::Left));
            }
            Some(subtree.value())
        } else {
            None
        }
    }
}

pub struct IntoIter<T>
    where T: NodeMut,
          T::NodePtr: Unbox<Target=T>
{
    stack: Vec<(T::NodePtr, IterAction)>,
}

impl<T> IntoIter<T>
    where T: NodeMut,
          T::NodePtr: Unbox<Target=T>
{
    pub fn new(root: Option<T::NodePtr>) -> IntoIter<T> {
        IntoIter { stack: root.map_or(vec![], |node| vec![(node, IterAction::Left)]) }
    }
}

impl<T> Iterator for IntoIter<T>
    where T: NodeMut,
          T::NodePtr: Unbox<Target=T>
{
    type Item = T::Value;

    fn next(&mut self) -> Option<T::Value> {
        if let Some((mut subtree, action)) = self.stack.pop() {
            if action == IterAction::Left {
                while let Some(left) = subtree.detach_left() {
                    self.stack.push((subtree, IterAction::Right));
                    subtree = left;
                }
            }
            let (value, _, right) = subtree.unbox().into_parts();
            if let Some(st) = right {
                self.stack.push((st, IterAction::Left));
            }
            Some(value)
        } else {
            None
        }
    }
}

impl<T> Drop for IntoIter<T>
    where T: NodeMut,
          T::NodePtr: Unbox<Target=T>
{
    fn drop(&mut self) {
        for _ in self {}
    }
}

#[cfg(test)]
mod tests {
    use NodeMut;
    use test::TestNode;
    use super::Iter;
    use super::IntoIter;

    #[test]
    fn iteration() {
        let mut ct = Box::new(TestNode::new(7));
        let mut ct_l = Box::new(TestNode::new(8));
        ct_l.insert_right(Some(Box::new(TestNode::new(12))));
        ct.insert_left(Some(ct_l));
        ct.insert_right(Some(Box::new(TestNode::new(5))));

        {
            let vals: Vec<_> = Iter::new(Some(&*ct)).collect();
            assert_eq!(vals, [&8, &12, &7, &5]);
        }

        let node_mi: IntoIter<TestNode<_>> = IntoIter::new(Some(ct));
        let vals: Vec<_> = node_mi.collect();
        assert_eq!(vals, [8, 12, 7, 5]);
    }
}