Binary trees in Rust: iterators

Question

This is a follow-up to Basic binary tree manipulation in Rust, where Shepmaster suggested that I implement the Iterator trait for the binary tree.

This ended up being less trivial than I thought. My original data structure design was that the iterator would simply contain a vector of binary tree nodes representing the current stack frame of the in-order traversal, had it been implemented recursively. However, I realized (after banging my head against the borrow checker) that this is not possible: it's not possible for both Branch(left, right) and left to be owned by the vector, because Branch(left, right) itself owns left.

So instead we only track the right-facing subtrees (because those are all that's necessary to continue the "recursion"), in addition to the actual value of the left-most node.

Other questions:

• I had originally written new_branch and new_leaf inside the impl<T> BTree<T> block. This worked, but I wasn't able to use BTree::new_branch to access them without qualification: I got E0432 "There is no new_branch in BTree." I even tried making it a pub fn, but that didn't help. What's the proper way to have this namespaced and yet useable?
• Is the shadowing of btree in the test case considered bad form?
• Stylistically, when are expr and return expr; preferred? Is there a technical difference?
• Is panic! the best way to handle invariant violations?

/// A binary tree with data only at the leaves.
#[derive(Debug, PartialEq, Clone)]
enum BTree<T> {
    Leaf(T),
    Branch(Box<BTree<T>>, Box<BTree<T>>),
}

/// This makes it slightly less painful to put things in boxes.
fn new_branch<T>(left: BTree<T>, right: BTree<T>) -> BTree<T> {
    BTree::Branch(Box::new(left), Box::new(right))
}

/// For consistency with `new_branch`.
fn new_leaf<T>(elem: T) -> BTree<T> {
    BTree::Leaf(elem)
}

#[test]
fn test_btree_creation() {
    new_leaf(10);

    let branch: BTree<i32> = new_branch(new_leaf(15), new_leaf(20));
    new_branch(branch.clone(), new_leaf(30));

    assert_eq!(branch, branch.clone());
}

impl<T> BTree<T> {
    // presumably there are useful methods here, like depth(&self) -> u32, etc.
}

impl<T> IntoIterator for BTree<T> {
    type Item = T;
    type IntoIter = BTreeIterator<T>;

    fn into_iter(self) -> BTreeIterator<T> {
        BTreeIterator::new(self)
    }
}

/// Iterator type for a binary tree.
/// This is a generator that progresses through an in-order traversal.
struct BTreeIterator<T> {
    right_nodes: Vec<BTree<T>>,
    current_node: Option<BTree<T>>,
}
impl<T> BTreeIterator<T> {

    fn new(node: BTree<T>) -> BTreeIterator<T> {
        let mut iter = BTreeIterator { right_nodes: vec![], current_node: None };
        iter.add_left_subtree(node);
        return iter;
    }

    /// Consume a binary tree node, traversing its left subtree and
    /// adding all branches to the right to the `right_nodes` field
    /// while setting the current node to the left-most child.
    fn add_left_subtree(&mut self, root: BTree<T>) {
        let mut node: BTree<T> = root;
        while let BTree::Branch(left, right) = node {
            self.right_nodes.push(*right);
            node = *left;
        }
        self.current_node = Some(node);
    }

}

impl<T> Iterator for BTreeIterator<T> {

    type Item = T;

    fn next(&mut self) -> Option<T> {
        // Get the item we're going to return.
        let result = self.current_node.take().map(|node|
            match node {
                BTree::Leaf(n) => n,
                _ => { panic!("invariant violation: expected Leaf") },
            });

        // Now add the next left subtree
        // (this is the "recursive call")
        match self.right_nodes.pop() {
            Some(node) => self.add_left_subtree(node),
            _ => {},
        };

        return result;
    }

}

#[test]
fn test_btree_iteration() {
    assert_eq!(new_leaf(123).into_iter().collect::<Vec<_>>(), vec![123]);

    let btree = new_branch(new_leaf(10), new_branch(new_leaf(20), new_leaf(30)));
    assert_eq!(btree.into_iter().collect::<Vec<_>>(), vec![10, 20, 30]);

    let btree = new_branch(
        new_branch(
            new_branch(new_leaf(9), new_leaf(8)), new_branch(new_leaf(7), new_leaf(6))),
        new_branch(
            new_branch(new_leaf(5), new_leaf(4)), new_branch(new_leaf(3), new_leaf(2))));
    assert_eq!(btree.into_iter().collect::<Vec<_>>(), vec![9, 8, 7, 6, 5, 4, 3, 2]);
}

Output:

$ cargo test
   Compiling btree v0.1.0 (file:///path/to/btree)
     Running target/debug/btree-76ef84b71534f406

running 2 tests
test test_btree_iteration ... ok
test test_btree_creation ... ok

test result: ok. 2 passed; 0 failed; 0 ignored; 0 measured

Answer 1

Again, this code looks pretty great! A few minor style points to start:

1. Avoid using explicit typing when you don't need it. Don't say let mut node: BTree<T> = root when just let mut node = root would work. This lets you avoid visual clutter and can make refactoring less tedious.

2. Adding to that last example, you can specify that a function argument is mut so you don't need to immediately re-assign it.

   // Don't do this
   fn foo(var: u8) {
       let mut var2 = var;
       // ...
   }
   
   // When you can do this
   fn foo(mut var: u8) {
       // ...
   }
   

3. Avoid explicit return statements at the end of functions. Almost every line of code is an expression and thus has a value. We use that everywhere else, such as the result of an if or a match statement, the function body itself is no different. However, it is fine to use return if you need to exit early from a function. I'd encourage you to limit this to "guard clauses" where possible, as returning from the middle of a long function can be hard to spot.

4. Use an if let statement when you only have a single useful match arm:

   // Don't do this
   match self.right_nodes.pop() {
       Some(node) => self.add_left_subtree(node),
       _ => {},
   };
   
   // When you can do this
   if let Some(node) = self.right_nodes.pop() {
       self.add_left_subtree(node);
   }
   

5. I prefer to group all of my tests in a test module. This allows me to easily have test-specific dependencies or helper methods.

6. I'd split your tests into more granular pieces. Having one test for "iteration" belies that there are actually a few separate things you are testing.

7. I'd remove the word test from each test name. The [test] annotation clearly marks that. I'd also remove btree from the test name as it's just a bit redundant.

8. I'd fill the newly-found space in the test names with more detail about what you are testing. This goes hand-in-hand with more granular tests. It also is nice that when the tests fail, the name of the test will help you understand exactly what broke.

9. I'd avoid extra collect calls when you don't need them. Even though it's just in tests, it is an extra allocation that is easy enough to avoid. In this case, I'd also bring in the itertools crate, specifically to get the assert_equal function.

The one functional change I'd make is to avoid the panic! in the iterator implementation. It's often nicer if we can encode requirements like that into the code itself. To do this, I'd suggest keeping current_node as an Option<T> instead. Unfortunately, this makes the while let loop a bit uglier. I think it's worth it to avoid the invalid state though.

And to your questions:

What's the proper way to have [BTree::new_leaf] namespaced and yet useable?

I'm not 100% sure I'm following, but making these constructors worked just fine. If you are asking about having a shorter form to access them, you can do something like let mk_leaf = BTree::new_leaf. See the last test for an example of this.

Is the shadowing of btree in the test case considered bad form?

In general, shadowing of variables isn't completely a bad thing in Rust. There are times where you are transforming a variable right at the beginning and that's fine. Shadowing gets more ugly when it surprises you.

In this case, I think that the shadowing is a side-effect of your tests not being granular enough.

Stylistically, when are expr and return expr; preferred? Is there a technical difference?

There is no technical difference other than return expr; works when it is not the final statement in the block. Prefer expr whenever it works.

Is panic! the best way to handle invariant violations?

The best way is to try to prevent them from happening. Beyond that, I'd normally use unreachable! to indicate these types of errors.

---

/// A binary tree with data only at the leaves.
#[derive(Debug, PartialEq, Clone)]
enum BTree<T> {
    Leaf(T),
    Branch(Box<BTree<T>>, Box<BTree<T>>),
}

impl<T> BTree<T> {
    /// This makes it slightly less painful to put things in boxes.
    fn new_branch(left: BTree<T>, right: BTree<T>) -> BTree<T> {
        BTree::Branch(Box::new(left), Box::new(right))
    }

    /// For consistency with `new_branch`.
    fn new_leaf(elem: T) -> BTree<T> {
        BTree::Leaf(elem)
    }
}

impl<T> IntoIterator for BTree<T> {
    type Item = T;
    type IntoIter = BTreeIterator<T>;

    fn into_iter(self) -> BTreeIterator<T> {
        BTreeIterator::new(self)
    }
}

/// Iterator type for a binary tree.
/// This is a generator that progresses through an in-order traversal.
struct BTreeIterator<T> {
    right_nodes: Vec<BTree<T>>,
    current_node: Option<T>,
}

impl<T> BTreeIterator<T> {
    fn new(node: BTree<T>) -> BTreeIterator<T> {
        let mut iter = BTreeIterator { right_nodes: vec![], current_node: None };
        iter.add_left_subtree(node);
        iter
    }

    /// Consume a binary tree node, traversing its left subtree and
    /// adding all branches to the right to the `right_nodes` field
    /// while setting the current node to the left-most child.
    fn add_left_subtree(&mut self, mut node: BTree<T>) {
        loop {
            match node {
                BTree::Branch(left, right) => {
                    self.right_nodes.push(*right);
                    node = *left;
                },
                BTree::Leaf(val) => {
                    self.current_node = Some(val);
                    break;
                }
            }
        }
    }
}

impl<T> Iterator for BTreeIterator<T> {
    type Item = T;

    fn next(&mut self) -> Option<T> {
        // Get the item we're going to return.
        let result = self.current_node.take();

        // Now add the next left subtree
        // (this is the "recursive call")
        if let Some(node) = self.right_nodes.pop() {
            self.add_left_subtree(node);
        }

        result
    }
}

mod test {
    extern crate itertools;
    use super::BTree;

    #[test]
    fn creation() {
        BTree::new_leaf(10);

        let branch = BTree::new_branch(BTree::new_leaf(15), BTree::new_leaf(20));
        BTree::new_branch(branch.clone(), BTree::new_leaf(30));

        assert_eq!(branch, branch.clone());
    }

    #[test]
    fn iteration_leaf() {
        itertools::assert_equal(BTree::new_leaf(123), vec![123]);
    }

    #[test]
    fn iteration_branch() {
        let btree = BTree::new_branch(BTree::new_leaf(10), BTree::new_branch(BTree::new_leaf(20), BTree::new_leaf(30)));
        itertools::assert_equal(btree, vec![10, 20, 30]);
    }

    #[test]
    fn iteration_something_else() {
        let new_branch = BTree::new_branch;
        let new_leaf = BTree::new_leaf;

        let btree = new_branch(
            new_branch(
                new_branch(new_leaf(9), new_leaf(8)), new_branch(new_leaf(7), new_leaf(6))),
            new_branch(
                new_branch(new_leaf(5), new_leaf(4)), new_branch(new_leaf(3), new_leaf(2))));

        itertools::assert_equal(btree, vec![9, 8, 7, 6, 5, 4, 3, 2]);
    }
}

---

Your next task, should you choose to accept it... a non-consuming iterator! ^_^