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

//! The JPEG-AI exact decoder's output, pinned byte for byte, so that a platform that
//! decodes differently fails.
//!
//! [`wai::jpegai_decode::decode`] is meant to give the same bytes on every platform:
//! its networks are integer, and its float stages use only IEEE-754 basic operations
//! in a fixed order, plus in-crate `exp` and rounding built from those operations
//! (`wai::float_repro`), so no platform math library is on the path. This test is the
//! check. It decodes the pinned `jpegai/v1` inputs and compares SHA-256 hashes, pinned
//! from a measured decode, of:
//!
//! - the float inputs as parsed (the eICCI weights and the LEF magnitudes), because
//!   the JSON-to-float conversion is part of the path;
//! - the float values that [`wai::jpegai_decode::decode_with_stages`] reports after
//!   each stage, bit for bit. A final byte can hide a float difference smaller than
//!   its rounding step, so these are checked as well. Values inside a stage are not
//!   reported, so a difference there that a later stage absorbs goes unseen;
//! - the decoded RGB8 bytes. This is the same value a reconstruction binding over the
//!   decode records (SHA-256 of the pixels).
//!
//! On a mismatch it prints every hash, so the first diverging stage names where two
//! platforms part. The pinned values were measured on the platforms listed at
//! [`PINNED`]; a platform not listed there has not been measured.
//!
//! env: WAI_JPEGAI_CODESTREAM = test256.bin
//!      WAI_JPEGAI_WEIGHTS_DIR = dir of {hs,hyp}_{y,uv}.json, mcm{0..3}_y.json,
//!                               synth_{y,uv}.json, icci.json
//!      WAI_JPEGAI_NORM        = norm.json
//! The inputs are the `jpegai/v1` bundle of WAI's models root; each file is checked
//! against [`INPUTS`] before decoding, because the pinned outputs hold only for those
//! bytes. Without the variables the test skips and says so on stderr; with
//! `WAI_REQUIRE_SOTA_VECTORS=1` a missing input fails it instead.
#![cfg(all(feature = "neural_int", feature = "provenance"))]

