Implementing Iterator for Vec

help
6

It is possible in safe Rust, by utilizing split_mut or split_first_mut or the like. I feel it's an illustrative example, too.

First, when I look at your current example, you're holding onto the entire &mut [T] while handing out mutable references to its contents. You should know that even if you never use both of them, having two &mut to the same memory is instant undefined behavior (UB). You're going to need to "give up" your &mut access to the parts of the slice you return somehow.

One way would be to use *mut and some unsafe when returning values. But split_mut and friends give you a way to do it in safe code, still using &mut (by soundly encapsulating the unsafe parts in a safe API themselves). The main idea is that your &mut[T] will need to shrink as you hand out pieces of it.

split_first_mut seems ideal for next, so let's see what happens when we try to use it:

    fn next(&mut self) -> Option<&'a mut T> {
        let (first, rest) = self.vec.split_first_mut()?;
        self.vec = rest;
        Some(first)
    }

Playground. It doesn't work -- for the same reasons as your original code! You're calling split_first_mut on self.vec by going through the mutable reference on the parameter. More explicitly:

    // `Self` = `VecIter<'a, T>`.  `'next` may be much smaller than `'a`
    fn next(&'next mut self) -> Option<&'a mut T> {

So you can't get a long enough lifetime to return by going through &'next mut self.

This is where the really illustrative parts comes in. You cannot move non-Copy values (like your &mut [T]) out from beneath an outer &mut reference, as that could leave the underlying struct in an invalid state. So in order for this to work in safe code, you'll have to first get the &mut [T] out from behind your &mut self.

And fortunately, there is a safe way to do so:

    fn next(&mut self) -> Option<&'a mut T> {
        // Temporarily remove the slice from behind &mut self
        let slice = std::mem::take(&mut self.vec);
        // Now we can get a long enough lifetime to return
        let (first, rest) = slice.split_first_mut()?;
        // And put the `rest` back where we found it
        self.vec = rest;
        // 🎉
        Some(first)
    }

Playground. There you go, an iterator over mutable references without using any unsafe yourself.

So, how does it work? The key line is

        let slice = std::mem::take(&mut self.vec);

Which could have also been written as

        let slice = std::mem::replace(&mut self.vec, &mut []);

And it's possible because this is a special case that Rust is aware of: a &mut to no memory at all (such as a &mut[T] of zero length) is completely safe, so you can conjure one out of thin air (like I did with replace), and &mut[T] implements Default to produce an empty slice as well (which is why std::mem::take also works).

You can read more about this "swap out, modify, replace" pattern in Niko's post, moving from borrowed data and the sentinel pattern.

12 Likes