Simple as possible example of returning a mutable reference from your own iterator

Question

This question is related, however it moreso covers the reason why the compiler cannot infer a safe lifetime when returning a mutable reference from Iterator::next, which I think I understand.

My question is:

What are the specific steps you can take when designing your own iterator so that it can produce mutable references? Ultimately, I'm hoping for a concise-as-possible, step-by-step, commented example of both an Iterator and its next implementation that I (and anyone) can go to as a clear reference when they run into this situation. unsafe examples are fine, I imagine they are probably necessary!

NOTE: I understand that MutItems is normally the recommended example, however its implementation can be difficult to follow as there isn't any documentation on 1. How the markers work in that situation and 2. What the iterator! macro expands to and how it works. If you use MutItems as your example could you please clarify these things?

Answer 1

Here's a way of having a mutable iterator over a hypothetical Point struct. I find it very useful to annotate every unsafe block in a similar fashion, since I'm only shooting myself in the foot if I get it wrong!

The Rust compiler does not know that you will get a different mutable reference every time you advance the iterator. This unsafe block is safe because the programmer guarantees that this iterator can never return the same mutable reference twice, or allow any other way to get to the same address.

#[derive(Debug)]
struct Point {
    x: u8,
    y: u8,
    z: u8,
}

impl Point {
    fn iter_mut(&mut self) -> IterMut {
        IterMut {
            point: self,
            idx: 0,
        }
    }
}

struct IterMut<'a> {
    point: &'a mut Point,
    idx: u8,
}

impl<'a> Iterator for IterMut<'a> {
    type Item = &'a mut u8;

    fn next(&mut self) -> Option<&'a mut u8> {
        let retval = match self.idx {
            0 => &mut self.point.x,
            1 => &mut self.point.y,
            2 => &mut self.point.z,
            _ => return None,
        };

        self.idx += 1;

        // I copied this code from Stack Overflow without paying attention to
        // the prose which described why this code is actually safe.
        unsafe { Some(&mut *(retval as *mut u8)) }
    }
}

fn main() {
    let mut p1 = Point { x: 1, y: 2, z: 3 };

    for x in p1.iter_mut() {
        *x += 1;
    }

    println!("{:?}", p1);
}

See also

• How can I create my own data structure with an iterator that returns mutable references?