rustiteratorclosures

Create mutable iterator over a set of indices?


I'm trying to implement an array where items can occupy multiple slots (details not relevant to question). I have created a method that allows me to iterate over a range of items:

pub fn get(&self, index: I) -> Option<&T>;

pub fn iter_range(&self, index: I, size: usize) -> impl Iterator<Item = &T> {
    std::iter::successors(Some(index), |i| i.get_next())
        .take(size)
        .filter_map(|i| self.get(i))
}

Note here that the index is a generic and the only way to get the "next" index in the range is to call the get_next method returning an Option<I>. I use std::iter::successors to generate the indices which is then used to get the items.

Next, I try to make a mutable version of this method:

pub fn get_mut(&mut self, index: I) -> Option<&mut T>;

pub fn iter_range_mut(&mut self, index: I, size: usize) -> impl Iterator<Item = &mut T> {
    std::iter::successors(Some(index), |i| i.get_next())
        .take(size)
        .filter_map(|i| self.get_mut(i))
}

But I get this compiler error:

error: captured variable cannot escape `FnMut` closure body
   --> foo\bar.rs:306:29
    |
303 |     pub fn iter_range_mut(&mut self, index: I, size: usize) -> impl Iterator<Item = &mut T)> {
    |                           --------- variable defined here
...
306 |             .filter_map(|i| self.get_mut(i))
    |                           - ----^^^^^^^^^^^
    |                           | |
    |                           | returns a reference to a captured variable which escapes the closure body       
    |                           | variable captured here
    |                           inferred to be a `FnMut` closure
    |
    = note: `FnMut` closures only have access to their captured variables while they are executing...
    = note: ...therefore, they cannot allow references to captured variables to escape

I understand why I get this error, but not how to circumvent it. I am making a mutable iterator so I expect having to write unsafe code, but how do I do so here? Preferably I wouldn't have to create a custom iterator type.


Solution

  • You cannot implement this iterator in terms of fn get_mut(&mut self, index: I) -> Option<&mut T>. The &mut self part of the signature implies exclusive access to the entire collection, which contradicts being able to have more than one of these at once (note that every Iterator must be collect()able).

    What you need to do is write code which fetches a slice reference, &mut [T], of the appropriate range of the underlying array(/slice/vector), then do an ordinary iteration of that slice reference. If the data isn't stored contiguously then you may be able to use slice::get_disjoint_mut() to get all the parts simultaneously, but this requires the number of parts to be known at compile time.

    If the number of parts cannot be known at compile time, then you have to use unsafe code and raw pointers — instead of &mut references. The details will depend on the organization of your data structure, but the key is that you can use an arbitrary number of possibly-overlapping *mut Ts and *mut [T]s, and convert them into non-overlapping &mut Ts, as long as you do not create any further &mut Selfs during this process.

    If you do this, then either the trait containing your get_next() function must be an unsafe trait (because implementors can cause unsoundness by returning overlapping indices), or you must implement your own check that none of the indices are the same.