use sha2::{Digest, Sha256};
use std::io::Write;
use std::path::Path;
use wai::jpegai_decode::{decode_with_stages, DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_icci::{Conv, IcciNet, IcciPath};
use wai::jpegai_synthesis::{IntNet, Weights};

/// SHA-256 of each `jpegai/v1` input file.
const INPUTS: [(&str, &str); 13] = [
    ("hs_uv.json", "f3b325c583c6b7613ccddc2d89bfaa50e0e8314ceecd8d82fc4df7fb159b6083"),
    ("hs_y.json", "1d193f15b5d8d1531ef9ffef8676904d76cfedbd1d2df931d581dfb26daa2688"),
    ("hyp_uv.json", "2c5f6f12fcf9a36489553e0e02e95684c339d82bab64cb7bae76a4635f9ec56d"),
    ("hyp_y.json", "32ab3a153e33ded902b1d8a6c413047634bc39e8b74525dfbb40eded0763433d"),
    ("icci.json", "bd3aa0c23294a462895416eda38399468781163858229c386deaba58def5a1c3"),
    ("mcm0_y.json", "562740aba18ba4dccdc007b3d54a669ecd18feac9d0bd1fd7a29e66a0ba498ee"),
    ("mcm1_y.json", "fbd4b7c05e81cd32df4db4e2f6dfa4d6e3fc2edb6c21150f80ef5ac04765228d"),
    ("mcm2_y.json", "b42d6fa96bf865dd2db47955aba14d8d65b9152c941be3ee3e92eae5643834a7"),
    ("mcm3_y.json", "e06766d0a7bff16188a0e96297d6e498e28fbaf7850f68b0a29447a9391c9d9d"),
    ("norm.json", "5353b76f25e59f52d85adc896c9dc6b4aa8552589e3e4e1d97b8f29692e37b13"),
    ("synth_uv.json", "8ddf8fe9dd15569910ad5755a5495526df959835bcf86a259d06ab6445f02b4e"),
    ("synth_y.json", "4f2c9ae807ea22197adaae0263c9d6d4d53ef7d4fccb81848b7bb4da9ef0ea69"),
    ("test256.bin", "99840e5f472423f5a1261c9ec219d12dc910e9623d752672555dfea69ea4a097"),
];

/// SHA-256 of each checked value, in decode order: `float_inputs`, then one entry
/// per stage ([`wai::jpegai_decode::Stage::name`]) over the f32 bit patterns
/// (little-endian, planes Y, U, V), then `rgb` over the decoded RGB8 bytes.
///
/// Measured identical on macOS arm64 (native), wasm32-wasip2 (a WASI runtime), Linux
/// arm64 (glibc) and Linux x86_64 (glibc, under user-mode emulation);
/// `extensions/determinism-tiers.md` §6 records the toolchain and the method for each,
/// and which state of the decode's source it ran.
const PINNED: [(&str, &str); 9] = [
    ("float_inputs", "693740148c1269adbf058339deec3ad3aefa8dc4f4c2f0777d0778f770625af6"),
    ("latent_y", "5889f3806308496c64850aee04f3706178f92d1e0edcda1e6975fe256779061e"),
    ("latent_uv", "be89a76f726ef362f267aba3dee8f61594cc512576a3864f52f665cc5b6ab77b"),
    ("synthesis", "d56695726bb6ebed403199f60a79a5698734eef1e5ca2813859d490c0882d2a4"),
    ("efl", "1783ea77bd4f6b0ac6a599137577a18db2a86753aa3eb3d79ca987e1cf25792e"),
    ("icci", "3566c51c4042a749d7826df812663c4e3ff7a9e807cbd0c3bb0545c9dbd79f2b"),
    ("efn", "09dc833c7a6b91ce08da20adb8bbc55e2eeedba7d8a2dc4b133ab0915b267de7"),
    ("lef", "e2b20c72ac6a59dd12067f75562c82715061f6d7a3444ca075fc0b11e3f56468"),
    ("rgb", "dbc494a5a670914edf55c07c48afc0dd348cd05d3a82fff293bf384b2bb69ea3"),
];

fn required() -> bool {
    std::env::var("WAI_REQUIRE_SOTA_VECTORS")
        .map(|v| !matches!(v.trim().to_ascii_lowercase().as_str(), "" | "0" | "false" | "no"))
        .unwrap_or(false)
}

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

fn sha_f32(planes: &[&[f32]]) -> String {
    let mut h = Sha256::new();
    for p in planes {
        for v in p.iter() {
            h.update(v.to_bits().to_le_bytes());
        }
    }
    hex(&h.finalize())
}

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_or_else(|e| panic!("parse {}: {e}", p.display()))
}

fn conv_floats(c: &Conv, out: &mut Vec<f32>) {
    out.extend_from_slice(&c.w);
    out.extend_from_slice(&c.b);
}

fn path_floats(p: &IcciPath, out: &mut Vec<f32>) {
    conv_floats(&p.conv_first, out);
    for b in &p.blocks {
        conv_floats(&b.conv1, out);
        conv_floats(&b.conv2, out);
        out.push(b.scale);
    }
    conv_floats(&p.conv_last, out);
    out.push(p.scale);
}

