Implementing an Iterator of mutable references

help
2

Your chaining approach was correct. Unfortunately, implementing this directly is impossible without using unsafe because rust can't statically guarantee that you won't call get_mut twice with the same index (getting two mutable references to the same item).

use std::iter;
use std::slice::IterMut;

struct NE<T> {
    head: T,
    rest: Vec<T>
}


impl<T> NE<T> {
    pub fn get(&self, idx: usize) -> Option<&T> {
        if idx == 0 {
            Some(&self.head)
        } else {
            self.rest.get(idx - 1)
        }
    }

    pub fn get_mut(&mut self, idx: usize) -> Option<&mut T> {
        if idx == 0 {
            Some(&mut self.head)
        } else {
            self.rest.get_mut(idx - 1)
        }
    }
    pub fn iter_mut(&mut self) -> NEIterMut<T> {
        NEIterMut {
            iter: iter::once(&mut self.head).chain(self.rest.iter_mut()),
        }
    }
}
pub struct NEIterMut<'a, T: 'a> {
    iter: iter::Chain<iter::Once<&'a mut T>, IterMut<'a, T>>,
}

impl<'a, T> Iterator for NEIterMut<'a, T> {
    type Item = &'a mut T;

    #[inline]
    fn next(&mut self) -> Option<&'a mut T> {
        self.iter.next()
    }
    #[inline]
    fn size_hint(&self) -> (usize, Option<usize>) {
        self.iter.size_hint()
    }
}

Note: You don't have to use chain, once, etc. You could also define NEIterMut as:

pub struct NEIterMut<'a, T: 'a> {
    head: Option<&'a mut T>,
    rest: IterMut<'a, T>,
}

and manually implement the iterator protocol.

3 Likes