Mutable slice iterator

The standard library has an iterator over &mut [T] which is implemented (as of this writing) in terms of pointer arithmetic, presumably for the sake of optimization. In this example, we'll show how one can implement their own mutable slice iterator with entirely safe code.

Here's the starting place for our implementation:

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> {
    slice: &'a mut [T],
    // ...maybe other fields for your needs...
}

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

Below are a few starting attempts at it. Spoilers, they don't compile.

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        // Eh, we'll worry about iterative logic later!
        self.slice.get_mut(0)
    }
}
}
#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        // Actually, this method looks perfect for our iteration logic
        let (first, rest) = self.slice.split_first_mut()?;
        self.slice = rest;
        Some(first)
    }
}
}
#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
   type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        // 🤔 Pattern matching??
        match &mut self.slice {
            [] => None,
            [first, rest @ ..] => Some(first),
        }
    }
}
}

Yeah, the compiler really doesn't like any of that. Let's take a minute to write out all the elided lifetimes. Some of them are in aliases, which we're also going to expand:

  • Item is &'a mut T
  • &mut self is short for self: &mut Self, and
    • Self is MyIterMut<'a, T>

Here's what it looks like with everything being explicit:

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next<'s>(self: &'s mut MyIterMut<'a, T>) -> Option<&'a mut T> {
        todo!()
    }
}
}

And remember that in MyIterMut<'a, T>, slice is a &'a mut [T].

Ah, yes. We have a nested exclusive borrow here.

You cannot get a &'long mut U through dereferencing a &'short mut &'long mut U.

  • You can only reborrow a &'short mut U.

There is no safe way to go through the &'s mut self and pull out a &'a mut T.

Are we stuck then? No, there is actually a way forward! As it turns out, slices are special. In particular, the compiler understands that an empty slice covers no actual data, so there can't be any memory aliasing concerns or data races, et cetera. So the compiler understands it's perfectly sound to pull an empty slice reference out of no where, with any lifetime at all. Even if it's an exclusive slice reference!

#![allow(unused)]
fn main() {
fn magic<T>() -> &'static mut [T] {
    &mut []
}
}

For our purposes, we don't even need the magic: the standard library has a Default implementation for &mut [T].

Why does this unstick us? With that implementation, we can conjure an empty &mut [T] out of nowhere and move our slice field out from behind &mut self:

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        let mut slice = std::mem::take(&mut self.slice);
        // Eh, we'll worry about iterative logic later!
        slice.get_mut(0)
    }
}
}

std::mem::take and swap and replace are very useful and safe functions; don't be thrown off by them being in std::mem along side the dangerous transmute and other low-level functions. Note how we passed &mut self.slice -- that's a &mut &mut [T]. take replaces everything inside of the outer &mut, which can have an arbitrarily short lifetime -- just long enough to move the memory around.

So we're done aside from iterative logic, right? This should just give us the first element forever?

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
  fn next(&mut self) -> Option<Self::Item> {
      let mut slice = std::mem::take(&mut self.slice);
      slice.get_mut(0)
  }
}
let mut arr = [0, 1, 2, 3];
let iter = MyIterMut { slice: &mut arr };
for x in iter.take(10) {
    println!("{x}");
}
}

Uh, it only gave us one item. Oh right -- when we're done with the slice, we need to move it back into our slice field. We only want to temporarily replace that field with an empty slice.

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        let mut slice = std::mem::take(&mut self.slice);
        // Eh, we'll worry about iterative logic later!
        let first = slice.get_mut(0);
        self.slice = slice;
        first
    }
}
}

Uh oh, now what.

#![allow(unused)]

fn main() {
error[E0499]: cannot borrow `*slice` as mutable more than once at a time
9  |         let first = slice.get_mut(0);
   |                     ----- first mutable borrow occurs here
10 |         self.slice = slice;
   |                      ^^^^^ second mutable borrow occurs here
11 |         first
   |         ----- returning this value requires that `*slice` is borrowed for `'a`
}

Oh, right! These are exclusive references. We can't return the same item multiple times -- that would mean someone could get multiple &mut to the same element if they collected the iterator, for example. Come to think of it, we can't punt on our iteration logic either -- if we try to hold on to the entire &mut [T] while handing out &mut T to the elements, that's also multiple &mut to the same memory!

This is what the error is telling us: We can't hold onto the entire slice and return first.

(There's a pattern called "leanding iterators" where you can hand out borrows of data you own in an iterator-like fashion, but it's not possible with the current Iterator trait; it is also a topic for another day.)

Alright, let's try split_first_mut again instead, that really did seem like a perfect fit for our iteration logic.

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        let mut slice = std::mem::take(&mut self.slice);
        let (first, rest) = slice.split_first_mut()?;
        self.slice = rest;
        Some(first)
    }
}

// ...

let mut arr = [0, 1, 2, 3];
let iter = MyIterMut { slice: &mut arr };
for x in iter {
    println!("{x}");
}
}

Finally, a working version! split_first_mut is a form of borrow splitting, which we briefly mentioned before.

And for the sake of completion, here's the pattern based approach to borrow splitting:

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        //    ....these are all that changed....
        //    vvvvvvvvvvvvvvvvvv               v
        match std::mem::take(&mut self.slice) {
            [] => None,
            [first, rest @ ..] => Some(first),
        }
     }
}

// ...

let mut arr = [0, 1, 2, 3];
let iter = MyIterMut { slice: &mut arr };
for x in iter {
    println!("{x}");
}
}

Oh whoops, just one element again. Right. We need to put the rest back in self.slice:

#![allow(unused)]
fn main() {
struct MyIterMut<'a, T> { slice: &'a mut [T], }
impl<'a, T> Iterator for MyIterMut<'a, T> {
  type Item = &'a mut T;
    fn next(&mut self) -> Option<Self::Item> {
        match std::mem::take(&mut self.slice) {
            [] => None,
            [first, rest @ ..] => {
                self.slice = rest;
                Some(first)
            }
        }
     }
}

// ...

let mut arr = [0, 1, 2, 3];
let iter = MyIterMut { slice: &mut arr };
for x in iter {
    println!("{x}");
}
}

👍