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

//! `wai.neural.int_mlicpp` on its committed conformance vector
//! (`tests/vectors/int_mlicpp/`, written by `tools/wai_int_mlicpp_export.py`). No
//! test here skips: the vector is in the repository.
//!
//! The vector's weights are synthetic (a fixed seed, quantised the way a trained
//! model's would be). These tests establish decode-equivalence — that this decoder
//! reproduces the export tool's independent integer reference byte for byte, derives
//! every scale bucket the encoder used, and does so under every kernel policy and
//! thread count — and nothing about compression quality.
//!
//! What they pin, in `expected.json` and again as literals below so that a changed
//! vector cannot change them unnoticed: FNV-1a of the decoded RGB, and FNV-1a over the
//! y stream's bucket indices (u16 LE) and residual symbols (i32 LE) in coding order.
//! The SHA-256 of the pixels, of the model bundle and of the payload are checked by
//! the `byte-exact-conformance` golden, whose platform matrix runs the same decode.
#![cfg(feature = "neural_int")]

use std::path::PathBuf;

use wai::container::Wai;
use wai::int_hyper_synth::{fnv1a, IntLayer};
use wai::int_kernels::{with_policy, Policy};
use wai::int_mlicpp::{decode, decode_traced, parse_header, Header, IntMlicImage, IntMlicModel, Trace};

const RGB_FNV1A: &str = "7c0b9a4b3423d775";
const BUCKET_FNV1A: &str = "2d30432f4fb206a0";
const SYMBOL_FNV1A: &str = "14ace1d4ad3a2446";
const MODEL_SHA256: &str = "7ef44f3b1fc4a2df6c3e3b68460f595cae3d31cb7f9018699e29e738340272a4";

fn dir() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/int_mlicpp")
}
fn read(name: &str) -> Vec<u8> {
    let p = dir().join(name);
    std::fs::read(&p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()))
}
fn model() -> IntMlicModel {
    IntMlicModel::from_bundle(&read("model.wmb")).expect("the committed model bundle")
}
fn envelope() -> Wai {
    Wai::from_bytes(&read("stream.wai")).expect("the committed envelope")
}
fn expected() -> serde_json::Value {
    serde_json::from_slice(&read("expected.json")).expect("expected.json")
}
fn hex(h: u64) -> String {
    format!("{h:016x}")
}
fn bucket_hash(t: &Trace) -> String {
    let b: Vec<u8> = t.buckets.iter().flat_map(|&b| u16::try_from(b).expect("bucket fits u16").to_le_bytes()).collect();
    hex(fnv1a(&b))
}
fn symbol_hash(t: &Trace) -> String {
    let b: Vec<u8> = t.symbols.iter().flat_map(|s| s.to_le_bytes()).collect();
    hex(fnv1a(&b))
}

#[test]
fn the_committed_vector_decodes_to_its_pinned_image() {
    let (env, want) = (envelope(), expected());
    let img = decode(&model(), &env.payload).expect("decode");
    assert_eq!((img.width, img.height), (64, 64));
    assert_eq!(img.rgb.len(), 64 * 64 * 3);
    let got = hex(fnv1a(&img.rgb));
    assert_eq!(got, want["rgb_fnv1a"].as_str().unwrap(), "the export tool's integer reference");
    assert_eq!(got, RGB_FNV1A, "the pinned value");
}

/// The property the float capability lacks: the decoder derives, from its own
/// integer networks, exactly the bucket the encoder used for every element — and
/// reads exactly the encoder's symbol with it.
#[test]
fn the_decoder_derives_every_bucket_the_encoder_used() {
    let (env, want, m) = (envelope(), expected(), model());
    let (trace, img) = decode_traced(&m, &env.payload);
    img.expect("decode");
    assert_eq!(trace.buckets.len() as u64, want["y_symbols"].as_u64().unwrap());
    assert_eq!(bucket_hash(&trace), want["bucket_fnv1a"].as_str().unwrap());
    assert_eq!(bucket_hash(&trace), BUCKET_FNV1A);
    assert_eq!(symbol_hash(&trace), want["symbol_fnv1a"].as_str().unwrap());
    assert_eq!(symbol_hash(&trace), SYMBOL_FNV1A);
    // What the vector exercises: every scale bucket, and the rANS escape (a residual
    // outside its table, coded in bypass).
    let used: std::collections::BTreeSet<_> = trace.buckets.iter().copied().collect();
    assert_eq!(used.len() as u64, want["buckets_used"].as_u64().unwrap());
    assert_eq!(used.len(), m.tables.edges.len(), "every bucket of the scale table is used");
    let escapes = trace.buckets.iter().zip(&trace.symbols)
        .filter(|&(&b, &s)| { let v = s - m.tables.offsets[b]; v < 0 || v >= m.tables.cdf_lengths[b] as i32 - 2 })
        .count();
    assert_eq!(escapes as u64, want["y_escapes"].as_u64().unwrap());
    assert!(escapes > 0, "the vector exercises the escape path");
}

