Binary tree with depth-first iterator in Rust.

binarytree.rs

// Binary tree and depth-first iterator in Rust.
//
// Pay close attention to the types, they are carefully chosen:
//
// enum BinaryTree:
// - We want a Leaf to own its value of type T, so its value is a T (not a reference or pointer to T).
// - We want an Internal node to own its two children, so its elements are Boxes (owned heap pointers).
// - The methods new_internal() and new_leaf() in its impl return a Box<BinaryTree<T>> and not a BinaryTree<T>
//      because we want to allocate nodes on the heap and it would be very inconvenient if the caller had to
//      create a Box for each node.
//
// trait BinaryTreeDepthFirstIter:
// - Interface that defines an iter() method that creates a DepthFirstIterator.
// - Implemented for Box<BinaryTree<T>>, so that we can call iter() on a node on the heap owned by a Box.
//
// struct DepthFirstIterator:
// - Contains the state of the iterator, which is a stack of references to nodes on the heap. These are references
//      because we want the iterator to only borrow, and not own, the boxes in the tree. If these were not references,
//      then a call to iter() would mean we would move the nodes of the tree into the iterator instead of just
//      borrowing them, which would destroy the tree (make it inaccessible after calling iter()).
// - The item type of the iterator is &T, because we also only want to borrow the values from the leaf nodes.
// - Note the construction "&**boxed" in the implementation of the next() method. It works like this:
//      - boxed is a &Box<BinaryTree<T>> (reference to a box that owns a node)
//      - *boxed dereferences the reference to Box<BinaryTree<T>>
//      - **boxed dereferences the box to BinaryTree<T>
//      - &**boxed takes a reference so we have &BinaryTree<T>, a reference to a node that we can match on
//      So we first have to remove two and then add a layer of indirection. The last & is necessary because, again,
//      we just want to borrow the node, not move it out of the box.
// - Finally, the 'ref' in the match patterns makes it so that we borrow the contained nodes and value
//      instead of moving them out.

#[derive(Debug)]
enum BinaryTree<T> {
    Internal(Box<BinaryTree<T>>, Box<BinaryTree<T>>),
    Leaf(T),
}

trait BinaryTreeDepthFirstIter<'a, T> {
    fn iter(&'a self) -> DepthFirstIterator<'a, T>;
}

#[derive(Debug)]
struct DepthFirstIterator<'a, T> {
    stack: Vec<&'a Box<BinaryTree<T>>>,
}

// ==== Implementation =================================================================================================

impl<T> BinaryTree<T> {
    fn new_internal(left: Box<BinaryTree<T>>, right: Box<BinaryTree<T>>) -> Box<BinaryTree<T>> {
        Box::new(BinaryTree::Internal(left, right))
    }

    fn new_leaf(value: T) -> Box<BinaryTree<T>> {
        Box::new(BinaryTree::Leaf(value))
    }
}

impl<'a, T> BinaryTreeDepthFirstIter<'a, T> for Box<BinaryTree<T>> {
    fn iter(&'a self) -> DepthFirstIterator<'a, T> {
        DepthFirstIterator { stack: vec![self] }
    }
}

impl<'a, T> Iterator for DepthFirstIterator<'a, T> {
    type Item = &'a T;

    fn next(&mut self) -> Option<&'a T> {
        match self.stack.pop() {
            Some(boxed) => match &**boxed {
                // Internal node: push its children and call next() recursively
                BinaryTree::Internal(ref left, ref right) => {
                    self.stack.push(right);
                    self.stack.push(left);
                    self.next()
                }

                // Leaf node
                BinaryTree::Leaf(ref value) => Some(value),
            },

            // Empty stack
            None => None,
        }
    }
}

// ===== main ==========================================================================================================

// Non-copy wrapper around an i32
#[derive(Debug)]
struct Value(i32);

fn main() {
    // Tree with non-copy values
    let tree = BinaryTree::new_internal(
        BinaryTree::new_internal(
            BinaryTree::new_leaf(Value(1)),
            BinaryTree::new_leaf(Value(2)),
        ),
        BinaryTree::new_leaf(Value(3)),
    );

    for value in tree.iter() {
        println!("{:?}", value);
    }

    println!("{:?}", tree);
}