Creating an Iterator in Rust

To be able to externally iterate over the tree, we need to implement the Iterator trait. It looks roughly like this:

trait Iterator { type Item; fn next(&mut self) -> Option<Self::Item>;
}

The iterator trait is usually not implemented for a collection directly. Instead, a new type is created that wraps the collection:

struct NodeIter<'a, It>(&'a Tree<It>);

As the immutable reference indicates, this struct can only iterate over immutable references of the items. To get an instance of this type, we add a .iter() method:

impl<It> Node<It> { fn iter(&self) -> NodeIter<'_, It> { NodeIter(self) }
}

Now we can start with the actual implementation!

impl<'a, It> Iterator for NodeIter<'a, It> { type Item = &'a It; fn next(&mut self) -> Option<Self::Item> { todo!() }
}

Now there’s a problem: Since Iterator provides external iteration, we have to produce one item at a time. This means that we can’t use the simple recursive algorithm we used in the traverse method. Instead, we have to keep track of the state of the iterator manually:

struct NodeIter<'a, It> { children: &'a [Node<It>], parent: Option<Box<NodeIter<'a, It>>>,
}

The children field contains the remaining children of a node, the parent field is the iterator of the parent node, if present. It must be wrapped in a Box because a struct in Rust can’t contain itself without indirection – otherwise, it would be impossible to compute its size on the stack.

So how does this work? When we create the iterator, we put the node into the children slice and set parent to None:

impl<It> Node<It> { fn iter(&self) -> NodeIter<'_, It> { NodeIter { children: std::slice::from_ref(self), parent: None, } }
}

When the iterator is advanced, we first check if children is empty. If that’s the case, we try to continue iterating the parent node. If there is no parent node, we return None.

If children is not empty, we remove the first child and check its variant. If it is a Node::Leaf, we return its content; if it is a Node::Children, we create a new iterator for the children. The parent field is set to self, and self is replaced with the newly created iterator:

use std::mem; impl<'a, It> Iterator for NodeIter<'a, It> { type Item = &'a It; fn next(&mut self) -> Option<Self::Item> { match self.children.get(0) { None => match self.parent.take() { Some(parent) => { // continue with the parent node *self = *parent; self.next() } None => None, }, Some(Node::Leaf(item)) => { self.children = &self.children[1..]; Some(item) } Some(Node::Children(children)) => { self.children = &self.children[1..]; // start iterating the child trees *self = NodeIter { children: children.as_slice(), parent: Some(Box::new(mem::take(self))), }; self.next() } } }
}

This doesn’t work yet, because mem::take() requires that NodeIter implements Default. But this can be fixed easily:

impl<It> Default for NodeIter<'_, It> { fn default() -> Self { NodeIter { children: &[], parent: None } }
}

The mem::take() function

mem::take() replaces a mutable reference with its default value and returns the previous value. The previous value is effectively moved out of the reference. We use it here to convert &mut self to an owned value, because parent must be owned.

Now let’s see if the iterator works!