Iterator of mutable references over a hashmap

help
1

Hi,

I have a HashMap<usize, Foo> that I want to iterate over in an arbitrary order that is set independently from the HashMap internal state. For context, this is for modelling, so I need the iteration order to be reproducible (for testing) and to be able to shuffle it with an Rng to avoid order effects when iterating over my Foos. I therefore cannot just use HashMap::values() and HashMap::values_mut() whose iteration order is not guaranteed nor can be shuffled.

For convenience, I was hoping to be able to bundle the HashMap together with a Vec<usize> that determines the iteration order. Implementing iter() following that vector is trivial, but iter_mut() is a different beast...

Here's what I attempted, and am not sure at all is sound:

use std::{collections::HashMap, mem};

use crate::foo::Foo;

mod foo {
    use std::sync::atomic::{AtomicUsize, Ordering};

    static ID: AtomicUsize = AtomicUsize::new(0);

    pub struct Foo(usize);

    impl Foo {
        /// Ensure each `Foo` has its own unique index
        pub fn new() -> Foo {
            Foo(ID.fetch_add(1, Ordering::Relaxed))
        }

        pub fn id(&self) -> usize {
            self.0
        }
    }
}

/// Mutable iterator over the collection of `Foo`s
struct BarIterMut<'a> {
    index: usize,
    bar: &'a mut Bar,
}

impl<'a> Iterator for BarIterMut<'a> {
    type Item = &'a mut Foo;

    fn next(&mut self) -> Option<Self::Item> {
        if self.index < self.bar.iter_ids.len() {
            let id = self.bar.iter_ids[self.index];
            self.index += 1;
            let foo = self.bar.foos.get_mut(&id).unwrap();
            let foo = unsafe {
                // SAFETY: we're holding a &'a mut Foo, and iter_ids
                // does not have any duplicate by construction.
                // Therefore, all "foo"s here are disjoint and live
                // at least for 'a.
                mem::transmute(foo)
            };
            Some(foo)
        } else {
            None
        }
    }
}

/// The collection of `Foo`s
#[derive(Default)]
pub struct Bar {
    /// The vector dictating the iteration order
    iter_ids: Vec<usize>,
    /// The map itself
    foos: HashMap<usize, Foo>,
}

impl Bar {
    pub fn insert(&mut self, foo: Foo) {
        let id = foo.id();
        self.iter_ids.push(id);
        self.foos.insert(id, foo);
    }

    pub fn foos(&self) -> impl Iterator<Item = &Foo> {
        self.iter_ids.iter().map(|&id| self.foos.get(&id).unwrap())
    }

    pub fn foos_mut(&mut self) -> impl Iterator<Item = &mut Foo> {
        BarIterMut {
            index: 0,
            bar: self,
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn iter_mut() {
        let mut bar = Bar::default();
        let nfoos = 100;
        for _ in 0..nfoos {
            bar.insert(Foo::new());
        }

        // for miri to catch potential UB if an &mut is somehow duplicated
        let mut_refs: Vec<_> = bar.foos_mut().enumerate().collect();
        assert_eq!(mut_refs.len(), nfoos);

        // check the iteration happens in the expected order
        for (i, foo) in mut_refs {
            assert_eq!(foo.id(), i);
        }
    }
}

Transmuting the &mut to an &'a mut feels dirty and I'm not 100% sure the reasoning in the safety comment is correct.

Running the test through miri yields the following error:

running 1 test
test test::iter_mut ... error: Undefined Behavior: trying to retag from <252046> for Unique permission at alloc63906[0x428], but that tag does not exist in the borrow stack for this location
   --> /home/adrien/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/alloc/src/vec/into_iter.rs:269:23
    |
269 |         Some(unsafe { ptr.read() })
    |                       ^^^^^^^^^^
    |                       |
    |                       this error occurs as part of retag at alloc63906[0x428..0x430]
    |                       while retagging field .1
    |
    = 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: <252046> was created by a Unique retag at offsets [0x428..0x430]
   --> src/lib.rs:86:32
    |
 86 |         let mut_refs: Vec<_> = bar.foos_mut().enumerate().collect();
    |                                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
help: <252046> was later invalidated at offsets [0x420..0x430] by a SharedReadOnly retag
   --> src/lib.rs:35:23
    |
 35 |             let foo = self.bar.foos.get_mut(&id).unwrap();
    |                       ^^^^^^^^^^^^^^^^^^^^^^^^^^
    = note: this is on thread `test::iter_mut`
    = note: stack backtrace:
            0: <std::vec::IntoIter<(usize, &mut foo::Foo)> as std::iter::Iterator>::next
                at /home/adrien/.rustup/toolchains/nightly-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/alloc/src/vec/into_iter.rs:269:23: 269:33
            1: test::iter_mut
                at src/lib.rs:88:25: 88:33
            2: test::iter_mut::{closure#0}
                at src/lib.rs:80:18: 80:18

I'm tempted to trust miri over myself on this, but this leaves me with the following questions:

  • If my code is actually sound, what's tripping miri here?
  • If miri is indeed correct that there is an UB, what's wrong with my reasoning?
  • Is there a way to achieve what I want besides keeping iter_ids separate from the hashmap and loop over iter_ids directly and use get_mut to obtain my &mut Foos? There is nothing wrong/blocking with that for my use case, it's just a little clunky and unergonomic (it also makes it slightly harder to keep consistency between the content of the map and the content of the Vec as the set of Foos changes...).

Thanks!

2

this is impossible to implement on top of HashMap public api.

each time you call HashMap::get_mut(), it invalidates previously yielded items, because of the &mut self parameter: it claims exclusiveness of the entire HashMap, ultimately, that's what this miri message means:

in rust, a mut borrow is not just permission of mutability, more importantly, it carries a much stronger guarantee: the reference is exclusive, a.k.a. noalias in LLVM term.

you can only implement an iterator that yields &mut Value if you have access to the underlying hash buckets, which you don't.

if this is only for testing, you can collect the entries iterator into an intermediate Vec then sort them according your desired order:

// if you prefer a named type instead of opaque
struct BarIterMut<'a> {
    Vec<hash_map::IterMut<'a, usize, Foo>>,
}

impl Bar {
    fn foos_mut(&mut self) -> impl Iterator<Item = &'_ mut Foo> + '_ {
        let mut entries = self.foos.iter_mut().collect::<Vec<_>>();
        // build a reverse key to index map from `self.iter_ids`
        let orders = todo!();
        entries.sort_by_key(|(&key, _value)| orders[key]);
        entries.into_iter().map(|(_key, value)| value)
    }
}
3

Ah right, of course... I don't know why I didn't see that.

Thinking a bit more about it, it's acceptable for self.bar.foos.get_mut(&id) to internally move stuff around anyway, so I have no right to assume the references are still valid...

I'll just keep my vector of indices dictating the iteration order separated from the hashmap then.

4

another options is to use an alternative hash map implementation which allows you to control the iteration order. for example, indexmap allows you to iterate through the map in its original insertion order.