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

// wai-conformance: int-kernels int-kernels-amd
//! `wai.neural.int_mlicpp` through the integer operator set: its lowering `wmb1/1`
//! (`extensions/integer-operators.md` §13.4; each of the bundle's three `WHM1` networks by
//! `whm1/1`, §13.2: `h_s`, `entropy`, then `g_s` with the RGB8 step of §13.0) and the
//! certificate of every committed set (§8.5), as this crate loads them (`int_ops`,
//! `int_mlicpp::ops`).
//!
//! The lowered graphs and certificates are committed under `ops-conformance/certificates/`,
//! where `wai-verify` and the independent checker (`tools/wai_ops_check.py`) each reproduce
//! them byte for byte (`wai-ops.yml`). Here the lowering this crate runs, from the networks
//! it parses (`int_ops::whm1_net`), must write those same bytes, and the load a sink makes
//! must give the verdict the certificate records.

use std::path::PathBuf;

use serde_json::Value;
use wai::int_hyper_synth::IntModel;
use wai::int_mlicpp::{self as mlic, IntMlicModel};
use wai_ops::Profile;

const CAP: &str = "wai.neural.int_mlicpp";

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

fn read(rel: &str) -> Vec<u8> {
    std::fs::read(wai().join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}

fn json(rel: &str) -> Value {
    serde_json::from_slice(&read(rel)).unwrap_or_else(|e| panic!("{rel}: {e}"))
}

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

/// A `WMB1` bundle's four sections, in order: `h_s`, `entropy`, `g_s` and the tables.
fn sections(b: &[u8]) -> [Vec<u8>; 4] {
    assert_eq!(&b[..4], b"WMB1");
    let mut o = 4;
    let mut s: [Vec<u8>; 4] = Default::default();
    for x in &mut s {
        let n = u32::from_le_bytes(b[o..o + 4].try_into().unwrap()) as usize;
        *x = b[o + 4..o + 4 + n].to_vec();
        o += 4 + n;
    }
    assert_eq!(o, b.len(), "nothing follows the tables");
    s
}

/// A committed set of the capability, as `ops-conformance/certificates/index.json` lists
/// it: its name, its one member's path below `wai/`, its digest and its lowered graphs
/// (`h_s`, `entropy`, then `g_s`).
struct Indexed {
    name: String,
    bundle: String,
    digest: String,
    roles: [String; 3],
}

fn indexed() -> Vec<Indexed> {
    json("ops-conformance/certificates/index.json")["sets"]
        .as_array()
        .unwrap()
        .iter()
        .filter(|s| s["capability"] == CAP)
        .map(|s| {
            assert_eq!(s["lowering"], "wmb1/1");
            let roles = s["roles"].as_array().unwrap();
            assert_eq!(roles.len(), 3, "wmb1/1 lowers a bundle to three graphs");
            let members = s["members"].as_object().unwrap();
            assert_eq!(members.keys().collect::<Vec<_>>(), ["model.wmb"], "a one-file set");
            let role = |k: usize| roles[k].as_str().unwrap().to_string();
            Indexed {
                name: s["name"].as_str().unwrap().to_string(),
                bundle: members["model.wmb"].as_str().unwrap().to_string(),
                digest: s["set_sha256"].as_str().unwrap().to_string(),
                roles: [role(0), role(1), role(2)],
            }
        })
        .collect()
}

/// Gate 1 of the migration: every committed set of `wai.neural.int_mlicpp` lowers to its
/// committed graphs, from its bundle's bytes (`lower::wmb1`) and from the networks this
/// crate parses (`int_mlicpp::ops::lower`), whose conversion (`int_ops::whm1_net`) gives the
/// networks `wai-ops` parses from each section; the certificate beside them is those
/// graphs'; and loading each graph under the capability's profile gives the certificate's
/// verdict. The sets carry `check` ranges (§13.2), so they meet `acc64` by construction and
/// not `acc32`; the first `acc32` failure is `h_s`'s input range, at 47 bits for the
/// conformance vector's set (Appendix C's prediction), 43 for `natural`, the set trained on
/// openly licensed natural pictures (`int_mlicpp_natural/natural_set.json`), and 54 for the
/// synthetic one. The
/// decode's own load (`int_mlicpp::ops::load`) holds those graphs.
#[test]
fn every_committed_set_lowers_to_its_committed_graphs_and_verdict() {
    let sets = indexed();
    let names: Vec<&str> = sets.iter().map(|s| s.name.as_str()).collect();
    assert_eq!(names, ["int_mlicpp", "mlicpp", "natural"], "the committed int_mlicpp sets");
    for s in &sets {
        let name = &s.name;
        let bundle = read(&s.bundle);
        assert_eq!(wai_ops::generate::sets::set_digest(&[("model.wmb", bundle.clone())]), s.digest, "{name}: the set's digest");
        assert_eq!(sha256(&bundle), s.digest, "{name}: a one-file set's digest is its file's");
        // The tables are the stage operators' data, not lowered; the bundle must load.
        let m = IntMlicModel::from_bundle(&bundle).unwrap_or_else(|e| panic!("{name}: {e}"));
        let lowered = wai_ops::lower::wmb1(&bundle).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
        let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
        assert_eq!(roles, ["h_s", "entropy", "g_s"], "{name}");
        // The conversion of each network this crate parses is the one wai-ops parses from
        // its section.
        let secs = sections(&bundle);
        for (k, net) in [&m.h_s, &m.entropy, &m.g_s].into_iter().enumerate() {
            let conv = wai::int_ops::whm1_net(net).expect("a network parsed from bytes converts");
            assert!(conv == wai_ops::lower::parse_whm1(&secs[k]).unwrap(), "{name} {}: the converted network is the one wai-ops parses from its section", roles[k]);
            let again = wai::int_ops::whm1_net(&IntModel::from_bytes(&secs[k]).unwrap()).unwrap();
            assert!(conv == again, "{name} {}: the bundle's network is its section's", roles[k]);
        }
        let (h, e, g) = mlic::ops::lower(&m).unwrap_or_else(|r| panic!("{name}: {r}"));
        let cert = json(&format!("ops-conformance/certificates/{CAP}/{}.json", s.digest));
        assert_eq!(cert["set_sha256"].as_str(), Some(s.digest.as_str()), "{name}");
        assert_eq!(cert["capability"].as_str(), Some(CAP), "{name}");
        assert_eq!(cert["lowering"], "wmb1/1", "{name}");
        for (k, (role, ours)) in [("h_s", &h), ("entropy", &e), ("g_s", &g)].into_iter().enumerate() {
            let from_bytes = lowered.get(role).unwrap();
            // The verdict first: the load a sink makes.
            let loaded = wai::int_ops::load(from_bytes).unwrap_or_else(|r| panic!("{name} {role}: refused under acc64: {r:?}"));
            assert!(!loaded.meets(Profile::Acc32), "{name} {role}: a graph with check ranges inside it does not meet acc32");
            // Then the committed bytes: the graph, both ways, and the certificate's record of it.
            let graph = from_bytes.to_bytes();
            assert!(graph == read(&s.roles[k]), "{name}: the lowering of the bytes differs from the committed {}", s.roles[k]);
            assert!(ours.to_bytes() == graph, "{name}: the lowering of the parsed networks differs from the committed {role}");
            let c = &cert["graphs"][k];
            assert_eq!(c["name"].as_str(), Some(role), "{name}");
            assert_eq!(c["graph_sha256"].as_str(), Some(sha256(&graph).as_str()), "{name} {role}: the certificate is the graph's");
            // The input's check range is the lowering's X_0 for the network's first layer (§13.2).
            let x0 = &c["inputs"][0];
            let t = &from_bytes.tensors[0];
            assert_eq!((x0["policy"].as_str(), x0["lo"].as_str(), x0["hi"].as_str()), (Some("check"), Some(t.lo.to_string().as_str()), Some(t.hi.to_string().as_str())), "{name} {role}");
            assert!(t.hi > 0 && t.lo == -t.hi, "{name} {role}: a symmetric check range");
        }
        let v = &cert["verdict"];
        assert_eq!((v["acc64"].as_bool(), v["acc32"].as_bool()), (Some(true), Some(false)), "{name}: the certificate's verdict");
        let bits = match name.as_str() {
            "int_mlicpp" => "47",
            "natural" => "43",
            _ => "54",
        };
        let first = &v["acc32_first_failure"];
        assert_eq!(
            (first["graph"].as_str(), first["where"].as_str(), first["index"].as_str(), first["key"].as_str(), first["bits"].as_str()),
            (Some("h_s"), Some("input"), Some("0"), Some("range"), Some(bits)),
            "{name}: the first acc32 failure"
        );
        // The decode's own load: the three networks, in that order.
        let nets = mlic::ops::load(&m).unwrap_or_else(|r| panic!("{name}: {r}"));
        assert!(
            nets.h_s.graph() == lowered.get("h_s").unwrap() && nets.entropy.graph() == lowered.get("entropy").unwrap() && nets.g_s.graph() == lowered.get("g_s").unwrap(),
            "{name}: the decode loads the committed graphs"
        );
    }
}

fn hex(h: u64) -> String {
    format!("{h:016x}")
}

/// FNV-1a over the `y` stream's buckets (u16 LE) and symbols (i32 LE), in coding order.
fn trace_hashes(t: &mlic::Trace) -> (String, String) {
    let b: Vec<u8> = t.buckets.iter().flat_map(|&b| u16::try_from(b).unwrap().to_le_bytes()).collect();
    let s: Vec<u8> = t.symbols.iter().flat_map(|s| s.to_le_bytes()).collect();
    (hex(wai::int_hyper_synth::fnv1a(&b)), hex(wai::int_hyper_synth::fnv1a(&s)))
}

/// The conformance vector's directory, below `wai/`.
const DIR: &str = "wai-rs/tests/vectors/int_mlicpp";

/// The committed vector's bundle, parsed.
fn model() -> IntMlicModel {
    IntMlicModel::from_bundle(&read(&format!("{DIR}/model.wmb"))).unwrap()
}

/// The payload of a committed envelope of the vector.
fn payload_of(env: &str) -> Vec<u8> {
    wai::container::Wai::from_bytes(&read(&format!("{DIR}/{env}"))).unwrap().payload
}

/// Gate 2 of the migration: golden equality. The committed vector of
/// `wai.neural.int_mlicpp` (`tests/vectors/int_mlicpp`, its form-less envelope and its
/// set-form twin, which carry the same payload) decodes through the lowering
/// (`int_mlicpp::ops`) to the image the direct decode gives (`int_mlicpp::legacy`), with
/// the same bucket and symbol for every element, and to every digest `expected.json`
/// commits for it: the image's FNV-1a and SHA-256, the buckets' and the symbols' FNV-1a,
/// and the count of `y` symbols.
#[test]
fn the_committed_vector_decodes_through_the_lowering_to_its_golden() {
    use sha2::Digest;
    let exp = json(&format!("{DIR}/expected.json"));
    let m = model();
    let (p, p_set) = (payload_of("stream.wai"), payload_of("stream_set.wai"));
    assert_eq!(sha256(&p), exp["payload_sha256"].as_str().unwrap());
    assert!(p == p_set, "the set-form twin carries the same payload");
    let (t_ops, img_ops) = mlic::ops::decode_traced(&m, &p);
    let (t_dir, img_dir) = mlic::legacy::decode_traced(&m, &p);
    let (img_ops, img_dir) = (img_ops.expect("the decode through the lowering"), img_dir.expect("the direct decode"));
    assert!(t_ops == t_dir, "every bucket and symbol, through the lowering and directly");
    assert!(img_ops == img_dir, "the decode through the lowering differs from the direct decode");
    assert_eq!((img_ops.width, img_ops.height), (exp["width"].as_u64().unwrap() as usize, exp["height"].as_u64().unwrap() as usize));
    assert_eq!(hex(wai::int_hyper_synth::fnv1a(&img_ops.rgb)), exp["rgb_fnv1a"].as_str().unwrap(), "the export's FNV-1a");
    let sha: String = sha2::Sha256::digest(&img_ops.rgb).iter().map(|x| format!("{x:02x}")).collect();
    assert_eq!(sha, exp["rgb_sha256"].as_str().unwrap(), "the image's SHA-256");
    let (b, s) = trace_hashes(&t_ops);
    assert_eq!((b.as_str(), s.as_str()), (exp["bucket_fnv1a"].as_str().unwrap(), exp["symbol_fnv1a"].as_str().unwrap()));
    assert_eq!(t_ops.buckets.len() as u64, exp["y_symbols"].as_u64().unwrap());
}

/// The code integer-payloads §8 gives a refusal of the decode, as `int_payload::decode_int`
/// reports it: the payload's, or the bundle's against it.
fn code_of(e: &mlic::IntMlicError) -> &'static str {
    match e {
        mlic::IntMlicError::Payload(_) => "payload_malformed",
        mlic::IntMlicError::Model(_) => "model_mismatch",
    }
}

