Implement iterators for BinaryTree example

bin_tree/src/main.rs

@@ -11,6 +11,14 @@ struct TreeNode<T> {
     right: BinaryTree<T>
 }
 
+impl<T> BinaryTree<T> {
+    fn iter(&self) -> TreeIter<T> {
+        let mut iter = TreeIter { unvisited: Vec::new() };
+        iter.push_left_edge(self);
+        iter
+    }
+}
+
 impl<T: Ord> BinaryTree<T> {
     fn add(&mut self, value: T) {
         match *self {
@@ -30,6 +38,61 @@ impl<T: Ord> BinaryTree<T> {
     }
 }
 
+impl<'a, T: 'a> IntoIterator for &'a BinaryTree<T> {
+    type Item = &'a T;
+    type IntoIter = TreeIter<'a, T>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.iter()
+    }
+}
+
+use self::BinaryTree::*;
+
+// The state of an in-order traversal of a `BinaryTree`
+struct TreeIter<'a, T: 'a> {
+    // A stack of references to tree nodes. Since we use `Vec`'s
+    // `push` and `pop` methods, the top of the stack is the end of the
+    // vector.
+    //
+    // The node the iterator will visit next is at the top of the stack,
+    // with those ancestors still unvisited below it. If the stack is empty,
+    // the iteration is over
+    unvisited: Vec<&'a TreeNode<T>>
+}
+
+impl<'a, T:'a> TreeIter<'a, T> {
+    fn push_left_edge(&mut self, mut tree: &'a BinaryTree<T>) {
+        while let NonEmpty(ref node) = *tree {
+            self.unvisited.push(node);
+            tree = &node.left;
+        }
+    }
+}
+
+impl<'a, T> Iterator for TreeIter<'a, T> {
+    type Item = &'a T;
+    fn next(&mut self) -> Option<&'a T> {
+        // Find the node this iteration must produce,
+        // or finish the iteration
+        let node = match self.unvisited.pop() {
+            None => return None,
+            Some(n) => n,
+        };
+
+        // The next node after this one is the leftmost child of
+        // this node's right child, so push the path from here down.
+        self.push_left_edge(&node.right);
+
+        // Produce a reference to this node's value.
+        Some(&node.element)
+    }
+}
+
+fn make_node<T>(left: BinaryTree<T>, element: T, right: BinaryTree<T>) -> BinaryTree<T> {
+    NonEmpty(Box::new(TreeNode { left, element, right }))
+}
+
 fn main() {
     let mut tree = BinaryTree::Empty;
     tree.add("Mercury");
@@ -43,4 +106,23 @@ fn main() {
     tree.add("Pluto");
 
     println!("Tree: {:?}", tree);
+
+    // Build a small tree
+    let subtree_l = make_node(Empty, "mecha", Empty);
+    let subtree_rl = make_node(Empty, "droid", Empty);
+    let subtree_r = make_node(subtree_rl, "robot", Empty);
+    let tree2 = make_node(subtree_l, "Jaeger", subtree_r);
+
+    // Iterate over it.
+    let mut v = Vec::new();
+    for kind in &tree2 {
+        v.push(*kind);
+    }
+    assert_eq!(v, ["mecha", "Jaeger", "droid", "robot"]);
+
+    assert_eq!(tree2.iter()
+        .map(|name| format!("mega-{}", name))
+        .collect::<Vec<_>>(),
+        vec!["mega-mecha", "mega-Jaeger", "mega-droid", "mega-robot"]
+    );
 }