diff --git a/.claude/CHACHA20_MATRYOSHKA_PLAN.md b/.claude/CHACHA20_MATRYOSHKA_PLAN.md new file mode 100644 index 00000000..bb6c961d --- /dev/null +++ b/.claude/CHACHA20_MATRYOSHKA_PLAN.md @@ -0,0 +1,106 @@ +# ChaCha20 SIMD via the `U32x16` ARX lane — matryoshka usage + open work + +> Handoff doc. What this branch shipped, what's deferred, and the **matryoshka** +> plan for *using* the `U32x16` lane to accelerate ChaCha20 without owning any +> crypto algorithm. Operator-ratified design (2026-07-11 session). + +## The doctrine (why this shape) + +- **RustCrypto owns the algorithm; ndarray owns the SIMD.** We do NOT hand-roll + ChaCha20/Poly1305/the XChaCha20-Poly1305 AEAD. Rolling your own AEAD (HChaCha20 + + Poly1305 + framing) is the footgun; it is forbidden. +- The `encryption` crate (`crates/encryption`, RustCrypto wrappers) stays + **unchanged**. See its `aead.rs` module doc for why the AEAD is not re-wired. +- The only crypto-algorithm SIMD we accelerate is the **ChaCha20 keystream ARX + core**, and even that is done by feeding RustCrypto's *own* backend a new + SIMD lane — never by re-implementing the cipher. + +## DONE on this branch (the lane is ready) + +`ndarray::simd::U32x16` now carries the full ARX triple — `Add` + `BitXor` + +`rotate_left` — on every tier, so a ChaCha20 backend can ride it: + +| tier | `U32x16` backing | `rotate_left` | verified | +|---|---|---|---| +| avx512 (server) | `__m512i` | `_mm512_rolv_epi32` (VPROLVD) | executed, parity-green | +| wasm128 (browser) | `[U32x4; 4]` (`v128`) | `v128_or(u32x4_shl, u32x4_shr)` | full-lib wasm build + **node parity** | +| avx2 | `[u32; 16]` polyfill (native 2×`__m256i` = deferred TD-SIMD-3) | `u32::rotate_left` loop | compiles; == reference | +| scalar / nightly | `[u32; 16]` / `core::simd` | `u32::rotate_left` / shift-or | completeness tier + reference | + +CI: `wasm_simd` job (`.github/workflows/ci.yaml` + `scripts/wasm-parity.sh` + +`crates/wasm-simd-parity/`) builds the real `ndarray::simd` wasm types and runs a +lane-by-lane parity selfcheck under **node** — the standing guard for the wasm +tier (invisible to the x86 `cargo test`). Extend its per-lane blocks as lanes land. + +The stand-alone `src/simd_crypto.rs` (`chacha20_block`/`chacha20_keystream`, +scalar+AVX-512+wasm128, RustCrypto-parity-proven) is the *interim* primitive; the +matryoshka below **supersedes it** (RustCrypto supplies the rounds, so the +from-scratch cipher retires once the fork lands). + +## DEFERRED — TO-DO (next session) + +1. **Native neon `U32x16 = [U32x4; 4]`.** `simd_neon.rs` has native + `U32x4(uint32x4_t)` with add/sub/min/max; add `bitxor` (`veorq_u32`) + + `rotate_left` (shift-or via `vshlq_u32` with variable count vectors, `n%32` + guard), then compose `U32x16([U32x4; 4])` with `impl Add`/`impl BitXor`/ + `rotate_left` (fan over 4 lanes, mirror the wasm shape in `simd_wasm.rs`), fix + `F32x16::to_bits/from_bits` to `from_array`/`to_array`, and swap the `simd.rs` + aarch64 arm's scalar `U32x16` re-export for the native one. +2. **aarch64 cross parity CI.** Generalize `wasm-simd-parity` into a shared + `simd-parity` harness (or add a sibling) run under **qemu** via `cross` in a + new `neon_simd` CI job — same selfcheck, closing NEON's identical x86-suite + blind spot (its `F32x16`/`I8x16` are unverified in CI today too). +3. **avx2 native `U32x16`** (2×`__m256i`) — the TD-SIMD-3 lowering; optional, the + scalar polyfill is correct meanwhile. + +## The MATRYOSHKA — how the lane gets USED (the finalization) + +Goal: ChaCha20 (and thus the whole `encryption` AEAD stack) accelerated on +server (AVX-512) + browser (wasm128), with **essentially zero owned crypto code** +and a **one-file** delta over RustCrypto that is trivial to re-sync on security +updates. + +**Structure (nesting):** + +``` +RustCrypto chacha20 (rounds / constants / counter / StreamBackend / TESTS — VERBATIM, vetted) + └─ one backend's round body expressed over → ndarray::simd::U32x16 (Add / BitXor / rotate_left — generic, NOT crypto) + └─ dispatched by ndarray's polyfill → avx512 / wasm128 / neon / scalar +``` + +**Steps:** + +1. **Fork `chacha20`** (the RustCrypto stream-cipher crate; `AdaWorldAPI/chacha20`, + or vendor). Its backends live in `src/backends/{soft,sse2,avx2,neon}.rs`, each a + `struct Backend` impl of `StreamBackend` with `gen_ks_block` / + `gen_par_ks_blocks`. **Clone `avx2.rs` once**; keep soft/sse2/avx2/neon + untouched. +2. **Rewire the clone over `ndarray::simd::U32x16`** — the word-sliced ChaCha + round (`state[a] = state[a] + state[b]; state[d] = (state[d] ^ state[a]).rotate_left(16);` + …) written against `U32x16`, `ParBlocksSize = U16`. `ndarray` lowers it to + AVX-512 (server) / wasm128 (browser) / NEON / scalar automatically. The + carried file has **no `unsafe`, no intrinsics** — all of that lives once in + `ndarray::simd` (audited once; only AMX+F16 are byte-asm, neither on this path). +3. **Dispatch:** add the generic backend to the fork's `lib.rs` selection. + Per the ndarray model this is **compile-time** (`cfg(target_feature)`), not + `is_x86_feature_detected!` — server built `x86-64-v4`, browser built `+simd128`. + Non-portable per-target binaries by design (SIGILL/validation-fail on a CPU + built-for-but-absent) — matches how the servers/browsers are actually built. +4. **`[patch]` the fork in** (the pattern already used in MedCare-rs for + `encryption`): `chacha20poly1305 → encryption → ogar-encryption → consumers` + all accelerate **transitively**, with **zero** change to the AEAD or any + consumer, because the AEAD just calls `apply_keystream` which now dispatches + to the new backend. HChaCha20 + Poly1305 + framing stay 100% RustCrypto. +5. **Gate bit-exact vs RustCrypto's own `soft` backend** + run its stock test + vectors — the same "reference + KAT, then vectorize with parity" discipline. +6. **Retire `src/simd_crypto.rs`** — the rounds now come from RustCrypto; only + the `U32x16` lane (not crypto) is ours. + +**Maintenance win:** a RustCrypto security advisory almost never touches the ARX +lane ops (fixes land in framing/counter/AEAD — all pristine upstream), so the +one-file delta re-applies with ~zero conflict: bump the vendored `chacha20`, +re-apply the one backend file, run the stock vectors + the soft parity. And if +the backend is ever **upstreamed** (RustCrypto would plausibly take an avx512 / +wasm128 backend), the fork evaporates → maintenance goes to zero. AMX has no +ChaCha/Poly backend (it is a matrix engine); the upstream targets are avx512 + +wasm128 (NEON already exists upstream). diff --git a/.claude/blackboard.md b/.claude/blackboard.md index d311673a..4aff62da 100644 --- a/.claude/blackboard.md +++ b/.claude/blackboard.md @@ -609,3 +609,24 @@ index was transiently stale this session (returned empty for `simd_int_ops.rs`, `vnni_gemm.rs`, `bf16_gemm_f32`); Bash ground-truth confirmed all present. Orchestrator should `cargo fmt`/`clippy`/`test` centrally (edits were edit-only, no compile performed here). + +--- + +## 2026-07-11 — U32x16 ARX lane + ChaCha20 matryoshka (see .claude/CHACHA20_MATRYOSHKA_PLAN.md) + +**DONE:** `ndarray::simd::U32x16` full ARX triple (Add/BitXor/**rotate_left**) on +every tier — avx512 native `_mm512_rolv_epi32`, native wasm `[U32x4;4]`, avx2/ +scalar/nightly. Node-run wasm parity CI gate (`wasm_simd` job + `scripts/wasm-parity.sh` ++ `crates/wasm-simd-parity/`, workspace-excluded, tests the real types, no drift). +Interim `src/simd_crypto.rs` (chacha20 scalar+avx512+wasm128, RustCrypto-parity-proven) +is superseded by the matryoshka once the fork lands. + +**TO-DO (deferred, token limit):** +1. Native neon `U32x16 = [U32x4;4]` (extend `U32x4(uint32x4_t)` w/ xor+rotl, compose, + fix F32x16 to_bits, swap simd.rs aarch64 arm). +2. aarch64 cross (qemu) parity CI job — generalize `wasm-simd-parity` to a shared + `simd-parity` harness; closes NEON's x86-suite blind spot. +3. avx2 native `U32x16` (2×__m256i, TD-SIMD-3) — optional. +4. **Matryoshka execution:** fork `chacha20`, clone `avx2.rs` backend → rewire over + `ndarray::simd::U32x16`, `[patch]` into the encryption stack (transitive accel), + gate vs RustCrypto `soft`, retire `simd_crypto.rs`. Full plan in the doc above. diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index f7ab91ca..0dd21577 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -112,6 +112,29 @@ jobs: # crates that have no business existing in the nostd build. cargo rustc -p ndarray "--target=${{ matrix.target }}" --no-default-features --features portable-atomic-critical-section + wasm_simd: + # The x86 `cargo test` suite never touches `simd_wasm.rs` (it is cfg'd to + # wasm32), so wasm SIMD correctness is otherwise unverified. Build the real + # `ndarray::simd` wasm types (+simd128) via the excluded `wasm-simd-parity` + # cdylib and run its numeric selfcheck under node — both integration (the + # lib compiles for wasm32+simd128) and parity (each lane == scalar). The + # standing version of the one-off node check from commit 500d57e6. + runs-on: ubuntu-latest + name: wasm-simd/parity-node + steps: + - uses: actions/checkout@v4 + - uses: dtolnay/rust-toolchain@stable + with: + targets: wasm32-unknown-unknown + # rust-toolchain.toml pins 1.95.0 — install the wasm target for the pinned + # toolchain too (dtolnay installs for `stable`, which may differ). + - run: rustup target add wasm32-unknown-unknown + - uses: actions/setup-node@v4 + with: + node-version: "22" + - name: wasm SIMD parity (build + node run) + run: ./scripts/wasm-parity.sh + tests: runs-on: ubuntu-latest needs: pass-msrv @@ -335,6 +358,7 @@ jobs: - clippy - format - nostd + - wasm_simd - tests - native-backend - hpc-stream-parallel diff --git a/Cargo.lock b/Cargo.lock index 48855aa3..bc704f14 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1157,6 +1157,7 @@ dependencies = [ "approx", "blake3", "cblas-sys", + "chacha20", "cranelift-codegen", "cranelift-frontend", "cranelift-jit", diff --git a/Cargo.toml b/Cargo.toml index 0bdb4e54..0e826fa7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -217,6 +217,11 @@ approx = { workspace = true, default-features = true } itertools = { workspace = true } ndarray-gen = { workspace = true } criterion = { version = "0.5", features = ["html_reports"] } +# The vetted RustCrypto ChaCha20 stream cipher — used ONLY as the independent +# oracle for the `chacha20_rustcrypto_parity` integration test that pins +# `ndarray::simd::chacha20_keystream` byte-for-byte against a trusted reference. +# Same version chacha20poly1305 0.10 pulls transitively (see Cargo.lock). +chacha20 = "0.9" [[bench]] name = "append" @@ -402,7 +407,7 @@ members = [ "ndarray-rand", "crates/*", ] -exclude = ["crates/burn"] +exclude = ["crates/burn", "crates/wasm-simd-parity"] default-members = [ ".", "ndarray-rand", diff --git a/crates/encryption/src/aead.rs b/crates/encryption/src/aead.rs index 25fbd452..777c987e 100644 --- a/crates/encryption/src/aead.rs +++ b/crates/encryption/src/aead.rs @@ -5,6 +5,19 @@ //! is no counter state to persist or synchronise between a browser tab //! and a server. Callers normally use [`crate::envelope`], which manages //! nonce generation and layout; this module is the raw primitive. +//! +//! ## Deliberately NOT hand-composed from `ndarray::simd::chacha20_keystream` +//! +//! `ndarray` now ships a hardware-accelerated ChaCha20 keystream +//! (`ndarray::simd::chacha20_keystream`, AVX-512 / wasm128, byte-parity-proven +//! against RustCrypto). It is tempting to swap this AEAD's keystream to that +//! primitive for speed — **do not.** XChaCha20-Poly1305 is not just a keystream: +//! it is HChaCha20 subkey derivation + the Poly1305 one-time-key + MAC framing + +//! the AAD/length encoding. Re-wiring only the keystream means re-implementing +//! that authenticated composition by hand, which is exactly the "roll your own +//! AEAD" footgun the stack forbids. The vetted RustCrypto `XChaCha20Poly1305` +//! construction stays here. The accelerated primitive is for *raw-stream* use +//! sites whose caller already owns a vetted MAC/framing — not this module. use chacha20poly1305::aead::{Aead, KeyInit, Payload}; use chacha20poly1305::{Key, XChaCha20Poly1305, XNonce}; diff --git a/crates/wasm-simd-parity/Cargo.lock b/crates/wasm-simd-parity/Cargo.lock new file mode 100644 index 00000000..f734f708 --- /dev/null +++ b/crates/wasm-simd-parity/Cargo.lock @@ -0,0 +1,192 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3fb67a6e08acf24fdeccbac2cb6ac4305825bd1f117462e0e6f2f193345ad56" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures", +] + +[[package]] +name = "cc" +version = "1.2.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e17dd265a7d0f31ef544e1b20e03add05d3b45b491b633b10d67145d2acc1a38" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "fractal" +version = "0.1.0" +dependencies = [ + "libm", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + +[[package]] +name = "matrixmultiply" +version = "0.3.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06de3016e9fae57a36fd14dba131fccf49f74b40b7fbdb472f96e361ec71a08" +dependencies = [ + "autocfg", + "rawpointer", +] + +[[package]] +name = "ndarray" +version = "0.17.2" +dependencies = [ + "blake3", + "fractal", + "matrixmultiply", + "num-complex", + "num-integer", + "num-traits", + "p64", + "paste", + "portable-atomic", + "portable-atomic-util", + "rawpointer", +] + +[[package]] +name = "num-complex" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73f88a1307638156682bada9d7604135552957b7818057dcef22705b4d509495" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-integer" +version = "0.1.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f" +dependencies = [ + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "p64" +version = "0.1.0" +dependencies = [ + "fractal", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "rawpointer" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60a357793950651c4ed0f3f52338f53b2f809f32d83a07f72909fa13e4c6c1e3" + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "wasm-simd-parity" +version = "0.0.0" +dependencies = [ + "ndarray", +] diff --git a/crates/wasm-simd-parity/Cargo.toml b/crates/wasm-simd-parity/Cargo.toml new file mode 100644 index 00000000..eb6d4249 --- /dev/null +++ b/crates/wasm-simd-parity/Cargo.toml @@ -0,0 +1,25 @@ +# wasm-simd-parity — the run-under-node gate for the wasm SIMD tier. +# +# The x86 `cargo test` suite never touches `simd_wasm.rs` (it is `cfg`'d to +# wasm32), so wasm SIMD correctness is otherwise unverified by CI. This crate +# is a cdylib that calls the REAL `ndarray::simd` wasm types (no copy → no +# drift) and self-checks each lane's arithmetic against scalar in-module; the +# `.wasm` is built with `+simd128` and run under node, asserting rc == 0. +# +# EXCLUDED from the workspace (see root Cargo.toml `exclude`) so it has zero +# effect on the x86 jobs; the CI `wasm_simd` job builds it via --manifest-path. +[package] +name = "wasm-simd-parity" +version = "0.0.0" +edition = "2021" +publish = false + +[lib] +crate-type = ["cdylib"] + +[dependencies] +ndarray = { path = "../.." } + +[profile.release] +panic = "abort" +opt-level = "z" diff --git a/crates/wasm-simd-parity/run.mjs b/crates/wasm-simd-parity/run.mjs new file mode 100644 index 00000000..74e631c0 --- /dev/null +++ b/crates/wasm-simd-parity/run.mjs @@ -0,0 +1,17 @@ +// Node harness for the wasm SIMD parity gate: instantiate the .wasm and assert +// selfcheck() == 0. A nonzero rc names the failing lane+op (see lib.rs codes). +import { readFileSync } from 'node:fs'; +const wasmPath = process.argv[2]; +if (!wasmPath) { + console.error('usage: node run.mjs '); + process.exit(2); +} +const { instance } = await WebAssembly.instantiate(readFileSync(wasmPath), {}); +const rc = instance.exports.selfcheck(); +if (rc === 0) { + console.log('wasm SIMD parity: OK (selfcheck rc=0)'); + process.exit(0); +} else { + console.error(`wasm SIMD parity: FAIL — selfcheck rc=${rc} (see crates/wasm-simd-parity/src/lib.rs return codes)`); + process.exit(1); +} diff --git a/crates/wasm-simd-parity/src/lib.rs b/crates/wasm-simd-parity/src/lib.rs new file mode 100644 index 00000000..c333a6ec --- /dev/null +++ b/crates/wasm-simd-parity/src/lib.rs @@ -0,0 +1,107 @@ +//! Run-under-node parity gate for the wasm SIMD tier. +//! +//! `#[no_mangle] selfcheck()` runs each wasm SIMD lane's arithmetic against the +//! scalar `u32`/`f32`/`i8` reference **in the same module**, returning `0` iff +//! every lane is bit-identical. The companion `run.mjs` instantiates the +//! `.wasm` and asserts `selfcheck() == 0`; `scripts/wasm-parity.sh` wires the +//! build + node run, and the CI `wasm_simd` job runs the script. This is the +//! standing version of the one-off node check from commit 500d57e6 — extend the +//! per-lane blocks below whenever a new `ndarray::simd` lane lands. +//! +//! The whole crate is gated to `wasm32 + simd128` so it compiles to nothing on +//! any other target (it is workspace-excluded and only built by the CI job). +#![cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + +use ndarray::simd::{F32x16, I8x16, U32x16}; + +/// Distinct nonzero return codes make a CI failure point at the exact lane+op. +#[no_mangle] +pub extern "C" fn selfcheck() -> u32 { + if let Err(code) = check_u32x16() { + return code; + } + if let Err(code) = check_f32x16() { + return code; + } + if let Err(code) = check_i8x16() { + return code; + } + 0 +} + +/// `U32x16` ARX triple (Add / BitXor / rotate_left) — the ChaCha20/BLAKE lane. +fn check_u32x16() -> Result<(), u32> { + let a_arr: [u32; 16] = [ + 0x0000_0000, 0xFFFF_FFFF, 0x0000_0001, 0x8000_0000, 0x1234_5678, 0x9ABC_DEF0, 0xDEAD_BEEF, 0xCAFE_BABE, + 0x0F0F_0F0F, 0xF0F0_F0F0, 0x5555_5555, 0xAAAA_AAAA, 0x0000_00FF, 0xFF00_0000, 0x0101_0101, 0x8080_8080, + ]; + let b_arr: [u32; 16] = [ + 0x9E37_79B9, 0x1111_1111, 0xDEAD_C0DE, 0x0BAD_F00D, 0x7FFF_FFFF, 0x0000_0000, 0xFFFF_FFFF, 0x1357_9BDF, + 0x2468_ACE0, 0xFEDC_BA98, 0x0000_0010, 0x0000_001F, 0xABCD_EF01, 0x1020_4080, 0x0F0F_F0F0, 0xC0DE_CAFE, + ]; + // from_array / to_array roundtrip (lane ordering). + if U32x16::from_array(a_arr).to_array() != a_arr { + return Err(10); + } + let a = U32x16::from_array(a_arr); + let b = U32x16::from_array(b_arr); + let add = (a + b).to_array(); + let xor = (a ^ b).to_array(); + for i in 0..16 { + if add[i] != a_arr[i].wrapping_add(b_arr[i]) { + return Err(11); + } + if xor[i] != a_arr[i] ^ b_arr[i] { + return Err(12); + } + } + // ARX rotate — ChaCha20 uses 16/12/8/7; edges included. + for &n in &[0u32, 1, 7, 8, 12, 16, 24, 31] { + let r = a.rotate_left(n).to_array(); + for i in 0..16 { + if r[i] != a_arr[i].rotate_left(n) { + return Err(13); + } + } + } + Ok(()) +} + +/// `F32x16` — the float hot-path lane (splat / roundtrip / add / reduce_sum). +fn check_f32x16() -> Result<(), u32> { + let data: [f32; 16] = [ + 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, + ]; + if F32x16::from_array(data).to_array() != data { + return Err(20); + } + let a = F32x16::from_array(data); + let b = F32x16::splat(2.0); + let sum = (a + b).to_array(); + for i in 0..16 { + if sum[i] != data[i] + 2.0 { + return Err(21); + } + } + // 0+1+…+15 = 120. + if F32x16::from_array(data).reduce_sum() != 120.0 { + return Err(22); + } + Ok(()) +} + +/// `I8x16` — the byte lane (roundtrip / add). +fn check_i8x16() -> Result<(), u32> { + let a_arr: [i8; 16] = [-128, -1, 0, 1, 127, 2, -2, 3, -3, 42, -42, 100, -100, 7, -7, 64]; + let b_arr: [i8; 16] = [1, 1, 1, 1, 1, -1, -1, -1, 5, -5, 10, -10, 25, -25, 63, -64]; + if I8x16::from_array(a_arr).to_array() != a_arr { + return Err(30); + } + let sum = I8x16::from_array(a_arr).add(I8x16::from_array(b_arr)).to_array(); + for i in 0..16 { + if sum[i] != a_arr[i].wrapping_add(b_arr[i]) { + return Err(31); + } + } + Ok(()) +} diff --git a/scripts/wasm-parity.sh b/scripts/wasm-parity.sh new file mode 100755 index 00000000..56d18abb --- /dev/null +++ b/scripts/wasm-parity.sh @@ -0,0 +1,25 @@ +#!/usr/bin/env bash +# wasm SIMD parity gate — build the real ndarray::simd wasm types (via the +# `wasm-simd-parity` cdylib) with +simd128, then run the numeric selfcheck +# under node. Exercises BOTH: +# - correctness/integration: the full ndarray lib compiles for wasm32+simd128 +# (the harness path-deps ndarray, so building it builds ndarray for wasm); +# - parity: each wasm SIMD lane is bit-identical to scalar, run under node. +# The x86 `cargo test` suite never touches simd_wasm.rs, so this is the only +# automated guard for the wasm SIMD tier. +set -euo pipefail + +ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +MANIFEST="$ROOT/crates/wasm-simd-parity/Cargo.toml" +TARGET="wasm32-unknown-unknown" + +echo "==> building wasm-simd-parity (ndarray::simd wasm types) for $TARGET +simd128" +RUSTFLAGS="-C target-feature=+simd128" cargo build --release \ + --manifest-path "$MANIFEST" --target "$TARGET" + +WASM="$ROOT/crates/wasm-simd-parity/target/$TARGET/release/wasm_simd_parity.wasm" +# Excluded crate → its own target dir under the crate; fall back to workspace target. +[ -f "$WASM" ] || WASM="$ROOT/target/$TARGET/release/wasm_simd_parity.wasm" + +echo "==> running selfcheck under node: $WASM" +node "$ROOT/crates/wasm-simd-parity/run.mjs" "$WASM" diff --git a/src/lib.rs b/src/lib.rs index c19d6fc7..098b92cc 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -246,6 +246,13 @@ pub(crate) mod simd_avx512; #[allow(clippy::all, missing_docs, dead_code, unused_variables, unused_imports)] pub mod simd_avx2; +// Crypto ARX primitives (ChaCha20 block; Poly1305 / Blake2 to follow) — the +// hardware-acceleration layer the Ada `encryption` crate draws its hot-path +// keystream from. Scalar reference + RFC 8439 KAT land first (arch-agnostic, +// no cfg gate); the AVX-512 / wasm128 / NEON vectorized backends slot in behind +// the same signature, parity-checked against the KAT (W1a contract). +pub mod simd_crypto; + // Portable-SIMD backend — nightly-only. Wraps `core::simd::*` so miri can // execute the polyfill paths (intrinsic-based backends are opaque to // miri). Gated behind `nightly-simd` feature; the file itself requires diff --git a/src/simd.rs b/src/simd.rs index 4d871963..4b1552d7 100644 --- a/src/simd.rs +++ b/src/simd.rs @@ -396,13 +396,17 @@ pub use scalar::{ // The `wasm32_simd` module only exists under `target_feature = "simd128"`, // so this arm is gated identically. #[cfg(all(target_arch = "wasm32", target_feature = "simd128", not(feature = "nightly-simd")))] -pub use crate::simd_wasm::wasm32_simd::{f32x16, f64x8, i8x16, F32Mask16, F32x16, F64Mask8, F64x8, I8x16}; +pub use crate::simd_wasm::wasm32_simd::{ + f32x16, f64x8, i8x16, u32x16, F32Mask16, F32x16, F64Mask8, F64x8, I8x16, U32x16, +}; +// `u32x16`/`U32x16` now come from the native `wasm32_simd` arm above (the ARX +// lane the ChaCha20 backend rides), so they are dropped from this scalar list. #[cfg(all(target_arch = "wasm32", target_feature = "simd128", not(feature = "nightly-simd")))] pub use scalar::{ batch_packed_i4_16, f32x8, f64x4, i16x16, i16x32, i32x16, i32x8, i64x4, i64x8, i8x32, i8x64, palette_lookup_u8x8, - prefetch_read_t0, prefetch_read_t1, prefetch_read_t2, u16x16, u16x8, u32x16, u32x8, u64x4, u64x8, u8x64, u8x8, - F32x8, F64x4, I16x16, I16x32, I32x16, I32x8, I64x4, I64x8, I8x32, I8x64, U16x16, U16x32, U16x8, U32x16, U32x8, - U64x4, U64x8, U8x64, U8x8, + prefetch_read_t0, prefetch_read_t1, prefetch_read_t2, u16x16, u16x8, u32x8, u64x4, u64x8, u8x64, u8x8, F32x8, + F64x4, I16x16, I16x32, I32x16, I32x8, I64x4, I64x8, I8x32, I8x64, U16x16, U16x32, U16x8, U32x8, U64x4, U64x8, + U8x64, U8x8, }; // Other non-x86 targets — wasm32 without simd128, riscv, etc.: full scalar @@ -674,6 +678,13 @@ pub use crate::simd_ops::{ sub_f32_inplace, }; +// ChaCha20 ARX keystream — the crypto hot path the `encryption` AEAD draws from. +// `chacha20_block` is the scalar reference (RFC 8439 KAT); `chacha20_keystream` +// is the runtime-dispatched accelerated fill (AVX-512 16-block strides + scalar +// tail), byte-parity-pinned to the reference. Consumers pull from +// `ndarray::simd::*` per the W1a contract. +pub use crate::simd_crypto::{chacha20_block, chacha20_keystream, chacha20_state}; + // ============================================================================ // Tests // ============================================================================ @@ -695,6 +706,41 @@ mod tests { assert_eq!(v.to_array(), data); } + /// ARX triple parity: `U32x16`'s `Add` / `BitXor` / `rotate_left` must be + /// bit-identical to per-lane `u32` semantics — the ChaCha20/BLAKE primitive + /// set. Runs on whichever tier this build compiled (`super::*` re-exports the + /// dispatched `U32x16`), so it gates avx512 / avx2 / neon / wasm / scalar + /// alike. `rotate_left` is the newly-added op; add/xor are locked with it so + /// the whole quarter-round vocabulary is proven together. + #[test] + fn u32x16_arx_ops_match_scalar() { + let a_arr: [u32; 16] = [ + 0x0000_0000, 0xFFFF_FFFF, 0x0000_0001, 0x8000_0000, 0x1234_5678, 0x9ABC_DEF0, 0xDEAD_BEEF, 0xCAFE_BABE, + 0x0F0F_0F0F, 0xF0F0_F0F0, 0x5555_5555, 0xAAAA_AAAA, 0x0000_00FF, 0xFF00_0000, 0x0101_0101, 0x8080_8080, + ]; + let b_arr: [u32; 16] = [ + 0x9E37_79B9, 0x1111_1111, 0xDEAD_C0DE, 0x0BAD_F00D, 0x7FFF_FFFF, 0x0000_0000, 0xFFFF_FFFF, 0x1357_9BDF, + 0x2468_ACE0, 0xFEDC_BA98, 0x0000_0010, 0x0000_001F, 0xABCD_EF01, 0x1020_4080, 0x0F0F_F0F0, 0xC0DE_CAFE, + ]; + let a = U32x16::from_array(a_arr); + let b = U32x16::from_array(b_arr); + + let add = (a + b).to_array(); + let xor = (a ^ b).to_array(); + for i in 0..16 { + assert_eq!(add[i], a_arr[i].wrapping_add(b_arr[i]), "lane {i} add"); + assert_eq!(xor[i], a_arr[i] ^ b_arr[i], "lane {i} xor"); + } + + // The ARX rotate — ChaCha20 uses 16/12/8/7; edges included. + for n in [0u32, 1, 7, 8, 12, 16, 24, 31] { + let got = a.rotate_left(n).to_array(); + for i in 0..16 { + assert_eq!(got[i], a_arr[i].rotate_left(n), "lane {i} rotate_left({n})"); + } + } + } + #[test] fn f32x16_add_sub_mul_div() { let a = F32x16::splat(6.0); diff --git a/src/simd_avx2.rs b/src/simd_avx2.rs index 10942b7f..1b00235e 100644 --- a/src/simd_avx2.rs +++ b/src/simd_avx2.rs @@ -1542,6 +1542,23 @@ avx2_int_type!(U16x32, u16, 32, 0u16); avx2_int_type!(U32x16, u32, 16, 0u32); avx2_int_type!(U64x8, u64, 8, 0u64); +impl U32x16 { + /// Lane-wise left-rotate by `n` bits — the ARX rotate (matches + /// `u32::rotate_left`), completing `Add` + `BitXor` for ChaCha20/BLAKE. + /// This arm's `U32x16` is the `[u32; 16]` polyfill (native `2× __m256i` is + /// the deferred TD-SIMD-3 lowering), so the rotate is a per-lane + /// `u32::rotate_left` loop — bit-identical to the scalar tier and to the + /// native `VPROLVD` path, the shared parity reference. + #[inline(always)] + pub fn rotate_left(self, n: u32) -> Self { + let mut out = [0u32; 16]; + for i in 0..16 { + out[i] = self.0[i].rotate_left(n); + } + Self(out) + } +} + // 256-bit int lanes — scalar polyfills filling the gap surfaced by the // 2026-05-20 matrix audit. None of these had wrappers anywhere except // for `U32x8` / `U64x4` in `simd_nightly`. Adding `U16x16`, `U32x8`, diff --git a/src/simd_avx512.rs b/src/simd_avx512.rs index 9525f341..54021524 100644 --- a/src/simd_avx512.rs +++ b/src/simd_avx512.rs @@ -1432,6 +1432,19 @@ impl U32x16 { pub fn reduce_sum(self) -> u32 { unsafe { _mm512_reduce_add_epi32(self.0) as u32 } } + + /// Lane-wise left-rotate by `n` bits — the ARX rotate (matches + /// `u32::rotate_left`), the third ChaCha20/BLAKE-family primitive alongside + /// `Add` + `BitXor`. Single `VPROLVD` (AVX-512F variable rotate). The rotate + /// amount in ARX ciphers is a public constant, never secret, so a runtime + /// `n` carries no timing concern; LLVM folds a constant `n` to the immediate + /// `VPROLD` form. + #[inline(always)] + pub fn rotate_left(self, n: u32) -> Self { + // SAFETY: `_mm512_rolv_epi32` is AVX-512F; this whole module is gated to + // an AVX-512F build. Reads only `self`, returns an owned register. + Self(unsafe { _mm512_rolv_epi32(self.0, _mm512_set1_epi32(n as i32)) }) + } } impl_bin_op!(U32x16, Add, add, _mm512_add_epi32); diff --git a/src/simd_crypto.rs b/src/simd_crypto.rs new file mode 100644 index 00000000..32d7d022 --- /dev/null +++ b/src/simd_crypto.rs @@ -0,0 +1,484 @@ +//! Crypto ARX primitives for the `ndarray::simd` polyfill — the hardware +//! acceleration layer the Ada `encryption` crate (Argon2id / XChaCha20-Poly1305 +//! / Ed25519 / SHA-384) draws its hot-path keystream from. +//! +//! # Why these live in `ndarray::simd` +//! +//! The workspace invariant (`AdaWorldAPI/lance-graph` `simd-savant`): **all SIMD +//! comes from `ndarray::simd` via the polyfill — `simd.rs` + `simd_ops.rs` > +//! `simd_{arch}.rs`**. The crypto keystream is a SIMD hot path like any other, so +//! its accelerated form belongs here (server AVX-512 / browser wasm128), not +//! re-hidden inside a per-consumer copy. The consumer-facing crypto surface +//! (`ogar-encryption` → the `encryption` crate) composes AEAD/KDF framing on top; +//! this module supplies the vectorizable core. +//! +//! # Why a SIMD ChaCha20 is *safe* to hand-vectorize (unlike, say, AES) +//! +//! ChaCha20 is an **ARX** cipher: every operation is a 32-bit `wrapping_add`, +//! `xor`, or fixed-distance `rotate_left`. There are **no secret-dependent +//! branches and no secret-dependent memory indices** (no S-boxes, no T-tables), +//! so the block function is **constant-time by construction** — a property a +//! straight-line SIMD implementation preserves automatically. That is exactly +//! why RFC 8439 chose ChaCha20 for constant-time software, and why an +//! `ndarray::simd` vectorization does not introduce a timing side channel the +//! scalar path lacked. (AES in software, by contrast, is NOT safe to hand-roll +//! this way — its table lookups are secret-indexed; that is deliberately out of +//! scope here, and browsers already expose hardware AES via WebCrypto.) +//! +//! # Backend ladder (W1a consumer contract) +//! +//! | Backend | ChaCha20 block strategy | +//! |------------------|----------------------------------------------------------| +//! | scalar | the RFC 8439 reference double-round — the CORRECTNESS anchor | +//! | AVX-512 (server) | 16 blocks in parallel, word-sliced `__m512i` lanes (DONE) | +//! | wasm128 (browser)| 4 blocks in parallel, word-sliced `v128` u32x4 lanes (DONE) | +//! | NEON (edge) | `uint32x4_t` lanes (next) | +//! +//! The scalar reference lands FIRST with the RFC 8439 §2.3.2 Known-Answer-Test: +//! every future vectorized backend is a drop-in that MUST reproduce this exact +//! keystream (parity test), so correctness is pinned before performance. This is +//! the "reference + KAT, then vectorize with parity" discipline the W1a contract +//! mandates for any new `ndarray::simd` primitive. + +/// The ChaCha20 constants — the ASCII of `"expand 32-byte k"` as four +/// little-endian `u32` words (RFC 8439 §2.3). +const CHACHA20_CONSTANTS: [u32; 4] = [0x6170_7865, 0x3320_646e, 0x7962_2d32, 0x6b20_6574]; + +/// One ChaCha20 quarter-round on four working words (RFC 8439 §2.1). Pure ARX: +/// `wrapping_add` / `xor` / `rotate_left` — no branch, no memory index, so it is +/// constant-time regardless of the (secret) word values. +#[inline(always)] +fn quarter_round(state: &mut [u32; 16], a: usize, b: usize, c: usize, d: usize) { + state[a] = state[a].wrapping_add(state[b]); + state[d] = (state[d] ^ state[a]).rotate_left(16); + state[c] = state[c].wrapping_add(state[d]); + state[b] = (state[b] ^ state[c]).rotate_left(12); + state[a] = state[a].wrapping_add(state[b]); + state[d] = (state[d] ^ state[a]).rotate_left(8); + state[c] = state[c].wrapping_add(state[d]); + state[b] = (state[b] ^ state[c]).rotate_left(7); +} + +/// Compute one 64-byte ChaCha20 keystream block from a fully-populated 16-word +/// input `state` (RFC 8439 §2.3.1): 20 rounds (10 column + diagonal +/// double-rounds), then add the original input state, then serialize each word +/// little-endian. +/// +/// `state` is the caller-assembled block state — words `0..4` the constants, +/// `4..12` the 256-bit key, word `12` the block counter, `13..16` the 96-bit +/// nonce. Building that state (and the XChaCha `HChaCha20` nonce extension) is +/// the caller's job; this is the pure, vectorizable block core. +/// +/// **Constant-time:** straight-line ARX, no data-dependent control flow. +/// This scalar form is the reference every SIMD backend is parity-checked +/// against (see the module KAT). +#[must_use] +pub fn chacha20_block(state: &[u32; 16]) -> [u8; 64] { + let mut w = *state; + // 10 double-rounds = 20 rounds. + for _ in 0..10 { + // Column round. + quarter_round(&mut w, 0, 4, 8, 12); + quarter_round(&mut w, 1, 5, 9, 13); + quarter_round(&mut w, 2, 6, 10, 14); + quarter_round(&mut w, 3, 7, 11, 15); + // Diagonal round. + quarter_round(&mut w, 0, 5, 10, 15); + quarter_round(&mut w, 1, 6, 11, 12); + quarter_round(&mut w, 2, 7, 8, 13); + quarter_round(&mut w, 3, 4, 9, 14); + } + let mut out = [0u8; 64]; + for (i, word) in w.iter().enumerate() { + let sum = word.wrapping_add(state[i]); + out[i * 4..i * 4 + 4].copy_from_slice(&sum.to_le_bytes()); + } + out +} + +/// Assemble a ChaCha20 block state from a 256-bit `key`, a 32-bit block +/// `counter`, and a 96-bit `nonce` (RFC 8439 §2.3). Little-endian word packing. +/// Convenience for callers and for the KAT; the vectorized backends operate on +/// the assembled `[u32; 16]` from [`chacha20_block`]. +#[must_use] +pub fn chacha20_state(key: &[u8; 32], counter: u32, nonce: &[u8; 12]) -> [u32; 16] { + let mut s = [0u32; 16]; + s[0..4].copy_from_slice(&CHACHA20_CONSTANTS); + for i in 0..8 { + s[4 + i] = u32::from_le_bytes([key[i * 4], key[i * 4 + 1], key[i * 4 + 2], key[i * 4 + 3]]); + } + s[12] = counter; + for i in 0..3 { + s[13 + i] = u32::from_le_bytes([nonce[i * 4], nonce[i * 4 + 1], nonce[i * 4 + 2], nonce[i * 4 + 3]]); + } + s +} + +/// Fill `out` with successive ChaCha20 keystream blocks: block `b` uses block +/// counter `counter.wrapping_add(b)`, key and nonce fixed (RFC 8439 CTR mode). +/// +/// Accelerated entry point for the `encryption` AEAD keystream. Compile-time +/// tier dispatch (the `ndarray::simd` model — no `is_x86_feature_detected!`): +/// the AVX-512 16-block backend when built `+avx512f`, the wasm128 4-block +/// backend when built `+simd128`, else the scalar [`chacha20_block`] reference, +/// which also fills the ragged tail. Every backend is byte-parity-pinned to the +/// scalar reference (`chacha20_keystream_dispatch_parity_scalar`). +/// +/// Constant-time: straight-line ARX (add / xor / rotate), no secret-dependent +/// control flow or memory indexing. +pub fn chacha20_keystream(key: &[u8; 32], counter: u32, nonce: &[u8; 12], out: &mut [[u8; 64]]) { + let mut i = 0usize; + + #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] + { + while i + 16 <= out.len() { + let state = chacha20_state(key, counter.wrapping_add(i as u32), nonce); + let blocks = chacha20_block16_avx512(&state); + out[i..i + 16].copy_from_slice(&blocks); + i += 16; + } + } + + // wasm128 (browser) backend: 4-block strides. `simd128` is a compile-time + // feature, so this arm is present only when the module is built for it. + #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] + { + while i + 4 <= out.len() { + let state = chacha20_state(key, counter.wrapping_add(i as u32), nonce); + let blocks = chacha20_block4_wasm(&state); + out[i..i + 4].copy_from_slice(&blocks); + i += 4; + } + } + + // Scalar reference: the ragged tail, and the whole slice on scalar targets. + while i < out.len() { + let state = chacha20_state(key, counter.wrapping_add(i as u32), nonce); + out[i] = chacha20_block(&state); + i += 1; + } +} + +/// AVX-512 backend: 16 ChaCha20 blocks in parallel, word-sliced (lane `l` = block +/// `l`), so each round is pure SIMD add/xor/rotate and only the final LE serialize +/// transposes. cfg-gated to an `+avx512f` build → the intrinsics are statically +/// available (no `#[target_feature]`, no runtime detect). +#[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))] +fn chacha20_block16_avx512(state: &[u32; 16]) -> [[u8; 64]; 16] { + use core::arch::x86_64::*; + // SAFETY: avx512f is statically enabled (cfg gate), so every AVX-512F + // intrinsic below is available; all lane indices are fixed constants. + unsafe { + // One AVX-512 vertical quarter-round: same word indices as the scalar + // `quarter_round`, applied across all 16 blocks at once. Pure ARX. + macro_rules! qr512 { + ($v:expr, $a:expr, $b:expr, $c:expr, $d:expr) => {{ + $v[$a] = _mm512_add_epi32($v[$a], $v[$b]); + $v[$d] = _mm512_rol_epi32::<16>(_mm512_xor_si512($v[$d], $v[$a])); + $v[$c] = _mm512_add_epi32($v[$c], $v[$d]); + $v[$b] = _mm512_rol_epi32::<12>(_mm512_xor_si512($v[$b], $v[$c])); + $v[$a] = _mm512_add_epi32($v[$a], $v[$b]); + $v[$d] = _mm512_rol_epi32::<8>(_mm512_xor_si512($v[$d], $v[$a])); + $v[$c] = _mm512_add_epi32($v[$c], $v[$d]); + $v[$b] = _mm512_rol_epi32::<7>(_mm512_xor_si512($v[$b], $v[$c])); + }}; + } + + // Build the 16 initial working vectors. Every word is broadcast identically + // across the 16 lanes EXCEPT the counter (word 12), which gets lane-index + // 0..=15 added (u32 wrapping) — the 16 consecutive block counters. + let lane_index = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15); + let mut orig = [_mm512_setzero_si512(); 16]; + for (w, o) in orig.iter_mut().enumerate() { + *o = _mm512_set1_epi32(state[w] as i32); + } + orig[12] = _mm512_add_epi32(_mm512_set1_epi32(state[12] as i32), lane_index); + + let mut v = orig; + // 10 double-rounds = 20 rounds, identical schedule to the scalar reference. + for _ in 0..10 { + qr512!(v, 0, 4, 8, 12); + qr512!(v, 1, 5, 9, 13); + qr512!(v, 2, 6, 10, 14); + qr512!(v, 3, 7, 11, 15); + qr512!(v, 0, 5, 10, 15); + qr512!(v, 1, 6, 11, 12); + qr512!(v, 2, 7, 8, 13); + qr512!(v, 3, 4, 9, 14); + } + + // Add the original input state back (RFC 8439 §2.3.1), then de-vectorize: + // `words[w][l]` = final word `w` of block `l`. + let mut words = [[0u32; 16]; 16]; + for (w, dst) in words.iter_mut().enumerate() { + let summed = _mm512_add_epi32(v[w], orig[w]); + let mut tmp = [0i32; 16]; + _mm512_storeu_si512(tmp.as_mut_ptr().cast(), summed); + for (l, slot) in dst.iter_mut().enumerate() { + *slot = tmp[l] as u32; + } + } + + // Serialize each block little-endian: block `l`, word `w` → bytes 4w..4w+4. + let mut out = [[0u8; 64]; 16]; + for (l, block) in out.iter_mut().enumerate() { + for w in 0..16 { + block[w * 4..w * 4 + 4].copy_from_slice(&words[w][l].to_le_bytes()); + } + } + out + } +} + +/// wasm128 (browser) backend: compute 4 consecutive ChaCha20 keystream blocks in +/// parallel using the WebAssembly SIMD `v128` register as four `u32x4` lanes. +/// Same word-sliced ("vertical") layout as the AVX-512 backend — lane `l` of +/// working vector `w` holds word `w` of block `l` — so every round is a pure +/// `u32x4` add/xor/rotate with no lane shuffles. +/// +/// Requires the `simd128` target feature (a compile-time WASM feature, not +/// runtime-detected). Uses **no `unsafe`**: the wasm arithmetic and +/// lane-extraction intrinsics are safe under the enabled feature; output is read +/// with `u32x4_extract_lane` rather than a raw `v128_store`. +#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))] +fn chacha20_block4_wasm(state: &[u32; 16]) -> [[u8; 64]; 4] { + use core::arch::wasm32::*; + + // Fixed-distance left-rotate of each u32 lane: `(x << n) | (x >> (32-n))`. + macro_rules! rotl { + ($x:expr, $n:expr) => { + v128_or(u32x4_shl($x, $n), u32x4_shr($x, 32 - $n)) + }; + } + // One vertical quarter-round across all 4 blocks. Same word indices as the + // scalar `quarter_round`. Pure ARX. + macro_rules! qr128 { + ($v:expr, $a:expr, $b:expr, $c:expr, $d:expr) => {{ + $v[$a] = u32x4_add($v[$a], $v[$b]); + $v[$d] = rotl!(v128_xor($v[$d], $v[$a]), 16); + $v[$c] = u32x4_add($v[$c], $v[$d]); + $v[$b] = rotl!(v128_xor($v[$b], $v[$c]), 12); + $v[$a] = u32x4_add($v[$a], $v[$b]); + $v[$d] = rotl!(v128_xor($v[$d], $v[$a]), 8); + $v[$c] = u32x4_add($v[$c], $v[$d]); + $v[$b] = rotl!(v128_xor($v[$b], $v[$c]), 7); + }}; + } + + // Initial working vectors: every word broadcast across the 4 lanes EXCEPT the + // counter (word 12), which gets lane-index 0..=3 added (u32 wrapping). + let lane_index = u32x4(0, 1, 2, 3); + let mut orig = [u32x4_splat(0); 16]; + for (w, o) in orig.iter_mut().enumerate() { + *o = u32x4_splat(state[w]); + } + orig[12] = u32x4_add(u32x4_splat(state[12]), lane_index); + + let mut v = orig; + for _ in 0..10 { + qr128!(v, 0, 4, 8, 12); + qr128!(v, 1, 5, 9, 13); + qr128!(v, 2, 6, 10, 14); + qr128!(v, 3, 7, 11, 15); + qr128!(v, 0, 5, 10, 15); + qr128!(v, 1, 6, 11, 12); + qr128!(v, 2, 7, 8, 13); + qr128!(v, 3, 4, 9, 14); + } + + // Add the original input state back, then de-vectorize via lane extraction. + let mut words = [[0u32; 4]; 16]; + for (w, dst) in words.iter_mut().enumerate() { + let summed = u32x4_add(v[w], orig[w]); + dst[0] = u32x4_extract_lane::<0>(summed); + dst[1] = u32x4_extract_lane::<1>(summed); + dst[2] = u32x4_extract_lane::<2>(summed); + dst[3] = u32x4_extract_lane::<3>(summed); + } + + let mut out = [[0u8; 64]; 4]; + for (l, block) in out.iter_mut().enumerate() { + for w in 0..16 { + block[w * 4..w * 4 + 4].copy_from_slice(&words[w][l].to_le_bytes()); + } + } + out +} + +#[cfg(test)] +mod tests { + use super::*; + + /// The RFC 8439 §2.3.2 canonical keystream block (key = `00,01,…,1f`, block + /// counter = 1, nonce = `00 00 00 09 00 00 00 4a 00 00 00 00`). The scalar + /// reference and every SIMD backend MUST reproduce this exact 64-byte block. + const RFC8439_KAT_BLOCK: [u8; 64] = [ + 0x10, 0xf1, 0xe7, 0xe4, 0xd1, 0x3b, 0x59, 0x15, 0x50, 0x0f, 0xdd, 0x1f, 0xa3, 0x20, 0x71, 0xc4, 0xc7, 0xd1, + 0xf4, 0xc7, 0x33, 0xc0, 0x68, 0x03, 0x04, 0x22, 0xaa, 0x9a, 0xc3, 0xd4, 0x6c, 0x4e, 0xd2, 0x82, 0x64, 0x46, + 0x07, 0x9f, 0xaa, 0x09, 0x14, 0xc2, 0xd7, 0x05, 0xd9, 0x8b, 0x02, 0xa2, 0xb5, 0x12, 0x9c, 0xd1, 0xde, 0x16, + 0x4e, 0xb9, 0xcb, 0xd0, 0x83, 0xe8, 0xa2, 0x50, 0x3c, 0x4e, + ]; + + /// The RFC 8439 §2.3.2 key material (key `00..1f`, counter 1, the §2.3.2 nonce). + fn rfc8439_kat_inputs() -> ([u8; 32], u32, [u8; 12]) { + let mut key = [0u8; 32]; + for (i, b) in key.iter_mut().enumerate() { + *b = i as u8; + } + let nonce: [u8; 12] = [0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00]; + (key, 1, nonce) + } + + /// RFC 8439 §2.3.2 Known-Answer-Test on the scalar reference block. This pins + /// the scalar reference; every SIMD backend added later MUST reproduce this + /// exact 64-byte keystream (the parity gate). + #[test] + fn chacha20_block_rfc8439_kat() { + let (key, counter, nonce) = rfc8439_kat_inputs(); + let state = chacha20_state(&key, counter, &nonce); + let block = chacha20_block(&state); + assert_eq!(block, RFC8439_KAT_BLOCK, "ChaCha20 block must match RFC 8439 §2.3.2"); + } + + /// The quarter-round test vector (RFC 8439 §2.1.1) — the ARX core in + /// isolation, so a backend can be debugged at the round level. + #[test] + fn quarter_round_rfc8439_vector() { + // §2.1.1: inputs a,b,c,d and expected outputs, placed in a 16-word state + // at indices 0..4 so `quarter_round` operates on them. + let mut s = [0u32; 16]; + s[0] = 0x1111_1111; + s[1] = 0x0102_0304; + s[2] = 0x9b8d_6f43; + s[3] = 0x0123_4567; + quarter_round(&mut s, 0, 1, 2, 3); + assert_eq!( + [s[0], s[1], s[2], s[3]], + [0xea2a_92f4, 0xcb1c_f8ce, 0x4581_472e, 0x5881_c4bb], + "quarter-round must match RFC 8439 §2.1.1" + ); + } + + /// The mandatory SIMD-crypto correctness gate: every block produced by the + /// `chacha20_keystream` dispatcher (AVX-512 backend when `avx512f` is present, + /// scalar tail always) MUST be byte-identical to the scalar + /// [`chacha20_block`] reference for the same key/counter/nonce. `n = 40` + /// spans two full 16-block AVX-512 strides plus an 8-block scalar tail, so + /// both the vectorized body and the ragged-tail path are exercised. If this + /// test is red, no SIMD crypto backend may be committed. + #[test] + fn chacha20_keystream_dispatch_parity_scalar() { + let mut key = [0u8; 32]; + for (i, b) in key.iter_mut().enumerate() { + *b = (i as u8).wrapping_mul(7).wrapping_add(3); + } + let nonce: [u8; 12] = [0xde, 0xad, 0xbe, 0xef, 0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66, 0x77]; + let base = 0x0102_0304u32; + + // 40 blocks: two AVX-512 strides (16 each) + an 8-block scalar tail. + // Fixed-size array (not `vec!`) so the test compiles under `no_std`. + let mut ks = [[0u8; 64]; 40]; + chacha20_keystream(&key, base, &nonce, &mut ks); + + for (b, got) in ks.iter().enumerate() { + let state = chacha20_state(&key, base.wrapping_add(b as u32), &nonce); + let want = chacha20_block(&state); + assert_eq!(*got, want, "keystream block {b} must equal the scalar reference"); + } + } + + /// The dispatcher must reproduce the RFC 8439 §2.3.2 KAT at block 0. When the + /// host CPU has AVX-512F this drives the RFC vector THROUGH the AVX-512 + /// backend (16 blocks per stride, block 0 uses counter 1), directly pinning + /// the vectorized path to the published answer — not just to the scalar path. + #[test] + fn chacha20_keystream_dispatch_matches_rfc8439_kat() { + let (key, counter, nonce) = rfc8439_kat_inputs(); + // 16 blocks forces the AVX-512 stride when available; block 0 == the KAT. + let mut ks = [[0u8; 64]; 16]; + chacha20_keystream(&key, counter, &nonce, &mut ks); + assert_eq!(ks[0], RFC8439_KAT_BLOCK, "dispatcher block 0 must match RFC 8439 §2.3.2"); + } + + /// A native scalar mirror of `chacha20_block4_wasm`: it emulates each wasm + /// `v128` as a `[u32; 4]` lane array and applies the *identical* operation + /// sequence (word-sliced layout, counter-lane `0..=3`, the 8-quarter-round + /// double-round schedule, add-original, per-lane LE serialization). The wasm + /// intrinsics (`u32x4_add`/`v128_xor`/`u32x4_shl`/`u32x4_shr`/`u32x4_extract_lane`) + /// compute exactly these lane-wise u32 ops, so this validates the wasm arm's + /// *logic* on a target that can't run wasm; the wasm intrinsic API itself is + /// covered by the `wasm32-unknown-unknown` + `simd128` compile-check, and the + /// same word-sliced algorithm is byte-parity-proven 16-wide by the AVX-512 + /// tests above. + fn chacha20_block4_wasm_logic_ref(state: &[u32; 16]) -> [[u8; 64]; 4] { + // Lane-wise emulations of the wasm ops. + let add = |a: [u32; 4], b: [u32; 4]| [0, 1, 2, 3].map(|l| a[l].wrapping_add(b[l])); + let xor = |a: [u32; 4], b: [u32; 4]| [0, 1, 2, 3].map(|l| a[l] ^ b[l]); + let rotl = |a: [u32; 4], n: u32| [0, 1, 2, 3].map(|l| a[l].rotate_left(n)); + let splat = |x: u32| [x; 4]; + + macro_rules! qr { + ($v:expr, $a:expr, $b:expr, $c:expr, $d:expr) => {{ + $v[$a] = add($v[$a], $v[$b]); + $v[$d] = rotl(xor($v[$d], $v[$a]), 16); + $v[$c] = add($v[$c], $v[$d]); + $v[$b] = rotl(xor($v[$b], $v[$c]), 12); + $v[$a] = add($v[$a], $v[$b]); + $v[$d] = rotl(xor($v[$d], $v[$a]), 8); + $v[$c] = add($v[$c], $v[$d]); + $v[$b] = rotl(xor($v[$b], $v[$c]), 7); + }}; + } + + let lane_index = [0u32, 1, 2, 3]; + let mut orig = [[0u32; 4]; 16]; + for (w, o) in orig.iter_mut().enumerate() { + *o = splat(state[w]); + } + orig[12] = add(splat(state[12]), lane_index); + + let mut v = orig; + for _ in 0..10 { + qr!(v, 0, 4, 8, 12); + qr!(v, 1, 5, 9, 13); + qr!(v, 2, 6, 10, 14); + qr!(v, 3, 7, 11, 15); + qr!(v, 0, 5, 10, 15); + qr!(v, 1, 6, 11, 12); + qr!(v, 2, 7, 8, 13); + qr!(v, 3, 4, 9, 14); + } + + let mut words = [[0u32; 4]; 16]; + for (w, dst) in words.iter_mut().enumerate() { + *dst = add(v[w], orig[w]); + } + let mut out = [[0u8; 64]; 4]; + for (l, block) in out.iter_mut().enumerate() { + for w in 0..16 { + block[w * 4..w * 4 + 4].copy_from_slice(&words[w][l].to_le_bytes()); + } + } + out + } + + /// The wasm arm's 4-block word-sliced logic must equal the scalar reference + /// for the 4 consecutive counters it covers — the native proxy for the wasm + /// backend's correctness (see `chacha20_block4_wasm_logic_ref`). + #[test] + fn chacha20_wasm_logic_matches_scalar() { + let mut key = [0u8; 32]; + for (i, b) in key.iter_mut().enumerate() { + *b = (i as u8).wrapping_mul(11).wrapping_add(5); + } + let nonce: [u8; 12] = [0x01, 0x23, 0x45, 0x67, 0x89, 0xab, 0xcd, 0xef, 0xfe, 0xdc, 0xba, 0x98]; + let base = 0x1000_0000u32; + + let state = chacha20_state(&key, base, &nonce); + let four = chacha20_block4_wasm_logic_ref(&state); + for (l, got) in four.iter().enumerate() { + let want = chacha20_block(&chacha20_state(&key, base + l as u32, &nonce)); + assert_eq!(*got, want, "wasm-logic block {l} must equal the scalar reference"); + } + } +} diff --git a/src/simd_nightly/u_word_types.rs b/src/simd_nightly/u_word_types.rs index bd34c01b..b8add767 100644 --- a/src/simd_nightly/u_word_types.rs +++ b/src/simd_nightly/u_word_types.rs @@ -364,6 +364,20 @@ impl U32x16 { pub fn cmpgt_mask(self, other: Self) -> u16 { self.0.simd_gt(other.0).to_bitmask() as u16 } + + /// Lane-wise left-rotate by `n` bits — the ARX rotate (matches + /// `u32::rotate_left`), completing `Add` + `BitXor` for ChaCha20/BLAKE. + /// `core::simd` has no bit-rotate, so this is the shift-or composition with + /// the `n % 32 == 0` guard (a `>> 32` on a `u32` lane is UB), matching + /// `u32::rotate_left`'s wrap-by-32 semantics. + #[inline(always)] + pub fn rotate_left(self, n: u32) -> Self { + let n = n % 32; + if n == 0 { + return self; + } + Self((self.0 << u32x16::splat(n)) | (self.0 >> u32x16::splat(32 - n))) + } } impl Default for U32x16 { diff --git a/src/simd_scalar.rs b/src/simd_scalar.rs index fbf4c33e..952825ce 100644 --- a/src/simd_scalar.rs +++ b/src/simd_scalar.rs @@ -1028,6 +1028,22 @@ impl Shl for U32x16 { } } +impl U32x16 { + /// Lane-wise left-rotate by `n` bits — the ARX rotate (matches + /// `u32::rotate_left`). The completeness tier: a full, correct + /// implementation held to the same bar as the SIMD tiers, and the bit-exact + /// reference they are parity-checked against. Delegates to `u32::rotate_left` + /// per lane (`rotate_left(0)` and `rotate_left(32)` both no-op, matching std). + #[inline(always)] + pub fn rotate_left(self, n: u32) -> Self { + let mut out = [0u32; 16]; + for i in 0..16 { + out[i] = self.0[i].rotate_left(n); + } + Self(out) + } +} + // Shift operators for U64x8 impl Shr for U64x8 { type Output = Self; diff --git a/src/simd_wasm.rs b/src/simd_wasm.rs index 273f0a98..5680ed90 100644 --- a/src/simd_wasm.rs +++ b/src/simd_wasm.rs @@ -66,13 +66,16 @@ pub mod wasm32_simd { use core::arch::wasm32::*; use core::fmt; - use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}; + use core::ops::{Add, AddAssign, BitXor, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}; // Long-tail integer types come from the scalar fallback — they are not on // the perf-critical f32/byte path this module accelerates, and reusing the // scalar types keeps the `to_bits` / `from_bits` / `cast_i32` return types // identical to every other backend (same choice `simd_neon` makes). - pub use crate::simd::scalar::{I32x16, U32x16, U64x8}; + // `U32x16` is the exception: it carries the ARX vocabulary (Add/BitXor/ + // rotate_left) the ChaCha20 lane needs, so it is native here (`[U32x4; 4]`, + // NEON-style) rather than the scalar fallback — see below. + pub use crate::simd::scalar::{I32x16, U64x8}; // ════════════════════════════════════════════════════════════════════ // F32x16 — 16 × f32 backed by 4 × v128 (f32x4 interpretation) @@ -274,13 +277,14 @@ pub mod wasm32_simd { for i in 0..16 { o[i] = a[i].to_bits(); } - U32x16(o) + U32x16::from_array(o) } #[inline(always)] pub fn from_bits(bits: U32x16) -> Self { + let b = bits.to_array(); let mut o = [0.0f32; 16]; for i in 0..16 { - o[i] = f32::from_bits(bits.0[i]); + o[i] = f32::from_bits(b[i]); } Self::from_array(o) } @@ -853,6 +857,148 @@ pub mod wasm32_simd { } } + // ════════════════════════════════════════════════════════════════════ + // U32x4 / U32x16 — u32 ARX lanes (ChaCha20 / BLAKE), NEON-style. + // + // `U32x4` is the dispatched native unit (one `v128`, u32x4 interpretation); + // `U32x16` is the 16-wide polyfill `[U32x4; 4]`, fanning each op over the 4 + // sub-lanes — the same composition `simd_neon` uses (`[U32x4; 4]`, where + // NEON's U32x4 wraps `uint32x4_t`). `U32x16`'s consumer API (`Add` / + // `BitXor` / `rotate_left`) mirrors `simd_avx512::U32x16` exactly, so one + // source (the ChaCha20 backend) compiles unchanged on every tier. The u32 + // ARX ops are all safe wasm intrinsics; only the v128 load/store are unsafe. + // ════════════════════════════════════════════════════════════════════ + + /// 4×u32 in one WASM `v128` (u32x4 interpretation) — the native ARX unit. + #[derive(Copy, Clone)] + #[repr(transparent)] + pub struct U32x4(pub v128); + + impl U32x4 { + pub const LANES: usize = 4; + + #[inline(always)] + pub fn splat(v: u32) -> Self { + Self(u32x4_splat(v)) + } + + #[inline(always)] + pub fn from_array(a: [u32; 4]) -> Self { + // SAFETY: a `[u32; 4]` is exactly 16 bytes; one unaligned v128 load. + Self(unsafe { v128_load(a.as_ptr() as *const v128) }) + } + + #[inline(always)] + pub fn to_array(self) -> [u32; 4] { + let mut a = [0u32; 4]; + // SAFETY: store writes exactly 16 bytes into the 16-byte array. + unsafe { v128_store(a.as_mut_ptr() as *mut v128, self.0) }; + a + } + + #[inline(always)] + pub fn add(self, o: Self) -> Self { + Self(u32x4_add(self.0, o.0)) + } + + #[inline(always)] + pub fn bitxor(self, o: Self) -> Self { + Self(v128_xor(self.0, o.0)) + } + + /// Lane-wise left-rotate by `n` bits — the ARX rotate (matches + /// `u32::rotate_left`). wasm has no rotate op, so this is the shift-or + /// with the `n % 32 == 0` guard (`u32x4_shr` by 32 would over-shift); + /// `u32x4_shl`/`u32x4_shr` already mask the count to `& 31`, so + /// non-zero `n` is exact. Rotate amount is a public ARX constant. + #[inline(always)] + pub fn rotate_left(self, n: u32) -> Self { + let n = n % 32; + if n == 0 { + return self; + } + Self(v128_or(u32x4_shl(self.0, n), u32x4_shr(self.0, 32 - n))) + } + } + + /// 16×u32 as `[U32x4; 4]` — the NEON-style 16-wide polyfill. Consumer API + /// (`Add` / `BitXor` / `rotate_left`) matches `simd_avx512::U32x16`. + #[derive(Copy, Clone)] + #[repr(align(64))] + pub struct U32x16(pub [U32x4; 4]); + + impl U32x16 { + pub const LANES: usize = 16; + + #[inline(always)] + pub fn splat(v: u32) -> Self { + Self([U32x4::splat(v); 4]) + } + + #[inline(always)] + pub fn from_array(a: [u32; 16]) -> Self { + Self([ + U32x4::from_array([a[0], a[1], a[2], a[3]]), + U32x4::from_array([a[4], a[5], a[6], a[7]]), + U32x4::from_array([a[8], a[9], a[10], a[11]]), + U32x4::from_array([a[12], a[13], a[14], a[15]]), + ]) + } + + #[inline(always)] + pub fn to_array(self) -> [u32; 16] { + let mut o = [0u32; 16]; + for i in 0..4 { + o[i * 4..i * 4 + 4].copy_from_slice(&self.0[i].to_array()); + } + o + } + + /// Lane-wise left-rotate by `n` bits (ARX rotate), fanned over 4 lanes. + #[inline(always)] + pub fn rotate_left(self, n: u32) -> Self { + Self([ + self.0[0].rotate_left(n), + self.0[1].rotate_left(n), + self.0[2].rotate_left(n), + self.0[3].rotate_left(n), + ]) + } + } + + impl Add for U32x16 { + type Output = Self; + #[inline(always)] + fn add(self, r: Self) -> Self { + Self([self.0[0].add(r.0[0]), self.0[1].add(r.0[1]), self.0[2].add(r.0[2]), self.0[3].add(r.0[3])]) + } + } + + impl BitXor for U32x16 { + type Output = Self; + #[inline(always)] + fn bitxor(self, r: Self) -> Self { + Self([ + self.0[0].bitxor(r.0[0]), + self.0[1].bitxor(r.0[1]), + self.0[2].bitxor(r.0[2]), + self.0[3].bitxor(r.0[3]), + ]) + } + } + + impl fmt::Debug for U32x16 { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "U32x16({:?})", self.to_array()) + } + } + + impl PartialEq for U32x16 { + fn eq(&self, other: &Self) -> bool { + self.to_array() == other.to_array() + } + } + // ════════════════════════════════════════════════════════════════════ // Lowercase aliases (consumer-API parity with the other backends) // ════════════════════════════════════════════════════════════════════ @@ -863,6 +1009,8 @@ pub mod wasm32_simd { pub type f64x8 = F64x8; #[allow(non_camel_case_types)] pub type i8x16 = I8x16; + #[allow(non_camel_case_types)] + pub type u32x16 = U32x16; // ════════════════════════════════════════════════════════════════════ // Free hot-kernel functions — v128 counterparts to the NEON kernels in diff --git a/src/tri.rs b/src/tri.rs index a98332fc..117b1693 100644 --- a/src/tri.rs +++ b/src/tri.rs @@ -342,7 +342,10 @@ mod tests { #[test] fn test_odd_k() { let x = array![[1, 2, 3], [4, 5, 6], [7, 8, 9]]; - let z = Array2::zeros([3, 3]); + // Explicit element type: `serde_json` (dev-dep via criterion) now ships + // `impl PartialEq for i32`, so bare `Array2::zeros` leaves the + // `assert_eq!` element type ambiguous. + let z = Array2::::zeros([3, 3]); assert_eq!(x.triu(isize::MIN), x); assert_eq!(x.tril(isize::MIN), z); assert_eq!(x.triu(isize::MAX), z); diff --git a/tests/chacha20_rustcrypto_parity.rs b/tests/chacha20_rustcrypto_parity.rs new file mode 100644 index 00000000..f3aa3b8e --- /dev/null +++ b/tests/chacha20_rustcrypto_parity.rs @@ -0,0 +1,138 @@ +//! Cross-implementation trust gate for `ndarray::simd::chacha20_keystream`. +//! +//! The module KAT (`src/simd_crypto.rs`) pins the scalar reference to the single +//! RFC 8439 §2.3.2 vector, and `chacha20_keystream_dispatch_parity_scalar` pins +//! the AVX-512 backend to that scalar reference. This integration test closes the +//! remaining gap: it proves the *whole dispatcher* is byte-for-byte identical to +//! an **independent, vetted implementation** — RustCrypto's `chacha20::ChaCha20` +//! (the IETF 96-bit-nonce / 32-bit-counter variant, RFC 8439) — across a wide +//! space of random keys, nonces, counters, and lengths. This is the +//! "reference + KAT, then vectorize with parity" discipline extended to a second +//! oracle: no consumer (the `encryption` AEAD, a future XChaCha vault) should +//! draw keystream from the accelerated primitive until it is proven equivalent +//! to the trusted implementation everywhere, not just at one KAT point. +//! +//! ChaCha20 is ARX, so the accelerated path is constant-time by construction; +//! this test guards *correctness*, which is the property a hand-vectorization +//! can actually get wrong. +//! +//! Gated to `std`: `ndarray::simd` is `#[cfg(feature = "std")]`, so this test is +//! a no-op in `--no-default-features` builds. +#![cfg(feature = "std")] + +use chacha20::cipher::{KeyIvInit, StreamCipher, StreamCipherSeek}; +use chacha20::ChaCha20; +use ndarray::simd::{chacha20_block, chacha20_keystream, chacha20_state}; + +/// Deterministic SplitMix64 — a stand-in for a PRNG so the vectors are fixed and +/// reproducible without a `rand` dev-dependency (and without `Math.random`-style +/// nondeterminism that would make a failure impossible to bisect). +struct SplitMix64(u64); + +impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + + fn fill(&mut self, buf: &mut [u8]) { + for chunk in buf.chunks_mut(8) { + let bytes = self.next_u64().to_le_bytes(); + chunk.copy_from_slice(&bytes[..chunk.len()]); + } + } +} + +/// The trusted oracle: `n_blocks` of RustCrypto ChaCha20 keystream starting at +/// block `counter` (produced by encrypting a zero buffer). +fn rustcrypto_keystream(key: &[u8; 32], counter: u32, nonce: &[u8; 12], n_blocks: usize) -> Vec { + let mut cipher = ChaCha20::new_from_slices(key, nonce).expect("valid key/nonce lengths"); + // Seek to block `counter` (position is in bytes; each block is 64 bytes). + cipher.seek(u64::from(counter) * 64); + let mut buf = vec![0u8; n_blocks * 64]; + cipher.apply_keystream(&mut buf); + buf +} + +/// The primitive under test, flattened to a byte stream for comparison. +fn ndarray_keystream(key: &[u8; 32], counter: u32, nonce: &[u8; 12], n_blocks: usize) -> Vec { + let mut blocks = vec![[0u8; 64]; n_blocks]; + chacha20_keystream(key, counter, nonce, &mut blocks); + blocks.concat() +} + +/// The dispatcher (AVX-512 backend when present, scalar otherwise) must equal the +/// RustCrypto oracle across a broad, deterministic sample of the parameter space. +#[test] +fn keystream_matches_rustcrypto_across_vectors() { + let mut rng = SplitMix64(0x0DDB_1A5E_5EED_1234); + + // Lengths chosen to exercise: sub-stride, exactly one AVX-512 stride (16), + // multi-stride, and every ragged-tail size in between. + let block_counts = [1usize, 2, 7, 15, 16, 17, 31, 33, 40, 64]; + + for &n in &block_counts { + // A handful of independent (key, nonce, counter) draws per length. + for _ in 0..8 { + let mut key = [0u8; 32]; + let mut nonce = [0u8; 12]; + rng.fill(&mut key); + rng.fill(&mut nonce); + // Keep counter + n well inside u32 so the RustCrypto oracle (which + // refuses to wrap its 32-bit counter) stays in its valid range; the + // wrapping edge is covered separately by the scalar-parity unit test. + let counter = (rng.next_u64() as u32) & 0x00FF_FFFF; + + let got = ndarray_keystream(&key, counter, &nonce, n); + let want = rustcrypto_keystream(&key, counter, &nonce, n); + assert_eq!( + got, want, + "keystream mismatch vs RustCrypto: n={n} blocks, counter={counter}, key[0]={}, nonce[0]={}", + key[0], nonce[0] + ); + } + } +} + +/// The single-block scalar primitive `chacha20_block` (fed via `chacha20_state`) +/// must equal RustCrypto's first block for the same inputs — pins the low-level +/// entry point, not just the batched dispatcher. +#[test] +fn single_block_matches_rustcrypto() { + let mut rng = SplitMix64(0xC0FF_EE00_1357_9BDF); + for _ in 0..64 { + let mut key = [0u8; 32]; + let mut nonce = [0u8; 12]; + rng.fill(&mut key); + rng.fill(&mut nonce); + let counter = (rng.next_u64() as u32) & 0x00FF_FFFF; + + let got = chacha20_block(&chacha20_state(&key, counter, &nonce)); + let want = rustcrypto_keystream(&key, counter, &nonce, 1); + assert_eq!(&got[..], &want[..], "single-block mismatch vs RustCrypto at counter={counter}"); + } +} + +/// Sanity-check the oracle itself is the IETF RFC 8439 variant (not the 64-bit +/// legacy nonce cipher) by reproducing the §2.3.2 KAT through RustCrypto. If this +/// fails, the oracle is misconfigured and the parity tests above prove nothing. +#[test] +fn rustcrypto_oracle_is_rfc8439_ietf() { + let mut key = [0u8; 32]; + for (i, b) in key.iter_mut().enumerate() { + *b = i as u8; + } + let nonce: [u8; 12] = [0x00, 0x00, 0x00, 0x09, 0x00, 0x00, 0x00, 0x4a, 0x00, 0x00, 0x00, 0x00]; + let block = rustcrypto_keystream(&key, 1, &nonce, 1); + + let expected: [u8; 64] = [ + 0x10, 0xf1, 0xe7, 0xe4, 0xd1, 0x3b, 0x59, 0x15, 0x50, 0x0f, 0xdd, 0x1f, 0xa3, 0x20, 0x71, 0xc4, 0xc7, 0xd1, + 0xf4, 0xc7, 0x33, 0xc0, 0x68, 0x03, 0x04, 0x22, 0xaa, 0x9a, 0xc3, 0xd4, 0x6c, 0x4e, 0xd2, 0x82, 0x64, 0x46, + 0x07, 0x9f, 0xaa, 0x09, 0x14, 0xc2, 0xd7, 0x05, 0xd9, 0x8b, 0x02, 0xa2, 0xb5, 0x12, 0x9c, 0xd1, 0xde, 0x16, + 0x4e, 0xb9, 0xcb, 0xd0, 0x83, 0xe8, 0xa2, 0x50, 0x3c, 0x4e, + ]; + assert_eq!(&block[..], &expected[..], "RustCrypto oracle must reproduce RFC 8439 §2.3.2"); +}