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

// wai-conformance: int-kernels int-kernels-amd
//! `wai.world.model`'s integer renderer through the integer operator set: its lowering
//! `world/1` (`extensions/integer-operators.md` §13.5, informative: the transition `T`
//! with its `A` action planes, then the synthesis `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_world::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 models it parses (`int_ops::whm1_net`), must write those same bytes, and the load a
//! sink makes must give the verdict the certificate records.
#![cfg(all(feature = "world_int", feature = "legacy_int"))]

use std::path::PathBuf;

use serde_json::Value;
use wai::int_hyper_synth::IntModel;
use wai::int_world;
use wai_ops::Profile;

const CAP: &str = "wai.world.model";

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 committed `wai.world.model` set, as `ops-conformance/certificates/index.json` lists
/// it: its name, whether it is synthetic, the latent's channel count `C`, its members
/// (name, path below `wai/`), its digest and its lowered graphs (`T`, then `g_s`).
struct Indexed {
    name: String,
    synthetic: bool,
    c: u32,
    members: Vec<(String, String)>,
    digest: String,
    roles: [String; 2],
}

impl Indexed {
    fn member(&self, name: &str) -> Vec<u8> {
        read(&self.members.iter().find(|(k, _)| k == name).unwrap_or_else(|| panic!("{}: no {name}", self.name)).1)
    }
}

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"], "world/1");
            let roles = s["roles"].as_array().unwrap();
            assert_eq!(roles.len(), 2, "world/1 lowers a set to two graphs, T and g_s");
            Indexed {
                name: s["name"].as_str().unwrap().to_string(),
                synthetic: s["synthetic"].as_bool().unwrap(),
                c: s["c"].as_str().unwrap().parse().unwrap(),
                members: s["members"].as_object().unwrap().iter().map(|(k, v)| (k.clone(), v.as_str().unwrap().to_string())).collect(),
                digest: s["set_sha256"].as_str().unwrap().to_string(),
                roles: [roles[0].as_str().unwrap().to_string(), roles[1].as_str().unwrap().to_string()],
            }
        })
        .collect()
}

