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

// wai-conformance: int-kernels int-kernels-amd
//! `WIR1` (`wai.{image,video,audio,splat}.int_refine`, staged-delivery Appendix A)
//! through the integer operator set: the stage operators of
//! `extensions/integer-operators.md` §13.5, E1 `rans64_decode`, E11 `up2_predict`, E12
//! `rate_predict` and E9 `residual_step`, with no graph.
//!
//! A layer runs no network, so it has no lowering and no bounds certificate. What it has
//! is its **stage form** (`staged::wir1::ops::stages`): the operands its decode hands
//! each stage operator. The operators' own conformance vectors are
//! `ops-conformance/stage/`, which `wai-verify vectors verify` and the independent
//! checker (`tools/wai_ops_check.py`) reproduce byte for byte.
#![cfg(all(feature = "staged", feature = "legacy_int"))]

use std::collections::BTreeSet;
use std::path::{Path, PathBuf};

use serde_json::Value;
use wai::staged::wir1::{self, ops, Below, Predictor, RefineParams, UnitKind};
use wai_ops::stage::rans64::{self, End, Escape, TableRule};

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

fn read(p: &Path) -> Vec<u8> {
    std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))
}

/// A `WAI1` envelope's capability and payload.
fn envelope(env: &[u8]) -> Option<(String, Vec<u8>)> {
    if env.get(..4)? != b"WAI1" {
        return None;
    }
    let ml = u32::from_le_bytes(env.get(4..8)?.try_into().ok()?) as usize;
    let man: Value = serde_json::from_slice(env.get(8..8 + ml)?).ok()?;
    let at = 8 + ml;
    let pl = u32::from_le_bytes(env.get(at..at + 4)?.try_into().ok()?) as usize;
    let cap = man["model_requirement"]["capability"].as_str()?.to_string();
    Some((cap, env.get(at + 4..at + 4 + pl)?.to_vec()))
}

/// Every committed `WIR1` layer of the staged corpus that its parser admits, with its
/// name and unit kind. (The corpus also commits layers each parse refuses, for the
/// codes they carry; those are the refusal fixtures of the later gates.)
fn committed() -> Vec<(String, UnitKind, Vec<u8>)> {
    let cases = wai().join("staged-conformance/cases");
    let mut dirs: Vec<_> = std::fs::read_dir(&cases).unwrap().map(|e| e.unwrap().path()).collect();
    dirs.sort();
    let mut out = Vec::new();
    for d in dirs {
        let mut files: Vec<_> = std::fs::read_dir(&d).unwrap().map(|e| e.unwrap().path()).filter(|p| p.extension().is_some_and(|x| x == "wai")).collect();
        files.sort();
        for f in files {
            let Some((cap, p)) = envelope(&read(&f)) else { continue };
            let Some(kind) = wai::staged::unit_kind(&cap) else { continue };
            if p.starts_with(b"WIR1") && wir1::parse(&p, kind).is_ok() {
                out.push((format!("{}/{}", d.file_name().unwrap().to_string_lossy(), f.file_name().unwrap().to_string_lossy()), kind, p));
            }
        }
    }
    out
}

/// E9's 32-bit form holds on `[lo, hi]` (§5E E9, *Width*): `hi − lo ≤ 65535`, and
/// `lo − 131070` and `hi + 131070` fit `i32`.
fn e9_width_holds((lo, hi): (i64, i64)) -> bool {
    lo < hi && hi - lo <= 65535 && lo - 131070 >= i32::MIN as i64 && hi + 131070 <= i32::MAX as i64
}

/// Gate 1. Every committed layer has its stage form, and the form is the payload's:
/// E1's modes `t4097`, `lenient`, `canonical`, every table one E1 admits, and the
/// layer's offset; the prediction its predictor names (identity, E11 for up2, E12 for
/// rate at the layer's rates); E9's step is `q` on the unit kind's range, `[0, 255]` or
/// `[−32768, 32767]`, on both of which E9's 32-bit form holds. The committed layers
/// reach every unit kind and every predictor.
#[test]
fn every_committed_layer_has_its_stage_form() {
    assert_eq!((ops::MODES.table, ops::MODES.escape, ops::MODES.end), (TableRule::T4097, Escape::Lenient, End::Canonical));
    assert_eq!((ops::range(UnitKind::Rgb8), ops::range(UnitKind::Attr), ops::range(UnitKind::Pcm)), ((0, 255), (0, 255), (-32768, 32767)));
    let all = committed();
    let mut seen = BTreeSet::new();
    for (name, kind, p) in &all {
        let w = wir1::parse(p, *kind).unwrap();
        let f = ops::stages(&w).unwrap_or_else(|e| panic!("{name}: {e}"));
        for t in 0..w.n_tables() {
            assert!(rans64::table_ok(w.table(t), TableRule::T4097), "{name}: table {t}");
        }
        assert_eq!(f.offset, w.offset as i64, "{name}: E1's offset");
        assert_eq!(f.q, w.q as i64, "{name}: E9's step");
        assert_eq!(f.range, ops::range(*kind), "{name}: E9's range");
        assert!(e9_width_holds(f.range), "{name}: E9's 32-bit form holds on {:?}", f.range);
        let want = match w.predictor {
            Predictor::Identity => ops::Prediction::Identity,
            Predictor::Up2 => ops::Prediction::Up2,
            Predictor::Rate => ops::Prediction::Rate { num: w.rate_num, den: w.rate_den },
        };
        assert_eq!(f.prediction, want, "{name}: the prediction");
        seen.insert(format!("{kind:?}/{}", w.predictor.token()));
    }
    let want: BTreeSet<String> = ["Attr/identity", "Pcm/identity", "Pcm/rate", "Rgb8/identity", "Rgb8/up2"].iter().map(|s| s.to_string()).collect();
    assert_eq!(seen, want, "the committed layers reach every unit kind and predictor");
    println!("{} committed WIR1 layers, each with its stage form", all.len());
}

/// Gate 1. E1's table rule `t4097` is `WIR1`'s: the parser admits a table exactly when
/// E1 does, on a coded layer's table and on each rule broken once; and the stage form
/// refuses a table E1 refuses.
#[test]
fn e1s_table_rule_is_the_payloads() {
    let below: Vec<u8> = (0..4 * 4 * 3).map(|i| (i * 5) as u8).collect();
    let target: Vec<u8> = (0..4 * 4 * 3).map(|i| (i * 7 % 251) as u8).collect();
    let p = RefineParams { q: 3, predictor: Predictor::Identity, rate: (0, 0), table_per_unit: false, max_symbols: 4095 };
    let (ok, _) = wir1::encode(UnitKind::Rgb8, &[Below::U8(&below)], (4, 4, 3), &[Below::U8(&target)], (4, 4, 3), &p).unwrap();
    let w = wir1::parse(&ok, UnitKind::Rgb8).unwrap();
    let n = w.cdf_len;
    let t = |payload: &[u8], k: usize| i32::from_le_bytes(payload[44 + 4 * k..48 + 4 * k].try_into().unwrap());
    let set = |payload: &mut Vec<u8>, k: usize, v: i32| payload[44 + 4 * k..48 + 4 * k].copy_from_slice(&v.to_le_bytes());
    let mut cases: Vec<(&str, Vec<u8>)> = vec![("the coded table", ok.clone())];
    let mut c = ok.clone();
    set(&mut c, 0, 1);
    cases.push(("a first entry of 1", c));
    let mut c = ok.clone();
    set(&mut c, n - 1, 65535);
    cases.push(("a last entry of 65535", c));
    let mut c = ok.clone();
    set(&mut c, n - 1, 65537);
    cases.push(("a last entry of 65537", c));
    let mut c = ok.clone();
    let v1 = t(&c, 1);
    set(&mut c, 2, v1);
    cases.push(("two equal entries", c));
    let mut c = ok.clone();
    let (v1, v2) = (t(&c, 1), t(&c, 2));
    set(&mut c, 1, v2);
    set(&mut c, 2, v1);
    cases.push(("a decrease", c));
    for (what, payload) in cases {
        let table: Vec<i32> = (0..n).map(|k| t(&payload, k)).collect();
        let e1 = rans64::table_ok(&table, TableRule::T4097);
        let parsed = wir1::parse(&payload, UnitKind::Rgb8);
        assert_eq!(parsed.is_ok(), e1, "{what}: the parser and E1 disagree");
        match parsed {
            Ok(w) => assert!(ops::stages(&w).is_ok(), "{what}"),
            Err(e) => assert_eq!(e.code(), "cdf_malformed", "{what}"),
        }
    }
    // Lengths: the header's cdf_len is 3..=4097 (field_range past it), E1's t4097 too.
    for (len, admits) in [(2usize, false), (3, true), (4097, true), (4098, false)] {
        let cdf: Vec<i32> = (0..len).map(|k| if k + 1 == len { 65536 } else { k as i32 }).collect();
        assert_eq!(rans64::table_ok(&cdf, TableRule::T4097), admits, "{len} entries");
        assert_eq!((wir1::CDF_LEN_MIN..=wir1::CDF_LEN_MAX).contains(&len), admits, "{len} entries: the header's bound");
    }
}