#[test]
fn decode_matches_pinned_hashes() {
    let vars = ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"];
    let got: Vec<Option<String>> = vars.iter().map(|v| std::env::var(v).ok()).collect();
    if got.iter().any(Option::is_none) {
        let msg = format!(
            "jpegai_cross_platform: NOT VERIFIED - set {} (the jpegai/v1 inputs)",
            vars.join(" + ")
        );
        if required() {
            panic!("{msg}; WAI_REQUIRE_SOTA_VECTORS is set, so a missing input is a failure");
        }
        // Past libtest's capture, so the skip shows in an ordinary run.
        let _ = writeln!(std::io::stderr(), "\n!!! SKIPPED - {msg}\n");
        return;
    }
    let [cs, dir, nf] = [0, 1, 2].map(|i| got[i].clone().unwrap());
    let dpath = Path::new(&dir);

    // The pinned outputs hold only for these exact input bytes.
    for (name, want) in INPUTS {
        let p = match name {
            "test256.bin" => Path::new(&cs).to_path_buf(),
            "norm.json" => Path::new(&nf).to_path_buf(),
            _ => dpath.join(name),
        };
        let bytes = std::fs::read(&p).unwrap_or_else(|e| panic!("read {}: {e}", p.display()));
        assert_eq!(hex(&Sha256::digest(&bytes)), want, "{} is not the pinned jpegai/v1 {name}", p.display());
    }

    let w = |n: &str| Weights::from_json(&jload(&dpath.join(n)));
    let net = |n: &str| IntNet::from_json(&jload(&dpath.join(n)));
    let icci = IcciNet::from_json(&jload(&dpath.join("icci.json")));
    let nv = jload(Path::new(&nf));
    let i64v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
    let u8v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u8).collect::<Vec<u8>>();
    let i32v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::<Vec<i32>>();
    let a3 = |k: &str| nv[k].as_array().unwrap().clone();
    let (thr, mag) = (a3("lef_thr0"), a3("lef_mag0"));
    let norm = NormTables {
        pmf: nv["pmf"].as_array().unwrap().iter()
            .map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
        bounds: u8v("bounds"),
        per_channel: nv["per_channel"].as_u64().unwrap() as usize,
        cdfs_y: u8v("cdfs_y"), cdfs_uv: u8v("cdfs_uv"),
        gain_y: i64v("gain_y"), gain_uv: i64v("gain_uv"),
        lsbs_t0: i32v("lsbs_t0"), lsbs_t1: i32v("lsbs_t1"),
        lef_thr: [0, 1, 2].map(|i| thr[i].as_i64().unwrap()),
        lef_mag: [0, 1, 2].map(|i| mag[i].as_f64().unwrap() as f32),
    };

    let mut float_inputs = Vec::new();
    path_floats(&icci.luma, &mut float_inputs);
    path_floats(&icci.chroma, &mut float_inputs);
    conv_floats(&icci.conv_last_v, &mut float_inputs);
    float_inputs.push(icci.scale_v);
    float_inputs.extend_from_slice(&norm.lef_mag);
    let mut hashes: Vec<(String, String)> = vec![("float_inputs".into(), sha_f32(&[&float_inputs]))];

    let weights = DecoderWeights {
        luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
        chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
        mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
        icci,
    };
    let data = std::fs::read(&cs).unwrap();
    let img = decode_with_stages(&data, &weights, &norm, &mut |stage, planes| {
        hashes.push((stage.name().into(), sha_f32(planes)));
    })
    .expect("the jpegai/v1 codestream is in scope and the FP environment is the default");
    assert_eq!((img.w, img.h, img.data.len()), (256, 256, 256 * 256 * 3));
    hashes.push(("rgb".into(), hex(&Sha256::digest(&img.data))));

    let refrgb = u8v("ref_rgb");
    let (maxd, off) = img.data.iter().zip(&refrgb).fold((0i32, 0usize), |(m, o), (g, r)| {
        let d = (*g as i32 - *r as i32).abs();
        (m.max(d), o + (d != 0) as usize)
    });

    let env = if cfg!(target_env = "gnu") { "-gnu" } else if cfg!(target_env = "musl") { "-musl" }
        else if cfg!(target_env = "p2") { "-p2" } else { "" };
    let os = if cfg!(target_os = "wasi") { "wasi" } else { std::env::consts::OS };
    let platform = format!("{}-{os}{env}", std::env::consts::ARCH);
    eprintln!("jpegai_cross_platform on {platform}: vs reference max|Δ| = {maxd}, {off}/{} samples differ",
              refrgb.len());
    let mut first_diff = None;
    for (i, (name, got)) in hashes.iter().enumerate() {
        let want = PINNED.iter().find(|(n, _)| n == name).map(|(_, h)| *h);
        let ok = want == Some(got.as_str());
        if !ok && first_diff.is_none() {
            first_diff = Some(name.clone());
        }
        eprintln!("  {:<2} {:<13} {got} {}", i, name, if ok { "ok" } else { "DIFFERS" });
    }
    let names: Vec<&str> = hashes.iter().map(|(n, _)| n.as_str()).collect();
    let pinned_names: Vec<&str> = PINNED.iter().map(|(n, _)| *n).collect();
    assert_eq!(names, pinned_names, "the decode reported a different set of stages than was pinned");
    assert!(first_diff.is_none(),
        "the decode on {platform} differs from the pinned decode, first at `{}`",
        first_diff.unwrap_or_default());
    eprintln!("  all {} hashes match the pinned decode", hashes.len());
}