/// A `WIM1` payload's fields, owned, so a fixture can change any of them: `dims` are
/// `Cz, Zh, Zw, Cy, Yh, Yw`.
#[derive(Clone)]
struct Wim1 {
    h: u32,
    w: u32,
    dims: [u16; 6],
    k: u16,
    reserved: u16,
    z: Vec<u8>,
    y: Vec<u8>,
}

impl Wim1 {
    fn committed() -> Wim1 {
        let p = payload_of("stream.wai");
        let u16at = |o: usize| u16::from_le_bytes([p[o], p[o + 1]]);
        let u32at = |o: usize| u32::from_le_bytes(p[o..o + 4].try_into().unwrap());
        let (lz, ly) = (u32at(28) as usize, u32at(32) as usize);
        Wim1 {
            h: u32at(4), w: u32at(8), dims: [12, 14, 16, 18, 20, 22].map(u16at), k: u16at(24), reserved: u16at(26),
            z: p[36..36 + lz].to_vec(), y: p[36 + lz..36 + lz + ly].to_vec(),
        }
    }
    fn bytes(&self) -> Vec<u8> {
        let mut p = b"WIM1".to_vec();
        p.extend(self.h.to_le_bytes());
        p.extend(self.w.to_le_bytes());
        for d in &self.dims {
            p.extend(d.to_le_bytes());
        }
        p.extend(self.k.to_le_bytes());
        p.extend(self.reserved.to_le_bytes());
        p.extend((self.z.len() as u32).to_le_bytes());
        p.extend((self.y.len() as u32).to_le_bytes());
        p.extend(&self.z);
        p.extend(&self.y);
        p
    }
}

