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

//! Every committed integer-decoder vector whose decode reaches the fast kernels
//! (int_hyper, int_mlicpp, int_audio, int_splat, int_world one and two agents, the
//! neural-keyframe video, the int_synth corpus), plus the real-footage video,
//! decoded under the scalar reference kernels AND under each fast-kernel policy
//! (serial, explicit multi-threaded tilings, automatic): the outputs must be
//! byte-identical to each other and must reproduce the vector's frozen hash. The
//! plain `int_video` vector (block motion compensation + rANS, no learned layer)
//! never reaches the kernels and is left to `tests/int_video.rs` and the
//! `byte-exact-conformance` goldens. This is the whole-decoder half of the
//! fast-vs-reference equivalence check; the per-kernel half (randomised shapes and
//! operand ranges, element by element) lives beside each kernel. CI runs both
//! halves on every platform leg of `wai-conformance`, with overflow checks on and
//! off.
#![cfg(all(feature = "neural_int", feature = "video_int", feature = "audio_int", feature = "world_int", feature = "splat_int"))]

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

use serde_json::Value;

use wai::int_kernels::{with_policy, Policy};

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

const RGB8: &str = "RGB8: rows top to bottom, pixels left to right, R, G, B bytes";

const POLICIES: [Policy; 5] = [Policy::reference(), Policy::fast(1), Policy::fast(2), Policy::fast(5), Policy::fast(0)];

fn vroot() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors")
}
fn jload(p: &Path) -> serde_json::Value {
    serde_json::from_slice(&std::fs::read(p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()))).unwrap()
}

/// Decode under every policy; all outputs must be identical. Returns the output.
fn same_under_every_policy<T: PartialEq + std::fmt::Debug>(what: &str, f: impl Fn() -> T) -> T {
    let want = with_policy(POLICIES[0], &f);
    for p in &POLICIES[1..] {
        let got = with_policy(*p, &f);
        assert!(got == want, "{what}: output under {p:?} differs from the reference kernels");
    }
    want
}

#[test]
fn int_hyper_tiny_conditional_codec() {
    use wai::int_hyper_synth::{decode_conditional_checked, fnv1a, CondBitstream, CondEntropy, IntModel};
    let d = vroot().join("int_hyper");
    let want = jload(&d.join("tiny.json"))["rgb_hash"].as_str().unwrap().to_string();
    let ent = CondEntropy::from_value(&jload(&d.join("tiny_entropy.json")));
    let bs = CondBitstream::from_value(&jload(&d.join("tiny_bitstream.json")));
    let g_s = IntModel::from_bytes(&std::fs::read(d.join("tiny_g_s.bin")).unwrap()).unwrap();
    let h_s = IntModel::from_bytes(&std::fs::read(d.join("tiny_h_s.bin")).unwrap()).unwrap();
    let (rgb, _, _) = same_under_every_policy("int_hyper tiny", || decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).unwrap());
    assert_eq!(format!("{:016x}", fnv1a(&rgb)), want);
    sha256_sidecar::one(&d, "tiny", RGB8, &rgb);
}

#[test]
fn int_mlicpp_conformance_vector() {
    use wai::int_hyper_synth::fnv1a;
    use wai::int_mlicpp::{decode, IntMlicModel};
    let d = vroot().join("int_mlicpp");
    let model = IntMlicModel::from_bundle(&std::fs::read(d.join("model.wmb")).unwrap()).unwrap();
    let env = wai::container::Wai::from_bytes(&std::fs::read(d.join("stream.wai")).unwrap()).unwrap();
    let img = same_under_every_policy("int_mlicpp", || decode(&model, &env.payload).unwrap());
    assert_eq!(format!("{:016x}", fnv1a(&img.rgb)), jload(&d.join("expected.json"))["rgb_fnv1a"].as_str().unwrap());
}

#[test]
fn int_audio_clip() {
    use wai::int_audio::{decode, pcm_fnv1a, AudioModel};
    let d = vroot().join("int_audio");
    let model = AudioModel::from_bytes(&std::fs::read(d.join("audio_model.bin")).unwrap()).unwrap();
    let bs = jload(&d.join("audio_bitstream.json"));
    let pcm = same_under_every_policy("int_audio", || decode(&model, &bs).unwrap());
    assert_eq!(format!("{:016x}", pcm_fnv1a(&pcm)), jload(&d.join("audio.json"))["pcm_fnv1a"].as_str().unwrap());
    let s16le: Vec<u8> = pcm.iter().flat_map(|s| s.to_le_bytes()).collect();
    sha256_sidecar::one(&d, "audio", "PCM: signed 16-bit little-endian samples in output order, channels interleaved", &s16le);
}

