Simultaneous mutable access to arbitrary indices of a large vector that are guaranteed to be disjoint
You can sort indices_to_update
and extract mutable references by calling split_*_mut
.
let len = big_vector_of_elements.len();
while has_things_to_do() {
let mut tail = big_vector_of_elements.as_mut_slice();
let mut indices_to_update = compute_indices();
// I assumed compute_indices() returns unsorted vector
// to highlight the importance of sorted order
indices_to_update.sort();
let mut elems = Vec::new();
for idx in indices_to_update {
// cut prefix, so big_vector[idx] will be tail[0]
tail = tail.split_at_mut(idx - (len - tail.len())).1;
// extract tail[0]
let (elem, new_tail) = tail.split_first_mut().unwrap();
elems.push(elem);
tail = new_tail;
}
}
Double check everything in this code; I didn't test it. Then you can call elems.par_iter(...)
or whatever.
When the compiler can't enforce that mutable references to a slice elements aren't exclusive, Cell
is pretty nice.
You can transform a &mut [T]
into a &Cell<[T]>
using Cell::from_mut
, and then a &Cell<[T]>
into a &[Cell<T>]
using Cell::as_slice_of_cells
. All of this is zero-cost: It's just there to guide the type-system.
A &[Cell<T>]
is like a &[mut T]
, if that were possible to write: A shared reference to a slice of mutable elements. What you can do with Cell
s is limited to read or replace — you can't get a reference, mutable or not, to the wrapped elements themselves. Rust also knows that Cell
isn't thread-safe (it does not implement Sync
). This guarantees that everything is safe, at no dynamic cost.
fn main() {
use std::cell::Cell;
let slice: &mut [i32] = &mut [1, 2, 3];
let cell_slice: &Cell<[i32]> = Cell::from_mut(slice);
let slice_cell: &[Cell<i32>] = cell_slice.as_slice_of_cells();
let two = &slice_cell[1];
let another_two = &slice_cell[1];
println!("This is 2: {:?}", two);
println!("This is also 2: {:?}", another_two);
two.set(42);
println!("This is now 42!: {:?}", another_two);
}
I think this is a reasonable place to use unsafe
code. The logic itself is safe but cannot be checked by the compiler because it relies on knowledge outside of the type system (the contract of BTreeSet
, which itself relies on the implementation of Ord
and friends for usize
).
In this sample, we preemptively bounds check all the indices via range
, so each call to add
is safe to use. Since we take in a set, we know that all the indices are disjoint, so we aren't introducing mutable aliasing. It's important to get the raw pointer from the slice to avoid aliasing between the slice itself and the returned values.
use std::collections::BTreeSet;
fn uniq_refs<'i, 'd: 'i, T>(
data: &'d mut [T],
indices: &'i BTreeSet<usize>,
) -> impl Iterator<Item = &'d mut T> + 'i {
let start = data.as_mut_ptr();
let in_bounds_indices = indices.range(0..data.len());
// I copied this from a Stack Overflow answer
// without reading the text that explains why this is safe
in_bounds_indices.map(move |&i| unsafe { &mut *start.add(i) })
}
use std::iter::FromIterator;
fn main() {
let mut scores = vec![1, 2, 3];
let selected_scores: Vec<_> = {
// The set can go out of scope after we have used it.
let idx = BTreeSet::from_iter(vec![0, 2]);
uniq_refs(&mut scores, &idx).collect()
};
for score in selected_scores {
*score += 1;
}
println!("{:?}", scores);
}
Once you have used this function to find all the separate mutable references, you can use Rayon to modify them in parallel:
use rayon::prelude::*; // 1.0.3
fn example(scores: &mut [i32], indices: &BTreeSet<usize>) {
let selected_scores: Vec<_> = uniq_refs(scores, indices).collect();
selected_scores.into_par_iter().for_each(|s| *s *= 2);
// Or
uniq_refs(scores, indices).par_bridge().for_each(|s| *s *= 2);
}
You may wish to consider using a bitset instead of a BTreeMap
to be more efficient, but this answer uses only the standard library.
See also:
- How do I use Rayon with an existing iterator?