to_vec_benchmark.rs

use criterion::{criterion_group, criterion_main, Criterion};

pub fn to_vec1(s: &[u32]) -> Vec<u32> {
    struct DropGuard<'a> {
        vec: &'a mut Vec<u32>,
        num_init: usize,
    }
    impl<'a> Drop for DropGuard<'a> {
        #[inline]
        fn drop(&mut self) {
            // SAFETY:
            // items were marked initialized in the loop below
            unsafe {
                self.vec.set_len(self.num_init);
            }
        }
    }
    let mut vec = Vec::with_capacity(s.len());
    let mut guard = DropGuard {
        vec: &mut vec,
        num_init: 0,
    };
    let slots = guard.vec.spare_capacity_mut();
    // .take(slots.len()) is necessary for LLVM to remove bounds checks
    // and has better codegen than zip.
    for (i, b) in s.iter().enumerate().take(slots.len()) {
        guard.num_init = i;
        slots[i].write(b.clone());
    }
    std::mem::forget(guard);
    // SAFETY:
    // the vec was allocated and initialized above to at least this length.
    unsafe {
        vec.set_len(s.len());
    }
    vec
}

pub fn to_vec2(s: &[u32]) -> Vec<u32> {
    struct DropGuard<'a> {
        vec: &'a mut Vec<u32>,
        num_init: usize,
    }
    impl<'a> Drop for DropGuard<'a> {
        #[inline]
        fn drop(&mut self) {
            // SAFETY:
            // items were marked initialized in the loop below
            unsafe {
                self.vec.set_len(self.num_init);
            }
        }
    }
    let mut vec = Vec::with_capacity(s.len());
    let mut guard = DropGuard {
        vec: &mut vec,
        num_init: 0,
    };
    let slots = guard.vec.spare_capacity_mut();
    // .take(slots.len()) is necessary for LLVM to remove bounds checks
    // and has better codegen than zip.
    for (i, (b, slot)) in s.iter().zip(slots.iter_mut()).enumerate() {
        guard.num_init = i;
        slot.write(b.clone());
    }
    std::mem::forget(guard);
    // SAFETY:
    // the vec was allocated and initialized above to at least this length.
    unsafe {
        vec.set_len(s.len());
    }
    vec
}

fn bench_to_vecs(c: &mut Criterion) {
    let mut group = c.benchmark_group("to_vec");
    let slice: Vec<u32> = (0..10000).collect();
    group.bench_function("original", |b| b.iter(|| to_vec1(&slice)));
    group.bench_function("proposed", |b| b.iter(|| to_vec2(&slice)));
    group.finish();
}

criterion_group!(benches, bench_to_vecs);
criterion_main!(benches);