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

//! The reference neural sink applies the declared cost of an integer payload before it
//! decodes a symbol of it (`extensions/integer-payloads.md` §7, *Memory and work*;
//! `cost_over_limit`, §8), and its C ABI does not decode an integer output it has no
//! kind for. Measured with a global allocator that tracks the bytes live and their
//! peak, so "before decoding" is what the allocator sees: nothing sized from the
//! header is allocated.
//!
//! The payload is a clip a few kilobytes can code: sixteen frames whose stream is one
//! symbol of frequency 65535/65536 over and over. At 4729 × 4729 it decodes to a clip at
//! the 2^30-byte cap (release builds); a debug build uses 1024 × 1024 frames, 48 MiB.
//!
//! The tests share the allocator's counters, so they run one at a time.
#![cfg(all(feature = "neural", feature = "ffi", feature = "video_int"))]

use std::alloc::{GlobalAlloc, Layout, System};
use std::collections::BTreeMap;
use std::ffi::{CStr, CString};
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};
use std::sync::Mutex;

use wai::codecs::{CAP_NEURAL_INT_HYPER, CAP_VIDEO_INT_MOTION};
use wai::container::{Conditioning, Manifest, ModelRequirement, Rendition, Wai, WaiMulti};
use wai::ffi::{wai_int_payload_declared, wai_last_error, WaiBuffer, WaiIntDeclared, WAI_OK};
use wai::ffi_neural::{wai_neural_decode_ex, wai_neural_decode_video, WaiDecoded};
use wai::int_payload::{wire, DeclaredCost, IntDecoded};
use wai::neural::{decode_envelope, decode_multi, DecodeError, Decoded, ModelRegistry};

/// 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;

static SERIAL: Mutex<()> = Mutex::new(());

/// The most bytes `f` holds at once beyond what was live when it started.
fn peak_during<R>(f: impl FnOnce() -> R) -> (R, usize) {
    let base = CUR.load(Relaxed);
    PEAK.store(base, Relaxed);
    let r = f();
    (r, PEAK.load(Relaxed) - base)
}

/// Refusing before decoding allocates no more than this: the parsed table, a few
/// motion vectors, the refusal's message.
const SMALL: usize = 64 << 10;

fn vectors() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors")
}

fn read(rel: &str) -> Vec<u8> {
    std::fs::read(vectors().join(rel)).unwrap()
}

fn last_error() -> String {
    let p = wai_last_error();
    if p.is_null() { String::new() } else { unsafe { CStr::from_ptr(p) }.to_string_lossy().into_owned() }
}

/// A rANS stream of `n` symbols, each coded with `(start, freq)`, from the state 2^31,
/// so it ends canonically: what an encoder writes.
fn encode(n: u64, (start, freq): (u64, u64)) -> Vec<u8> {
    const L: u64 = 1 << 31;
    let mut state = L;
    let mut renorm: Vec<u32> = Vec::new();
    for _ in 0..n {
        let x_max = ((L >> 16) << 32) * freq;
        while state >= x_max {
            renorm.push(state as u32);
            state >>= 32;
        }
        state = ((state / freq) << 16) + (state % freq) + start;
    }
    renorm.reverse();
    let mut words = vec![state as u32, (state >> 32) as u32];
    words.extend(renorm);
    words.iter().flat_map(|w| w.to_le_bytes()).collect()
}

/// A keyframe-0 clip of sixteen frames whose every residual is 0 at frequency
/// 65535/65536: a valid stream of a few bytes per 2^20 symbols.
fn bomb() -> Vec<u8> {
    let side = if cfg!(debug_assertions) { 1024 } else { 4729 };
    let n = 16u32;
    let rans = encode(n as u64 * side as u64 * side as u64 * 3, (0, 65_535));
    wire::write_wiv1(&wire::Wiv1 {
        mv_precision: 1, keyframe: 0, n_frames: n, h: side, w: side, block: side, q: 1, offset: 0, cdf: vec![0, 65_535, 65_536],
        mvs: vec![(0, 0); n as usize - 1], rans: &rans, wih1: None,
    })
}