/// Gate 1 of the migration: every committed `wai.world.model` set lowers to its committed
/// graphs, from its members' bytes (`lower::world`) and from the models this crate parses
/// (`int_world::ops::lower`); the certificate beside them is those graphs'; and loading
/// each graph under `acc64` gives the certificate's verdict, as `int_world::ops::load`
/// holds them.
///
/// - `A`, the action planes, is `T`'s input channel count less `C` (§13.5): 3 for the one
///   agent set, 4 for the two agent set. `T` has one scalar per plane, each over all of
///   `i64`, and its input `z` is `check` over all of `i64` (§13.0's driver contract).
/// - `g_s`'s input range is the lowering's `X_0` for its first layer (§13.2).
/// - The sets meet `acc64` and not `acc32`: the first `acc32` failure is `T`'s input range,
///   at 64 bits, as Appendix C of the design predicted.
#[test]
fn every_committed_set_lowers_to_its_committed_graphs_and_verdict() {
    let sets = indexed();
    let names: Vec<(&str, bool, u32)> = sets.iter().map(|s| (s.name.as_str(), s.synthetic, s.c)).collect();
    assert_eq!(names, [("world", false, 24), ("world2", false, 24), ("world", true, 2)], "the committed world sets");
    for s in &sets {
        let name = format!("{}{}", s.name, if s.synthetic { " (synthetic)" } else { "" });
        let (t, g) = (s.member("T.whm"), s.member("g_s.whm"));
        assert_eq!(wai_ops::generate::sets::set_digest(&[("T.whm", t.clone()), ("g_s.whm", g.clone())]), s.digest, "{name}: the set's digest");
        let lowered = wai_ops::lower::world(&t, &g, s.c).unwrap_or_else(|r| panic!("{name}: does not lower: {r:?}"));
        let roles: Vec<&str> = lowered.roles.iter().map(|(r, _)| *r).collect();
        assert_eq!(roles, ["T", "g_s"], "{name}");
        // The models this crate parses convert to the networks wai-ops parses from the bytes.
        let (tm, gm) = (IntModel::from_bytes(&t).unwrap(), IntModel::from_bytes(&g).unwrap());
        for (m, b) in [(&tm, &t), (&gm, &g)] {
            let net = wai::int_ops::whm1_net(m).expect("a model parsed from bytes converts");
            assert!(net == wai_ops::lower::parse_whm1(b).unwrap(), "{name}: the converted model is the one wai-ops parses from the bytes");
        }
        let (ot, og) = int_world::ops::lower(&tm, &gm, s.c as usize).unwrap_or_else(|r| panic!("{name}: {r}"));
        let world = int_world::ops::load(&tm, &gm, s.c as usize).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"], "world/1", "{name}");
        for (k, (role, ours, held)) in [("T", &ot, world.t()), ("g_s", &og, world.g_s())].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, as the renderer holds it, 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 models differs from the committed {role}");
            assert!(held.graph().to_bytes() == graph, "{name}: the renderer holds another {role} than the committed one");
            assert!(!held.meets(Profile::Acc32), "{name} {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");
            let x0 = &c["inputs"][0];
            let t0 = &from_bytes.tensors[0];
            assert_eq!((x0["policy"].as_str(), x0["lo"].as_str(), x0["hi"].as_str()), (Some("check"), Some(t0.lo.to_string().as_str()), Some(t0.hi.to_string().as_str())), "{name} {role}");
        }
        // T: z over all of i64 with C channels, and A scalars over all of i64, A being T's
        // input channels less C.
        let tg = lowered.get("T").unwrap();
        let cin = match &wai_ops::lower::parse_whm1(&t).unwrap().layers[0] {
            wai_ops::lower::NetLayer::Conv { inc, .. } | wai_ops::lower::NetLayer::ConvT { inc, .. } => *inc,
            l => panic!("{name}: T's first layer {l:?}"),
        };
        let a = (cin - s.c) as usize;
        assert_eq!((tg.tensors[0].channels, tg.tensors[0].lo, tg.tensors[0].hi), (s.c, i64::MIN, i64::MAX), "{name}: T's input z");
        assert_eq!(tg.scalars.len(), a, "{name}: one scalar per action plane");
        assert!(tg.scalars.iter().all(|r| (r.lo, r.hi) == (i64::MIN, i64::MAX)), "{name}: every plane value over all of i64");
        assert_eq!((world.channels(), world.planes()), (s.c as usize, a), "{name}: the renderer's C and A");
        if !s.synthetic {
            assert_eq!(a, if s.name == "world2" { 4 } else { 3 }, "{name}: A");
        }
        // g_s: a symmetric check range on z.
        let g0 = &lowered.get("g_s").unwrap().tensors[0];
        assert!(g0.channels == s.c && g0.hi > 0 && g0.lo == -g0.hi, "{name}: g_s reads z of C channels within 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 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("T"), Some("input"), Some("0"), Some("range"), Some("64")),
            "{name}: the first acc32 failure Appendix C predicts"
        );
    }
}

/// A committed world vector: its two networks, `C`, the latent's side, `z_0`, its JSON
/// (the export's FNV-1a frame and clip hashes) and its SHA-256 sidecar entry.
struct Vector {
    name: &'static str,
    t: IntModel,
    g_s: IntModel,
    c: usize,
    lat: usize,
    z0: Vec<i64>,
    j: Value,
    sha: Value,
}

/// The two committed vectors: one agent (`A = 3`) and two agents (`A = 4`).
fn vectors() -> [Vector; 2] {
    let v = |name: &'static str, dir: &str, js: &str, t: &str, g: &str| {
        let j = json(&format!("wai-rs/tests/vectors/{dir}/{js}"));
        let sha = json(&format!("wai-rs/tests/vectors/{dir}/sha256.json"))[name].clone();
        Vector {
            name,
            t: IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{t}"))).unwrap(),
            g_s: IntModel::from_bytes(&read(&format!("wai-rs/tests/vectors/{dir}/{g}"))).unwrap(),
            c: j["c"].as_u64().unwrap() as usize,
            lat: j["lat"].as_u64().unwrap() as usize,
            z0: j["z0"].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect(),
            j,
            sha,
        }
    };
    [v("world", "int_world", "world.json", "world_T.bin", "world_gs.bin"), v("world2", "int_world2", "world2.json", "world2_T.bin", "world2_gs.bin")]
}

