Mutable iterator witchcraft

1

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
2

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.

3

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
4

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
5

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
6

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

7

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.