dsi-bitstream 0.10.1

/*
 * SPDX-FileCopyrightText: 2025 Tommaso Fontana
 * SPDX-FileCopyrightText: 2025 Sebastiano Vigna
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 */

//! Comparative Criterion benchmarks for dsi-bitstream codes.
//!
//! Compares a selection of codes side by side using both implied and universal
//! Zipf distributions.
//!
//! Use Criterion's built-in regex filter for ad-hoc selection:
//!
//! ```bash
//! # Only gamma benchmarks
//! cargo bench --bench comparative -- 'gamma'
//! # Only big-endian reads
//! cargo bench --bench comparative -- '/BE/.*read'
//! ```

mod common;

use common::N;
use common::data::*;
use criterion::{Criterion, Throughput, criterion_group, criterion_main};
use dsi_bitstream::prelude::*;
use std::hint::black_box;

#[cfg(feature = "bench-u16")]
type ReadWord = u16;
#[cfg(all(feature = "bench-u64", not(feature = "bench-u16")))]
type ReadWord = u64;
#[cfg(not(any(feature = "bench-u16", feature = "bench-u64")))]
type ReadWord = u32;

#[cfg(feature = "bench-reads")]
type WriteWord = u64;

#[cfg(all(not(feature = "bench-reads"), feature = "bench-u16"))]
type WriteWord = u16;
#[cfg(all(
    not(feature = "bench-reads"),
    feature = "bench-u64",
    not(feature = "bench-u16")
))]
type WriteWord = u64;
#[cfg(all(
    not(feature = "bench-reads"),
    not(any(feature = "bench-u16", feature = "bench-u64"))
))]
type WriteWord = u32;

/// Macro to register a comparative benchmark for a code (both endiannesses,
/// both distributions, read + write). Write methods receive `(value, args...)`
/// and read methods receive `(args...)`, where `args` comes from the
/// parenthesized list after the method name (empty for non-parametric codes).
/// Set `implied_only` to `true` for codes whose codeword length is proportional
/// to the value (e.g., unary), making the universal Zipf distribution
/// impractical.
macro_rules! bench_comp {
    ($group:expr, $name:expr, $write_method:ident($($wargs:expr),*), $read_method:ident($($rargs:expr),*), $len_fn:expr) => {
        bench_comp!($group, $name, $write_method($($wargs),*), $read_method($($rargs),*), $len_fn, false)
    };
    ($group:expr, $name:expr, $write_method:ident($($wargs:expr),*), $read_method:ident($($rargs:expr),*), $len_fn:expr, $implied_only:expr) => {{
        let name_str = $name;

        let dists: &[(&str, bool)] = if $implied_only {
            &[("implied", false)]
        } else {
            &[("implied", false), ("univ", true)]
        };

        for &(dist_name, univ) in dists {
            let data = gen_data($len_fn, univ);

            // Write benchmarks
            {
                let bench_id = format!("{}/BE/{}/write", name_str, dist_name);
                let data_ref = &data;
                $group.bench_function(&bench_id, |b| {
                    let mut buffer: Box<[WriteWord]> = vec![0; 10 * N].into_boxed_slice();
                    b.iter(|| {
                        let mut w = BufBitWriter::<BE, _>::new(
                            MemWordWriterSlice::<WriteWord, _>::new(&mut *buffer),
                        );
                        for &value in data_ref {
                            black_box(w.$write_method(value, $($wargs),*).unwrap());
                        }
                    });
                });
            }
            {
                let bench_id = format!("{}/LE/{}/write", name_str, dist_name);
                let data_ref = &data;
                $group.bench_function(&bench_id, |b| {
                    let mut buffer: Box<[WriteWord]> = vec![0; 10 * N].into_boxed_slice();
                    b.iter(|| {
                        let mut w = BufBitWriter::<LE, _>::new(
                            MemWordWriterSlice::<WriteWord, _>::new(&mut *buffer),
                        );
                        for &value in data_ref {
                            black_box(w.$write_method(value, $($wargs),*).unwrap());
                        }
                    });
                });
            }

            // Read benchmarks — encode into Box<[u64]> to guarantee alignment
            // for reinterpretation as &[ReadWord] via align_to.
            {
                let encoded = {
                    let mut buffer: Box<[u64]> = vec![0u64; 10 * N].into_boxed_slice();
                    {
                        let mut w = BufBitWriter::<BE, _>::new(
                            MemWordWriterSlice::<u64, _>::new(&mut *buffer),
                        );
                        for &value in &data {
                            w.$write_method(value, $($wargs),*).unwrap();
                        }
                    }
                    buffer
                };
                let n = data.len();
                let bench_id = format!("{}/BE/{}/read", name_str, dist_name);
                $group.bench_function(&bench_id, |b| {
                    b.iter(|| {
                        // SAFETY: Box<[u64]> is aligned to 8 bytes, which
                        // satisfies alignment for ReadWord (u16/u32/u64).
                        let slice: &[ReadWord] = unsafe { encoded.align_to::<ReadWord>().1 };
                        let mut r =
                            BufBitReader::<BE, _>::new(MemWordReader::new(slice));
                        for _ in 0..n {
                            black_box(r.$read_method($($rargs),*).unwrap());
                        }
                    });
                });
            }
            {
                let encoded = {
                    let mut buffer: Box<[u64]> = vec![0u64; 10 * N].into_boxed_slice();
                    {
                        let mut w = BufBitWriter::<LE, _>::new(
                            MemWordWriterSlice::<u64, _>::new(&mut *buffer),
                        );
                        for &value in &data {
                            w.$write_method(value, $($wargs),*).unwrap();
                        }
                    }
                    buffer
                };
                let n = data.len();
                let bench_id = format!("{}/LE/{}/read", name_str, dist_name);
                $group.bench_function(&bench_id, |b| {
                    b.iter(|| {
                        // SAFETY: Box<[u64]> is aligned to 8 bytes, which
                        // satisfies alignment for ReadWord (u16/u32/u64).
                        let slice: &[ReadWord] = unsafe { encoded.align_to::<ReadWord>().1 };
                        let mut r =
                            BufBitReader::<LE, _>::new(MemWordReader::new(slice));
                        for _ in 0..n {
                            black_box(r.$read_method($($rargs),*).unwrap());
                        }
                    });
                });
            }
        }
    }};
}

