Implementing an Iterator of mutable references

help
1

I have a custom data structure which is essentially a non-empty vector:

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)
        }
    }
}

I was able to create an Iterator for this type:

pub struct NEIter<'a, T: 'a> {
    ne: &'a NE<T>,
    current: usize,
}

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

    #[inline]
    fn next(&mut self) -> Option<&'a T> {
        let r = self.ne.get(self.current);
        self.current += 1;
        r
    }
}

The problem comes when I want to implement an iterable of mutable references, similar to how Vec provides .iter_mut(). That is, I want my iterator items to be &'a mut T.

I've tried a number of solutions, including by trying to implement my own custom NEIterMut with associated Iterator impl, as well as trying to implement it by chaining together a Some::iter_mut with a Vec::iter_mut, but I can't seem to get it to work. I have tried many different things and they all seemed like dead ends, so right now I'm just looking for general advice: is this even possible, given my current data structure?

I see unsafe blocks in the implementation of the Iterator for Vec in the stdlib; is it possible to implement a mutable iterator for this datatype using safe rust? I have not been able to find any applicable documentation for how to implement mutable iterators.

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

Thank you very much, @stebalien. I didn't know about iter::once, I was trying Some::iter / iter_mut and was struggling with it... though it might work too. I need to spend some time figuring out what the differences are between my attempts and your successful version. I'll also probably implement the manual version to see what I need to do.

I guess your strategy here is to take advantage of Vec's Iterator so we don't need to implement the unsafe stuff ourselves. My own non-chaining attempts were using my get_mut directly instead of using Vec's iterator, so that was probably my main problem.