/// The committed payload, changed by `f`.
fn wim1(f: impl FnOnce(&mut Wim1)) -> Vec<u8> {
    let mut o = Wim1::committed();
    f(&mut o);
    o.bytes()
}

/// The committed bundle, parsed, changed by `f`.
fn bundle(f: impl FnOnce(&mut IntMlicModel)) -> IntMlicModel {
    let mut m = model();
    f(&mut m);
    m
}

/// Both decodes of a payload against a bundle, with their traces: through the lowering,
/// then directly. A refusal is its code.
type Decoded = (mlic::Trace, Result<Vec<u8>, &'static str>);
fn both(m: &IntMlicModel, p: &[u8]) -> (Decoded, Decoded) {
    let f = |(t, r): (mlic::Trace, Result<mlic::IntMlicImage, mlic::IntMlicError>)| (t, r.map(|i| i.rgb).map_err(|e| code_of(&e)));
    (f(mlic::ops::decode_traced(m, p)), f(mlic::legacy::decode_traced(m, p)))
}

/// A layer's biases, in place.
fn bias(l: &mut wai::int_hyper_synth::IntLayer) -> &mut Vec<i64> {
    use wai::int_hyper_synth::IntLayer;
    match l {
        IntLayer::Conv { bias, .. } | IntLayer::ConvT { bias, .. } => bias,
        _ => panic!("not a convolution"),
    }
}

/// Gate 3 of the migration: refusal parity. A fixture for each refusal the direct decode
/// makes (integer-payloads §14, in its order: the header, the bundle against the header
/// and the networks' geometry, then each stream and each network as it runs), each made
/// from the committed bundle and payload by one change. Each is refused by both decodes
/// with the same code (integer-payloads §8: `payload_malformed` for the payload and the
/// decode, `model_mismatch` for a bundle that does not fit the payload), as far as the
/// same `y` symbol. The committed bundle and payload decode alike. Then the sink's codes
/// for two committed cases of `int_payload`: a payload that is not one, and a bundle that
/// is not one.
#[test]
fn every_refusal_of_the_direct_decode_is_refused_alike_through_the_lowering() {
    use wai::int_hyper_synth::IntLayer;
    let (ops, direct) = both(&model(), &Wim1::committed().bytes());
    assert!(ops.1.is_ok() && ops == direct, "the fixtures' base decodes alike");
    let same = model;
    let as_is = || Wim1::committed().bytes();
    // A stream moved: z cut to `lz` bytes, y to `ly` (each from its start).
    let cut = |lz: usize, ly: usize| wim1(|o| { o.z.truncate(lz); o.y.truncate(ly); });
    let (lz, ly) = { let o = Wim1::committed(); (o.z.len(), o.y.len()) };
    // A synthesis that spreads the latent over a plane past the cap and brings it back:
    // the right output, through a buffer over 2^26 elements.
    let wide = || bundle(|m| {
        let s = 1usize << 20;
        let unit = |out: usize| (vec![1i64; 8 * out], vec![1i64; out], vec![0i64; out]);
        let (weight, mult, b) = unit(8);
        m.g_s.layers.insert(0, IntLayer::ConvT { out_ch: 8, k: 1, stride: s, pad: 0, opad: 0, weight, mult, bias: b });
        let (weight, mult, b) = unit(8);
        m.g_s.layers.insert(1, IntLayer::Conv { out_ch: 8, k: 1, stride: s, pad: 0, weight, mult, bias: b });
    });
    let fixtures: Vec<(&str, &str, IntMlicModel, Vec<u8>)> = vec![
        ("a payload shorter than its header", "payload_malformed", same(), as_is()[..35].to_vec()),
        ("not WIM1", "payload_malformed", same(), { let mut p = as_is(); p[3] = b'2'; p }),
        ("a reserved field that is not 0", "payload_malformed", same(), wim1(|o| o.reserved = 1)),
        ("a dimension of 0", "payload_malformed", same(), wim1(|o| o.dims[3] = 0)),
        ("Cy not a multiple of K", "payload_malformed", same(), wim1(|o| o.k = 3)),
        ("an image past 2^26 elements", "payload_malformed", same(), wim1(|o| (o.h, o.w) = (1 << 14, 1 << 14))),
        ("lengths that are not the payload's", "payload_malformed", same(), [as_is(), vec![0]].concat()),
        ("more hyper-latent channels than z tables", "model_mismatch", bundle(|m| {
            let t = &mut m.tables;
            (t.z_cdfs.pop(), t.z_cdf_lengths.pop(), t.z_offsets.pop(), t.z_median_q.pop());
        }), as_is()),
        ("an h_s that does not give the declared latent", "model_mismatch", same(), wim1(|o| o.dims[4] = 9)),
        ("an entropy network that does not give 2·Cy channels", "model_mismatch", bundle(|m| {
            let to4 = IntLayer::Conv { out_ch: 4, k: 1, stride: 1, pad: 0, weight: vec![1; 16 * 4], mult: vec![1; 4], bias: vec![0; 4] };
            m.entropy.layers.push(to4);
        }), as_is()),
        ("a g_s that does not give the declared image", "model_mismatch", same(), wim1(|o| o.h = 65)),
        ("a synthesis through a buffer past 2^26 elements", "model_mismatch", wide(), as_is()),
        ("a z stream too short to be one", "payload_malformed", same(), cut(4, ly)),
        ("a z value that leaves i64", "payload_malformed", bundle(|m| m.tables.z_median_q = vec![i64::MAX; 4]), as_is()),
        ("a z stream that runs out", "payload_malformed", same(), cut(lz - 4, ly)),
        ("hyper-latents past h_s's range", "payload_malformed", bundle(|m| m.tables.z_median_q = vec![1 << 46; 4]), as_is()),
        ("a y stream too short to be one", "payload_malformed", same(), cut(lz, 4)),
        ("hyper past the entropy network's range", "payload_malformed", bundle(|m| bias(&mut m.h_s.layers[2]).iter_mut().for_each(|b| *b = 1 << 50)), as_is()),
        ("means near i64's end", "payload_malformed", bundle(|m| bias(&mut m.entropy.layers[2])[8..].iter_mut().for_each(|b| *b = i64::MAX - (1 << 20))), as_is()),
        ("a y stream that runs out", "payload_malformed", same(), cut(lz, ly - 4)),
        ("latents past g_s's range", "payload_malformed", bundle(|m| bias(&mut m.g_s.layers[0])[0] = i64::MAX - (1 << 20)), as_is()),
    ];
    for (what, want, m, p) in &fixtures {
        let (ops, direct) = both(m, p);
        assert_eq!((ops.1.as_ref().err(), direct.1.as_ref().err()), (Some(want), Some(want)), "{what}: through the lowering, then directly");
        assert!(ops.0 == direct.0, "{what}: the two decodes read different symbols before the refusal");
    }
    // The sink's codes (`int_payload/tests.rs`' committed cases).
    let cap = "wai.neural.int_mlicpp";
    let bundle_bytes = read(&format!("{DIR}/model.wmb"));
    let sink = |p: &[u8], set: &[u8]| {
        let params = wai::int_payload::resolve(cap, None, &[wai::pin::HeldBytes::File(set)], None).expect("an unpinned set resolves");
        wai::int_payload::decode_int(cap, p, &params).map(|_| ()).map_err(|e| e.code())
    };
    assert_eq!(sink(b"WIM1", &bundle_bytes), Err("payload_malformed"));
    assert_eq!(sink(&as_is(), b"WHM1"), Err("model_mismatch"));
    assert_eq!(sink(&as_is(), &bundle_bytes), Ok(()));
}

/// Gate 3, the end rule `final_only` (integer-payloads §14; integer-operators E1): a
/// stream that runs out is refused once, after its last symbol. The `y` stream cut short by
/// every multiple of 4 bytes that leaves it a stream runs out part of the way through, and
/// both decodes still decode every
/// symbol after that point, from the state as it stands, feed each to the passes after it
/// (the entropy network's input, in which every element of `y` not yet decoded is 0), and
/// refuse at the end, with the same buckets and symbols. Where a symbol decoded after the
/// run-out gives a refusal of its own before the end (a value that leaves `i64`, or one
/// outside the entropy network's range), both give that one, at the same symbol. The
/// `z` stream cut short is refused alike too.
#[test]
fn a_stream_that_runs_out_is_refused_at_its_end_alike() {
    let m = model();
    let full = { let (t, r) = mlic::ops::decode_traced(&m, &Wim1::committed().bytes()); r.unwrap(); t.buckets.len() };
    let msg = |p: &[u8]| match mlic::ops::decode(&m, p) {
        Err(mlic::IntMlicError::Payload(e)) => e,
        other => panic!("{:?}", other.map(|_| ())),
    };
    let (mut at_end, mut before) = (0, 0);
    for cut in (4..=Wim1::committed().y.len() - 8).step_by(4) {
        let p = wim1(|o| { let n = o.y.len() - cut; o.y.truncate(n); });
        let (ops, direct) = both(&m, &p);
        assert!(ops == direct, "y cut by {cut}: the two decodes differ");
        assert_eq!(ops.1, Err("payload_malformed"), "y cut by {cut}");
        if ops.0.buckets.len() == full {
            assert_eq!(msg(&p), "y stream ended before its last symbol", "y cut by {cut}: every symbol decoded, then refused at the end");
            at_end += 1;
        } else {
            before += 1;
        }
    }
    eprintln!("y cut short: {at_end} decoded to the last symbol and refused at the end, {before} refused before it");
    assert!(at_end > 0, "a stream that runs out is decoded to its last symbol before it is refused");
    for cut in [4, 8] {
        let p = wim1(|o| { let n = o.z.len() - cut; o.z.truncate(n); });
        let (ops, direct) = both(&m, &p);
        assert!(ops == direct && ops.1 == Err("payload_malformed"), "z cut by {cut}");
    }
}

/// Gate 3, at the edge of every `check` range: for each layer of the committed `h_s`,
/// `entropy` and `g_s` that reads one (§13.2: a layer of tag 0, 1, 2 or 4; a ReLU's input
/// has none), a network of that layer alone (with the bundle's header; a leaky ReLU followed
/// by a zero 1×1 transposed convolution, since a network with no layer of tag 0, 1 or 2 does
/// not lower) is lowered, and its input's `check` range is `[−X_ℓ, X_ℓ]`, the range the full
/// graph gives tensor `ℓ`. An input whose largest magnitude is `X_ℓ` runs, through the
/// lowering and directly, to the same output; one at `X_ℓ + 1` is refused by both: by the
/// graph (`field_range`, at input 0) and by the direct decode's check of that layer
/// (`run_net_checked`, a value out of range at layer 0). So the ranges refuse exactly what
/// the direct decode refuses, at the same layer.
#[test]
fn every_check_range_refuses_exactly_what_the_direct_decode_refuses() {
    use wai::int_hyper_synth::{legacy::run_net_checked, IntLayer};
    use wai_ops::{Code, Site, Tensor};
    let secs = sections(&read(&format!("{DIR}/model.wmb")));
    let mut seen = 0;
    for (k, (net, rgb8)) in [("h_s", false), ("entropy", false), ("g_s", true)].into_iter().enumerate() {
        let set = &secs[k];
        let full = wai_ops::lower::whm1(set, rgb8).unwrap();
        let n = wai_ops::lower::parse_whm1(set).unwrap().layers.len();
        for l in 0..n {
            let mut one = IntModel::from_bytes(set).unwrap();
            let layer = one.layers.swap_remove(l);
            if matches!(layer, IntLayer::Relu) {
                continue;
            }
            let leaky = matches!(layer, IntLayer::Leaky { .. });
            one.layers = vec![layer];
            let c = full.tensors[l].channels;
            if leaky {
                one.layers.push(IntLayer::ConvT { out_ch: 1, k: 1, stride: 1, pad: 0, opad: 0, weight: vec![0; c as usize], mult: vec![1], bias: vec![0] });
            }
            let g = wai_ops::lower::whm1_net(&wai::int_ops::whm1_net(&one).unwrap(), false).unwrap();
            let (lo, hi) = (g.tensors[0].lo, g.tensors[0].hi);
            assert_eq!((g.tensors[0].channels, lo, hi), (c, full.tensors[l].lo, full.tensors[l].hi), "{net} layer {l}: the range of tensor {l}");
            assert!(hi > 0 && lo == -hi, "{net} layer {l}: a symmetric check range [{lo}, {hi}]");
            let loaded = wai::int_ops::load(&g).unwrap();
            let x = |v: i64| {
                let mut d = vec![0i64; c as usize * 4];
                d[c as usize * 2 + 1] = v;
                d
            };
            for v in [hi, -hi] {
                let ops = wai::int_ops::run1(&loaded, Tensor::new(c, 2, 2, x(v)), &[]).unwrap_or_else(|r| panic!("{net} layer {l}: {v} runs: {r:?}"));
                let (direct, ..) = run_net_checked(&one, &x(v), c as usize, 2, 2).unwrap_or_else(|r| panic!("{net} layer {l}: {v} runs directly: {r}"));
                assert!(ops.data == direct, "{net} layer {l}: at {v} the two differ");
            }
            for v in [hi + 1, -hi - 1] {
                let r = wai::int_ops::run1(&loaded, Tensor::new(c, 2, 2, x(v)), &[]).expect_err("one past the range");
                assert_eq!((r.code, r.at), (Code::FieldRange, Site::Input(0)), "{net} layer {l}: {v}");
                let d = run_net_checked(&one, &x(v), c as usize, 2, 2).expect_err("one past the range, directly");
                assert!(d.out_of_range && d.layer == 0, "{net} layer {l}: {v}: {d}");
            }
            seen += 1;
        }
    }
    assert_eq!(seen, 10, "every layer of the committed h_s, entropy and g_s that reads a check range");
}

/// Gate 5 of the migration, the whole-decoder half: the decode through the lowering under
/// the reference kernels (the interpreter computes every operator) and under each fast
/// policy (the convolutions, the transposed convolutions and the inverse GDNs on this
/// crate's fast kernels, serial and threaded) reads the same buckets and symbols, and gives
/// the same image or the same refusal, on the committed vector (both envelopes), on a `y`
/// stream that runs out, and on bundles whose hyper-latent channel 0 lies about the edge
/// of `h_s`'s input `check` range (its median at `±X_0` and one past it). The per-operator
/// half is `int_ops`' equivalence tests, one of which counts the fast kernels' runs on this
/// bundle.
#[test]
fn every_kernel_policy_decodes_through_the_lowering_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 x0 = wai_ops::lower::whm1(&sections(&read(&format!("{DIR}/model.wmb")))[0], false).unwrap().tensors[0].hi;
    let base = Wim1::committed();
    let z0 = base.z.clone();
    let mut cases: Vec<(String, IntMlicModel, Vec<u8>)> = vec![
        ("stream.wai".into(), model(), payload_of("stream.wai")),
        ("stream_set.wai".into(), model(), payload_of("stream_set.wai")),
        ("y cut by 8".into(), model(), wim1(|o| { let n = o.y.len() - 8; o.y.truncate(n); })),
    ];
    // Channel 0's median moved to ±X_0 or one past it: its elements, the median plus each
    // symbol at 2^12, lie about the edge of h_s's input check range.
    let syms = {
        let t = model();
        let mut d = wai::rans::RansDecoder::new(&z0).unwrap();
        let per = (base.dims[1] * base.dims[2]) as usize;
        (0..per).map(|_| d.decode(&t.tables.z_cdfs[0], t.tables.z_cdf_lengths[0], t.tables.z_offsets[0])).collect::<Vec<i32>>()
    };
    let (lo, hi) = (*syms.iter().min().unwrap() as i64, *syms.iter().max().unwrap() as i64);
    for (what, med) in [("channel 0 at X_0", x0 - (hi << 12)), ("at X_0 + 1", x0 + 1 - (hi << 12)), ("at -X_0", -x0 - (lo << 12)), ("at -X_0 - 1", -x0 - 1 - (lo << 12))] {
        cases.push((what.into(), bundle(move |m| m.tables.z_median_q[0] = med), base.bytes()));
    }
    let mut refused = 0;
    for (k, (what, m, p)) in cases.iter().enumerate() {
        let want = with_policy(POLICIES[0], || mlic::ops::decode_traced(m, p));
        assert!(k >= 2 || want.1.is_ok(), "{what}: the committed vector decodes");
        refused += want.1.is_err() as u32;
        for pol in &POLICIES[1..] {
            let got = with_policy(*pol, || mlic::ops::decode_traced(m, p));
            assert!(got == want, "{what}: {pol:?} differs from the interpreter alone");
        }
    }
    assert!(refused >= 3, "the cut stream and the two medians one past the edge are refused: {refused}");
}

