bench.rs

use criterion::*;
use rand::distributions::Distribution;
use std::collections::BinaryHeap;
use std::cmp::Reverse;

fn top_n_vec_sort(dist_num: usize, n: usize) -> Vec<usize> {
    let dist = rand::distributions::Uniform::from(100000000..1000000000)
        .sample_iter(rand::thread_rng())
        .take(dist_num);
    let mut v = dist.collect::<Vec<_>>();
    v.sort_unstable();
    v.into_iter().take(n).collect::<Vec<_>>()
}

fn top_n_heap_push_pop(dist_num: usize, n: usize) -> Vec<usize> {
    let dist = rand::distributions::Uniform::from(100000000..1000000000)
        .sample_iter(rand::thread_rng())
        .take(dist_num);
    dist.fold(BinaryHeap::<Reverse<_>>::with_capacity(n), |mut heap, e| {
        if heap.len() == n && e <= heap.peek().unwrap().0 {
            return heap;
        }
        heap.push(Reverse(e));
        if heap.len() > n {
            heap.pop();
        }
        heap
    })
        .into_sorted_vec()
        .into_iter()
        .map(|e| e.0)
        .collect()
}

fn top_n_heap_peek_1(dist_num: usize, n: usize) -> Vec<usize> {
    let dist = rand::distributions::Uniform::from(100000000..1000000000)
        .sample_iter(rand::thread_rng())
        .take(dist_num);
    dist.fold(BinaryHeap::<Reverse<_>>::with_capacity(n), |mut heap, e| {
        if heap.len() == n {
            let mut top = heap.peek_mut().unwrap();
            if e > top.0 {
                *top = Reverse(e);
            }
        } else {
            heap.push(Reverse(e));
        }
        heap
    })
        .into_sorted_vec()
        .into_iter()
        .map(|e| e.0)
        .collect()
}

//fn top_n_heap_peek_2(dist_num: usize, n: usize) -> Vec<usize> {
//    let dist = rand::distributions::Uniform::from(100000000..1000000000)
//        .sample_iter(rand::thread_rng())
//        .take(dist_num);
//    dist.fold(BinaryHeap::<Reverse<_>>::with_capacity(n), |mut heap, e| {
//        if heap.len() == n && e > heap.peek().unwrap().0 {
//            *(heap.peek_mut().unwrap()) = Reverse(e);
//        } else {
//            heap.push(Reverse(e));
//        }
//        heap
//    })
//        .into_sorted_vec()
//        .into_iter()
//        .map(|e| e.0)
//        .collect()
//}

fn top_n_iter_peekmut_unguarded(dist_num: usize, n: usize) -> Vec<usize> {
    let mut dist = rand::distributions::Uniform::from(100000000..1000000000)
        .sample_iter(rand::thread_rng())
        .take(dist_num);
    let mut heap: BinaryHeap<Reverse<_>> = (&mut dist).take(n).map(Reverse).collect();
    for el in dist {
        let mut top = heap.peek_mut().unwrap();
        if el > top.0 {
            *top = Reverse(el);
        }
    }
    heap
        .into_sorted_vec()
        .into_iter()
        .map(|e| e.0)
        .collect()
}


fn top_n_iter_peekmut_guarded(dist_num: usize, n: usize) -> Vec<usize> {
    let mut dist = rand::distributions::Uniform::from(100000000..1000000000)
        .sample_iter(rand::thread_rng())
        .take(dist_num);
    let mut heap: BinaryHeap<Reverse<_>> = (&mut dist).take(n).map(Reverse).collect();
    for el in dist {
        if el > heap.peek().unwrap().0 {
            *heap.peek_mut().unwrap() = Reverse(el);
        }
    }
    heap
        .into_sorted_vec()
        .into_iter()
        .map(|e| e.0)
        .collect()
}

fn top_n_iter_pushpop(dist_num: usize, n: usize) -> Vec<usize> {
    let mut dist = rand::distributions::Uniform::from(100000000..1000000000)
        .sample_iter(rand::thread_rng())
        .take(dist_num);
    let mut heap: BinaryHeap<Reverse<_>> = (&mut dist).take(n).map(Reverse).collect();
    for el in dist {
        heap.push(Reverse(el));
        heap.pop();
    }
    heap
        .into_sorted_vec()
        .into_iter()
        .map(|e| e.0)
        .collect()
}


fn bench(c: &mut Criterion) {
    c.bench(
        "top_n",
        ParameterizedBenchmark::new(
            "vec_sort",
            |b, i| b.iter(|| top_n_vec_sort(*i, 100)),
            vec![
                100usize, 1000usize,
                10000usize,
                100000usize
            ],
        )
            .with_function("heap_push_pop", |b, i| b.iter(|| top_n_heap_push_pop(*i, 100)))
            .with_function("heap_peek_1", |b, i| b.iter(|| top_n_heap_peek_1(*i, 100)))
//            .with_function("heap_peek_2", |b, i| b.iter(|| top_n_heap_peek_2(*i, 100)))
            .with_function("heap_top_n_iter_pushpop", |b, i| b.iter(|| top_n_iter_pushpop(*i, 100)))
            .with_function("heap_top_n_iter_peekmut_guarded", |b, i| b.iter(|| top_n_iter_peekmut_guarded(*i, 100)))
            .with_function("heap_top_n_iter_peekmut_unguarded", |b, i| b.iter(|| top_n_iter_peekmut_unguarded(*i, 100)))
    );
}

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