#[test]
fn int_splat_cloud() {
    use wai::int_hyper_synth::IntModel;
    use wai::int_splat::decode;
    let d = vroot().join("int_splat");
    let gs = IntModel::from_bytes(&std::fs::read(d.join("splat_gs.bin")).unwrap()).unwrap();
    let bs = jload(&d.join("splat.json"));
    let attr = same_under_every_policy("int_splat", || decode(&gs, &bs).unwrap().attr);
    let hash = wai::int_hyper_synth::fnv1a(&attr);
    assert_eq!(format!("{hash:016x}"), bs["attr_hash"].as_str().unwrap());
    sha256_sidecar::one(&d, "splat", "DecodedSplat::attr: the decoded attribute tensor's bytes in decoder output order", &attr);
}

#[test]
fn int_world_clips() {
    use wai::int_hyper_synth::IntModel;
    use wai::int_world::{decode, decode_planes, ACTIONS};
    // single agent
    let d = vroot().join("int_world");
    let t = IntModel::from_bytes(&std::fs::read(d.join("world_T.bin")).unwrap()).unwrap();
    let gs = IntModel::from_bytes(&std::fs::read(d.join("world_gs.bin")).unwrap()).unwrap();
    let j = jload(&d.join("world.json"));
    let (c, lat) = (j["c"].as_u64().unwrap() as usize, j["lat"].as_u64().unwrap() as usize);
    let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|v| v.as_i64().unwrap()).collect();
    let actions: Vec<usize> = j["actions"].as_array().unwrap().iter().map(|v| v.as_u64().unwrap() as usize).collect();
    let frames = same_under_every_policy("int_world", || decode(&t, &gs, &z0, &actions, c, lat).map(|c| c.frames).map_err(|r| r.kind)).expect("the committed clip renders");
    let clip = wai::int_world::WorldClip { frames, frame_px: 0 };
    assert_eq!(format!("{:016x}", clip.clip_hash()), j["clip_hash"].as_str().unwrap());
    sha256_sidecar::frames(&d, "world", RGB8, &clip.frames);

    // two agents
    let d = vroot().join("int_world2");
    let t = IntModel::from_bytes(&std::fs::read(d.join("world2_T.bin")).unwrap()).unwrap();
    let gs = IntModel::from_bytes(&std::fs::read(d.join("world2_gs.bin")).unwrap()).unwrap();
    let j = jload(&d.join("world2.json"));
    let (c, lat) = (j["c"].as_u64().unwrap() as usize, j["lat"].as_u64().unwrap() as usize);
    let one = 1i64 << j["afrac"].as_u64().unwrap();
    let z0: Vec<i64> = j["z0"].as_array().unwrap().iter().map(|v| v.as_i64().unwrap()).collect();
    let rows: Vec<Vec<i64>> = j["pairs"]
        .as_array()
        .unwrap()
        .iter()
        .map(|p| {
            let (a, b) = (ACTIONS[p[0].as_u64().unwrap() as usize], ACTIONS[p[1].as_u64().unwrap() as usize]);
            vec![a.0 * one, a.1 * one, b.0 * one, b.1 * one]
        })
        .collect();
    let frames = same_under_every_policy("int_world2", || decode_planes(&t, &gs, &z0, &rows, c, lat).map(|c| c.frames).map_err(|r| r.kind)).expect("the committed clip renders");
    let want = j["frame_hashes"].as_array().unwrap();
    assert_eq!(frames.len(), want.len());
    for (f, w) in frames.iter().zip(want) {
        assert_eq!(format!("{:016x}", wai::int_hyper_synth::fnv1a(f)), w.as_str().unwrap());
    }
    sha256_sidecar::frames(&d, "world2", RGB8, &frames);
}

