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

//! End-to-end: a **real** byte-exact JPEG-AI decode → a reproducible-reconstruction
//! binding over the decoded pixels. The [`wai::reconstruction_binding`] unit tests cover
//! the binding on synthetic RGB; this closes the loop on real content — the byte-exact
//! decoder produces the pixels, and the C2PA/JPEG-Trust hard binding is taken over them.
//!
//! It exercises the property the binding rests on: the exact decoder is deterministic,
//! so the binding taken over one decode verifies a second decode, and a single flipped
//! pixel breaks it. Both decodes run in this process; that the same pixels come out on
//! another machine rests on the conditions and measurements in
//! `extensions/determinism-tiers.md` §5–§6. This test compares with other machines
//! only through `WAI_JPEGAI_EXPECT_SHA256` (below). (`wai.image.jpegai` is registered
//! at entropy-consistency; the pinned reconstruction is a property of the exact
//! decoder.)
//!
//! env: WAI_JPEGAI_CODESTREAM, WAI_JPEGAI_WEIGHTS_DIR, WAI_JPEGAI_NORM
//!   (the bundle from `tools/wai_jpegai_fetch_e2e_inputs.sh`; see jpegai_e2e_full).
//!      WAI_JPEGAI_EXPECT_SHA256 (optional) = the SHA-256 the decoded pixels must
//!   have. Set, it also compares this machine's decode with a value decoded on
//!   other machines, which the two in-process decodes cannot do. That value rests
//!   on a measurement, not on the registry: the codec registry classes
//!   `wai.image.jpegai` as `NeuralEntropyExact`, a conservative class that
//!   guarantees only the entropy decode (the decoded symbols and their entropy
//!   parameters) to be byte-identical across machines, not the reconstruction.
//!   The exact decoder's pinned output for the `jpegai/v1` inputs was measured
//!   identical on four platforms: macOS arm64, wasm32-wasip2, Linux arm64, and
//!   Linux x86_64 under user-mode emulation (`extensions/determinism-tiers.md`
//!   §6). `wai-jpegai-e2e.yml` pins that value, so a mismatch there is a departure
//!   from that measurement.
//!
//! The inputs are not in the repository; see jpegai_e2e_full for where they come
//! from and how CI fetches them. Without the three variables this test skips
//! loudly; with `WAI_REQUIRE_SOTA_VECTORS=1` it fails instead; a variable naming a
//! missing path, or a weights directory missing a file, fails either way.
#![cfg(all(feature = "neural_int", feature = "provenance"))]

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

use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_icci::IcciNet;
use wai::jpegai_synthesis::{IntNet, Weights};
use wai::reconstruction_binding::{ReconstructionBinding, ReconstructionDescriptor};

fn jload(p: &std::path::Path) -> serde_json::Value {
    serde_json::from_slice(&std::fs::read(p).unwrap()).unwrap()
}

#[test]
fn real_decode_to_reconstruction_binding() {
    let Some([cs, dir, nf]) = require_vectors::env_paths(
        "jpegai_reconstruction_binding::real_decode_to_reconstruction_binding",
        ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"],
        "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)",
    ) else {
        return;
    };
    let data = std::fs::read(&cs).unwrap_or_else(|e| panic!("read {cs}: {e}"));

    // --- load the model bundle (same shape as jpegai_e2e_full) ---
    let dpath = std::path::Path::new(&dir);
    let w = |n: &str| Weights::from_json(&jload(&dpath.join(n))).unwrap();
    let net = |n: &str| IntNet::from_json(&jload(&dpath.join(n))).unwrap();
    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: IcciNet::from_json(&jload(&dpath.join("icci.json"))).unwrap(),
    };
    let nv = jload(std::path::Path::new(&nf));
    let i64v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect::<Vec<i64>>();
    let i32v = |k: &str| nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap() as i32).collect::<Vec<i32>>();
    let i3 = |k: &str| { let a = nv[k].as_array().unwrap(); [a[0].as_i64().unwrap(), a[1].as_i64().unwrap(), a[2].as_i64().unwrap()] };
    let f3 = |k: &str| { let a = nv[k].as_array().unwrap(); [a[0].as_f64().unwrap() as f32, a[1].as_f64().unwrap() as f32, a[2].as_f64().unwrap() as f32] };
    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: i64v("bounds"), per_channel: nv["per_channel"].as_u64().unwrap() as usize,
        cdfs_y: i64v("cdfs_y"), cdfs_uv: i64v("cdfs_uv"), gain_y: i64v("gain_y"), gain_uv: i64v("gain_uv"),
        lsbs_t0: i32v("lsbs_t0"), lsbs_t1: i32v("lsbs_t1"), lef_thr: i3("lef_thr0"), lef_mag: f3("lef_mag0"),
    };

    // --- decode the real codestream → RGB, then bind over the decoded pixels ---
    let img = decode(&data, &weights, &norm).expect("supported geometry");
    assert_eq!((img.h, img.w), (256, 256));
    assert_eq!(img.data.len(), 256 * 256 * 3);

    // The descriptor names the synthesis branch the decode actually ran. The exact
    // decoder implements only the high-operation-point synthesis
    // (`synthesis_transform_id` 2) and refuses codestreams that select another.
    let segs = wai::jpegai_bitstream::parse_container(&data);
    let pih = wai::jpegai_bitstream::find_segment(&segs, wai::jpegai_bitstream::Marker::Pih).unwrap();
    let ph = wai::jpegai_bitstream::parse_picture_header(&data, pih.start_bit());
    assert_eq!(ph.synthesis_transform_ids[0], 2, "test256 selects the high-operation-point synthesis");
    let desc = ReconstructionDescriptor {
        decoder_id: ph.synthesis_transform_ids[0] as u8,
        profile: "tools_0".into(),
        level: "L1".into(),
        width: img.w as u32,
        height: img.h as u32,
        format: "RGB8".into(),
        color_space: "bt709".into(),
    };
    let binding = ReconstructionBinding::bind(&img.data, desc.clone());
    assert_eq!(binding.alg, "sha256");
    assert_eq!(binding.hash_hex().len(), 64);

    // The binding verifies its own decoded pixels.
    assert!(binding.verify(&img.data), "the binding must verify the pixels it was taken over");

    // A second decode in this process is byte-identical, and the binding — taken over
    // the first — verifies the second.
    let img2 = decode(&data, &weights, &norm).expect("supported geometry");
    assert_eq!(img2.data, img.data, "the exact decoder is deterministic within a process");
    assert!(binding.verify(&img2.data), "the binding verifies any faithful reproduction");

    // A single flipped pixel breaks the hard binding.
    let mut tampered = img.data.clone();
    tampered[12_345] ^= 1;
    assert!(!binding.verify(&tampered), "a single flipped byte must fail the binding");

    eprintln!("reconstruction binding over the real JPEG-AI decode:");
    eprintln!("  assertion  = {}", ReconstructionBinding::C2PA_LABEL);
    eprintln!("  sha256     = {}", binding.hash_hex());
    eprintln!("  json       = {}", binding.to_json());

    // Against other machines: the pixels this machine decoded must hash to the value
    // measured from other machines' decodes of the same inputs (see the header).
    if let Ok(want) = std::env::var("WAI_JPEGAI_EXPECT_SHA256") {
        assert_eq!(
            binding.hash_hex(),
            want.trim().to_ascii_lowercase(),
            "the decoded pixels differ from the pinned decode of these inputs"
        );
        eprintln!("  matches the pinned WAI_JPEGAI_EXPECT_SHA256");
    }
}