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?

Thank you!

Answer 1

Here's a way of having a mutable iterator over a hypothetical Point struct. Pay special attention to the large comment that I put on the unsafe block. 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!

use std::mem;

#[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 => Some(&mut self.point.x),
            1 => Some(&mut self.point.y),
            2 => Some(&mut self.point.z),
            _ => None
        };

        if retval.is_some() {
            self.idx += 1;
        }

        // This is safe because...
        // (from http://stackoverflow.com/questions/25730586):
        // The Rust compiler does not know that when you ask a mutable
        // iterator for the next element, that you get a different
        // reference every time and never the same reference twice. Of
        // course, we know that such an iterator won't give you the
        // same reference twice.
        unsafe { mem::transmute(retval) }
    }
}

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

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

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