How to correctly build a mutable iterator over borrowed data

Hi everyone,
I'm new to this forum and semi-new to Rust and currently battling with the borrow checker and Miri.
My final goal is to build a tree data structure where nodes are stored in a central shared Vec and refer to each other with index-like handle objects. Now I want to build the iter() and iter_mut() functions for said container, i.e. I need to hand out references into the vec based on indices within the iterator. iter() is fine, but of course iter_mut() is the hard part.

Coming from C++, my first draft was as follows: (using an iterator built via a .map() function here as a proof of concept before building the proper type)

fn main() {
    let mut source = vec![1, 2, 3, 4];
    let index = vec![0, 3, 1, 2];

    let it = index.iter()
        .map(|i| &mut source[*i]); // fails borrowck
    //...
}

This fails the borrow checker, and I believe I by now even understood why (correct me if I'm wrong):
The Iterator trait requires the Item type to have a named lifetime 'i that is different from the one next<'b>(&'b self) is called with, so the body of next(), or in this example my lambda, must ensure the returned reference outlives its source due to the &'short&'long -> &'short lifetime coercion. (Correct?)

To fix it I tried the following, basically masking the original borrow's lifetime with a raw pointer deref:

.map(|i| unsafe { &mut *(&raw mut source[*i]) }); // passes borrowck, but not Miri

(I guess I could have used std::mem::transmute() as well, what is the more idiomatic way to do it?)
These pointer shenanigans at least make the borrow checker happy, but they still aren't enough for Miri. When using a special function I found online that collect()s the mutable iterator and then std::mem::swap()s the returned references a bunch, Miri reports the following error:

error: Undefined Behavior: trying to retag from <930> for SharedReadWrite permission at alloc238[0x0], but that tag does not exist in the borrow stack for this location
  --> src/main.rs:19:31
   |
19 |         std::mem::swap(dummy, r);
   |                               ^ this error occurs as part of two-phase retag at alloc238[0x0..0x4]
   |
   = help: this indicates a potential bug in the program: it performed an invalid operation, but the Stacked Borrows rules it violated are still experimental
   = help: see https://github.com/rust-lang/unsafe-code-guidelines/blob/master/wip/stacked-borrows.md for further information
help: <930> was created by a Unique retag at offsets [0x0..0x4]
  --> src/main.rs:17:28
   |
17 |     let mut refs: Vec<_> = iterator.collect();
   |                            ^^^^^^^^^^^^^^^^^^
help: <930> was later invalidated at offsets [0x0..0x10] by a Unique retag
  --> src/main.rs:7:49
   |
 7 |         .map(|i| unsafe { &mut *(&raw mut source[*i]) }); // fails Miri
   |                                                 ^^^^
   = note: BACKTRACE (of the first span):
   = note: inside `check_all_refs::<'_, i32, std::iter::Map<std::slice::Iter<'_, usize>, {closure@src/main.rs:7:14: 7:17}>>` at src/main.rs:19:31: 19:32
note: inside `main`
  --> src/main.rs:11:5
   |
11 |     check_all_refs(&mut 0, it);
   |     ^^^^^^^^^^^^^^^^^^^^^^^^^^

note: some details are omitted, run with `MIRIFLAGS=-Zmiri-backtrace=full` for a verbose backtrace

error: aborting due to 1 previous error

(Playground link at the end if you want to see exactly what I did)
Now I have to admit, I'm still pretty bad at understanding Miri errors, but it appears to me that it's taking issue with me borrowing the entire slice of the vec each time in order to take out an item's reference in the lambda in line 7. (Is that actually what's happening here?)
Now the question is, how to avoid this? Or is this "safe enough" for production and Miri is just overzealous? Note that the usual solutions to this I found online, such as using split_off_first_mut() would be quite complicated to implement, since I require access by index in a random order (no duplicates of course, so it's still safe to do at all without a true lending iterator).
I tried the following workaround to get Miri to shut up, but this seems unsafer and especially less maintainable to me than the version above:

.map(|i| unsafe { 
    assert!(*i < source.len());
    &mut *source.as_mut_ptr().add(*i) 
});

So what would be the cleanest way to do this? Am I overlooking an obvious simple solution?
Since all three of the above functions would probably compile to exactly the same bytecode, I wonder if the tools are perhaps a bit too eager in this case.

I'm looking forward to your enlightening answers!
/ J. Aetherwing

As promised, the playground link to my experiment:

It is possible to know that this will not borrow check while not knowing anything about the specific properties of Iterator. Here are two reasons for that:

  • If |i| &mut source[*i] were called twice with the same i, then it would return two overlapping exclusive borrows. That is prohibited, therefore this function can only be FnOnce, callable only once — but we know that map() calls the function it is given more than once, so this can't compile.

  • In general, when writing a iterator that produces mutable references, the solution will never involve use of &mut x[y], because &mut x[y]:

    • exclusively borrows all of x, and
    • does not know that the ys are non-duplicate keys

The way I would describe this is that the Iterator trait requires that the item type be a single type per implementation. That is, whenever you impl Iterator for ..., you must specify type Item = <some specific type>. That, by itself, is sufficient to completely exclude the possibility of borrowing from &'b mut Self, because 'b is not a single lifetime which exists before / outside the scope of the impl Iterator, so it cannot possibly be part of a single Item type.

For any Iterator, you can always ask what its ::Item type is, and the answer cannot depend on any borrowing of the iterator because no borrowing of the iterator was involved in asking the question.

This is correct. In particular, source[*i] calls DerefMut::deref_mut() which requires a mutable reference to source. (This is arguably a bit of a gap in Rust’s foundational traits — there is no way to offer indexing that is not a borrow.)

There are very, very few practical situations where ‘Miri is overzealous’. Outside of very particular situations of which this is not one, if Miri fails, you should assume your code is incorrect.

For random access, you should:

  1. construct a raw *mut [T] pointer when you construct the itterator
  2. use get_unchecked() to get an element pointer without involving any references
  3. convert that pointer to &mut T
3 Likes

Thank you for the detailled explanation, I found the alternative view on why Iterator is defined the way it is very insightful.

Looks like overall I got close to the right thing, but just needed to find the "best" std methods to do it :slight_smile: