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

// wai-conformance: int-kernels int-kernels-amd
//! `wai.image.jpegai`'s scale path and entropy decodes through the integer operator set:
//! its lowering `jpegai-int/1` (`extensions/integer-operators.md` §13.6.1) and the
//! certificate of each bundle (§8.5), as this crate loads them (`jpegai_ops`).
//!
//! The committed bundle (`ops-conformance/sets/jpegai`, synthetic) has its lowered graphs
//! and certificate under `ops-conformance/certificates/`, where `wai-verify` and the
//! independent checker (`tools/wai_ops_check.py`) reproduce them byte for byte. The
//! `jpegai/v1` bundle of the models root is not committed; its certificate is, by digest,
//! and the tests that read it take its files from `WAI_JPEGAI_WEIGHTS_DIR` and
//! `WAI_JPEGAI_NORM` (`tools/wai_jpegai_fetch_e2e_inputs.sh`), with
//! `WAI_REQUIRE_SOTA_VECTORS` deciding whether their absence is a skip or a failure.
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]

#[path = "support/require_vectors.rs"]
mod require_vectors;

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

use serde_json::Value;
use wai::jpegai_decode::NormTables;
use wai::jpegai_ops;
use wai::jpegai_synthesis::IntNet;
use wai_ops::Profile;

const CAP: &str = "wai.image.jpegai";
const SYNTHETIC: &str = "1bc0daae512797a8817f2e77e57d009ace0d22acf4fe6945af8f4e0ee95244b9";
const V1: &str = "3e243dc11566922f5cb7f4c628afd440b18d56a08b99b07cc33d5c8522aeec37";
const PRODUCE: &str = "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)";

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()))
}

fn json(b: &[u8]) -> Value {
    serde_json::from_slice(b).unwrap()
}

fn sha256(b: &[u8]) -> String {
    wai_ops::sha256::sha256(b).iter().map(|x| format!("{x:02x}")).collect()
}

fn i64s(v: &Value) -> Vec<i64> {
    v.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()
}

/// The tables of `norm.json` as the decoders parse them. A synthetic set has only the
/// members §13.6.1 reads; the others are left empty.
fn norm_tables(nv: &Value) -> NormTables {
    let opt = |k: &str| nv.get(k).map(i64s).unwrap_or_default();
    let three = |k: &str, d: f64| -> [f64; 3] {
        nv.get(k).map(|a| std::array::from_fn(|i| a[i].as_f64().unwrap())).unwrap_or([d; 3])
    };
    let t = three("lef_thr0", 0.0);
    let m = three("lef_mag0", 0.0);
    NormTables {
        pmf: nv["pmf"].as_array().unwrap().iter().map(i64s).collect(),
        bounds: i64s(&nv["bounds"]),
        per_channel: nv["per_channel"].as_u64().unwrap() as usize,
        cdfs_y: i64s(&nv["cdfs_y"]),
        cdfs_uv: i64s(&nv["cdfs_uv"]),
        gain_y: i64s(&nv["gain_y"]),
        gain_uv: i64s(&nv["gain_uv"]),
        lsbs_t0: opt("lsbs_t0").iter().map(|&x| x as i32).collect(),
        lsbs_t1: opt("lsbs_t1").iter().map(|&x| x as i32).collect(),
        lef_thr: t.map(|x| x as i64),
        lef_mag: m.map(|x| x as f32),
    }
}

/// A bundle's three scale-path members, as bytes: `norm.json`, `hs_y.json`, `hs_uv.json`.
struct Bundle {
    norm: Vec<u8>,
    hs_y: Vec<u8>,
    hs_uv: Vec<u8>,
}

impl Bundle {
    fn synthetic() -> Bundle {
        let d = wai().join("ops-conformance/sets/jpegai");
        Bundle { norm: read(&d.join("norm.json")), hs_y: read(&d.join("weights/hs_y.json")), hs_uv: read(&d.join("weights/hs_uv.json")) }
    }
    /// The `jpegai/v1` bundle, from the variables the whole-decode tests read.
    fn pinned(test: &str) -> Option<(Bundle, PathBuf)> {
        let [dir, norm] = require_vectors::env_paths(test, ["WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
        let d = PathBuf::from(dir);
        Some((Bundle { norm: read(Path::new(&norm)), hs_y: read(&d.join("hs_y.json")), hs_uv: read(&d.join("hs_uv.json")) }, d))
    }
    fn parsed(&self) -> (NormTables, IntNet, IntNet) {
        (norm_tables(&json(&self.norm)), IntNet::from_json(&json(&self.hs_y)).unwrap(), IntNet::from_json(&json(&self.hs_uv)).unwrap())
    }
}

/// Gate 1, on one bundle: the lowering of its bytes; the bundle as this crate parses it,
/// converted (`jpegai_ops::net`), is the bundle `wai-ops` parses from those bytes and
/// lowers to the same graphs; the scale path this crate loads (`jpegai_ops::load`) holds
/// those graphs and meets `acc32`; and the committed certificate `digest` is those graphs',
/// each of whose nodes meets `acc32` with no accelerator flag (the scale path's
/// `role_profiles`, §11).
fn lowers_to_its_certificate(name: &str, b: &Bundle, digest: &str) {
    let lowered = wai_ops::lower::jpegai_scale(&b.norm, &b.hs_y, &b.hs_uv).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
    let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
    assert_eq!(roles, ["hs_y", "hs_uv"], "{name}");
    let (n, y, uv) = b.parsed();
    let net = jpegai_ops::net(&n, &y, &uv).unwrap_or_else(|r| panic!("{name}: does not convert: {r:?}"));
    assert!(net == wai_ops::lower::parse_jpegai(&b.norm, &b.hs_y, &b.hs_uv).unwrap(), "{name}: the converted bundle is the one wai-ops parses from the bytes");
    assert!(wai_ops::lower::jpegai_net(&net).unwrap() == lowered, "{name}: the lowering of the parsed bundle differs from the lowering of its bytes");
    let sp = jpegai_ops::load(&n, &y, &uv).unwrap_or_else(|r| panic!("{name}: refused under acc32: {r:?}"));
    let cert = json(&read(&wai().join(format!("ops-conformance/certificates/{CAP}/{digest}.json"))));
    assert_eq!(cert["set_sha256"].as_str(), Some(digest), "{name}");
    assert_eq!(cert["lowering"], jpegai_ops::LOWERING, "{name}");
    for (i, (role, plane)) in [("hs_y", jpegai_ops::Plane::Luma), ("hs_uv", jpegai_ops::Plane::Chroma)].into_iter().enumerate() {
        let g = lowered.get(role).unwrap();
        assert!(*sp.graph(plane).graph() == *g, "{name}: {role} as loaded");
        assert!(sp.graph(plane).meets(Profile::Acc32), "{name}: {role} meets acc32");
        assert_eq!(cert["graphs"][i]["name"], role, "{name}");
        assert_eq!(cert["graphs"][i]["graph_sha256"].as_str(), Some(sha256(&g.to_bytes()).as_str()), "{name}: the certificate is {role}'s");
        let x = &cert["graphs"][i]["inputs"][0];
        assert_eq!((x["policy"].as_str(), x["lo"].as_str(), x["hi"].as_str()), (Some("check"), Some("-31"), Some("32")), "{name}: z's range");
    }
    // The set's verdict covers its other networks too (§13.6.2), which meet acc64 only;
    // the scale path's two graphs meet acc32, in the accelerator subset, with no flags
    // (§11's `role_profiles`).
    assert_eq!(cert["verdict"]["acc64"].as_bool(), Some(true), "{name}: the certificate's verdict");
    for i in 0..2 {
        for n in cert["graphs"][i]["nodes"].as_array().unwrap() {
            assert!(n["fits"].as_array().unwrap().iter().any(|p| p == "acc32"), "{name}: node {} of the scale path fits acc32", n["index"]);
            assert_eq!(n["npu_flags"].as_array().map(Vec::len), Some(0), "{name}: no accelerator flags");
        }
    }
}

/// The committed synthetic bundle: its lowered scale graphs are the committed ones, and
/// its set digest is the listing of its eleven members.
#[test]
fn the_committed_bundle_lowers_to_its_committed_graphs_and_verdict() {
    let b = Bundle::synthetic();
    let dir = wai().join(format!("ops-conformance/certificates/{CAP}/{SYNTHETIC}"));
    let lowered = wai_ops::lower::jpegai_scale(&b.norm, &b.hs_y, &b.hs_uv).unwrap();
    for role in ["hs_y", "hs_uv"] {
        assert!(lowered.get(role).unwrap().to_bytes() == read(&dir.join(format!("{role}.whg1"))), "the lowering of the bytes differs from the committed {role}.whg1");
    }
    let nets: Vec<(String, Vec<u8>)> = wai_ops::lower::jpegai_nets::NET_MEMBERS.iter().map(|n| (n.to_string(), read(&wai().join("ops-conformance/sets/jpegai").join(n)))).collect();
    let mut members: BTreeMap<String, &[u8]> =
        [("norm.json", &b.norm), ("weights/hs_y.json", &b.hs_y), ("weights/hs_uv.json", &b.hs_uv)].into_iter().map(|(k, v)| (k.to_string(), v.as_slice())).collect();
    members.extend(nets.iter().map(|(k, v)| (k.clone(), v.as_slice())));
    let digest = wai::int_set_cache::SetKey::listing_digest(&members).unwrap();
    assert_eq!(digest.iter().map(|x| format!("{x:02x}")).collect::<String>(), SYNTHETIC, "the set's digest");
    lowers_to_its_certificate("synthetic", &b, SYNTHETIC);
}

/// The `jpegai/v1` bundle: its set digest is the listing of its twelve members, the one
/// its set-form pin and its committed certificate name.
#[test]
fn the_pinned_bundle_lowers_to_its_certificate_and_meets_acc32() {
    let Some((b, dir)) = Bundle::pinned("jpegai_ops::the_pinned_bundle_lowers_to_its_certificate_and_meets_acc32") else {
        return;
    };
    let names = ["hs_y", "hs_uv", "hyp_y", "hyp_uv", "synth_y", "synth_uv", "mcm0_y", "mcm1_y", "mcm2_y", "mcm3_y", "icci"];
    let files: Vec<(String, Vec<u8>)> = names.iter().map(|n| (format!("weights/{n}.json"), read(&dir.join(format!("{n}.json"))))).collect();
    let mut members: BTreeMap<String, &[u8]> = files.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
    members.insert("norm.json".into(), &b.norm);
    let digest = wai::int_set_cache::SetKey::listing_digest(&members).unwrap();
    assert_eq!(digest.iter().map(|x| format!("{x:02x}")).collect::<String>(), V1, "the set's digest");
    lowers_to_its_certificate("jpegai/v1", &b, V1);
    eprintln!("jpegai/v1: hs_y and hs_uv lower to the committed certificate {V1} and meet acc32");
}

/// The whole `jpegai/v1` bundle and the test codestream, as the decoders parse them.
struct Whole {
    data: Vec<u8>,
    weights: wai::jpegai_decode::DecoderWeights,
    norm: NormTables,
}

impl Whole {
    fn pinned(test: &str) -> Option<Whole> {
        use wai::jpegai_decode::{DecoderWeights, PlaneWeights};
        use wai::jpegai_synthesis::Weights;
        let [cs, dir, nf] = require_vectors::env_paths(test, ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"], PRODUCE)?;
        let d = PathBuf::from(dir);
        let j = |n: &str| json(&read(&d.join(n)));
        let w = |n: &str| Weights::from_json(&j(n)).unwrap();
        let net = |n: &str| IntNet::from_json(&j(n)).unwrap();
        let weights = DecoderWeights {
            luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
            chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
            mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
            icci: wai::jpegai_icci::IcciNet::from_json(&j("icci.json")).unwrap(),
        };
        Some(Whole { data: read(Path::new(&cs)), weights, norm: norm_tables(&json(&read(Path::new(&nf)))) })
    }
}

/// A codestream's segment payloads and a plane's GRFS flags, as the decoder reads them.
fn segments(data: &[u8], plane: jpegai_ops::Plane) -> (Vec<u8>, Vec<u8>, Vec<bool>) {
    use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
    let segs = parse_container(data);
    let pay = |m| {
        let s = find_segment(&segs, m).unwrap();
        data[s.start..s.start + s.len].to_vec()
    };
    let lat = if plane == jpegai_ops::Plane::Luma { Marker::Sorp } else { Marker::Sors };
    let ph = parse_picture_header(data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
    let grfs = ph.components[plane as usize].grfs_channel_flags.clone();
    (pay(Marker::Soz), pay(lat), grfs)
}

/// Gate 2, the integers: on the test codestream, each plane's hyper-latent, scale,
/// mask, `likely`, RVS addition, `σ_idx` and latent through the operator set are the
/// direct computation's, every value. The chroma plane's hyper-latent comes after the 160
/// luma rows its decode reads and discards.
#[test]
fn the_planes_decode_through_the_lowering_as_directly() {
    let Some(wh) = Whole::pinned("jpegai_ops::the_planes_decode_through_the_lowering_as_directly") else {
        return;
    };
    let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap();
    for (plane, pw, gain) in [
        (jpegai_ops::Plane::Luma, &wh.weights.luma, &wh.norm.gain_y),
        (jpegai_ops::Plane::Chroma, &wh.weights.chroma, &wh.norm.gain_uv),
    ] {
        let (soz, lat, grfs) = segments(&wh.data, plane);
        let ops = jpegai_ops::plane(&sp, plane, &soz, &lat, &grfs).unwrap();
        let direct = jpegai_ops::legacy::plane(&wh.norm, &pw.hyper_scale, gain, plane, &soz, &lat, &grfs);
        for (what, a, b) in [("z", &ops.z, &direct.z), ("v", &ops.v, &direct.v), ("likely", &ops.likely, &direct.likely), ("rvs", &ops.rvs, &direct.rvs), ("σ_idx", &ops.sigma, &direct.sigma), ("y_q", &ops.y_q, &direct.y_q)] {
            assert!(a == b, "{plane:?}: {what} through the lowering differs from the direct computation");
        }
        assert!(ops.mask == direct.mask, "{plane:?}: the mask");
        let coded = ops.mask.iter().filter(|&&m| m).count();
        eprintln!("{plane:?}: z, v, mask ({coded} coded), likely, rvs, σ_idx and y_q equal the direct computation's");
    }
}

/// Gate 2, the decode: the whole decode through the operator set gives the direct
/// decode's float values after every stage, bit for bit, and its RGB bytes, which hash to
/// `WAI_JPEGAI_EXPECT_SHA256` where that is set.
#[test]
fn the_whole_decode_through_the_lowering_is_the_direct_decode() {
    let Some(wh) = Whole::pinned("jpegai_ops::the_whole_decode_through_the_lowering_is_the_direct_decode") else {
        return;
    };
    let bits = |planes: &[&[f32]]| -> String { sha256(&planes.iter().flat_map(|p| p.iter().flat_map(|v| v.to_bits().to_le_bytes())).collect::<Vec<u8>>()) };
    let mut direct = vec![];
    let a = wai::jpegai_decode::legacy::decode_with_stages(&wh.data, &wh.weights, &wh.norm, &mut |s, p| direct.push((s.name(), bits(p)))).expect("the direct decode");
    let mut ops = vec![];
    let b = wai::jpegai_decode::decode_with_stages(&wh.data, &wh.weights, &wh.norm, &mut |s, p| ops.push((s.name(), bits(p)))).expect("the decode through the operator set");
    assert_eq!(ops, direct, "the stages through the lowering");
    assert!(a.data == b.data, "the RGB bytes through the lowering");
    let rgb = sha256(&b.data);
    match std::env::var("WAI_JPEGAI_EXPECT_SHA256") {
        Ok(want) if !want.is_empty() => {
            assert_eq!(rgb, want, "the decode through the operator set is not the pinned decode");
            eprintln!("the decode through the operator set matches the pinned WAI_JPEGAI_EXPECT_SHA256 ({} stages equal)", ops.len());
        }
        _ => eprintln!("the decode through the operator set equals the direct decode ({} stages), rgb {rgb}", ops.len()),
    }
}

/// One member of the committed bundle changed by `edit`: the bundle's three members as
/// bytes, and whether this crate's parser can read it at all (a non-integer value it
/// cannot).
fn edited(edit: impl Fn(&mut Value, &mut Value)) -> (Bundle, bool) {
    let b = Bundle::synthetic();
    let (mut n, mut y) = (json(&b.norm), json(&b.hs_y));
    edit(&mut n, &mut y);
    let ints = |v: &Value| v.as_array().is_none_or(|a| a.iter().all(|x| x.is_i64() || x.is_array() || x.is_object()));
    let parseable = y.as_object().unwrap().values().all(|m| ints(&m["v"]));
    (Bundle { norm: serde_json::to_vec(&n).unwrap(), hs_y: serde_json::to_vec(&y).unwrap(), hs_uv: b.hs_uv }, parseable)
}

/// Gate 3, the bundle: each load check of §13.6.1, broken one change at a time from the
/// committed bundle, is refused `model_mismatch`, by the load from the members' bytes and
/// by the load from the bundle as this crate parses it; and so is a bundle whose scale
/// graphs do not meet `acc32`. The direct decode made none of these checks
/// (integer-operators §15.3 item 4): it read a table value into a byte, shifted by
/// whatever a shift said, and accumulated in `i64`.
#[test]
fn every_load_check_of_the_bundle_refuses_model_mismatch() {
    let arr = |v: &mut Value, k: &str| -> Vec<Value> { v[k].as_array().unwrap().clone() };
    let mut cases: Vec<(&str, (Bundle, bool))> = vec![
        ("pmf: 31 rows", edited(|n, _| { let mut p = arr(n, "pmf"); p.truncate(31); n["pmf"] = Value::from(p); })),
        ("pmf: a row summing to 257", edited(|n, _| { let r = &mut n["pmf"][0][0]; *r = Value::from(r.as_i64().unwrap() + 1); })),
        ("pmf: a row of another width", edited(|n, _| { n["pmf"][1].as_array_mut().unwrap().push(Value::from(0)); })),
        ("pmf: a negative count", edited(|n, _| { let a = n["pmf"][0][0].as_i64().unwrap(); let b = n["pmf"][0][1].as_i64().unwrap(); n["pmf"][0][0] = Value::from(-1); n["pmf"][0][1] = Value::from(a + b + 1); })),
        ("bounds: 31", edited(|n, _| { let mut b = arr(n, "bounds"); b.truncate(31); n["bounds"] = Value::from(b); })),
        ("bounds: one of 256", edited(|n, _| { n["bounds"][0] = Value::from(256); })),
        ("bounds: one of -1", edited(|n, _| { n["bounds"][0] = Value::from(-1); })),
        ("per_channel: 15", edited(|n, _| { n["per_channel"] = Value::from(15); })),
        ("cdfs_y: one short", edited(|n, _| { let mut c = arr(n, "cdfs_y"); c.pop(); n["cdfs_y"] = Value::from(c); })),
        ("cdfs_y: a count of 256", edited(|n, _| { n["cdfs_y"][62] = Value::from(256); })),
        ("cdfs_y: a row that decreases", edited(|n, _| { let v = n["cdfs_y"][1].as_i64().unwrap(); n["cdfs_y"][0] = Value::from(v + 1); n["cdfs_y"][1] = Value::from(v); })),
        ("cdfs_uv: one long", edited(|n, _| { let mut c = arr(n, "cdfs_uv"); c.push(Value::from(255)); n["cdfs_uv"] = Value::from(c); })),
        ("gain_y: 159", edited(|n, _| { let mut g = arr(n, "gain_y"); g.pop(); n["gain_y"] = Value::from(g); })),
        ("gain_uv: 97", edited(|n, _| { let mut g = arr(n, "gain_uv"); g.push(Value::from(0)); n["gain_uv"] = Value::from(g); })),
        ("hs_y: a member missing", edited(|_, y| { y.as_object_mut().unwrap().remove("onnx::Pow_91"); })),
        ("hs_y: weights of other dims", edited(|_, y| { y["conv1.weight"]["dims"] = serde_json::json!([160, 2, 1, 1]); })),
        ("hs_y: a shift of 128", edited(|_, y| { y["onnx::Pow_85"]["v"][0] = Value::from(128); })),
        ("hs_y: a shift of -1", edited(|_, y| { y["onnx::Pow_85"]["v"][0] = Value::from(-1); })),
        ("hs_y: a bias that is not an integer", edited(|_, y| { y["conv1.bias"]["v"][0] = Value::from(1.5); })),
    ];
    // Past acc32 and within acc64: block 3's sums exceed i32 with one weight of 2^30 over
    // inputs in [0, 127].
    let big = edited(|_, y| { y["pointwise.weight"]["v"][0] = Value::from(1i64 << 30); });
    let g = wai_ops::lower::jpegai_scale(&big.0.norm, &big.0.hs_y, &big.0.hs_uv).expect("the bundle lowers");
    assert!(wai_ops::graph::load(&g.get("hs_y").unwrap().to_bytes(), Profile::Acc64).is_ok(), "it meets acc64");
    assert!(wai_ops::graph::load(&g.get("hs_y").unwrap().to_bytes(), Profile::Acc32).is_err(), "and not acc32");
    cases.push(("hs_y: a sum past i32 (acc64, not acc32)", big));
    for (what, (b, parseable)) in &cases {
        let r = jpegai_ops::load_bytes(&b.norm, &b.hs_y, &b.hs_uv).err().unwrap_or_else(|| panic!("{what}: the load from bytes accepts it"));
        assert_eq!(r.code, wai_ops::Code::ModelMismatch, "{what}");
        if *parseable {
            let (n, y, uv) = b.parsed();
            let r = jpegai_ops::load(&n, &y, &uv).err().unwrap_or_else(|| panic!("{what}: the load from the parsed bundle accepts it"));
            assert_eq!(r.code, wai_ops::Code::ModelMismatch, "{what}");
        }
    }
    assert!(jpegai_ops::load_bytes(&Bundle::synthetic().norm, &Bundle::synthetic().hs_y, &Bundle::synthetic().hs_uv).is_ok(), "the committed bundle loads");
    eprintln!("{} load checks refuse model_mismatch from bytes and from the parsed bundle", cases.len());
}

/// How the direct front half came out on one input: its integers, or a panic.
fn direct(n: &NormTables, hs: &IntNet, gain: &[i64], p: jpegai_ops::Plane, soz: &[u8], lat: &[u8], grfs: &[bool]) -> Option<jpegai_ops::PlaneInts> {
    std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| jpegai_ops::legacy::plane(n, hs, gain, p, soz, lat, grfs))).ok()
}

/// The refusal parity of the entropy decodes on one bundle: each plane, with the SOZ
/// payload cut at every length, then the plane's latent payload cut at every length, then
/// each lengthened at its front by 16, 256 and 4096 bytes.
/// Where the operator set decodes, the direct decode gives the same integers. Where the
/// operator set refuses, the refusal is `stream_exhausted`, and the direct decode
/// panicked: in a build with overflow checks it panics at the read or the refill that
/// E8 refuses; in one without, it may instead compute on from a window that wrapped, and
/// those cases are counted. Returns (decoded alike, refused and panicked, refused where
/// the unchecked build computed on).
fn cut_streams(sp: &jpegai_ops::ScalePath, n: &NormTables, hs: [&IntNet; 2], data: &[u8]) -> (usize, usize, usize) {
    let mut tally = (0, 0, 0);
    for (k, p) in [jpegai_ops::Plane::Luma, jpegai_ops::Plane::Chroma].into_iter().enumerate() {
        let gain = if k == 0 { &n.gain_y } else { &n.gain_uv };
        let (soz, lat, grfs) = segments(data, p);
        // Bytes put before a payload are the last its backward window reaches: a payload
        // so lengthened decodes further before it runs out.
        let filler = |n: usize, seed: u64| -> Vec<u8> { (0..n as u64).map(|i| ((i.wrapping_add(seed)).wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 56) as u8).collect() };
        let longer = |pre: Vec<u8>, p: &[u8]| -> Vec<u8> { let mut v = pre; v.extend_from_slice(p); v };
        let cuts = (0..=soz.len()).map(|c| (soz[..c].to_vec(), lat.clone()))
            .chain((0..lat.len()).map(|c| (soz.clone(), lat[..c].to_vec())))
            .chain([16usize, 256, 4096].into_iter().flat_map(|n| [
                (longer(filler(n, 1), &soz), lat.clone()),
                (soz.clone(), longer(filler(n, 2), &lat)),
                (longer(filler(n, 3), &soz), longer(filler(n, 4), &lat)),
            ]))
            .collect::<Vec<_>>();
        for (s, l) in cuts {
            let ops = jpegai_ops::plane(sp, p, &s, &l, &grfs);
            match (ops, direct(n, hs[k], gain, p, &s, &l, &grfs)) {
                (Ok(a), Some(b)) => { assert!(a == b, "{p:?} ({}, {}): the decodes differ", s.len(), l.len()); tally.0 += 1 }
                (Err(r), None) => { assert_eq!(r.code, wai_ops::Code::StreamExhausted, "{p:?}"); tally.1 += 1 }
                (Err(r), Some(_)) => {
                    assert_eq!(r.code, wai_ops::Code::StreamExhausted, "{p:?}");
                    if cfg!(debug_assertions) {
                        panic!("{p:?} ({}, {}): refused, and the checked build's direct decode did not panic", s.len(), l.len());
                    }
                    tally.2 += 1
                }
                (Ok(_), None) => panic!("{p:?} ({}, {}): the direct decode panicked where the operator set decodes", s.len(), l.len()),
            }
        }
    }
    tally
}

/// Gate 3, the streams, on the committed bundle and the committed test codestream. The
/// committed bundle's tables are not the ones the codestream was coded with, so the
/// decodes run out early unless a payload is lengthened; the integers they give differ
/// from the pinned bundle's, and both decodes compute them alike.
#[test]
fn a_stream_cut_short_is_refused_where_the_direct_decode_panics() {
    let b = Bundle::synthetic();
    let (n, y, uv) = b.parsed();
    let sp = jpegai_ops::load(&n, &y, &uv).unwrap();
    let data = read(&wai().join("wai-rs/tests/vectors/jpegai/test256.bin"));
    let t = cut_streams(&sp, &n, [&y, &uv], &data);
    eprintln!("committed bundle: {} decoded alike, {} refused where the direct decode panicked, {} refused where the unchecked build computed on", t.0, t.1, t.2);
    assert!(t.0 > 0 && t.1 > 0, "both outcomes are reached");
}

/// Gate 3, the streams, on the pinned bundle.
#[test]
fn a_stream_cut_short_is_refused_where_the_direct_decode_panics_pinned() {
    let Some(wh) = Whole::pinned("jpegai_ops::a_stream_cut_short_is_refused_where_the_direct_decode_panics_pinned") else {
        return;
    };
    let sp = wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap();
    let t = cut_streams(&sp, &wh.norm, [&wh.weights.luma.hyper_scale, &wh.weights.chroma.hyper_scale], &wh.data);
    eprintln!("jpegai/v1: {} decoded alike, {} refused where the direct decode panicked, {} refused where the unchecked build computed on", t.0, t.1, t.2);
    assert!(t.0 > 0 && t.1 > 0, "both outcomes are reached");
}

/// Gate 3, the codestream: a codestream outside the decoder's configuration is refused
/// through the operator set with the reason `check_scope` gives, before the bundle is
/// read, as the direct decode refused it.
#[test]
fn a_codestream_out_of_scope_is_refused_with_its_reason() {
    let Some(wh) = Whole::pinned("jpegai_ops::a_codestream_out_of_scope_is_refused_with_its_reason") else {
        return;
    };
    use wai::jpegai_bitstream::{find_segment, parse_container, Marker};
    let pih = find_segment(&parse_container(&wh.data), Marker::Pih).unwrap().start_bit();
    let set = |d: &mut [u8], pos: usize, n: usize, value: u64| {
        for i in 0..n {
            let bit = (value >> (n - 1 - i)) & 1;
            let (byte, sh) = ((pos + i) / 8, 7 - (pos + i) % 8);
            d[byte] = (d[byte] & !(1 << sh)) | ((bit as u8) << sh);
        }
    };
    // model_id (4 bits at 85), the colour transform (2 bits at 83, the reserved 3), β
    // (12 bits at 90, past the encoder's range), the first synthesis transform (4 bits at
    // 12, the base-operation-point one).
    for (at, n, v) in [(85usize, 4usize, 1u64), (83, 2, 3), (90, 12, 2048 + 703), (12, 4, 1)] {
        let mut d = wh.data.clone();
        set(&mut d, pih + at, n, v);
        let why = wai::jpegai_decode::check_scope(&d).unwrap_err();
        assert!(wai::jpegai_decode::legacy::decode(&d, &wh.weights, &wh.norm).is_none(), "the direct decode refuses it");
        assert_eq!(wai::jpegai_decode::decode(&d, &wh.weights, &wh.norm).err(), Some(wai::jpegai_decode::Refusal::Scope(why)), "through the operator set");
    }
    // Out of scope and with a bundle that fails a load check: the scope is checked first.
    let mut d = wh.data.clone();
    set(&mut d, pih + 85, 4, 1);
    let mut bad = NormTables { bounds: vec![256; 32], ..norm_tables(&json(&read(Path::new(&std::env::var("WAI_JPEGAI_NORM").unwrap())))) };
    assert!(matches!(wai::jpegai_decode::decode(&d, &wh.weights, &bad), Err(wai::jpegai_decode::Refusal::Scope(_))));
    bad.bounds = wh.norm.bounds.clone();
    bad.gain_y.pop();
    assert_eq!(wai::jpegai_decode::decode(&wh.data, &wh.weights, &bad).err(), Some(wai::jpegai_decode::Refusal::Ops(wai_ops::Code::ModelMismatch)), "in scope, the bundle's load check");
}

/// Gate 5: under every kernel policy, the reference kernels and the fast ones at each
/// thread count, both planes' integers are the same, on the committed bundle and, where
/// it is supplied, the pinned one, with the test codestream's payloads whole and cut.
#[test]
fn every_kernel_policy_decodes_the_planes_alike() {
    use wai::int_kernels::{with_policy, Policy};
    static POLICIES: std::sync::LazyLock<Vec<Policy>> = std::sync::LazyLock::new(wai::int_kernels::equivalence_policies);
    let mut sets: Vec<(&str, jpegai_ops::ScalePath, Vec<u8>)> = vec![];
    let b = Bundle::synthetic();
    let (n, y, uv) = b.parsed();
    sets.push(("committed", jpegai_ops::load(&n, &y, &uv).unwrap(), read(&wai().join("wai-rs/tests/vectors/jpegai/test256.bin"))));
    if let Some(wh) = Whole::pinned("jpegai_ops::every_kernel_policy_decodes_the_planes_alike") {
        sets.push(("jpegai/v1", wai::jpegai_decode::load(&wh.weights, &wh.norm).unwrap(), wh.data));
    }
    let mut compared = 0;
    for (name, sp, data) in &sets {
        for p in [jpegai_ops::Plane::Luma, jpegai_ops::Plane::Chroma] {
            let (soz, lat, grfs) = segments(data, p);
            for (s, l) in [(soz.clone(), lat.clone()), (soz[..soz.len() / 2].to_vec(), lat.clone()), (soz.clone(), lat[..lat.len() / 2].to_vec())] {
                let want = with_policy(POLICIES[0], || jpegai_ops::plane(sp, p, &s, &l, &grfs));
                for pol in &POLICIES[1..] {
                    let got = with_policy(*pol, || jpegai_ops::plane(sp, p, &s, &l, &grfs));
                    assert!(got == want, "{name} {p:?}: {pol:?} differs from the interpreter alone");
                    compared += 1;
                }
            }
        }
    }
    eprintln!("{compared} plane decodes under the fast policies equal the reference policy's");
}