/// Gate 1. The operators' conformance vectors cover what this stage form uses: E1 under
/// `t4097`, `lenient` and `canonical`; E11 and E12; E9 on both ranges, at `r = ±R` and
/// `±(R + 1)`.
#[test]
fn the_stage_vectors_cover_the_stage_form() {
    let lines = |op: &str| -> Vec<Value> {
        let t = String::from_utf8(read(&wai().join(format!("ops-conformance/stage/{op}.jsonl")))).unwrap();
        t.lines().map(|l| serde_json::from_str(l).unwrap()).collect()
    };
    let e1 = lines("rans64_decode");
    assert!(e1.iter().any(|v| v["modes"]["table"] == "t4097" && v["modes"]["escape"] == "lenient" && v["modes"]["end"] == "canonical"));
    assert!(!lines("up2_predict").is_empty() && !lines("rate_predict").is_empty(), "E11 and E12");
    let e9 = lines("residual_step");
    for (lo, hi) in [(0i64, 255i64), (-32768, 32767)] {
        let on: Vec<&Value> = e9.iter().filter(|v| v["in"]["lo"] == lo.to_string().as_str() && v["in"]["hi"] == hi.to_string().as_str()).collect();
        let at = |k: i64| on.iter().any(|v| {
            let q: i64 = v["in"]["q"].as_str().unwrap().parse().unwrap();
            let r: i64 = v["in"]["r"].as_str().unwrap().parse().unwrap();
            let big_r = (hi - lo) / q + 1;
            r == k * big_r || r == k * (big_r + 1)
        });
        assert!(at(1) && at(-1), "E9 on [{lo}, {hi}] at ±R and ±(R + 1)");
    }
}

/// A unit's current output for `w`'s predictor: the shape it reads (for rate, the
/// fewest source samples whose length `⌊n′·rate_num/rate_den⌋` is `d0`, if any), filled
/// with a fixed pattern that reaches both ends of the element range.
fn below_for(w: &wir1::Wir1, seed: u64) -> Option<((u32, u32, u32), Buf)> {
    let shape = match w.predictor {
        Predictor::Identity => (w.d0, w.d1, w.d2),
        Predictor::Up2 => (w.d0.div_ceil(2), w.d1.div_ceil(2), w.d2),
        Predictor::Rate => {
            let (num, den) = (w.rate_num as u64, w.rate_den as u64);
            let mut n = (w.d0 as u64 * den).div_ceil(num).max(1);
            while n * num / den < w.d0 as u64 {
                n += 1;
            }
            if n * num / den != w.d0 as u64 || n > u32::MAX as u64 {
                return None;
            }
            (n as u32, 1, w.d2)
        }
    };
    if !w.accepts_below(shape) {
        return None;
    }
    let len = shape.0 as u64 * shape.1 as u64 * shape.2 as u64;
    let at = |i: u64| -> u64 { (i.wrapping_add(seed)).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 17 };
    Some((shape, match w.kind {
        UnitKind::Pcm => Buf::I16((0..len).map(|i| match i % 23 { 0 => i16::MIN, 1 => i16::MAX, _ => at(i) as i16 }).collect()),
        _ => Buf::U8((0..len).map(|i| match i % 23 { 0 => 0, 1 => 255, _ => at(i) as u8 }).collect()),
    }))
}