/// Every committed `WIM1` payload, alone or inside another file (an envelope, an HTTP
/// conformance object): each occurrence of the magic in the corpora that hold integer
/// payloads whose bytes from there parse as a `WIM1` of exactly the length its header's
/// stream lengths give. Returns `(where, payload)` pairs.
fn committed_wim1_payloads() -> Vec<(String, Vec<u8>)> {
    fn files(d: &std::path::Path, out: &mut Vec<PathBuf>) {
        let Ok(rd) = std::fs::read_dir(d) else { return };
        let mut es: Vec<_> = rd.filter_map(|e| e.ok()).map(|e| e.path()).collect();
        es.sort();
        for p in es {
            if p.is_dir() { files(&p, out) } else { out.push(p) }
        }
    }
    let mut all = Vec::new();
    for d in ["int-payload-conformance", "staged-conformance", "staged-measure-conformance", "http-conformance", "pin-conformance", "signer-conformance", "wai-rs/tests/vectors"] {
        files(&wai().join(d), &mut all);
    }
    let u32_at = |b: &[u8], i: usize| b.get(i..i + 4).map(|s| u32::from_le_bytes(s.try_into().unwrap()) as usize);
    let mut v = Vec::new();
    for f in all {
        let b = std::fs::read(&f).unwrap();
        for i in (0..b.len().saturating_sub(3)).filter(|&i| &b[i..i + 4] == b"WIM1") {
            let (Some(lz), Some(ly)) = (u32_at(&b, i + 28), u32_at(&b, i + 32)) else { continue };
            let Some(end) = (i + 36).checked_add(lz).and_then(|e| e.checked_add(ly)) else { continue };
            if let Some(p) = b.get(i..end).filter(|p| mlic::parse_header(p).is_ok()) {
                v.push((format!("{}@{i}", f.strip_prefix(wai()).unwrap().display()), p.to_vec()));
            }
        }
    }
    v
}

