open-standards/wai/wai-web/src/neural_int.rs

//! In-browser **integer-exact** neural decode — zero JS ML runtime.
//!
//! The rest of this crate is parse-and-dispatch: it names a capability and
//! hands the payload to whatever backend the deployer registered (ONNX,
//! WebGPU, transformers.js). This module is the exception the integer-entropy
//! contract earns. `int_synth` (integer synthesis) and `int_entropy` (the
//! deterministic scale→bucket step) are pure i64 — no `ort`, no float — so the
//! SAME Rust that decodes natively compiles straight to wasm32 and runs in the
//! browser with **no ML runtime at all**. The bytes it produces are identical
//! to the native decode on every machine: a learned-codec reconstruction you
//! can verify in-page against a receipt.
//!
//! `decodeIntSynth(payload, specJson)` → RGB. The host supplies the model spec
//! bytes (the `int_synth.json` weights); WAI carries the payload. The
//! `int_entropy` primitives are exposed too, so a future conditional decode can
//! run its bucketing in-browser deterministically.

use wasm_bindgen::prelude::*;

/// An integer-exact decoded image: RGB row-major, `width * height * 3` bytes.
#[wasm_bindgen]
pub struct DecodedImage {
    rgb: Vec<u8>,
    width: u32,
    height: u32,
}

#[wasm_bindgen]
impl DecodedImage {
    /// RGB bytes, row-major, length `width * height * 3`. Copy into an
    /// `ImageData` (after widening to RGBA) to paint a canvas.
    #[wasm_bindgen(getter)]
    pub fn rgb(&self) -> Box<[u8]> {
        self.rgb.clone().into_boxed_slice()
    }
    #[wasm_bindgen(getter)]
    pub fn width(&self) -> u32 {
        self.width
    }
    #[wasm_bindgen(getter)]
    pub fn height(&self) -> u32 {
        self.height
    }
}

/// Decode a `wai.neural.int_synth` payload to RGB in the browser, integer-exact.
/// `payload` is the envelope's payload bytes; `spec_json` is the model spec (the
/// `int_synth.json` weights) the host has installed. No float, no ONNX — the
/// output is byte-identical to the native reference decode. It runs
/// `wai::int_payload::decode_int`, as the other payload entry points do, so a set it
/// has loaded is kept for the next decode (`loadedSetStats`), and a refusal's message
/// starts with its code (`payload_malformed`).
#[wasm_bindgen(js_name = decodeIntSynth)]
pub fn decode_int_synth(payload: &[u8], spec_json: &[u8]) -> Result<DecodedImage, JsError> {
    match crate::int_payload::decode("wai.neural.int_synth", payload, &[("int_synth.json", spec_json)]).map_err(|e| JsError::new(&e))? {
        wai::int_payload::IntDecoded::Image { rgb, width, height } => Ok(DecodedImage { rgb, width, height }),
        _ => Err(JsError::new("unsupported_capability: not an image")),
    }
}

/// Decode a `WIH1` payload (`wai.neural.int_hyper`, `extensions/integer-payloads.md`
/// §4) → RGB in the browser, byte-identical to the native reference, with the
/// conditional codec's set: `g_s.whm` and `h_s.whm` (the compact `WHM1` networks) and
/// `entropy.json` (the tables). A refusal's message starts with its code
/// (`bad_magic` for a JSON form). No float, no ONNX.
#[wasm_bindgen(js_name = decodeWih1)]
pub fn decode_wih1(payload: &[u8], g_s: &[u8], h_s: &[u8], entropy: &[u8]) -> Result<DecodedImage, JsError> {
    let set = crate::int_payload::int_hyper_set(g_s, h_s, entropy);
    match crate::int_payload::decode("wai.neural.int_hyper", payload, &set).map_err(|e| JsError::new(&e))? {
        wai::int_payload::IntDecoded::Image { rgb, width, height } => Ok(DecodedImage { rgb, width, height }),
        _ => Err(JsError::new("unsupported_capability: not an image")),
    }
}