/// The float MLIC++ decode desynchronised when its encoder evaluated the entropy
/// network with a different thread count from the decoder's. Here the fast kernels
/// split the work across 1, 2, 3, 5 or all available threads, or the scalar reference
/// runs: every bucket, every symbol and every pixel is the same.
#[test]
fn thread_count_cannot_move_a_bucket() {
    let (env, m) = (envelope(), model());
    let run = |p: Policy| -> (Trace, IntMlicImage) {
        let (t, img) = with_policy(p, || decode_traced(&m, &env.payload));
        (t, img.expect("decode"))
    };
    let want = run(Policy::reference());
    assert_eq!(hex(fnv1a(&want.1.rgb)), RGB_FNV1A);
    for p in [Policy::fast(1), Policy::fast(2), Policy::fast(3), Policy::fast(5), Policy::fast(0)] {
        assert!(run(p) == want, "{p:?} differs from the reference kernels");
    }
}

fn bump(l: &mut IntLayer, i: usize, by: i64, bias: bool) {
    match l {
        IntLayer::Conv { weight, bias: b, .. } | IntLayer::ConvT { weight, bias: b, .. } => {
            if bias { b[i] += by } else { weight[i] += by }
        }
        _ => panic!("not a convolution"),
    }
}

/// The pinned image is not vacuous: each of four model parameters, one on each path
/// that feeds a latent value or a pixel — a weight of the hyper-decoder's first layer,
/// of the entropy network's shared layer, a bias of its mean head, and a weight of the
/// synthesis's first layer — changed by one integer unit, and the decode no longer
/// reproduces it (it decodes a different image, or refuses the stream once its symbols
/// stop making sense). Four chosen parameters, not every one: a weight that multiplies
/// only zeros, or a unit a later rounding absorbs, can leave the image as it was. The
/// scale head is different by design; see `a_scale_matters_only_through_its_bucket`.
#[test]
fn one_unit_on_each_value_path_changes_the_image() {
    let env = envelope();
    type Change = fn(&mut IntMlicModel);
    let cases: [(&str, Change); 4] = [
        ("h_s weight", |m| bump(&mut m.h_s.layers[0], 7, 1, false)),
        ("entropy weight", |m| bump(&mut m.entropy.layers[0], 1234, 1, false)),
        ("entropy mean bias", |m| bump(&mut m.entropy.layers[2], 8, 1, true)),
        ("g_s weight", |m| bump(&mut m.g_s.layers[0], 99, 1, false)),
    ];
    for (what, change) in cases {
        let mut m = model();
        change(&mut m);
        let same = decode(&m, &env.payload).is_ok_and(|img| hex(fnv1a(&img.rgb)) == RGB_FNV1A);
        assert!(!same, "{what}: a one-unit change left the pinned image unchanged");
    }
}

/// How the float decode failed, and why the integer one cannot: a scale selects a
/// distribution only through its bucket. Moved by one unit without crossing a table
/// edge, channel 0's scales change nothing — same buckets, same symbols, same image.
/// Moved across edges (by one whole unit of scale, `2^afrac`), the first element
/// whose bucket moves is where the decode leaves the encoder: its symbol is read
/// against another distribution, and so is everything after it. Two float decoders
/// that compute a scale on either side of an edge do exactly this (the float MLIC++
/// validation run diverged at symbol 416 of 12,288); two integer decoders compute
/// the same scale, so they cannot.
#[test]
fn a_scale_matters_only_through_its_bucket() {
    let env = envelope();
    let (truth, img) = decode_traced(&model(), &env.payload);
    img.expect("decode");
    let sigma_bias_ch0 = |by: i64| {
        let mut m = model();
        bump(&mut m.entropy.layers[2], 0, by, true);
        decode_traced(&m, &env.payload)
    };

    let (t, img) = sigma_bias_ch0(1);
    assert_eq!(t, truth, "one unit within the buckets: the same buckets and symbols");
    assert_eq!(hex(fnv1a(&img.expect("decode").rgb)), RGB_FNV1A);

    let (t, img) = sigma_bias_ch0(1 << model().afrac());
    let first = (0..truth.buckets.len())
        .find(|&i| t.buckets.get(i) != Some(&truth.buckets[i]) || t.symbols.get(i) != Some(&truth.symbols[i]))
        .expect("the stream leaves the encoder somewhere");
    assert!(first < t.buckets.len(), "the divergence is within what was decoded");
    assert_ne!(t.buckets[first], truth.buckets[first], "the first difference is a bucket, not a symbol");
    // Channel 0's anchors are the first 32 elements coded (slice 0, anchor pass).
    assert!(first < 32, "the first moved bucket is one of channel 0's, at element {first}");
    let outcome = match &img {
        Ok(img) => format!("decodes to {}", hex(fnv1a(&img.rgb))),
        Err(e) => format!("is refused ({e})"),
    };
    println!("scale +1.0 on channel 0: first moved bucket at element {first} ({} -> {}); the stream then {outcome}",
             truth.buckets[first], t.buckets[first]);
    assert!(img.map_or(true, |img| hex(fnv1a(&img.rgb)) != RGB_FNV1A), "the image is not the pinned one");
}