impl Vector {
    /// The vector's rows at `2^afrac` (§13.5): one agent, each action's move `(dx, dy)` as
    /// `dequantize(dx, afrac, 0)`, `dequantize(dy, afrac, 0)` and 0; two agents, each pair's
    /// two moves.
    fn rows(&self) -> Vec<Vec<i64>> {
        let afrac = self.j["afrac"].as_u64().unwrap() as u32;
        let q = |d: i64| wai_ops::stage::dequantize(d, afrac, 0).unwrap();
        let mv = |a: &Value| int_world::ACTIONS[a.as_u64().unwrap() as usize];
        match self.j.get("pairs") {
            Some(p) => p.as_array().unwrap().iter().map(|p| {
                let ((xa, ya), (xb, yb)) = (mv(&p[0]), mv(&p[1]));
                vec![q(xa), q(ya), q(xb), q(yb)]
            }).collect(),
            None => self.j["actions"].as_array().unwrap().iter().map(|a| {
                let (dx, dy) = mv(a);
                vec![q(dx), q(dy), 0]
            }).collect(),
        }
    }

    /// The direct renderer's clip (`int_world::legacy`).
    fn direct(&self) -> int_world::WorldClip {
        int_world::legacy::decode_planes(&self.t, &self.g_s, &self.z0, &self.rows(), self.c, self.lat)
    }

    /// The vector's goldens against `frames`: each frame's FNV-1a (the export's) and
    /// SHA-256 (the sidecar's), then the clip's. Every frame that differs is named.
    fn check_goldens(&self, what: &str, frames: &[Vec<u8>]) {
        let fnv = self.j["frame_hashes"].as_array().unwrap();
        let sha = self.sha["frames_sha256"].as_array().unwrap();
        assert_eq!((frames.len(), fnv.len(), sha.len()), (fnv.len(), sha.len(), self.rows().len() + 1), "{} {what}: the frame count", self.name);
        let off: Vec<usize> = (0..frames.len())
            .filter(|&i| format!("{:016x}", wai::int_hyper_synth::fnv1a(&frames[i])) != fnv[i].as_str().unwrap() || sha256(&frames[i]) != sha[i].as_str().unwrap())
            .collect();
        assert!(off.is_empty(), "{} {what}: frames {off:?} of 0..{} differ from their committed FNV-1a and SHA-256", self.name, frames.len());
        let clip = int_world::WorldClip { frames: frames.to_vec(), frame_px: 0 };
        assert_eq!(Some(format!("{:016x}", clip.clip_hash()).as_str()), self.j["clip_hash"].as_str(), "{} {what}: the clip's FNV-1a", self.name);
        let all: Vec<u8> = frames.iter().flat_map(|f| wai_ops::sha256::sha256(f)).collect();
        assert_eq!(Some(sha256(&all).as_str()), self.sha["clip_sha256"].as_str(), "{} {what}: the clip's SHA-256", self.name);
    }
}

/// Gate 2 of the migration: golden equality. Each committed vector renders through the
/// lowering (`int_world::ops`) to the frames the direct renderer gives
/// (`int_world::legacy`), every byte, and to every golden the corpus commits for it: the
/// nine frame SHA-256s and the clip SHA-256 of `sha256.json`, and the export's FNV-1a
/// frame and clip hashes. The two-agent vector (`A = 4`) renders from its rows through
/// `decode_planes`; the one-agent vector (`A = 3`) from its rows, its moves
/// (`decode_deltas`) and its action indices (`decode`) alike.
#[test]
fn every_committed_vector_renders_through_the_lowering_to_its_goldens() {
    for v in vectors() {
        let direct = v.direct();
        let world = int_world::ops::load(&v.t, &v.g_s, v.c).unwrap_or_else(|r| panic!("{}: {r}", v.name));
        assert_eq!(world.planes(), if v.name == "world2" { 4 } else { 3 }, "{}: A", v.name);
        let ops = world.decode_planes(&v.z0, &v.rows(), v.lat).unwrap_or_else(|r| panic!("{}: {r}", v.name));
        let off: Vec<usize> = (0..ops.frames.len().max(direct.frames.len())).filter(|&i| ops.frames.get(i) != direct.frames.get(i)).collect();
        assert!(off.is_empty(), "{}: frames {off:?} through the lowering differ from the direct renderer's", v.name);
        assert_eq!((ops.frame_px, direct.frame_px), (v.j["frame"].as_u64().unwrap() as usize, v.j["frame"].as_u64().unwrap() as usize), "{}", v.name);
        v.check_goldens("through the lowering", &ops.frames);
        v.check_goldens("direct", &direct.frames);
        let free = int_world::ops::decode_planes(&v.t, &v.g_s, &v.z0, &v.rows(), v.c, v.lat).unwrap();
        assert!(free.frames == ops.frames, "{}: the set loaded per call renders alike", v.name);
        if v.name == "world" {
            let actions: Vec<usize> = v.j["actions"].as_array().unwrap().iter().map(|a| a.as_u64().unwrap() as usize).collect();
            let deltas: Vec<(i64, i64)> = actions.iter().map(|&a| int_world::ACTIONS[a]).collect();
            assert!(world.decode_deltas(&v.z0, &deltas, v.lat).unwrap().frames == ops.frames, "world: decode_deltas");
            assert!(int_world::ops::decode(&v.t, &v.g_s, &v.z0, &actions, v.c, v.lat).unwrap().frames == ops.frames, "world: decode");
            assert!(int_world::legacy::decode(&v.t, &v.g_s, &v.z0, &actions, v.c, v.lat).frames == direct.frames, "world: the direct decode of the actions");
        }
        // The live path: one transition and one frame at a time.
        let rows = v.rows();
        let (f0, h, w) = world.render(&v.z0, v.lat).unwrap();
        assert_eq!((h, w), (direct.frame_px, direct.frame_px));
        let mut frames = vec![f0];
        let mut z = v.z0.clone();
        for r in &rows {
            z = world.step(&z, r, v.lat).unwrap();
            frames.push(world.render(&z, v.lat).unwrap().0);
        }
        assert!(frames == ops.frames, "{}: step and render, one at a time, give the clip", v.name);
    }
}

