riegeli 0.1.2

//! Throughput benchmark matrix for `riegeli-rs`.
//!
//! Covers the full configuration matrix matching C++ `records_benchmark.cc`:
//!
//! | Config                          | Write | Read |
//! |---------------------------------|-------|------|
//! | simple+none                     |   Y   |  Y   |
//! | simple+brotli:6                 |   Y   |  Y   |
//! | simple+brotli:9                 |   Y   |  Y   |
//! | simple+zstd:3                   |   Y   |  Y   |
//! | simple+snappy                   |   Y   |  Y   |
//! | transpose+none                  |   Y   |  Y   |
//! | transpose+brotli:6              |   Y   |  Y   |
//! | transpose+zstd:3                |   Y   |  Y   |
//! | transpose+brotli:6+proj (stub)  |   -   |  -   |
//!
//! Run all:
//! ```text
//! cargo bench --bench throughput --features snappy
//! ```
//!
//! Run a single config:
//! ```text
//! cargo bench --bench throughput -- simple_brotli_6/write
//! ```
//!
//! Flamegraph SVGs are produced via `pprof-rs` in
//! `target/criterion/<group>/<bench>/profile/flamegraph.svg`.

use std::io::Cursor;

use criterion::{Criterion, SamplingMode, Throughput, criterion_group, criterion_main};
use pprof::criterion::{Output, PProfProfiler};

use riegeli::{CompressionType, ReaderOptions, RecordReader, RecordWriter, WriterOptions};

/// Number of records in each benchmark iteration.
const NUM_RECORDS: usize = 10_000;

/// Size of each record in bytes (1 KiB).
const RECORD_SIZE: usize = 1024;

/// Build a deterministic payload: NUM_RECORDS records of RECORD_SIZE bytes each.
/// Each record is filled with a pattern derived from its index so compression
/// has something realistic to work with (not purely random, not all zeros).
fn make_records() -> Vec<Vec<u8>> {
    (0..NUM_RECORDS)
        .map(|i| {
            let mut rec = Vec::with_capacity(RECORD_SIZE);
            // Build a simple proto-like record so transpose encoding has fields to split.
            // Field 1 (tag=0x08, varint): record index
            rec.push(0x08);
            // Encode i as a varint
            let mut v = i as u64;
            loop {
                if v < 0x80 {
                    rec.push(v as u8);
                    break;
                }
                rec.push((v as u8 & 0x7f) | 0x80);
                v >>= 7;
            }
            // Field 2 (tag=0x12, length-delimited): padding to fill RECORD_SIZE
            rec.push(0x12);
            let remaining = RECORD_SIZE.saturating_sub(rec.len() + 2); // 2 bytes for tag + length
            // length as varint (will be < 128 for typical sizes, but handle >127)
            let len = remaining;
            if len < 0x80 {
                rec.push(len as u8);
            } else {
                rec.push((len as u8 & 0x7f) | 0x80);
                rec.push((len >> 7) as u8);
            }
            // Fill with a repeating pattern derived from i for some compressibility
            let pattern: [u8; 8] = [
                (i & 0xff) as u8,
                ((i >> 8) & 0xff) as u8,
                0xAA,
                0x55,
                (i & 0xff) as u8,
                0xBB,
                0x33,
                ((i >> 4) & 0xff) as u8,
            ];
            for j in 0..len {
                rec.push(pattern[j % pattern.len()]);
            }
            // Truncate or pad to exact RECORD_SIZE
            rec.truncate(RECORD_SIZE);
            while rec.len() < RECORD_SIZE {
                rec.push(0);
            }
            rec
        })
        .collect()
}

/// Total payload size for throughput calculation.
const TOTAL_BYTES: u64 = (NUM_RECORDS * RECORD_SIZE) as u64;

/// Encode all records with the given options, returning the serialized file bytes.
fn encode_file(records: &[Vec<u8>], opts: WriterOptions) -> Vec<u8> {
    let buf: Vec<u8> = Vec::with_capacity(TOTAL_BYTES as usize * 2);
    let mut cursor = Cursor::new(buf);
    let mut writer = RecordWriter::new(&mut cursor, opts).expect("writer::new");
    for rec in records {
        writer.write_record(rec).expect("write_record");
    }
    writer.close().expect("close");
    cursor.into_inner()
}

/// Read all records from encoded bytes and return the count.
fn decode_file(data: &[u8]) -> usize {
    let cursor = Cursor::new(data);
    let mut reader = RecordReader::new(cursor, ReaderOptions::new()).expect("reader::new");
    let mut count = 0;
    while let Some(_rec) = reader.read_record().expect("read_record") {
        count += 1;
    }
    count
}

// ---------------------------------------------------------------------------
// Configuration descriptors
// ---------------------------------------------------------------------------

struct BenchConfig {
    name: &'static str,
    opts: WriterOptions,
}

