diff --git a/.github/workflows/wai-web-wasm.yml b/.github/workflows/wai-web-wasm.yml index 6470617c8..cdc544145 100644 --- a/.github/workflows/wai-web-wasm.yml +++ b/.github/workflows/wai-web-wasm.yml @@ -13,6 +13,9 @@ name: wai-web wasm32 # - demo/verify_pack.mjs: the browser writers `pack` and `packMulti` refuse # every candidate the export lists, in every place an envelope can name one, # and a manifest that names a member twice in one object; +# - demo/verify_staged_wasm.mjs against the staged-delivery corpus +# wai/staged-conformance/: every case replayed through the wasm32 StageSink +# must give the reports, final units and stage hashes the native sink wrote; # - a byte comparison of the demo's staged int_synth copies with the corpus; # - the crate's wasm_tests under wasm-bindgen-test-runner, and its native unit # tests on the host. @@ -28,6 +31,7 @@ on: - 'wai/wai-rs/**' - 'wai/wai-quantum/**' - 'wai/neural-conformance/**' + - 'wai/staged-conformance/**' - 'wai/registry/**' - '.github/workflows/wai-web-wasm.yml' pull_request: @@ -36,6 +40,7 @@ on: - 'wai/wai-rs/**' - 'wai/wai-quantum/**' - 'wai/neural-conformance/**' + - 'wai/staged-conformance/**' - 'wai/registry/**' - '.github/workflows/wai-web-wasm.yml' @@ -91,7 +96,9 @@ jobs: # bindings, which wasm_tests does not run. They decode the committed integer # payloads to their sidecars' SHA-256 (int_payload::tests) and check that the # writers `pack` and `packMulti` refuse a payload that is not the format SPEC §5 - # pins its integer capability to, and a manifest that names a member twice (tests). The count check turns a test that + # pins its integer capability to, and a manifest that names a member twice (tests); the + # staged-delivery binding's pure half replays the staged corpus and reads a layer's + # header and cost (stage::tests). The count check turns a test that # stopped being compiled into a failure instead of a silent "0 passed". - name: wai-web native unit tests shell: bash @@ -99,7 +106,7 @@ jobs: run: | set -euo pipefail cargo test --locked --lib 2>&1 | tee "$RUNNER_TEMP/wai-web-native.log" - grep -q "test result: ok. 9 passed; 0 failed" "$RUNNER_TEMP/wai-web-native.log" + grep -q "test result: ok. 14 passed; 0 failed" "$RUNNER_TEMP/wai-web-native.log" - name: wasm32 int_synth decode equals the frozen native corpus shell: bash working-directory: wai/wai-web @@ -110,6 +117,17 @@ jobs: shell: bash working-directory: wai/wai-web run: node demo/verify_int_payload_wasm.mjs ./pkg-node/wai_web.js ../wai-rs/tests/vectors + # Every staged-delivery case replayed through the wasm32 StageSink: each integer + # base decoded in wasm32 to its stage0.bin with its pin, and every report, final + # unit, stage hash and clip hash equal to what the native sink wrote in case.json. + # Then the binding's guards: bases that declare huge clips refused under a cost + # limit with linear memory checked, a base past the chain budget, base facts a + # page cannot raise (fallback, tier, claimed), addBase spending its base, and a + # seeded fuzz that must never trap. + - name: wasm32 staged-delivery sink equals the native corpus + shell: bash + working-directory: wai/wai-web + run: node demo/verify_staged_wasm.mjs ./pkg-node/wai_web.js ../staged-conformance - name: The browser menu is registered and names no candidate shell: bash working-directory: wai/wai-web diff --git a/wai/CHANGELOG.md b/wai/CHANGELOG.md index c7bea8e52..9c4a4881d 100644 --- a/wai/CHANGELOG.md +++ b/wai/CHANGELOG.md @@ -6,6 +6,82 @@ wire format may change and such changes are recorded here. ## v0.4.0 — current +### Staged delivery in the browser, a base writer, a trace and a bench + +- **A base is bounded before it is decoded, and its chain from the start** + ([staged-delivery](extensions/staged-delivery.md) §4.5, §4.2, §14). A base whose units' + output, with 32 bytes for each stage its kept record holds, passes the chain budget is + not added: its envelope is reported `chain-over-limit`, no chain is kept, its kept + record is untouched, and its layers stay held. Where a base's capability declares a + cost before decoding (the integer payloads SPEC §5 pins), a sink applies its cost + limit to that declared cost first, and SPEC §4 treats a base past it as one it does + not decode, moving on to the fallback; a capability that declares no cost + (`wai.neural.int_synth`, `wai.neural.int_mlicpp`) decodes under its format's caps. The corpus gains `base_over_budget`; + `staged_verify.py` implements the rule independently. +- **The sink in the browser.** `wai-web` binds `wai::staged::StageSink` for JS + (`src/stage.rs`). `StageBase.decode` decodes a base envelope by SPEC §4 through the + integer path, with its pin verified against the parameter set the page holds in + memory and under the page's cost limit, which is a required argument (`null` for the + format's caps only; it bounds a base whose capability declares a cost); + `StageBase.fromStage0` takes a base the page decoded elsewhere, which is **claimed**: + its capability must be the manifest's or its fallback and one SPEC §4 would take (an + inert envelope is refused), its tier is never above what the capability gives an + unpinned base, and every report, unit and stage hash of its chain says `claimed`. A base's facts cannot + be raised from JS (no `viaFallback` setter; `attested` only rises), and `addBase` + takes the base. `StageSink` adds bases, offers layers, expires holds, removes, attests + and forgets bases, and reads back every unit's stage, producer, frozen flag, tier, + shape, output, stage hash (with its `claimed` mark) and squared error against a + reference. `stageLayerInfo` reads a layer's `refines`, `WIR1` header, declared cost and + byte split (framing, manifest, payload), and a base's declared decode cost, before + either is offered. Reports and units come back in the corpus's own form. Media reach + the page as `WAI1` bytes only. +- **The browser equals the native sink.** `wai-web/demo/verify_staged_wasm.mjs` + replays every case of `staged-conformance/` through the wasm32 build: each integer + base is decoded in wasm32 and must give its `stage0.bin`, its tier and its fallback + flag, and every report, final unit, stage hash and clip hash must equal what the + native sink wrote. `wai-web-wasm.yml` runs it and watches the corpus; the crate's + native tests replay the corpus through the binding's pure-Rust half as well. +- **One base path.** `staged::decode_base` runs SPEC §4 for a base through + `Wai::dispatch_verified`: its capability, then its fallback, each only where it is + dispatchable, so a refinement-only or companion-only capability makes the envelope + inert (`missing_capability`) and never reaches its fallback. It takes the sink's cost + limit (checked before any symbol is decoded or output allocated, `cost_over_limit`) + and its key log (SPEC §3.1 check 5 over a pin and over an unpinned set; none skips + it). `decode_base_via` asks for one of the two capabilities and refuses any other + (`not-offered`). The refusals come in a stated order: `not-wai1`, then `encrypted` + (SPEC §6.4 before §4), then `not-offered`, then §4's; a parameter set is read only + for a capability that reads one. Its own tokens are staged-delivery's kebab-case sink + vocabulary; SPEC and payload codes pass through in snake_case. A pin counts toward + the base's tier only where it resolved. A + one-file set pinned without a form offers every held file to the pin, the member the + capability reads first. The CLI and the browser share this path. +- **A keyframe-0 base writer.** `int_payload::encode_wiv1_zero_motion` codes RGB8 + frames as a `WIV1` keyframe-0 payload (zero motion, one block per frame, a closed + group of pictures): the reference encoder the corpus generator used, now in the + library. The corpus regenerates byte for byte through it. It refuses a clip past the + format's caps before allocating (`too_large`). `wai stage base` writes a + `wai.video.int_motion` base with it from raw frames, with its dimensions multiplied + checked, and decodes it back before writing. +- **A trace to check against.** `wai stage apply` is now a sink replay: envelopes are + taken in order, routed by content (one with `refines`, or a refinement-only + capability, is a layer; any other is a base, decoded by SPEC §4), and a base given + after its layers releases the layers held for it. `--json` prints each step's reports + in the corpus's form (bytes that are not a `WAI1` this sink reads, and a base it + cannot decode, are the step's error by their code, and the replay goes on; a base + read with `--stage0` is claimed, and its step, units, chain's reports and final entry + say so), and every stage a unit reached in the step, in the + order the sink applied the layers (`StageSink::set_trace`), so a batch that takes a + unit up several stages lists each; every stage hash a delivery in that order reaches + is in it. A test replays the corpus through it and requires the cases' reports and + final units. +- **`wai_stage_bench`.** Measures one staged set natively: per stage, the envelopes' + bytes (framing, manifest and payload, and cumulative), PSNR or SNR per unit and over + the clip from the exact integer squared error against the source at the units' final + shape (a unit of another shape has no figure at that stage), the time each stage took + (median and quartiles over at least 20 runs), the direct comparator (one layer per + base at the stage's step, straight from stage 0, or why it does not apply) and every + unit's stage hash. Energy is reported as unmetered. + ### Staged delivery: refinement layers, base admission, and WAI2 refinement components - **The extension** [staged-delivery](extensions/staged-delivery.md): a stage set is a diff --git a/wai/extensions/staged-delivery.md b/wai/extensions/staged-delivery.md index 3b5feda14..da16e7b3a 100644 --- a/wai/extensions/staged-delivery.md +++ b/wai/extensions/staged-delivery.md @@ -140,6 +140,7 @@ unit. | 4 | `base-not-admitted` | envelope | discarded | | 5 | `refinement-unattested` | envelope | discarded | | 6 | `payload-malformed`, `cost-over-limit`, `chain-over-limit` | envelope | discarded | +| §4.5 | `chain-over-limit`, for a base | the base's envelope | the base is not added; its layers stay held | | 7 | `layer-conflict` | unit | not applied; the unit is frozen | | 8 | `prev-mismatch` | unit | not applied to the unit | | 8 | `layer-gap` | unit | held for the unit, then not applied | @@ -246,6 +247,16 @@ leaves the layer's: an application lowers the state by the unit's old output (an state rises only when a layer is admitted, so it never passes the budget, and the bytes a chain holds never pass the state. +A chain starts with its base. A base whose units' output, with 32 bytes for each stage +the records the sink keeps of it hold (§4.4), passes the budget is not added: the sink +reports `chain-over-limit` for the base's envelope, keeps no chain for it, leaves the +records it keeps of it as they were, and holds the layers that name it as for a base it +has not decoded (§4.3). A receipt verified with the base is recorded all the same +(§10). The base's own decode is bounded before it starts: where its capability declares +a cost before decoding (`extensions/integer-payloads.md` §7), a sink applies its cost +limit to that declared cost first, and SPEC §4 treats a base past it as one this sink +does not decode (`cost_over_limit`), so dispatch moves on to the fallback. + ## 5. The tier of a stage, and the stage hash - **The base step**: `decode-equivalence` when the capability the base was decoded @@ -403,7 +414,11 @@ ignores the header. - A hold bound (§4.3), a cost limit (Appendix A.5) and a chain budget (§4.5) are resource limits. A layer's header declares its output, its symbols and its working bytes before any symbol is decoded, so a sink refuses a layer past its limit or its - budget without decoding it. A sink's own work for a layer is linear in the + budget without decoding it. A base of an integer codec declares its cost the same way + (`extensions/integer-payloads.md` §7), so a sink refuses a base past its limit before + decoding it, and one whose output passes the chain budget before keeping it (§4.5): + a few bytes of envelope can declare a large output, and only the declared cost bounds + the work. A sink's own work for a layer is linear in the layer's payload bytes and declared symbols, up to a factor logarithmic in its units: it parses the payload once, at step 6; decides each unit the layer covers once for each offer or release of the layer; finds a unit's producer in `prev` by diff --git a/wai/staged-conformance/README.md b/wai/staged-conformance/README.md index 112002b05..29676af69 100644 --- a/wai/staged-conformance/README.md +++ b/wai/staged-conformance/README.md @@ -52,8 +52,9 @@ receipt for it, so that step 5 requires `delivered-verified` of its layers. **The content is synthetic**: pictures, frames and PCM computed by integer functions in the generator. The bases of the integer codecs are the committed vectors under [`../wai-rs/tests/vectors/`](../wai-rs/tests/vectors/), decoded against their -parameter sets with their pins; the `wai.video.int_motion` bases are coded by a -keyframe-0 encoder in the generator (zero motion, one block per frame), and the PNG +parameter sets with their pins; the `wai.video.int_motion` bases are coded by the +library's keyframe-0 reference encoder (`int_payload::encode_wiv1_zero_motion`, which +`wai stage base` also runs: zero motion, one block per frame), and the PNG and FLAC bases are written by small writers in the generator (stored deflate with the five row filters; VERBATIM subframes). diff --git a/wai/staged-conformance/cases/base_over_budget/L1.wai b/wai/staged-conformance/cases/base_over_budget/L1.wai new file mode 100644 index 000000000..84238136c Binary files /dev/null and b/wai/staged-conformance/cases/base_over_budget/L1.wai differ diff --git a/wai/staged-conformance/cases/base_over_budget/base.wai b/wai/staged-conformance/cases/base_over_budget/base.wai new file mode 100644 index 000000000..a5feb5542 Binary files /dev/null and b/wai/staged-conformance/cases/base_over_budget/base.wai differ diff --git a/wai/staged-conformance/cases/base_over_budget/case.json b/wai/staged-conformance/cases/base_over_budget/case.json new file mode 100644 index 000000000..2f08ed86a --- /dev/null +++ b/wai/staged-conformance/cases/base_over_budget/case.json @@ -0,0 +1,115 @@ +{ + "format": "wai-staged-case/1", + "name": "base_over_budget", + "about": "a chain budget one byte short of the base's units: the base is not added (chain-over-limit, §4.5), and the layer held for it stays held until the bound (base-absent)", + "bases": [ + { + "file": "base.wai", + "object": "551af715c03c0216f9e79a6b0bafcce5c31c0c5e573cf4cc7bbab34e1dc8d4c7", + "decoded_via": "wai.video.int_motion", + "via_fallback": false, + "attested": false, + "first_unit": 0, + "pin_resolved": false, + "tier": "decode-equivalence", + "units": [ + { + "shape": [ + 8, + 8, + 3 + ], + "elem": "u8" + }, + { + "shape": [ + 8, + 8, + 3 + ], + "elem": "u8" + } + ], + "stage0": "stage0.bin", + "stage0_sha256": "acb4b066fd922bf6b50128457ea7132f123f18eac82b3d98ac49e145b091742c", + "decode": { + "kind": "int", + "set": {} + } + } + ], + "files": [ + { + "file": "L1.wai", + "object": "9f91cd9d50204f57545e42014eef169c50f9167b038037a83705f7a50aee4a4d", + "bytes": 587 + } + ], + "config": { + "bound": { + "max_envelopes": 16, + "max_age_ms": 1000 + }, + "cost_limit": null, + "budget": 383, + "record_budget": 1048576, + "live": null + }, + "steps": [ + { + "op": "offer", + "file": "L1.wai", + "at_ms": 0, + "evidence": "none" + }, + { + "op": "base", + "index": 0, + "at_ms": 1 + }, + { + "op": "expire", + "at_ms": 1001 + } + ], + "expect": { + "reports": [ + [ + { + "object": "9f91cd9d50204f57545e42014eef169c50f9167b038037a83705f7a50aee4a4d", + "envelope": null, + "detail": null, + "held_for_base": true, + "applied": [], + "held": [], + "refused": [] + } + ], + [ + { + "object": "551af715c03c0216f9e79a6b0bafcce5c31c0c5e573cf4cc7bbab34e1dc8d4c7", + "envelope": "chain-over-limit", + "detail": null, + "held_for_base": false, + "applied": [], + "held": [], + "refused": [] + } + ], + [ + { + "object": "9f91cd9d50204f57545e42014eef169c50f9167b038037a83705f7a50aee4a4d", + "envelope": "base-absent", + "detail": null, + "held_for_base": false, + "applied": [], + "held": [], + "refused": [] + } + ] + ], + "bases": [ + null + ] + } +} diff --git a/wai/staged-conformance/cases/base_over_budget/stage0.bin b/wai/staged-conformance/cases/base_over_budget/stage0.bin new file mode 100644 index 000000000..d74be3a35 --- /dev/null +++ b/wai/staged-conformance/cases/base_over_budget/stage0.bin @@ -0,0 +1 @@ + ( (@0@H0@X8PXPXp`px`hx08H8@X8PXPPhPhx`xxppp8PXHP``pxhxpxxPXpxȰx萘и𨰸ȸ 0 (@0@H0@`@PXPXp`px`px08H8@X8PXPPhPhx`xxpxp8X`HX`P`p`pxpxpPXp`pp`ppxhxhȰx萘𨰸ȸ \ No newline at end of file diff --git a/wai/staged-conformance/index.json b/wai/staged-conformance/index.json index a73311241..a210b272e 100644 --- a/wai/staged-conformance/index.json +++ b/wai/staged-conformance/index.json @@ -319,6 +319,13 @@ "b139902d21228247e2a7901db2077ca0e346c6f34c398b62be71d4516a1dcdf5" ] }, + { + "name": "base_over_budget", + "about": "a chain budget one byte short of the base's units: the base is not added (chain-over-limit, §4.5), and the layer held for it stays held until the bound (base-absent)", + "clip_sha256": [ + null + ] + }, { "name": "chain_budget_retained", "about": "a budget that holds the base and either layer, not both: layer 2, held, is retained, so layer 1 is chain-over-limit; once layer 2 expires and is dropped, layer 1 applies", diff --git a/wai/staged-conformance/staged_verify.py b/wai/staged-conformance/staged_verify.py index 350903489..f2c0efa8f 100644 --- a/wai/staged-conformance/staged_verify.py +++ b/wai/staged-conformance/staged_verify.py @@ -893,6 +893,18 @@ class Sink: if info["attested"]: self.bases[info["hash"]].info["attested"] = True return [] + # §4.5: a chain starts with its base. Its units' output, and 32 bytes for each + # stage the record kept of it holds, past the budget: the base is not added, the + # record is left as it was, and its layers stay held. A receipt verified with it + # is recorded all the same (§10). + k0 = self.kept.get(info["hash"]) + recorded = sum(32 * reached for u, (_, reached) in k0["units"].items() if u < len(units)) if k0 else 0 + if recorded + sum(len(d) * (2 if e == "i16" else 1) for (_, e, d) in units) > self.cfg["budget"]: + if info["attested"]: + self.mark_attested(info["hash"]) + out = Outcome() + out.refuse(info["hash"], "chain-over-limit") + return out.reports() # Attestation belongs to the base hash: a base attested before stays attested. k = self.kept.pop(info["hash"], None) info = dict(info, attested=info["attested"] or (k is not None and k["attested"])) diff --git a/wai/wai-rs/Cargo.toml b/wai/wai-rs/Cargo.toml index 6854bb57d..f34d7a72d 100644 --- a/wai/wai-rs/Cargo.toml +++ b/wai/wai-rs/Cargo.toml @@ -511,6 +511,11 @@ name = "wai_staged_vectors" path = "src/bin/wai_staged_vectors.rs" required-features = ["staged", "video_int", "audio_int", "splat_int"] +[[bin]] +name = "wai_stage_bench" +path = "src/bin/wai_stage_bench.rs" +required-features = ["staged", "video_int", "audio_int", "splat_int"] + [[bin]] name = "wai_film_vectors" path = "src/bin/wai_film_vectors.rs" diff --git a/wai/wai-rs/src/bin/wai_cli/stage.rs b/wai/wai-rs/src/bin/wai_cli/stage.rs index c8cdd2da2..78f41cf47 100644 --- a/wai/wai-rs/src/bin/wai_cli/stage.rs +++ b/wai/wai-rs/src/bin/wai_cli/stage.rs @@ -1,9 +1,10 @@ //! `wai stage`: write, apply, inspect, pack and explode staged-delivery layers -//! (`extensions/staged-delivery.md`; SPEC §7.1). Built with the `staged` feature and -//! the integer codecs' features, which decode the bases. +//! (`extensions/staged-delivery.md`; SPEC §7.1), and write a `wai.video.int_motion` +//! base. Built with the `staged` feature and the integer codecs' features, which decode +//! the bases. //! -//! A base of an integer capability is decoded through `int_payload::decode_int`: with -//! no parameter set for `wai.video.int_motion`, otherwise against the set in the +//! A base of an integer capability is decoded through `staged::decode_base`: with no +//! parameter set for `wai.video.int_motion`, otherwise against the set in the //! directory `--set` names (its files by the names SPEC §5 gives them), with the base's //! pin verified when it carries one. Any other base (PNG, FLAC: this tool has no //! `codecs` decoder for them) is read from `--stage0`, the units' canonical buffers @@ -11,16 +12,27 @@ use std::collections::BTreeMap; +use serde_json::{json, Value}; use wai::container::{Prev, PrevRun, Refines, Span, Wai, WaiMulti}; -use wai::int_payload::{decode_int, resolve}; -use wai::pin::HeldBytes; use wai::staged::wir1::{self, Below, Next, Predictor, RefineParams}; -use wai::staged::{tier, BaseInfo, ChainConfig, HoldBound, ReceiptEvidence, Samples, StageChain, UnitBuf}; +use wai::staged::{tier, BaseInfo, ChainConfig, HoldBound, OfferReport, ReceiptEvidence, Samples, StageChain, StageSink, UnitBuf}; fn hex(b: &[u8]) -> String { b.iter().map(|x| format!("{x:02x}")).collect() } +fn unhex32(s: &str) -> Option<[u8; 32]> { + let b = s.as_bytes(); + if b.len() != 64 { + return None; + } + let mut out = [0u8; 32]; + for (i, o) in out.iter_mut().enumerate() { + *o = u8::from_str_radix(std::str::from_utf8(&b[2 * i..2 * i + 2]).ok()?, 16).ok()?; + } + Some(out) +} + fn opt<'a>(args: &'a [String], flag: &str) -> Option<&'a str> { args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(String::as_str) } @@ -29,6 +41,9 @@ fn opts<'a>(args: &'a [String], flag: &str) -> Vec<&'a str> { args.iter().enumerate().filter(|(_, a)| *a == flag).filter_map(|(i, _)| args.get(i + 1)).map(String::as_str).collect() } +/// The flags that take no value. +const SWITCHES: &[&str] = &["--json", "--table-per-unit"]; + /// The positional arguments: those that are neither a flag nor a flag's value. fn positional(args: &[String]) -> Vec<&str> { let mut out = Vec::new(); @@ -36,6 +51,7 @@ fn positional(args: &[String]) -> Vec<&str> { for a in args { if skip { skip = false; + } else if SWITCHES.contains(&a.as_str()) { } else if a.starts_with("--") || a == "-o" { skip = true; } else { @@ -57,56 +73,59 @@ fn dims(s: &str) -> Result<(u32, u32, u32), String> { } } +/// The parameter set `--set` names, by file name; empty without `--set`. +fn param_set(args: &[String]) -> Result>, String> { + match opt(args, "--set") { + Some(dir) => wai::pin::read_dir_set(std::path::Path::new(dir)).map_err(|e| format!("--set {dir}: {e}")), + None => Ok(BTreeMap::new()), + } +} + +/// An integer base, decoded through `staged::decode_base` against `--set`. +fn int_base(path: &str, env: &[u8], args: &[String]) -> Result<(BaseInfo, Vec), String> { + decode_int_base(env, args)?.map_err(|e| format!("{path}: {e}")) +} + +/// A base decoded by SPEC §4 against `--set`: the outer error is the tool's (an +/// unreadable `--set`), the inner one the base's refusal. +fn decode_int_base(env: &[u8], args: &[String]) -> Result), wai::staged::BaseError>, String> { + let set = param_set(args)?; + let files: BTreeMap = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); + // No cost limit and no key log: a tool over files its user chose (check 5 skipped). + Ok(wai::staged::decode_base(env, &files, None, None)) +} + /// The base's info and units, decoded as the module doc states. fn base(path: &str, args: &[String]) -> Result<(Vec, BaseInfo, Vec), String> { let env = read(path)?; + let Some(s0) = opt(args, "--stage0") else { + let (info, units) = int_base(path, &env, args)?; + return Ok((env, info, units)); + }; let w = Wai::from_bytes(&env).map_err(|e| format!("{path}: {e}"))?; let cap = w.capability().to_string(); - let hash = wai::jwp_hash::content_hash(&env); - let (units, pin_resolved) = if let Some(s0) = opt(args, "--stage0") { - let bytes = read(s0)?; - let n: usize = opt(args, "--units").unwrap_or("1").parse().map_err(|e| format!("--units: {e}"))?; - let d = dims(opt(args, "--dims").ok_or("--stage0 needs --dims")?)?; - let per = d.0 as usize * d.1 as usize * d.2 as usize; - let i16 = opt(args, "--elem") == Some("i16"); - let need = n * per * if i16 { 2 } else { 1 }; - if bytes.len() != need { - return Err(format!("{s0}: {} bytes, the units need {need}", bytes.len())); - } - let units = (0..n) - .map(|u| { - let samples = if i16 { - Samples::I16(bytes[2 * u * per..2 * (u + 1) * per].chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect()) - } else { - Samples::U8(bytes[u * per..(u + 1) * per].to_vec()) - }; - UnitBuf { shape: d, samples } - }) - .collect(); - (units, false) - } else { - let set: BTreeMap> = match opt(args, "--set") { - Some(dir) => wai::pin::read_dir_set(std::path::Path::new(dir)).map_err(|e| format!("--set {dir}: {e}"))?, - None => BTreeMap::new(), - }; - let files: BTreeMap = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); - let listing = if files.is_empty() { vec![] } else { wai::pin::listing_from(&files).map_err(|e| e.to_string())? }; - let one_file = wai::container::pin_shape(&cap) == wai::codecs::ParameterSetShape::OneFile; - let held: Vec = if files.is_empty() { - vec![] - } else if one_file && w.prior().is_some_and(|p| p.form().is_none()) { - vec![HeldBytes::File(files.values().next().copied().unwrap_or_default())] - } else { - vec![HeldBytes::Set { listing: &listing, files: &files }] - }; - let params = resolve(&cap, w.prior(), &held, None).map_err(|e| format!("{path}: {e}"))?; - let pinned = params.is_pinned(); - let d = decode_int(&cap, &w.payload, ¶ms).map_err(|e| format!("{path}: {e}"))?; - (wai::staged::units_from_decoded(d).ok_or("an output kind with no units")?, pinned) - }; + let bytes = read(s0)?; + let n: usize = opt(args, "--units").unwrap_or("1").parse().map_err(|e| format!("--units: {e}"))?; + let d = dims(opt(args, "--dims").ok_or("--stage0 needs --dims")?)?; + let per = d.0 as usize * d.1 as usize * d.2 as usize; + let i16 = opt(args, "--elem") == Some("i16"); + let need = n * per * if i16 { 2 } else { 1 }; + if bytes.len() != need { + return Err(format!("{s0}: {} bytes, the units need {need}", bytes.len())); + } + let units = (0..n) + .map(|u| { + let samples = if i16 { + Samples::I16(bytes[2 * u * per..2 * (u + 1) * per].chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect()) + } else { + Samples::U8(bytes[u * per..(u + 1) * per].to_vec()) + }; + UnitBuf { shape: d, samples } + }) + .collect(); let info = BaseInfo { - content_hash: hash, - tier: tier::base_tier(&cap, if pin_resolved { w.prior() } else { None }), + content_hash: wai::jwp_hash::content_hash(&env), + tier: tier::base_tier(&cap, None), capability: cap, via_fallback: false, attested: false, @@ -139,35 +158,294 @@ fn print_units(c: &StageChain) { } } -/// `wai stage apply …`: offer each layer in order; print every -/// report and the units at the end. -fn apply(args: &[String]) -> Result<(), String> { - let pos = positional(args); - let (_, info, units) = base(pos.first().ok_or("wai stage apply …")?, args)?; - let mut c = StageChain::new(info, units, cfg()); - for (i, l) in pos[1..].iter().enumerate() { - let bytes = read(l)?; - for r in c.offer(&bytes, i as u64, ReceiptEvidence::None).map_err(|e| format!("{l}: {e}"))? { - print!("{l}: {}", &hex(&r.content_hash)[..16]); - if let Some(e) = r.envelope { - print!(" refused {}", e.code()); - if let Some(d) = r.detail { - print!(" ({d})"); - } - } - if !r.applied.is_empty() { - print!(" applied {:?}", r.applied); - } - if !r.held.is_empty() { - print!(" held {:?}", r.held); - } - for (u, w) in &r.refused { - print!(" unit {u}: {}", w.code()); - } - println!(); +/// A report as the corpus states one (`staged-conformance/README.md`). +fn report_json(r: &OfferReport) -> Value { + json!({ + "object": hex(&r.content_hash), + "envelope": r.envelope.map(|x| x.code()), + "detail": r.detail, + "held_for_base": r.held_for_base, + "applied": r.applied, + "held": r.held, + "refused": r.refused.iter().map(|(u, w)| json!([u, w.code()])).collect::>(), + }) +} + +/// Unit `u` of the chain as the corpus states a final unit, with its base. +fn unit_json(c: &StageChain, u: u32) -> Value { + let s = c.unit(u).expect("in range"); + json!({ + "base": hex(&c.base().content_hash), + "unit": u, + "stage": s.stage, + "producer": hex(&s.producer), + "frozen": s.frozen, + "tier": tier::token(s.tier), + "shape": [s.buf.shape.0, s.buf.shape.1, s.buf.shape.2], + "stage_sha256": c.unit_sha256(u).map(|h| hex(&h)), + }) +} + +fn print_report(file: &str, r: &OfferReport) { + print!("{file}: {}", &hex(&r.content_hash)[..16]); + if let Some(e) = r.envelope { + print!(" refused {}", e.code()); + if let Some(d) = r.detail { + print!(" ({d})"); } } - print_units(&c); + if r.held_for_base { + print!(" held for its base"); + } + if !r.applied.is_empty() { + print!(" applied {:?}", r.applied); + } + if !r.held.is_empty() { + print!(" held {:?}", r.held); + } + for (u, w) in &r.refused { + print!(" unit {u}: {}", w.code()); + } + println!(); +} + +/// A traced application as the JSON trace lists a unit: its base, index, stage, +/// producer, frozen flag, tier, shape and stage hash. +fn applied_json(base: &[u8; 32], a: &wai::staged::Applied) -> Value { + json!({ + "base": hex(base), + "unit": a.unit, + "stage": a.stage, + "producer": hex(&a.producer), + "frozen": a.frozen, + "tier": tier::token(a.tier), + "shape": [a.shape.0, a.shape.1, a.shape.2], + "stage_sha256": a.sha256.map(|h| hex(&h)), + }) +} + +/// The base's units at stage 0, as the JSON trace lists them. +fn stage0_json(info: &BaseInfo, units: &[UnitBuf]) -> Vec { + let de = info.tier == Some(wai::codecs::DeterminismTier::DecodeEquivalence); + units + .iter() + .enumerate() + .map(|(u, b)| json!({ + "base": hex(&info.content_hash), + "unit": u, + "stage": 0, + "producer": hex(&info.content_hash), + "frozen": false, + "tier": tier::token(info.tier), + "shape": [b.shape.0, b.shape.1, b.shape.2], + "stage_sha256": de.then(|| hex(&wai::pin::sha256(&b.samples.canonical_bytes()))), + })) + .collect() +} + +/// `wai stage apply … [--json]`: a sink that takes each envelope in order +/// (SPEC §4 and staged-delivery §4). An envelope that carries `refines`, or whose +/// capability is refinement-only, is a layer, offered at the next millisecond; any +/// other is a base, decoded by SPEC §4 (its capability, then its fallback) against +/// `--set` and added, which releases the layers held for it (§4.4). With `--stage0`, +/// the first envelope is a base read from that file. Bytes that are not a `WAI1` this +/// sink reads are that step's error (`not-wai1`, `encrypted`), and the replay goes on. +/// +/// With `--stage0`, the first base's units are that file's, not a decode: the base is +/// **claimed**, and its step, its units, the reports of its chain and its final entry +/// say `"claimed": true` (as the browser binding marks a claimed base's chain); its +/// stage hashes are of what the file holds and the layers applied to it. A base this +/// sink cannot decode (encrypted, inert, refused by its pin or its payload) is its +/// step's error, by its code, and the replay goes on. +/// +/// Prints every report and the units at the end. With `--json`, one JSON document +/// instead: for each step its reports in the corpus's form (or its error), then every +/// stage a unit reached in that step, in the order the sink applied the layers (a +/// base's step starts with its units at stage 0, and a batch that takes a unit up +/// several stages lists each), and the units whose frozen flag the step changed; then +/// every base's units and clip hash. Every stage hash a delivery in this order reaches +/// is in it, for a sink elsewhere to check its own against. +fn apply(args: &[String]) -> Result<(), String> { + let pos = positional(args); + let as_json = args.iter().any(|a| a == "--json"); + if pos.is_empty() { + return Err("wai stage apply … [--json]".into()); + } + let mut sink = StageSink::new(cfg()); + sink.set_trace(true); + let mut bases: Vec<[u8; 32]> = Vec::new(); + let mut frozen: BTreeMap<([u8; 32], u32), bool> = BTreeMap::new(); + let mut steps: Vec = Vec::new(); + // With --stage0 the first base's units are the file's, not a decode: the base is + // claimed, and so is everything its chain reports. + let mut claimed: Option<[u8; 32]> = None; + for (i, path) in pos.iter().enumerate() { + // The base whose chain this step concerns, where it is one. + let mut chain_of: Option<[u8; 32]> = None; + let (op, object, outcome, mut units): (&str, [u8; 32], Result, &str>, Vec) = if i == 0 && opt(args, "--stage0").is_some() { + let (env, info, units) = base(path, args)?; + let first = stage0_json(&info, &units); + bases.push(info.content_hash); + claimed = Some(info.content_hash); + chain_of = Some(info.content_hash); + ("base", wai::jwp_hash::content_hash(&env), Ok(sink.add_base(info, units)), first) + } else { + let bytes = read(path)?; + let object = wai::jwp_hash::content_hash(&bytes); + match Wai::from_bytes(&bytes) { + Err(_) => ("offer", object, Err("not-wai1"), vec![]), + Ok(w) if w.manifest.refines.is_some() || wai::codecs::is_refinement_only(w.capability()) => { + chain_of = w.refines().and_then(Result::ok).and_then(|r| unhex32(&r.base)); + match sink.offer(&bytes, i as u64, ReceiptEvidence::None) { + Ok(r) => ("offer", object, Ok(r), vec![]), + Err(wai::staged::OfferError::NotWai1(_)) => ("offer", object, Err("not-wai1"), vec![]), + Err(_) => ("offer", object, Err("encrypted"), vec![]), + } + } + Ok(_) => { + if i > 0 && opt(args, "--stage0").is_some() { + return Err(format!("{path}: a second base; --stage0 reads the first base only")); + } + match decode_int_base(&bytes, args)? { + // A base this sink cannot decode (encrypted, inert, refused by + // its pin or its payload) is the step's error; the replay goes on. + Err(e) => ("base", object, Err(e.code()), vec![]), + Ok((info, units)) => { + chain_of = Some(info.content_hash); + let first = if sink.chain(&info.content_hash).is_some() { vec![] } else { stage0_json(&info, &units) }; + if !bases.contains(&info.content_hash) { + bases.push(info.content_hash); + } + let r = sink.add_base(info, units); + // A base refused (chain-over-limit) has no stage 0. + let first = if sink.chain(&object).is_some() { first } else { vec![] }; + ("base", object, Ok(r), first) + } + } + } + } + }; + units.extend(sink.take_trace().iter().map(|(b, a)| applied_json(b, a))); + for b in &bases { + let Some(c) = sink.chain(b) else { continue }; + for u in 0..c.n_units() { + let f = c.unit(u).expect("in range").frozen; + if frozen.insert((*b, u), f).unwrap_or(false) != f { + units.push(unit_json(c, u)); + } + } + } + let on_claimed = claimed.is_some() && chain_of == claimed; + if as_json { + let mut step = json!({"file": path, "op": op, "object": hex(&object)}); + match &outcome { + Ok(r) => step["reports"] = json!(r.iter().map(|x| { + let mut v = report_json(x); + if on_claimed { + v["claimed"] = json!(true); + } + v + }).collect::>()), + Err(e) => { + step["error"] = json!(e); + step["reports"] = json!([]); + } + } + if on_claimed { + step["claimed"] = json!(true); + } + if let Some(c) = claimed { + for u in units.iter_mut().filter(|u| u["base"] == hex(&c)) { + u["claimed"] = json!(true); + } + } + step["units"] = Value::Array(units); + steps.push(step); + } else { + match &outcome { + Err(e) => println!("{path}: {} {e}", &hex(&object)[..16]), + Ok(r) => { + if op == "base" && on_claimed { + println!("{path}: base {} claimed: its units are --stage0's, not a decode", &hex(&object)[..16]); + } else if op == "base" { + println!("{path}: base {} decoded", &hex(&object)[..16]); + } + for x in r { + print_report(path, x); + } + } + } + } + } + if as_json { + let finals: Vec = bases + .iter() + .filter_map(|b| sink.chain(b)) + .map(|c| { + let is_claimed = claimed == Some(c.base().content_hash); + let mut v = json!({ + "object": hex(&c.base().content_hash), + "capability": c.base().capability, + "via_fallback": c.base().via_fallback, + "units": (0..c.n_units()).map(|u| { + let mut x = unit_json(c, u); + if is_claimed { + x["claimed"] = json!(true); + } + x + }).collect::>(), + "clip_sha256": c.clip_sha256().map(|h| hex(&h)), + }); + if is_claimed { + v["claimed"] = json!(true); + } + v + }) + .collect(); + let doc = json!({"format": "wai-stage-apply/1", "steps": steps, "bases": finals}); + println!("{}", serde_json::to_string_pretty(&doc).map_err(|e| e.to_string())?); + } else { + for b in &bases { + if let Some(c) = sink.chain(b) { + if bases.len() > 1 || claimed == Some(*b) { + println!("base {}:{}", &hex(b)[..16], if claimed == Some(*b) { " claimed (--stage0): these hashes are of the file's units and the layers on them, not of a decode" } else { "" }); + } + print_units(c); + } + } + } + Ok(()) +} + +/// `wai stage base --frames --dims N,H,W --q Q -o base.wai`: a +/// `wai.video.int_motion` base (`WIV1`, keyframe 0) of the `N` frames of `H × W` RGB8 +/// in `--frames` (concatenated, rows top to bottom), coded by the reference encoder +/// `int_payload::encode_wiv1_zero_motion` at step `Q`: zero motion, one block per frame, +/// a closed group of pictures. The envelope is decoded back before it is written, and +/// must give the encoder's reconstruction. +fn base_cmd(args: &[String]) -> Result<(), String> { + let out = opt(args, "-o").ok_or("-o ")?; + let (n, h, w) = dims(opt(args, "--dims").ok_or("--dims N,H,W")?)?; + let q: u32 = opt(args, "--q").ok_or("--q Q")?.parse().map_err(|e| format!("--q: {e}"))?; + let src = read(opt(args, "--frames").ok_or("--frames ")?)?; + let per = (h as usize).checked_mul(w as usize).and_then(|x| x.checked_mul(3)).ok_or("--dims: H × W × 3 overflows")?; + let need = (n as usize).checked_mul(per).ok_or("--dims: N × H × W × 3 overflows")?; + if n == 0 || per == 0 || src.len() != need { + return Err(format!("--frames holds {} bytes; {n} frames of {h}x{w} RGB8 need {need}", src.len())); + } + let frames: Vec> = src.chunks_exact(per).map(<[u8]>::to_vec).collect(); + let (payload, recon) = wai::int_payload::encode_wiv1_zero_motion(&frames, h, w, q).map_err(|e| e.to_string())?; + let cap = wai::codecs::CAP_VIDEO_INT_MOTION; + let m = wai::Manifest::new("video", "replicate", wai::ModelRequirement::new(cap), wai::Conditioning::new("int_motion")); + let bytes = Wai::new(m, payload).to_emit_bytes().map_err(|e| e.to_string())?; + let (info, units) = wai::staged::decode_base(&bytes, &BTreeMap::new(), None, None).map_err(|e| e.to_string())?; + if units.len() != recon.len() || units.iter().zip(&recon).any(|(u, r)| u.samples != Samples::U8(r.clone())) { + return Err("the base does not decode to the encoder's reconstruction".into()); + } + std::fs::write(out, &bytes).map_err(|e| format!("write {out}: {e}"))?; + let err = frames.iter().flatten().zip(recon.iter().flatten()).map(|(a, b)| a.abs_diff(*b)).max().unwrap_or(0); + println!("base → {out} ({} B) {cap} {n} frames of {h}x{w} q {q} tier {} largest error {err}", bytes.len(), tier::token(info.tier)); + println!(" content hash {}", hex(&info.content_hash)); Ok(()) } @@ -359,6 +637,7 @@ fn explode(args: &[String]) -> Result<(), String> { pub fn main(args: &[String]) -> Result<(), String> { match args.first().map(String::as_str) { Some("apply") => apply(&args[1..]), + Some("base") => base_cmd(&args[1..]), Some("refine") => refine(&args[1..]), Some("inspect") => inspect(&args[1..]), Some("pack") => pack(&args[1..]), @@ -369,7 +648,11 @@ pub fn main(args: &[String]) -> Result<(), String> { pub const USAGE: &str = "wai stage refine --base [--on ]… --target --q N \ [--predictor identity|up2|rate:NUM/DEN] [--span F:C] [--out-dims D0,D1,D2] [--max-symbols N] [--table-per-unit] -o L.wai -wai stage apply … offer the layers in order; print each report and the units +wai stage apply … [--json] take each envelope in order: a layer is offered, a base + decoded (SPEC §4) and added; print each report and the + units, or a JSON trace of every stage reached +wai stage base --frames --dims N,H,W --q Q -o base.wai a wai.video.int_motion base (WIV1, + keyframe 0, zero motion) of N frames of HxW RGB8 wai stage inspect a layer's refines, WIR1 header and declared cost wai stage pack … -o x.wai pack WAI1 envelopes into a WAI2 (manifests verbatim) wai stage explode [-o out.wai] a WAI2 component as its WAI1 envelope diff --git a/wai/wai-rs/src/bin/wai_stage_bench.rs b/wai/wai-rs/src/bin/wai_stage_bench.rs new file mode 100644 index 000000000..1a6ee8573 --- /dev/null +++ b/wai/wai-rs/src/bin/wai_stage_bench.rs @@ -0,0 +1,404 @@ +//! `wai_stage_bench` — measure one staged set natively (`extensions/staged-delivery.md`; +//! the measurement method of the staged-delivery design, §9). +//! +//! wai_stage_bench --source [--set ] [--runs N] [--out report.json] … +//! +//! The envelopes are a stage set: bases (no `refines`), decoded through +//! `staged::decode_base` against the parameter set in `--set` (none for +//! `wai.video.int_motion`), and their refinement layers. Stage 0 is every base decoded; +//! stage `k` is every layer whose `refines.layer` is `k` applied, in the order given. +//! `--source` holds what the units approximate, at their final shape: each unit's +//! canonical buffer as the last stage shapes it, base by base in the order given and +//! unit by unit. A unit whose stage holds another number of samples (below an up2 or a +//! rate layer) has no quality figure at that stage, and the clip's figure counts the +//! units that have one (`units_measured`). +//! +//! For each stage it reports: +//! - **bytes**: the envelopes' bytes as delivered, split into `WAI1` framing, manifest +//! and payload, and the cumulative bytes through that stage. JSON and base64 are +//! never counted; there is no receipt, index or transport byte in a set on disk. +//! - **quality**: for 8-bit units, PSNR = 10·log10(255²·N / SSE) per unit and over the +//! whole clip (N samples, SSE the exact integer sum of squared differences); for +//! 16-bit units, SNR = 10·log10(Σx² / Σ(x−y)²) on the samples. A unit equal to its +//! source has no finite figure and is reported as `null` with `"exact": true`. +//! - **time**: the microseconds the stage took (decoding the bases; applying the +//! layers), median and interquartile range over `--runs` repetitions (at least 20), +//! single-threaded, in this process. +//! - **the direct comparator** (§9.3): for stage `k ≥ 1` whose layers share one step +//! `q` and the identity predictor, one layer per base coded straight from the base's +//! stage 0 to the source at that `q` — the cost of reaching stage `k` in one layer +//! instead of `k`. Its bytes and quality are reported beside the stage's. +//! +//! Every unit's stage hash at every stage is reported too, so a run can be checked +//! against `wai stage apply --json`. Energy is reported as `unmetered`: this tool reads +//! no meter. The figures describe this set on this machine and build, nothing wider. + +use std::collections::BTreeMap; +use std::process::ExitCode; +use std::time::Instant; + +use serde_json::{json, Value}; +use wai::container::{Prev, Refines, Wai}; +use wai::staged::wir1::{self, Below, Next, Predictor, RefineParams}; +use wai::staged::{BaseInfo, ChainConfig, HoldBound, ReceiptEvidence, Samples, StageSink, UnitBuf}; + +fn hex(b: &[u8]) -> String { + b.iter().map(|x| format!("{x:02x}")).collect() +} + +fn opt<'a>(args: &'a [String], flag: &str) -> Option<&'a str> { + args.iter().position(|a| a == flag).and_then(|i| args.get(i + 1)).map(String::as_str) +} + +fn positional(args: &[String]) -> Vec<&str> { + let mut out = Vec::new(); + let mut skip = false; + for a in args { + if skip { + skip = false; + } else if a.starts_with("--") { + skip = true; + } else { + out.push(a.as_str()); + } + } + out +} + +fn read(path: &str) -> Result, String> { + std::fs::read(path).map_err(|e| format!("read {path}: {e}")) +} + +fn cfg() -> ChainConfig { + ChainConfig { bound: HoldBound { max_envelopes: 4096, max_age_ms: u64::MAX }, cost_limit: None, budget: u64::MAX, record_budget: u64::MAX, live: None } +} + +/// One envelope of the set, read once. +struct Env { + file: String, + bytes: Vec, + /// 0 for a base, `refines.layer` for a layer. + stage: u8, + manifest: usize, + payload: usize, + /// For a layer: its step and predictor. + q: Option<(u32, Predictor)>, +} + +fn read_env(file: &str) -> Result { + let bytes = read(file)?; + let w = Wai::from_bytes(&bytes).map_err(|e| format!("{file}: {e}"))?; + let manifest = u32::from_le_bytes(bytes[4..8].try_into().map_err(|_| format!("{file}: no manifest length"))?) as usize; + let (stage, q) = match w.refines() { + None => (0, None), + Some(Err(e)) => return Err(format!("{file}: {e}")), + Some(Ok(r)) => { + let kind = wai::staged::unit_kind(w.capability()).ok_or_else(|| format!("{file}: not a residual refinement"))?; + let p = wir1::parse(&w.payload, kind).map_err(|e| format!("{file}: {e}"))?; + (r.layer, Some((p.q, p.predictor))) + } + }; + Ok(Env { file: file.to_string(), stage, manifest, payload: w.payload.len(), q, bytes }) +} + +fn quantiles(mut v: Vec) -> Value { + v.sort_by(|a, b| a.partial_cmp(b).expect("finite")); + let at = |p: f64| { + let x = p * (v.len() - 1) as f64; + let (lo, hi) = (x.floor() as usize, x.ceil() as usize); + v[lo] + (v[hi] - v[lo]) * (x - lo as f64) + }; + json!({"median": round2(at(0.5)), "p25": round2(at(0.25)), "p75": round2(at(0.75))}) +} + +fn round2(x: f64) -> f64 { + (x * 100.0).round() / 100.0 +} + +/// The squared error of a unit against its source, and its sample count and signal +/// energy (for SNR). +fn sse(s: &Samples, src: &[u8]) -> (u64, u64, u64) { + match s { + Samples::U8(v) => (v.iter().zip(src).map(|(a, b)| (*a as i64 - *b as i64).pow(2) as u64).sum(), v.len() as u64, 0), + Samples::I16(v) => { + let x: Vec = src.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]]) as i64).collect(); + let e = v.iter().zip(&x).map(|(a, b)| (*a as i64 - b).pow(2) as u64).sum(); + (e, v.len() as u64, x.iter().map(|b| (b * b) as u64).sum()) + } + } +} + +/// PSNR (8-bit) or SNR (16-bit) in dB, `None` for an exact match. +fn quality(err: u64, n: u64, energy: u64, pcm: bool) -> Option { + if err == 0 { + return None; + } + let db = if pcm { 10.0 * (energy as f64 / err as f64).log10() } else { 10.0 * (255.0 * 255.0 * n as f64 / err as f64).log10() }; + Some(round2(db)) +} + +fn quality_json(err: u64, n: u64, energy: u64, pcm: bool) -> Value { + let mut v = json!({"sse": err, "samples": n}); + v[if pcm { "snr_db" } else { "psnr_db" }] = json!(quality(err, n, energy, pcm)); + v["exact"] = json!(err == 0); + v +} + +/// The refinement capability over a base's units (as `wai stage refine` picks it). +fn refine_cap(info: &BaseInfo, units: &[UnitBuf]) -> &'static str { + use wai::codecs::{capability_media, Media}; + match units.first().map(|u| &u.samples) { + Some(Samples::I16(_)) => wai::codecs::CAP_AUDIO_INT_REFINE, + _ if capability_media(&info.capability) == Some(Media::Video) => wai::codecs::CAP_VIDEO_INT_REFINE, + _ if capability_media(&info.capability) == Some(Media::Volumetric) => wai::codecs::CAP_SPLAT_INT_REFINE, + _ => wai::codecs::CAP_IMAGE_INT_REFINE, + } +} + +fn below(s: &Samples) -> Below<'_> { + match s { + Samples::U8(v) => Below::U8(v), + Samples::I16(v) => Below::I16(v), + } +} + +fn run(args: &[String]) -> Result<(), String> { + let source = read(opt(args, "--source").ok_or("--source ")?)?; + let runs: usize = opt(args, "--runs").unwrap_or("20").parse().map_err(|e| format!("--runs: {e}"))?; + if runs < 20 { + return Err("--runs: at least 20, so the quartiles mean something".into()); + } + let set: BTreeMap> = match opt(args, "--set") { + Some(dir) => wai::pin::read_dir_set(std::path::Path::new(dir)).map_err(|e| format!("--set {dir}: {e}"))?, + None => BTreeMap::new(), + }; + let files: BTreeMap = set.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); + let envs: Vec = positional(args).iter().map(|f| read_env(f)).collect::>()?; + let n_stages = envs.iter().map(|e| e.stage).max().ok_or("no envelopes")?; + if !envs.iter().any(|e| e.stage == 0) { + return Err("no base among the envelopes".into()); + } + + // A replay, untimed: every stage's units, each base's in order. Every layer must + // apply whole: the bench measures a set delivered whole and in order. + let mut bases: Vec<(BaseInfo, Vec)> = Vec::new(); + for e in envs.iter().filter(|e| e.stage == 0) { + bases.push(wai::staged::decode_base(&e.bytes, &files, None, None).map_err(|x| format!("{}: {x}", e.file))?); + } + let hashes: Vec<[u8; 32]> = bases.iter().map(|(i, _)| i.content_hash).collect(); + // stages[k][b][u]: the unit's samples, stage and stage hash at stage k. + let mut stages_seen: Vec)>>> = Vec::new(); + { + let mut sink = StageSink::new(cfg()); + for (info, units) in &bases { + sink.add_base(info.clone(), units.clone()); + } + let snap = |sink: &StageSink| -> Vec)>> { + hashes.iter().map(|h| { + let c = sink.chain(h).expect("added"); + c.units().iter().enumerate().map(|(u, s)| (s.buf.samples.clone(), s.stage, c.unit_sha256(u as u32))).collect() + }).collect() + }; + stages_seen.push(snap(&sink)); + for k in 1..=n_stages { + for (i, e) in envs.iter().enumerate().filter(|(_, e)| e.stage == k) { + let reports = sink.offer(&e.bytes, i as u64, ReceiptEvidence::None).map_err(|x| format!("{}: {x}", e.file))?; + if let Some(bad) = reports.iter().find(|x| x.envelope.is_some() || !x.refused.is_empty() || !x.held.is_empty() || x.held_for_base) { + return Err(format!("{}: not applied whole ({:?}); the bench measures a set delivered whole and in order", e.file, bad)); + } + } + stages_seen.push(snap(&sink)); + } + } + + // The source, at the units' final shape: each unit's canonical buffer at the last + // stage, base by base and unit by unit. A stage whose unit holds another number of + // samples (below an up2 or rate layer) has no quality figure for that unit. + let last = stages_seen.last().expect("stage 0"); + let mut src_of: Vec> = Vec::new(); + let mut at = 0usize; + for units in last { + let mut v = Vec::new(); + for (samples, ..) in units { + let n = samples.byte_len() as usize; + v.push(source.get(at..at + n).ok_or("--source is shorter than the units at their last stage")?); + at += n; + } + src_of.push(v); + } + if at != source.len() { + return Err(format!("--source holds {} bytes; the units at their last stage hold {at}", source.len())); + } + let pcm = matches!(bases[0].1[0].samples, Samples::I16(_)); + let mut snapshots: Vec> = vec![Vec::new(); n_stages as usize + 1]; + let mut clip: Vec<(u64, u64, u64, usize)> = vec![(0, 0, 0, 0); n_stages as usize + 1]; + for (k, per_base) in stages_seen.iter().enumerate() { + for (b, units) in per_base.iter().enumerate() { + for (u, (samples, stage, sha)) in units.iter().enumerate() { + let src = src_of[b][u]; + let mut q = if samples.byte_len() as usize == src.len() { + let (err, n, energy) = sse(samples, src); + clip[k].0 += err; + clip[k].1 += n; + clip[k].2 += energy; + clip[k].3 += 1; + quality_json(err, n, energy, pcm) + } else { + json!({"sse": null, "samples": null, "exact": null, "why": "the unit's shape at this stage is not the source's"}) + }; + q["base"] = json!(hex(&hashes[b])); + q["unit"] = json!(u); + q["stage"] = json!(stage); + q["stage_sha256"] = json!(sha.map(|x| hex(&x))); + snapshots[k].push(q); + } + } + } + + // Timed replays. + let mut times: Vec> = vec![Vec::new(); n_stages as usize + 1]; + for _ in 0..runs { + let mut sink = StageSink::new(cfg()); + let t = Instant::now(); + for e in envs.iter().filter(|e| e.stage == 0) { + let (info, units) = wai::staged::decode_base(&e.bytes, &files, None, None).map_err(|x| format!("{}: {x}", e.file))?; + sink.add_base(info, units); + } + times[0].push(t.elapsed().as_secs_f64() * 1e6); + for k in 1..=n_stages { + let t = Instant::now(); + for (i, e) in envs.iter().enumerate().filter(|(_, e)| e.stage == k) { + sink.offer(&e.bytes, i as u64, ReceiptEvidence::None).map_err(|x| format!("{}: {x}", e.file))?; + } + times[k as usize].push(t.elapsed().as_secs_f64() * 1e6); + } + } + + // The direct comparator for each stage whose layers share one step and identity. + let direct = |k: u8| -> Result { + let qs: Vec<(u32, Predictor)> = envs.iter().filter(|e| e.stage == k).filter_map(|e| e.q).collect(); + let Some(&(q, p)) = qs.first() else { return Ok(Value::Null) }; + if p != Predictor::Identity || qs.iter().any(|x| *x != (q, p)) { + return Ok(json!({"q": null, "why": "the stage's layers do not share one step and the identity predictor"})); + } + let (mut bytes, mut err, mut n, mut energy) = (0u64, 0u64, 0u64, 0u64); + for (b, (info, units)) in bases.iter().enumerate() { + let shape = units[0].shape; + if units.iter().any(|u| u.shape != shape) { + return Ok(json!({"q": q, "why": "a base's units differ in shape"})); + } + // The target is the source at the units' final shape; an identity layer + // from stage 0 reaches it only where stage 0 has that shape. + if units.iter().zip(&src_of[b]).any(|(u, s)| u.samples.byte_len() as usize != s.len()) { + return Ok(json!({"q": q, "why": "the units change shape between stage 0 and the source; an identity layer from stage 0 cannot reach it"})); + } + let belows: Vec = units.iter().map(|u| below(&u.samples)).collect(); + let targets: Vec = src_of[b] + .iter() + .map(|s| if pcm { Samples::I16(s.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect()) } else { Samples::U8(s.to_vec()) }) + .collect(); + let tb: Vec = targets.iter().map(below).collect(); + let params = RefineParams { q, predictor: Predictor::Identity, rate: (0, 0), table_per_unit: false, max_symbols: 4095 }; + let cap = refine_cap(info, units); + let kind = wai::staged::unit_kind(cap).expect("a refinement"); + let (payload, recon) = wir1::encode(kind, &belows, shape, &tb, shape, ¶ms).map_err(|e| e.to_string())?; + let refines = Refines { base: hex(&info.content_hash), prev: Prev::One(hex(&info.content_hash)), layer: 1, span: None }; + let media = wai::codecs::capability_media(cap).map(|m| m.label()).unwrap_or("binary"); + let m = wai::Manifest::new(media, "replicate", wai::ModelRequirement::new(cap), wai::Conditioning::new("wir1")).with_refines(&refines); + bytes += Wai::new(m, payload).to_emit_bytes().map_err(|e| e.to_string())?.len() as u64; + for (r, t) in recon.iter().zip(&targets) { + let s = match r { + Next::U8(v) => Samples::U8(v.clone()), + Next::I16(v) => Samples::I16(v.clone()), + }; + let (e, m, x) = sse(&s, &t.canonical_bytes()); + err += e; + n += m; + energy += x; + } + } + let mut v = quality_json(err, n, energy, pcm); + v["q"] = json!(q); + v["bytes"] = json!(bytes); + Ok(v) + }; + + let mut stages = Vec::new(); + let mut cumulative = 0u64; + for k in 0..=n_stages { + let es: Vec<&Env> = envs.iter().filter(|e| e.stage == k).collect(); + let bytes: u64 = es.iter().map(|e| e.bytes.len() as u64).sum(); + cumulative += bytes; + let (err, n, energy, measured) = clip[k as usize]; + let mut quality = if measured == 0 { json!({"sse": null, "samples": null, "exact": null, "why": "no unit has the source's shape at this stage"}) } else { quality_json(err, n, energy, pcm) }; + quality["units_measured"] = json!(measured); + stages.push(json!({ + "stage": k, + "envelopes": es.len(), + "bytes": bytes, + "framing_bytes": es.iter().map(|e| (e.bytes.len() - e.manifest - e.payload) as u64).sum::(), + "manifest_bytes": es.iter().map(|e| e.manifest as u64).sum::(), + "payload_bytes": es.iter().map(|e| e.payload as u64).sum::(), + "cumulative_bytes": cumulative, + "clip": quality, + "time_us": quantiles(times[k as usize].clone()), + "direct": if k == 0 { Value::Null } else { direct(k)? }, + "units": snapshots[k as usize], + })); + } + let report = json!({ + "format": "wai-stage-bench/1", + "conditions": { + "runs": runs, + "profile": if cfg!(debug_assertions) { "debug" } else { "release" }, + "arch": std::env::consts::ARCH, + "os": std::env::consts::OS, + "threads": 1, + "energy": "unmetered", + "receipts": "none: a set on disk carries no delivery receipts", + "quality": if pcm { "SNR = 10·log10(Σx² / Σ(x−y)²) on the 16-bit samples" } else { "PSNR = 10·log10(255²·N / SSE), SSE the exact integer sum over N samples" }, + }, + "envelopes": envs.iter().map(|e| json!({ + "file": e.file, "object": hex(&wai::jwp_hash::content_hash(&e.bytes)), "stage": e.stage, "bytes": e.bytes.len(), + "manifest_bytes": e.manifest, "payload_bytes": e.payload, "q": e.q.map(|x| x.0), "predictor": e.q.map(|x| x.1.token()), + })).collect::>(), + "stages": stages, + }); + + println!("stage envelopes bytes cumulative clip quality time µs (median [p25, p75]) direct: bytes quality"); + for s in report["stages"].as_array().expect("stages") { + let qk = if pcm { "snr_db" } else { "psnr_db" }; + let q = |v: &Value| match v[qk].as_f64() { + Some(x) => format!("{x:.2} dB"), + None if v["exact"] == true => "exact".to_string(), + None => "-".to_string(), + }; + let d = &s["direct"]; + let direct = if d.is_null() || d["bytes"].is_null() { "-".to_string() } else { format!("{:>9} {}", d["bytes"], q(d)) }; + println!("{:>5} {:>9} {:>9} {:>10} {:>16} {:>9} [{}, {}] {}", + s["stage"], s["envelopes"], s["bytes"], s["cumulative_bytes"], q(&s["clip"]), + s["time_us"]["median"], s["time_us"]["p25"], s["time_us"]["p75"], direct); + } + if let Some(out) = opt(args, "--out") { + let mut text = serde_json::to_string_pretty(&report).map_err(|e| e.to_string())?; + text.push('\n'); + std::fs::write(out, text).map_err(|e| format!("write {out}: {e}"))?; + println!("report → {out}"); + } + Ok(()) +} + +fn main() -> ExitCode { + let args: Vec = std::env::args().skip(1).collect(); + if args.is_empty() { + eprintln!("usage: wai_stage_bench --source [--set ] [--runs N≥20] [--out report.json] …"); + return ExitCode::from(2); + } + match run(&args) { + Ok(()) => ExitCode::SUCCESS, + Err(e) => { + eprintln!("error: {e}"); + ExitCode::from(1) + } + } +} diff --git a/wai/wai-rs/src/bin/wai_staged_vectors.rs b/wai/wai-rs/src/bin/wai_staged_vectors.rs index 1edae28d1..0d18cf8a4 100644 --- a/wai/wai-rs/src/bin/wai_staged_vectors.rs +++ b/wai/wai-rs/src/bin/wai_staged_vectors.rs @@ -24,8 +24,9 @@ //! //! The content is synthetic: pictures, frames and PCM computed by integer functions in //! this file. The base payloads of the integer codecs are the committed vectors; the -//! `wai.video.int_motion` bases are coded by a keyframe-0 encoder confined to this file -//! (zero motion, one block per frame), which the decode checks. +//! `wai.video.int_motion` bases are coded by the library's keyframe-0 reference encoder +//! (`int_payload::encode_wiv1_zero_motion`: zero motion, one block per frame), which +//! the decode checks. use std::collections::BTreeMap; use std::path::{Path, PathBuf}; @@ -117,59 +118,6 @@ fn envelope(media: &str, cap: &str, fallback: Option<&str>, kind: &str, prior: O Wai::new(m, payload).to_bytes().expect("serialise") } -// ───────────────────────── a keyframe-0 WIV1 encoder (this generator only) ───────────────────────── - -/// Code `frames` (each `h × w` RGB8) as `wai.video.int_motion` (`WIV1`, keyframe 0) -/// with zero motion and one block per frame: frame 0 predicted from grey, each later -/// frame from the previous decoded frame. Returns the payload and its decode. -fn encode_wiv1(frames: &[Vec], h: u32, w: u32, q: u32) -> (Vec, Vec>) { - let mut recon: Vec> = Vec::new(); - let mut rs: Vec = Vec::new(); - for (f, src) in frames.iter().enumerate() { - let mut cur = Vec::with_capacity(src.len()); - for (i, &t) in src.iter().enumerate() { - let pred = if f == 0 { 128 } else { recon[f - 1][i] as i64 }; - let r = wai::staged::predict::rdiv(t as i64 - pred, q as i64); - rs.push(r); - cur.push((pred + r * q as i64).clamp(0, 255) as u8); - } - recon.push(cur); - } - let lo = *rs.iter().min().unwrap(); - let hi = *rs.iter().max().unwrap(); - let width = ((hi - lo + 1) as u64).min(4095); - let mut hist = vec![0u64; width as usize]; - for &r in &rs { - if r >= lo && r < lo + width as i64 { - hist[(r - lo) as usize] += 1; - } - } - let cdf = wai::int_entropy::hist_to_cdf(&hist, 16).expect("table"); - let mut e = wai::rans::RansEncoder::new(); - for &r in &rs { - assert!(e.encode(r as i32, &cdf, cdf.len(), lo as i32)); - } - let stream = e.flush(); - let n = frames.len() as u32; - let block = h.max(w); - let v = wai::int_wire::Wiv1 { - mv_precision: 1, - keyframe: 0, - n_frames: n, - h, - w, - block: if n == 1 { 1 } else { block }, - q, - offset: lo as i32, - cdf, - mvs: vec![(0, 0); (n as usize - 1) * if n == 1 { 0 } else { 1 }], - rans: &stream, - wih1: None, - }; - let payload = wai::int_wire::write_wiv1(&v); - (payload, recon) -} - // ───────────────────────── PNG and FLAC writers (fixtures) ───────────────────────── fn crc32(data: &[u8]) -> u32 { @@ -526,7 +474,7 @@ fn splat_base(ctx: &Ctx) -> Base { /// `seed`, starting at time `t0`. fn gop_base(ctx: &Ctx, h: u32, w: u32, n: u32, seed: u32, t0: u32, q: u32) -> Base { let frames: Vec> = (0..n).map(|f| picture(h, w, seed, t0 + f)).collect(); - let (payload, recon) = encode_wiv1(&frames, h, w, q); + let (payload, recon) = wai::int_payload::encode_wiv1_zero_motion(&frames, h, w, q).expect("a clip the encoder codes"); let units = decode_with::<&str>(ctx, CAP_VIDEO_INT_MOTION, &payload, &[], None); let mine: Vec = recon.into_iter().map(|f| UnitBuf { shape: (h, w, 3), samples: Samples::U8(f) }).collect(); assert_eq!(units, mine, "the WIV1 decode is the encoder's reconstruction"); @@ -1143,6 +1091,14 @@ fn cases(ctx: &Ctx) -> Vec { out.push(Case { name: "chain_over_limit", about: "a chain budget one byte short of the base's units and the layer's payload, declared output and working bytes and 32 bytes a unit: chain-over-limit, before any stream is decoded", bases: vec![a], files: vec![("L1.wai".into(), l1.clone())], cfg: ChainConfig { budget: state0 + reserve(&l1) - 1, ..cfg() }, steps: vec![base_at(0, 0), offer("L1.wai", 1)] }); } + { + let a = small(141); + let (l1, _) = l1_of(&a, 142, 4); + let state0: u64 = a.units.iter().map(|u| u.samples.byte_len()).sum(); + out.push(Case { name: "base_over_budget", about: "a chain budget one byte short of the base's units: the base is not added (chain-over-limit, §4.5), and the layer held for it stays held until the bound (base-absent)", bases: vec![a], + files: vec![("L1.wai".into(), l1)], cfg: ChainConfig { budget: state0 - 1, ..cfg() }, + steps: vec![offer("L1.wai", 0), base_at(0, 1), Step::Expire { at_ms: 1001 }] }); + } { let a = small(117); let (l1, mid) = l1_of(&a, 118, 4); diff --git a/wai/wai-rs/src/int_payload/mod.rs b/wai/wai-rs/src/int_payload/mod.rs index d8905d011..ffd49ffb3 100644 --- a/wai/wai-rs/src/int_payload/mod.rs +++ b/wai/wai-rs/src/int_payload/mod.rs @@ -827,3 +827,101 @@ pub fn wis1_from_json(bs: &serde_json::Value) -> Result, IntPayloadError wire::parse_wis1(&out)?; Ok(out) } + +// ───────────────────────── a reference keyframe-0 encoder ───────────────────────── + +/// `n / d` rounded to the nearest integer, halves away from zero, for `d ≥ 1`: the +/// quantizer of [`encode_wiv1_zero_motion`] (the rounding of +/// `staged::predict::rdiv`, which this module cannot reach without `staged`). +#[cfg(feature = "video_int")] +fn round_div(n: i64, d: u64) -> i64 { + let m = n.unsigned_abs(); + let (q, r) = (m / d, m % d); + let mag = (q + (r >= d - r) as u64) as i64; + if n < 0 { -mag } else { mag } +} + +/// Code `frames` (each `h × w` RGB8, rows top to bottom) as a `WIV1` keyframe-0 +/// payload (`wai.video.int_motion`) with zero motion and one block per frame: frame 0 +/// is predicted from grey (128), each later frame from the previous frame as decoded, +/// and every residual is quantized at step `q` rounded half away from zero, so the +/// decode adds `r·q` back and clamps. The residual table is fitted to the clip's +/// residuals (`int_entropy::hist_to_cdf`, at most 4,095 symbols wide; a residual past +/// it is coded through the bypass escape). Returns the payload and its decode, each +/// frame as `decode_int` gives it. +/// +/// A reference encoder for the format, not a rate-distortion encoder: it searches no +/// motion. It writes a closed group of pictures, since frame 0 refers to nothing. A +/// clip past the format's caps (`MAX_TENSOR` elements in a frame, `MAX_FRAMES` frames, +/// `MAX_CLIP` elements) is refused `too_large` before anything is allocated. +#[cfg(feature = "video_int")] +pub fn encode_wiv1_zero_motion(frames: &[Vec], h: u32, w: u32, q: u32) -> Result<(Vec, Vec>), IntPayloadError> { + let field = |why: &str| IntPayloadError::new(Refusal::FieldRange, why.to_string()); + if !(1..=wire::MAX_Q).contains(&q) { + return Err(field("q is outside 1..=65535")); + } + if frames.is_empty() || h == 0 || w == 0 { + return Err(field("a clip needs a frame of at least one pixel")); + } + // The format's caps, checked before anything is allocated: a frame of at most + // MAX_TENSOR elements, at most MAX_FRAMES frames and MAX_CLIP elements in the clip. + let too_large = |why: &str| IntPayloadError::new(Refusal::TooLarge, why.to_string()); + let len = (h as u64).checked_mul(w as u64).and_then(|p| p.checked_mul(3)).filter(|&l| l <= wire::MAX_TENSOR).ok_or_else(|| too_large("h × w × 3 exceeds the tensor cap"))?; + if frames.len() as u64 > wire::MAX_FRAMES { + return Err(too_large("more frames than the frame cap")); + } + if (frames.len() as u64).checked_mul(len).is_none_or(|c| c > wire::MAX_CLIP) { + return Err(too_large("the clip exceeds the clip cap")); + } + let n = frames.len() as u32; + if frames.iter().any(|f| f.len() as u64 != len) { + return Err(field("a frame is not h × w × 3 bytes")); + } + let mut recon: Vec> = Vec::with_capacity(frames.len()); + let mut rs: Vec = Vec::with_capacity(frames.len() * len as usize); + for (f, src) in frames.iter().enumerate() { + let mut cur = Vec::with_capacity(src.len()); + for (i, &t) in src.iter().enumerate() { + let pred = if f == 0 { 128 } else { recon[f - 1][i] as i64 }; + let r = round_div(t as i64 - pred, q as u64); + rs.push(r); + cur.push((pred + r * q as i64).clamp(0, 255) as u8); + } + recon.push(cur); + } + let lo = *rs.iter().min().expect("at least one sample"); + let hi = *rs.iter().max().expect("at least one sample"); + let width = ((hi - lo + 1) as u64).min(4095); + let mut hist = vec![0u64; width as usize]; + for &r in &rs { + if r >= lo && r < lo + width as i64 { + hist[(r - lo) as usize] += 1; + } + } + let cdf = crate::int_entropy::hist_to_cdf(&hist, 16).ok_or_else(|| field("no residual table fits"))?; + let mut e = crate::rans::RansEncoder::new(); + for &r in &rs { + if !e.encode(r as i32, &cdf, cdf.len(), lo as i32) { + return Err(field("a residual the stream cannot code")); + } + } + let stream = e.flush(); + let one = n == 1; + let v = wire::Wiv1 { + mv_precision: 1, + keyframe: 0, + n_frames: n, + h, + w, + block: if one { 1 } else { h.max(w) }, + q, + offset: lo as i32, + cdf, + mvs: vec![(0, 0); if one { 0 } else { n as usize - 1 }], + rans: &stream, + wih1: None, + }; + let out = wire::write_wiv1(&v); + wire::parse_wiv1(&out, Some(0))?; + Ok((out, recon)) +} diff --git a/wai/wai-rs/src/int_payload/tests.rs b/wai/wai-rs/src/int_payload/tests.rs index 445557452..4b124310d 100644 --- a/wai/wai-rs/src/int_payload/tests.rs +++ b/wai/wai-rs/src/int_payload/tests.rs @@ -775,3 +775,54 @@ fn a_frame_entry_that_is_not_a_list_is_refused() { assert_eq!(wiv1_from_json(&js, None).unwrap_err().code(), "field_range", "{bad}"); } } + +/// `encode_wiv1_zero_motion` writes a `WIV1` keyframe-0 payload that `decode_int` +/// decodes to the encoder's reconstruction, frame by frame; at `q = 1` that is the +/// source itself, and at a coarser step every sample is within `q / 2` of it. A +/// one-frame clip carries no motion vector; a step of 0, an empty clip and a frame of +/// the wrong size are refused. +#[cfg(feature = "video_int")] +#[test] +fn encode_wiv1_zero_motion_round_trips() { + use crate::codecs::CAP_VIDEO_INT_MOTION; + let (h, w) = (9u32, 14u32); + let frames: Vec> = (0..4u32) + .map(|t| (0..h * w * 3).map(|i| ((i * 7 + t * 29 + (i % 5) * (i % 3) * 40) % 256) as u8).collect()) + .collect(); + for q in [1u32, 3, 8, 64] { + for n in [1usize, 4] { + let (payload, recon) = encode_wiv1_zero_motion(&frames[..n], h, w, q).unwrap(); + let v = wire::parse_wiv1(&payload, Some(0)).unwrap(); + assert_eq!((v.n_frames as usize, v.h, v.w, v.q), (n, h, w, q)); + assert_eq!(v.mvs.len(), n - 1, "zero motion, one block per frame"); + match decode_int(CAP_VIDEO_INT_MOTION, &payload, &IntParams::none()).unwrap() { + IntDecoded::Video { frames: got, width, height } => { + assert_eq!((width, height), (w, h)); + assert_eq!(got, recon, "q {q}, {n} frames: the decode is the encoder's reconstruction"); + } + other => panic!("{other:?}"), + } + for (src, r) in frames[..n].iter().zip(&recon) { + let worst = src.iter().zip(r).map(|(a, b)| a.abs_diff(*b) as u32).max().unwrap(); + assert!(worst <= q / 2, "q {q}: an error of {worst}"); + } + if q == 1 { + assert_eq!(recon, frames[..n].to_vec(), "q = 1 is lossless"); + } + } + } + let code = |r: Result<(Vec, Vec>), IntPayloadError>| r.unwrap_err().code(); + assert_eq!(code(encode_wiv1_zero_motion(&frames, h, w, 0)), "field_range"); + assert_eq!(code(encode_wiv1_zero_motion(&[], h, w, 4)), "field_range"); + assert_eq!(code(encode_wiv1_zero_motion(&frames, h, w + 1, 4)), "field_range"); + // The caps, before anything is read or allocated: a frame past the tensor cap (its + // bytes are never looked at), more frames than the frame cap, a clip past the clip + // cap, and dimensions whose product overflows. + let one = vec![vec![0u8; 3]]; + assert_eq!(code(encode_wiv1_zero_motion(&one, 4097, 5462, 4)), "too_large"); + assert_eq!(code(encode_wiv1_zero_motion(&one, u32::MAX, u32::MAX, 4)), "too_large"); + let many = vec![Vec::new(); (wire::MAX_FRAMES + 1) as usize]; + assert_eq!(code(encode_wiv1_zero_motion(&many, 1, 1, 4)), "too_large"); + let clip = vec![Vec::new(); 17]; + assert_eq!(code(encode_wiv1_zero_motion(&clip, 4096, 5461, 4)), "too_large", "17 frames of a frame within the tensor cap pass the clip cap"); +} diff --git a/wai/wai-rs/src/staged/mod.rs b/wai/wai-rs/src/staged/mod.rs index d91fc1857..298af3a08 100644 --- a/wai/wai-rs/src/staged/mod.rs +++ b/wai/wai-rs/src/staged/mod.rs @@ -50,6 +50,267 @@ pub fn units_from_decoded(d: crate::int_payload::IntDecoded) -> Option &'static str { + match self { + BaseError::NotWai1(_) => "not-wai1", + BaseError::Encrypted => "encrypted", + BaseError::Pin(p) => p.code(), + BaseError::Decode(e) => e.code(), + BaseError::NoUnits => "no-units", + BaseError::NotOffered(_) => "not-offered", + BaseError::Inert => "missing_capability", + } + } +} + +impl std::fmt::Display for BaseError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + BaseError::NotWai1(e) => write!(f, "not-wai1: {e}"), + BaseError::Encrypted => f.write_str("encrypted: an encrypted envelope this sink cannot open"), + BaseError::Pin(p) => write!(f, "{}: {p}", p.code()), + BaseError::Decode(e) => write!(f, "{e}"), + BaseError::NoUnits => f.write_str("no-units: the decode gave an output with no units"), + BaseError::NotOffered(c) => write!(f, "not-offered: {c} is neither the manifest's capability nor its fallback"), + BaseError::Inert => f.write_str( + "missing_capability: SPEC §4 takes no capability for this envelope here (a refinement-only or \ + companion-only capability, or none this sink decodes through the integer path)", + ), + } + } +} + +impl std::error::Error for BaseError {} + +/// Whether `cap` decodes through `int_payload::decode_int`: the capabilities SPEC §5 +/// pins an integer payload format to, and `wai.neural.int_synth` and +/// `wai.neural.int_mlicpp`. What a base decode advertises to SPEC §4. +fn int_path(cap: &str) -> bool { + crate::codecs::pinned_int_payload(cap).is_some() || matches!(cap, crate::codecs::CAP_NEURAL_INT_SYNTH | crate::codecs::CAP_NEURAL_INT_MLICPP) +} + +/// The sets a sink holding `set` offers a decode through `cap` (SPEC §3.1): none for a +/// capability that reads no parameter set (`wai.video.int_motion`), whatever `set` +/// holds; for a capability whose set is one file, pinned without a form, every file as +/// its own candidate (the pin resolves when any one satisfies it), the member the +/// capability reads first; otherwise the files as one set, which an unpinned decode +/// reads by its members' names, with their listing (SPEC §3.1), made from `set` the +/// first time a capability needs it. A set whose listing cannot be made (a file name +/// SPEC §3.1 refuses) is that capability's refusal. +fn held_for<'a>( + cap: &str, + prior: Option<&crate::container::Prior>, + set: &'a BTreeMap, + listing: &'a std::cell::OnceCell, crate::container::PinRefusal>>, +) -> Result>, crate::container::PinRefusal> { + use crate::pin::HeldBytes; + if set.is_empty() || (crate::codecs::is_known_capability(cap) && !crate::codecs::decodes_against_parameter_set(cap)) { + return Ok(vec![]); + } + let one_file = crate::container::pin_shape(cap) == crate::codecs::ParameterSetShape::OneFile; + if one_file && prior.is_some_and(|p| p.form().is_none()) { + let member = crate::int_payload::set_members(cap).first().copied(); + let mut v: Vec = member.and_then(|m| set.get(m)).map(|b| HeldBytes::File(b)).into_iter().collect(); + v.extend(set.iter().filter(|(k, _)| Some(k.as_str()) != member).map(|(_, b)| HeldBytes::File(b))); + Ok(v) + } else { + let listing = listing.get_or_init(|| crate::pin::listing_from(set)).as_ref().map_err(Clone::clone)?; + Ok(vec![HeldBytes::Set { listing, files: set }]) + } +} + +/// Whether a refusal of the capability §4 picked makes it unsupported for this +/// envelope, so dispatch moves on to the fallback: no decoder in this build +/// (`unsupported_capability`), a set that lacks a member (`set_incomplete`), or a +/// declared cost past the sink's limit (`cost_over_limit`), each refused before any +/// symbol was decoded. Any other refusal is the payload's, and final. +fn routes_on(e: &crate::int_payload::IntPayloadError) -> bool { + use crate::int_payload::Refusal; + matches!(e.refusal(), Refusal::UnsupportedCapability | Refusal::SetIncomplete | Refusal::CostOverLimit) +} + +/// Decode a base envelope for a sink (§2) by SPEC §4: its capability first, then its +/// fallback, each taken only where [`crate::codecs::is_dispatchable`] accepts it, as +/// `Wai::dispatch_verified` takes them (a refinement-only or companion-only capability +/// makes the envelope inert, and its fallback is never reached), through +/// `int_payload::decode_int`. +/// +/// **The order of refusals.** Bytes that are not a `WAI1` envelope (`not-wai1`); then +/// an envelope whose manifest declares an encryption this sink cannot open +/// (`encrypted`: SPEC §6.4 step 2 comes before §4); then, for [`decode_base_via`], a +/// capability the manifest does not offer (`not-offered`); then §4's dispatch, whose +/// refusal is the first kept for the capability before the one kept for the fallback +/// (a pin or set refusal, a payload code, or `missing_capability`); then an output with +/// no units (`no-units`). A parameter set is read only for a capability that reads one: +/// a file name in `set` that SPEC §3.1 would refuse is no refusal for +/// `wai.video.int_motion`. +/// +/// **The codes.** This function's own refusals (`not-wai1`, `encrypted`, +/// `not-offered`, `no-units`) are tokens of staged-delivery's sink vocabulary, in +/// kebab-case. Codes from elsewhere pass through verbatim, in snake_case: SPEC §4's +/// `missing_capability`, SPEC §3.1's pin refusals (`prior_revoked`, `set_incomplete`, +/// …) and the integer payloads' codes (`cost_over_limit`, `stream_exhausted`, …, +/// `extensions/integer-payloads.md` §8). +/// +/// - `set` is the capability's parameter set, each file by the name SPEC §5 gives it +/// (empty for `wai.video.int_motion`, which reads none). The envelope's pin, where +/// it carries one, is verified against those bytes (SPEC §3.1). +/// - `limit` is the sink's cost limit, for a base whose capability declares a cost +/// before decoding: the integer payloads SPEC §5 pins a format to (`WIV1`, `WIH1`, +/// `WIA1`, `WIS1`; `int_payload::check_cost`). Their declared cost is checked +/// before any symbol is decoded or any output allocated, and a payload past it is +/// `cost_over_limit`. That refusal, like a capability this build has no decoder +/// for (`unsupported_capability`) or a set that lacks a member (`set_incomplete`), +/// makes the capability unsupported for this envelope, and dispatch moves on to the +/// fallback. A capability that declares no cost (`wai.neural.int_synth`, +/// `wai.neural.int_mlicpp`) is decoded under its format's own caps whatever the +/// limit, as `ModelRegistry::with_cost_limit` does (staged-delivery §4.5: "where its +/// capability declares a cost"). With no limit, every decode is bounded by its +/// format's caps alone. +/// - `authority` is the sink's key log: where given, SPEC §3.1's check 5 runs over the +/// pin and over an unpinned decode's set (`prior_revoked`, `prior_not_authorised`). +/// `None` skips check 5: a sink that holds no key log (SPEC §4, *Module +/// authorisation*), such as the browser binding. +/// +/// The base's info states its content hash, the capability it was decoded through, +/// whether that was the fallback, and its tier (§5, [`tier::base_tier`]: the pin +/// counts only where it resolved); it is not attested, and its unit 0 starts the live +/// clock. When nothing decodes, the error is the first refusal kept, the capability's +/// before the fallback's, else `missing_capability`. +pub fn decode_base( + envelope: &[u8], + set: &BTreeMap, + limit: Option<&crate::int_payload::DeclaredCost>, + authority: Option<&dyn crate::pin::PinAuthority>, +) -> Result<(BaseInfo, Vec), BaseError> { + decode_base_via(envelope, None, set, limit, authority) +} + +/// [`decode_base`] at a sink that decodes through `via` alone: the manifest's +/// capability or its fallback, which §4 takes only where it would at a sink that +/// advertised nothing else (`missing_capability` otherwise, so a refinement-only or +/// companion-only capability never reaches its fallback). Any other string is +/// `not-offered`. `None` is [`decode_base`]. +pub fn decode_base_via( + envelope: &[u8], + via: Option<&str>, + set: &BTreeMap, + limit: Option<&crate::int_payload::DeclaredCost>, + authority: Option<&dyn crate::pin::PinAuthority>, +) -> Result<(BaseInfo, Vec), BaseError> { + use crate::container::Inert; + use crate::int_payload::{check_cost, decode_int, resolve}; + let w = Wai::from_bytes(envelope).map_err(|e| BaseError::NotWai1(e.to_string()))?; + // SPEC §6.4 step 2 comes before §4: an encrypted payload reaches no capability. + if w.manifest.security.as_ref().is_some_and(|s| !s.is_plaintext()) { + return Err(BaseError::Encrypted); + } + let primary = w.capability().to_string(); + if let Some(c) = via { + if c != primary && w.fallback() != Some(c) { + return Err(BaseError::NotOffered(c.to_string())); + } + } + let listing = std::cell::OnceCell::new(); + let mut excluded: Vec = Vec::new(); + // The first refusal that made a capability unsupported, with that capability. + let mut kept: Option<(String, BaseError)> = None; + loop { + let d = w.dispatch_verified( + |c| int_path(c) && via.is_none_or(|v| v == c) && !excluded.iter().any(|x| x == c), + |c, p| held_for(c, Some(p), set, &listing).and_then(|h| resolve(c, Some(p), &h, authority)), + ); + let (cap, params) = match d { + Ok(d) => { + let cap = d.capability.to_string(); + match d.verified { + Some(p) => (cap, p), + // Unpinned (no pin, or one that does not apply): the set the sink + // holds for the capability, with check 5 over what it reads. + None => match held_for(&cap, None, set, &listing).and_then(|h| resolve(&cap, None, &h, authority)) { + Ok(p) => (cap, p), + Err(r) => { + kept.get_or_insert((cap.clone(), BaseError::Pin(r))); + excluded.push(cap); + continue; + } + }, + } + } + Err(Inert::Encrypted { .. }) => return Err(BaseError::Encrypted), + // Step 6: the reason kept for the capability (step 3) before the one kept for + // the fallback (step 5). + Err(Inert::PinRefused { capability, refusal }) => { + return Err(match kept { + Some((c, e)) if c == primary || capability != primary => e, + _ => BaseError::Pin(refusal), + }) + } + Err(Inert::MissingCapability) => return Err(kept.map(|(_, e)| e).unwrap_or(BaseError::Inert)), + }; + if let Some(limit) = limit { + if let Err(e) = check_cost(&cap, &w.payload, ¶ms, limit) { + if routes_on(&e) { + kept.get_or_insert((cap.clone(), BaseError::Decode(e))); + excluded.push(cap); + continue; + } + return Err(BaseError::Decode(e)); + } + } + let decoded = match decode_int(&cap, &w.payload, ¶ms) { + Ok(d) => d, + Err(e) if routes_on(&e) => { + kept.get_or_insert((cap.clone(), BaseError::Decode(e))); + excluded.push(cap); + continue; + } + Err(e) => return Err(BaseError::Decode(e)), + }; + let units = units_from_decoded(decoded).filter(|u| !u.is_empty()).ok_or(BaseError::NoUnits)?; + let pinned = params.is_pinned(); + let info = BaseInfo { + content_hash: crate::jwp_hash::content_hash(envelope), + tier: tier::base_tier(&cap, if pinned { w.prior() } else { None }), + via_fallback: cap != primary, + capability: cap, + attested: false, + first_unit: 0, + }; + return Ok((info, units)); + } +} + /// Why a layer, or a layer for one unit, was not applied: the closed list of /// staged-delivery §4.2, in the order of §4.1's steps. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] @@ -450,6 +711,22 @@ pub struct StageChain { /// The chain's state (§4.5), in bytes. state: u64, seq: u64, + /// Every application, when the sink traces them ([`StageSink::set_trace`]). + trace: Option>, +} + +/// One application of a layer to a unit, as a traced sink records it +/// ([`StageSink::set_trace`]): the unit's state just after it, with its stage hash where +/// the unit holds `decode-equivalence` (§5). +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Applied { + pub unit: u32, + pub stage: u8, + pub producer: [u8; 32], + pub frozen: bool, + pub tier: Option, + pub shape: (u32, u32, u32), + pub sha256: Option<[u8; 32]>, } fn hex(h: &[u8; 32]) -> String { @@ -494,6 +771,7 @@ impl StageChain { held_layers: 0, state, seq: 0, + trace: None, } } @@ -862,6 +1140,11 @@ impl StageChain { if !recorded { self.passed.remove(&u); } + if let Some(t) = self.trace.as_mut() { + let s = &self.units[u as usize]; + let sha256 = (s.tier == Some(DeterminismTier::DecodeEquivalence)).then(|| crate::pin::sha256(&s.buf.samples.canonical_bytes())); + t.push(Applied { unit: u, stage: s.stage, producer: s.producer, frozen: s.frozen, tier: s.tier, shape: s.buf.shape, sha256 }); + } Ok(()) } @@ -1104,6 +1387,8 @@ pub struct StageSink { kept_order: BTreeMap, kept_bytes: u64, tick: u64, + /// Whether every chain records its applications ([`Self::set_trace`]). + trace: bool, } /// What a sink keeps of one base it does not hold. @@ -1125,7 +1410,7 @@ impl StageSink { /// A sink holding no base, with the bound, cost limit, budget and schedule every /// chain it creates uses. pub fn new(cfg: ChainConfig) -> Self { - StageSink { cfg, chains: BTreeMap::new(), pending: vec![], kept: BTreeMap::new(), kept_order: BTreeMap::new(), kept_bytes: 0, tick: 0 } + StageSink { cfg, chains: BTreeMap::new(), pending: vec![], kept: BTreeMap::new(), kept_order: BTreeMap::new(), kept_bytes: 0, tick: 0, trace: false } } /// The chain of the base whose content hash is `base`. @@ -1133,12 +1418,42 @@ impl StageSink { self.chains.get(base) } + /// Record every application of a layer to a unit from now on, in every chain the + /// sink holds or creates, for [`Self::take_trace`]: a batch that takes a unit up + /// several stages in one offer or release records each stage it passes. Off by + /// default; a traced sink hashes each unit it applies a layer to. + pub fn set_trace(&mut self, on: bool) { + self.trace = on; + for c in self.chains.values_mut() { + if on { + c.trace.get_or_insert_with(Vec::new); + } else { + c.trace = None; + } + } + } + + /// The applications recorded since the last call, with their base's content hash: + /// chain by chain in base order, each in the order the sink applied them. + pub fn take_trace(&mut self) -> Vec<([u8; 32], Applied)> { + self.chains + .iter_mut() + .flat_map(|(h, c)| c.trace.as_mut().map(std::mem::take).unwrap_or_default().into_iter().map(move |a| (*h, a))) + .collect() + } + /// Record a decoded base and release the layers held for it, together (§4.4). A /// base the sink already holds is kept as it is: decoding it again does not reset /// its units, and reports nothing; if this decode comes with the base's delivery /// receipt verified, the base is attested from now on (§10), and never the reverse. /// A base the sink attested before, held, dropped or not yet decoded, is attested /// whatever `base.attested` says. + /// + /// A base whose chain would start past the chain budget (§4.5: its units' output and + /// 32 bytes for each stage the records the sink keeps of it hold) is not added: the + /// one report is the base envelope's, refused `chain-over-limit`; the sink keeps no + /// chain for it, the records it keeps of it are untouched, and the layers held for + /// it stay held (§4.3). A receipt verified with it is still recorded (§10). pub fn add_base(&mut self, mut base: BaseInfo, units: Vec) -> Vec { let hash = base.content_hash; if let Some(c) = self.chains.get_mut(&hash) { @@ -1147,10 +1462,23 @@ impl StageSink { } return vec![]; } + let recorded: u64 = self.kept.get(&hash).map_or(0, |k| { + k.records.iter().filter(|(u, _)| (*u as usize) < units.len()).map(|(_, r)| 32 * r.producers.len() as u64).sum() + }); + let state = units.iter().fold(recorded, |s, u| s.saturating_add(u.samples.byte_len())); + if state > self.cfg.budget { + if base.attested { + self.mark_attested(&hash); + } + return vec![OfferReport::refuse(hash, StageRefusal::ChainOverLimit, None)]; + } let kept = self.take_kept(&hash).unwrap_or_default(); base.attested |= kept.attested; let records = kept.records; let mut chain = StageChain::with_records(base, units, self.cfg.clone(), records); + if self.trace { + chain.trace = Some(Vec::new()); + } let hexed = hex(&hash); let (mine, rest): (Vec, Vec) = std::mem::take(&mut self.pending).into_iter().partition(|p| p.refines.base == hexed); self.pending = rest; diff --git a/wai/wai-rs/src/staged/tests.rs b/wai/wai-rs/src/staged/tests.rs index 9e2baf05d..32e022466 100644 --- a/wai/wai-rs/src/staged/tests.rs +++ b/wai/wai-rs/src/staged/tests.rs @@ -1534,3 +1534,253 @@ fn refinement_roles() { tampered.renditions[2][50] ^= 1; assert_eq!(tampered.refinements_of(0).iter().map(|(i, ..)| *i).collect::>(), vec![1]); } + +/// `decode_base` decodes a base for a sink through the integer path: a pinned base +/// whose set the sink holds gives its stage-0 units and `decode-equivalence`; the same +/// payload unpinned gives no tier; a pin whose set is absent or altered is refused +/// before anything is decoded; a `wai.video.int_motion` base needs no set; bytes that +/// are not an envelope are `not-wai1`. The corpus states the units and tiers +/// (`staged-conformance/`), and `wai_staged_vectors verify` replays them. +#[cfg(feature = "video_int")] +#[test] +fn decode_base_reads_the_pin_and_the_corpus_units() { + use super::{decode_base, BaseError}; + use std::collections::BTreeMap; + let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(".."); + let rd = |rel: &str| std::fs::read(wai.join(rel)).unwrap_or_else(|e| panic!("{rel}: {e}")); + let stage0 = |units: &[UnitBuf]| units.iter().flat_map(|u| u.samples.canonical_bytes()).collect::>(); + let set_bytes: Vec<(String, Vec)> = [("entropy.json", "tiny_entropy.json"), ("g_s.whm", "tiny_g_s.bin"), ("h_s.whm", "tiny_h_s.bin")] + .iter() + .map(|(m, f)| (m.to_string(), rd(&format!("wai-rs/tests/vectors/int_hyper/{f}")))) + .collect(); + let set: BTreeMap = set_bytes.iter().map(|(m, b)| (m.clone(), b.as_slice())).collect(); + + let env = rd("staged-conformance/cases/image_int_hyper_2l/base.wai"); + let (info, units) = decode_base(&env, &set, None, None).expect("the pinned base decodes"); + assert_eq!(info.tier, Some(DeterminismTier::DecodeEquivalence)); + assert_eq!(info.capability, crate::codecs::CAP_NEURAL_INT_HYPER); + assert_eq!(info.content_hash, crate::jwp_hash::content_hash(&env)); + assert!(!info.via_fallback && !info.attested && info.first_unit == 0); + assert_eq!(stage0(&units), rd("staged-conformance/cases/image_int_hyper_2l/stage0.bin")); + + assert!(matches!(decode_base(&env, &BTreeMap::new(), None, None), Err(BaseError::Pin(_))), "no set held"); + let mut altered = set_bytes.clone(); + altered[1].1[0] ^= 1; + let altered: BTreeMap = altered.iter().map(|(m, b)| (m.clone(), b.as_slice())).collect(); + assert!(matches!(decode_base(&env, &altered, None, None), Err(BaseError::Pin(_))), "an altered set"); + + let (info, _) = decode_base(&rd("staged-conformance/cases/image_int_hyper_unpinned/base.wai"), &set, None, None).expect("unpinned decodes"); + assert_eq!(info.tier, None, "an unpinned base holds no tier"); + + let env = rd("staged-conformance/cases/video_int_motion_gop_perframe/base_0.wai"); + let (info, units) = decode_base(&env, &BTreeMap::new(), None, None).expect("int_motion reads no set"); + assert_eq!(info.tier, Some(DeterminismTier::DecodeEquivalence)); + assert_eq!(units.len(), 4); + assert_eq!(stage0(&units), rd("staged-conformance/cases/video_int_motion_gop_perframe/stage0_0.bin")); + + let e = decode_base(b"not an envelope", &BTreeMap::new(), None, None).unwrap_err(); + assert_eq!(e.code(), "not-wai1"); + + // A base offered with a fallback, decoded through it, is recorded so; a capability + // the manifest does not offer is refused. + let env = rd("staged-conformance/cases/base_via_fallback/base.wai"); + let (info, _) = super::decode_base_via(&env, Some(crate::codecs::CAP_VIDEO_INT_MOTION), &BTreeMap::new(), None, None).expect("the fallback decodes"); + assert!(info.via_fallback); + assert_eq!(info.capability, crate::codecs::CAP_VIDEO_INT_MOTION); + let e = super::decode_base_via(&env, Some(crate::codecs::CAP_VIDEO_INT_HYPER), &BTreeMap::new(), None, None).unwrap_err(); + assert_eq!(e.code(), "not-offered"); +} + +/// A small `wai.video.int_motion` payload (two 4 × 4 frames) in an envelope naming +/// `cap`, with `fallback`. +#[cfg(feature = "video_int")] +fn motion_envelope(cap: &str, fallback: Option<&str>) -> Vec { + let frames: Vec> = (0..2u32).map(|t| (0..4 * 4 * 3u32).map(|i| ((i * 7 + t * 13) % 256) as u8).collect()).collect(); + let (payload, _) = crate::int_payload::encode_wiv1_zero_motion(&frames, 4, 4, 4).unwrap(); + let mut mr = ModelRequirement::new(cap); + if let Some(f) = fallback { + mr = mr.with_fallback(f); + } + Wai::new(Manifest::new("video", "replicate", mr, Conditioning::new("int_motion")), payload).to_bytes().unwrap() +} + +/// SPEC §4 for a base, as `Wai::dispatch_verified` runs it: a refinement-only or a +/// companion-only capability makes the envelope inert, so its fallback is never +/// reached, by `decode_base` or by `decode_base_via` asked for that fallback +/// (`missing_capability`); a capability this build does not decode moves dispatch on +/// to the fallback, recorded as decoded through it. +#[cfg(feature = "video_int")] +#[test] +fn decode_base_runs_spec_4() { + use super::{decode_base, decode_base_via}; + use crate::codecs::{CAP_VIDEO_INT_MOTION as MOTION, CAP_VIDEO_LCEVC}; + let none = std::collections::BTreeMap::new(); + // The same two envelopes are committed for the wasm32 verifier + // (wai-web/demo/verify_staged_wasm.mjs); they must be what this builds. + // WAI_STAGED_BOUNDS_WRITE=1 rewrites them. + let dir = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/staged_bounds"); + for (inert_primary, file) in [(CAP_VIDEO_INT_REFINE, "inert_refinement_primary.wai"), (CAP_VIDEO_LCEVC, "inert_companion_primary.wai")] { + let env = motion_envelope(inert_primary, Some(MOTION)); + if std::env::var_os("WAI_STAGED_BOUNDS_WRITE").is_some() { + std::fs::write(dir.join(file), &env).unwrap(); + } + assert_eq!(std::fs::read(dir.join(file)).unwrap(), env, "{file} is what this test builds"); + assert_eq!(decode_base(&env, &none, None, None).unwrap_err().code(), "missing_capability", "{inert_primary}"); + assert_eq!(decode_base_via(&env, Some(MOTION), &none, None, None).unwrap_err().code(), "missing_capability", "{inert_primary}"); + } + for (primary, via_fallback) in [(MOTION, false), ("x.unregistered.cap", true), (crate::codecs::CAP_VIDEO_AV1, true)] { + let env = motion_envelope(primary, if via_fallback { Some(MOTION) } else { None }); + let (info, units) = decode_base(&env, &none, None, None).unwrap_or_else(|e| panic!("{primary}: {e}")); + assert_eq!((info.capability.as_str(), info.via_fallback, units.len()), (MOTION, via_fallback, 2), "{primary}"); + } + let env = motion_envelope(crate::codecs::CAP_VIDEO_AV1, None); + assert_eq!(decode_base(&env, &none, None, None).unwrap_err().code(), "missing_capability", "no capability decodes here"); +} + +/// A sink's cost limit applies to a base before it is decoded: a payload whose +/// declared output passes the limit is `cost_over_limit`, at the limit it decodes, and +/// a capability refused for its cost moves dispatch on to the fallback. +#[cfg(feature = "video_int")] +#[test] +fn decode_base_applies_the_cost_limit() { + use super::decode_base; + use crate::int_payload::DeclaredCost; + let none = std::collections::BTreeMap::new(); + let env = motion_envelope(crate::codecs::CAP_VIDEO_INT_MOTION, None); + let out = 2 * 4 * 4 * 3; + let at = DeclaredCost::new(out, u64::MAX, u64::MAX, u128::MAX); + assert!(decode_base(&env, &none, Some(&at), None).is_ok()); + let under = DeclaredCost::new(out - 1, u64::MAX, u64::MAX, u128::MAX); + assert_eq!(decode_base(&env, &none, Some(&under), None).unwrap_err().code(), "cost_over_limit"); + let few_symbols = DeclaredCost::new(u64::MAX, 3, u64::MAX, u128::MAX); + assert_eq!(decode_base(&env, &none, Some(&few_symbols), None).unwrap_err().code(), "cost_over_limit"); +} + +/// SPEC §3.1 check 5 on a base: a key log that revokes the pinned set refuses the base +/// (`prior_revoked`), and one that revokes the set an unpinned base would read refuses +/// that too; with no key log both decode. +#[test] +fn decode_base_runs_check_5() { + use super::decode_base; + use crate::container::{PinRefusal, Prior}; + #[derive(Debug)] + struct Revokes(Vec); + impl crate::pin::PinAuthority for Revokes { + fn check(&self, capability: &str, prior: &Prior) -> Result<(), PinRefusal> { + self.check_unpinned(capability, std::slice::from_ref(&prior.sha256)) + } + fn check_unpinned(&self, capability: &str, digests: &[String]) -> Result<(), PinRefusal> { + if digests.iter().any(|d| self.0.contains(d)) { + return Err(PinRefusal::Revoked { capability: capability.into(), revoked_by: None }); + } + Ok(()) + } + } + let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(".."); + let rd = |rel: &str| std::fs::read(wai.join(rel)).unwrap(); + let set_bytes: Vec<(String, Vec)> = [("entropy.json", "tiny_entropy.json"), ("g_s.whm", "tiny_g_s.bin"), ("h_s.whm", "tiny_h_s.bin")] + .iter() + .map(|(m, f)| (m.to_string(), rd(&format!("wai-rs/tests/vectors/int_hyper/{f}")))) + .collect(); + let set: std::collections::BTreeMap = set_bytes.iter().map(|(m, b)| (m.clone(), b.as_slice())).collect(); + let listing_digest = crate::pin::sha256_hex(&crate::pin::listing_from(&set).unwrap()); + let revoked = Revokes(vec![listing_digest.clone()]); + let pinned = rd("staged-conformance/cases/image_int_hyper_2l/base.wai"); + assert_eq!(Wai::from_bytes(&pinned).unwrap().prior().unwrap().sha256, listing_digest, "the base pins this set"); + assert!(decode_base(&pinned, &set, None, None).is_ok()); + assert_eq!(decode_base(&pinned, &set, None, Some(&revoked)).unwrap_err().code(), "prior_revoked"); + assert!(decode_base(&pinned, &set, None, Some(&Revokes(vec![]))).is_ok()); + let unpinned = rd("staged-conformance/cases/image_int_hyper_unpinned/base.wai"); + assert!(decode_base(&unpinned, &set, None, None).is_ok()); + assert_eq!(decode_base(&unpinned, &set, None, Some(&revoked)).unwrap_err().code(), "prior_revoked"); +} + +/// A capability whose set is one file, pinned without a form, resolves against +/// whichever held file satisfies the pin (SPEC §3.1), not the first by name: a README +/// beside the model does not stop the base decoding, pinned or unpinned. +#[cfg(feature = "audio_int")] +#[test] +fn decode_base_one_file_pin_ignores_other_files() { + use super::decode_base; + let wai = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join(".."); + let rd = |rel: &str| std::fs::read(wai.join(rel)).unwrap(); + let model = rd("wai-rs/tests/vectors/int_audio/audio_model.bin"); + let readme = b"notes about this model".to_vec(); + let env = rd("staged-conformance/cases/audio_int_codec_1l/base.wai"); + let w = Wai::from_bytes(&env).unwrap(); + assert!(w.prior().is_some_and(|p| p.form().is_none()), "the corpus base pins one file, form-less"); + let set: std::collections::BTreeMap = [("README".to_string(), readme.as_slice()), ("model.wam".to_string(), model.as_slice())].into_iter().collect(); + let (info, units) = decode_base(&env, &set, None, None).expect("the model satisfies the pin"); + assert_eq!(info.tier, Some(DeterminismTier::DecodeEquivalence)); + let stage0: Vec = units.iter().flat_map(|u| u.samples.canonical_bytes()).collect(); + assert_eq!(stage0, rd("staged-conformance/cases/audio_int_codec_1l/stage0.bin")); + let mut unpinned = Wai::from_bytes(&env).unwrap(); + unpinned.manifest.conditioning.prior = None; + assert!(decode_base(&unpinned.to_bytes().unwrap(), &set, None, None).is_ok(), "unpinned, the member is read by name"); +} + +/// §4.5 for a base: a base whose units' output passes the chain budget is not added; +/// its envelope is refused `chain-over-limit`, the sink keeps no chain, and a layer held +/// for it stays held, to the hold bound. A base exactly at the budget is added. +#[test] +fn add_base_counts_its_state_against_the_budget() { + let (_env, info, units) = base4(); + let out: u64 = units.iter().map(|u| u.samples.byte_len()).sum(); + let l1 = layer(&info.content_hash, one(&info.content_hash), 1, Some(Span { from: 0, count: 1 }), &[frame(0)], 50); + let mut small = StageSink::new(ChainConfig { budget: out - 1, ..cfg() }); + assert!(small.offer(&l1, 0, ReceiptEvidence::None).unwrap()[0].held_for_base); + let r = small.add_base(info.clone(), units.clone()); + assert_eq!(r.len(), 1); + assert_eq!((r[0].content_hash, r[0].envelope), (info.content_hash, Some(StageRefusal::ChainOverLimit))); + assert!(small.chain(&info.content_hash).is_none()); + let expired = small.expire(2000); + assert_eq!(expired.iter().map(|r| r.envelope).collect::>(), vec![Some(StageRefusal::BaseAbsent)], "the layer stayed held for its base"); + let mut room = StageSink::new(ChainConfig { budget: out, ..cfg() }); + assert!(room.add_base(info.clone(), units).is_empty(), "a base exactly at the budget is added"); + assert_eq!(room.chain(&info.content_hash).map(|c| c.state_bytes()), Some(out)); +} + +/// `decode_base` refuses in its stated order: `not-wai1`, then `encrypted` (SPEC §6.4 +/// before §4), then `not-offered`, then §4; and a parameter set is read only for a +/// capability that reads one, so a stray file name SPEC §3.1 refuses is no refusal for +/// `wai.video.int_motion`. +#[cfg(feature = "video_int")] +#[test] +fn decode_base_refuses_in_its_stated_order() { + use super::{decode_base, decode_base_via}; + use crate::codecs::{CAP_VIDEO_INT_HYPER, CAP_VIDEO_INT_MOTION as MOTION}; + let none = std::collections::BTreeMap::new(); + let plain = motion_envelope(MOTION, None); + let mut w = Wai::from_bytes(&plain).unwrap(); + w.manifest.security = Some(crate::container::Security::new("aes-256-gcm", 1)); + let enc = w.to_bytes().unwrap(); + let bad: std::collections::BTreeMap = [("../x".to_string(), b"a".as_slice())].into_iter().collect(); + assert_eq!(decode_base(b"WAI1", &bad, None, None).unwrap_err().code(), "not-wai1"); + assert_eq!(decode_base(&enc, &none, None, None).unwrap_err().code(), "encrypted"); + assert_eq!(decode_base_via(&enc, Some(CAP_VIDEO_INT_HYPER), &none, None, None).unwrap_err().code(), "encrypted", "encrypted before not-offered"); + assert_eq!(decode_base(&enc, &bad, None, None).unwrap_err().code(), "encrypted", "encrypted before the set is read"); + assert_eq!(decode_base_via(&plain, Some(CAP_VIDEO_INT_HYPER), &bad, None, None).unwrap_err().code(), "not-offered", "not-offered before the set is read"); + let (info, _) = decode_base(&plain, &bad, None, None).expect("int_motion reads no set, so a stray name is no refusal"); + assert_eq!(info.capability, MOTION); + let notes: std::collections::BTreeMap = [("notes.txt".to_string(), b"a".as_slice())].into_iter().collect(); + assert!(decode_base(&plain, ¬es, None, None).is_ok()); +} + +/// The cost limit binds a base whose capability declares a cost before decoding (the +/// integer payloads SPEC §5 pins); `wai.neural.int_synth` declares none, so it decodes +/// under its format's caps whatever the limit, as the neural registry's limit does. +#[test] +fn decode_base_limit_binds_declared_costs_only() { + use super::decode_base; + let v = std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("../neural-conformance/vectors/gradient"); + let latent = std::fs::read(v.join("latent.bin")).unwrap(); + let spec = std::fs::read(v.join("int_synth.json")).unwrap(); + let m = Manifest::new("image", "replicate", ModelRequirement::new(crate::codecs::CAP_NEURAL_INT_SYNTH), Conditioning::new("int_synth")); + let env = Wai::new(m, latent).to_bytes().unwrap(); + let set: std::collections::BTreeMap = [("int_synth.json".to_string(), spec.as_slice())].into_iter().collect(); + let zero = crate::int_payload::DeclaredCost::new(0, 0, 0, 0); + let (info, units) = decode_base(&env, &set, Some(&zero), None).expect("int_synth declares no cost"); + assert_eq!(info.capability, crate::codecs::CAP_NEURAL_INT_SYNTH); + assert!(units.iter().map(|u| u.samples.byte_len()).sum::() > 0); + assert_eq!(crate::int_payload::check_cost(crate::codecs::CAP_NEURAL_INT_SYNTH, &[], &crate::int_payload::IntParams::none(), &zero).unwrap(), None, "no declared cost to check"); +} diff --git a/wai/wai-rs/tests/cli_stage.rs b/wai/wai-rs/tests/cli_stage.rs index a2a2dcc49..13e24f393 100644 --- a/wai/wai-rs/tests/cli_stage.rs +++ b/wai/wai-rs/tests/cli_stage.rs @@ -137,3 +137,403 @@ fn wrap_refuses_a_refinement_capability() { assert!(text(&o).contains("refinement-only"), "{}", text(&o)); assert!(!out.exists()); } + +fn sha256_hex(b: &[u8]) -> String { + use sha2::{Digest, Sha256}; + Sha256::digest(b).iter().map(|x| format!("{x:02x}")).collect() +} + +/// `wai stage base` writes a `wai.video.int_motion` base that decodes; `wai stage +/// refine` at q = 1 over it writes a layer per frame; `wai stage apply --json` traces +/// the base's units at stage 0, then each layer's unit at stage 1 with the SHA-256 of +/// the source frame, in the corpus's report form. +#[test] +fn base_refine_apply_json_trace() { + let (n, h, w) = (3usize, 6usize, 10usize); + let per = h * w * 3; + let frames: Vec = (0..n * per).map(|i| ((i * 13 + (i / per) * 50) % 256) as u8).collect(); + let fpath = scratch("frames.rgb"); + std::fs::write(&fpath, &frames).unwrap(); + let bpath = scratch("gop.wai"); + let o = wai(&[os("stage"), os("base"), os("--frames"), fpath.as_os_str(), os("--dims"), os("3,6,10"), os("--q"), os("40"), os("-o"), bpath.as_os_str()]); + let t = text(&o); + assert!(o.status.success(), "{t}"); + assert!(t.contains("wai.video.int_motion") && t.contains("tier decode-equivalence"), "{t}"); + let mut layers = Vec::new(); + for f in 0..n { + let tpath = scratch(&format!("t{f}.rgb")); + std::fs::write(&tpath, &frames[f * per..(f + 1) * per]).unwrap(); + let lpath = scratch(&format!("L1_f{f}.wai")); + let span = format!("{f}:1"); + let o = wai(&[os("stage"), os("refine"), os("--base"), bpath.as_os_str(), os("--span"), os(&span), os("--target"), tpath.as_os_str(), os("--q"), os("1"), os("-o"), lpath.as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + layers.push(lpath); + } + let mut args = vec![os("stage"), os("apply"), os("--json"), bpath.as_os_str()]; + // Out of order: the last frame's layer first. + args.extend(layers.iter().rev().map(|p| p.as_os_str())); + let o = wai(&args); + assert!(o.status.success(), "{}", text(&o)); + let doc: serde_json::Value = serde_json::from_slice(&o.stdout).expect("one JSON document"); + assert_eq!(doc["format"], "wai-stage-apply/1"); + let steps = doc["steps"].as_array().unwrap(); + assert_eq!(steps.len(), 1 + n); + assert_eq!(steps[0]["op"], "base"); + assert_eq!(steps[0]["units"].as_array().unwrap().len(), n, "a base's step lists every unit"); + for u in steps[0]["units"].as_array().unwrap() { + assert_eq!((u["stage"].as_u64(), u["tier"].as_str()), (Some(0), Some("decode-equivalence"))); + } + for (i, s) in steps[1..].iter().enumerate() { + let f = n - 1 - i; + assert_eq!(s["op"], "offer"); + let r = &s["reports"][0]; + assert_eq!(r["applied"], serde_json::json!([f]), "{s}"); + assert_eq!(r["envelope"], serde_json::Value::Null); + let changed = s["units"].as_array().unwrap(); + assert_eq!(changed.len(), 1); + assert_eq!(changed[0]["unit"].as_u64(), Some(f as u64)); + assert_eq!(changed[0]["stage"].as_u64(), Some(1)); + assert_eq!(changed[0]["producer"], r["object"]); + assert_eq!(changed[0]["stage_sha256"].as_str().unwrap(), sha256_hex(&frames[f * per..(f + 1) * per]), "q = 1 reaches the source"); + } + assert!(doc["bases"][0]["clip_sha256"].is_string()); +} + +/// `wai stage apply` takes several bases: a layer whose base comes later is held for +/// it and released when it arrives, and the corpus's three-GOP clip reaches the clip +/// hash `index.json` states for each GOP. +#[test] +fn apply_holds_layers_for_a_later_base() { + let d = corpus("cases/video_int_motion_gop_perframe"); + let index: serde_json::Value = serde_json::from_slice(&std::fs::read(corpus("index.json")).unwrap()).unwrap(); + let want = index["cases"].as_array().unwrap().iter().find(|c| c["name"] == "video_int_motion_gop_perframe").unwrap()["clip_sha256"].clone(); + let mut files: Vec = vec![d.join("base_0.wai")]; + // Every layer before bases 1 and 2. + for g in 0..3 { + files.extend((0..4).map(|f| d.join(format!("g{g}_L1_f{f}.wai")))); + files.push(d.join(format!("g{g}_L2.wai"))); + } + files.push(d.join("base_1.wai")); + files.push(d.join("base_2.wai")); + let mut args = vec![os("stage"), os("apply"), os("--json")]; + args.extend(files.iter().map(|p| p.as_os_str())); + let o = wai(&args); + assert!(o.status.success(), "{}", text(&o)); + let doc: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap(); + let steps = doc["steps"].as_array().unwrap(); + let held: usize = steps.iter().filter(|s| s["reports"].as_array().is_some_and(|r| r.iter().any(|r| r["held_for_base"] == true))).count(); + assert_eq!(held, 10, "GOPs 1 and 2's layers wait for their bases"); + let last = &steps[steps.len() - 1]; + assert_eq!(last["op"], "base"); + assert_eq!(last["reports"].as_array().unwrap().len(), 5, "base 2 releases its five layers"); + // Every stage reached in that one step is listed: each unit at stage 0, then at 1, + // then at 2, the last equal to the base's final unit. + let b2 = last["object"].as_str().unwrap(); + let listed = last["units"].as_array().unwrap(); + for u in 0..4u64 { + let stages: Vec = listed.iter().filter(|x| x["base"] == b2 && x["unit"] == u).map(|x| x["stage"].as_u64().unwrap()).collect(); + assert_eq!(stages, vec![0, 1, 2], "unit {u}: {listed:?}"); + let fin = doc["bases"].as_array().unwrap().iter().find(|b| b["object"] == b2).unwrap()["units"][u as usize]["stage_sha256"].clone(); + let at2 = listed.iter().find(|x| x["base"] == b2 && x["unit"] == u && x["stage"] == 2).unwrap()["stage_sha256"].clone(); + assert_eq!(at2, fin); + } + assert!(listed.iter().all(|x| x["stage_sha256"].is_string()), "every stage reached has its hash"); + let got: Vec = doc["bases"].as_array().unwrap().iter().map(|b| b["clip_sha256"].clone()).collect(); + assert_eq!(serde_json::Value::Array(got), want); +} + +/// `wai_stage_bench` measures a set made with `wai stage`: stage 0 is the base, +/// stage 1 its layers; the bytes are the envelopes' and sum to them; a layer at q = 1 +/// reaches the source, which the report states as exact (no finite PSNR) with each +/// unit's stage hash that of its source frame; the direct comparator at the same step +/// is exact too; fewer than 20 runs are refused. +#[test] +fn stage_bench_reports_bytes_quality_and_hashes() { + let (n, h, w) = (2usize, 5usize, 7usize); + let per = h * w * 3; + let frames: Vec = (0..n * per).map(|i| ((i * 29 + (i / per) * 70) % 256) as u8).collect(); + let fpath = scratch("bench_frames.rgb"); + std::fs::write(&fpath, &frames).unwrap(); + let bpath = scratch("bench_base.wai"); + let o = wai(&[os("stage"), os("base"), os("--frames"), fpath.as_os_str(), os("--dims"), os("2,5,7"), os("--q"), os("30"), os("-o"), bpath.as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + let lpath = scratch("bench_L1.wai"); + let o = wai(&[os("stage"), os("refine"), os("--base"), bpath.as_os_str(), os("--target"), fpath.as_os_str(), os("--q"), os("1"), os("-o"), lpath.as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + let out = scratch("bench.json"); + let bench = |runs: &str| Command::new(env!("CARGO_BIN_EXE_wai_stage_bench")) + .args([os("--source"), fpath.as_os_str(), os("--runs"), os(runs), os("--out"), out.as_os_str(), bpath.as_os_str(), lpath.as_os_str()]) + .output() + .unwrap(); + assert!(!bench("5").status.success(), "fewer than 20 runs are refused"); + let o = bench("20"); + assert!(o.status.success(), "{}", text(&o)); + let r: serde_json::Value = serde_json::from_slice(&std::fs::read(&out).unwrap()).unwrap(); + assert_eq!(r["format"], "wai-stage-bench/1"); + assert_eq!(r["conditions"]["energy"], "unmetered"); + let s = r["stages"].as_array().unwrap(); + assert_eq!(s.len(), 2); + let sizes = [std::fs::metadata(&bpath).unwrap().len(), std::fs::metadata(&lpath).unwrap().len()]; + for (k, st) in s.iter().enumerate() { + assert_eq!(st["bytes"].as_u64(), Some(sizes[k])); + assert_eq!(st["framing_bytes"].as_u64().unwrap() + st["manifest_bytes"].as_u64().unwrap() + st["payload_bytes"].as_u64().unwrap(), sizes[k]); + assert!(st["time_us"]["median"].as_f64().is_some()); + } + assert_eq!(s[1]["cumulative_bytes"].as_u64(), Some(sizes[0] + sizes[1])); + assert!(s[0]["clip"]["psnr_db"].as_f64().is_some_and(|db| db > 0.0), "the base is lossy at q = 30"); + assert_eq!((s[1]["clip"]["exact"].as_bool(), s[1]["clip"]["psnr_db"].is_null()), (Some(true), true)); + for (f, u) in s[1]["units"].as_array().unwrap().iter().enumerate() { + assert_eq!(u["stage"].as_u64(), Some(1)); + assert_eq!(u["stage_sha256"].as_str().unwrap(), sha256_hex(&frames[f * per..(f + 1) * per])); + } + assert_eq!((s[1]["direct"]["q"].as_u64(), s[1]["direct"]["exact"].as_bool()), (Some(1), Some(true))); +} + +/// `wai stage apply --json` is the corpus's sink: every case it can replay (bases +/// decoded through the integer path or, one given first, read with `--stage0`, offers +/// with no receipt, no expiry, removal or attestation, at the corpus's default +/// configuration, since the tool's sink applies no budget or bound, or at a budget that +/// never binds) gives, step by step, the reports `case.json` states, with a step +/// whose bytes are not a `WAI1` this sink reads as that step's error, and ends with the +/// final units it states. Ported from the review's cli_vs_corpus.py. +#[test] +fn apply_json_replays_the_corpus() { + let root = corpus("cases"); + let mut names: Vec = std::fs::read_dir(&root).unwrap().flatten().map(|e| e.file_name().to_string_lossy().into_owned()).collect(); + names.sort(); + let default_cfg = serde_json::json!({"bound": {"max_envelopes": 16, "max_age_ms": 1000}, "cost_limit": null, "budget": 1u64 << 30, "record_budget": 1u64 << 20, "live": null}); + let (mut same, mut skipped, mut extra) = (0, 0, 0); + let mut differ: Vec = Vec::new(); + for name in names { + let d = root.join(&name); + let c = case(&name); + let steps = c["steps"].as_array().unwrap(); + let bases = c["bases"].as_array().unwrap(); + let sets: std::collections::BTreeSet = bases.iter().map(|b| b["decode"]["set"].to_string()).collect(); + // A base the tool cannot decode (PNG, FLAC) is read from its stage0.bin with + // --stage0, where the case has that one base, given first, in units of one shape. + let stage0_base = bases.len() == 1 + && bases[0]["decode"]["kind"] != "int" + && steps[0]["op"] == "base" + && steps.iter().filter(|s| s["op"] == "base").count() == 1 + && bases[0]["units"].as_array().unwrap().windows(2).all(|w| w[0]["shape"] == w[1]["shape"]); + // budget_partly_applied states a budget with no byte to spare that never refuses + // anything: the tool's sink, with no budget, reports what it reports. + let config_ok = c["config"] == default_cfg || name == "budget_partly_applied"; + let skip = steps.iter().any(|s| !["base", "offer"].contains(&s["op"].as_str().unwrap())) + || steps.iter().any(|s| s.get("evidence").is_some_and(|e| e != "none")) + || bases.iter().any(|b| b["attested"] == true || b["first_unit"] != 0 || (b["decode"]["kind"] != "int" && !stage0_base)) + || !config_ok + || sets.len() > 1; + if skip { + skipped += 1; + continue; + } + if stage0_base || name == "budget_partly_applied" { + extra += 1; + } + // The parameter set, by the names SPEC §5 gives its files. + let set: serde_json::Map = bases[0]["decode"]["set"].as_object().cloned().unwrap_or_default(); + let set_dir = Path::new(env!("CARGO_TARGET_TMPDIR")).join("cli_stage/corpus_sets").join(&name); + let mut args: Vec = vec!["stage".into(), "apply".into(), "--json".into()]; + if !set.is_empty() { + std::fs::create_dir_all(&set_dir).unwrap(); + for (m, p) in &set { + std::fs::copy(wai_dir().join(p.as_str().unwrap()), set_dir.join(m)).unwrap(); + } + args.push("--set".into()); + args.push(set_dir.clone().into()); + } + if stage0_base { + let b = &bases[0]; + let shape: Vec = b["units"][0]["shape"].as_array().unwrap().iter().map(|x| x.to_string()).collect(); + args.extend(["--stage0".into(), d.join(b["stage0"].as_str().unwrap()).into_os_string(), "--dims".into(), shape.join(",").into(), + "--units".into(), b["units"].as_array().unwrap().len().to_string().into()]); + if b["units"][0]["elem"] == "i16" { + args.extend(["--elem".into(), "i16".into()]); + } + } + for s in steps { + let f = if s["op"] == "base" { bases[s["index"].as_u64().unwrap() as usize]["file"].as_str().unwrap() } else { s["file"].as_str().unwrap() }; + args.push(d.join(f).into()); + } + let refs: Vec<&std::ffi::OsStr> = args.iter().map(|a| a.as_os_str()).collect(); + let o = wai(&refs); + if !o.status.success() { + differ.push(format!("{name}: the tool fails: {}", text(&o))); + continue; + } + let doc: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap(); + // A claimed base's chain says so (--stage0); the mark is set aside to compare, + // and must be on every report of that chain. + let got: Vec = doc["steps"].as_array().unwrap().iter() + .map(|st| if let Some(e) = st.get("error") { serde_json::json!({"error": e}) } else { + serde_json::Value::Array(st["reports"].as_array().unwrap().iter().map(|r| { + let mut r = r.clone(); + let mark = r.as_object_mut().unwrap().remove("claimed"); + if stage0_base && mark != Some(serde_json::json!(true)) { + differ.push(format!("{name}: a report of the claimed chain is not marked")); + } + r + }).collect()) + }) + .collect(); + if serde_json::Value::Array(got.clone()) != c["expect"]["reports"] { + let i = got.iter().zip(c["expect"]["reports"].as_array().unwrap()).position(|(g, w)| g != w).unwrap_or(got.len()); + differ.push(format!("{name}: step {i}: the tool {} / the case {}", got.get(i).map(|v| v.to_string()).unwrap_or_default(), c["expect"]["reports"][i])); + continue; + } + let finals: std::collections::BTreeMap = + doc["bases"].as_array().unwrap().iter().map(|b| (b["object"].as_str().unwrap().to_string(), b)).collect(); + for (i, b) in bases.iter().enumerate() { + let want = &c["expect"]["bases"][i]; + let got = finals.get(b["object"].as_str().unwrap()); + match (want.is_null(), got) { + (true, None) => {} + (false, Some(g)) => { + let units: Vec = g["units"].as_array().unwrap().iter().map(|u| { + serde_json::json!({"stage": u["stage"], "producer": u["producer"], "frozen": u["frozen"], "tier": u["tier"], "shape": u["shape"], "stage_sha256": u["stage_sha256"]}) + }).collect(); + if serde_json::Value::Array(units) != want["units"] || g["clip_sha256"] != want["clip_sha256"] { + differ.push(format!("{name}: base {i}'s final units differ")); + } + } + _ => differ.push(format!("{name}: base {i} held {} by the tool", if got.is_some() { "" } else { "not" })), + } + } + same += 1; + } + println!("apply --json against the corpus: {same} cases agree ({extra} through --stage0 or a budget that never binds), {} differ, {skipped} need what the tool's sink does not do", differ.len()); + assert!(differ.is_empty(), "{}", differ.join("\n")); + assert!(same >= 30, "only {same} cases compared"); + assert_eq!(extra, 5, "the four PNG and FLAC cases and budget_partly_applied replay"); +} + +/// `wai_stage_bench` on a set whose layer 1 changes the units' shape (up2) and whose +/// layer 2 is identity: the source is taken at the units' final shape, so stage 0 has +/// no quality figure (its units are half the size), stages 1 and 2 are measured against +/// it, and the direct comparator says why it does not apply (an up2 stage; an identity +/// stage whose stage 0 has another shape); a source sized for stage 0 is refused. +#[test] +fn stage_bench_takes_the_source_at_the_final_shape() { + let (n, h, w) = (2usize, 4usize, 6usize); + let small: Vec = (0..n * h * w * 3).map(|i| ((i * 11) % 256) as u8).collect(); + let big: Vec = (0..n * 4 * h * w * 3).map(|i| ((i * 5 + 40) % 256) as u8).collect(); + let (sp, bp) = (scratch("up2_small.rgb"), scratch("up2_big.rgb")); + std::fs::write(&sp, &small).unwrap(); + std::fs::write(&bp, &big).unwrap(); + let base = scratch("up2_base.wai"); + let o = wai(&[os("stage"), os("base"), os("--frames"), sp.as_os_str(), os("--dims"), os("2,4,6"), os("--q"), os("16"), os("-o"), base.as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + let l1 = scratch("up2_L1.wai"); + let o = wai(&[os("stage"), os("refine"), os("--base"), base.as_os_str(), os("--predictor"), os("up2"), os("--target"), bp.as_os_str(), os("--q"), os("1"), os("-o"), l1.as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + let l2 = scratch("up2_L2.wai"); + let o = wai(&[os("stage"), os("refine"), os("--base"), base.as_os_str(), os("--on"), l1.as_os_str(), os("--target"), bp.as_os_str(), os("--q"), os("1"), os("-o"), l2.as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + let out = scratch("up2_bench.json"); + let bench = |src: &Path| Command::new(env!("CARGO_BIN_EXE_wai_stage_bench")) + .args([os("--source"), src.as_os_str(), os("--runs"), os("20"), os("--out"), out.as_os_str(), base.as_os_str(), l1.as_os_str(), l2.as_os_str()]) + .output() + .unwrap(); + let o = bench(&sp); + assert!(!o.status.success() && text(&o).contains("last stage"), "a source at stage 0's shape is refused: {}", text(&o)); + let o = bench(&bp); + assert!(o.status.success(), "{}", text(&o)); + let r: serde_json::Value = serde_json::from_slice(&std::fs::read(&out).unwrap()).unwrap(); + let s = r["stages"].as_array().unwrap(); + assert_eq!(s[0]["clip"]["units_measured"], 0); + assert!(s[0]["clip"]["psnr_db"].is_null() && s[0]["units"].as_array().unwrap().iter().all(|u| u["why"].is_string())); + assert_eq!((s[1]["clip"]["units_measured"].as_u64(), s[1]["clip"]["exact"].as_bool()), (Some(2), Some(true))); + assert!(s[1]["direct"]["why"].as_str().is_some_and(|w| w.contains("identity predictor")), "{}", s[1]["direct"]); + // Layer 2 is identity, but stage 0 has another shape than the source: the identity + // comparator from stage 0 cannot reach it. + assert!(s[2]["direct"]["why"].as_str().is_some_and(|w| w.contains("shape")), "{}", s[2]["direct"]); + assert_eq!(s[2]["clip"]["units_measured"], 2); +} + +/// `wai stage base` refuses dimensions whose product overflows, with a message, not a +/// panic. +#[test] +fn stage_base_refuses_overflowing_dims() { + let f = scratch("tiny.rgb"); + std::fs::write(&f, [0u8; 3]).unwrap(); + let out = scratch("never.wai"); + for dims in ["2,4294967295,4294967295", "4294967295,65536,65536"] { + let o = wai(&[os("stage"), os("base"), os("--frames"), f.as_os_str(), os("--dims"), os(dims), os("--q"), os("4"), os("-o"), out.as_os_str()]); + let t = text(&o); + assert_eq!(o.status.code(), Some(2), "{dims}: {t}"); + assert!(t.contains("overflows") && !t.contains("panicked"), "{dims}: {t}"); + assert!(!out.exists()); + } +} + + +/// `--stage0` publishes a claimed base: the units are the file's bytes, not a decode, +/// so the base's step, each of its units, every report of its chain and its final entry +/// say `"claimed": true`, and the stage-0 hashes are those of the file's bytes. +#[test] +fn apply_stage0_marks_the_base_claimed() { + let d = corpus("cases/out_of_order"); + // 384 bytes that are no decode of anything: two 8 × 8 RGB8 units. + let mut x = 0x9e37_79b9u32; + let junk: Vec = (0..384).map(|_| { x ^= x << 13; x ^= x >> 17; x ^= x << 5; x as u8 }).collect(); + let jpath = scratch("junk384.bin"); + std::fs::write(&jpath, &junk).unwrap(); + let o = wai(&[os("stage"), os("apply"), os("--json"), os("--stage0"), jpath.as_os_str(), os("--units"), os("2"), os("--dims"), os("8,8,3"), + d.join("base.wai").as_os_str(), d.join("L1.wai").as_os_str()]); + assert!(o.status.success(), "{}", text(&o)); + let doc: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap(); + let steps = doc["steps"].as_array().unwrap(); + assert_eq!(steps[0]["claimed"], true); + let s0 = steps[0]["units"].as_array().unwrap(); + assert_eq!(s0.len(), 2); + for (u, e) in s0.iter().enumerate() { + assert_eq!(e["claimed"], true); + assert_eq!(e["stage_sha256"].as_str().unwrap(), sha256_hex(&junk[u * 192..(u + 1) * 192]), "the hash is the file's bytes'"); + } + assert_eq!(steps[1]["claimed"], true, "the layer is on the claimed chain"); + assert!(steps[1]["reports"].as_array().unwrap().iter().all(|r| r["claimed"] == true)); + assert!(steps[1]["units"].as_array().unwrap().iter().all(|u| u["claimed"] == true)); + let b = &doc["bases"][0]; + assert_eq!(b["claimed"], true); + assert!(b["units"].as_array().unwrap().iter().all(|u| u["claimed"] == true)); + // Without --stage0 the decoded base claims nothing. + let o = wai(&[os("stage"), os("apply"), os("--json"), d.join("base.wai").as_os_str(), d.join("L1.wai").as_os_str()]); + let doc: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap(); + assert!(!o.stdout.windows(9).any(|w| w == b"\"claimed\""), "{doc}"); +} + +/// A base the tool cannot decode is that step's error, by its code, and the replay goes +/// on: an encrypted copy of a base (`encrypted`) and an inert envelope +/// (`missing_capability`); the layer for the encrypted base's plaintext is held for it. +#[test] +fn apply_records_a_refused_base_and_goes_on() { + let d = corpus("cases/out_of_order"); + let mut w = wai::container::Wai::from_bytes(&std::fs::read(d.join("base.wai")).unwrap()).unwrap(); + w.manifest.security = Some(wai::container::Security::new("aes-256-gcm", 1)); + let enc = scratch("encrypted_base.wai"); + std::fs::write(&enc, w.to_bytes().unwrap()).unwrap(); + let inert = wai_dir().join("wai-rs/tests/vectors/staged_bounds/inert_companion_primary.wai"); + let l1 = d.join("L1.wai"); + for json in [true, false] { + let mut args = vec![os("stage"), os("apply")]; + if json { + args.push(os("--json")); + } + args.extend([enc.as_os_str(), inert.as_os_str(), l1.as_os_str()]); + let o = wai(&args); + let t = text(&o); + assert!(o.status.success(), "{t}"); + if json { + let doc: serde_json::Value = serde_json::from_slice(&o.stdout).unwrap(); + let steps = doc["steps"].as_array().unwrap(); + assert_eq!((steps[0]["op"].as_str(), steps[0]["error"].as_str()), (Some("base"), Some("encrypted"))); + assert_eq!(steps[1]["error"], "missing_capability"); + assert_eq!(steps[2]["reports"][0]["held_for_base"], true); + assert!(doc["bases"].as_array().unwrap().is_empty()); + } else { + assert!(t.contains("encrypted") && t.contains("missing_capability") && t.contains("held for its base"), "{t}"); + } + } +} diff --git a/wai/wai-rs/tests/int_kernels_vectors.rs b/wai/wai-rs/tests/int_kernels_vectors.rs index fcccbec42..a7da2a40f 100644 --- a/wai/wai-rs/tests/int_kernels_vectors.rs +++ b/wai/wai-rs/tests/int_kernels_vectors.rs @@ -226,6 +226,7 @@ fn every_decoded_output_states_a_sha256() { ("staged-conformance", "each case.json states the SHA-256 of every stage-0 buffer and of every unit's stage (SPEC §7 Digests)"), ("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 that declare clips past a sink's cost limit, refused before any symbol is decoded, and inert envelopes: nothing here is decoded"), ]; let wai = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(".."); let mut dirs: Vec<(String, PathBuf)> = Vec::new(); diff --git a/wai/wai-rs/tests/staged_resources.rs b/wai/wai-rs/tests/staged_resources.rs index b89c75913..b91bff0e6 100644 --- a/wai/wai-rs/tests/staged_resources.rs +++ b/wai/wai-rs/tests/staged_resources.rs @@ -4,8 +4,10 @@ //! //! - Work per layer is linear in its units: an offer of a layer over 4096, 16384 and //! 65536 units (one layer alone; a layer whose `prev` is a run per unit, held; a -//! layer held and then released by the layer below it) allocates, holds and takes -//! time in proportion to the units, within a margin a quadratic sink is far outside. +//! layer held and then released by the layer below it) allocates, holds and spends +//! CPU time in proportion to the units, within a margin a quadratic sink is far +//! outside. The time is the thread's CPU time, not wall-clock time, so a loaded host +//! does not decide it. //! - The chain budget: a layer that would take the chain's state past it is refused //! (`chain-over-limit`) before any stream is decoded, and the state the chain //! reports is the sum §4.5 defines. @@ -19,7 +21,7 @@ use std::alloc::{GlobalAlloc, Layout, System}; use std::sync::atomic::{AtomicUsize, Ordering::Relaxed}; -use std::time::{Duration, Instant}; +use std::time::Duration; use wai::codecs::{DeterminismTier, CAP_VIDEO_INT_MOTION, CAP_VIDEO_INT_REFINE}; use wai::container::{Conditioning, Manifest, ModelRequirement, Prev, PrevRun, Refines, Wai}; @@ -66,7 +68,7 @@ unsafe impl GlobalAlloc for Count { static ALLOC: Count = Count; /// What a window measured: bytes allocated, allocations, the most bytes live at once -/// beyond the start, and the time. +/// beyond the start, and the time the calling thread spent on the CPU. #[derive(Debug, Clone, Copy)] struct Window { allocated: usize, @@ -75,13 +77,45 @@ struct Window { time: Duration, } +/// The calling thread's CPU time. Wall-clock time also counts the time other processes +/// held the core, so on a loaded host it grows with the load rather than the work: a +/// CI runner shared with other jobs made a linear offer look 64–183× slower for 16× the +/// units. The thread's own CPU time grows only with the work it did. The offer runs on +/// the calling thread, so its CPU time is the offer's. +#[cfg(unix)] +fn now() -> Duration { + // `time_t` and `long` are 64 bits on every 64-bit Unix target. + #[repr(C)] + struct Timespec { + tv_sec: i64, + tv_nsec: i64, + } + unsafe extern "C" { + fn clock_gettime(clock: i32, tp: *mut Timespec) -> i32; + } + #[cfg(target_vendor = "apple")] + const CLOCK_THREAD_CPUTIME_ID: i32 = 16; + #[cfg(not(target_vendor = "apple"))] + const CLOCK_THREAD_CPUTIME_ID: i32 = 3; + let mut ts = Timespec { tv_sec: 0, tv_nsec: 0 }; + assert_eq!(unsafe { clock_gettime(CLOCK_THREAD_CPUTIME_ID, &mut ts) }, 0, "clock_gettime"); + Duration::new(ts.tv_sec as u64, ts.tv_nsec as u32) +} + +/// Elsewhere, wall-clock time from a fixed origin. +#[cfg(not(unix))] +fn now() -> Duration { + static ORIGIN: std::sync::OnceLock = std::sync::OnceLock::new(); + ORIGIN.get_or_init(std::time::Instant::now).elapsed() +} + fn measure(f: impl FnOnce() -> T) -> (T, Window) { let start = CUR.load(Relaxed); PEAK.store(start, Relaxed); let (a0, c0) = (ALLOCATED.load(Relaxed), CALLS.load(Relaxed)); - let t = Instant::now(); + let t = now(); let out = f(); - let time = t.elapsed(); + let time = now() - t; let w = Window { allocated: ALLOCATED.load(Relaxed) - a0, calls: CALLS.load(Relaxed) - c0, peak: PEAK.load(Relaxed) - start, time }; (out, w) } @@ -392,4 +426,40 @@ fn staged_sink_resources() { let c = w.declared_cost(); println!("parse: {} units, one table each: holds {} bytes, declares {} working bytes", n, parse.peak, c.working_bytes); assert!(parse.peak as u64 <= c.working_bytes, "the parse held {} bytes, past the declared {}", parse.peak, c.working_bytes); + + // ── A base is bounded before it is decoded (§4.5; integer-payloads §7) ── + // A few bytes of base envelope declare a clip of about 2^30 or 2^26 bytes. With a cost + // limit, decode_base refuses each (cost_over_limit) holding no more than a few + // times the envelope: nothing of the declared output is allocated. + use wai::int_payload::{declared_cost, DeclaredCost, IntParams}; + let empty = std::collections::BTreeMap::new(); + let limit = DeclaredCost::new(1 << 20, u64::MAX, u64::MAX, u128::MAX); + for (file, out) in [("clip_1gib.wai", 16 * 4096 * 5461 * 3u64), ("clip_64mib.wai", 4096 * 5461 * 3)] { + let env = std::fs::read(std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("tests/vectors/staged_bounds").join(file)).unwrap(); + let w = Wai::from_bytes(&env).unwrap(); + assert_eq!(declared_cost(CAP_VIDEO_INT_MOTION, &w.payload, &IntParams::none()).unwrap().output_bytes, out, "{file} declares what its README states"); + let (r, win) = measure(|| wai::staged::decode_base(&env, &empty, Some(&limit), None)); + println!("bounded base: {file} ({} bytes, declares {out}): {:?} after {} bytes live at peak", env.len(), r.as_ref().err().map(|e| e.code()), win.peak); + // Never print a decoded clip: it may be a gigabyte. + match r { + Err(e) => assert_eq!(e.code(), "cost_over_limit"), + Ok((_, units)) => panic!("{file}: decoded {} units instead of refusing", units.len()), + } + assert!(win.peak < 64 * env.len() + (64 << 10), "{file}: {} bytes live while refusing it", win.peak); + } + // A decoded base past the chain budget is not kept: the sink holds no chain for it, + // and what it allocated for the base is released with the refusal. + let bh = [0x5e; 32]; + let (info, units) = base(16, 256, bh); + let out: u64 = units.iter().map(|u| u.samples.byte_len()).sum(); + let mut sink = StageSink::new(cfg(16, out / 2)); + let before = CUR.load(Relaxed); + let r = sink.add_base(info, units); + assert_eq!(r.len(), 1); + assert_eq!(r[0].envelope, Some(StageRefusal::ChainOverLimit)); + drop(r); + assert!(sink.chain(&bh).is_none()); + let left = CUR.load(Relaxed).saturating_sub(before); + println!("budget: a base of {out} output bytes in a {}-byte budget is refused; the sink holds {left} bytes of it", out / 2); + assert!(left < 4096, "the sink kept {left} bytes of a base it refused"); } diff --git a/wai/wai-rs/tests/vectors/staged_bounds/README.md b/wai/wai-rs/tests/vectors/staged_bounds/README.md new file mode 100644 index 000000000..6a0235530 --- /dev/null +++ b/wai/wai-rs/tests/vectors/staged_bounds/README.md @@ -0,0 +1,26 @@ +# staged_bounds — bases that declare far more than they carry + +Two `wai.video.int_motion` (`WIV1`, keyframe 0) base envelopes whose few bytes declare a +large decoded clip: grey frames from a residual stream of zeros, coded with a table +that makes the zero symbol nearly free (`{0, 65535, 65536}`), so the stream stays short +while its symbol count, and the declared output, are large. + +| file | envelope bytes | declared output | +|---|---|---| +| `clip_1gib.wai` | 3,243 | 16 frames of 4096 × 5461: 1,073,676,288 bytes, just under the clip cap (2^30) | +| `clip_64mib.wai` | 415 | 1 frame of 4096 × 5461: 67,104,768 bytes, just under the tensor cap (2^26 elements) | + +A sink that applies a cost limit to a base's declared cost refuses them before any +symbol is decoded or any output allocated (`cost_over_limit`; +`extensions/staged-delivery.md` §4.5, `extensions/integer-payloads.md` §7). +`tests/staged_resources.rs` checks that natively with a counting allocator, and +`wai-web/demo/verify_staged_wasm.mjs` checks it in wasm32 by the growth of linear memory. +Both also check the declared output stated above, so the files describe themselves. + +Two more envelopes carry a small valid `WIV1` payload under a capability SPEC §4 never +dispatches, with `wai.video.int_motion` as the fallback: `inert_refinement_primary.wai` +(`wai.video.int_refine`, refinement-only) and `inert_companion_primary.wai` +(`wai.video.lcevc`, companion-only). Each is inert: a sink decodes neither the +capability nor, through it, the fallback (`missing_capability`). They are what +`staged::tests::decode_base_runs_spec_4` builds, which checks them byte for byte +(`WAI_STAGED_BOUNDS_WRITE=1` rewrites them). diff --git a/wai/wai-rs/tests/vectors/staged_bounds/clip_1gib.wai b/wai/wai-rs/tests/vectors/staged_bounds/clip_1gib.wai new file mode 100644 index 000000000..b398da24b Binary files /dev/null and b/wai/wai-rs/tests/vectors/staged_bounds/clip_1gib.wai differ diff --git a/wai/wai-rs/tests/vectors/staged_bounds/clip_64mib.wai b/wai/wai-rs/tests/vectors/staged_bounds/clip_64mib.wai new file mode 100644 index 000000000..3ab84b0c5 Binary files /dev/null and b/wai/wai-rs/tests/vectors/staged_bounds/clip_64mib.wai differ diff --git a/wai/wai-rs/tests/vectors/staged_bounds/inert_companion_primary.wai b/wai/wai-rs/tests/vectors/staged_bounds/inert_companion_primary.wai new file mode 100644 index 000000000..697418fa3 Binary files /dev/null and b/wai/wai-rs/tests/vectors/staged_bounds/inert_companion_primary.wai differ diff --git a/wai/wai-rs/tests/vectors/staged_bounds/inert_refinement_primary.wai b/wai/wai-rs/tests/vectors/staged_bounds/inert_refinement_primary.wai new file mode 100644 index 000000000..779a96fe5 Binary files /dev/null and b/wai/wai-rs/tests/vectors/staged_bounds/inert_refinement_primary.wai differ diff --git a/wai/wai-web/Cargo.toml b/wai/wai-web/Cargo.toml index e85a6b578..bf63cadff 100644 --- a/wai/wai-web/Cargo.toml +++ b/wai/wai-web/Cargo.toml @@ -19,7 +19,7 @@ js-sys = "0.3" # with the codec chain off. No second implementation: replay-equivalence # in the browser is the identical Rust that generated the conformance # corpus hashes. -wai = { path = "../wai-rs", default-features = false, features = ["world", "audio_scene", "haptic", "splat4d", "provenance", "score", "film", "neural_int", "generate", "video_int", "audio_int", "world_model", "world_int", "splat_int", "quantum", "quantum_ops", "quantum_cal", "quantum_mitigate", "quantum_qec", "quantum_control", "quantum_noise", "quantum_compile", "quantum_atom", "quantum_qbom", "quantum_receipt", "quantum_toolchain", "quantum_mps", "quantum_stabilizer", "quantum_pauli", "quantum_vml", "quantum_phasor", "quantum_qfhrr", "quantum_kernel", "quantum_phasor_meter", "quantum_qdata", "video_workflow"] } +wai = { path = "../wai-rs", default-features = false, features = ["world", "audio_scene", "haptic", "splat4d", "provenance", "score", "film", "neural_int", "generate", "video_int", "audio_int", "world_model", "world_int", "splat_int", "quantum", "quantum_ops", "quantum_cal", "quantum_mitigate", "quantum_qec", "quantum_control", "quantum_noise", "quantum_compile", "quantum_atom", "quantum_qbom", "quantum_receipt", "quantum_toolchain", "quantum_mps", "quantum_stabilizer", "quantum_pauli", "quantum_vml", "quantum_phasor", "quantum_qfhrr", "quantum_kernel", "quantum_phasor_meter", "quantum_qdata", "video_workflow", "staged"] } ed25519-dalek = "2" # actor keys for browser-side op authoring (same version as wai-rs) blake3 = "1" # in-page receipt check: hash a decoded image against a pinned digest diff --git a/wai/wai-web/demo/README.md b/wai/wai-web/demo/README.md index 649261951..f90dc1e47 100644 --- a/wai/wai-web/demo/README.md +++ b/wai/wai-web/demo/README.md @@ -50,3 +50,32 @@ nodejs-target pkg and checks all four vectors' hashes): ``` ./demo/verify_int_synth.sh ``` + +## Staged delivery — `StageSink` in the browser + +`src/stage.rs` binds `wai::staged::StageSink`, the staged-delivery sink a native +player runs (`extensions/staged-delivery.md`), for JS: `StageBase.decode` decodes a +base envelope by SPEC §4 through the integer path, with its pin and under the page's +cost limit; `StageBase.fromStage0` takes a base the page decoded elsewhere, which is +*claimed* (its chain's reports, units and stage hashes say so). `StageSink` adds bases +(taking them), offers refinement layers, expires holds and removes bases in any order, +and every unit's stage, producer, tier, output and stage hash reads back. +`stageLayerInfo` reads a layer's `refines`, `WIR1` header and declared cost, and a +base's declared decode cost, before either is offered. Reports and units come back in +the corpus's own form (`staged-conformance/README.md`). + +`verify_staged_wasm.mjs` replays every case of `staged-conformance/` through the +wasm32 build and requires every report, final unit, stage hash and clip hash the +native sink wrote; each integer base is decoded in wasm32 and must give its +`stage0.bin`. It then checks the binding's guards: bases that declare huge clips +(`wai-rs/tests/vectors/staged_bounds/`) refused under a cost limit with linear memory +watched, a base past the chain budget, base facts a page cannot raise, and a seeded +fuzz that must not trap: + +``` +# from wai-web/ +cargo build --release --target wasm32-unknown-unknown +wasm-bindgen --target nodejs --out-dir demo/pkg-node --out-name wai_web \ + target/wasm32-unknown-unknown/release/wai_web.wasm +node demo/verify_staged_wasm.mjs ./pkg-node/wai_web.js ../staged-conformance +``` diff --git a/wai/wai-web/demo/verify_staged_wasm.mjs b/wai/wai-web/demo/verify_staged_wasm.mjs new file mode 100644 index 000000000..33b84f09c --- /dev/null +++ b/wai/wai-web/demo/verify_staged_wasm.mjs @@ -0,0 +1,319 @@ +// The wasm32 staged-delivery sink equals the native one: every case of the corpus +// wai/staged-conformance/ is replayed through the browser crate's StageSink, built +// for wasm32, and every report of every step and every base's final units (stage, +// producer, frozen flag, tier, shape and stage hash) and clip hash must equal what +// case.json states, which wai_staged_vectors wrote from the native StageSink. +// +// A base whose decode is the integer path (case.json `decode.kind` "int") is decoded +// in wasm32 too, through StageBase.decode with its parameter set, its pin and a cost +// limit, and must give its stage0.bin byte for byte, the case's tier and the case's +// fallback flag. A PNG or FLAC base (no decoder in this build) is claimed from +// stage0.bin (StageBase.fromStage0), whose SHA-256 is checked; every report, unit and +// stage hash of its chain must then say claimed, and with that mark set aside equal +// the case's. Each base then takes the case's facts (attested, first unit). +// +// Then the binding's guards: +// - bounds: the bases in wai-rs/tests/vectors/staged_bounds/ declare clips of about +// 2^30 and 2^26 bytes in a few bytes; stageLayerInfo reads that cost, and StageBase.decode +// with a cost limit refuses them (cost_over_limit) while linear memory grows by +// less than a few MiB; a sink whose chain budget a base passes refuses the base +// (chain-over-limit); +// - facts: a base's viaFallback cannot be set, so its layers stay refused +// base-via-fallback; an inert envelope is neither decoded nor claimed, through its +// capability or its fallback; fromStage0 takes only the manifest's capability or +// its fallback, lowers a claimed tier to what the capability gives an unpinned +// base, refuses a misspelt tier, and marks the chain claimed; StageBase.decode +// refuses a call with no limit argument (null is the format's caps); +// - addBase moves the base into the sink: the JS object is spent; +// - a seeded fuzz of every corpus base and layer through stageLayerInfo, +// StageBase.decode (with a limit) and offer never traps. +// +// node verify_staged_wasm.mjs [path/to/wai_web.js] [path/to/staged-conformance] +// +// Exit 0 iff every check passes; the number of cases replayed is printed, and a +// corpus with no cases fails. + +import { readFileSync, readdirSync } from 'node:fs'; +import { createHash } from 'node:crypto'; +import { dirname, resolve } from 'node:path'; +import { fileURLToPath, pathToFileURL } from 'node:url'; + +const here = dirname(fileURLToPath(import.meta.url)); +const pkg = resolve(here, process.argv[2] ?? './pkg-node/wai_web.js'); +const corpus = resolve(process.argv[3] ?? resolve(here, '../../staged-conformance')); +const waiDir = resolve(corpus, '..'); +// A page's cost limit for a base decode: far above every corpus base. +const LIMIT = { output_bytes: 64 << 20, symbols: 64 << 20, working_bytes: 256 << 20 }; +const wai = await import(pathToFileURL(pkg).href); + +const sha = (b) => createHash('sha256').update(b).digest('hex'); +const bytes = (p) => new Uint8Array(readFileSync(p)); +// Key order matters: the reports and units are compared as the corpus writes them. +const same = (a, b) => JSON.stringify(a) === JSON.stringify(b); + +let fails = 0; +const fail = (m) => { fails++; console.log(`FAIL ${m}`); }; + +const names = readdirSync(`${corpus}/cases`, { withFileTypes: true }).filter((d) => d.isDirectory()).map((d) => d.name).sort(); +let replayed = 0; +let stageHashes = 0; +for (const name of names) { + const dir = `${corpus}/cases/${name}`; + const c = JSON.parse(readFileSync(`${dir}/case.json`, 'utf8')); + const bases = []; + let ok = true; + for (const b of c.bases) { + const env = bytes(`${dir}/${b.file}`); + const stage0 = bytes(`${dir}/${b.stage0}`); + if (sha(stage0) !== b.stage0_sha256) { fail(`${name}: ${b.stage0} does not have its stated SHA-256`); ok = false; } + if (wai.stageContentHash(env) !== b.object) { fail(`${name}: ${b.file} does not have its stated content hash`); ok = false; } + const elems = new Set(b.units.map((u) => u.elem)); + if (elems.size !== 1) { fail(`${name}: mixed element types`); ok = false; continue; } + const shapes = new Uint32Array(b.units.flatMap((u) => u.shape)); + let base; + if (b.decode.kind === 'int') { + const set = new Map(Object.entries(b.decode.set).map(([m, p]) => [m, bytes(`${waiDir}/${p}`)])); + try { + base = wai.StageBase.decode(env, set.size ? set : undefined, b.decoded_via, LIMIT); + } catch (e) { + fail(`${name}: ${b.file} does not decode: ${e.message ?? e}`); ok = false; continue; + } + if (!Buffer.from(base.stage0()).equals(Buffer.from(stage0))) { fail(`${name}: the wasm decode of ${b.file} is not ${b.stage0}`); ok = false; } + if (base.tier !== b.tier) { fail(`${name}: the base tier is ${base.tier}, the case states ${b.tier}`); ok = false; } + if (base.viaFallback !== b.via_fallback) { fail(`${name}: via fallback ${base.viaFallback}, the case states ${b.via_fallback}`); ok = false; } + if (base.capability !== b.decoded_via) { fail(`${name}: decoded via ${base.capability}`); ok = false; } + } else { + base = wai.StageBase.fromStage0(env, b.decoded_via, b.tier, shapes, [...elems][0], stage0); + if (!base.claimed || base.viaFallback !== b.via_fallback || base.tier !== b.tier) { fail(`${name}: the claimed base's facts differ from the case's`); ok = false; } + } + base.attested = b.attested; + base.setFirstUnit(b.first_unit); + const remake = () => { + const x = b.decode.kind === 'int' + ? wai.StageBase.decode(env, b.decode.set && Object.keys(b.decode.set).length ? new Map(Object.entries(b.decode.set).map(([m, p]) => [m, bytes(`${waiDir}/${p}`)])) : undefined, b.decoded_via, LIMIT) + : wai.StageBase.fromStage0(env, b.decoded_via, b.tier, shapes, [...elems][0], stage0); + x.attested = b.attested; + x.setFirstUnit(b.first_unit); + return x; + }; + bases.push({ base, hash: base.contentHash, claimed: base.claimed, remake, spent: false }); + } + if (!ok) continue; + const files = new Map(c.files.map((f) => [f.file, bytes(`${dir}/${f.file}`)])); + const sink = new wai.StageSink(c.config); + const reports = []; + // A claimed chain's reports say so; set the mark aside to compare with the case. + const unmark = (r) => { + if (!Array.isArray(r)) return r; + return r.map((x) => { + const { claimed, ...rest } = x; + if (claimed !== undefined && claimed !== true) fail(`${name}: claimed is ${claimed}`); + return rest; + }); + }; + const anyClaimed = bases.some((b) => b.claimed); + for (const s of c.steps) { + let got; + switch (s.op) { + case 'offer': + try { got = sink.offer(files.get(s.file), s.at_ms, s.evidence); } + catch (e) { got = { error: String(e.message ?? e).split(':')[0] }; } + break; + case 'base': { + // addBase spends the StageBase; a case that adds a base again gets a fresh one. + const b = bases[s.index]; + if (b.spent) b.base = b.remake(); + got = sink.addBase(b.base); + b.spent = true; + break; + } + case 'expire': got = sink.expire(s.at_ms); break; + case 'remove': got = sink.removeBase(bases[s.index].hash) ?? []; break; + case 'attest': sink.markAttested(bases[s.index].hash); got = []; break; + default: fail(`${name}: unknown op ${s.op}`); + } + if (Array.isArray(got) && !anyClaimed && got.some((x) => 'claimed' in x)) fail(`${name}: a report says claimed with no claimed base`); + reports.push(unmark(got)); + } + c.expect.reports.forEach((want, i) => { + if (!same(reports[i], want)) fail(`${name}: step ${i} (${c.steps[i].op}) reports ${JSON.stringify(reports[i])}, the case states ${JSON.stringify(want)}`); + }); + if (reports.length !== c.expect.reports.length) fail(`${name}: ${reports.length} steps replayed, the case states ${c.expect.reports.length}`); + c.expect.bases.forEach((want, i) => { + const h = bases[i].hash; + const raw = sink.units(h); + const claimed = sink.isClaimed(h); + if (raw !== null && claimed !== bases[i].claimed) fail(`${name}: base ${i}: isClaimed ${claimed}`); + if (raw !== null && claimed && !(raw.claimed === true && raw.units.every((u) => u.claimed === true))) fail(`${name}: base ${i}: a claimed chain's units do not say so`); + const got = raw === null ? null : { ...raw, units: raw.units.map(({ claimed, ...u }) => u) }; + if (got !== null) delete got.claimed; + if (!same(got, want)) { fail(`${name}: base ${i}'s final units differ: ${JSON.stringify(got)}`); return; } + if (got === null) return; + // The per-unit accessors agree with the summary, unit by unit. + got.units.forEach((u, k) => { + const sh = sink.unitSha256(h, k); + const viaAccessor = sh?.sha256 ?? null; + if (viaAccessor !== u.stage_sha256) fail(`${name}: base ${i} unit ${k}: unitSha256 ${viaAccessor}`); + if (sh && sh.claimed !== claimed) fail(`${name}: base ${i} unit ${k}: unitSha256 claimed ${sh.claimed}`); + if (sink.unitStage(h, k) !== u.stage || sink.unitProducer(h, k) !== u.producer || sink.unitTier(h, k) !== u.tier || sink.unitFrozen(h, k) !== u.frozen) { + fail(`${name}: base ${i} unit ${k}: an accessor disagrees with units()`); + } + if (u.stage_sha256 !== null) { + stageHashes++; + const out = u.shape[2] && sink.unitShape(h, k); + const buf = (() => { try { return sink.unitBytes(h, k); } catch { const p = sink.unitPcm(h, k); return new Uint8Array(p.buffer, p.byteOffset, p.byteLength); } })(); + if (sha(buf) !== u.stage_sha256) fail(`${name}: base ${i} unit ${k}: the output read back does not hash to its stage hash`); + if (!out || out.length !== 3) fail(`${name}: base ${i} unit ${k}: unitShape`); + } + }); + if ((sink.clipSha256(h)?.sha256 ?? null) !== want.clip_sha256) fail(`${name}: base ${i}: clipSha256`); + }); + sink.free(); + bases.forEach((b) => { if (!b.spent) b.base.free(); }); + replayed++; +} + +const expectThrow = (what, f, code) => { + try { f(); fail(`${what}: no error`); } + catch (e) { + const msg = String(e?.message ?? e); + if (e instanceof WebAssembly.RuntimeError) fail(`${what}: trapped: ${msg}`); + else if (code && !msg.startsWith(code)) fail(`${what}: ${msg.slice(0, 120)}, want ${code}`); + } +}; + +// ── bounds: a base is bounded before it is decoded ── +{ + const B = resolve(waiDir, 'wai-rs/tests/vectors/staged_bounds'); + const small = { output_bytes: 1 << 20, symbols: 1 << 20, working_bytes: 1 << 22 }; + for (const [file, out] of [['clip_1gib.wai', 16 * 4096 * 5461 * 3], ['clip_64mib.wai', 4096 * 5461 * 3]]) { + const env = bytes(`${B}/${file}`); + const info = wai.stageLayerInfo(env); + if (info.declared_cost?.output_bytes !== out) fail(`bounds: ${file} declares ${info.declared_cost?.output_bytes}, want ${out}`); + const before = wai.wasmMemoryBytes(); + expectThrow(`bounds: ${file} under a limit`, () => wai.StageBase.decode(env, undefined, undefined, small), 'cost_over_limit'); + const grew = wai.wasmMemoryBytes() - before; + if (!(before > 0) || grew > 4 * 2 ** 20) fail(`bounds: refusing ${file} grew linear memory by ${grew} bytes`); + } + // A base whose units pass the chain budget: refused, and not held. + const d = `${corpus}/cases/video_int_motion_gop_perframe`; + const env = bytes(`${d}/base_0.wai`); + const base = wai.StageBase.decode(env, undefined, undefined, small); + const s = new wai.StageSink({ bound: { max_envelopes: 4, max_age_ms: 1000 }, cost_limit: null, budget: 100, record_budget: 1000, live: null }); + const r = s.addBase(base); + if (r.length !== 1 || r[0].envelope !== 'chain-over-limit' || s.hasBase(wai.stageContentHash(env))) fail(`bounds: a base past the chain budget: ${JSON.stringify(r)}`); + // addBase spent the StageBase. + expectThrow('addBase spends the base', () => base.nUnits); + s.free(); +} + +// ── facts a page cannot raise ── +{ + const d = `${corpus}/cases/base_via_fallback`; + const c = JSON.parse(readFileSync(`${d}/case.json`, 'utf8')); + const base = wai.StageBase.decode(bytes(`${d}/${c.bases[0].file}`), undefined, c.bases[0].decoded_via, LIMIT); + if (base.viaFallback !== true) fail('facts: the fallback decode is not marked via fallback'); + try { base.viaFallback = false; } catch { /* no setter: a TypeError in a module */ } + if (base.viaFallback !== true) fail('facts: viaFallback was lowered'); + base.attested = true; + base.attested = false; + if (base.attested !== true) fail('facts: attested was lowered'); + const s = new wai.StageSink(c.config); + s.addBase(base); + const r = s.offer(bytes(`${d}/L1.wai`), 1, 'delivered-verified'); + if (r[0]?.envelope !== 'base-via-fallback') fail(`facts: a fallback base's layer: ${JSON.stringify(r)}`); + s.free(); + + // Two inert envelopes: a refinement-only and a companion-only capability with an + // int_motion fallback never reach it. + const fb = bytes(`${corpus}/cases/out_of_order/base.wai`); + for (const inert of ['inert_refinement_primary.wai', 'inert_companion_primary.wai']) { + const ienv = bytes(`${resolve(waiDir, 'wai-rs/tests/vectors/staged_bounds')}/${inert}`); + const icap = wai.stageLayerInfo(ienv).capability; + expectThrow(`facts: ${inert}`, () => wai.StageBase.decode(ienv, undefined, undefined, LIMIT), 'missing_capability'); + expectThrow(`facts: ${inert} via its fallback`, () => wai.StageBase.decode(ienv, undefined, 'wai.video.int_motion', LIMIT), 'missing_capability'); + // Nor can a page claim it: SPEC §4 never takes its capability or its fallback. + for (const cap of [icap, 'wai.video.int_motion']) { + expectThrow(`facts: ${inert} claimed as ${cap}`, () => wai.StageBase.fromStage0(ienv, cap, 'decode-equivalence', new Uint32Array([4, 4, 3, 4, 4, 3]), 'u8', new Uint8Array(96).fill(5)), 'missing_capability'); + } + } + // The limit is required: left out, the decode is refused; null is the format's caps. + expectThrow('facts: a decode with no limit argument', () => wai.StageBase.decode(fb, undefined, undefined), 'limit'); + wai.StageBase.decode(fb, undefined, undefined, null).free(); + + // A claimed base: only its manifest's capability or fallback; its tier no higher than + // an unpinned base's; a misspelt tier refused; its chain marked claimed. + const u = `${corpus}/cases/image_int_hyper_unpinned`; + const uc = JSON.parse(readFileSync(`${u}/case.json`, 'utf8')); + const uenv = bytes(`${u}/${uc.bases[0].file}`); + const shapes = new Uint32Array(uc.bases[0].units.flatMap((x) => x.shape)); + const junk = new Uint8Array(shapes[0] * shapes[1] * shapes[2]).fill(77); + expectThrow('facts: a capability the manifest does not name', () => wai.StageBase.fromStage0(uenv, 'wai.image.png', 'decode-equivalence', shapes, 'u8', junk), 'not-offered'); + expectThrow('facts: a misspelt tier', () => wai.StageBase.fromStage0(uenv, 'wai.neural.int_hyper', 'decode-equivalance', shapes, 'u8', junk), 'tier'); + const claimed = wai.StageBase.fromStage0(uenv, 'wai.neural.int_hyper', 'decode-equivalence', shapes, 'u8', junk); + if (claimed.tier !== 'none' || !claimed.claimed) fail(`facts: a claimed int_hyper base has tier ${claimed.tier}`); + const ch = claimed.contentHash; + const s2 = new wai.StageSink(uc.config); + s2.addBase(claimed); + const units = s2.units(ch); + if (!s2.isClaimed(ch) || units.claimed !== true || units.units.some((x) => x.claimed !== true)) fail('facts: the claimed chain is not marked'); + if (s2.unitSha256(ch, 0) !== undefined) fail('facts: a claimed base with no tier has a stage hash'); + const r2 = s2.offer(bytes(`${u}/${uc.files[0].file}`), 1, 'none'); + if (!r2.every((x) => x.claimed === true)) fail(`facts: a claimed chain's report is not marked: ${JSON.stringify(r2)}`); + s2.free(); +} + +// ── a seeded fuzz: nothing traps ── +{ + let seed = 99; + const rnd = () => { seed = (seed * 1103515245 + 12345) >>> 0; return seed; }; + let calls = 0; + const guard = (what, f) => { + calls++; + try { f(); } catch (e) { if (e instanceof WebAssembly.RuntimeError) fail(`fuzz: ${what} trapped: ${e}`); } + }; + for (const name of names) { + const d = `${corpus}/cases/${name}`; + const c = JSON.parse(readFileSync(`${d}/case.json`, 'utf8')); + const ints = c.bases.filter((b) => b.decode.kind === 'int'); + for (const b of ints) { + const env = bytes(`${d}/${b.file}`); + const set = new Map(Object.entries(b.decode.set).map(([m, p]) => [m, bytes(`${waiDir}/${p}`)])); + const pl = wai.stageLayerInfo(env, set.size ? set : undefined).payload_bytes; + for (let k = 0; k < 12; k++) { + const m = env.slice(); + for (let j = 0; j < 1 + (rnd() % 3); j++) m[m.length - pl + (rnd() % pl)] ^= 1 << (rnd() % 8); + guard(`${name}/${b.file}`, () => { wai.stageLayerInfo(m, set.size ? set : undefined); wai.StageBase.decode(m, set.size ? set : undefined, b.decoded_via, LIMIT).free(); }); + } + } + for (const f of c.files.map((x) => x.file)) { + const env = bytes(`${d}/${f}`); + for (let k = 0; k < 12; k++) { + const m = env.slice(); + for (let j = 0; j < 1 + (rnd() % 4); j++) m[rnd() % m.length] ^= 1 << (rnd() % 8); + guard(`${name}/${f}`, () => { + wai.stageLayerInfo(m); + wai.stageContentHash(m); + const s = new wai.StageSink(c.config); + try { s.offer(m, 0, 'none'); } finally { s.free(); } + }); + } + } + } + console.log(`fuzz: ${calls} mutated envelopes, no trap`); +} + +// A layer's header and declared cost read before it is offered. +{ + const d = `${corpus}/cases/video_int_motion_gop_perframe`; + const env = bytes(`${d}/g0_L2.wai`); + const info = wai.stageLayerInfo(env); + if (info.framing_bytes + info.manifest_bytes + info.payload_bytes !== env.length) fail('stageLayerInfo: the bytes do not sum to the envelope'); + if (info.refines?.layer !== 2 || !Array.isArray(info.refines?.prev) || info.wir1?.n_units !== 4) fail('stageLayerInfo: refines or WIR1 header'); + if (!(info.declared_cost?.output_bytes > 0)) fail('stageLayerInfo: declared cost'); +} + +console.log(`${replayed} of ${names.length} cases replayed through the wasm32 StageSink; ${stageHashes} unit stage hashes equal the native ones`); +if (names.length === 0 || replayed !== names.length) fail('not every case replayed'); +console.log(fails === 0 ? 'staged-conformance (wasm32): OK' : `staged-conformance (wasm32): ${fails} failure(s)`); +process.exit(fails === 0 ? 0 : 1); diff --git a/wai/wai-web/src/lib.rs b/wai/wai-web/src/lib.rs index 11819025a..2ed3e5e9c 100644 --- a/wai/wai-web/src/lib.rs +++ b/wai/wai-web/src/lib.rs @@ -152,6 +152,12 @@ pub mod quantum_phasor; // a total-energy roll-up, and show a swapped stage break the bill. pub mod quantum_qbom; +// Staged delivery (extensions/staged-delivery.md): wai::staged::StageSink, the sink a +// native player runs, bound for JS — bases added, refinement layers offered, holds +// expired and bases removed in any order, every unit's stage, producer, tier, output +// and stage hash read back, and a layer's declared cost read before it is offered. +pub mod stage; + // ---- Capability menu ---- // The capabilities this browser sink advertises: a subset of the registry diff --git a/wai/wai-web/src/stage.rs b/wai/wai-web/src/stage.rs new file mode 100644 index 000000000..91d9a5f2a --- /dev/null +++ b/wai/wai-web/src/stage.rs @@ -0,0 +1,962 @@ +//! Staged delivery in the browser (`extensions/staged-delivery.md`): a base first, +//! refinement layers after, each its own `WAI1` object bound to what it refines by +//! content hash. +//! +//! [`StageSink`] is `wai::staged::StageSink`, the sink a native player runs, bound for +//! JS: bases are added, layers offered, holds expired and bases removed in any order, +//! and every outcome is the reference implementation's report (§4.1, §4.2). A page +//! reads each unit's stage, producer, tier, output and stage hash back from it, and a +//! layer's declared cost before offering it ([`stage_layer_info`]). Nothing here is a +//! second implementation: the same Rust decides on every machine, so the stage hashes +//! a page computes are the native ones (`demo/verify_staged_wasm.mjs` replays the +//! corpus `staged-conformance/` and requires every report and hash). +//! +//! Media cross the network as `WAI1` bytes only. Reports, units and layer headers come +//! back as plain objects in the corpus's own form (`staged-conformance/README.md`): +//! snake_case members, content hashes as lowercase hex, reasons as §4.2's tokens. +//! +//! The pure-Rust half (`config_from`, `report_value`, `units_value`, `layer_info`) +//! is what the native tests run; the `#[wasm_bindgen]` wrappers only convert. + +use std::collections::BTreeMap; + +use serde_json::{json, Value}; +use wai::container::{Prev, Refines, Wai}; +use wai::staged::schedule::UnitClock; +use wai::staged::wir1::{self, Wir1Cost}; +use wai::staged::{self, tier, BaseInfo, ChainConfig, HoldBound, Live, OfferReport, ReceiptEvidence, Samples, StageChain, UnitBuf}; +use wasm_bindgen::prelude::*; + +fn hex(b: &[u8]) -> String { + b.iter().map(|x| format!("{x:02x}")).collect() +} + +fn unhex32(s: &str) -> Option<[u8; 32]> { + let b = s.as_bytes(); + if b.len() != 64 { + return None; + } + let nib = |c: u8| match c { + b'0'..=b'9' => Some(c - b'0'), + b'a'..=b'f' => Some(c - b'a' + 10), + _ => None, + }; + let mut out = [0u8; 32]; + for (i, o) in out.iter_mut().enumerate() { + *o = (nib(b[2 * i])? << 4) | nib(b[2 * i + 1])?; + } + Some(out) +} + +// ───────────────────────── the pure-Rust half ───────────────────────── + +/// A count, a byte figure or a time in milliseconds, from a JSON number: a +/// non-negative integer. A figure at or past `2^64` (JavaScript reads +/// 18446744073709551615 as `2^64`) is the largest `u64`. +pub(crate) fn count(v: &Value, what: &str) -> Result { + if let Some(n) = v.as_u64() { + return Ok(n); + } + match v.as_f64() { + Some(f) if f >= 0.0 && f.fract() == 0.0 => Ok(if f >= 18_446_744_073_709_551_616.0 { u64::MAX } else { f as u64 }), + _ => Err(format!("{what}: a non-negative integer")), + } +} + +/// A sink's configuration, in the form a corpus case states it (`case.json`'s +/// `config`): `{bound: {max_envelopes, max_age_ms}, cost_limit: null | {output_bytes, +/// symbols, working_bytes}, budget, record_budget, live: null | {t0_ms, unit_ms, +/// target_latency_ms, lead_ms: [..]}}`. `cost_limit` and `live` may be absent. +pub(crate) fn config_from(v: &Value) -> Result { + let obj = v.as_object().ok_or("the configuration is an object")?; + for k in obj.keys() { + if !["bound", "cost_limit", "budget", "record_budget", "live"].contains(&k.as_str()) { + return Err(format!("the configuration has no member {k}")); + } + } + let n = |v: &Value, what: &str| count(v, what); + let bound = &v["bound"]; + let max_envelopes = n(&bound["max_envelopes"], "bound.max_envelopes")?; + let cost_limit = match &v["cost_limit"] { + Value::Null => None, + c => Some(Wir1Cost { + output_bytes: n(&c["output_bytes"], "cost_limit.output_bytes")?, + symbols: n(&c["symbols"], "cost_limit.symbols")?, + working_bytes: n(&c["working_bytes"], "cost_limit.working_bytes")?, + }), + }; + let live = match &v["live"] { + Value::Null => None, + l => Some(Live { + clock: UnitClock { + t0_ms: n(&l["t0_ms"], "live.t0_ms")?, + unit_ms: n(&l["unit_ms"], "live.unit_ms")?, + target_latency_ms: n(&l["target_latency_ms"], "live.target_latency_ms")?, + }, + lead_ms: l["lead_ms"].as_array().ok_or("live.lead_ms: an array")?.iter().map(|x| n(x, "live.lead_ms")).collect::>()?, + }), + }; + Ok(ChainConfig { + bound: HoldBound { max_envelopes: usize::try_from(max_envelopes).unwrap_or(usize::MAX), max_age_ms: n(&bound["max_age_ms"], "bound.max_age_ms")? }, + cost_limit, + budget: n(&v["budget"], "budget")?, + record_budget: n(&v["record_budget"], "record_budget")?, + live, + }) +} + +/// A report as the corpus states one. +pub(crate) fn report_value(r: &OfferReport) -> Value { + json!({ + "object": hex(&r.content_hash), + "envelope": r.envelope.map(|x| x.code()), + "detail": r.detail, + "held_for_base": r.held_for_base, + "applied": r.applied, + "held": r.held, + "refused": r.refused.iter().map(|(u, w)| json!([u, w.code()])).collect::>(), + }) +} + + +/// A chain's units and clip hash as the corpus states a base at the end of a case. +pub(crate) fn units_value(c: &StageChain) -> Value { + json!({ + "units": c.units().iter().enumerate().map(|(u, s)| json!({ + "stage": s.stage, + "producer": hex(&s.producer), + "frozen": s.frozen, + "tier": tier::token(s.tier), + "shape": [s.buf.shape.0, s.buf.shape.1, s.buf.shape.2], + "stage_sha256": c.unit_sha256(u as u32).map(|h| hex(&h)), + })).collect::>(), + "clip_sha256": c.clip_sha256().map(|h| hex(&h)), + "n_units": c.n_units(), + }) +} + +/// Units from their canonical buffers concatenated (a corpus case's `stage0.bin`): +/// `shapes` holds three dimensions per unit, `elem` is `u8` or `i16`. +pub(crate) fn units_from_stage0(shapes: &[u32], elem: &str, stage0: &[u8]) -> Result, String> { + if shapes.is_empty() || shapes.len() % 3 != 0 { + return Err("shapes: three dimensions per unit".into()); + } + let width = match elem { + "u8" => 1usize, + "i16" => 2, + e => return Err(format!("elem {e}: u8 or i16")), + }; + let mut at = 0usize; + let mut out = Vec::with_capacity(shapes.len() / 3); + for s in shapes.chunks_exact(3) { + let n = (s[0] as usize).checked_mul(s[1] as usize).and_then(|x| x.checked_mul(s[2] as usize)).and_then(|x| x.checked_mul(width)).ok_or("a unit's size overflows")?; + let end = at.checked_add(n).filter(|&e| e <= stage0.len()).ok_or("stage0 is shorter than its units")?; + let b = &stage0[at..end]; + let samples = if width == 1 { Samples::U8(b.to_vec()) } else { Samples::I16(b.chunks_exact(2).map(|c| i16::from_le_bytes([c[0], c[1]])).collect()) }; + out.push(UnitBuf { shape: (s[0], s[1], s[2]), samples }); + at = end; + } + if at != stage0.len() { + return Err("stage0 is longer than its units".into()); + } + Ok(out) +} + +fn prev_value(p: &Prev) -> Value { + match p { + Prev::One(h) => json!(h), + Prev::Runs(runs) => Value::Array(runs.iter().map(|r| json!({"producer": r.producer, "count": r.count})).collect()), + } +} + +/// What an envelope is, read before anything is decoded: its content hash, its bytes +/// (the `WAI1` framing, the manifest and the payload, which sum to the whole), its +/// capability; for a layer its `refines`, its `WIR1` header and its declared cost +/// (Appendix A.5: the output bytes, symbols and working bytes a sink applies its cost +/// limit to); for a base of an integer codec, the declared cost of its decode +/// (`extensions/integer-payloads.md` §7; `macs` as a decimal string), against `set` +/// where the codec reads a parameter set. A member that does not parse is `null`, with +/// the reason under `error`. +pub(crate) fn layer_info(env: &[u8], set: &BTreeMap) -> Result { + let w = Wai::from_bytes(env).map_err(|e| format!("not-wai1: {e}"))?; + let manifest_bytes = u32::from_le_bytes(env[4..8].try_into().map_err(|_| "not-wai1: no manifest length")?) as usize; + let payload_bytes = w.payload.len(); + let cap = w.capability().to_string(); + let mut out = json!({ + "object": hex(&wai::jwp_hash::content_hash(env)), + "bytes": env.len(), + "framing_bytes": env.len() - manifest_bytes - payload_bytes, + "manifest_bytes": manifest_bytes, + "payload_bytes": payload_bytes, + "capability": cap, + "refinement_only": wai::codecs::is_refinement_only(&cap), + "refines": null, + "wir1": null, + "declared_cost": null, + "error": null, + }); + match w.refines() { + None => {} + Some(Err(e)) => out["error"] = json!(format!("refines-malformed: {e}")), + Some(Ok(r)) => { + let Refines { base, prev, layer, span } = &r; + out["refines"] = json!({ + "base": base, + "prev": prev_value(prev), + "layer": layer, + "span": span.map(|s| json!({"from": s.from, "count": s.count})), + }); + } + } + // A base: the declared cost of the capability SPEC §4 would take first through the + // integer path (its own, else its fallback; none for a refinement-only or + // companion-only capability), read from its header against `set`, before anything + // is decoded (`extensions/integer-payloads.md` §7). + if staged::unit_kind(&cap).is_none() && !wai::codecs::is_companion_only(&cap) && !wai::codecs::is_refinement_only(&cap) { + let via = [Some(cap.as_str()), w.fallback()].into_iter().flatten().find(|c| wai::codecs::is_dispatchable(c) && wai::codecs::pinned_int_payload(c).is_some()); + if let Some(via) = via { + match crate::int_payload::cost(via, &w.payload, &set.iter().map(|(k, v)| (k.as_str(), *v)).collect::>()) { + Ok(c) => { + out["declared_cost"] = json!({"output_bytes": c.output_bytes, "symbols": c.symbols, "working_bytes": c.working_bytes, "macs": c.macs.to_string()}); + out["cost_capability"] = json!(via); + } + Err(e) => out["error"] = json!(e), + } + } + } + if let Some(kind) = staged::unit_kind(&cap) { + match wir1::parse(&w.payload, kind) { + Ok(p) => { + let c = p.declared_cost(); + out["wir1"] = json!({ + "n_units": p.n_units, + "shape": [p.d0, p.d1, p.d2], + "predictor": p.predictor.token(), + "q": p.q, + "tables": p.n_tables(), + }); + out["declared_cost"] = json!({"output_bytes": c.output_bytes, "symbols": c.symbols, "working_bytes": c.working_bytes}); + } + Err(e) => out["error"] = json!(format!("payload-malformed: {}", e.code())), + } + } + Ok(out) +} + +/// A cost limit, in the form the binding takes it: `{output_bytes, symbols, +/// working_bytes, macs}` (`macs` may be absent, for no bound on it), each figure read +/// as [`count`] reads one; `null` for no limit. +pub(crate) fn limit_from(v: &Value) -> Result, String> { + if v.is_null() { + return Ok(None); + } + let obj = v.as_object().ok_or("the limit is an object, or null")?; + for k in obj.keys() { + if !["output_bytes", "symbols", "working_bytes", "macs"].contains(&k.as_str()) { + return Err(format!("the limit has no member {k}")); + } + } + let macs = match &v["macs"] { + Value::Null => u128::MAX, + m => { + let n = count(m, "limit.macs")?; + if n == u64::MAX { u128::MAX } else { n as u128 } + } + }; + Ok(Some(wai::int_payload::DeclaredCost::new( + count(&v["output_bytes"], "limit.output_bytes")?, + count(&v["symbols"], "limit.symbols")?, + count(&v["working_bytes"], "limit.working_bytes")?, + macs, + ))) +} + +/// A base decoded elsewhere, as the binding takes one ([`StageBase::from_stage0`]): the +/// envelope must parse, `capability` must be its manifest's capability or its fallback +/// (decoded through the fallback where it is not the capability) that SPEC §4 would +/// take (`missing_capability` for an inert envelope, whose capability is +/// refinement-only or companion-only, or a capability that is not dispatchable), and +/// the claimed tier +/// is a tier token, lowered to what the capability's registry class gives a base with +/// no resolved pin ([`tier::base_tier`]): a page never raises a base's tier by claiming +/// it. +pub(crate) fn claim_base(envelope: &[u8], capability: &str, tier_token: &str, shapes: &[u32], elem: &str, stage0: &[u8]) -> Result<(BaseInfo, Vec), String> { + let w = Wai::from_bytes(envelope).map_err(|e| format!("not-wai1: {e}"))?; + let primary = w.capability(); + if capability != primary && w.fallback() != Some(capability) { + return Err(format!("not-offered: {capability} is neither the manifest's capability nor its fallback")); + } + // SPEC §4, as decode_base_via applies it: a refinement-only or companion-only + // capability makes the envelope inert, so neither it nor its fallback is ever + // taken, and only a dispatchable capability decodes a payload. + if wai::codecs::is_refinement_only(primary) || wai::codecs::is_companion_only(primary) || !wai::codecs::is_dispatchable(capability) { + return Err(format!("missing_capability: SPEC §4 takes no capability for this envelope through {capability}")); + } + let claimed = match tier_token { + "decode-equivalence" | "entropy-consistency" | "none" => tier::from_token(tier_token), + t => return Err(format!("tier: {t} is not decode-equivalence, entropy-consistency or none")), + }; + let ceiling = tier::base_tier(capability, None); + let units = units_from_stage0(shapes, elem, stage0)?; + let info = BaseInfo { + content_hash: wai::jwp_hash::content_hash(envelope), + capability: capability.to_string(), + via_fallback: capability != primary, + tier: claimed.min(ceiling), + attested: false, + first_unit: 0, + }; + Ok((info, units)) +} + +/// The sum of squared differences of two equal-length byte buffers, exact. +pub(crate) fn sse_u8(a: &[u8], b: &[u8]) -> Result { + if a.len() != b.len() { + return Err(format!("the buffers differ in length: {} and {}", a.len(), b.len())); + } + Ok(a.iter().zip(b).map(|(&x, &y)| { + let d = x as i64 - y as i64; + (d * d) as u64 + }).sum()) +} + +/// The same for 16-bit samples. +pub(crate) fn sse_i16(a: &[i16], b: &[i16]) -> Result { + if a.len() != b.len() { + return Err(format!("the buffers differ in length: {} and {}", a.len(), b.len())); + } + Ok(a.iter().zip(b).map(|(&x, &y)| { + let d = x as i64 - y as i64; + (d * d) as u64 + }).sum()) +} + +// ───────────────────────── the JS surface ───────────────────────── + +fn js_err(e: impl std::fmt::Display) -> JsError { + JsError::new(&e.to_string()) +} + +fn to_js(v: &Value) -> Result { + use serde::Serialize; + v.serialize(&serde_wasm_bindgen::Serializer::json_compatible()).map_err(js_err) +} + +/// A time or figure from JS: a non-negative integer, or `Infinity` for the largest. +fn ms(v: f64, what: &str) -> Result { + if v == f64::INFINITY { + return Ok(u64::MAX); + } + if !(v >= 0.0) || v.fract() != 0.0 { + return Err(JsError::new(&format!("{what}: a non-negative integer"))); + } + Ok(if v >= 18_446_744_073_709_551_616.0 { u64::MAX } else { v as u64 }) +} + +/// A JS `Map` of file name → bytes, as owned pairs. +fn set_from(set: Option) -> Result)>, JsError> { + let mut held: Vec<(String, Vec)> = Vec::new(); + if let Some(m) = set { + for entry in m.entries() { + let pair = js_sys::Array::from(&entry.map_err(|_| JsError::new("set: unreadable entry"))?); + let name = pair.get(0).as_string().ok_or_else(|| JsError::new("set: a file name is a string"))?; + let bytes = pair.get(1).dyn_into::().map_err(|_| JsError::new("set: a file's bytes are a Uint8Array"))?.to_vec(); + held.push((name, bytes)); + } + } + Ok(held) +} + +/// A base, ready for a sink: its info (content hash, the capability it was decoded +/// through, whether through a fallback, its tier, whether the sink verified a delivery +/// receipt for it, the live position of its unit 0) and its units. +/// +/// A base is **decoded** ([`StageBase::decode`]: the reference decode ran here, so its +/// units and tier are the decode's) or **claimed** ([`StageBase::from_stage0`]: the page +/// supplied its units, and nothing here checked them against the envelope). A claimed +/// base's tier is never above what its capability gives a base with no resolved pin, +/// and every unit, report and stage hash of its chain says `claimed`: a stage hash of a +/// claimed chain is the hash of what the page supplied and the layers applied to it, +/// not evidence that the envelope decodes to it. +#[wasm_bindgen] +pub struct StageBase { + info: BaseInfo, + units: Vec, + claimed: bool, +} + +#[wasm_bindgen] +impl StageBase { + /// Decode a base envelope by SPEC §4 through the integer path + /// (`wai::staged::decode_base_via`). `set` maps each parameter file's name (as SPEC + /// §5 names it) to its bytes, and is absent for a capability that reads none + /// (`wai.video.int_motion`); the envelope's pin is verified against those bytes. + /// `via` is the capability to decode through: absent to run §4 (the manifest's + /// capability, then its fallback), or one of those two. `limit` is required: the + /// page's cost limit, `{output_bytes, symbols, working_bytes, macs}` (`macs` + /// optional), or `null` for the payload format's caps only; leaving it out is an + /// error, so a page never decodes unbounded by accident. The limit binds a base + /// whose capability declares a cost before decoding (the integer payloads SPEC §5 + /// pins: `WIV1`, `WIH1`, `WIA1`, `WIS1`), checked before any symbol is decoded or + /// any output allocated (`cost_over_limit`); a capability that declares no cost + /// (`wai.neural.int_synth`, `wai.neural.int_mlicpp`) decodes under its format's caps + /// whatever the limit. No key log is held here, so SPEC §3.1's check 5 is not run. + /// A refusal's message starts with its code. + pub fn decode(envelope: &[u8], set: Option, via: Option, limit: JsValue) -> Result { + if limit.is_undefined() { + return Err(JsError::new("limit: required; pass the page's cost limit, or null for the payload format's caps only")); + } + let limit_v: Value = if limit.is_null() { Value::Null } else { serde_wasm_bindgen::from_value(limit).map_err(js_err)? }; + let limit = limit_from(&limit_v).map_err(js_err)?; + let held = set_from(set)?; + let files: BTreeMap = held.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); + let (info, units) = staged::decode_base_via(envelope, via.as_deref(), &files, limit.as_ref(), None).map_err(js_err)?; + Ok(StageBase { info, units, claimed: false }) + } + + /// A base the page decoded elsewhere (a codec this build does not run, or a + /// fixture): its envelope, the capability it was decoded through (the manifest's + /// capability, or its fallback, which makes it decoded through a fallback), the + /// tier the page claims for that step (`decode-equivalence`, `entropy-consistency` + /// or `none`; lowered to what the capability gives a base with no resolved pin), + /// and its units as their canonical buffers concatenated, `shapes` holding three + /// dimensions per unit and `elem` `u8` or `i16`. The base is **claimed**. + #[wasm_bindgen(js_name = fromStage0)] + pub fn from_stage0(envelope: &[u8], capability: &str, tier_token: &str, shapes: &[u32], elem: &str, stage0: &[u8]) -> Result { + let (info, units) = claim_base(envelope, capability, tier_token, shapes, elem, stage0).map_err(js_err)?; + Ok(StageBase { info, units, claimed: true }) + } + + /// The base envelope's content hash, lowercase hex. + #[wasm_bindgen(getter, js_name = contentHash)] + pub fn content_hash(&self) -> String { + hex(&self.info.content_hash) + } + #[wasm_bindgen(getter)] + pub fn capability(&self) -> String { + self.info.capability.clone() + } + #[wasm_bindgen(getter)] + pub fn tier(&self) -> String { + tier::token(self.info.tier).to_string() + } + #[wasm_bindgen(getter, js_name = nUnits)] + pub fn n_units(&self) -> u32 { + self.units.len() as u32 + } + /// Whether the base was decoded through its manifest's fallback (§4.1 step 3 + /// refuses its layers). Read-only: the decode, or the capability a claimed base + /// names, decides it. + #[wasm_bindgen(getter, js_name = viaFallback)] + pub fn via_fallback(&self) -> bool { + self.info.via_fallback + } + /// Whether the page supplied this base's units ([`StageBase::from_stage0`]). + #[wasm_bindgen(getter)] + pub fn claimed(&self) -> bool { + self.claimed + } + /// Whether the sink verified a delivery receipt for the base (§10). + #[wasm_bindgen(getter)] + pub fn attested(&self) -> bool { + self.info.attested + } + /// Record that the sink verified a delivery receipt for the base (§10): from then + /// on its layers need `delivered-verified` evidence. A verified receipt is never + /// unverified, so `false` changes nothing. + #[wasm_bindgen(setter)] + pub fn set_attested(&mut self, v: bool) { + self.info.attested |= v; + } + /// The live clock position of the base's unit 0 (§8). + #[wasm_bindgen(getter, js_name = firstUnit)] + pub fn first_unit(&self) -> f64 { + self.info.first_unit as f64 + } + /// Set the live clock position of the base's unit 0: a non-negative integer. + #[wasm_bindgen(js_name = setFirstUnit)] + pub fn set_first_unit(&mut self, v: f64) -> Result<(), JsError> { + self.info.first_unit = ms(v, "firstUnit")?; + Ok(()) + } + /// The units' canonical buffers concatenated, in unit order (SPEC §7 *Digests*). + #[wasm_bindgen(js_name = stage0)] + pub fn stage0(&self) -> Vec { + self.units.iter().flat_map(|u| u.samples.canonical_bytes()).collect() + } +} + +/// A staged-delivery sink: `wai::staged::StageSink`. Every method that decides returns +/// the reports of that step, an array in the corpus's form; a unit is read back by its +/// base's content hash (hex) and its index. A report, a base's units and a stage hash +/// from the chain of a claimed base carry `claimed: true` (see [`StageBase`]). +#[wasm_bindgen] +pub struct StageSink { + sink: staged::StageSink, + /// The bases whose chain started from a claimed base. + claimed: std::collections::BTreeSet<[u8; 32]>, +} + +/// Mark the reports of `rs` whose envelope belongs to a claimed chain: those `of` +/// says so of. +fn marked(rs: &[OfferReport], of: impl Fn(&OfferReport) -> bool) -> Value { + Value::Array(rs.iter().map(|r| { + let mut v = report_value(r); + if of(r) { + v["claimed"] = json!(true); + } + v + }).collect()) +} + +impl StageSink { + fn hash(base: &str) -> Result<[u8; 32], JsError> { + unhex32(base).ok_or_else(|| JsError::new("base: a content hash, 64 lowercase hex digits")) + } + fn chain(&self, base: &str) -> Result<&StageChain, JsError> { + self.sink.chain(&Self::hash(base)?).ok_or_else(|| JsError::new("base-absent: the sink holds no such base")) + } + fn state(&self, base: &str, u: u32) -> Result<&staged::UnitState, JsError> { + self.chain(base)?.unit(u).ok_or_else(|| JsError::new("unit: out of range")) + } + fn is_claimed(&self, base: &str) -> bool { + unhex32(base).is_some_and(|h| self.claimed.contains(&h)) + } +} + +/// A stage hash and whether it is a claimed chain's. +fn stage_hash(h: Option<[u8; 32]>, claimed: bool) -> Result { + match h { + None => Ok(JsValue::UNDEFINED), + Some(h) => to_js(&json!({"sha256": hex(&h), "claimed": claimed})), + } +} + +#[wasm_bindgen] +impl StageSink { + /// A sink with the configuration `config`, in the form a corpus case states it: + /// `{bound: {max_envelopes, max_age_ms}, cost_limit, budget, record_budget, live}` + /// (`cost_limit` and `live` null or absent for none). + #[wasm_bindgen(constructor)] + pub fn new(config: JsValue) -> Result { + let v: Value = serde_wasm_bindgen::from_value(config).map_err(js_err)?; + Ok(StageSink { sink: staged::StageSink::new(config_from(&v).map_err(js_err)?), claimed: Default::default() }) + } + + /// Add a base, which releases the layers held for it (§4.4). The base is moved + /// into the sink: the JS object is spent. A base the sink already holds is kept as + /// it is. A base whose units' output passes the chain budget is not added: its one + /// report refuses its envelope `chain-over-limit` (§4.5). + #[wasm_bindgen(js_name = addBase)] + pub fn add_base(&mut self, base: StageBase) -> Result { + let StageBase { info, units, claimed } = base; + let h = info.content_hash; + let held_before = self.sink.chain(&h).is_some(); + let rs = self.sink.add_base(info, units); + if claimed && !held_before && self.sink.chain(&h).is_some() { + self.claimed.insert(h); + } + let mark = claimed || self.claimed.contains(&h); + to_js(&marked(&rs, |_| mark)) + } + + /// Offer the envelope `envelope` at `now_ms` with the receipt evidence `evidence` + /// (`delivered-verified`, `inclusion-verified` or `none`). Bytes that are not a + /// `WAI1` envelope, or one this sink cannot open, are an error whose message + /// starts with `not-wai1` or `encrypted`. + pub fn offer(&mut self, envelope: &[u8], now_ms: f64, evidence: &str) -> Result { + let e = ReceiptEvidence::from_token(evidence).ok_or_else(|| JsError::new("evidence: delivered-verified, inclusion-verified or none"))?; + // An offer decides only the chain of the base the layer names. + let base = Wai::from_bytes(envelope).ok().and_then(|w| w.refines().and_then(Result::ok)).and_then(|r| unhex32(&r.base)); + let mark = base.is_some_and(|b| self.claimed.contains(&b)); + match self.sink.offer(envelope, ms(now_ms, "now_ms")?, e) { + Ok(r) => to_js(&marked(&r, |_| mark)), + Err(staged::OfferError::NotWai1(m)) => Err(JsError::new(&format!("not-wai1: {m}"))), + Err(e) => Err(JsError::new(&format!("encrypted: {e}"))), + } + } + + /// Discard what has been held past the bound at `now_ms` (§4.3). + pub fn expire(&mut self, now_ms: f64) -> Result { + // The envelopes claimed chains hold, before they expire. + let held: std::collections::BTreeSet<[u8; 32]> = + self.claimed.iter().filter_map(|b| self.sink.chain(b)).flat_map(|c| c.held().into_iter().map(|(h, _)| h)).collect(); + let rs = self.sink.expire(ms(now_ms, "now_ms")?); + to_js(&marked(&rs, |r| held.contains(&r.content_hash))) + } + + /// Drop a base's chain (§4.4); its units' records are kept. `null` when the sink + /// holds no such base. + #[wasm_bindgen(js_name = removeBase)] + pub fn remove_base(&mut self, base: &str) -> Result { + let h = Self::hash(base)?; + let mark = self.claimed.remove(&h); + match self.sink.remove_base(&h) { + Some(r) => to_js(&marked(&r, |_| mark)), + None => Ok(JsValue::NULL), + } + } + + /// Record that the sink verified a delivery receipt for the base (§10). + #[wasm_bindgen(js_name = markAttested)] + pub fn mark_attested(&mut self, base: &str) -> Result<(), JsError> { + self.sink.mark_attested(&Self::hash(base)?); + Ok(()) + } + + /// Discard what the sink keeps of a base it does not hold (§14). + #[wasm_bindgen(js_name = forgetBase)] + pub fn forget_base(&mut self, base: &str) -> Result { + Ok(self.sink.forget_base(&Self::hash(base)?)) + } + + /// Whether the sink holds the base. + #[wasm_bindgen(js_name = hasBase)] + pub fn has_base(&self, base: &str) -> bool { + unhex32(base).is_some_and(|h| self.sink.chain(&h).is_some()) + } + + /// Whether the base's chain started from a claimed base. + #[wasm_bindgen(js_name = isClaimed)] + pub fn is_claimed_js(&self, base: &str) -> bool { + self.is_claimed(base) + } + + /// The base's units and clip hash, as the corpus states a base at the end of a case + /// (`{units: [{stage, producer, frozen, tier, shape, stage_sha256}], clip_sha256, + /// n_units}`), with `claimed: true` on the base and on each unit where its chain + /// started from a claimed base; `null` when the sink holds no such base. + pub fn units(&self, base: &str) -> Result { + match unhex32(base).and_then(|h| self.sink.chain(&h)) { + Some(c) => { + let mut v = units_value(c); + if self.is_claimed(base) { + v["claimed"] = json!(true); + for u in v["units"].as_array_mut().into_iter().flatten() { + u["claimed"] = json!(true); + } + } + to_js(&v) + } + None => Ok(JsValue::NULL), + } + } + + #[wasm_bindgen(js_name = nUnits)] + pub fn n_units(&self, base: &str) -> Result { + Ok(self.chain(base)?.n_units()) + } + #[wasm_bindgen(js_name = unitStage)] + pub fn unit_stage(&self, base: &str, u: u32) -> Result { + Ok(self.state(base, u)?.stage as u32) + } + /// The content hash of the envelope that produced the unit's stage (the base at + /// stage 0). + #[wasm_bindgen(js_name = unitProducer)] + pub fn unit_producer(&self, base: &str, u: u32) -> Result { + Ok(hex(&self.state(base, u)?.producer)) + } + /// Whether the unit is frozen by a conflict (§4.1 step 7). + #[wasm_bindgen(js_name = unitFrozen)] + pub fn unit_frozen(&self, base: &str, u: u32) -> Result { + Ok(self.state(base, u)?.frozen) + } + #[wasm_bindgen(js_name = unitTier)] + pub fn unit_tier(&self, base: &str, u: u32) -> Result { + Ok(tier::token(self.state(base, u)?.tier).to_string()) + } + /// The unit's shape: rows, columns and channels of a picture. + #[wasm_bindgen(js_name = unitShape)] + pub fn unit_shape(&self, base: &str, u: u32) -> Result, JsError> { + let s = self.state(base, u)?.buf.shape; + Ok(vec![s.0, s.1, s.2]) + } + /// The unit's current output, 8-bit samples (RGB8 for a picture). An error for a + /// unit of 16-bit samples. + #[wasm_bindgen(js_name = unitBytes)] + pub fn unit_bytes(&self, base: &str, u: u32) -> Result, JsError> { + match &self.state(base, u)?.buf.samples { + Samples::U8(v) => Ok(v.clone()), + Samples::I16(_) => Err(JsError::new("the unit holds 16-bit samples: unitPcm")), + } + } + /// The unit's current output, 16-bit samples, channels interleaved. + #[wasm_bindgen(js_name = unitPcm)] + pub fn unit_pcm(&self, base: &str, u: u32) -> Result, JsError> { + match &self.state(base, u)?.buf.samples { + Samples::I16(v) => Ok(v.clone()), + Samples::U8(_) => Err(JsError::new("the unit holds 8-bit samples: unitBytes")), + } + } + /// The unit's stage hash (SHA-256 of its canonical buffer) as `{sha256, claimed}`, + /// or `undefined` unless the unit holds `decode-equivalence` (§5). `claimed` is true + /// where the chain started from a claimed base. + #[wasm_bindgen(js_name = unitSha256)] + pub fn unit_sha256(&self, base: &str, u: u32) -> Result { + let c = self.chain(base)?; + if u >= c.n_units() { + return Err(JsError::new("unit: out of range")); + } + stage_hash(c.unit_sha256(u), self.is_claimed(base)) + } + /// The base's clip hash as `{sha256, claimed}`, or `undefined` unless every unit + /// holds `decode-equivalence`. + #[wasm_bindgen(js_name = clipSha256)] + pub fn clip_sha256(&self, base: &str) -> Result { + stage_hash(self.chain(base)?.clip_sha256(), self.is_claimed(base)) + } + /// The squared error of the unit's current 8-bit output against `reference`, + /// summed exactly (a `BigInt`): the numerator of the page's PSNR. + #[wasm_bindgen(js_name = unitSseU8)] + pub fn unit_sse_u8(&self, base: &str, u: u32, reference: &[u8]) -> Result { + match &self.state(base, u)?.buf.samples { + Samples::U8(v) => sse_u8(v, reference).map_err(js_err), + Samples::I16(_) => Err(JsError::new("the unit holds 16-bit samples")), + } + } + /// The envelopes held for the base's units, `[[object, unit], …]`. + pub fn held(&self, base: &str) -> Result { + let v: Vec = self.chain(base)?.held().iter().map(|(h, u)| json!([hex(h), u])).collect(); + to_js(&Value::Array(v)) + } + /// The chain's state in bytes (§4.5). + #[wasm_bindgen(js_name = stateBytes)] + pub fn state_bytes(&self, base: &str) -> Result { + Ok(self.chain(base)?.state_bytes() as f64) + } + /// The bytes the sink keeps of bases it does not hold (§4.4). + #[wasm_bindgen(js_name = keptBytes)] + pub fn kept_bytes(&self) -> f64 { + self.sink.kept_bytes() as f64 + } +} + +/// What an envelope is, before it is offered or decoded: `{object, bytes, +/// framing_bytes, manifest_bytes, payload_bytes, capability, refinement_only, refines, +/// wir1, declared_cost, error}`, where `wir1` is a layer's header (`n_units`, `shape`, +/// `predictor`, `q`, `tables`) and `declared_cost` a layer's output bytes, symbols and +/// working bytes (Appendix A.5), or a base's declared decode cost +/// (`extensions/integer-payloads.md` §7, with `macs` as a decimal string and the +/// capability it was read for as `cost_capability`), against the parameter set `set` +/// where the base's codec reads one. Not a `WAI1` envelope: an error starting +/// `not-wai1`. +#[wasm_bindgen(js_name = stageLayerInfo)] +pub fn stage_layer_info(envelope: &[u8], set: Option) -> Result { + let held = set_from(set)?; + let files: BTreeMap = held.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); + to_js(&layer_info(envelope, &files).map_err(js_err)?) +} + +/// The content hash of an envelope (jwp-receipts §2.1), lowercase hex: the identity a +/// layer's `refines` names its base and producers by. +#[wasm_bindgen(js_name = stageContentHash)] +pub fn stage_content_hash(envelope: &[u8]) -> String { + hex(&wai::jwp_hash::content_hash(envelope)) +} + +/// The bytes of linear memory this instance holds: what a page or a test reads to see +/// what a call allocated (linear memory only grows). 0 off wasm32. +#[wasm_bindgen(js_name = wasmMemoryBytes)] +pub fn wasm_memory_bytes() -> f64 { + #[cfg(target_arch = "wasm32")] + { + (core::arch::wasm32::memory_size(0) as f64) * 65536.0 + } + #[cfg(not(target_arch = "wasm32"))] + { + 0.0 + } +} + +/// The exact sum of squared differences of two equal-length byte buffers (a `BigInt`). +#[wasm_bindgen(js_name = sseU8)] +pub fn sse_u8_js(a: &[u8], b: &[u8]) -> Result { + sse_u8(a, b).map_err(js_err) +} + +/// The same for 16-bit samples. +#[wasm_bindgen(js_name = sseI16)] +pub fn sse_i16_js(a: &[i16], b: &[i16]) -> Result { + sse_i16(a, b).map_err(js_err) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn wai_dir() -> std::path::PathBuf { + std::path::Path::new(env!("CARGO_MANIFEST_DIR")).join("..") + } + + /// Every corpus case, replayed through the binding's pure-Rust half the way + /// `demo/verify_staged_wasm.mjs` drives the wasm build: each base decoded through + /// `decode_base_via` where its decode is the integer path (and required to give its + /// stage-0 buffer and tier), from `stage0.bin` otherwise, then given the case's + /// facts; the sink configured from `case.json`'s `config`; every step's reports and + /// every base's final units in the corpus's form, equal to what the case states. + #[test] + fn the_corpus_replays_through_the_binding() { + let root = wai_dir().join("staged-conformance/cases"); + let mut names: Vec<_> = std::fs::read_dir(&root).unwrap().flatten().map(|e| e.file_name().to_string_lossy().into_owned()).collect(); + names.sort(); + assert!(names.len() >= 60, "the corpus is present"); + for name in names { + let dir = root.join(&name); + let case: Value = serde_json::from_slice(&std::fs::read(dir.join("case.json")).unwrap()).unwrap(); + let mut bases = Vec::new(); + for b in case["bases"].as_array().unwrap() { + let env = std::fs::read(dir.join(b["file"].as_str().unwrap())).unwrap(); + let stage0 = std::fs::read(dir.join(b["stage0"].as_str().unwrap())).unwrap(); + let shapes: Vec = b["units"].as_array().unwrap().iter().flat_map(|u| u["shape"].as_array().unwrap().iter().map(|x| x.as_u64().unwrap() as u32).collect::>()).collect(); + let elem = b["units"][0]["elem"].as_str().unwrap(); + let via = b["decoded_via"].as_str().unwrap(); + let units = if b["decode"]["kind"] == "int" { + let held: Vec<(String, Vec)> = b["decode"]["set"].as_object().unwrap().iter() + .map(|(m, p)| (m.clone(), std::fs::read(wai_dir().join(p.as_str().unwrap())).unwrap())) + .collect(); + let files: BTreeMap = held.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect(); + let (info, units) = staged::decode_base_via(&env, Some(via), &files, None, None).unwrap_or_else(|e| panic!("{name}: {e}")); + let got: Vec = units.iter().flat_map(|u| u.samples.canonical_bytes()).collect(); + assert_eq!(got, stage0, "{name}: the base decodes to its stage 0"); + assert_eq!(tier::token(info.tier), b["tier"], "{name}: the base tier"); + assert_eq!(info.via_fallback, b["via_fallback"] == true, "{name}: via fallback"); + units + } else { + units_from_stage0(&shapes, elem, &stage0).unwrap() + }; + let info = BaseInfo { + content_hash: wai::jwp_hash::content_hash(&env), + capability: via.to_string(), + via_fallback: b["via_fallback"] == true, + tier: tier::from_token(b["tier"].as_str().unwrap()), + attested: b["attested"] == true, + first_unit: b["first_unit"].as_u64().unwrap(), + }; + assert_eq!(hex(&info.content_hash), b["object"], "{name}: the base's content hash"); + bases.push((info, units)); + } + let mut sink = staged::StageSink::new(config_from(&case["config"]).unwrap()); + let files: BTreeMap> = case["files"].as_array().unwrap().iter() + .map(|f| { let n = f["file"].as_str().unwrap().to_string(); (n.clone(), std::fs::read(dir.join(&n)).unwrap()) }) + .collect(); + let mut reports = Vec::new(); + for s in case["steps"].as_array().unwrap() { + let at = s["at_ms"].as_u64().unwrap(); + let idx = || s["index"].as_u64().unwrap() as usize; + let got = match s["op"].as_str().unwrap() { + "offer" => match sink.offer(&files[s["file"].as_str().unwrap()], at, ReceiptEvidence::from_token(s["evidence"].as_str().unwrap()).unwrap()) { + Ok(r) => marked(&r, |_| false), + Err(staged::OfferError::NotWai1(_)) => json!({"error": "not-wai1"}), + Err(_) => json!({"error": "encrypted"}), + }, + "base" => { let (i, u) = bases[idx()].clone(); marked(&sink.add_base(i, u), |_| false) } + "expire" => marked(&sink.expire(at), |_| false), + "remove" => marked(&sink.remove_base(&bases[idx()].0.content_hash).unwrap_or_default(), |_| false), + "attest" => { sink.mark_attested(&bases[idx()].0.content_hash); json!([]) } + o => panic!("{name}: op {o}"), + }; + reports.push(got); + } + assert_eq!(Value::Array(reports), case["expect"]["reports"], "{name}: the reports"); + let finals: Vec = bases.iter().map(|(i, _)| sink.chain(&i.content_hash).map(units_value).unwrap_or(Value::Null)).collect(); + assert_eq!(Value::Array(finals), case["expect"]["bases"], "{name}: the final units"); + } + } + + /// The configuration reader takes the corpus's form, maps a figure at or past + /// `2^64` (as JavaScript reads `u64::MAX`) to the largest `u64`, and refuses a + /// negative or fractional figure and a member it does not know. + #[test] + fn config_reads_the_corpus_form() { + let c = config_from(&json!({"bound": {"max_envelopes": 4, "max_age_ms": 1000}, "cost_limit": {"output_bytes": 100, "symbols": 18446744073709551616.0_f64, "working_bytes": 18446744073709551615u64}, "budget": 1, "record_budget": 2, "live": {"t0_ms": 1, "unit_ms": 40, "target_latency_ms": 3, "lead_ms": [10, 30]}})).unwrap(); + assert_eq!((c.bound.max_envelopes, c.bound.max_age_ms, c.budget, c.record_budget), (4, 1000, 1, 2)); + let l = c.cost_limit.unwrap(); + assert_eq!((l.output_bytes, l.symbols, l.working_bytes), (100, u64::MAX, u64::MAX)); + assert_eq!(c.live.unwrap().lead_ms, vec![10, 30]); + assert!(config_from(&json!({"bound": {"max_envelopes": 4, "max_age_ms": -1}, "budget": 1, "record_budget": 2})).is_err()); + assert!(config_from(&json!({"bound": {"max_envelopes": 4.5, "max_age_ms": 1}, "budget": 1, "record_budget": 2})).is_err()); + assert!(config_from(&json!({"bound": {"max_envelopes": 4, "max_age_ms": 1}, "budget": 1, "record_budget": 2, "extra": 0})).is_err()); + let ok = config_from(&json!({"bound": {"max_envelopes": 4, "max_age_ms": 1}, "budget": 1, "record_budget": 2})).unwrap(); + assert!(ok.cost_limit.is_none() && ok.live.is_none()); + } + + /// A layer's header and declared cost are read before it is offered, and its bytes + /// are its framing, manifest and payload; a base has no `refines` or `WIR1`. + #[test] + fn layer_info_reads_header_cost_and_bytes() { + let d = wai_dir().join("staged-conformance/cases/video_int_motion_gop_perframe"); + let l2 = std::fs::read(d.join("g0_L2.wai")).unwrap(); + let none = BTreeMap::new(); + let v = layer_info(&l2, &none).unwrap(); + assert_eq!(v["bytes"].as_u64().unwrap() as usize, l2.len()); + assert_eq!(v["framing_bytes"].as_u64().unwrap() + v["manifest_bytes"].as_u64().unwrap() + v["payload_bytes"].as_u64().unwrap(), l2.len() as u64); + assert_eq!(v["framing_bytes"], 12); + assert_eq!(v["refines"]["layer"], 2); + assert_eq!(v["refines"]["prev"].as_array().map(Vec::len), Some(4), "a run per frame's producer"); + assert_eq!(v["wir1"]["n_units"], 4); + assert!(v["declared_cost"]["output_bytes"].as_u64().unwrap() > 0); + assert_eq!(v["error"], Value::Null); + let b = layer_info(&std::fs::read(d.join("base_0.wai")).unwrap(), &none).unwrap(); + assert_eq!((b["refines"].is_null(), b["wir1"].is_null(), b["refinement_only"].as_bool()), (true, true, Some(false))); + // A base's declared decode cost, read before it is decoded: four 16 × 16 frames. + assert_eq!(b["declared_cost"]["output_bytes"], 4 * 16 * 16 * 3); + assert_eq!((b["cost_capability"].as_str(), b["declared_cost"]["macs"].as_str()), (Some("wai.video.int_motion"), Some("0"))); + assert!(layer_info(b"nope", &none).unwrap_err().starts_with("not-wai1")); + assert_eq!(sse_u8(&[0, 10, 255], &[1, 7, 0]).unwrap(), 1 + 9 + 255 * 255); + assert!(sse_u8(&[0], &[]).is_err()); + assert_eq!(sse_i16(&[-32768], &[32767]).unwrap(), 65535 * 65535); + } + + /// A claimed base: the capability must be the manifest's or its fallback (decoded + /// through the fallback where it is not the capability); the claimed tier is a + /// tier token, lowered to what the capability gives a base with no resolved pin. + #[test] + fn a_claimed_base_cannot_raise_its_facts() { + let c = wai_dir().join("staged-conformance/cases"); + let unpinned = std::fs::read(c.join("image_int_hyper_unpinned/base.wai")).unwrap(); + let shapes = [16u32, 16, 3]; + let junk = vec![77u8; 16 * 16 * 3]; + let err = claim_base(&unpinned, "wai.image.png", "decode-equivalence", &shapes, "u8", &junk).unwrap_err(); + assert!(err.starts_with("not-offered"), "{err}"); + let (info, _) = claim_base(&unpinned, "wai.neural.int_hyper", "decode-equivalence", &shapes, "u8", &junk).unwrap(); + assert_eq!(info.tier, None, "an integer neural base claims no tier without a resolved pin"); + assert!(!info.via_fallback); + assert!(claim_base(&unpinned, "wai.neural.int_hyper", "decode-equivalance", &shapes, "u8", &junk).unwrap_err().starts_with("tier")); + let png = std::fs::read(c.join("image_png_up2/base.wai")).unwrap(); + let (info, _) = claim_base(&png, "wai.image.png", "decode-equivalence", &[16, 16, 3], "u8", &junk).unwrap(); + assert_eq!(info.tier, Some(wai::codecs::DeterminismTier::DecodeEquivalence), "a standard-defined capability keeps the tier claimed"); + let (info, _) = claim_base(&png, "wai.image.png", "none", &[16, 16, 3], "u8", &junk).unwrap(); + assert_eq!(info.tier, None, "a page may lower the tier"); + let fb = std::fs::read(c.join("base_via_fallback/base.wai")).unwrap(); + let w = Wai::from_bytes(&fb).unwrap(); + let shape: Vec = vec![8, 8, 3, 8, 8, 3]; + let (info, _) = claim_base(&fb, w.fallback().unwrap(), "decode-equivalence", &shape, "u8", &vec![0u8; 2 * 8 * 8 * 3]).unwrap(); + assert!(info.via_fallback, "named through its fallback, decoded through a fallback"); + assert!(claim_base(b"nope", "x", "none", &shapes, "u8", &junk).unwrap_err().starts_with("not-wai1")); + // SPEC §4: an inert envelope (a refinement-only or companion-only capability) + // is claimed neither through its capability nor through its fallback. + let b = wai_dir().join("wai-rs/tests/vectors/staged_bounds"); + let two = [4u32, 4, 3, 4, 4, 3]; + for f in ["inert_refinement_primary.wai", "inert_companion_primary.wai"] { + let env = std::fs::read(b.join(f)).unwrap(); + let w = Wai::from_bytes(&env).unwrap(); + for cap in [w.capability(), w.fallback().unwrap()] { + let err = claim_base(&env, cap, "decode-equivalence", &two, "u8", &[5u8; 96]).unwrap_err(); + assert!(err.starts_with("missing_capability"), "{f} as {cap}: {err}"); + } + } + } + + /// The cost limit's form: four figures, `macs` optional; `null` for none. + #[test] + fn limit_reads_its_form() { + assert_eq!(limit_from(&Value::Null).unwrap(), None); + let l = limit_from(&json!({"output_bytes": 10, "symbols": 20, "working_bytes": 30})).unwrap().unwrap(); + assert_eq!((l.output_bytes, l.symbols, l.working_bytes, l.macs), (10, 20, 30, u128::MAX)); + let l = limit_from(&json!({"output_bytes": 1, "symbols": 2, "working_bytes": 3, "macs": 4})).unwrap().unwrap(); + assert_eq!(l.macs, 4); + assert!(limit_from(&json!({"output_bytes": 1, "symbols": 2, "working_bytes": 3, "extra": 0})).is_err()); + assert!(limit_from(&json!({"output_bytes": -1, "symbols": 2, "working_bytes": 3})).is_err()); + assert!(limit_from(&json!(5)).is_err()); + } +}