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::

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?