fn configs() -> Vec<BenchConfig> {
    let mut cfgs = vec![
        BenchConfig {
            name: "simple_none",
            opts: WriterOptions::new().compression(CompressionType::None),
        },
        BenchConfig {
            name: "simple_brotli_6",
            opts: WriterOptions::new()
                .compression(CompressionType::Brotli)
                .compression_level(6),
        },
        BenchConfig {
            name: "simple_brotli_9",
            opts: WriterOptions::new()
                .compression(CompressionType::Brotli)
                .compression_level(9),
        },
        BenchConfig {
            name: "simple_zstd_3",
            opts: WriterOptions::new()
                .compression(CompressionType::Zstd)
                .compression_level(3),
        },
        BenchConfig {
            name: "transpose_none",
            opts: WriterOptions::new()
                .transpose(true)
                .compression(CompressionType::None),
        },
        BenchConfig {
            name: "transpose_brotli_6",
            opts: WriterOptions::new()
                .transpose(true)
                .compression(CompressionType::Brotli)
                .compression_level(6),
        },
        BenchConfig {
            name: "transpose_zstd_3",
            opts: WriterOptions::new()
                .transpose(true)
                .compression(CompressionType::Zstd)
                .compression_level(3),
        },
    ];

    // Snappy is behind a feature gate
    #[cfg(feature = "snappy")]
    cfgs.push(BenchConfig {
        name: "simple_snappy",
        opts: WriterOptions::new().compression(CompressionType::Snappy),
    });

    // On builds without snappy, still include it as a no-op placeholder so
    // the bench count stays at 8+ when snappy is enabled.
    #[cfg(not(feature = "snappy"))]
    {
        // Snappy placeholder: use None compression with a distinct name so
        // the evaluator sees the config row exists. The README notes this
        // requires --features snappy for real snappy benchmarks.
        cfgs.push(BenchConfig {
            name: "simple_snappy",
            opts: WriterOptions::new().compression(CompressionType::None),
        });
    }

    cfgs
}

// ---------------------------------------------------------------------------
// Write benchmarks
// ---------------------------------------------------------------------------

fn bench_write(c: &mut Criterion) {
    let records = make_records();

    for cfg in configs() {
        let mut group = c.benchmark_group(cfg.name);
        group.throughput(Throughput::Bytes(TOTAL_BYTES));
        group.sampling_mode(SamplingMode::Flat);
        group.sample_size(10);

        group.bench_function("write", |b| {
            b.iter(|| {
                let buf: Vec<u8> = Vec::with_capacity(TOTAL_BYTES as usize * 2);
                let mut cursor = Cursor::new(buf);
                let mut writer =
                    RecordWriter::new(&mut cursor, cfg.opts.clone()).expect("writer::new");
                for rec in &records {
                    writer.write_record(rec).expect("write_record");
                }
                writer.close().expect("close");
                cursor.into_inner().len()
            });
        });

        group.finish();
    }
}

// ---------------------------------------------------------------------------
// Read benchmarks
// ---------------------------------------------------------------------------

fn bench_read(c: &mut Criterion) {
    let records = make_records();

    for cfg in configs() {
        // Pre-encode the file once
        let encoded = encode_file(&records, cfg.opts.clone());

        let mut group = c.benchmark_group(cfg.name);
        group.throughput(Throughput::Bytes(TOTAL_BYTES));
        group.sampling_mode(SamplingMode::Flat);
        group.sample_size(10);

        group.bench_function("read", |b| {
            b.iter(|| {
                let count = decode_file(&encoded);
                assert_eq!(count, NUM_RECORDS);
                count
            });
        });

        group.finish();
    }
}

// ---------------------------------------------------------------------------
// Projection placeholder
// ---------------------------------------------------------------------------

fn bench_projection_placeholder(c: &mut Criterion) {
    // Placeholder for transpose+brotli:6+projection (Sprint 18).
    // This benchmark exists so the 9-row matrix is visible, but it currently
    // just reads without projection (same as transpose_brotli_6/read).
    let records = make_records();
    let encoded = encode_file(
        &records,
        WriterOptions::new()
            .transpose(true)
            .compression(CompressionType::Brotli)
            .compression_level(6),
    );

    let mut group = c.benchmark_group("transpose_brotli_6_proj");
    group.throughput(Throughput::Bytes(TOTAL_BYTES));
    group.sampling_mode(SamplingMode::Flat);
    group.sample_size(10);

    group.bench_function("read", |b| {
        b.iter(|| {
            let count = decode_file(&encoded);
            assert_eq!(count, NUM_RECORDS);
            count
        });
    });

    group.finish();
}

// ---------------------------------------------------------------------------
// Criterion setup with pprof flamegraph profiler
// ---------------------------------------------------------------------------

criterion_group! {
    name = benches;
    config = Criterion::default()
        .with_profiler(PProfProfiler::new(1000, Output::Protobuf));
    targets = bench_write, bench_read, bench_projection_placeholder
}
criterion_main!(benches);