/// Decode the tool's JSON conditional bitstream → RGB in the browser. For debugging
/// and conversion only: the payload `wai.neural.int_hyper` carries is `WIH1`
/// ([`decode_wih1`]). The rest of this note describes the decode both share:
/// integer-exact (`wai.neural.int_hyper`). The host installs the model
/// (`int_synth_sota_model.json` = g_s, `int_hs_model.json` = h_s,
/// `cond_entropy.json` = the entropy tables); WAI carries the bitstream
/// (`cond_bitstream.json`). The WHOLE decode — factorized z, integer h_s,
/// integer-σ entropy, integer g_s — is pure i64, no float, no ONNX, so the RGB
/// is byte-identical to the native reference on every machine: a pretrained
/// learned codec decoding deterministically in-page, zero JS ML runtime.
#[wasm_bindgen(js_name = decodeIntHyper)]
pub fn decode_int_hyper(
    g_s_json: &[u8], h_s_json: &[u8], entropy_json: &[u8], bitstream_json: &[u8],
) -> Result<DecodedImage, JsError> {
    use wai::int_hyper_synth::IntModel;
    let parse = |b: &[u8]| -> Result<serde_json::Value, JsError> {
        serde_json::from_slice(b).map_err(|e| JsError::new(&e.to_string()))
    };
    let net = |b: &[u8]| IntModel::try_from_json(&parse(b)?).map_err(|k| JsError::new(&format!("model_mismatch: a model's `{k}`")));
    cond_decode(net(g_s_json)?, net(h_s_json)?, entropy_json, bitstream_json)
}

/// Same decode from the **compact binary** model (`*_model.bin`, int12 weights as
/// `i16`), of the tool's JSON bitstream. For debugging and conversion only: the
/// payload is `WIH1` ([`decode_wih1`]). Byte-identical to the JSON model and to the
/// native reference.
#[wasm_bindgen(js_name = decodeIntHyperBin)]
pub fn decode_int_hyper_bin(
    g_s_bin: &[u8], h_s_bin: &[u8], entropy_json: &[u8], bitstream_json: &[u8],
) -> Result<DecodedImage, JsError> {
    use wai::int_hyper_synth::IntModel;
    let g_s = IntModel::from_bytes(g_s_bin).map_err(JsError::new)?;
    let h_s = IntModel::from_bytes(h_s_bin).map_err(JsError::new)?;
    cond_decode(g_s, h_s, entropy_json, bitstream_json)
}

fn cond_decode(
    g_s: wai::int_hyper_synth::IntModel, h_s: wai::int_hyper_synth::IntModel,
    entropy_json: &[u8], bitstream_json: &[u8],
) -> Result<DecodedImage, JsError> {
    let (rgb, height, width) = cond_decode_checked(&g_s, &h_s, entropy_json, bitstream_json).map_err(|e| JsError::new(&e))?;
    Ok(DecodedImage { rgb, width: width as u32, height: height as u32 })
}

/// The debug JSON decode, through the checked readers (`CondEntropy::try_from_value`,
/// `CondBitstream::try_from_value`) and `decode_conditional_checked`, so a JSON the
/// export tool did not write is refused, never a panic: the pure-Rust half the native
/// tests run.
pub(crate) fn cond_decode_checked(
    g_s: &wai::int_hyper_synth::IntModel, h_s: &wai::int_hyper_synth::IntModel, entropy_json: &[u8], bitstream_json: &[u8],
) -> Result<(Vec<u8>, usize, usize), String> {
    use wai::int_hyper_synth::{decode_conditional_checked, CondBitstream, CondEntropy};
    let parse = |b: &[u8]| serde_json::from_slice::<serde_json::Value>(b).map_err(|e| e.to_string());
    let ent = CondEntropy::try_from_value(&parse(entropy_json)?).map_err(|e| format!("model_mismatch: {e}"))?;
    let bs = CondBitstream::try_from_value(&parse(bitstream_json)?).map_err(|k| format!("field_range: `{k}`"))?;
    decode_conditional_checked(g_s, h_s, &ent, &bs.stream()).map_err(|r| r.to_string())
}

#[cfg(test)]
mod tests {
    use serde_json::json;
    use wai::int_hyper_synth::IntModel;

    fn read(rel: &str) -> Vec<u8> {
        let p = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../wai-rs/tests/vectors").join(rel);
        std::fs::read(&p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))
    }

    /// The debug JSON decode refuses, rather than panics on, a JSON the export tool
    /// did not write: an empty model or bitstream, a hyper-latent of 2^61 elements,
    /// more hyper-latent channels than tables. The committed tiny codec decodes.
    #[test]
    fn the_debug_json_decode_refuses_instead_of_panicking() {
        assert!(IntModel::try_from_json(&json!({})).is_err());
        let (g, h) = (IntModel::from_bytes(&read("int_hyper/tiny_g_s.bin")).unwrap(), IntModel::from_bytes(&read("int_hyper/tiny_h_s.bin")).unwrap());
        let ent = read("int_hyper/tiny_entropy.json");
        let got = super::cond_decode_checked(&g, &h, &ent, &read("int_hyper/tiny_bitstream.json")).unwrap();
        assert_eq!((got.1, got.2), (16, 16));
        // The JSON models read through the checked reader decode alike.
        let net = |rel: &str| IntModel::try_from_json(&serde_json::from_slice(&read(rel)).unwrap()).unwrap();
        let from_json = super::cond_decode_checked(&net("int_hyper/tiny_g_s.json"), &net("int_hyper/tiny_h_s.json"), &ent, &read("int_hyper/tiny_bitstream.json")).unwrap();
        assert_eq!(from_json, got);
        let bs = |v: serde_json::Value| serde_json::to_vec(&v).unwrap();
        let zeros = vec![0u8; 8];
        for (what, b) in [
            ("{}", bs(json!({}))),
            ("n_z·zh·zw = 2^61", bs(json!({"z_rans": zeros, "rans": zeros, "n_z": 1u64 << 20, "zh": 1u64 << 20, "zw": 1u64 << 21, "m": 1, "sh": 1, "sw": 1}))),
            ("more z channels than tables", bs(json!({"z_rans": zeros, "rans": zeros, "n_z": 4096, "zh": 1, "zw": 1, "m": 1, "sh": 1, "sw": 1}))),
        ] {
            let r = std::panic::catch_unwind(|| super::cond_decode_checked(&g, &h, &ent, &b));
            assert!(matches!(r, Ok(Err(_))), "{what}: {r:?}");
        }
    }
}

// ---- int_entropy: the deterministic scale→bucket primitive, in the browser ----

/// Precompute a CompressAI-style ascending `scale_table` as fixed-point integer
/// bucket edges (Q`frac`). Returns a `BigInt64Array` to JS.
#[wasm_bindgen(js_name = fixedScaleEdges)]
pub fn fixed_scale_edges(scale_table: &[f32], frac: u32) -> Vec<i64> {
    wai::int_entropy::fixed_scale_edges(scale_table, frac)
}

/// Snap a float σ to the fixed-point grid (Q`frac`), deterministic floor.
#[wasm_bindgen(js_name = quantizeScale)]
pub fn quantize_scale(sigma: f32, frac: u32) -> i64 {
    wai::int_entropy::quantize_scale(sigma, frac)
}

/// Integer `build_index`: bucket a fixed-point σ to a scale-table index — pure
/// integer, identical on every machine given the same `sigma_q`. `edges_q` from
/// [`fixed_scale_edges`].
#[wasm_bindgen(js_name = detBuildIndex)]
pub fn det_build_index(sigma_q: i64, edges_q: &[i64]) -> usize {
    wai::int_entropy::det_build_index(sigma_q, edges_q)
}

/// Lowercase-hex BLAKE3 of `bytes` — the in-page receipt check. Hash the RGB of
/// an integer-exact decode and compare it to the digest a publisher pinned: the
/// reconstruction either matches on this machine or it does not, verified
/// client-side with no server round-trip. This is the whole point of an
/// integer-exact decode — the hash is portable.
#[wasm_bindgen(js_name = blake3Hex)]
pub fn blake3_hex(bytes: &[u8]) -> String {
    blake3::hash(bytes).to_hex().to_string()
}

// ---- JPEG-AI byte-exact mANS: the ISO entropy decode, in the browser ----

/// Decode the ISO/IEC 6048 JPEG-AI **mANS** entropy stream in the browser, byte-exact —
/// the same pure-i64 decoder native sinks run (`wai::jpegai_mans`), compiled to wasm32.
///
/// A **reproduction vector**: given the reference stream `mem`, the integer PMF rows
/// (`pmf_flat`, `n_rows × 256` row-major), the per-distribution `bounds`, and the
/// per-position `sigma` distribution indices + coded-position `masks` (0/1), it returns
/// the decoded residual symbols. The decode is entirely integer, so the symbols — and
/// their hash — are identical on every machine. The page hashes them and compares to the
/// digest the native reference produced; a byte-identical match across host + browser is
/// the whole point of the integer-entropy contract. This is the same stream and
/// symbols the cross-architecture byte-equality matrix checks natively.
#[wasm_bindgen(js_name = decodeJpegAiMans)]
pub fn decode_jpegai_mans(
    mem: &[u8],
    pmf_flat: &[i64],
    bounds: &[u8],
    sigma: &[u8],
    masks: &[u8],
) -> Vec<i16> {
    use wai::jpegai_mans::{build_decode_transitions, cdf_from_pmf, MansDecoder};
    const WIDTH: usize = 256;
    let rows = pmf_flat.len() / WIDTH;
    let pmf: Vec<Vec<i64>> =
        (0..rows).map(|r| pmf_flat[r * WIDTH..(r + 1) * WIDTH].to_vec()).collect();
    let flat: Vec<u32> = build_decode_transitions(&pmf, &cdf_from_pmf(&pmf))
        .iter()
        .flatten()
        .copied()
        .collect();
    let mut dec = MansDecoder::new(mem, &[mem.len() as u32]);
    dec.set_sgm_transitions(&flat, bounds);
    let masks_b: Vec<bool> = masks.iter().map(|&m| m != 0).collect();
    let mut out = vec![0i16; sigma.len()];
    dec.decode_sgm(sigma, &mut out, &masks_b);
    out
}

