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, and each kernel set the
//! processor runs, `int_kernels::equivalence_policies`): 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";

static POLICIES: std::sync::LazyLock<Vec<Policy>> = std::sync::LazyLock::new(wai::int_kernels::equivalence_policies);

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"),
        ("platform-decode-conformance", "platform queries: codec strings and sizes read from payload headers; no case decodes a payload"),
        ("format-probe-conformance", "header probes: each vector states the parameters a probe reads from its file's headers or the code it refuses with, and its file's SHA-256; no case decodes a payload, and provenance.json states the SHA-256 of what each external decoder made of the encoders' vectors"),
        ("encode-conformance", "encoded AVC and HEVC elementary streams and their signed records: streams.json states each stream's SHA-256 and each record binds its stream; no case decodes a stream"),
        ("wgsl-conformance", "WGSL's integer semantics: every case states its result word in full, never a digest of decoded output"),
        ("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"),
        ("conform-result-conformance", "conformance results (conformance-results §7): signed records of conformance runs, each case's expected and actual answer stated as text; no decoded output"),
        ("edge-conformance", "derivation records of a deriving edge: each replicate-tier output is recomputed and compared byte for byte by its verifier, and every output is named by its content hash in the record"),
        ("live-conformance", "the corpora of the live, provenance and delivery extensions, one directory each (live-conformance/README.md lists its extension, generator and verifier): signed records, commitments and statements authored in code from fixed seeds and recomputed byte for byte by each corpus's verifier; payloads are opaque test bytes no case decodes, and a picture a corpus generates is named in it by the SHA-256 of its RGB8 samples; aggregate/vdaf-22/ holds the Prio3 test vectors of draft-irtf-cfrg-vdaf-22 as published, field elements and seeds with no decoded output, each file pinned by its SHA-256 in tests/telemetry_aggregate.rs; nothing decoded against a parameter set but splat/'s wai.splat.int_codec objects, whose decoded attribute tensors are stated by their SHA-256 (splat-exact §3.3's attribute digest) in its vectors.json and its commitments, as the frames they render are, and dred/'s payloads and streams, decoded under the weights file it names by SHA-256, whose decoded integers and feature rows are stated by their SHA-256 in its vectors.json"),
        ("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"),
        ("pq-conformance", "hybrid signatures: signed group receipts, claims, records and key logs, and the ACVP and composite-draft vectors their signatures are held to, as published: signed content, 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);
}