How to implement Iterator yielding mutable references [duplicate]

I am trying to implement a simple lookup iterator:

pub struct LookupIterMut<'a, D> {
    data : &'a mut [D],
    indices : &'a [usize],
    i: usize
}

impl<'a, D> Iterator for LookupIterMut<'a, D> {
    type Item = &'a mut D;

    fn next(&mut self) -> Option<Self::Item> {
        if self.i >= self.indices.len() {
            None
        } else {
            let index = self.indices[self.i] as usize;
            self.i += 1;
            Some(&mut self.data[index]) // error here
        }
    }
}

The idea was to allow a caller consecutive mutable access to an internal storage. However I am getting the error cannot infer an appropriate lifetime for lifetime parameter in function call due to conflicting requirements.

As far as I understand I would have to change the function signature to next(&'a mut self) -> .. but this would not be an Iterator anymore.

I also discovered that I could simply use raw pointers, though I am not sure if this is appropriate here:

// ...
type Item = *mut D;
// ...

Thanks for your help

Solution 1:

Your code is invalid because you try to return multiple mutable references to the same slice with the same lifetime 'a.

For such a thing to work, you would need a different lifetime for each returned Item so that you wouldn't hold 2 mutable references to the same slice. You cannot do that for now because it requires Generic Associated Types:

type Item<'item> = &'item mut D; // Does not work today

One solution is to check that the indices are unique and to rebind the lifetime of the referenced item to 'a in an unsafe block. This is safe because all the indices are unique, so the user cannot hold 2 mutable references to the same item.

Don't forget to encapsulate the whole code inside a module, so that the struct cannot be build without the check in new:

mod my_mod {
    pub struct LookupIterMut<'a, D> {
        data: &'a mut [D],
        indices: &'a [usize],
        i: usize,
    }

    impl<'a, D> LookupIterMut<'a, D> {
        pub fn new(data: &'a mut [D], indices: &'a [usize]) -> Result<Self, ()> {
            let mut uniq = std::collections::HashSet::new();
            let all_distinct = indices.iter().all(move |&x| uniq.insert(x));

            if all_distinct {
                Ok(LookupIterMut {
                    data,
                    indices,
                    i: 0,
                })
            } else {
                Err(())
            }
        }
    }

    impl<'a, D> Iterator for LookupIterMut<'a, D> {
        type Item = &'a mut D;

        fn next(&mut self) -> Option<Self::Item> {
            self.indices.get(self.i).map(|&index| {
                self.i += 1;

                unsafe { std::mem::transmute(&mut self.data[index]) }
            })
        }
    }
}

Note that your code will panic if one index is out of bounds.

Solution 2:

Using unsafe

Reminder: it is unsound to have, at any time, two accessible mutable references to the same underlying value.

The crux of the problem is that the language cannot guarantee that the code abides by the above rule, should indices contain any duplicate, then the iterator as implemented would allow obtaining concurrently two mutable references to the same item in the slice, which is unsound.

When the language cannot make the guarantee on its own, then you either need to find an alternative approach or you need to do your due diligence and then use unsafe.

In this case, on the Playground:

impl<'a, D> LookupIterMut<'a, D> {
    pub fn new(data: &'a mut [D], indices: &'a [usize]) -> Self {
        let set: HashSet<usize> = indices.iter().copied().collect();
        assert!(indices.len() == set.len(), "Duplicate indices!");

        Self { data, indices, i: 0 }
    }
}

impl<'a, D> Iterator for LookupIterMut<'a, D> {
    type Item = &'a mut D;

    fn next(&mut self) -> Option<Self::Item> {
        if self.i >= self.indices.len() {
            None
        } else {
            let index = self.indices[self.i];
            assert!(index < self.data.len());

            self.i += 1;

            //  Safety:
            //  -   index is guaranteed to be within bounds.
            //  -   indices is guaranteed not to contain duplicates.
            Some(unsafe { &mut *self.data.as_mut_ptr().offset(index as isize) })
        }
    }
}

Performance wise, the construction of a HashSet in the constructor is rather unsatisfying but cannot really be avoided in general. If indices was guaranteed to be sorted for example, then the check could be performed without allocation.