stream-vbyte 0.4.1

#[cfg(feature = "x86_ssse3")]
use std::arch::x86_64::__m128i;

use criterion::{criterion_group, criterion_main, Bencher, BenchmarkId, Criterion, Throughput};

use rand::{
    distributions::{Distribution as _, Uniform},
    Rng,
};

use std::iter;

#[cfg(feature = "x86_ssse3")]
use stream_vbyte::x86::{self};
use stream_vbyte::{
    decode::{
        cursor::DecodeCursor, decode, DecodeQuadSink, DecodeSingleSink, Decoder, WriteQuadToSlice,
    },
    decode_quad_scalar,
    encode::{encode, Encoder},
    scalar::Scalar,
};

fn bench_encode(c: &mut Criterion) {
    let sizes = &[1_000, 1_000_000];

    bench_group_sizes(c, "encode_scalar_zeros", sizes, |b, &size| {
        do_encode_bench(b, iter::repeat(0).take(size as usize), Scalar)
    });

    bench_group_sizes(c, "encode_scalar_rand", sizes, |b, &size| {
        do_encode_bench(
            b,
            RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
            Scalar,
        )
    });
}

fn bench_decode(c: &mut Criterion) {
    let sizes = &[1_000, 1_000_000];

    bench_group_sizes(c, "decode_scalar_zeros", sizes, |b, &size| {
        do_decode_bench(b, iter::repeat(0).take(size as usize), Scalar)
    });

    bench_group_sizes(c, "decode_scalar_rand", sizes, |b, &size| {
        do_decode_bench(
            b,
            RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
            Scalar,
        )
    });

    bench_group_sizes(c, "decode_cursor_slice_scalar_rand", sizes, |b, &size| {
        do_decode_cursor_slice_bench(
            b,
            RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
            Scalar,
        )
    });

    bench_group_sizes(
        c,
        "decode_cursor_sink_no_op_scalar_rand",
        sizes,
        |b, &size| {
            do_decode_cursor_sink_no_op_bench(
                b,
                RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
                Scalar,
            )
        },
    );

    bench_group_sizes(c, "decode_cursor_skip_all", sizes, |b, &size| {
        let mut nums: Vec<u32> = Vec::new();
        let mut encoded = Vec::new();
        let mut decoded = Vec::new();

        for i in RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize) {
            nums.push(i);
        }

        encoded.resize(nums.len() * 5, 0);
        let bytes_written = encode::<Scalar>(&nums, &mut encoded);

        decoded.resize(nums.len(), 0);
        b.iter(|| {
            DecodeCursor::new(&encoded[0..bytes_written], size as usize).skip(size as usize);
        });
    });
}

#[cfg(not(feature = "x86_sse41"))]
fn bench_encode_x86(_c: &mut Criterion) {}

#[cfg(not(feature = "x86_ssse3"))]
fn bench_decode_x86(_c: &mut Criterion) {}

#[cfg(feature = "x86_sse41")]
fn bench_encode_x86(c: &mut Criterion) {
    let sizes = &[1_000, 1_000_000];

    bench_group_sizes(c, "encode_sse41_zeros", sizes, |b, &size| {
        do_encode_bench(b, iter::repeat(0).take(size as usize), x86::Sse41);
    });

    bench_group_sizes(c, "encode_sse41_rand", sizes, |b, &size| {
        do_encode_bench(
            b,
            RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
            x86::Sse41,
        );
    });
}

#[cfg(feature = "x86_ssse3")]
fn bench_decode_x86(c: &mut Criterion) {
    let sizes = &[1_000, 1_000_000];
    bench_group_sizes(c, "decode_ssse3_zeros", sizes, |b, &size| {
        do_decode_bench(b, iter::repeat(0).take(size as usize), x86::Ssse3);
    });

    bench_group_sizes(c, "decode_ssse3_rand", sizes, |b, &size| {
        do_decode_bench(
            b,
            RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
            x86::Ssse3,
        );
    });

    bench_group_sizes(c, "decode_cursor_slice_ssse3_rand", sizes, |b, &size| {
        do_decode_cursor_slice_bench(
            b,
            RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
            x86::Ssse3,
        );
    });

    bench_group_sizes(
        c,
        "decode_cursor_sink_no_op_ssse3_rand",
        sizes,
        |b, &size| {
            do_decode_cursor_sink_no_op_bench(
                b,
                RandomVarintEncodedLengthIter::new(rand::thread_rng()).take(size as usize),
                x86::Ssse3,
            );
        },
    );
}

criterion_group!(
    benches,
    bench_encode,
    bench_decode,
    bench_encode_x86,
    bench_decode_x86
);
criterion_main!(benches);

