Skip to content

Instantly share code, notes, and snippets.

@jesperdj
Last active June 29, 2020 22:31
Show Gist options
  • Save jesperdj/466ff014adce89a1405ce033ff92231a to your computer and use it in GitHub Desktop.
Save jesperdj/466ff014adce89a1405ce033ff92231a to your computer and use it in GitHub Desktop.
Binary tree with depth-first iterator in Rust.
// 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);
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment