clawsync/crates/clawhdf5-onion/benches/tdt_bench.rs

//! Criterion benchmarks for the TDT byte-interleaving transform + ZstdTdt codec.
//!
//! Measures compression ratio and throughput for common HDF5 data types,
//! comparing plain Zstd against TDT-transformed Zstd.
//!
//! Run:
//!   cargo bench -p clawhdf5-onion -- tdt_compress

use clawhdf5_onion::compress::{compress_page, decompress_page};
use clawhdf5_onion::format::Codec;
use clawhdf5_onion::tdt;
use criterion::{BenchmarkId, Criterion, Throughput, black_box, criterion_group, criterion_main};

// ─────────────────────────────────────────────────────────────────────────────
// Test data generators
// ─────────────────────────────────────────────────────────────────────────────

/// Smooth f32: values from a sine wave — good compressibility, representative
/// of neural network activation or weight tensors.
fn make_f32_smooth(n_bytes: usize) -> Vec<u8> {
    let n_floats = n_bytes / 4;
    let mut out = Vec::with_capacity(n_bytes);
    for i in 0..n_floats {
        let v = ((i as f64 / 256.0).sin() as f32).to_le_bytes();
        out.extend_from_slice(&v);
    }
    // Pad to exact size if n_bytes % 4 != 0.
    out.resize(n_bytes, 0);
    out
}

/// Random f32: worst case for compression.
fn make_f32_random(n_bytes: usize) -> Vec<u8> {
    // Deterministic LCG so benchmarks are reproducible.
    let mut state = 0x_dead_beef_u64;
    let mut out = Vec::with_capacity(n_bytes);
    while out.len() < n_bytes {
        state = state
            .wrapping_mul(6364136223846793005)
            .wrapping_add(1442695040888963407);
        out.extend_from_slice(&(state as u32).to_le_bytes());
    }
    out.truncate(n_bytes);
    out
}

/// Random i32 data.
fn make_int32_random(n_bytes: usize) -> Vec<u8> {
    make_f32_random(n_bytes) // same byte distribution
}

/// Sequential i8 data (highly compressible, TDT should be neutral).
fn make_int8_seq(n_bytes: usize) -> Vec<u8> {
    (0u8..=255).cycle().take(n_bytes).collect()
}

// ─────────────────────────────────────────────────────────────────────────────
// Helpers
// ─────────────────────────────────────────────────────────────────────────────

/// Run compress + decompress roundtrip and return compressed size.
fn bench_compress(data: &[u8], codec: Codec) -> usize {
    let compressed = compress_page(black_box(data), codec).unwrap();
    let size = compressed.len();
    let _ = decompress_page(black_box(&compressed), codec, data.len() as u32).unwrap();
    size
}

// ─────────────────────────────────────────────────────────────────────────────
// Benchmark groups
// ─────────────────────────────────────────────────────────────────────────────

fn bench_tdt_compress(c: &mut Criterion) {
    let sizes = [4096usize, 65536];

    type DataGen = fn(usize) -> Vec<u8>;
    let datasets: &[(&str, DataGen, usize)] = &[
        ("f32_smooth", make_f32_smooth, 4),
        ("f32_random", make_f32_random, 4),
        ("int32_random", make_int32_random, 4),
        ("int8_seq", make_int8_seq, 1),
    ];

    let codecs = [("zstd", Codec::Zstd), ("zstd_tdt", Codec::ZstdTdt)];

    // ── Throughput benchmark ──────────────────────────────────────────────────
    let mut group = c.benchmark_group("tdt_compress/throughput");
    for &size in &sizes {
        for &(dtype, make, _width) in datasets {
            let data = make(size);
            group.throughput(Throughput::Bytes(size as u64));
            for &(codec_name, codec) in &codecs {
                let id = BenchmarkId::new(format!("{codec_name}/{dtype}"), format!("{size}B"));
                group.bench_with_input(id, &data, |b, d| {
                    b.iter(|| bench_compress(d, codec));
                });
            }
        }
    }
    group.finish();

    // ── TDT transform only (encode + decode, without compression) ────────────
    let mut group = c.benchmark_group("tdt_compress/transform_only");
    for &size in &sizes {
        for &(dtype, make, width) in datasets {
            let data = make(size);
            group.throughput(Throughput::Bytes(size as u64));

            let id = BenchmarkId::new(format!("encode/{dtype}"), format!("{size}B"));
            group.bench_with_input(id, &data, |b, d| {
                b.iter(|| tdt::encode(black_box(d), width));
            });

            let encoded = tdt::encode(&data, width);
            let id = BenchmarkId::new(format!("decode/{dtype}"), format!("{size}B"));
            group.bench_with_input(id, &encoded, |b, enc| {
                b.iter(|| tdt::decode(black_box(enc), width, size));
            });
        }
    }
    group.finish();

    // ── Compression ratio summary (printed, not timed) ───────────────────────
    // Run once outside Criterion to print ratio comparison.
    println!("\n─── TDT compression ratio summary ───");
    println!(
        "{:<20} {:>8} {:>10} {:>10} {:>8}",
        "dataset/size", "orig", "zstd", "zstd_tdt", "savings"
    );
    for &size in &sizes {
        for &(dtype, make, _w) in datasets {
            let data = make(size);
            let zstd_size = compress_page(&data, Codec::Zstd).unwrap().len();
            let tdt_size = compress_page(&data, Codec::ZstdTdt).unwrap().len();
            let savings_pct = 100.0 * (1.0 - tdt_size as f64 / zstd_size as f64);
            println!(
                "{:<20} {:>8} {:>10} {:>10} {:>7.1}%",
                format!("{dtype}/{size}B"),
                size,
                zstd_size,
                tdt_size,
                savings_pct
            );
        }
    }
    println!();
}

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