/// The layer whose input check a refusal of a world graph is (§13.2's `x_ℓ`), from where
/// the graph refused: for `T` (`A` fills, then the concatenation that writes `x_0`, then
/// the layers), the output of operator `A + ℓ`; for `g_s`, its input for layer 0 and the
/// output of operator `ℓ − 1` for layer `ℓ`.
fn layer_of(net: &str, a: usize, at: wai_ops::Site) -> usize {
    match (net, at) {
        ("T", wai_ops::Site::Op(i)) if i as usize >= a => i as usize - a,
        ("g_s", wai_ops::Site::Input(0)) => 0,
        ("g_s", wai_ops::Site::Op(i)) => i as usize + 1,
        _ => panic!("{net} refused at {at:?}, which is no layer's check"),
    }
}

/// The renderer through the operator set, as `World::decode_planes` runs it, one graph at
/// a time, so that a refusal names its network, its step and its layer
/// ([`layer_of`]). It must agree with `World::decode_planes` (checked by the caller).
fn traced(w: &int_world::ops::World, z0: &[i64], rows: &[Vec<i64>], lat: usize) -> Result<Vec<Vec<u8>>, (&'static str, usize, usize, wai_ops::Code)> {
    use wai_ops::Tensor;
    let (c, a) = (w.channels() as u32, w.planes());
    let g = |z: &Tensor, step: usize| {
        let out = wai::int_ops::run1(w.g_s(), z.clone(), &[]).map_err(|r| ("g_s", step, layer_of("g_s", a, r.at), r.code))?;
        let plane = (out.h * out.w) as usize;
        let mut hwc = vec![0u8; out.data.len()];
        for (k, &v) in out.data.iter().enumerate() {
            hwc[(k % plane) * out.c as usize + k / plane] = v as u8;
        }
        Ok(hwc)
    };
    let mut z = Tensor::new(c, lat as u32, lat as u32, z0.to_vec());
    let mut frames = vec![g(&z, 0)?];
    for (k, row) in rows.iter().enumerate() {
        z = wai::int_ops::run1(w.t(), z, row).map_err(|r| ("T", k + 1, layer_of("T", a, r.at), r.code))?;
        frames.push(g(&z, k + 1)?);
    }
    Ok(frames)
}