#[test]
fn int_video_neural_keyframe_and_real_footage() {
    use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy, IntModel};
    use wai::int_video::{decode, decode_pframes, fnv1a};
    let (ih, iv) = (vroot().join("int_hyper"), vroot().join("int_video"));
    let g_s = IntModel::from_json(&jload(&ih.join("tiny_g_s.json")));
    let h_s = IntModel::from_json(&jload(&ih.join("tiny_h_s.json")));
    let ent = CondEntropy::from_value(&jload(&ih.join("tiny_entropy.json")));
    let bs = CondBitstream::from_value(&jload(&ih.join("tiny_bitstream.json")));
    let vbs = jload(&iv.join("video_nk_bitstream.json"));
    let frames = same_under_every_policy("neural-keyframe video", || {
        let (f0, h, w) = decode_conditional_checked(&g_s, &h_s, &ent, &bs.stream()).unwrap();
        decode_pframes(&f0, h, w, &vbs).unwrap().frames
    });
    let want = jload(&iv.join("video_nk.json"));
    let want = want["frame_hashes"].as_array().unwrap();
    assert_eq!(frames.len(), want.len(), "neural-keyframe video: frame count");
    for (f, w) in frames.iter().zip(want) {
        assert_eq!(format!("{:016x}", fnv1a(f)), w.as_str().unwrap());
    }
    sha256_sidecar::frames(&iv, "video_nk", RGB8, &frames);
    // integer block-MC only (no learned kernels) — unaffected, checked for completeness
    let rbs = jload(&iv.join("video_real_bitstream.json"));
    let frames = same_under_every_policy("real-footage video", || decode(&rbs).unwrap().frames);
    let want = jload(&iv.join("video_real.json"));
    let want = want["frame_hashes"].as_array().unwrap();
    assert_eq!(frames.len(), want.len(), "real-footage video: frame count");
    for (f, w) in frames.iter().zip(want) {
        assert_eq!(format!("{:016x}", fnv1a(f)), w.as_str().unwrap());
    }
    sha256_sidecar::frames(&iv, "video_real", RGB8, &frames);
    // The plain block-motion clip (wai.video.int_motion), for completeness.
    let pbs = jload(&iv.join("int_video_bitstream.json"));
    let frames = decode(&pbs).unwrap().frames;
    let want = jload(&iv.join("int_video.json"));
    for (f, w) in frames.iter().zip(want["frame_hashes"].as_array().unwrap()) {
        assert_eq!(format!("{:016x}", fnv1a(f)), w.as_str().unwrap());
    }
    sha256_sidecar::frames(&iv, "int_video", RGB8, &frames);
}

#[test]
fn neural_conformance_int_synth_corpus() {
    use wai::int_synth::decode_from_spec;
    let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../neural-conformance");
    let manifest = jload(&dir.join("expected_hashes.json"));
    let mut n = 0;
    for (name, spec) in manifest.as_object().unwrap() {
        let vd = dir.join("vectors").join(name);
        let payload = std::fs::read(vd.join("latent.bin")).unwrap();
        let model = std::fs::read(vd.join("int_synth.json")).unwrap();
        let rgb = same_under_every_policy(name, || decode_from_spec(&payload, &model).unwrap().rgb);
        let mut h = blake3::Hasher::new();
        h.update(b"wai:int-synth-rgb\x01");
        h.update(&rgb);
        let got: String = h.finalize().as_bytes().iter().map(|b| format!("{b:02x}")).collect();
        assert_eq!(got, spec["rgb_hash"].as_str().unwrap(), "{name}");
        sha256_sidecar::one(&dir, name, RGB8, &rgb);
        if let Ok(r) = std::fs::read(vd.join("ref_rgb.bin")) {
            assert_eq!(rgb, r, "{name}: must equal the python integer-forward reference");
        }
        n += 1;
    }
    assert!(n >= 4, "expected the four committed int_synth vectors, found {n}");
}