#[test]
fn the_envelope_pins_the_model_as_an_integer_exact_prior() {
    let (env, want) = (envelope(), expected());
    assert_eq!(env.capability(), "wai.neural.int_mlicpp");
    assert_eq!(env.manifest.intent, "replicate");
    let prior = env.prior().expect("a pinned prior");
    assert_eq!(prior.sha256, MODEL_SHA256);
    assert_eq!(prior.sha256, want["model_sha256"].as_str().unwrap());
    assert!(prior.is_cross_hardware_exact(), "an integer-exact contract, in no declared form");
    // SPEC §7: the registry class (NeuralIntegerExact) and the prior both back replicate.
    assert_eq!(env.replicate_refusal(), None);
    let h = parse_header(&env.payload).expect("header");
    assert_eq!((h.cz, h.zh, h.zw, h.cy, h.yh, h.yw, h.slices), (4, 2, 2, 8, 8, 8, 4));
    assert_eq!((h.width, h.height), (64, 64));
}

/// Damage to the committed payload — each header byte three ways, and random changes
/// in both streams — returns an image or a refusal, never a panic, in this
/// overflow-checked test build and in release alike.
#[test]
fn damaging_the_committed_payload_never_panics() {
    let (pl, m) = (envelope().payload, model());
    let mut state = 0x1A7E_4C10u64;
    let mut next = move || {
        state = state.wrapping_add(0x9E37_79B9_7F4A_7C15);
        let mut z = state;
        z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
        z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
        z ^ (z >> 31)
    };
    let (mut decoded, mut refused) = (0, 0);
    let mut tally = |r: Result<IntMlicImage, _>| if r.is_ok() { decoded += 1 } else { refused += 1 };
    for i in 0..Header::LEN {
        for x in [0x01u8, 0x80, 0xFF] {
            let mut p = pl.clone();
            p[i] ^= x;
            tally(decode(&m, &p));
        }
    }
    for _ in 0..200 {
        let mut p = pl.clone();
        for _ in 0..1 + next() % 4 {
            let i = Header::LEN + (next() as usize) % (p.len() - Header::LEN);
            p[i] = next() as u8;
        }
        tally(decode(&m, &p));
    }
    assert!(decoded > 0 && refused > 0, "{decoded} decoded, {refused} refused");
}

// ---- the parameter-set pin over the committed vector (SPEC §3.1, §7.1) ----

const LISTING_SHA256: &str = "85a05b0dac15566c5a657ddeee253eaa9f54ce852ae04924f5b04151901d9d1b";
const STREAM_SET_SHA256: &str = "8f23e85cd3d2690817c8935c2473dc2ba3af3cb302e3dca3cf3feb60f1c263d0";
const RGB_SHA256: &str = "45040d17b4178b91a0547b435e57a50fa1021101e290a94d1c57e6fbb4583a10";

fn scratch(tag: &str) -> PathBuf {
    let d = std::env::temp_dir().join(format!("wai-int-mlicpp-{tag}-{}", std::process::id()));
    std::fs::create_dir_all(&d).unwrap();
    d
}

