nordisk/znippy-zoomies

lgz/tests/alloc_count.rs

//! Zero allocation in the hot path: decoding a bigger input must not allocate
//! more. A counting global allocator measures the whole decode of two inputs,
//! one 4× the other; the count may not grow with the size (no allocation per
//! read window after warm-up). The compressed input is a slice (an mmap in
//! the CLI) and is never copied: the decoder reads it by offsets.

use std::alloc::{GlobalAlloc, Layout, System};
use std::io::Write;
use std::sync::atomic::{AtomicUsize, Ordering};

struct Counting;
static ALLOCS: AtomicUsize = AtomicUsize::new(0);
static BYTES: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        ALLOCS.fetch_add(1, Ordering::Relaxed);
        BYTES.fetch_add(l.size(), Ordering::Relaxed);
        unsafe { System.alloc(l) }
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) }
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, n: usize) -> *mut u8 {
        ALLOCS.fetch_add(1, Ordering::Relaxed);
        BYTES.fetch_add(n, Ordering::Relaxed);
        unsafe { System.realloc(p, l, n) }
    }
}

#[global_allocator]
static A: Counting = Counting;

fn corpus(len: usize, seed: u32) -> Vec<u8> {
    let words: [&[u8]; 8] = [
        b"usr",
        b"lib",
        b"python3.12",
        b"site-packages",
        b"share",
        b"etc",
        b"bin",
        b"README",
    ];
    let mut out = Vec::with_capacity(len + 64);
    let mut x: u32 = seed | 1;
    while out.len() < len {
        x ^= x << 13;
        x ^= x >> 17;
        x ^= x << 5;
        out.extend_from_slice(words[(x % 8) as usize]);
        out.push(if x >> 29 == 0 { b'\n' } else { b'/' });
        if x & 0x3f00 == 0 {
            out.extend_from_slice(&x.to_le_bytes());
        }
    }
    out.truncate(len);
    out
}

/// Counts bytes, keeps nothing (so the writer does not allocate either).
struct Sink(u64);
impl Write for Sink {
    fn write(&mut self, b: &[u8]) -> std::io::Result<usize> {
        self.0 += b.len() as u64;
        Ok(b.len())
    }
    fn flush(&mut self) -> std::io::Result<()> {
        Ok(())
    }
}

fn allocs_for(gz: &[u8]) -> (usize, usize, u64) {
    let mut w = Sink(0);
    let (a0, b0) = (
        ALLOCS.load(Ordering::Relaxed),
        BYTES.load(Ordering::Relaxed),
    );
    lgz::stream::decode_slice(gz, &mut w, 1).expect("decode");
    (
        ALLOCS.load(Ordering::Relaxed) - a0,
        BYTES.load(Ordering::Relaxed) - b0,
        w.0,
    )
}

#[test]
fn one_thread_allocates_nothing_per_window() {
    let small = corpus(80 << 20, 3);
    let big = corpus(320 << 20, 4);
    let gz = |d: &[u8]| {
        let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::new(6));
        e.write_all(d).unwrap();
        e.finish().unwrap()
    };
    let (gs, gb) = (gz(&small), gz(&big));
    let _ = allocs_for(&gs); // warm-up (tables, lazy statics)
    let (na, nb_bytes, outa) = allocs_for(&gs);
    let (nb, _, outb) = allocs_for(&gb);
    assert_eq!(outa, small.len() as u64);
    assert_eq!(outb, big.len() as u64);
    assert!(
        nb <= na + 2,
        "allocations grow with the input: {na} for 80 MB, {nb} for 320 MB ({nb_bytes} bytes for the small one)"
    );
    assert!(na < 64, "{na} allocations for one decode");
}

fn alloc_bytes_parallel(gz: &[u8], threads: usize) -> (usize, usize, u64) {
    let mut w = Sink(0);
    let (a0, b0) = (
        ALLOCS.load(Ordering::Relaxed),
        BYTES.load(Ordering::Relaxed),
    );
    lgz::stream::decode_slice(gz, &mut w, threads).expect("decode");
    (
        ALLOCS.load(Ordering::Relaxed) - a0,
        BYTES.load(Ordering::Relaxed) - b0,
        w.0,
    )
}

#[test]
#[ignore = "open: per-window bookkeeping (chunk record, split starts, scoped split threads, job/channel nodes) still allocates; 4x input = 1528 vs 460 allocs, 237 vs 136 MB (loki, 2026-11-02)"]
fn eight_threads_allocate_no_more_bytes_for_a_bigger_input() {
    // the parallel path: segment buffers, marker buffers, resolved heads and
    // Huffman tables come from pools, so a 4× input allocates (almost) no more
    // bytes; what is left per window is small bookkeeping, counted below
    let gz = |d: &[u8]| {
        let mut e = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::new(6));
        e.write_all(d).unwrap();
        e.finish().unwrap()
    };
    // the window plan sizes segments from the input (S / (N × k)): pin them at
    // min_seg for both inputs, so the pools see the same buffer sizes and only
    // the number of windows differs
    // SAFETY: set before any decode thread starts; read once per decode
    unsafe { std::env::set_var("LGZ_K", "64") };
    let small = corpus(160 << 20, 5);
    let big = corpus(640 << 20, 6);
    let (gs, gb) = (gz(&small), gz(&big));
    let _ = alloc_bytes_parallel(&gs, 8); // warm-up: the pools fill
    let (ns, bs, os) = alloc_bytes_parallel(&gs, 8);
    let (nb, bb, ob) = alloc_bytes_parallel(&gb, 8);
    assert_eq!(os, small.len() as u64);
    assert_eq!(ob, big.len() as u64);
    eprintln!(
        "8 threads: {ns} allocs / {bs} bytes for 160 MB, {nb} allocs / {bb} bytes for 640 MB"
    );
    assert!(
        bb < bs + (32 << 20) && nb < ns * 2,
        "allocated bytes grow with the input: {bs} for 160 MB of output, {bb} for 640 MB"
    );
}