/// Comparative benchmarks: all codes compared side by side.
fn bench_comparative(c: &mut Criterion) {
    #[cfg(target_os = "linux")]
    common::utils::pin_to_core(2);

    let mut group = c.benchmark_group("comparative");
    group.throughput(Throughput::Elements(N as u64));

    // Fixed-parameter codes
    bench_comp!(
        group,
        "unary",
        write_unary(),
        read_unary(),
        |x: u64| x as usize + 1,
        true
    );
    bench_comp!(group, "gamma", write_gamma(), read_gamma(), len_gamma);
    bench_comp!(group, "delta", write_delta(), read_delta(), len_delta);
    bench_comp!(group, "omega", write_omega(), read_omega(), len_omega);
    bench_comp!(
        group,
        "vbytebe",
        write_vbyte_be(),
        read_vbyte_be(),
        bit_len_vbyte
    );
    bench_comp!(
        group,
        "vbytele",
        write_vbyte_le(),
        read_vbyte_le(),
        bit_len_vbyte
    );

    // Specialized (table-using) variants
    bench_comp!(group, "zeta3", write_zeta3(), read_zeta3(), |x| len_zeta(
        x, 3
    ));
    bench_comp!(group, "pi2", write_pi2(), read_pi2(), |x| len_pi(x, 2));

    // Parametric codes with k = 2..4 or 2..5
    for k in 2..4usize {
        bench_comp!(
            group,
            format!("zeta_{}", k),
            write_zeta(k),
            read_zeta(k),
            |x| len_zeta(x, k)
        );
    }
    for k in 2..5usize {
        bench_comp!(group, format!("pi_{}", k), write_pi(k), read_pi(k), |x| {
            len_pi(x, k)
        });
        bench_comp!(
            group,
            format!("rice_{}", k),
            write_rice(k),
            read_rice(k),
            |x| len_rice(x, k),
            true
        );
        bench_comp!(
            group,
            format!("exgol_{}", k),
            write_exp_golomb(k),
            read_exp_golomb(k),
            |x| len_exp_golomb(x, k)
        );
        bench_comp!(
            group,
            format!("gol_{}", k),
            write_golomb(k as u64),
            read_golomb(k as u64),
            |x| len_golomb(x, k as u64),
            true
        );
    }

    group.finish();
}

criterion_group! {
    name = comparative;
    config = Criterion::default()
        .sample_size(10)
        .warm_up_time(std::time::Duration::from_millis(500))
        .measurement_time(std::time::Duration::from_secs(1));
    targets = bench_comparative
}

criterion_main!(comparative);