Mutable iterator witchcraft

I wanted to share an implementation of a mutable iterator that I'm particularly happy with, I'm more used to having these things turn into inscrutable spaghetti piles:

pub struct ObjectTreeContentsMutIterator<'t, O: CircleBoundedObject> {
    unvisited_objects_in_current_node: IterMut<'t, O>,
    unvisited_nodes: Vec<&'t mut ObjectTree<O>>,
}

impl<'t, O: CircleBoundedObject> Iterator for ObjectTreeContentsMutIterator<'t, O> {
    type Item = &'t mut O;
    fn next(&mut self) -> Option<Self::Item> {
        loop {
            if let Some(next_object) = self.unvisited_objects_in_current_node.next() {
                // If we still have unvisited objects in the current node, return one.
                return Some(next_object);
            } else if let Some(next_node) = self.unvisited_nodes.pop() {
                // Otherwise, move on to the next unvisited node.
                self.unvisited_objects_in_current_node = next_node.data.0.iter_mut();
                let child_refs = next_node.children.iter_mut().filter_map(Option::as_mut);
                self.unvisited_nodes.extend(child_refs);
            } else {
                // If there are no more nodes to visit, the iterator is done.
                return None;
            }
        }
    }
}

The end result is an iterator over all objects in a tree structure, where each node in the tree can contain a vector of objects and can Optionally have children in each of a fixed number of slots. It's so satisfying to me that this requires no unsafe code.

4 Likes

I was also able to turn this into an improved version of the corresponding immutable iterator because making this mutable version forced me to find more elegant solutions.

More compiler wizardry, the method that creates this iterator for a particular tree:

pub fn iter_mut(&mut self) -> ObjectTreeContentsMutIterator<O> {
    ObjectTreeContentsMutIterator {
        unvisited_objects_in_current_node: [].iter_mut(),
        unvisited_nodes: vec![self],
    }
}

[] can be static data so it automatically outlives &mut self!

2 Likes

The fact that you can create static mutable references to empty slices is indeed some occasionally really useful compiler magic :grinning_face_with_smiling_eyes:

4 Likes

Indeed! Although I've always found it weird that that doesn't work for other ZST literals:

let _: &'static mut _ = &mut []; // Works
let _: &'static mut _ = &mut (); // Doesn't work

With ::alloc one can emulate that without unsafe by doing Box::leak(Box::new( Zst )), but without ::alloc one needs to use unsafe at some point :frowning:

4 Likes

Considered in RFC 1414 but ultimately not implemented as part of that RFC.