/// The scan for the load rules the switch brings (integer-operators §15.3 item 3, and the
/// rules `int_mlicpp`'s networks now meet before either stream): no committed bundle or
/// payload is affected. Every committed bundle (the sets indexed with their certificates,
/// and the HTTP corpus's copy of the vector's) has, in each network, a layer
/// of tag 0, 1 or 2, header shifts under 128, no bias of magnitude `2^63 − 1` or more,
/// layers that fit the channels before them and no leaky shift of 64 or more, and its
/// three networks lower and load under `acc64`. Every committed `WIM1` payload, wherever it
/// is carried, is planned through the lowering as directly, against the vector's bundle,
/// then decodes through the lowering to the bytes the direct decode gives, with the same
/// buckets and symbols, or is refused by both with the same code. (The other file that
/// holds the magic is code: the browser crate's recorded build.)
#[test]
fn the_new_load_rules_affect_no_committed_set_or_payload() {
    use wai_ops::lower::NetLayer as N;
    let mut bundles: Vec<(String, Vec<u8>)> = indexed().into_iter().map(|s| (s.name.clone(), read(&s.bundle))).collect();
    let copy = format!("http-conformance/sets/{}", indexed()[0].digest);
    bundles.push((copy.clone(), read(&copy)));
    assert_eq!(sha256(&bundles[bundles.len() - 1].1), indexed()[0].digest, "the HTTP corpus's copy is the vector's bundle");
    for (name, b) in &bundles {
        let secs = sections(b);
        for (k, (role, rgb8)) in [("h_s", false), ("entropy", false), ("g_s", true)].into_iter().enumerate() {
            let net = wai_ops::lower::parse_whm1(&secs[k]).unwrap();
            assert!(net.layers.iter().any(|l| matches!(l, N::ConvT { .. } | N::Conv { .. } | N::Igdn { .. })), "{name} {role}: no layer of tag 0, 1 or 2");
            assert!(net.shift < 128 && net.gf < 128, "{name} {role}: a header shift of 128 or more");
            let mut c: Option<u32> = None;
            for l in &net.layers {
                match l {
                    N::ConvT { inc, out, bias, .. } | N::Conv { inc, out, bias, .. } => {
                        assert!(bias.iter().all(|b| b.unsigned_abs() < i64::MAX as u64), "{name} {role}: a bias of magnitude 2^63 - 1 or more");
                        assert!(c.is_none_or(|c| c == *inc), "{name} {role}: a layer that does not fit the channels before it");
                        c = Some(*out);
                    }
                    N::Igdn { c: cc, .. } => assert!(c.is_none_or(|c| c == *cc), "{name} {role}: an inverse GDN that does not fit"),
                    N::Leaky { slope_frac, .. } => assert!(*slope_frac < 64, "{name} {role}: a leaky shift of 64 or more"),
                    N::Relu => {}
                }
            }
            let g = wai_ops::lower::whm1(&secs[k], rgb8).unwrap_or_else(|r| panic!("{name} {role}: does not lower: {r:?}"));
            assert!(wai::int_ops::load(&g).is_ok(), "{name} {role}: does not load under acc64");
        }
        mlic::ops::load(&IntMlicModel::from_bundle(b).unwrap()).unwrap_or_else(|e| panic!("{name}: {e}"));
    }
    let m = model();
    let nets = mlic::ops::load(&m).unwrap();
    let (mut decoded, mut refused) = (0, 0);
    let payloads = committed_wim1_payloads();
    for (at, p) in &payloads {
        if let Ok(hd) = mlic::ops::check(&m, p) {
            mlic::ops::fit(&nets, &hd).unwrap_or_else(|e| panic!("{at}: the bundle fits the payload directly, and not through the lowering: {e}"));
        }
        let (ops, direct) = both(&m, p);
        assert!(ops == direct, "{at}: the decode through the lowering differs from the direct decode");
        if ops.1.is_ok() { decoded += 1 } else { refused += 1 }
    }
    let http = payloads.iter().filter(|(at, _)| at.starts_with("http-conformance")).count();
    assert!(decoded >= 4 && http >= 2, "the committed WIM1 payloads, at least those of this revision: {decoded} decoded, {refused} refused, {http} in the HTTP corpus");
}