open-standards/wai/wai-vector/tests/peak_alloc.rs

//! Parse memory is O(input) plus fixed: nothing is reserved by a header's counts until
//! the exact length shows the bytes back them (§2.0, §2.7 steps 2–3).
//!
//! One test in its own binary, so the counting allocator sees only it.

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

use wai_vector::vectors::{all_vectors, corpus_store};
use wai_vector::wvg1::{self, Header};

struct Counting;

static NOW: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Counting {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        let p = unsafe { System.alloc(l) };
        if !p.is_null() {
            let n = NOW.fetch_add(l.size(), Ordering::SeqCst) + l.size();
            PEAK.fetch_max(n, Ordering::SeqCst);
        }
        p
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        unsafe { System.dealloc(p, l) };
        NOW.fetch_sub(l.size(), Ordering::SeqCst);
    }
}

#[global_allocator]
static A: Counting = Counting;

/// The peak bytes allocated while `f` runs, above what was live before it.
fn peak(f: impl FnOnce()) -> usize {
    let base = NOW.load(Ordering::SeqCst);
    PEAK.store(base, Ordering::SeqCst);
    f();
    PEAK.load(Ordering::SeqCst) - base
}

const C: usize = 6;
const FIXED: usize = 64 << 10;

#[test]
fn peak_allocation_is_linear_in_the_input() {
    let mut store = corpus_store();
    let held: usize = wai_vector::vectors::sets::corpus_sets().iter().map(|(_, b)| b.len()).max().unwrap();
    let mut worst = 0f64;
    for v in all_vectors().into_iter().filter(|v| v.ext == "wvg1") {
        let input = v.bytes.len() + held;
        let p = peak(|| {
            let _ = wvg1::load(&v.bytes, &mut store);
        });
        assert!(p <= C * input + FIXED, "{}: peak {p} bytes for {input} bytes of input", v.name);
        worst = worst.max(p as f64 / input as f64);
    }
    // Headers that claim the caps over a payload of 80 bytes: refused at the exact
    // length, having reserved nothing the counts name.
    let mut h = Header::read(&all_vectors().into_iter().find(|v| v.name == "minimal").unwrap().bytes);
    h.n_keys = 1 << 22;
    h.n_glyphs = 1 << 20;
    h.n_items = 65_536;
    let mut b = Vec::new();
    h.write(&mut b);
    let p = peak(|| {
        assert!(wvg1::parse(&b).is_err());
    });
    assert!(p <= FIXED, "a header claiming the caps made the parser reserve {p} bytes");
    eprintln!("worst peak / input over the corpus: {worst:.2}");
}