iterators.rs

///////////////////////////////////////////////////////////////////////////
// The iterator trait

trait MutIterator<T> {
    fn has_next(&self) -> bool;
    fn next<'a>(&'a mut self) -> &'a mut T;
}

///////////////////////////////////////////////////////////////////////////
// Implementation for vectors (easy)

struct MutSliceIterator<'self, T> {
    v: &'self mut [T],
    pos: uint,
}

impl<'self, T> MutIterator<T> for MutSliceIterator<'self, T> {
    fn has_next(&self) -> bool {
        self.pos == self.v.len()
    }

    fn next<'a>(&'a mut self) -> &'a mut T {
        if self.pos == self.v.len() {
            fail!("next() called when has_next() is false");
        }
        let p = self.pos;
        self.pos += 1;
        &mut self.v[p]
    }
}

///////////////////////////////////////////////////////////////////////////
// Implementation for a binary tree (pre- and post-order, trickier)

struct BinaryTree<T> {
    value: T,
    left: Option<~BinaryTree<T>>,
    right: Option<~BinaryTree<T>>,
}

impl<T> BinaryTree<T> {
    fn iterator<'a>(&'a mut self) -> BinaryTreePreorderMutIterator<'a, T> {
        BinaryTreePreorderMutIterator {
            stack: ~[self]
        }
    }
    fn postorder_iterator<'a>(&'a mut self) -> BinaryTreePostorderMutIterator<'a, T> {
        let mut iter = BinaryTreePostorderMutIterator {
            stack: ~[]
        };
        iter.push_state(self);
        return iter;
    }
}

struct BinaryTreePreorderMutIterator<'self,T> {
    stack: ~[&'self mut BinaryTree<T>],
}

impl<'self,T> MutIterator<T> for BinaryTreePreorderMutIterator<'self,T> {
    fn has_next(&self) -> bool {
        !self.stack.is_empty()
    }

    fn next<'a>(&'a mut self) -> &'a mut T {
        if !self.has_next() {
            fail!("next() called when has_next() is false");
        }

        let cur = self.stack.pop();

        match cur.right {
            None => {}
            Some(~ref mut r) => {
                self.stack.push(r);
            }
        }

        match cur.left {
            None => {}
            Some(~ref mut l) => {
                self.stack.push(l);
            }
        }

        &mut cur.value
    }
}

/*

Postorder: visit node *after* its children.

       A    visits: C, D, B, E, A
      / \
     B   E
    / \
   C   D

Iterator maintains a stack of nodes. When we visit left child, we null
it out. When we visit right child, we null *it* out. When both
children are null, we visit node.

To generalize, you might have a slice of children and just divide the
slice to separate out the next children from the rest.

*/

struct BinaryTreePostorderMutIterator<'self,T> {
    stack: ~[PostOrderState<'self, T>]
}

struct PostOrderState<'self,T> {
    value: &'self mut T,
    left: Option<&'self mut BinaryTree<T>>,
    right: Option<&'self mut BinaryTree<T>>,
}

impl<'self,T> BinaryTreePostorderMutIterator<'self,T> {
    fn push_state(&mut self,
                  v: &'self mut BinaryTree<T>) {
        let value = &mut v.value;
        let left = match v.left {None => None,
                                  Some(~ref mut l) => Some(l)};
        let right = match v.right {None => None,
                                   Some(~ref mut r) => Some(r)};
        self.stack.push(PostOrderState {value: value,
                                        left: left,
                                        right: right});
    }
}

impl<'self,T> MutIterator<T> for BinaryTreePostorderMutIterator<'self,T> {
    fn has_next(&self) -> bool {
        !self.stack.is_empty()
    }

    fn next<'a>(&'a mut self) -> &'a mut T {
        if !self.has_next() {
            fail!("next() called when has_next() is false");
        }

        loop {
            let next = {
                let top = &mut self.stack[self.stack.len() - 1];
                if top.left.is_some() {
                    top.left.swap_unwrap()
                } else if top.right.is_some() {
                    top.right.swap_unwrap()
                } else {
                    break;
                }
            };
            self.push_state(next);
        }

        let PostOrderState {value, _} = self.stack.pop();
        return value;
    }
}

#[cfg(test)]
fn check_iter<V:Eq,I:MutIterator<V>>(mut iter: I,
                                     values: &[V]) {
    let mut counter = 0_u;
    while iter.has_next() {
        let next = iter.next();
        if *next != values[counter] {
            fail!(fmt!("Index %u yielded %? not %?",
                       counter, *next, values[counter]));
        }
        counter += 1_u;
    }
    assert_eq!(counter, values.len());
}

#[test]
pub fn preorder1() {
    let mut the_tree = BinaryTree {
        value: "r",
        left: Some(~BinaryTree {
            value: "r.l",
            left: None,
            right: Some(~BinaryTree {
                value: "r.l.r",
                left: None,
                right: None
            }),
        }),
        right: Some(~BinaryTree {
            value: "r.r",
            left: Some(~BinaryTree {
                value: "r.r.l",
                left: None,
                right: None,
            }),
            right: Some(~BinaryTree {
                value: "r.r.r",
                left: None,
                right: None
            }),
        }),
    };

    check_iter(the_tree.iterator(),
               ["r", "r.l", "r.l.r", "r.r", "r.r.l", "r.r.r"]);
}

#[test]
pub fn postorder1() {
    let mut the_tree = BinaryTree {
        value: "r",
        left: Some(~BinaryTree {
            value: "r.l",
            left: None,
            right: Some(~BinaryTree {
                value: "r.l.r",
                left: None,
                right: None
            }),
        }),
        right: Some(~BinaryTree {
            value: "r.r",
            left: Some(~BinaryTree {
                value: "r.r.l",
                left: None,
                right: None,
            }),
            right: Some(~BinaryTree {
                value: "r.r.r",
                left: None,
                right: None
            }),
        }),
    };

    check_iter(the_tree.postorder_iterator(),
               ["r.l.r", "r.l", "r.r.l", "r.r.r", "r.r", "r"]);
}

fn main() {}