/// Where two decodes of a unit first differ: the element index, or the refusals.
fn first_difference(a: &Result<wir1::Next, wir1::Wir1Error>, b: &Result<wir1::Next, wir1::Wir1Error>) -> String {
    let wide = |n: &wir1::Next| -> Vec<i64> {
        match n {
            wir1::Next::U8(v) => v.iter().map(|&x| x as i64).collect(),
            wir1::Next::I16(v) => v.iter().map(|&x| x as i64).collect(),
        }
    };
    match (a, b) {
        (Ok(x), Ok(y)) => {
            let (x, y) = (wide(x), wide(y));
            match x.iter().zip(&y).position(|(p, q)| p != q) {
                Some(i) => format!("element {i}: {} through the operator set, {} directly", x[i], y[i]),
                None => format!("lengths {} and {}", x.len(), y.len()),
            }
        }
        _ => format!("{:?} and {:?}", a.as_ref().err(), b.as_ref().err()),
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
enum Buf {
    U8(Vec<u8>),
    I16(Vec<i16>),
}

impl Buf {
    fn below(&self) -> Below<'_> {
        match self {
            Buf::U8(v) => Below::U8(v),
            Buf::I16(v) => Below::I16(v),
        }
    }
}

/// Gate 2: golden equality. Every committed `WIR1` layer, every one of its units, decodes
/// through the stage operators (`wir1::ops`) to the elements the direct decode
/// (`wir1::legacy`) gives, over a current output of the shape its predictor reads that
/// reaches both ends of the element range; and every layer's streams are checked alike.
/// (The corpus's own stage digests through the operator set are byte-exact-conformance's
/// `staged_sets_through_the_stage_operators_are_byte_exact` on every leg of the matrix.)
#[test]
fn every_committed_layer_applies_through_the_stage_operators_as_directly() {
    let (mut units, mut layers) = (0, 0);
    for (name, kind, p) in committed() {
        let w = wir1::parse(&p, kind).unwrap();
        assert_eq!(ops::check_streams(&w, &p), wir1::legacy::check_streams(&w, &p), "{name}: the streams' check");
        let Some((shape, below)) = below_for(&w, layers as u64) else { continue };
        for u in 0..w.n_units {
            let through = ops::apply_unit(&w, &p, u, below.below(), shape);
            let direct = wir1::legacy::apply_unit(&w, &p, u, below.below(), shape);
            if through != direct {
                panic!("{name} unit {u}: the decode through the stage operators differs from the direct decode, first at {:?}", first_difference(&through, &direct));
            }
            units += 1;
        }
        layers += 1;
    }
    assert!(layers >= 100, "{layers} layers applied");
    println!("{layers} committed layers, {units} units, through the stage operators as directly");
}

/// A one-unit layer coded by the reference encoder: `kind`, `predictor`, over a current
/// output of `below_shape` toward `dims`.
fn coded(kind: UnitKind, predictor: Predictor, rate: (u32, u32), below_shape: (u32, u32, u32), dims: (u32, u32, u32)) -> (Vec<u8>, Buf) {
    let n_below = (below_shape.0 * below_shape.1 * below_shape.2) as usize;
    let n = (dims.0 * dims.1 * dims.2) as usize;
    let p = RefineParams { q: 3, predictor, rate, table_per_unit: false, max_symbols: 4095 };
    match kind {
        UnitKind::Pcm => {
            let below: Vec<i16> = (0..n_below).map(|i| ((i * 2654) % 65536) as i16).collect();
            let target: Vec<i16> = (0..n).map(|i| ((i * 7919) % 65536) as i16).collect();
            let (out, _) = wir1::encode(kind, &[Below::I16(&below)], below_shape, &[Below::I16(&target)], dims, &p).unwrap();
            (out, Buf::I16(below))
        }
        _ => {
            let below: Vec<u8> = (0..n_below).map(|i| (i * 5) as u8).collect();
            let target: Vec<u8> = (0..n).map(|i| (i * 7 % 251) as u8).collect();
            let (out, _) = wir1::encode(kind, &[Below::U8(&below)], below_shape, &[Below::U8(&target)], dims, &p).unwrap();
            (out, Buf::U8(below))
        }
    }
}

/// `p`, a one-unit layer of `kind`, with its stream replaced by `s` (its length field
/// rewritten).
fn with_stream(kind: UnitKind, p: &[u8], s: &[u8]) -> Vec<u8> {
    let w = wir1::parse(p, kind).unwrap();
    let at = 44 + 4 * w.cdf_len * w.n_tables();
    let mut out = p[..at].to_vec();
    out.extend((s.len() as u32).to_le_bytes());
    out.extend(s);
    out
}

/// Both decodes of unit 0 of `p` refuse alike or decode alike, and so do both checks of
/// its streams; returns the unit's refusal code, if any.
fn alike(what: &str, kind: UnitKind, p: &[u8], below: &Buf, shape: (u32, u32, u32)) -> Option<&'static str> {
    let Ok(w) = wir1::parse(p, kind) else { return Some("parse") };
    assert_eq!(ops::check_streams(&w, p), wir1::legacy::check_streams(&w, p), "{what}: the streams' check");
    let through = ops::apply_unit(&w, p, 0, below.below(), shape);
    let direct = wir1::legacy::apply_unit(&w, p, 0, below.below(), shape);
    if through != direct {
        panic!("{what}: through the stage operators and directly differ, {}", first_difference(&through, &direct));
    }
    through.err().map(|e| e.code())
}

