open-standards/wai/wai-rs/tests/int_video_enc_memory.rs

// wai-conformance: int-kernels
//! What the integer video encoder holds (`int_video_enc::working_bytes`) against what it
//! holds in fact, measured with a global allocator that tracks the bytes live and their
//! peak, as `tests/int_payload_cost.rs` measures the decoder.
//!
//! The encoder codes in two passes and keeps no residual or symbol per sample, so beyond
//! the frames it is given it holds the reconstruction it returns, the motion vectors and
//! the stream (as emitted, finished and copied into the payload): a small multiple of
//! the clip at worst, about one clip for content that codes well. The clips: noise at
//! step 1 (a wide residual table, a long stream), a smooth pan, block 1 (a vector for
//! every pixel), one frame; tiny clips whose residual table is wide, where the table's
//! fit is most of what is held (1 × 3 noise, 2 frames; 2 × 5, 4 frames, block 5); and a
//! seeded sweep of small random clips over every block, precision, search and step.
//!
//! Its own test binary with a single test, so that no other test's allocations land in
//! the window it measures.
#![cfg(feature = "video_int")]

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

use wai::int_video_enc::{encode, working_bytes, EncParams};

/// The system allocator, tracking the bytes live and the most live at once.
struct Live;
static CUR: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Live {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        PEAK.fetch_max(CUR.fetch_add(l.size(), Relaxed) + l.size(), Relaxed);
        unsafe { System.alloc(l) }
    }
    unsafe fn alloc_zeroed(&self, l: Layout) -> *mut u8 {
        PEAK.fetch_max(CUR.fetch_add(l.size(), Relaxed) + l.size(), Relaxed);
        unsafe { System.alloc_zeroed(l) }
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, new: usize) -> *mut u8 {
        if new >= l.size() {
            PEAK.fetch_max(CUR.fetch_add(new - l.size(), Relaxed) + (new - l.size()), Relaxed);
        } else {
            CUR.fetch_sub(l.size() - new, Relaxed);
        }
        unsafe { System.realloc(p, l, new) }
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        CUR.fetch_sub(l.size(), Relaxed);
        unsafe { System.dealloc(p, l) }
    }
}

#[global_allocator]
static ALLOC: Live = Live;

fn clip(h: usize, w: usize, n: usize, f: impl Fn(usize, usize, usize, usize) -> u8) -> Vec<Vec<u8>> {
    (0..n).map(|t| (0..h * w * 3).map(|i| f(t, i / 3 / w, i / 3 % w, i % 3)).collect()).collect()
}

/// A clip: its name, height, width, frames and the encoder's parameters.
type Case = (&'static str, u32, u32, Vec<Vec<u8>>, EncParams);

#[test]
fn the_encoder_holds_at_most_its_working_bytes() {
    let mut x: u64 = 0x9e37_79b9_7f4a_7c15;
    let mut noise = move || {
        x ^= x << 13;
        x ^= x >> 7;
        x ^= x << 17;
        (x >> 24) as u8
    };
    let noisy: Vec<Vec<u8>> = (0..10).map(|_| (0..40 * 48 * 3).map(|_| noise()).collect()).collect();
    let cases: Vec<Case> = vec![
        ("noise, q 1", 40, 48, noisy, EncParams { block: 4, search: 2, mv_precision: 2, q: 1 }),
        ("smooth pan", 36, 64, clip(36, 64, 16, |t, y, x, c| ((x * 3 + y + 2 * t + c * 40) % 200) as u8), EncParams { block: 8, search: 3, mv_precision: 1, q: 6 }),
        ("block 1", 16, 16, clip(16, 16, 6, |t, y, x, c| ((x * 7 + y * 5 + t * 3 + c) % 256) as u8), EncParams { block: 1, search: 1, mv_precision: 1, q: 2 }),
        ("one frame", 50, 70, clip(50, 70, 1, |_, y, x, c| ((x ^ y) as u8).wrapping_mul(c as u8 + 1)), EncParams { block: 16, search: 4, mv_precision: 2, q: 3 }),
        ("1 x 3 noise, 2 frames, q 1", 1, 3, vec![vec![0, 255, 7, 250, 3, 128, 255, 0, 90], vec![255, 0, 250, 5, 255, 0, 0, 255, 1]], EncParams { block: 1, search: 1, mv_precision: 1, q: 1 }),
        ("2 x 5 noise, 4 frames, block 5, q 1", 2, 5, clip(2, 5, 4, |t, y, x, c| ((x * 97 + y * 31 + t * 151 + c * 67) * 2654435761 % 256) as u8), EncParams { block: 5, search: 1, mv_precision: 2, q: 1 }),
    ];
    for (name, h, w, frames, p) in cases {
        let base = CUR.load(Relaxed);
        PEAK.store(base, Relaxed);
        let e = encode(&frames, h, w, &p).expect("a clip the encoder codes");
        let peak = (PEAK.load(Relaxed) - base) as u64;
        let bound = working_bytes(frames.len() as u64, h, w, p.block);
        let clip_bytes = frames.len() as u64 * h as u64 * w as u64 * 3;
        assert!(peak <= bound, "{name}: held {peak} bytes, past the stated {bound}");
        // The reconstruction it returns is held, so at least one clip.
        assert!(peak >= clip_bytes, "{name}: held {peak} bytes, under the clip's {clip_bytes}");
        drop(e);
    }
    // The review's sweep, seeded: small clips of noise, ramps and extremes, every block
    // size to 8, both precisions, search to 2, steps from 1 to 65535.
    let mut s: u64 = 0x1234_5678_9abc_def1;
    let mut next = move || {
        s ^= s << 13;
        s ^= s >> 7;
        s ^= s << 17;
        s
    };
    let count = 3000;
    for i in 0..count {
        let h = 1 + (next() % 12) as u32;
        let w = 1 + (next() % 12) as u32;
        let n = 1 + (next() % 30) as usize;
        let block = 1 + (next() % 8) as u32;
        let mv_precision = 1 + (next() % 2) as u8;
        let search = (next() % 3) as u32;
        let q = [1u32, 2, 5, 64, 255, 65535][(next() % 6) as usize];
        let kind = next() % 3;
        let len = h as usize * w as usize * 3;
        let frames: Vec<Vec<u8>> = (0..n)
            .map(|t| {
                (0..len)
                    .map(|j| match kind {
                        0 => (next() >> 24) as u8,
                        1 => ((j * 7 + t * 13) % 251) as u8,
                        _ => if (j + t) % 2 == 0 { 0 } else { 255 },
                    })
                    .collect()
            })
            .collect();
        let p = EncParams { block, search, mv_precision, q };
        let base = CUR.load(Relaxed);
        PEAK.store(base, Relaxed);
        let e = encode(&frames, h, w, &p).expect("a clip the encoder codes");
        let peak = (PEAK.load(Relaxed) - base) as u64;
        let bound = working_bytes(n as u64, h, w, block);
        assert!(peak <= bound, "sweep case {i} ({h} x {w}, {n} frames, {p:?}): held {peak} bytes, past the stated {bound}");
        drop(e);
    }
}