Alternative to swapping vector elements in rust - rust

Question

I'm experimenting with rust by porting some c++ code. I write a lot of code that uses vectors as object pools by moving elements to the back in various ways and then resizing. Here's a ported function:
use rand::{thread_rng, Rng};
fn main() {
for n in 1..11 {
let mut a: Vec<u8> = (1..11).collect();
keep_n_rand(&mut a, n);
println!("{}: {:?}", n, a);
}
}
fn keep_n_rand<T>(x: &mut Vec<T>, n: usize) {
let mut rng = thread_rng();
for i in n..x.len() {
let j = rng.gen_range(0..i);
if j < n {
x.swap(i, j);
}
}
x.truncate(n);
}
It keeps n elements chosen at random. It is done this way because it does not reduce the capacity of the vector so that more objects can be added later without allocating (on average). This might be iterated millions of times.
In c++, I would use x[j] = std::move(x[i]); because I am about to truncate the vector. While it has no impact in this example, if the swap was expensive, it would make sense to move. Is that possible and desirable in rust? I can live with a swap. I'm just curious.

Answer 1

Correct me if I'm wrong: you're looking for a way to retain n random elements in a Vec and discard the rest. In that case, the easiest way would be to use partial_shuffle(), a rand function implemented for slices.
Shuffle a slice in place, but exit early.
Returns two mutable slices from the source slice. The first contains amount elements randomly permuted. The second has the remaining elements that are not fully shuffled.
use rand::{thread_rng, seq::SliceRandom};
fn main() {
let mut rng = thread_rng();
// Use the `RangeInclusive` (`..=`) syntax at times like this.
for n in 1..=10 {
let mut elements: Vec<u8> = (1..=10).collect();
let (elements, _rest) = elements.as_mut_slice().partial_shuffle(&mut rng, n);
println!("{n}: {elements:?}");
}
}
Run this snippet on Rust Playground.
elements is shadowed, going from a Vec to a &mut [T]. If you're only going to use it inside the function, that's probably all you'll need. However, since it's a reference, you can't return it; the data it's pointing to is owned by the original vector, which will be dropped when it goes out of scope. If that's what you need, you'll have to turn the slice into a Vec.
While you can simply construct a new one from it using Vec::from, I suspect (but haven't tested) that it's more efficient to use Vec::split_off.
Splits the collection into two at the given index.
Returns a newly allocated vector containing the elements in the range [at, len). After the call, the original vector will be left containing the elements [0, at) with its previous capacity unchanged.
use rand::{thread_rng, seq::SliceRandom};
fn main() {
let mut rng = thread_rng();
for n in 1..=10 {
let mut elements: Vec<u8> = (1..=10).collect();
elements.as_mut_slice().partial_shuffle(&mut rng, n);
let elements = elements.split_off(elements.len() - n);
// `elements` is still a `Vec`; this time, containing only
// the shuffled elements. You can use it as the return value.
println!("{n}: {elements:?}");
}
}
Run this snippet on Rust Playground.
Since this function lives on a performance-critical path, I'd recommend benchmarking it against your current implementation. At the time of writing this, criterion is the most popular way to do that. That said, rand is an established library, so I imagine it will perform as well or better than a manual implementation.
Sample Benchmark
I don't know what kind of numbers you're working with, but here's a sample benchmark with for n in 1..=100 and (1..=100).collect() (i.e. 100 instead of 10 in both places) without the print statements:
manual time: [73.683 µs 73.749 µs 73.821 µs]
rand with slice time: [68.074 µs 68.147 µs 68.226 µs]
rand with vec time: [54.147 µs 54.213 µs 54.288 µs]
Bizarrely, splitting off a Vec performed vastly better than not. Unless I made an error in my benchmarks, the compiler is probably doing something under the hood that you'll need a more experienced Rustacean than me to explain.
Benchmark Implementation
Cargo.toml
[dependencies]
rand = "0.8.5"
[dev-dependencies]
criterion = "0.4.0"
[[bench]]
name = "rand_benchmark"
harness = false
[[bench]]
name = "rand_vec_benchmark"
harness = false
[[bench]]
name = "manual_benchmark"
harness = false
benches/manual_benchmark.rs
use criterion::{criterion_group, criterion_main, Criterion};
fn manual_solution() {
for n in 1..=100 {
let mut elements: Vec<u8> = (1..=100).collect();
keep_n_rand(&mut elements, n);
}
}
fn keep_n_rand<T>(elements: &mut Vec<T>, n: usize) {
use rand::{thread_rng, Rng};
let mut rng = thread_rng();
for i in n..elements.len() {
let j = rng.gen_range(0..i);
if j < n {
elements.swap(i, j);
}
}
elements.truncate(n);
}
fn benchmark(c: &mut Criterion) {
c.bench_function("manual", |b| b.iter(manual_solution));
}
criterion_group!(benches, benchmark);
criterion_main!(benches);
benches/rand_benchmark.rs
use criterion::{criterion_group, criterion_main, Criterion};
fn rand_solution() {
use rand::{seq::SliceRandom, thread_rng};
let mut rng = thread_rng();
for n in 1..=100 {
let mut elements: Vec<u8> = (1..=100).collect();
let (_elements, _) = elements.as_mut_slice().partial_shuffle(&mut rng, n);
}
}
fn benchmark(c: &mut Criterion) {
c.bench_function("rand with slice", |b| b.iter(rand_solution));
}
criterion_group!(benches, benchmark);
criterion_main!(benches);
benches/rand_vec_benchmark.rs
use criterion::{criterion_group, criterion_main, Criterion};
fn rand_solution() {
use rand::{seq::SliceRandom, thread_rng};
let mut rng = thread_rng();
for n in 1..=100 {
let mut elements: Vec<u8> = (1..=100).collect();
elements.as_mut_slice().partial_shuffle(&mut rng, n);
let _elements = elements.split_off(elements.len() - n);
}
}
fn benchmark(c: &mut Criterion) {
c.bench_function("rand with vec", |b| b.iter(rand_solution));
}
criterion_group!(benches, benchmark);
criterion_main!(benches);

Answer 2

Is that possible and desirable in rust?
It is not possible unless you constrain T: Copy or T: Clone: while C++ uses non-destructive moves (the source is in a valid but unspecified state) Rust uses destructive moves (the source is gone).
There are ways around it using unsafe but they require being very careful and it's probably not worth the hassle (you can look at Vec::swap_remove for a taste, it basically does what you're doing here except only between j and the last element of the vec).
I'd also recommend verified_tinker's solution, as I'm not convinced your shuffle is unbiased.