fn bench_group_sizes<F>(c: &mut Criterion, group_name: &str, sizes: &[u64], mut f: F)
where
    F: FnMut(&mut Bencher<'_, criterion::measurement::WallTime>, &u64),
{
    let mut group = c.benchmark_group(group_name);
    for &size in sizes {
        group.throughput(Throughput::Elements(size));
        group.bench_with_input(BenchmarkId::from_parameter(size), &size, |b, size| {
            f(b, size)
        });
    }
    group.finish();
}

fn do_encode_bench<I: Iterator<Item = u32>, E: Encoder>(b: &mut Bencher, iter: I, _encoder: E) {
    let mut nums: Vec<u32> = Vec::new();
    let mut encoded = Vec::new();

    for i in iter {
        nums.push(i);
    }

    encoded.resize(nums.len() * 5, 0);

    b.iter(|| {
        let _ = encode::<E>(&nums, &mut encoded);
    });
}

// take a decoder param to save us some typing -- type inference won't work if you only specify some
// of the generic types
fn do_decode_bench<I, D>(b: &mut Bencher, iter: I, _decoder: D)
where
    I: Iterator<Item = u32>,
    D: Decoder + WriteQuadToSlice,
{
    let mut nums: Vec<u32> = Vec::new();
    let mut encoded = Vec::new();
    let mut decoded = Vec::new();

    for i in iter {
        nums.push(i);
    }

    encoded.resize(nums.len() * 5, 0);
    let bytes_written = encode::<Scalar>(&nums, &mut encoded);

    decoded.resize(nums.len(), 0);
    b.iter(|| {
        decode::<D>(&encoded[0..bytes_written], nums.len(), &mut decoded);
    });
}

fn do_decode_cursor_slice_bench<I, D>(b: &mut Bencher, iter: I, _decoder: D)
where
    I: Iterator<Item = u32>,
    D: Decoder + WriteQuadToSlice,
{
    let mut nums: Vec<u32> = Vec::new();
    let mut encoded = Vec::new();
    let mut decoded = Vec::new();

    for i in iter {
        nums.push(i);
    }

    encoded.resize(nums.len() * 5, 0);
    let _ = encode::<Scalar>(&nums, &mut encoded);

    decoded.resize(nums.len(), 0);
    b.iter(|| {
        let mut cursor = DecodeCursor::new(&encoded, nums.len());
        cursor.decode_slice::<D>(&mut decoded);
    })
}

fn do_decode_cursor_sink_no_op_bench<I: Iterator<Item = u32>, D: Decoder>(
    b: &mut Bencher,
    iter: I,
    _decoder: D,
) where
    NoOpSink: DecodeQuadSink<D>,
{
    let mut nums: Vec<u32> = Vec::new();
    let mut encoded = Vec::new();

    for i in iter {
        nums.push(i);
    }

    encoded.resize(nums.len() * 5, 0);
    let _ = encode::<Scalar>(&nums, &mut encoded);

    b.iter(|| {
        let mut cursor = DecodeCursor::new(&encoded, nums.len());
        let mut sink = NoOpSink;
        cursor.decode_sink::<D, _>(&mut sink, nums.len());
    })
}

// copied from tests because it's handy here too
struct RandomVarintEncodedLengthIter<R: Rng> {
    ranges: [Uniform<u32>; 4],
    range_for_picking_range: Uniform<usize>,
    rng: R,
}

impl<R: Rng> RandomVarintEncodedLengthIter<R> {
    fn new(rng: R) -> RandomVarintEncodedLengthIter<R> {
        RandomVarintEncodedLengthIter {
            ranges: [
                Uniform::new(0, 1 << 8),
                Uniform::new(1 << 8, 1 << 16),
                Uniform::new(1 << 16, 1 << 24),
                Uniform::new(1 << 24, u32::MAX), // this won't ever emit the max value, sadly
            ],
            range_for_picking_range: Uniform::new(0, 4),
            rng,
        }
    }
}

impl<R: Rng> Iterator for RandomVarintEncodedLengthIter<R> {
    type Item = u32;

    fn next(&mut self) -> Option<Self::Item> {
        // pick the range we'll use
        let value_range = self.ranges[self.range_for_picking_range.sample(&mut self.rng)];

        Some(value_range.sample(&mut self.rng))
    }
}

struct NoOpSink;

impl DecodeSingleSink for NoOpSink {
    fn on_number(&mut self, _num: u32, _nums_decoded: usize) {}
}

decode_quad_scalar!(NoOpSink);

#[cfg(feature = "x86_ssse3")]
impl DecodeQuadSink<x86::Ssse3> for NoOpSink {
    fn on_quad(&mut self, _quad: __m128i, _nums_decoded: usize) {}
}