fn video(cap: &str, payload: Vec<u8>) -> Wai {
    Wai::new(Manifest::new("video", "improve", ModelRequirement::new(cap), Conditioning::new("bytes")), payload)
}

/// A registry with a cost limit refuses a payload over it before decoding anything
/// (`cost_over_limit`), and dispatch moves on: a `WAI2` decodes its next entry. The
/// cost it applied is `neural::declared_cost` over the parameters it resolves.
#[test]
fn a_payload_over_the_cost_limit_is_refused_before_any_symbol_is_decoded() {
    let _serial = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
    let bomb = bomb();
    assert!(bomb.len() < 4096, "the clip is coded in {} bytes", bomb.len());
    let limit = DeclaredCost::new(16 << 20, u64::MAX, u64::MAX, u128::MAX);
    let reg = ModelRegistry::new().builtin(CAP_VIDEO_INT_MOTION).with_cost_limit(limit);

    let env = video(CAP_VIDEO_INT_MOTION, bomb.clone());
    let (r, peak) = peak_during(|| decode_envelope(&env, &reg));
    match r {
        Err(e @ DecodeError::OverCostLimit(_)) => assert!(e.to_string().starts_with("cost_over_limit: ") && e.to_string().contains("output bytes"), "{e}"),
        Err(e) => panic!("expected cost_over_limit, got {e}"),
        Ok(_) => panic!("expected cost_over_limit, the clip decoded"),
    }
    assert!(peak < SMALL, "refusing it held {peak} bytes");

    let small = read("int_video/int_video_bitstream.wiv1");
    let n_frames = wire::parse_wiv1(&small, Some(0)).unwrap().n_frames as usize;
    let m = WaiMulti::new("video", "improve", serde_json::Value::Null,
                          vec![Rendition::media(CAP_VIDEO_INT_MOTION, "bytes", bomb.clone()), Rendition::media(CAP_VIDEO_INT_MOTION, "bytes", small)]);
    let (r, peak) = peak_during(|| decode_multi(&m, &reg).map(|(i, d)| (i, d.output)));
    match r {
        Ok((1, Decoded::Int(IntDecoded::Video { frames, .. }))) => assert_eq!(frames.len(), n_frames),
        Ok((i, _)) => panic!("decoded entry {i}"),
        Err(e) => panic!("the entry within the limit must decode: {e}"),
    }
    assert!(peak < 1 << 20, "skipping the first entry and decoding the second held {peak} bytes");

    let params = reg.resolve_verified(CAP_VIDEO_INT_MOTION, None).unwrap();
    let cost = wai::neural::declared_cost(CAP_VIDEO_INT_MOTION, &bomb, &params).unwrap();
    assert_eq!(cost.exceeds(&limit), Some("output bytes"));
    assert_eq!(cost.output_bytes, 16 * 3 * if cfg!(debug_assertions) { 1024 * 1024 } else { 4729 * 4729 });
}