/// Gate 3: refusal parity. (a) Every committed layer whose streams are refused (the
/// corpus's `stream_exhausted` and `stream_noncanonical` cases among them) is refused by
/// both checks alike, with the same unit and code. (b) A.4 step 1, each rule broken
/// once for each predictor: both decodes refuse `geometry_mismatch` before any symbol.
/// (c) A coded layer of each unit kind and predictor, its stream cut at every word, a
/// word appended, and bytes changed across it: both decode alike or refuse alike, the
/// streams' check too.
#[test]
fn every_refusal_is_the_direct_decodes() {
    // (a)
    let mut refused = std::collections::BTreeMap::new();
    for (name, kind, p) in committed() {
        let w = wir1::parse(&p, kind).unwrap();
        let (through, direct) = (ops::check_streams(&w, &p), wir1::legacy::check_streams(&w, &p));
        assert_eq!(through, direct, "{name}");
        if let Err((_, e)) = through {
            *refused.entry(e.code()).or_insert(0) += 1;
            for case in ["stream_exhausted", "stream_noncanonical"] {
                if name.starts_with(&format!("{case}/")) {
                    assert_eq!(e.code(), case, "{name}");
                }
            }
        }
    }
    assert!(refused.contains_key("stream_exhausted") && refused.contains_key("stream_noncanonical"), "{refused:?}");
    // (b)
    let layers = [
        (UnitKind::Rgb8, Predictor::Identity, (0, 0), (5, 4, 3), (5, 4, 3)),
        (UnitKind::Rgb8, Predictor::Up2, (0, 0), (3, 2, 3), (5, 4, 3)),
        (UnitKind::Attr, Predictor::Identity, (0, 0), (2, 3, 4), (2, 3, 4)),
        (UnitKind::Pcm, Predictor::Identity, (0, 0), (40, 1, 2), (40, 1, 2)),
        (UnitKind::Pcm, Predictor::Rate, (3, 2), (20, 1, 2), (30, 1, 2)),
    ];
    let mut geometry = 0;
    for &(kind, predictor, rate, shape, dims) in &layers {
        let (p, below) = coded(kind, predictor, rate, shape, dims);
        let what = format!("{kind:?} {}", predictor.token());
        assert_eq!(alike(&what, kind, &p, &below, shape), None, "{what}: the coded layer");
        let w = wir1::parse(&p, kind).unwrap();
        let other = match &below {
            Buf::U8(v) => Buf::I16(v.iter().map(|&x| x as i16).collect()),
            Buf::I16(v) => Buf::U8(v.iter().map(|&x| x as u8).collect()),
        };
        let short = match &below {
            Buf::U8(v) => Buf::U8(v[1..].to_vec()),
            Buf::I16(v) => Buf::I16(v[1..].to_vec()),
        };
        let bad_shape = (shape.0 + 1, shape.1, shape.2);
        for (rule, b, s, unit) in [
            ("the unit outside the payload", &below, shape, 1u32),
            ("the other element type", &other, shape, 0),
            ("an element short", &short, shape, 0),
            ("a shape the predictor does not read", &below, bad_shape, 0),
            ("a shape past u64", &below, (u32::MAX, u32::MAX, u32::MAX), 0),
        ] {
            let through = ops::apply_unit(&w, &p, unit, b.below(), s);
            assert_eq!(through, wir1::legacy::apply_unit(&w, &p, unit, b.below(), s), "{what}: {rule}");
            assert_eq!(through.unwrap_err().code(), "geometry_mismatch", "{what}: {rule}");
            geometry += 1;
        }
    }
    // (c)
    let mut n = 0;
    for &(kind, predictor, rate, shape, dims) in &layers {
        let (p, below) = coded(kind, predictor, rate, shape, dims);
        let w = wir1::parse(&p, kind).unwrap();
        let s = w.stream(&p, 0).to_vec();
        let what = format!("{kind:?} {}", predictor.token());
        for cut in (8..s.len()).step_by(4) {
            let r = alike(&format!("{what}: the stream cut to {cut} bytes"), kind, &with_stream(kind, &p, &s[..cut]), &below, shape);
            assert!(r.is_some(), "{what}: the stream cut to {cut} bytes decodes");
            n += 1;
        }
        let mut long = s.clone();
        long.extend([0, 0, 0, 0]);
        assert_eq!(alike(&format!("{what}: a word appended"), kind, &with_stream(kind, &p, &long), &below, shape), Some("stream_noncanonical"));
        for at in 0..s.len() {
            let mut c = s.clone();
            c[at] ^= 1 << (at % 8);
            alike(&format!("{what}: byte {at} changed"), kind, &with_stream(kind, &p, &c), &below, shape);
            n += 1;
        }
    }
    println!("{refused:?} among the committed layers; {geometry} geometry fixtures and {n} changed streams refused alike");
}