/// SPEC §7 *Digests*: every committed vector or corpus that states a digest of the
/// decoded output of a capability that decodes against a parameter set states its
/// SHA-256, in a `sha256.json` sidecar that names the buffer it covers. The test
/// enumerates every `wai/*-conformance/` directory and every `tests/vectors/`
/// directory: each has a sidecar, or is listed below with the reason it has none.
#[test]
fn every_decoded_output_states_a_sha256() {
    const NATIVE: &str = "a WAI-native capability that decodes against no parameter set: its state and \
                          rendered-output digests are the algorithm its extension names (SPEC §7 *Digests*)";
    let excluded: &[(&str, &str)] = &[
        ("audio-conformance", NATIVE), ("avatar-conformance", NATIVE), ("feed-conformance", NATIVE),
        ("film-conformance", NATIVE), ("haptic-conformance", NATIVE), ("score-conformance", NATIVE),
        ("volumetric-conformance", NATIVE), ("world-conformance", NATIVE),
        ("quantum-conformance", NATIVE), ("quantum-ops-conformance", NATIVE),
        ("claim-conformance", "receipt claims: receipt content, no decoded output"),
        ("jwp-conformance", "JWP delivery receipts: receipt content, no decoded output"),
        ("oer-conformance", "OER/2 receipts: receipt content, no decoded output"),
        ("pack-conformance", "encode/expected.json states the SHA-256 of every decoded frame (RGB8); the rest are policies, refusals, conversions and scene cuts, with no decoded output"),
        ("pin-conformance", "parameter-set pins: every digest there is a SHA-256 of a parameter set"),
        ("ops-conformance", "integer operators: every operator, stage and graph vector, the lower_* decodes of the sets under sets/ among them, states its output integers in full or its refusal code, never a digest of them; certificates/ states each lowered set's SHA-256 (a parameter-set identity) and holds the lowered graphs as bytes, compared byte for byte; the output digests a capability decoded through its lowering reproduces (integer-operators §14.4) are stated by that capability's own corpus"),
        ("signer-conformance", "key logs, grants and statuses: signed documents, no decoded output"),
        ("live-conformance", "live events, rooms, takes, compositions, recordings, captions, metering, steering, moderation, watermark sessions and ads: signed records, commitments, packages, clips, cues and their renderings, per-hop records, ledgers and session receipts with their units and joules recomputed, steering policies and decisions with each pathway's evidence and the steering manifests recomputed byte for byte, review policies, decisions, notices and the items' timelines with each statement-of-reasons export written byte for byte, watermark tokens, session records, recovered sequences and disclosures, ad policies, breaks with their interstitial tags, events and asset lists written byte for byte and their impression receipts, and their verdicts, and layouts and mixes stated in full, rectangle for rectangle and sample for sample; nothing decoded against a parameter set"),
        ("http-conformance", "objects and head proofs for transport: no case decodes a payload, and the component a fetch case picks is stated by its SHA-256"),
        ("cache-conformance", "cache rules: a driver, its controls and their reports, with no objects of its own"),
        ("byte-exact-conformance", "goldens over tests/vectors, whose sidecars state the SHA-256"),
        ("int-payload-conformance", "refusal fixtures and the index of the payloads under tests/vectors, whose sidecars state the SHA-256"),
        ("staged-conformance", "each case.json states the SHA-256 of every stage-0 buffer and of every unit's stage (SPEC §7 Digests)"),
        ("staged-measure-conformance", "each fixture's expect.json and goldens.json state the SHA-256 of every unit's output at every stage (SPEC §7 Digests)"),
        ("telemetry-conformance", "advisory observations and their CMCD and CMSD strings: no decoded output"),
        ("vectors/int_mlicpp", "expected.json states rgb_sha256 (RGB8)"),
        ("vectors/jpegai", "inputs only; the decoded pixels' SHA-256 is pinned by the JPEG AI e2e job"),
        ("vectors/staged_bounds", "bases and a layer that declare outputs past a sink's cost limit, refused before any symbol is decoded, and inert envelopes; the one base decoded is wai.video.int_motion at keyframe 0, which decodes against no parameter set"),
    ];
    let wai = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("..");
    let mut dirs: Vec<(String, PathBuf)> = Vec::new();
    for e in std::fs::read_dir(&wai).unwrap().flatten() {
        let n = e.file_name().to_string_lossy().into_owned();
        if n.ends_with("-conformance") && e.path().is_dir() {
            dirs.push((n, e.path()));
        }
    }
    for e in std::fs::read_dir(vroot()).unwrap().flatten() {
        if e.path().is_dir() {
            dirs.push((format!("vectors/{}", e.file_name().to_string_lossy()), e.path()));
        }
    }
    dirs.sort();
    let mut with_sidecar = 0;
    for (name, path) in &dirs {
        if let Some((_, why)) = excluded.iter().find(|(n, _)| n == name) {
            assert!(!why.is_empty(), "{name}: an exclusion names its reason");
            continue;
        }
        let side = path.join("sha256.json");
        let doc = jload(&side);
        let entries = doc.as_object().unwrap_or_else(|| panic!("{}: an object", side.display()));
        assert!(!entries.is_empty(), "{}: no entries", side.display());
        for (k, e) in entries {
            assert!(e["buffer"].as_str().is_some_and(|b| !b.is_empty()), "{name}/{k}: names its buffer");
            let mut digests = 0;
            for (field, v) in e.as_object().unwrap() {
                if !field.ends_with("sha256") {
                    continue;
                }
                let all: Vec<&Value> = match v { Value::Array(a) => a.iter().collect(), other => vec![other] };
                for d in all {
                    let d = d.as_str().unwrap_or_default();
                    assert!(d.len() == 64 && d.bytes().all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b)),
                            "{name}/{k}/{field}: {d} is not 64 lowercase hexadecimal digits");
                    digests += 1;
                }
            }
            assert!(digests > 0, "{name}/{k}: states a SHA-256");
        }
        with_sidecar += 1;
    }
    for (n, _) in excluded {
        assert!(dirs.iter().any(|(d, _)| d == n), "exclusion {n} names no directory");
    }
    assert!(with_sidecar > 0);
}