/// Gate 3, at the edge of every `check` range: for each layer of `T` and of `g_s` in every
/// committed set, the layer's operator in the full `world/1` graph is run alone, from a
/// graph of that one operator whose input has the record the full graph gives the layer's
/// input (for `T`, layer 0 reads `x_0`, the concatenation of `z` and the action planes).
/// That range is §13.2's `[−X_ℓ, X_ℓ]`: the range of the layer lowered alone by `whm1/1`
/// for a convolution or an inverse GDN, and `⌊(2^63 − 1) / |slope|⌋` for a leaky ReLU. An
/// input whose largest magnitude is `X_ℓ` runs, through the operator set and through the
/// direct guard (`run_net_checked`), to the same output; one at `X_ℓ + 1` is refused by
/// both: by the graph (`field_range`, at its input) and by the guard (a value out of range
/// at that layer). So the ranges refuse exactly what the guard refuses, at the same layer.
#[test]
fn every_check_range_refuses_exactly_what_the_direct_guard_refuses() {
    use wai::int_hyper_synth::legacy::run_net_checked;
    use wai::int_hyper_synth::IntLayer;
    use wai_ops::graph::{Graph, InputRec, Op, Policy, TensorRec, CAP};
    use wai_ops::{Code, Site, Tensor};
    let mut checked = 0;
    for s in indexed() {
        let (tb, gb) = (s.member("T.whm"), s.member("g_s.whm"));
        let full = wai_ops::lower::world(&tb, &gb, s.c).unwrap();
        for (net, bytes) in [("T", &tb), ("g_s", &gb)] {
            let g = full.get(net).unwrap();
            // T's layers follow z, its A fills and the concatenation; g_s's follow its input.
            let (first_t, first_op) = if net == "T" { (1 + g.scalars.len(), g.scalars.len() + 1) } else { (0, 0) };
            let m = IntModel::from_bytes(bytes).unwrap();
            for l in 0..m.layers.len() {
                let what = format!("{} {net} layer {l}", s.name);
                let rec = g.tensors[first_t + l];
                let op = &g.ops[first_op + l];
                if rec.policy != Policy::Check {
                    assert!(matches!(m.layers[l], IntLayer::Relu), "{what}: only a ReLU's input is derived");
                    continue;
                }
                let mut one = IntModel::from_bytes(bytes).unwrap();
                let layer = one.layers.swap_remove(l);
                one.layers = vec![layer];
                match &one.layers[0] {
                    IntLayer::Leaky { slope_q, .. } => {
                        let x = i64::MAX as u64 / slope_q.unsigned_abs();
                        assert_eq!((rec.lo, rec.hi), (-(x as i64), x as i64), "{what}: X·|slope| <= 2^63 - 1");
                    }
                    _ => {
                        let alone = wai_ops::lower::whm1_net(&wai::int_ops::whm1_net(&one).unwrap(), false).unwrap();
                        assert_eq!(alone.tensors[0], rec, "{what}: the range of the layer lowered alone");
                    }
                }
                let (c, hi) = (rec.channels, rec.hi);
                assert!(hi > 0 && rec.lo == -hi, "{what}: a symmetric check range");
                let out = TensorRec { policy: Policy::Derived, lo: 0, hi: 0, ..g.tensors[op.outs[0] as usize] };
                let alone = Graph {
                    max_elems: CAP as u32,
                    scalars: vec![],
                    inputs: vec![InputRec { tensor: 0, max_h: 0, max_w: 0, kind: 0, src: 0, init: 0 }],
                    outputs: vec![1],
                    tensors: vec![rec, out],
                    ops: vec![Op { code: op.code, ins: vec![0], outs: vec![1], p: op.p.clone() }],
                };
                let loaded = wai::int_ops::load(&alone).unwrap_or_else(|r| panic!("{what}: the operator alone loads: {r:?}"));
                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!("{what}: {v} runs: {r:?}"));
                    let (direct, ..) = run_net_checked(&one, &x(v), c as usize, 2, 2).unwrap_or_else(|r| panic!("{what}: {v} runs directly: {r}"));
                    assert!(ops.data == direct, "{what}: 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)), "{what}: {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, "{what}: {v}: {d}");
                }
                checked += 1;
            }
        }
    }
    // world and world2: T's 3 layers less none, g_s's 5; the synthetic set: 3 and 2.
    assert_eq!(checked, 8 + 8 + 5, "every layer's range was checked");
}

