bin_search_slice_benchmark.rs

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

/// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The
/// `key_fn` extracts a key of type `K` from the data, and this
/// function finds the range of elements that match the key. `data`
/// must have been sorted as if by a call to `sort_by_key` for this to
/// work.
pub fn binary_search_slice<'d, E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E]
where
    K: Ord,
{
    let Ok(mid) = data.binary_search_by_key(key, &key_fn) else {
        return &[];
    };
    let size = data.len();

    // We get back *some* element with the given key -- so do
    // a galloping search backwards to find the *first* one.
    let mut start = mid;
    let mut previous = mid;
    let mut step = 1;
    loop {
        start = start.saturating_sub(step);
        if start == 0 || key_fn(&data[start]) != *key {
            break;
        }
        previous = start;
        step *= 2;
    }
    step = previous - start;
    while step > 1 {
        let half = step / 2;
        let mid = start + half;
        if key_fn(&data[mid]) != *key {
            start = mid;
        }
        step -= half;
    }
    // adjust by one if we have overshot
    if start < size && key_fn(&data[start]) != *key {
        start += 1;
    }

    // Now search forward to find the *last* one.
    let mut end = mid;
    let mut previous = mid;
    let mut step = 1;
    loop {
        end = end.saturating_add(step).min(size);
        if end == size || key_fn(&data[end]) != *key {
            break;
        }
        previous = end;
        step *= 2;
    }
    step = end - previous;
    while step > 1 {
        let half = step / 2;
        let mid = end - half;
        if key_fn(&data[mid]) != *key {
            end = mid;
        }
        step -= half;
    }

    &data[start..end]
}

/// Uses a sorted slice `data: &[E]` as a kind of "multi-map". The
/// `key_fn` extracts a key of type `K` from the data, and this
/// function finds the range of elements that match the key. `data`
/// must have been sorted as if by a call to `sort_by_key` for this to
/// work.
pub fn binary_search_slice_new<'d, E, K>(data: &'d [E], key_fn: impl Fn(&E) -> K, key: &K) -> &'d [E]
where
    K: Ord,
{
    let size = data.len();
    let start = data.partition_point(|x| key_fn(x) < *key);
    // At this point `start` either points at the first entry with equal or
    // greater key or is equal to `size` in case all elements have smaller keys
    // Invariant: start == size || key_fn(&data[start]) >= *key
    if start == size || key_fn(&data[start]) != *key {
        return &[];
    };
    // Invariant: start < size && key_fn(&data[start]) == *key

    // Find the first entry with key > `key`. Skip `start` entries since
    // key_fn(&data[start]) == *key
    // Invariant: offset == size ||  key_fn(&data[offset]) >= *key
    let offset = start + 1;
    let end = data[offset..].partition_point(|x| key_fn(x) <= *key) + offset;
    // Invariant: end == size || key_fn(&data[end]) > *key

    &data[start..end]
}

fn bench_inserts(c: &mut Criterion) {
    let mut group = c.benchmark_group("multiply add");
    let data = [(0, "zero"), (3, "three-a"), (3, "three-b"), (22, "twenty-two")];
    let keys = [-1, 0, 1, 2, 3, 22, 23];
    group.bench_function("binary_search_slice", |b| {
        b.iter(|| {
            let mut total_len = 0;
            for key in keys {
                total_len += binary_search_slice(&data, |x| x.0, &key).len();
            }
            total_len
        })
    });
    group.bench_function("binary_search_slice_new", |b| {
        b.iter(|| {
            let mut total_len = 0;
            for key in keys {
                total_len += binary_search_slice_new(&data, |x| x.0, &key).len();
            }
            total_len
        })
    });
    group.finish();
}

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