sum-segment-tree 0.2.0

A fixed-capacity sum-segment tree for weighted sampling and prioritized experience replay.
Documentation
use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
use std::hint::black_box;
use sum_segment_tree::SumTree;

const SIZES: [usize; 3] = [1_024, 16_384, 131_072];

fn filled_tree(capacity: usize) -> SumTree {
    let mut tree = SumTree::new(capacity);
    for i in 0..capacity {
        tree.push((i % 97) as f32 + 1.0);
    }
    tree
}

fn bench_push(c: &mut Criterion) {
    let mut group = c.benchmark_group("push");
    for size in SIZES {
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            b.iter(|| {
                let mut tree = SumTree::new(size);
                for i in 0..size {
                    tree.push(black_box((i % 97) as f32 + 1.0));
                }
                tree
            });
        });
    }
    group.finish();
}

fn bench_update(c: &mut Criterion) {
    let mut group = c.benchmark_group("update");
    for size in SIZES {
        let mut tree = filled_tree(size);
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &size| {
            let mut i = 0;
            b.iter(|| {
                i = (i + 1) % size;
                tree.update(black_box(i), black_box(1.0 + (i % 13) as f32));
            });
        });
    }
    group.finish();
}

fn bench_get(c: &mut Criterion) {
    let mut group = c.benchmark_group("get");
    for size in SIZES {
        let tree = filled_tree(size);
        let total = tree.total();
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, &_size| {
            let mut k = 0.0f32;
            b.iter(|| {
                k = (k + 0.618_034) % 1.0;
                tree.get(black_box(k * total))
            });
        });
    }
    group.finish();
}

criterion_group!(benches, bench_push, bench_update, bench_get);
criterion_main!(benches);