Implement iterators for BinaryTree example

This commit is contained in:
Timothy Warren 2019-03-12 11:09:31 -04:00
parent d4b3bb172a
commit 1be31f3b53

View File

@ -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"]
);
}