/// The C ABI has no output kind for integer video, so it does not decode it: the
/// clip is refused before anything is allocated for it, by the registry entry point
/// and by the one-model video entry point.
#[test]
fn the_c_abi_refuses_an_integer_video_before_decoding_it() {
    let _serial = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
    let env = video(CAP_VIDEO_INT_MOTION, bomb()).to_emit_bytes().unwrap();
    let reg = CString::new(r#"{"wai.video.int_motion": ["/nonexistent"]}"#).unwrap();
    let mut out = std::mem::MaybeUninit::<WaiDecoded>::uninit();
    let (rc, peak) = peak_during(|| unsafe { wai_neural_decode_ex(env.as_ptr(), env.len(), reg.as_ptr(), 0, out.as_mut_ptr()) });
    let why = last_error();
    assert_ne!(rc, WAI_OK);
    assert!(why.contains("no wai_decoded_t kind"), "{why}");
    assert!(peak < SMALL, "decode_ex held {peak} bytes refusing it");
    unsafe { wai::ffi::wai_buffer_free(out.assume_init().data) };

    let model = CString::new("/nonexistent").unwrap();
    let mut frames = WaiBuffer { data: std::ptr::null_mut(), len: 0, cap: 0 };
    let null = std::ptr::null_mut();
    let (rc, peak) = peak_during(|| unsafe { wai_neural_decode_video(env.as_ptr(), env.len(), model.as_ptr(), &mut frames, null, null, null, null) });
    let why = last_error();
    assert_ne!(rc, WAI_OK);
    assert!(why.contains("no wai_decoded_t kind"), "{why}");
    assert!(peak < SMALL, "decode_video held {peak} bytes refusing it");
}

/// `neural::declared_cost` over the set a registry resolves is the library's over the
/// same bytes, and so is the C ABI's given the set's directory: every figure, the
/// multiply-accumulates in two halves. A set without a member is `set_incomplete`.
#[test]
fn the_declared_cost_over_a_registrys_set_and_through_the_c_abi() {
    let _serial = SERIAL.lock().unwrap_or_else(|e| e.into_inner());
    let dir = std::env::temp_dir().join(format!("wai-cost-set-{}", std::process::id()));
    std::fs::create_dir_all(&dir).unwrap();
    let set: BTreeMap<String, Vec<u8>> = [("entropy.json", "int_hyper/tiny_entropy.json"), ("g_s.whm", "int_hyper/tiny_g_s.bin"), ("h_s.whm", "int_hyper/tiny_h_s.bin")]
        .into_iter().map(|(k, f)| (k.to_string(), read(f))).collect();
    for (k, b) in &set {
        std::fs::write(dir.join(k), b).unwrap();
    }
    let payload = read("int_hyper/tiny_bitstream.wih1");
    let view: BTreeMap<String, &[u8]> = set.iter().map(|(k, b)| (k.clone(), b.as_slice())).collect();
    let held = wai::int_payload::resolve(CAP_NEURAL_INT_HYPER, None, &[wai::pin::HeldBytes::Set { listing: b"", files: &view }], None).unwrap();
    let want = wai::int_payload::declared_cost(CAP_NEURAL_INT_HYPER, &payload, &held).unwrap();
    assert!(want.macs > 0);

    let reg = ModelRegistry::new().register(CAP_NEURAL_INT_HYPER, dir.clone());
    let params = reg.resolve_verified(CAP_NEURAL_INT_HYPER, None).unwrap();
    assert_eq!(wai::neural::declared_cost(CAP_NEURAL_INT_HYPER, &payload, &params).unwrap(), want);

    let cap = CString::new(CAP_NEURAL_INT_HYPER).unwrap();
    let path = CString::new(dir.to_str().unwrap()).unwrap();
    let declared = |set: *const std::ffi::c_char| {
        let mut out = WaiIntDeclared::default();
        let rc = unsafe { wai_int_payload_declared(cap.as_ptr(), payload.as_ptr(), payload.len(), set, &mut out) };
        (rc, out)
    };
    let (rc, out) = declared(path.as_ptr());
    assert_eq!(rc, WAI_OK, "{}", last_error());
    assert_eq!((out.output_known, out.output_bytes, out.symbols, out.cost_known, out.working_bytes),
               (1, want.output_bytes, want.symbols, 1, want.working_bytes));
    assert_eq!(((out.macs_hi as u128) << 64) | out.macs_lo as u128, want.macs);
    let (rc, out) = declared(std::ptr::null());
    assert_eq!((rc, out.output_known, out.cost_known), (WAI_OK, 0, 0), "the header alone");

    std::fs::remove_file(dir.join("h_s.whm")).unwrap();
    let (rc, _) = declared(path.as_ptr());
    assert_ne!(rc, WAI_OK);
    assert!(last_error().contains("set_incomplete"), "{}", last_error());
    std::fs::remove_dir_all(&dir).ok();
}