/// Gate 3, through the whole renderer: on each committed vector, with a latent value or an
/// action plane value placed at the edge of a `check` range or past it (`g_s`'s input
/// range for `z_0`, `T`'s `x_0` range for a row's value, the ends of `i64`), the renderer
/// through the operator set renders the frames the direct renderer gives, or refuses where
/// the direct guard refuses: `field_range`, in the same network, at the same step, at the
/// same layer.
#[test]
fn the_renderer_refuses_where_the_direct_guard_refuses_at_the_same_layer() {
    use int_world::legacy::checked;
    use wai::int_hyper_synth::RefusalKind;
    for v in vectors() {
        let world = int_world::ops::load(&v.t, &v.g_s, v.c).unwrap();
        let (tg, gg) = int_world::ops::lower(&v.t, &v.g_s, v.c).unwrap();
        let (xt, xg) = (tg.tensors[1 + world.planes()].hi, gg.tensors[0].hi);
        let k = 5 * v.lat * v.lat + 10;
        let mut cases: Vec<(String, Vec<i64>, Vec<Vec<i64>>)> = vec![("committed".into(), v.z0.clone(), v.rows())];
        for z in [xg, xg + 1, -xg, -xg - 1, i64::MIN, i64::MAX] {
            let mut z0 = v.z0.clone();
            z0[k] = z;
            cases.push((format!("z_0[{k}] = {z}"), z0, v.rows()));
        }
        for (row, p, x) in [(0usize, 0usize, xt), (0, 0, xt + 1), (0, 1, -xt), (0, 1, -xt - 1), (3, 2, i64::MAX), (7, 0, i64::MIN)] {
            let mut rows = v.rows();
            rows[row][p] = x;
            cases.push((format!("row {row} plane {p} = {x}"), v.z0.clone(), rows));
        }
        let (mut ok, mut refused) = (0, 0);
        for (what, z0, rows) in &cases {
            let ops = world.decode_planes(z0, rows, v.lat).map(|c| c.frames);
            let tr = traced(&world, z0, rows, v.lat);
            let direct = checked::decode_planes(&v.t, &v.g_s, z0, rows, v.c, v.lat).map(|c| c.frames);
            match (&ops, &tr, &direct) {
                (Ok(a), Ok(b), Ok(d)) => {
                    assert!(a == d && b == d, "{} {what}: the frames differ from the direct renderer's", v.name);
                    let plain = int_world::legacy::decode_planes(&v.t, &v.g_s, z0, rows, v.c, v.lat).frames;
                    assert!(plain == *d, "{} {what}: the guard changed the direct renderer's frames", v.name);
                    ok += 1;
                }
                (Err(a), Err((net, step, layer, code)), Err(d)) => {
                    assert_eq!((a.kind, *code), (RefusalKind::FieldRange, wai_ops::Code::FieldRange), "{} {what}: {a}", v.name);
                    assert!(d.out_of_range, "{} {what}: the guard refused a value, {d:?}", v.name);
                    assert_eq!((*net, *step, *layer), (d.net, d.step, d.layer), "{} {what}: the network, step and layer of the refusal", v.name);
                    refused += 1;
                }
                _ => panic!("{} {what}: through the operator set {ops:?}, traced {tr:?}, directly {direct:?}", v.name, ops = ops.as_ref().map(|f| f.len()), tr = tr.as_ref().map(|f| f.len()), direct = direct.as_ref().map(|f| f.len())),
            }
        }
        assert!(ok >= 1 && refused >= 8, "{}: {ok} rendered and {refused} refused alike", v.name);
    }
}

