Benchmark of taking Rust Vec items by predicate

bench.rs

// Results:
// |  algorithm  | Mixed | Odds first | Evens first |
// |-------------|-------|------------|-------------|
// |sort-split   | 465us | 35us       | 10us        |
// |drain_filter |  26us | 24us       | 22.5us      |
// |retain-uninit|  17us | 21us       | 19us        |
// 
// See also explanation on StackOverflow: https://stackoverflow.com/a/73005333/8195987

#![feature(drain_filter)]
#![deny(unsafe_op_in_unsafe_fn)]

use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
use rand::{thread_rng, Rng};

const NUM_ITEMS: usize = 16384;

mod data_gen {
    use super::*;
    fn get_even(rng: &mut impl Rng) -> u32 {
        let r: u32 = rng.gen();
        if r & 1 == 0 {
            r
        } else {
            r.wrapping_add(1)
        }
    }

    fn get_odd(rng: &mut impl Rng) -> u32 {
        let r: u32 = rng.gen();
        if r & 1 == 0 {
            r.wrapping_add(1)
        } else {
            r
        }
    }

    pub(super) fn generate_mixed() -> Vec<u32> {
        let mut v = vec![0; NUM_ITEMS];
        let mut rng = thread_rng();
        for (i, v) in v.iter_mut().enumerate() {
            let need_even = i % 2 == 0;
            if need_even {
                *v = get_even(&mut rng);
            } else {
                *v = get_odd(&mut rng);
            }
        }
        v
    }

    pub(super) fn generate_odds_first() -> Vec<u32> {
        let mut v = vec![0; NUM_ITEMS];
        let mut rng = thread_rng();
        for v in v[..NUM_ITEMS / 2].iter_mut() {
            *v = get_odd(&mut rng);
        }
        for v in v[NUM_ITEMS / 2..].iter_mut() {
            *v = get_even(&mut rng);
        }
        v
    }

    pub(super) fn generate_evens_first() -> Vec<u32> {
        let mut v = vec![0; NUM_ITEMS];
        let mut rng = thread_rng();
        for v in v[..NUM_ITEMS / 2].iter_mut() {
            *v = get_even(&mut rng);
        }
        for v in v[NUM_ITEMS / 2..].iter_mut() {
            *v = get_odd(&mut rng);
        }
        v
    }
}

mod algorithms {
    use std::ops::Deref;

    #[inline]
    fn predicate(x: impl Deref<Target = u32>) -> bool {
        (*x) & 1 != 0
    }

    // Returns slice with evens and slice with odds.
    #[inline(never)]
    pub(super) fn sort(v: &mut [u32]) -> (&mut [u32], &mut [u32]) {
        v.sort_by_key(|x| predicate(x));
        let split_pos = v.partition_point(|x| predicate(x));
        v.split_at_mut(split_pos)
    }

    // Takes odds from vect and put into another.
    #[inline(never)]
    pub(super) fn drain_filter(v: &mut Vec<u32>) -> Vec<u32> {
        let mut res = Vec::with_capacity(v.len());
        res.extend(v.drain_filter(|x| predicate(x)));
        res
    }

    /// Returns removed values.
    fn retain_unsafe_generic<T: Sized>(
        v: &mut Vec<T>,
        mut which_to_keep: impl FnMut(&T) -> bool,
    ) -> Vec<T> {
        use std::mem::{transmute, MaybeUninit};

        /// # Safety
        /// Caller must ensure that if it makes living copies of inner items,
        /// those items is removed from original vec before original reference became usable again.
        unsafe fn as_uninits<T: Sized>(v: &mut Vec<T>) -> &mut Vec<MaybeUninit<T>> {
            let orig_ptr = v.as_ptr();
            let orig_cap = v.capacity();
            let orig_size = v.len();

            let v: &mut Vec<MaybeUninit<T>> = unsafe {
                // Safety: since `MaybeUninit` has same memory layout
                // as wrapped type, we assume that we can treat vec with T
                // as MaybeUninit<T>. This assumption checked by asserts below.
                //
                // Lifetimes of elements must be correctly enforced by caller.
                transmute(v)
            };
            // Check if layout of Vec with different element type remains same.
            assert_eq!(v.len(), orig_size);
            assert_eq!(v.capacity(), orig_cap);
            assert_eq!(v.as_ptr(), orig_ptr.cast());

            v
        }

        let mut res: Vec<T> = Vec::with_capacity(v.len());
        let v = unsafe {
            // Safety: we keep result reference only in `retain` call.
            // We would remove all moved elements using retain.
            as_uninits(v)
        };
        v.retain(
            // Safety: `Vec::retain` would remove all items which values we moved into `res`.
            // It wouldn call `drop::<T>` for removed values
            // because `MaybeUninit` never drops wrapped values.
            |x| unsafe {
                // Safety: it is safe because `Vec::retain` visits elements sequentally
                // so we don't moved value from `x` yet.
                // https://doc.rust-lang.org/std/vec/struct.Vec.html#method.retain
                let val = &*x.as_ptr();
                if which_to_keep(val) {
                    return true;
                }
                res.reserve(1);
                // Any panic before this place is safe because
                // 1. We didn't moved value from `x` yet;
                // 2. In case of panic in predicate, `Vec::retain` preserve current value.

                // Here we could probably use `Vec::push`
                // but compiler currently fails to remove capacity check in `Vec::push`
                // so this function became slower than `Vec::drain_filter`
                // https://godbolt.org/z/7fhnnMh46
                // And `Vec::push(x.assume_init_read())` is unsafe operation too anyway.

                let old_len = res.len();
                // Safety: We just allocated memory for this place.
                let dst = res.as_mut_ptr().add(old_len);
                // Safety: since we cannot panic until the end of closure
                // and `Vec::retain` wouldn't panic and would remove `x`,
                // making bitwise copy of `x` is safe.
                x.as_ptr().copy_to_nonoverlapping(dst, 1);
                // Safety: we just wrote additional value.
                res.set_len(old_len + 1);

                false
            },
        );

        res
    }

    #[inline(never)]
    pub(super) fn retain_unsafe(v: &mut Vec<u32>) -> Vec<u32> {
        retain_unsafe_generic(v, |x| !predicate(x))
    }
}

fn compare_vec_item_extractors(c: &mut Criterion) {
    let mut group = c.benchmark_group("Vec element extractors");

    let generators: &[(&str, fn() -> Vec<u32>)] = &[
        ("Mixed", data_gen::generate_mixed),
        ("Odds first", data_gen::generate_odds_first),
        ("Evens first", data_gen::generate_evens_first),
    ];
    for (gen_name, provider_fn) in generators.iter().copied() {
        group.bench_function(format!("Sorts - {}", gen_name), |b| {
            b.iter_batched(
                provider_fn,
                |mut orig_vec| {
                    algorithms::sort(&mut orig_vec);
                    orig_vec
                },
                BatchSize::LargeInput,
            );
        });

        group.bench_function(format!("drain_filter - {}", gen_name), |b| {
            b.iter_batched(
                provider_fn,
                |mut orig_vec| {
                    let odds = algorithms::drain_filter(&mut orig_vec);
                    (orig_vec, odds)
                },
                BatchSize::LargeInput,
            );
        });

        group.bench_function(format!("retain_unsafe - {}", gen_name), |b| {
            b.iter_batched(
                provider_fn,
                |mut orig_vec| {
                    let odds = algorithms::retain_unsafe(&mut orig_vec);
                    (orig_vec, odds)
                },
                BatchSize::LargeInput,
            );
        });
    }

    group.finish();
}

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