// ---- JPEG-AI full decode → RGB, in the browser (PTQ value/synthesis path) ----

/// Decode a full ISO/IEC 6048 JPEG-AI codestream to RGB in the browser
/// (`wai.image.jpegai`). The same exact decoder a native sink runs
/// (`wai::jpegai_decode::decode`), compiled to wasm32: an integer entropy decode and
/// integer PTQ networks, then dequantisation, post-filters and colour conversion in
/// IEEE-754 float with a fixed evaluation order. It pins one reconstruction, which
/// reproduces across machines under the conditions in
/// `extensions/determinism-tiers.md` §5; that is a property of this decoder, not of
/// the capability, which is registered at entropy-consistency and does not back
/// `intent = replicate`. It decodes one configuration and refuses any other
/// codestream with the reason (`wai::jpegai_decode::check_scope`). The 13 byte
/// inputs are the model weight bundle (`tools/wai_jpegai_build_e2e_inputs.py`):
/// per-plane hyper-scale / hyper-decoder / synthesis, the four luma MCM stages, the
/// eICCI net, and the normative tables (`norm`).
#[wasm_bindgen(js_name = decodeJpegAiFull)]
#[allow(clippy::too_many_arguments)]
pub fn decode_jpegai_full(
    codestream: &[u8],
    hs_y: &[u8], hyp_y: &[u8], synth_y: &[u8],
    hs_uv: &[u8], hyp_uv: &[u8], synth_uv: &[u8],
    mcm0: &[u8], mcm1: &[u8], mcm2: &[u8], mcm3: &[u8],
    icci: &[u8], norm: &[u8],
) -> Result<DecodedImage, JsError> {
    use wai::jpegai_decode::{decode, DecoderWeights, NormTables, PlaneWeights};
    use wai::jpegai_icci::IcciNet;
    use wai::jpegai_synthesis::{IntNet, Weights};

    let parse = |b: &[u8]| -> Result<serde_json::Value, JsError> {
        serde_json::from_slice(b).map_err(|e| JsError::new(&e.to_string()))
    };
    let w = |b: &[u8]| -> Result<Weights, JsError> { Ok(Weights::from_json(&parse(b)?)) };
    let net = |b: &[u8]| -> Result<IntNet, JsError> { Ok(IntNet::from_json(&parse(b)?)) };

    let weights = DecoderWeights {
        luma: PlaneWeights { hyper_scale: net(hs_y)?, hyper_decoder: w(hyp_y)?, synthesis: w(synth_y)? },
        chroma: PlaneWeights { hyper_scale: net(hs_uv)?, hyper_decoder: w(hyp_uv)?, synthesis: w(synth_uv)? },
        mcm_y: [w(mcm0)?, w(mcm1)?, w(mcm2)?, w(mcm3)?],
        icci: IcciNet::from_json(&parse(icci)?),
    };

    let nv = parse(norm)?;
    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 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_t = NormTables {
        pmf: nv["pmf"].as_array().ok_or_else(|| JsError::new("norm.pmf"))?.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: i3("lef_thr0"), lef_mag: f3("lef_mag0"),
    };

    if let Err(why) = wai::jpegai_decode::check_scope(codestream) {
        return Err(JsError::new(&format!("wai.image.jpegai: codestream refused: {why}")));
    }
    let img = decode(codestream, &weights, &norm_t).ok_or_else(|| {
        JsError::new("wai.image.jpegai: decode refused (the weight bundle or tables do not fit the codestream)")
    })?;
    Ok(DecodedImage { rgb: img.data, width: img.w as u32, height: img.h as u32 })
}