/// The renderer's own checks (`extensions/integer-operators.md` §13.5, informative; §15.3
/// item 2), each refused through the operator set with its code before any frame is
/// rendered, where the direct renderer, which checked nothing, rendered a clip from them or
/// panicked: a network's `afrac` outside `1..=30` or unequal to the other's, `A` of 0, a row
/// that does not carry `A` values (a single agent's moves for the two agent set among them),
/// `T` that does not give the latent's shape back, `g_s` that does not read `C` channels
/// (`model_mismatch`); a latent that is not `C·lat·lat` values (`geometry_mismatch`), a side
/// of 0 or a move whose plane value leaves `i64` (`field_range`), and a latent past the cap
/// (`too_large`).
#[test]
fn the_renderers_own_checks_refuse_before_any_frame() {
    use wai::int_hyper_synth::{IntLayer, RefusalKind::*};
    let [w1, w2] = vectors();
    let direct = |t: &IntModel, g: &IntModel, z0: &[i64], rows: &[Vec<i64>], c: usize, lat: usize| {
        Some(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| int_world::legacy::decode_planes(t, g, z0, rows, c, lat).frames.len())).ok())
    };
    // `d`: what the direct renderer did with the same input (`Some(None)`: it panicked),
    // or `None` where it takes no such input.
    let refuse = |what: &str, want, got: Result<int_world::WorldClip, wai::int_hyper_synth::Refused>, d: Option<Option<usize>>| {
        assert_eq!(got.map(|c| c.frames.len()).map_err(|r| r.kind), Err(want), "{what}");
        let d = match d {
            None => "takes no such input".to_string(),
            Some(None) => "panicked".to_string(),
            Some(Some(n)) => format!("rendered {n} frames"),
        };
        println!("{what}: refused {want:?}; the direct renderer {d}");
    };
    let model = |m: &IntModel, f: &dyn Fn(&mut IntModel)| {
        let mut m = IntModel { afrac: m.afrac, gf: m.gf, shift: m.shift, layers: m.layers.iter().map(clone_layer).collect() };
        f(&mut m);
        m
    };
    let (t, g, z0, rows, c, lat) = (&w1.t, &w1.g_s, &w1.z0, w1.rows(), w1.c, w1.lat);
    for (what, tm, gm) in [
        ("T's afrac 0", model(t, &|m| m.afrac = 0), model(g, &|_| {})),
        ("g_s's afrac 31", model(t, &|_| {}), model(g, &|m| m.afrac = 31)),
        ("both afrac 31", model(t, &|m| m.afrac = 31), model(g, &|m| m.afrac = 31)),
        ("T's afrac 13, g_s's 12", model(t, &|m| m.afrac = 13), model(g, &|_| {})),
    ] {
        let d = direct(&tm, &gm, z0, &rows, c, lat);
        refuse(what, ModelMismatch, int_world::ops::decode_planes(&tm, &gm, z0, &rows, c, lat), d);
    }
    // A = 0: C is all of T's input channels.
    refuse("A = 0", ModelMismatch, int_world::ops::decode_planes(t, g, &vec![0; 27 * lat * lat], &[], 27, lat), direct(t, g, &vec![0; 27 * lat * lat], &[], 27, lat));
    // Rows without A values: short, long, and a single agent's moves for the two agent set.
    let mut short = rows.clone();
    short[4].pop();
    refuse("a row of A - 1 values", ModelMismatch, int_world::ops::decode_planes(t, g, z0, &short, c, lat), direct(t, g, z0, &short, c, lat));
    let mut long = rows.clone();
    long[0].push(0);
    refuse("a row of A + 1 values", ModelMismatch, int_world::ops::decode_planes(t, g, z0, &long, c, lat), direct(t, g, z0, &long, c, lat));
    let moves: Vec<Vec<i64>> = w2.rows().iter().map(|r| vec![r[0], r[1], 0]).collect();
    refuse("three planes for the two agent T", ModelMismatch, int_world::ops::decode_planes(&w2.t, &w2.g_s, &w2.z0, &moves, w2.c, w2.lat), direct(&w2.t, &w2.g_s, &w2.z0, &moves, w2.c, w2.lat));
    refuse("a single agent's moves for the two agent T", ModelMismatch, int_world::ops::decode_deltas(&w2.t, &w2.g_s, &w2.z0, &[(1, 0)], w2.c, w2.lat), None);
    // T that halves the latent's side; g_s that reads 23 channels.
    let halving = model(t, &|m| if let IntLayer::Conv { stride, .. } = &mut m.layers[2] { *stride = 2 });
    refuse("T that does not give the latent's shape back", ModelMismatch, int_world::ops::decode_planes(&halving, g, z0, &rows, c, lat), direct(&halving, g, z0, &rows, c, lat));
    let narrow = model(g, &|m| if let IntLayer::ConvT { weight, .. } = &mut m.layers[0] { weight.truncate(23 * 24 * 16) });
    refuse("g_s that reads 23 channels", ModelMismatch, int_world::ops::decode_planes(t, &narrow, z0, &rows, c, lat), direct(t, &narrow, z0, &rows, c, lat));
    // The latent's shape and the moves' values.
    refuse("z_0 one value short", GeometryMismatch, int_world::ops::decode_planes(t, g, &z0[1..], &rows, c, lat), direct(t, g, &z0[1..], &rows, c, lat));
    refuse("a side of 0", FieldRange, int_world::ops::decode_planes(t, g, &[], &rows, c, 0), direct(t, g, &[], &rows, c, 0));
    refuse("a latent past 2^26 elements", TooLarge, int_world::ops::decode_planes(t, g, &[], &rows, c, 1 << 11), None);
    refuse("a move past i64 at 2^afrac", FieldRange, int_world::ops::decode_deltas(t, g, z0, &[(1, 0), (i64::MAX >> 11, 0)], c, lat), None);
    // A refused set or row renders no frame: the checks precede frame 0.
    let world = int_world::ops::load(t, g, c).unwrap();
    assert_eq!(world.step(z0, &[0, 0], lat).map(|_| ()).map_err(|r| r.kind), Err(ModelMismatch), "step: a row of two values");
    assert_eq!(world.render(&z0[1..], lat).map(|_| ()).map_err(|r| r.kind), Err(GeometryMismatch), "render: z_0 one value short");
    assert_eq!(int_world::ops::load(&w2.t, &w2.g_s, w2.c).unwrap().step_action(&w2.z0, 1, w2.lat).map(|_| ()).map_err(|r| r.kind), Err(ModelMismatch), "an action index for the two agent T");
}