/// Gate 5: the 32-bit form. Every committed layer's every unit, over gate 2's current
/// outputs that reach both ends of the range, and gate 3's coded layers with every byte
/// of their streams changed, apply with E9 in its 32-bit form (`ops::apply_unit_i32`, the
/// arithmetic a 32-bit fabric runs) as with the exact form.
#[cfg(feature = "i32_forms")]
#[test]
fn the_32bit_form_applies_every_committed_layer_alike() {
    let mut n = 0;
    for (name, kind, p) in committed() {
        let w = wir1::parse(&p, kind).unwrap();
        let Some((shape, below)) = below_for(&w, n as u64) else { continue };
        for u in 0..w.n_units {
            let (narrow, exact) = (ops::apply_unit_i32(&w, &p, u, below.below(), shape), ops::apply_unit(&w, &p, u, below.below(), shape));
            if narrow != exact {
                panic!("{name} unit {u}: the 32-bit form differs, {}", first_difference(&narrow, &exact));
            }
            n += 1;
        }
    }
    for (kind, predictor, rate, shape, dims) in [
        (UnitKind::Rgb8, Predictor::Up2, (0, 0), (3, 2, 3), (5, 4, 3)),
        (UnitKind::Pcm, Predictor::Rate, (3, 2), (20, 1, 2), (30, 1, 2)),
    ] {
        let (p, below) = coded(kind, predictor, rate, shape, dims);
        let s = wir1::parse(&p, kind).unwrap().stream(&p, 0).to_vec();
        for at in 0..s.len() {
            let mut c = s.clone();
            c[at] ^= 1 << (at % 8);
            let q = with_stream(kind, &p, &c);
            let w = wir1::parse(&q, kind).unwrap();
            assert_eq!(ops::apply_unit_i32(&w, &q, 0, below.below(), shape), ops::apply_unit(&w, &q, 0, below.below(), shape), "{kind:?}: byte {at} changed");
            n += 1;
        }
    }
    println!("{n} units through the 32-bit form, each as through the exact form");
}