/// The set-form twin of `stream.wai` (`tools/wai_int_mlicpp_set_form.py`): the same
/// payload, pinned through the listing of the one-file set. The pin verifies against
/// the committed listing and bundle, and the decode from the very bytes it verified
/// is the pinned image. A flipped listing is not held; a flipped bundle under an
/// intact listing is a part mismatch; a listing that names another file verifies, and
/// the decode then needs a file the set does not contain.
#[test]
fn set_form_pin_resolves_and_decodes_the_same_image() {
    use wai::pin::{sha256_hex, verify, Held, READABLE_FORMS};
    let set_bytes = read("stream_set.wai");
    assert_eq!(sha256_hex(&set_bytes), STREAM_SET_SHA256);
    assert_eq!(sha256_hex(&read("model.sha256")), LISTING_SHA256);
    let set = Wai::from_bytes(&set_bytes).unwrap();
    let legacy = envelope();
    assert_eq!(set.payload, legacy.payload, "only the pin differs");
    let pin = set.prior().expect("a pin");
    assert_eq!((pin.sha256.as_str(), pin.form()), (LISTING_SHA256, Some("wai.prior.set")));
    assert_eq!(set.replicate_refusal(), None, "a set-form pin declaring decode-equivalence backs replicate");

    let v = verify("wai.neural.int_mlicpp", pin, &[Held::Listing(dir().join("model.sha256"))], READABLE_FORMS)
        .expect("the committed listing and bundle verify");
    let m = IntMlicModel::from_bundle(v.file("model.wmb").unwrap()).unwrap();
    let img = decode(&m, &set.payload).unwrap();
    assert_eq!(hex(fnv1a(&img.rgb)), RGB_FNV1A);
    assert_eq!(sha256_hex(&img.rgb), RGB_SHA256);

    // Twins on temporary copies.
    let code = |d: &PathBuf| verify("wai.neural.int_mlicpp", pin, &[Held::Listing(d.join("model.sha256"))], READABLE_FORMS)
        .map(|_| "ok").unwrap_or_else(|e| e.code());
    let flipped_listing = scratch("flipped-listing");
    let mut l = read("model.sha256");
    l[0] ^= 1;
    std::fs::write(flipped_listing.join("model.sha256"), &l).unwrap();
    std::fs::write(flipped_listing.join("model.wmb"), read("model.wmb")).unwrap();
    assert_eq!(code(&flipped_listing), "prior_not_held");
    let flipped_bundle = scratch("flipped-bundle");
    let mut b = read("model.wmb");
    let at = b.len() - 2;
    b[at] ^= 1;
    std::fs::write(flipped_bundle.join("model.sha256"), read("model.sha256")).unwrap();
    std::fs::write(flipped_bundle.join("model.wmb"), &b).unwrap();
    assert_eq!(code(&flipped_bundle), "part_mismatch");
    let other = scratch("other-name");
    std::fs::write(other.join("other.wmb"), read("model.wmb")).unwrap();
    let other_listing = wai::pin::write_listing(&other, &["other.wmb".to_string()]).unwrap();
    std::fs::write(other.join("model.sha256"), &other_listing).unwrap();
    let mut other_pin = pin.clone();
    other_pin.sha256 = sha256_hex(&other_listing);
    let v = verify("wai.neural.int_mlicpp", &other_pin, &[Held::Listing(other.join("model.sha256"))], READABLE_FORMS)
        .expect("a listing of other.wmb verifies");
    assert_eq!(v.file("model.wmb").unwrap_err().code(), "set_incomplete");
    for d in [flipped_listing, flipped_bundle, other] {
        std::fs::remove_dir_all(d).ok();
    }
}

/// SPEC §7.1: one envelope, the same payload twice, pinned to two parameter sets — a
/// successor (which names the committed set in `supersedes`) and the committed set.
/// A sink holding only the committed bundle skips the first entry
/// (`digest_superseded`), picks the second, and decodes the pinned image.
#[test]
fn wai2_pins_select_the_held_model() {
    use wai::codecs::DeclaredTier;
    use wai::container::{Prior, Rendition, Supersedes, WaiMulti};
    use wai::pin::{verify, Held, READABLE_FORMS};
    let pl = envelope().payload;
    let successor = "5ccc".repeat(16);
    let m = WaiMulti::new("image", "replicate", serde_json::Value::Null, vec![
        Rendition::media("wai.neural.int_mlicpp", "int_mlicpp", pl.clone())
            .with_prior(Prior::new("next", successor, "int-rans/p16+int-mlicpp/i64")
                .with_tier(DeclaredTier::DecodeEquivalence)
                .with_supersedes(Supersedes::new(MODEL_SHA256))),
        Rendition::media("wai.neural.int_mlicpp", "int_mlicpp", pl)
            .with_prior(Prior::new("committed", MODEL_SHA256, "int-rans/p16+int-mlicpp/i64")
                .with_tier(DeclaredTier::DecodeEquivalence)),
    ]);
    let m = WaiMulti::from_bytes(&m.to_bytes().unwrap()).unwrap();
    let held = [Held::File(dir().join("model.wmb"))];
    let p = m.pick(|_| true, |cap, prior| verify(cap, prior, &held, READABLE_FORMS)).expect("a rendition");
    assert_eq!(p.index, 1);
    assert_eq!(p.skipped.iter().map(|(i, s)| (*i, s.code())).collect::<Vec<_>>(), [(0, "digest_superseded")]);
    let bundle = p.verified.as_ref().unwrap().one().unwrap();
    let img = decode(&IntMlicModel::from_bundle(bundle).unwrap(), p.payload).unwrap();
    assert_eq!(hex(fnv1a(&img.rgb)), RGB_FNV1A);
    assert_eq!(m.replicate_refusal(1), None);
}