fn clone_layer(l: &wai::int_hyper_synth::IntLayer) -> wai::int_hyper_synth::IntLayer {
    use wai::int_hyper_synth::IntLayer::*;
    match l {
        ConvT { out_ch, k, stride, pad, opad, weight, mult, bias } => ConvT { out_ch: *out_ch, k: *k, stride: *stride, pad: *pad, opad: *opad, weight: weight.clone(), mult: mult.clone(), bias: bias.clone() },
        Conv { out_ch, k, stride, pad, weight, mult, bias } => Conv { out_ch: *out_ch, k: *k, stride: *stride, pad: *pad, weight: weight.clone(), mult: mult.clone(), bias: bias.clone() },
        Igdn { gamma, beta } => Igdn { gamma: gamma.clone(), beta: beta.clone() },
        Relu => Relu,
        Leaky { slope_q, slope_frac } => Leaky { slope_q: *slope_q, slope_frac: *slope_frac },
    }
}

/// The new rules affect no committed set or vector: every `wai.world.model` set the
/// corpus indexes loads (one agent, two agents, the synthetic set), and the committed
/// vectors render (gate 2 holds them to their goldens). `git grep -l WHM1` names no other
/// world set: the renderer's sets are the two under `wai-rs/tests/vectors/int_world{,2}`
/// and the synthetic one under `ops-conformance/sets/world`.
#[test]
fn the_new_rules_affect_no_committed_set() {
    for s in indexed() {
        let (t, g) = (IntModel::from_bytes(&s.member("T.whm")).unwrap(), IntModel::from_bytes(&s.member("g_s.whm")).unwrap());
        let w = int_world::ops::load(&t, &g, s.c as usize).unwrap_or_else(|r| panic!("{}: {r}", s.name));
        // The synthetic set's T keeps a latent of any side; it renders a frame from zeros.
        let lat = 3;
        let z = vec![0i64; s.c as usize * lat * lat];
        let (f, h, wd) = w.render(&z, lat).unwrap_or_else(|r| panic!("{}: {r}", s.name));
        assert_eq!(f.len(), w.plan(lat).map(|(c, h, w)| c * h * w).unwrap(), "{}", s.name);
        assert!(h > 0 && wd > 0);
        let row = vec![0i64; w.planes()];
        assert_eq!(w.step(&z, &row, lat).unwrap().len(), z.len(), "{}: T gives the latent's shape back", s.name);
    }
}

/// Gate 5 of the migration, the whole-renderer half: the renderer through the lowering
/// under the reference kernels (the interpreter computes every operator) and under each
/// fast policy (the convolutions and inverse GDNs on this crate's fast kernels, serial and
/// threaded) renders the same frames, or refuses alike, on each committed vector, and with
/// a latent value at `g_s`'s input range, one past it, and a row value at `T`'s `x_0` range
/// and one past it. The per-operator half is `int_ops`' equivalence tests, one of which
/// counts the fast kernels' runs on these sets.
#[test]
fn every_kernel_policy_renders_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);
    for v in vectors() {
        let world = int_world::ops::load(&v.t, &v.g_s, v.c).unwrap();
        let (tg, gg) = int_world::ops::lower(&v.t, &v.g_s, v.c).unwrap();
        let (xt, xg) = (tg.tensors[1 + world.planes()].hi, gg.tensors[0].hi);
        let mut cases: Vec<(String, Vec<i64>, Vec<Vec<i64>>)> = vec![("committed".into(), v.z0.clone(), v.rows())];
        for z in [xg, xg + 1] {
            let mut z0 = v.z0.clone();
            z0[7] = z;
            cases.push((format!("z_0[7] = {z}"), z0, v.rows()));
        }
        for x in [xt, -xt - 1] {
            let mut rows = v.rows();
            rows[2][1] = x;
            cases.push((format!("row 2 plane 1 = {x}"), v.z0.clone(), rows));
        }
        for (what, z0, rows) in &cases {
            let render = |p: Policy| with_policy(p, || world.decode_planes(z0, rows, v.lat).map(|c| c.frames).map_err(|r| r.kind));
            let want = render(POLICIES[0]);
            assert!(what != "committed" || want.is_ok(), "{}: the committed clip renders", v.name);
            for p in &POLICIES[1..] {
                assert!(render(*p) == want, "{} {what}: {p:?} renders other than the interpreter alone", v.name);
            }
        }
    }
}