Performing a reverse split_off operation in Rust efficiently

user_2181238 0
Viewed 357

My use case involves two mutable vectors a and b and a usize parameter x. I want to make the following change:

  1. take the elements b[0..x] and append them to the end of a (changing capacity of a as required)
  2. transform b into b[x..], without changing the original capacity of b

Currently I do the following:

while a.len() < x && !b.is_empty() {
    a.push(b.pop_front().unwrap());
    // here `b` is a VecDeque but I am happy to use a Vec if pop_front was not required
}

Obviously this seems a very slow operation, checking two conditions and calling unwrap at every iteration. It would be great if there was a rev_split_off operation such that:

let mut front = b.rev_split_off(x);
a.append(&mut front);

Here rev_split_off returns a newly allocated vector for the slice b[0..x] and transforms b into the remaining slice with unchanged capacity.

Question: How to perform my use case efficiently, with or without using such a thing as rev_split_off?

rust
3 Answers

Well, I think you will have to implement the rev_split_off yourself (even though I would probably call it split_off_back but it's the same).

Here is how I would implement it:

/// Moves the `i` first elements of `vec` at the end of `buffer`.
fn split_off_back<T>(vec: &mut Vec<T>, i: usize, buffer: &mut Vec<T>) {
    // We have to make sure vec has enough elements.
    // You could make the function unsafe and ask the caller to ensure
    // this condition.
    assert!(vec.len() >= i);

    // Reserve enough memory in the target buffer
    buffer.reserve(i);
    // Now we know `buffer.capacity() >= buffer.len() + i`.

    unsafe {
        // SAFETY:
        //  * `vec` and `buffer` are two distinct vectors (they come from mutable references)
        //     so their allocations cannot overlap.
        //  * `vec` is valid for reads because we have an exclusive reference to it and we
        //     checked the value of `i`.
        //  * `buffer` is valid for writes because we ensured we had enough memory to store
        //     `i` additional elements.
        std::ptr::copy_nonoverlapping(vec.as_ptr(), buffer.as_mut_ptr().add(buffer.len()), i);

        // Now the memory is moved.
        // we are not allowed to use it again from the `vec` vector.

        // We just extanded `buffer`, we need to update its length.
        // SAFEFY:
        //  * We ensured that the new length is less than the capacity (with `Vec::reserved`)
        //  * The vector is initialized for this new length (we moved the values).
        buffer.set_len(buffer.len() + i);

        // Now we need to update the first vector. The values from index `i` to its end
        // need to be moved at the begining of the vector.
        // SAFETY:
        //  * We have an exclusive reference to the vector. It is both valid for reads and writes.
        std::ptr::copy(vec.as_ptr().add(i), vec.as_mut_ptr(), i);

        // And update the length of `vec`.
        // SAFETY: This subtraction is safe because we previously checked that `vec.len() >= i`.
        vec.set_len(vec.len() - i);
    }
}

Note that I put buffer in the parameters of the function to avoid allocating a vector. If you want the same semantic as split_off, you can just do the following.

fn split_of_back<T>(vec: &mut Vec<T>, i: usize) -> Vec<T> {
    assert!(vec.len() >= i);
    
    let mut buffer = Vec::with_capacity(i);

    unsafe { /* same thing */ }

    buffer
}
user_133622760

This is very simple using drain and extend.

a.extend(b.drain(..x));

If your values are Copy, then you can get optimal speed using as_slice. IIUC, using extend_from_slice should be optional due to specialization, but aids clarity.

a.extend_from_slice(b.drain(..x).as_slice());
user_17633560

I have added some benchmarks to show that the unsafe version of this code is significantly faster.

#![feature(test)]

extern crate test;

#[cfg(test)]
mod tests {
    extern crate test;
    use test::{black_box, Bencher};

    /// Moves the `i` first elements of `vec` at the end of `buffer`.
    fn split_off_back<T>(vec: &mut Vec<T>, i: usize, buffer: &mut Vec<T>) {
        assert!(vec.len() >= i);

        buffer.reserve(i);

        unsafe {
            std::ptr::copy_nonoverlapping(vec.as_ptr(), buffer.as_mut_ptr().add(buffer.len()), i);
            buffer.set_len(buffer.len() + i);
            std::ptr::copy(vec.as_ptr().add(i), vec.as_mut_ptr(), i);
            vec.set_len(vec.len() - i);
        }
    }

    fn split_off_back_two<T>(vec: &mut Vec<T>, i: usize, buffer: &mut Vec<T>) {
        buffer.extend(vec.drain(..i));
    }

    const VEC_SIZE: usize = 100000;
    const SPLIT_POINT: usize = 20000;

    #[bench]
    fn run_v1(b: &mut Bencher) {
        b.iter(|| {
            let mut a = black_box(vec![0; VEC_SIZE]);
            let mut b = black_box(vec![0; VEC_SIZE]);
            split_off_back(&mut a, SPLIT_POINT, &mut b);
        });
    }

    #[bench]
    fn run_v2(b: &mut Bencher) {
        b.iter(|| {
            let mut a = black_box(vec![0; VEC_SIZE]);
            let mut b = black_box(vec![0; VEC_SIZE]);
            split_off_back_two(&mut a, SPLIT_POINT, &mut b);
        });
    }
}

This is the output of cargo bench on my machine:

running 2 tests
test tests::run_v1 ... bench:      98,863 ns/iter (+/- 2,058)
test tests::run_v2 ... bench:     230,665 ns/iter (+/- 6,093)

test result: ok. 0 passed; 0 failed; 0 ignored; 2 measured; 0 filtered out; finished in 0.48s
user_21812380
Related