From 398a2ae07c3a9177de07e75b43ab4b6f030b018c Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 17:33:58 +0000 Subject: [PATCH 01/34] feat: cut the fixed overhead of single-tensor output Single-tensor Compressor.compress output carried about 2.7 KB of fixed bytes, so small tensors came out larger than their input (issue #5). - The container writer interns only the plane codecs that the planes reference, not all 16 built-in codecs. A single-tensor container drops from about 2.5 KB to 0.5-1 KB of non-payload bytes. Readers resolve codec indices by position, so existing readers decode the new files. The explode/implode transcoder rebuilds whichever table layout matches the recorded hash, using the archive-wide codec list so a one-tensor implode keeps its codec indices. - New compact single-tensor frame (magic \x89PTWS): varint header, optional shape/hash/CRC/delta fields, the chain, and per-plane codec records. It stores the raw bytes when no codec helps. Fixed cost is 10-130 B. Compressor uses it by default (Framing.COMPACT) and falls back to the container for streaming or shared-state planes. Framing.CONTAINER keeps the old layout; the CLI file writers and the ablation harness pin it. The decoder bounds the chain's declared size by orig_size before the descriptor pass, so a corrupt frame cannot request gigabytes. - CompressionConfig exposes emit_payload_hash and emit_plane_crc. - benchmarks/container_overhead.py reports the size sweep; a CI test enforces the issue's acceptance limits. New compact_decode fuzz target. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- .github/workflows/fuzz.yml | 6 +- benchmarks/README.md | 1 + benchmarks/container_overhead.py | 157 ++++ crates/ptwm-core/src/compact.rs | 749 ++++++++++++++++++ crates/ptwm-core/src/compressor.rs | 118 ++- crates/ptwm-core/src/container.rs | 145 +++- crates/ptwm-core/src/lib.rs | 1 + crates/ptwm-core/src/transcode/mod.rs | 34 +- crates/ptwm-core/src/transcode/tests.rs | 72 ++ crates/ptwm-py/src/lib.rs | 120 +++ docs/architecture/wire-format.md | 55 ++ docs/site/content/30.guides/02.python-api.md | 3 + fuzz/Cargo.lock | 2 +- fuzz/Cargo.toml | 6 + fuzz/fuzz_targets/compact_decode.rs | 42 + python/ptwm/__init__.py | 4 +- python/ptwm/_config.py | 30 + python/ptwm/_rust/__init__.py | 63 ++ python/ptwm/ablation/_bench.py | 6 +- python/ptwm/ablation/_replay.py | 8 +- python/ptwm/cli/compress.py | 6 + python/ptwm/core/_compressor.py | 33 +- python/ptwm/core/_decompressor.py | 39 +- .../benchmarks/test_container_overhead.py | 72 ++ tests/ptwm/core/test_compact_frame.py | 165 ++++ tests/ptwm/core/test_zstd.py | 5 +- .../ptwm/wire_format/test_extension_table.py | 21 +- tests/ptwm/wire_format/test_v1_header.py | 6 +- 28 files changed, 1869 insertions(+), 100 deletions(-) create mode 100644 benchmarks/container_overhead.py create mode 100644 crates/ptwm-core/src/compact.rs create mode 100644 fuzz/fuzz_targets/compact_decode.rs create mode 100644 tests/ptwm/benchmarks/test_container_overhead.py create mode 100644 tests/ptwm/core/test_compact_frame.py diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 28fc5fe..67aee02 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -36,12 +36,12 @@ jobs: fail-fast: false matrix: # PRs run a 3-target smoke set (both entropy decoders + the wire - # parser); the full 11-target sweep runs on the weekly cron and on - # manual dispatch. Per-job whole-minute round-up × 11 targets on + # parser); the full 12-target sweep runs on the weekly cron and on + # manual dispatch. Per-job whole-minute round-up × 12 targets on # every core PR dominated this repo's Actions spend — the smoke set # keeps fast regression signal without paying for the full matrix on # every push. Edit both lists together when adding/removing targets. - target: ${{ github.event_name == 'pull_request' && fromJSON('["huffman_decode","rans_decode","parse_wire"]') || fromJSON('["huffman_decode","rans_decode","parse_header","parse_wire","container_open","order1_scale_ac_decode","order1_scale_ac_state_parse","extension_table","ptwx_wasm_extract","policy_resolve","raw_binary_schema"]') }} + target: ${{ github.event_name == 'pull_request' && fromJSON('["huffman_decode","rans_decode","parse_wire"]') || fromJSON('["huffman_decode","rans_decode","parse_header","parse_wire","container_open","order1_scale_ac_decode","order1_scale_ac_state_parse","extension_table","ptwx_wasm_extract","policy_resolve","raw_binary_schema","compact_decode"]') }} steps: - name: Checkout repository diff --git a/benchmarks/README.md b/benchmarks/README.md index 64d828c..d69c99d 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -49,6 +49,7 @@ BF16 tensor ≥ 50 MB; the full-model scripts use BF16 tensors up to a 3 GB cap. | Module | Needs a shard? | What it measures | | --- | --- | --- | +| `container_overhead` | no (synthetic) | Fixed per-output bytes of single-tensor `Compressor.compress` output, compact frame vs full container, bf16/fp32, constant/normal, 64 B to 1 MB, against zstd-19. `tests/ptwm/benchmarks/test_container_overhead.py` enforces the size limits in CI. | | `validate_dtype` | no (synthetic) | Losslessness + ratio across BF16/FP16/FP32/FP8 through the production streaming path. | | `validate_scale {small,large,edge}` | `large` only | `small`: many tiny tensors; `large`: a 72B-class shard (scaling + RSS); `edge`: single-tensor files, non-float dtypes, degenerate shapes. Every mode verifies a bit-exact roundtrip. | | `validate_realmodel` | yes (real shards) | Real-model validation matrix: ratio / encode / decode / losslessness across Mistral-7B (BF16), gpt-oss-20b (MXFP4), Nemotron (FP8), Qwen2.5-72B (BF16). | diff --git a/benchmarks/container_overhead.py b/benchmarks/container_overhead.py new file mode 100644 index 0000000..85b1158 --- /dev/null +++ b/benchmarks/container_overhead.py @@ -0,0 +1,157 @@ +# ruff: noqa: T201 +"""Fixed per-output overhead of single-tensor ``Compressor.compress`` output. + +Sweeps bf16 and fp32 tensors from 64 B to 1 MB, constant and random-normal, +and reports for each framing (compact frame and full container): output +size, non-payload bytes, ratio (raw / output), and zstd-19 on the same raw +bytes. Every output is decoded and checked bit-exact, including shape and +dtype. Needs no shards. + + python -m benchmarks.container_overhead + +"Non-payload" is the output size minus the codec payload. For the compact +frame the payload includes inline codec state; for the container it is the +sum of the plane payload lengths. A stored compact frame counts the raw bytes +as payload. + +``tests/ptwm/benchmarks/test_container_overhead.py`` runs the same sweep in +CI and asserts the size limits, so a regression in the fixed overhead fails +the build. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch +import zstandard +from ptwm import _rust +from ptwm._config import CompressionConfig, Format, Framing +from ptwm.core import Compressor, Decompressor + +SIZES = (64, 256, 1024, 4096, 16384, 65536, 1 << 20) +DTYPES = (("bf16", torch.bfloat16), ("fp32", torch.float32)) +KINDS = ("constant", "normal") + + +@dataclass(frozen=True) +class Row: + framing: str + dtype: str + kind: str + raw: int + out: int + non_payload: int + zstd19: int + + @property + def ratio(self) -> float: + return self.raw / self.out + + @property + def zstd_ratio(self) -> float: + return self.raw / self.zstd19 + + @property + def payload(self) -> int: + return self.out - self.non_payload + + +def make_tensor( + dtype: torch.dtype, kind: str, nbytes: int, seed: int = 1 +) -> torch.Tensor: + """A 1-D tensor of ``nbytes`` raw bytes: constant 0.5, or seeded normal.""" + n = nbytes // torch.empty(0, dtype=dtype).element_size() + if kind == "constant": + return torch.full((n,), 0.5, dtype=dtype) + g = torch.Generator().manual_seed(seed) + return torch.randn(n, generator=g).to(dtype) + + +def tiled_rows(rows: int, seed: int = 1) -> torch.Tensor: + """``rows`` copies of one random 1,024-element bf16 vector.""" + g = torch.Generator().manual_seed(seed) + row = torch.randn(1024, generator=g).to(torch.bfloat16) + return row.unsqueeze(0).repeat(rows, 1).contiguous() + + +def raw_bytes(t: torch.Tensor) -> bytes: + return t.contiguous().view(torch.uint8).numpy().tobytes() + + +def non_payload_bytes(blob: bytes) -> int: + """Output size minus codec payload (see the module docstring).""" + if blob[:5] == b"\x89PTWS": + return len(blob) - _rust.decode_compact(blob)[6] + payload = sum(s[5] for s in _rust.list_plane_summaries(blob)) + return len(blob) - payload + + +def compress_checked(t: torch.Tensor, framing: Framing) -> bytes: + """Compress ``t`` and assert a bit-exact round trip with shape and dtype.""" + blob = Compressor( + CompressionConfig(input_format=Format.TORCH, framing=framing) + ).compress(t) + back = Decompressor().decompress(blob) + assert isinstance(back, torch.Tensor) + assert back.dtype == t.dtype, (back.dtype, t.dtype) + assert back.shape == t.shape, (back.shape, t.shape) + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + return blob + + +def measure(t: torch.Tensor, framing: Framing, dtype: str, kind: str) -> Row: + raw = raw_bytes(t) + blob = compress_checked(t, framing) + z = zstandard.ZstdCompressor(level=19).compress(raw) + return Row( + framing=framing.name.lower(), + dtype=dtype, + kind=kind, + raw=len(raw), + out=len(blob), + non_payload=non_payload_bytes(blob), + zstd19=len(z), + ) + + +def sweep( + framings: tuple[Framing, ...] = (Framing.COMPACT, Framing.CONTAINER), +) -> list[Row]: + rows = [] + for framing in framings: + for dname, dt in DTYPES: + for nbytes in SIZES: + for kind in KINDS: + t = make_tensor(dt, kind, nbytes) + rows.append(measure(t, framing, dname, kind)) + return rows + + +def main() -> None: + header = ( + f"{'framing':<10} {'dtype':<5} {'data':<9} {'raw B':>9} {'out B':>9} " + f"{'non-pay B':>9} {'ratio':>8} {'zstd19 B':>9} {'zstd ratio':>10}" + ) + print(header) + print("-" * len(header)) + for r in sweep(): + print( + f"{r.framing:<10} {r.dtype:<5} {r.kind:<9} {r.raw:>9,} {r.out:>9,} " + f"{r.non_payload:>9,} {r.ratio:>8.3f} {r.zstd19:>9,} {r.zstd_ratio:>10.2f}" + ) + + print( + f"\n{'framing':<10} {'rows':>5} {'raw B':>9} {'out B':>7} {'ratio':>7} {'zstd19 B':>9}" + ) + for framing in (Framing.COMPACT, Framing.CONTAINER): + for rows in (1, 8, 64): + r = measure(tiled_rows(rows), framing, "bf16", "tiled") + print( + f"{r.framing:<10} {rows:>5} {r.raw:>9,} {r.out:>7,} " + f"{r.ratio:>7.1f} {r.zstd19:>9,}" + ) + + +if __name__ == "__main__": + main() diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs new file mode 100644 index 0000000..862f68a --- /dev/null +++ b/crates/ptwm-core/src/compact.rs @@ -0,0 +1,749 @@ +//! Compact single-tensor frame (magic `\x89PTWS`). +//! +//! A full `.ptwm` container carries a fixed header, an Extension Table, a +//! prelude, a chain registry, a tensor index and an EOF sentinel. For one +//! small tensor those sections cost more than the payload. The compact frame +//! holds one tensor with only the fields its decoder needs: +//! +//! ```text +//! magic: [u8; 5] \x89PTWS +//! flags: u8 see FLAG_*; unknown bits are rejected +//! dtype_code: varint +//! input_format: u8 +//! orig_size: varint raw byte count +//! [shape] varint ndim, then varint per dimension (FLAG_SHAPE) +//! [dtype_name] varint len, then UTF-8 bytes (FLAG_DTYPE_NAME) +//! [delta_ref] [u8; 32] BLAKE3 of the delta reference (FLAG_DELTA_REF) +//! [payload_hash] u64 xxh64 of the raw bytes (FLAG_PAYLOAD_HASH) +//! body +//! ``` +//! +//! When `FLAG_STORED` is set, the body is the `orig_size` raw bytes. Else the +//! body is a coded tensor: +//! +//! ```text +//! n_ops: u8 +//! op_ids: [u16; n_ops] OpId wire ids; the chain's u16 op +//! indices refer to this list +//! chain: chain::wire::write_chain bytes +//! per terminal plane, in chain terminal order: +//! codec_id: varint CodecId wire id +//! state_source: u8 0 = none, 1 = inline +//! [state_version] u8 (inline only) +//! [state_len] varint (inline only) +//! [state] bytes (inline only) +//! payload_len: varint +//! [crc32] u32 CRC32 of the raw plane (FLAG_PLANE_CRC) +//! payload: bytes +//! ``` +//! +//! All integers are little-endian; varints are unsigned LEB128. The frame +//! ends exactly after the body; trailing bytes are rejected. +//! +//! The plane layout, the decoded plane length and the plane count all come +//! from the chain, so the frame does not store them. Planes that need shared +//! state (a prelude) or chunking cannot use this frame; the caller falls back +//! to the full container for those. +//! +//! The magic differs from `\x89PTWM`, so readers that predate this frame +//! reject it as "not a .ptwm container" instead of misparsing it. Like the +//! container, the frame has no numeric version field: a reader rejects +//! unknown flag bits, and an incompatible layout gets a new magic. + +use crate::chain::Chain; +use crate::chain::dispatch::op_from_id; +use crate::chain::runtime::{InverseContext, chain_terminal_descriptors, inverse_chain}; +use crate::chain::validate::validate_chain; +use crate::chain::wire::{read_chain, write_chain}; +use crate::codec::{CodecId, StateSource, codec_for}; +use crate::compressor::{CompressorOptions, InputTensor, legacy_plane_layout, plan_model}; +use crate::dtype::Dtype; +use crate::error::PtwmCoreError; +use crate::extension::dispatch::{BuiltinKind, dispatch_builtin}; +use crate::extension::{ExtensionTable, ExtensionTableBuilder, builtin_entries}; +use crate::transforms::op::{OpId, Plane}; + +use std::sync::Arc; +use xxhash_rust::xxh64::xxh64; + +/// Compact-frame magic. Shares the `\x89PTW` prefix with the full container. +pub const MAGIC_COMPACT: [u8; 5] = [0x89, b'P', b'T', b'W', b'S']; + +/// The frame carries an xxh64 of the raw bytes. +pub const FLAG_PAYLOAD_HASH: u8 = 0x01; +/// Every coded plane carries a CRC32 of its raw bytes. +pub const FLAG_PLANE_CRC: u8 = 0x02; +/// The frame carries the tensor shape. +pub const FLAG_SHAPE: u8 = 0x04; +/// The frame carries a dtype name that differs from the canonical name of +/// `dtype_code`. +pub const FLAG_DTYPE_NAME: u8 = 0x08; +/// The frame carries the BLAKE3 hash of a delta reference. +pub const FLAG_DELTA_REF: u8 = 0x10; +/// The body is the raw bytes, not a coded tensor. +pub const FLAG_STORED: u8 = 0x20; +const KNOWN_FLAGS: u8 = FLAG_PAYLOAD_HASH + | FLAG_PLANE_CRC + | FLAG_SHAPE + | FLAG_DTYPE_NAME + | FLAG_DELTA_REF + | FLAG_STORED; + +/// Upper bound on the dimension count, so a corrupt `ndim` cannot drive a +/// large allocation. +const MAX_NDIM: u64 = 64; + +/// A tensor decoded from a compact frame. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CompactTensor { + pub dtype_code: u16, + pub input_format: u8, + /// Shape, when the writer supplied one. + pub shape: Option>, + /// Dtype name paired with `shape`. When the frame does not store a name, + /// this is the canonical name of `dtype_code` (or `None` if the code is + /// unknown). + pub dtype_name: Option, + /// BLAKE3 hash of the delta reference, when the tensor is delta-coded. + pub delta_reference_blake3: Option<[u8; 32]>, + /// Decoded raw bytes. + pub raw: Vec, + /// Bytes of codec payload and inline codec state in the frame (the raw + /// byte count for a stored frame). The frame length minus this value is + /// the frame's fixed overhead. + pub payload_len: u64, +} + +/// Return `true` when `buf` starts with the compact-frame magic. +pub fn is_compact(buf: &[u8]) -> bool { + buf.len() >= MAGIC_COMPACT.len() && buf[..MAGIC_COMPACT.len()] == MAGIC_COMPACT +} + +fn canonical_dtype_name(dtype_code: u16) -> Option<&'static str> { + u8::try_from(dtype_code) + .ok() + .and_then(|c| Dtype::from_code(c).ok()) + .map(Dtype::canonical_name) +} + +/// Encode one tensor as a compact frame. +/// +/// Runs the same chain and codec search as +/// [`crate::compressor::compress_model`], then writes the smaller of the +/// coded body and the raw bytes. Returns `Ok(None)` when the tensor cannot use +/// this frame: chunked encoding (`opts.chunk_size`) or a plane that selected +/// shared state. The caller then writes a full container. +pub fn compress_compact( + input: &InputTensor<'_>, + opts: &CompressorOptions, +) -> Result>, PtwmCoreError> { + if opts.chunk_size.is_some() { + return Ok(None); + } + let raw = input.raw_bytes; + + let mut out = Vec::with_capacity(64); + out.extend_from_slice(&MAGIC_COMPACT); + let flags_pos = out.len(); + out.push(0); + let mut flags = 0u8; + write_varint(&mut out, u64::from(input.dtype_code)); + out.push(input.input_format); + write_varint(&mut out, raw.len() as u64); + // Shape metadata is recorded only when both shape and dtype name are + // known, matching the full container's tensor metadata. + if let (Some(shape), Some(name)) = (&input.shape, &input.dtype_name) { + flags |= FLAG_SHAPE; + write_varint(&mut out, shape.len() as u64); + for &d in shape { + write_varint(&mut out, d); + } + if canonical_dtype_name(input.dtype_code) != Some(name.as_str()) { + flags |= FLAG_DTYPE_NAME; + write_varint(&mut out, name.len() as u64); + out.extend_from_slice(name.as_bytes()); + } + } + if let Some(h) = &input.delta_reference_blake3 { + flags |= FLAG_DELTA_REF; + out.extend_from_slice(h); + } + if opts.emit_payload_hash { + flags |= FLAG_PAYLOAD_HASH; + out.extend_from_slice(&xxh64(raw, 0).to_le_bytes()); + } + + let coded = encode_body(input, opts)?; + match coded { + None => return Ok(None), + Some(body) if body.len() < raw.len() => { + if opts.emit_plane_crc { + flags |= FLAG_PLANE_CRC; + } + out.extend_from_slice(&body); + } + Some(_) => { + flags |= FLAG_STORED; + out.extend_from_slice(raw); + } + } + out[flags_pos] = flags; + Ok(Some(out)) +} + +/// Encode the coded body (ops, chain, planes). `Ok(None)` when a plane needs +/// shared state. +fn encode_body( + input: &InputTensor<'_>, + opts: &CompressorOptions, +) -> Result>, PtwmCoreError> { + let plan = plan_model(std::slice::from_ref(input), opts)?; + if !plan.prelude_entries.is_empty() { + return Ok(None); + } + let chain: &Chain = &input.candidate_chains[plan.pass1[0].chain_idx]; + let planes = &plan.chosen_per_tensor[0]; + + let mut builder = ExtensionTableBuilder::new(); + let mut chain_bytes = Vec::new(); + write_chain(chain, &mut builder, &mut chain_bytes)?; + let entries = builder.finish(); + let n_ops = u8::try_from(entries.len()).map_err(|_| { + PtwmCoreError::InvalidContainer(format!( + "compact frame: chain uses {} ops (max 255)", + entries.len() + )) + })?; + + let mut body = Vec::with_capacity( + 16 + chain_bytes.len() + + planes + .iter() + .map(|p| p.payload_bytes.len() + 16) + .sum::(), + ); + body.push(n_ops); + for entry in &entries { + let op = match dispatch_builtin(&entry.canonical_id) { + Some(BuiltinKind::Op(op)) => op, + _ => { + return Err(PtwmCoreError::InvalidContainer( + "compact frame: chain op is not a built-in transform".into(), + )); + } + }; + body.extend_from_slice(&op.as_u16().to_le_bytes()); + } + body.extend_from_slice(&chain_bytes); + + for cp in planes { + write_varint(&mut body, u64::from(cp.codec_id.as_u16())); + match cp.state_source { + StateSource::None => body.push(0), + StateSource::Inline => { + body.push(1); + body.push(cp.encoded_state_version); + write_varint(&mut body, cp.encoded_state_bytes.len() as u64); + body.extend_from_slice(&cp.encoded_state_bytes); + } + StateSource::Shared | StateSource::External => return Ok(None), + } + write_varint(&mut body, cp.payload_bytes.len() as u64); + if opts.emit_plane_crc { + body.extend_from_slice(&crc32fast::hash(&cp.raw_plane_for_crc).to_le_bytes()); + } + body.extend_from_slice(&cp.payload_bytes); + } + Ok(Some(body)) +} + +/// Decode a compact frame. Verifies the payload hash and plane CRCs when the +/// frame carries them. +pub fn decode_compact(buf: &[u8]) -> Result { + let mut r = Reader { buf, pos: 0 }; + if !is_compact(buf) { + return Err(bad("bad magic")); + } + r.pos = MAGIC_COMPACT.len(); + let flags = r.u8()?; + if flags & !KNOWN_FLAGS != 0 { + return Err(bad(&format!( + "unknown flag bits {:#04x}", + flags & !KNOWN_FLAGS + ))); + } + let dtype_code = u16::try_from(r.varint()?).map_err(|_| bad("dtype_code exceeds u16"))?; + let input_format = r.u8()?; + let orig_size = r.varint()?; + + let mut shape = None; + let mut dtype_name = None; + if flags & FLAG_SHAPE != 0 { + let ndim = r.varint()?; + if ndim > MAX_NDIM { + return Err(bad(&format!("ndim {ndim} exceeds {MAX_NDIM}"))); + } + let mut dims = Vec::with_capacity(ndim as usize); + for _ in 0..ndim { + dims.push(r.varint()?); + } + shape = Some(dims); + dtype_name = canonical_dtype_name(dtype_code).map(str::to_owned); + } + if flags & FLAG_DTYPE_NAME != 0 { + if flags & FLAG_SHAPE == 0 { + return Err(bad("dtype name without shape")); + } + let len = r.len_varint()?; + let bytes = r.take(len)?; + let name = std::str::from_utf8(bytes).map_err(|_| bad("dtype name is not UTF-8"))?; + dtype_name = Some(name.to_owned()); + } + let delta_reference_blake3 = if flags & FLAG_DELTA_REF != 0 { + let mut h = [0u8; 32]; + h.copy_from_slice(r.take(32)?); + Some(h) + } else { + None + }; + let payload_hash = if flags & FLAG_PAYLOAD_HASH != 0 { + Some(u64::from_le_bytes(r.take(8)?.try_into().unwrap())) + } else { + None + }; + + let mut payload_len = 0u64; + let raw = if flags & FLAG_STORED != 0 { + if flags & FLAG_PLANE_CRC != 0 { + return Err(bad("stored frame with plane CRC flag")); + } + let rest = r.rest(); + r.pos = buf.len(); + if rest.len() as u64 != orig_size { + return Err(bad(&format!( + "stored body is {} bytes, header says {orig_size}", + rest.len() + ))); + } + payload_len = orig_size; + rest.to_vec() + } else { + decode_body( + &mut r, + flags & FLAG_PLANE_CRC != 0, + orig_size, + &mut payload_len, + )? + }; + if r.pos != buf.len() { + return Err(bad("trailing bytes after body")); + } + + if let Some(expected) = payload_hash + && xxh64(&raw, 0) != expected + { + return Err(bad("payload hash mismatch")); + } + + Ok(CompactTensor { + dtype_code, + input_format, + shape, + dtype_name, + delta_reference_blake3, + raw, + payload_len, + }) +} + +fn decode_body( + r: &mut Reader<'_>, + has_crc: bool, + orig_size: u64, + payload_len_out: &mut u64, +) -> Result, PtwmCoreError> { + let n_ops = r.u8()? as usize; + let all = builtin_entries(); + let mut table = ExtensionTable::empty(); + for _ in 0..n_ops { + let wire = u16::from_le_bytes(r.take(2)?.try_into().unwrap()); + let op = OpId::from_u16(wire)?; + let id = op.canonical_id(); + let entry = all + .iter() + .find(|e| e.canonical_id == id) + .ok_or_else(|| bad(&format!("op {wire:#06x} has no built-in entry")))?; + table.entries.push(entry.clone()); + } + let (chain, consumed) = read_chain(r.rest(), &table)?; + r.pos += consumed; + validate_chain(&chain, 0)?; + // `chain_terminal_descriptors` runs the chain forward on a zero buffer + // of the size the Source node declares, and codecs pre-allocate the + // plane lengths that follow from it. Bound the declared size by the + // frame's `orig_size` first, so a corrupt chain cannot request gigabytes + // for a small tensor. The factor covers a source declared in elements of + // up to 8 bytes for byte input; the constant covers block padding. + let max_len = orig_size.saturating_mul(8).saturating_add(4096); + let first = chain + .nodes + .first() + .ok_or_else(|| bad("chain has no nodes"))?; + if first.op != OpId::Source { + return Err(bad("chain node 0 is not Source")); + } + let source = op_from_id(OpId::Source, &first.params)?.propagate_descriptors(&[])?; + if source.iter().any(|d| d.length_bytes > max_len) { + return Err(bad(&format!( + "chain declares a source larger than {max_len} bytes for a \ + {orig_size}-byte tensor" + ))); + } + let descriptors = chain_terminal_descriptors(&chain)?; + if descriptors.iter().any(|d| d.length_bytes > max_len) { + return Err(bad("chain declares a plane larger than its bound")); + } + + let mut terminal_planes: Vec = Vec::with_capacity(descriptors.len()); + for descriptor in descriptors { + let wire = u16::try_from(r.varint()?).map_err(|_| bad("codec id exceeds u16"))?; + let codec_id = + CodecId::from_u16(wire).ok_or_else(|| bad(&format!("unknown codec {wire:#06x}")))?; + let (state_version, state) = match r.u8()? { + 0 => (0u8, &[][..]), + 1 => { + let v = r.u8()?; + let len = r.len_varint()?; + (v, r.take(len)?) + } + s => return Err(bad(&format!("unknown state source {s}"))), + }; + let payload_len = r.len_varint()?; + let crc = if has_crc { + Some(u32::from_le_bytes(r.take(4)?.try_into().unwrap())) + } else { + None + }; + let payload = r.take(payload_len)?; + *payload_len_out += (state.len() + payload.len()) as u64; + + let codec = codec_for(codec_id) + .ok_or_else(|| bad(&format!("no codec implementation for {codec_id:?}")))?; + let layout = legacy_plane_layout(descriptor.layout); + let decoded_len = usize::try_from(descriptor.length_bytes) + .map_err(|_| bad("plane length overflows usize"))?; + // Self-delimiting codecs may return a length other than the chain + // descriptor's (the descriptor can overstate it for byte input with a + // multi-byte dtype). The container reader accepts this too; the + // `orig_size` check and the payload hash below cover the result. + let bytes = codec.decode(state_version, state, payload, &layout, decoded_len)?; + if let Some(expected) = crc + && crc32fast::hash(&bytes) != expected + { + return Err(bad("plane CRC mismatch")); + } + terminal_planes.push(Plane { + bytes: Arc::from(bytes.into_boxed_slice()), + descriptor, + }); + } + + let ctx = InverseContext { + terminal_planes: &terminal_planes, + deps: &[], + }; + let source = inverse_chain(&chain, &ctx)?; + if source.bytes.len() as u64 != orig_size { + return Err(bad(&format!( + "decoded {} bytes, header says {orig_size}", + source.bytes.len() + ))); + } + Ok(source.bytes.as_ref().to_vec()) +} + +fn bad(msg: &str) -> PtwmCoreError { + PtwmCoreError::InvalidContainer(format!("compact frame: {msg}")) +} + +fn write_varint(out: &mut Vec, mut v: u64) { + while v >= 0x80 { + out.push((v as u8) | 0x80); + v >>= 7; + } + out.push(v as u8); +} + +struct Reader<'a> { + buf: &'a [u8], + pos: usize, +} + +impl<'a> Reader<'a> { + fn take(&mut self, n: usize) -> Result<&'a [u8], PtwmCoreError> { + let end = self + .pos + .checked_add(n) + .filter(|&e| e <= self.buf.len()) + .ok_or_else(|| bad("truncated"))?; + let s = &self.buf[self.pos..end]; + self.pos = end; + Ok(s) + } + + fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + fn rest(&self) -> &'a [u8] { + &self.buf[self.pos..] + } + + fn varint(&mut self) -> Result { + let mut v: u64 = 0; + for i in 0..10 { + let b = self.u8()?; + let bits = u64::from(b & 0x7F); + if i == 9 && b > 1 { + return Err(bad("varint overflows u64")); + } + v |= bits << (7 * i); + if b & 0x80 == 0 { + return Ok(v); + } + } + Err(bad("varint overflows u64")) + } + + /// A varint used as a byte length; must fit in the remaining buffer. + fn len_varint(&mut self) -> Result { + let v = self.varint()?; + let remaining = (self.buf.len() - self.pos) as u64; + if v > remaining { + return Err(bad("length exceeds remaining bytes")); + } + Ok(v as usize) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::chain::graph::{ChainEdge, ChainNode, TerminalRef}; + use crate::types::role::Role; + + const FLAGS_AT: usize = MAGIC_COMPACT.len(); + + fn source_params(n_elements: u32, dtype_code: u16) -> Vec { + // 1-D shape encoding (matches `Source::write_params`). + let mut p = vec![1u8]; + p.extend_from_slice(&n_elements.to_le_bytes()); + p.extend_from_slice(&dtype_code.to_le_bytes()); + p + } + + /// Source(int8) → BytePassthrough → one Value terminal. + fn passthrough_chain(n_bytes: u32) -> Chain { + let role = Role::Value { + format: crate::types::role::ValueFormat::IntN { bits: 8 }, + }; + Chain { + nodes: vec![ + ChainNode { + op: OpId::Source, + params: source_params(n_bytes, 17), + }, + ChainNode { + op: OpId::BytePassthrough, + params: Vec::new(), + }, + ], + edges: vec![ChainEdge { + src_node: 0, + src_output_idx: 0, + dst_node: 1, + dst_input_idx: 0, + role_override: Some(role.clone()), + vendor_bytes: Vec::new(), + }], + terminals: vec![TerminalRef { + node_idx: 1, + output_idx: 0, + role, + }], + } + } + + fn opts(hash: bool, crc: bool) -> CompressorOptions { + CompressorOptions { + method_hint: 1, + emit_payload_hash: hash, + emit_plane_crc: crc, + forced_codec: None, + allow_codec_ids: None, + chunk_size: None, + } + } + + fn input<'a>(raw: &'a [u8], shape: Option>, name: Option<&str>) -> InputTensor<'a> { + InputTensor { + name: "_".into(), + candidate_chains: vec![passthrough_chain(raw.len() as u32)], + dtype_code: 0, + input_format: 1, + raw_bytes: raw, + shape, + dtype_name: name.map(str::to_owned), + delta_reference_blake3: None, + } + } + + #[test] + fn constant_tensor_round_trips_coded() { + let raw = vec![7u8; 4096]; + for (h, c) in [(true, true), (false, false), (true, false), (false, true)] { + let blob = compress_compact(&input(&raw, None, None), &opts(h, c)) + .unwrap() + .unwrap(); + assert!(blob.len() < 128, "constant frame is {} bytes", blob.len()); + assert_eq!( + blob[FLAGS_AT] & FLAG_STORED, + 0, + "constant data must be coded" + ); + let t = decode_compact(&blob).unwrap(); + assert_eq!(t.raw, raw); + assert_eq!(t.shape, None); + } + } + + #[test] + fn incompressible_tensor_is_stored() { + // xorshift bytes: no codec beats the raw size. + let mut x: u64 = 0x9E37_79B9_7F4A_7C15; + let raw: Vec = (0..300) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x as u8 + }) + .collect(); + let blob = compress_compact( + &input(&raw, Some(vec![300]), Some("uint8")), + &opts(true, true), + ) + .unwrap() + .unwrap(); + assert_ne!(blob[FLAGS_AT] & FLAG_STORED, 0); + assert!(blob.len() <= raw.len() + 32); + let t = decode_compact(&blob).unwrap(); + assert_eq!(t.raw, raw); + assert_eq!(t.shape, Some(vec![300])); + assert_eq!(t.dtype_name.as_deref(), Some("uint8")); + } + + #[test] + fn noncanonical_dtype_name_is_kept() { + let raw = vec![1u8; 64]; + let blob = compress_compact( + &input(&raw, Some(vec![8, 8]), Some("custom")), + &opts(true, true), + ) + .unwrap() + .unwrap(); + assert_ne!(blob[FLAGS_AT] & FLAG_DTYPE_NAME, 0); + let t = decode_compact(&blob).unwrap(); + assert_eq!(t.shape, Some(vec![8, 8])); + assert_eq!(t.dtype_name.as_deref(), Some("custom")); + } + + #[test] + fn empty_tensor_round_trips() { + let raw: Vec = Vec::new(); + let blob = compress_compact( + &input(&raw, Some(vec![0]), Some("uint8")), + &opts(true, true), + ) + .unwrap() + .unwrap(); + let t = decode_compact(&blob).unwrap(); + assert!(t.raw.is_empty()); + } + + #[test] + fn chunked_request_is_not_eligible() { + let raw = vec![0u8; 64]; + let mut o = opts(true, true); + o.chunk_size = Some(16); + assert!( + compress_compact(&input(&raw, None, None), &o) + .unwrap() + .is_none() + ); + } + + #[test] + fn corruption_is_detected() { + let raw: Vec = (0..2048u32).map(|i| (i % 7) as u8).collect(); + let blob = compress_compact(&input(&raw, None, None), &opts(true, true)) + .unwrap() + .unwrap(); + // Every truncation fails cleanly. + for cut in 0..blob.len() { + assert!( + decode_compact(&blob[..cut]).is_err(), + "cut at {cut} decoded" + ); + } + // Trailing bytes are rejected. + let mut longer = blob.clone(); + longer.push(0); + assert!(decode_compact(&longer).is_err()); + // Unknown flag bits are rejected. + let mut flagged = blob.clone(); + flagged[FLAGS_AT] |= 0x80; + assert!(decode_compact(&flagged).is_err()); + // A flipped payload byte fails a check or decodes to different bytes + // that the payload hash rejects; it must never return wrong data. + for i in FLAGS_AT + 1..blob.len() { + let mut b = blob.clone(); + b[i] ^= 0x55; + if let Ok(t) = decode_compact(&b) { + assert_eq!(t.raw, raw, "byte {i} flip returned wrong data"); + } + } + } + + #[test] + fn oversized_chain_source_is_rejected() { + let raw: Vec = (0..3000u32).map(|i| (i % 5) as u8).collect(); + let blob = compress_compact(&input(&raw, None, None), &opts(false, false)) + .unwrap() + .unwrap(); + assert_eq!(blob[FLAGS_AT] & FLAG_STORED, 0); + // Rewrite the Source dimension (1-D, 3000) to about 4 G elements. + let needle: Vec = [&[1u8][..], &3000u32.to_le_bytes()].concat(); + let at = blob + .windows(needle.len()) + .position(|w| w == needle.as_slice()) + .unwrap(); + let mut b = blob.clone(); + b[at + 1..at + 5].copy_from_slice(&0xFFFF_FFF0u32.to_le_bytes()); + let t0 = std::time::Instant::now(); + let err = decode_compact(&b).unwrap_err().to_string(); + assert!(err.contains("source larger than"), "{err}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(1)); + } + + #[test] + fn varint_round_trips() { + for v in [0u64, 1, 127, 128, 300, u32::MAX as u64, u64::MAX] { + let mut out = Vec::new(); + write_varint(&mut out, v); + let mut r = Reader { buf: &out, pos: 0 }; + assert_eq!(r.varint().unwrap(), v); + assert_eq!(r.pos, out.len()); + } + } +} diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 32e43a6..7d602b1 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -38,7 +38,7 @@ use crate::types::role::{NibbleKind, Role}; /// Map a builtin [`CodecId`] to its [`CanonicalId`] so the compressor can /// intern the codec into the Extension Table and fill `codec_table_idx`. -fn codec_canonical_id(id: CodecId) -> CanonicalId { +pub(crate) fn codec_canonical_id(id: CodecId) -> CanonicalId { use crate::extension::builtin_canonical_id; match id { CodecId::Identity => builtin_canonical_id("identity"), @@ -139,7 +139,7 @@ fn legacy_plane_role(role: &Role) -> PlaneRole { } /// Map a typed [`Layout`] to the on-disk [`PlaneLayout`]. -fn legacy_plane_layout(layout: Layout) -> PlaneLayout { +pub(crate) fn legacy_plane_layout(layout: Layout) -> PlaneLayout { match layout { Layout::Flat => PlaneLayout::Flat, Layout::Rows { row_len } => PlaneLayout::rows(row_len).unwrap_or(PlaneLayout::Flat), @@ -288,18 +288,18 @@ fn append_shared_menu_items( } /// One terminal plane chosen by the trial-encoder, ready to be written. -struct ChosenTerminal { - role: PlaneRole, - codec_id: CodecId, - state_source: StateSource, - shared_state_id: Option, - encoded_state_bytes: Vec, - encoded_state_version: u8, - payload_bytes: Vec, +pub(crate) struct ChosenTerminal { + pub(crate) role: PlaneRole, + pub(crate) codec_id: CodecId, + pub(crate) state_source: StateSource, + pub(crate) shared_state_id: Option, + pub(crate) encoded_state_bytes: Vec, + pub(crate) encoded_state_version: u8, + pub(crate) payload_bytes: Vec, // Arc-shared with the terminal plane: an O(1) clone, avoiding a full copy // of every plane (~the whole model) in the hot Pass-1 loop. - raw_plane_for_crc: Arc<[u8]>, - layout: PlaneLayout, + pub(crate) raw_plane_for_crc: Arc<[u8]>, + pub(crate) layout: PlaneLayout, } /// Frame a single encoded chunk for storage inside a plane's payload_bytes. @@ -460,8 +460,8 @@ fn trial_encode_terminal( /// terminal planes, and the per-terminal `TrialResult`s under inline /// menus only — Pass 2 (shared-state fit) re-trials terminals that may /// benefit from a `Shared` variant. -struct Pass1Outcome { - chain_idx: usize, +pub(crate) struct Pass1Outcome { + pub(crate) chain_idx: usize, terminals: Vec<(Plane, Role)>, trials: Vec, /// Forced+chunked skip path only: the per-terminal `(payload, chunk_table)` @@ -637,25 +637,21 @@ fn build_chain_registry(chains: &[Chain]) -> (ChainRegistry, Vec) { // compress_model // --------------------------------------------------------------------------- -/// PPG entry point. Each input tensor carries one or more candidate -/// [`Chain`]s; the compressor selects the chain that produces the smallest -/// total compressed payload (multi-chain trial-encode), then writes a -/// chain-registry-backed `.ptwm` container. -/// -/// Returns the per-tensor chosen chain *registry index* (after dedup) for -/// audit. The order matches the input slice. -/// -/// **Pass-2 ordering note.** Shared-state fitting is performed *after* -/// Pass-1 multi-chain selection. This means the multi-chain ranker sees only -/// inline trial-encode payloads. For monotone shared-state codecs (PGC, -/// O1SAC), Inline is the lower bound on what the chain can achieve, so the -/// multi-chain ranking remains correct; per-terminal final encoding may -/// later switch to a `Shared` variant if it wins the re-trial. -pub fn compress_model( - writer: W, +/// Output of passes 0–3: the chosen chain and per-plane codec encoding of +/// every tensor, plus the shared-state prelude entries those planes use. +pub(crate) struct ModelPlan { + pub(crate) pass1: Vec, + pub(crate) chosen_per_tensor: Vec>, + pub(crate) prelude_entries: Vec, +} + +/// Passes 0–3 of [`compress_model`]: validate the candidate chains, select a +/// chain per tensor, fit shared state, and pick the final codec per plane. +/// Writes nothing; the caller serializes the plan. +pub(crate) fn plan_model( inputs: &[InputTensor<'_>], - opts: CompressorOptions, -) -> Result, PtwmCoreError> { + opts: &CompressorOptions, +) -> Result { // ---- Pass 0: structural validation of every candidate chain. ---------- // Reject malformed chains at the API boundary so downstream stages can // assume well-formed graphs (Source at 0, no cycles, all output indices @@ -732,7 +728,7 @@ pub fn compress_model( for (i, o) in computed { pass1_slots[i] = Some(o); } - let mut pass1: Vec = pass1_slots + let pass1: Vec = pass1_slots .into_iter() .map(|o| o.expect("every tensor index computed in pass1")) .collect(); @@ -846,7 +842,9 @@ pub fn compress_model( chosen_per_tensor.push(chosen_planes); } - prof!(prof_t, "pass2+3_fit+retrial"); + if profile { + eprintln!("[ptwm-profile] pass2+3_fit+retrial: {:?}", prof_t.elapsed()); + } // Drop unreferenced prelude entries. let prelude_entries: Vec = prelude_candidates @@ -854,6 +852,47 @@ pub fn compress_model( .filter(|e| used_shared_ids.contains(&e.shared_state_id)) .collect(); + Ok(ModelPlan { + pass1, + chosen_per_tensor, + prelude_entries, + }) +} + +/// PPG entry point. Each input tensor carries one or more candidate +/// [`Chain`]s; the compressor selects the chain that produces the smallest +/// total compressed payload (multi-chain trial-encode), then writes a +/// chain-registry-backed `.ptwm` container. +/// +/// Returns the per-tensor chosen chain *registry index* (after dedup) for +/// audit. The order matches the input slice. +/// +/// **Pass-2 ordering note.** Shared-state fitting is performed *after* +/// Pass-1 multi-chain selection. This means the multi-chain ranker sees only +/// inline trial-encode payloads. For monotone shared-state codecs (PGC, +/// O1SAC), Inline is the lower bound on what the chain can achieve, so the +/// multi-chain ranking remains correct; per-terminal final encoding may +/// later switch to a `Shared` variant if it wins the re-trial. +pub fn compress_model( + writer: W, + inputs: &[InputTensor<'_>], + opts: CompressorOptions, +) -> Result, PtwmCoreError> { + let profile = std::env::var_os("PTWM_PROFILE").is_some(); + let ModelPlan { + mut pass1, + chosen_per_tensor, + prelude_entries, + } = plan_model(inputs, &opts)?; + let mut prof_t = std::time::Instant::now(); + macro_rules! prof { + ($t:expr, $label:expr) => { + if profile { + eprintln!("[ptwm-profile] {}: {:?}", $label, $t.elapsed()); + $t = std::time::Instant::now(); + } + }; + } // ---- Pass 4: write container. ----------------------------------------- // Build a deduped chain registry from the chosen chain per tensor. let chosen_chains: Vec = pass1 @@ -865,7 +904,16 @@ pub fn compress_model( let mut writer = ContainerWriter::new(writer)?; writer.set_prelude(&prelude_entries)?; - writer.set_chain_registry(registry)?; + // Intern only the codecs the planes reference: interning every built-in + // codec costs about 125 B per unused entry, which dominates small files. + let mut used_codecs: Vec = Vec::new(); + for cp in chosen_per_tensor.iter().flatten() { + let id = codec_canonical_id(cp.codec_id); + if !used_codecs.contains(&id) { + used_codecs.push(id); + } + } + writer.set_chain_registry_with_codecs(registry, &used_codecs)?; // Pre-compute the chunked-path payloads in parallel before the sequential // write loop. Otherwise encode_plane_chunked runs inside the writer's diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index dc84825..0d6a320 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -105,6 +105,82 @@ impl ChainRegistry { } } +/// Which plane codecs an Extension Table interns besides the chain ops. +#[derive(Clone, Copy, Debug)] +pub enum CodecSet<'a> { + /// Every built-in plane codec (legacy layout, v1.1 and earlier). + AllBuiltins, + /// Only the listed codecs, interned in built-in registry order. + Only(&'a [CanonicalId]), +} + +/// Output of [`build_extension_table`]. +pub struct BuiltExtensionTable { + pub table: ExtensionTable, + /// Serialized chain registry whose op indices refer to `table`. + pub registry_bytes: Vec, + /// Table index of every interned plane codec. + pub codec_idx_map: HashMap, +} + +/// Build the Extension Table and the serialized chain registry for a +/// container. Chain ops are interned first (in registry order), then the +/// plane codecs selected by `codecs`, in the order of +/// [`crate::extension::builtin_entries`]. +/// +/// The result is a pure function of `registry` and the codec set, so a +/// reader-side tool (the transcoder) can rebuild the same bytes. +pub fn build_extension_table( + registry: &ChainRegistry, + codecs: CodecSet<'_>, +) -> Result { + use crate::extension::Kind; + use crate::extension::builtin_entries; + + let mut builder = ExtensionTableBuilder::new(); + let mut registry_bytes = Vec::new(); + registry.write(&mut builder, &mut registry_bytes)?; + + // Use `builtin_entries()` (the curated list with `is_builtin` coverage) + // filtered by `Kind::PlaneCodec` rather than hard-coded names. This + // excludes `huffman_nibble`, which `dispatch_builtin` keeps for legacy + // wire compat but which has no `builtin_entry`, so `is_builtin` and + // `collect_missing_flavors` do not recognise it. + let codec_entries: Vec<_> = builtin_entries() + .into_iter() + .filter(|e| e.kind == Kind::PlaneCodec) + .collect(); + if let CodecSet::Only(wanted) = codecs { + for id in wanted { + if !codec_entries.iter().any(|e| e.canonical_id == *id) { + return Err(PtwmCoreError::InvalidContainer(format!( + "codec {id:?} has no built-in Extension Table entry" + ))); + } + } + } + let mut codec_idx_map = HashMap::new(); + for entry in codec_entries { + let keep = match codecs { + CodecSet::AllBuiltins => true, + CodecSet::Only(wanted) => wanted.contains(&entry.canonical_id), + }; + if keep { + let canonical_id = entry.canonical_id; + let idx = builder.intern(entry); + codec_idx_map.insert(canonical_id, idx); + } + } + + Ok(BuiltExtensionTable { + table: ExtensionTable { + entries: builder.finish(), + }, + registry_bytes, + codec_idx_map, + }) +} + /// Streaming writer for a `.ptwm` container. pub struct ContainerWriter { writer: W, @@ -191,49 +267,54 @@ impl ContainerWriter { /// second table-building pass. Any previously set extension table is /// replaced. /// + /// This is the legacy (v1.1 and earlier) table layout. Writers that know + /// the set of codecs in use should call + /// [`set_chain_registry_with_codecs`](Self::set_chain_registry_with_codecs) + /// instead, which omits the unreferenced codec entries (about 125 B each). + /// /// After this call, use [`codec_table_idx_for`](Self::codec_table_idx_for) /// to obtain the stable index for a codec's canonical id. /// /// Must be called before any `append_tensor` call. pub fn set_chain_registry(&mut self, registry: ChainRegistry) -> Result<(), PtwmCoreError> { + self.install_registry(registry, CodecSet::AllBuiltins) + } + + /// Set the chain registry and build an Extension Table that holds only + /// the chain ops plus the plane codecs in `codecs`. + /// + /// Codec entries are interned in the order of + /// [`crate::extension::builtin_entries`], not in the order of `codecs`, + /// so the table is a pure function of the registry and the codec *set*. + /// The transcoder relies on this to rebuild the same table bytes. + /// + /// Returns `InvalidContainer` if a codec in `codecs` has no built-in + /// Extension Table entry. + /// + /// Must be called before any `append_tensor` call. + pub fn set_chain_registry_with_codecs( + &mut self, + registry: ChainRegistry, + codecs: &[CanonicalId], + ) -> Result<(), PtwmCoreError> { + self.install_registry(registry, CodecSet::Only(codecs)) + } + + fn install_registry( + &mut self, + registry: ChainRegistry, + codecs: CodecSet<'_>, + ) -> Result<(), PtwmCoreError> { if self.prelude_written || !self.index_entries.is_empty() { return Err(PtwmCoreError::InvalidContainer( "chain registry must be set before appending tensors".into(), )); } - // Build the extension table from the registry chains, producing the - // chain registry bytes in the same pass so table indices are consistent. - let mut builder = ExtensionTableBuilder::new(); - let mut registry_bytes = Vec::new(); - registry.write(&mut builder, &mut registry_bytes)?; - - // Additionally intern all builtin plane-codec entries so every codec - // the compressor selects gets a stable u16 table index for - // `PlaneRecord::codec_table_idx`. Ops that are already in the table - // from chain interning just get their existing index back. - // - // Note: we use `builtin_entries()` (the curated list with `is_builtin` - // coverage) and filter by `Kind::PlaneCodec` rather than hard-coding - // names. This avoids interning `huffman_nibble`, which is present in - // `dispatch_builtin` for legacy wire compat but has no `builtin_entry` - // and therefore is NOT recognised by `is_builtin` / `collect_missing_flavors`. - use crate::extension::Kind; - use crate::extension::builtin_entries; - let mut codec_idx_map = HashMap::new(); - for entry in builtin_entries() - .into_iter() - .filter(|e| e.kind == Kind::PlaneCodec) - { - let canonical_id = entry.canonical_id; - let idx = builder.intern(entry); - codec_idx_map.insert(canonical_id, idx); - } - - let entries = builder.finish(); - self.extension_table = ExtensionTable { entries }; - self.chain_registry_bytes = registry_bytes; + let built = build_extension_table(®istry, codecs)?; + self.extension_table = built.table; + self.chain_registry_bytes = built.registry_bytes; self.chain_registry = registry; - self.codec_idx_map = codec_idx_map; + self.codec_idx_map = built.codec_idx_map; Ok(()) } diff --git a/crates/ptwm-core/src/lib.rs b/crates/ptwm-core/src/lib.rs index 41c504c..aa232c3 100644 --- a/crates/ptwm-core/src/lib.rs +++ b/crates/ptwm-core/src/lib.rs @@ -19,6 +19,7 @@ pub mod chain; pub mod codec; pub mod codec_tagged; pub mod codecs; +pub mod compact; pub mod compressor; pub mod container; pub mod delta; diff --git a/crates/ptwm-core/src/transcode/mod.rs b/crates/ptwm-core/src/transcode/mod.rs index a8bd0a7..988eebf 100644 --- a/crates/ptwm-core/src/transcode/mod.rs +++ b/crates/ptwm-core/src/transcode/mod.rs @@ -23,8 +23,11 @@ use std::io::Cursor; use crate::chain::graph::{Chain, ChainEdge, ChainNode, TerminalRef}; use crate::codec::{CodecId, StateSource}; -use crate::container::{ChainRegistry, ContainerReader, ContainerWriter}; +use crate::container::{ + ChainRegistry, CodecSet, ContainerReader, ContainerWriter, build_extension_table, +}; use crate::error::PtwmCoreError; +use crate::extension::ExtensionTable; use crate::header::FLAG_PTWX; use crate::layout::PlaneLayout; use crate::metadata::{decode_shape_metadata, encode_shape_metadata}; @@ -557,7 +560,34 @@ pub fn implode( { let mut writer = ContainerWriter::new(&mut cursor)?; writer.set_prelude(&prelude_entries)?; - writer.set_chain_registry(chain_registry)?; + // Containers written by v1.1 and earlier intern every built-in codec; + // later writers intern only the codecs the planes reference. Rebuild + // whichever layout reproduces the recorded table hash, so the planes' + // `codec_table_idx` values resolve to the same entries. + // + // Take the codec set from the archive-wide registry, not from + // `records`: a caller may implode a subset of the tensors (the + // exploded store layouts rebuild one tensor at a time). + let used_codecs: Vec = registry + .registry + .codecs + .iter() + .filter_map(|c| CodecId::from_u16(c.id)) + .map(crate::compressor::codec_canonical_id) + .collect(); + let trimmed = build_extension_table(&chain_registry, CodecSet::Only(&used_codecs)).ok(); + let trimmed_matches = trimmed.is_some_and(|built| { + built + .table + .to_bytes() + .map(|b| hex::encode(ExtensionTable::hash(&b)) == registry.registry_hash_blake3) + .unwrap_or(false) + }); + if trimmed_matches { + writer.set_chain_registry_with_codecs(chain_registry, &used_codecs)?; + } else { + writer.set_chain_registry(chain_registry)?; + } for record in &records { writer.append_tensor(record)?; } diff --git a/crates/ptwm-core/src/transcode/tests.rs b/crates/ptwm-core/src/transcode/tests.rs index f4ed981..b16a6e5 100644 --- a/crates/ptwm-core/src/transcode/tests.rs +++ b/crates/ptwm-core/src/transcode/tests.rs @@ -286,3 +286,75 @@ fn implode_rejects_missing_member() { let empty: HashMap = HashMap::new(); assert!(implode(&exploded.registry_json, &tensor_jsons(&exploded), empty).is_err()); } + +/// Containers whose Extension Table holds only the referenced codecs must +/// round-trip byte-exactly, and imploding one tensor must keep its codec +/// index pointing at the same codec. +#[test] +fn round_trip_trimmed_codec_table() { + use crate::compressor::codec_canonical_id; + + let registry = ChainRegistry { + chains: vec![(0u16, passthrough_chain())], + }; + let rec = |name: &str, codec: CodecId| TensorRecord { + chain_ref: 0, + dtype_code: 6, + input_format: 2, + flags: 0, + orig_size: 8, + name: name.to_string(), + payload_hash: None, + inline_chain: None, + dependencies: Vec::new(), + tensor_metadata: None, + terminals: vec![flat_plane(PlaneRole::Value, codec, vec![5; 8])], + }; + let codecs = [ + codec_canonical_id(CodecId::Zstd), + codec_canonical_id(CodecId::Rans), + ]; + + let mut cursor = Cursor::new(Vec::::new()); + { + let mut writer = ContainerWriter::new(&mut cursor).unwrap(); + writer + .set_chain_registry_with_codecs(registry, &codecs) + .unwrap(); + let mut a = rec("a.weight", CodecId::Zstd); + a.terminals[0].codec_table_idx = writer.codec_table_idx_for(&codecs[0]).unwrap(); + let mut b = rec("b.weight", CodecId::Rans); + b.terminals[0].codec_table_idx = writer.codec_table_idx_for(&codecs[1]).unwrap(); + writer.append_tensor(&a).unwrap(); + writer.append_tensor(&b).unwrap(); + writer.finalize().unwrap(); + } + let blob = cursor.into_inner(); + + let reader = ContainerReader::open(&blob).unwrap(); + // Two chain ops plus two codecs; no unreferenced built-in codecs. + assert_eq!(reader.extension_table.entries.len(), 4); + assert_round_trip(&blob); + + // Implode only "b.weight" with the archive-wide registry. + let exploded = explode(&blob).unwrap(); + let b = exploded + .tensors + .iter() + .find(|t| t.name == "b.weight") + .unwrap(); + let rebuilt = implode( + &exploded.registry_json, + &[b.manifest_json.clone()], + members_map(&exploded), + ) + .unwrap(); + let reader = ContainerReader::open(&rebuilt).unwrap(); + let rec_bytes = reader.get_tensor_record_bytes("b.weight").unwrap(); + let (parsed, _) = crate::tensor_record::parse_tensor_record(rec_bytes).unwrap(); + let idx = parsed.terminals[0].codec_table_idx as usize; + assert_eq!( + reader.extension_table.entries[idx].canonical_id, codecs[1], + "imploded subset must resolve b.weight's plane to rANS" + ); +} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 7a56c71..a6f3d4c 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -45,6 +45,8 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_blake3_hash, m)?)?; m.add_function(wrap_pyfunction!(py_dtype_element_size, m)?)?; m.add_function(wrap_pyfunction!(compress_model, m)?)?; + m.add_function(wrap_pyfunction!(compress_compact, m)?)?; + m.add_function(wrap_pyfunction!(decode_compact, m)?)?; m.add_function(wrap_pyfunction!(decode_tensor, m)?)?; m.add_function(wrap_pyfunction!(decode_model, m)?)?; m.add_function(wrap_pyfunction!(decode_tensor_shape, m)?)?; @@ -446,6 +448,124 @@ fn compress_model<'py>( Ok(pyo3::types::PyBytes::new(py, &out)) } +/// Compress one tensor into a compact single-tensor frame. +/// +/// Arguments mirror one element of :func:`compress_model` (the tensor name is +/// not stored). Returns ``None`` when the tensor cannot use the compact frame +/// (chunked encoding, or a plane that selects shared state); the caller then +/// writes a full container with :func:`compress_model`. +#[allow(clippy::too_many_arguments)] +#[pyfunction] +#[pyo3(signature = (dtype_code, input_format, raw_bytes, shape, dtype_name, delta_reference_blake3, chains, emit_payload_hash, emit_plane_crc, method_hint=3, codec_menu=None))] +fn compress_compact<'py>( + py: Python<'py>, + dtype_code: u16, + input_format: u8, + raw_bytes: PyBuffer, + shape: Option>, + dtype_name: Option, + delta_reference_blake3: Option>, + chains: Vec>, + emit_payload_hash: bool, + emit_plane_crc: bool, + method_hint: u16, + codec_menu: Option>, +) -> PyResult>> { + let delta_reference_blake3 = match delta_reference_blake3 { + Some(h) => Some(<[u8; 32]>::try_from(h.as_slice()).map_err(|_| { + pyo3::exceptions::PyValueError::new_err(format!( + "delta_reference_blake3 must be 32 bytes, got {}", + h.len() + )) + })?), + None => None, + }; + if chains.is_empty() { + return Err(pyo3::exceptions::PyValueError::new_err( + "chains: at least one candidate chain required", + )); + } + let mut candidate_chains = Vec::with_capacity(chains.len()); + for (chain_idx, blob) in chains.iter().enumerate() { + let (chain, _) = read_chain_blob(blob).map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!("chains[{chain_idx}]: {e}")) + })?; + candidate_chains.push(chain); + } + let raw = contiguous_slice(&raw_bytes, "raw_bytes")?; + let forced_codec = match method_hint { + 2 => Some(ptwm_core::codec::CodecId::Zstd), + 4 => Some(ptwm_core::codec::CodecId::Rans), + 5 => Some(ptwm_core::codec::CodecId::Identity), + _ => None, + }; + let allow_codec_ids: Option> = codec_menu.map(|ids| { + ids.into_iter() + .filter_map(ptwm_core::codec::CodecId::from_u16) + .collect() + }); + let out = py + .allow_threads(|| { + let input = InputTensor { + name: "_".to_string(), + candidate_chains, + dtype_code, + input_format, + raw_bytes: raw, + shape, + dtype_name, + delta_reference_blake3, + }; + let opts = CompressorOptions { + method_hint, + emit_payload_hash, + emit_plane_crc, + forced_codec, + allow_codec_ids, + chunk_size: None, + }; + ptwm_core::compact::compress_compact(&input, &opts) + }) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + Ok(out.map(|b| pyo3::types::PyBytes::new(py, &b))) +} + +/// Decode a compact single-tensor frame. +/// +/// Returns ``(raw_bytes, shape, dtype_name, dtype_code, input_format, +/// delta_reference_blake3, payload_len)``; ``shape``, ``dtype_name`` and +/// ``delta_reference_blake3`` are ``None`` when the frame does not carry them. +/// ``payload_len`` counts codec payload plus inline codec state (the raw size +/// for a stored frame). +#[allow(clippy::type_complexity)] +#[pyfunction] +fn decode_compact<'py>( + py: Python<'py>, + blob: &[u8], +) -> PyResult<( + Bound<'py, pyo3::types::PyBytes>, + Option>, + Option, + u16, + u8, + Option>, + u64, +)> { + let t = py + .allow_threads(|| ptwm_core::compact::decode_compact(blob)) + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + Ok(( + pyo3::types::PyBytes::new(py, &t.raw), + t.shape, + t.dtype_name, + t.dtype_code, + t.input_format, + t.delta_reference_blake3 + .map(|h| pyo3::types::PyBytes::new(py, &h)), + t.payload_len, + )) +} + /// Decode a single tensor from a container blob. /// /// `blob` is the full container ``bytes`` returned by :func:`compress_model`. diff --git a/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index 91bb620..b4c473a 100644 --- a/docs/architecture/wire-format.md +++ b/docs/architecture/wire-format.md @@ -142,6 +142,16 @@ The capabilities are CBOR rather than TOML because the format is parsed in Rust without an extra dependency on a TOML parser path that the reader would otherwise not need. +The writer interns the ops of every chain in the chain registry, +then only the plane codecs that the plane records reference, in the +order of the built-in registry. Writers up to v1.1 interned all 16 +built-in plane codecs, about 2 KB per file, whether or not a plane +used them. Readers resolve `codec_table_idx` by position and check +the table hash against the bytes present, so they decode both +layouts. The table is a pure function of the chain registry and the +codec set; the explode/implode transcoder relies on this to rebuild +the same bytes from the codec list in the registry manifest. + ## Plane records Each tensor's record holds a fixed header plus a sequence of plane @@ -179,6 +189,51 @@ shared across planes within a tensor, a state shared across tensors the writer record exactly one of those four, which the reader uses to find the bytes to feed back into the codec. +## Compact single-tensor frame + +`Compressor.compress` writes one tensor. The container sections +above cost about 500 to 1,000 bytes for one tensor, which is more +than the payload of a small tensor. By default the compressor writes +a compact frame instead, with its own magic `\x89PTWS`: + +``` +magic [u8;5] \x89PTWS +flags u8 payload hash, plane CRC, shape, dtype name, + delta reference, stored; unknown bits rejected +dtype_code varint +input_format u8 +orig_size varint +shape varint ndim + varint dims (if flagged) +dtype_name varint length + utf8 (if flagged) +delta_ref [u8;32] blake3 of the reference (if flagged) +payload_hash u64 xxh64 of the raw bytes (if flagged) +body +``` + +The body is either the raw bytes (the stored flag, written when no +codec makes the tensor smaller) or a coded tensor: a list of op wire +ids, the chain in the wire form the chain registry uses (its op +indices refer to that list), and one record per terminal plane +(codec id, inline state, payload length, optional CRC32, payload). +The frame does not store plane layouts, plane lengths or the plane +count, because the decoder derives them from the chain. The dtype +name is stored only when it differs from the canonical name of +`dtype_code`. + +The fixed cost is 10 to 130 bytes, depending on the chain and on +which integrity fields are on. A tensor whose planes need shared +state (a prelude) or chunked encoding cannot use the frame; the +compressor writes a full container for it. The frame holds no +extension table, so it can reference built-in ops and codecs only. + +The frame follows the same versioning rule as the container: no +numeric version field, and a new magic for an incompatible layout. +Readers that predate the frame reject `\x89PTWS` as "not a .ptwm +container" rather than misparse it. `CompressionConfig(framing= +Framing.CONTAINER)` keeps the full container for consumers pinned to +such readers, and the CLI writes `.ptwm` files as full containers so +the random-access and store tools can open them. + ## Wire-format invariants The reader assumes the following without checking, since checks are diff --git a/docs/site/content/30.guides/02.python-api.md b/docs/site/content/30.guides/02.python-api.md index 80166af..750c7d3 100644 --- a/docs/site/content/30.guides/02.python-api.md +++ b/docs/site/content/30.guides/02.python-api.md @@ -65,6 +65,9 @@ assert torch.equal(tensor, restored) | `threads` | `None` (up to 16 cores) | Maximum threads for parallel encode | | `compression_threshold` | `0.95` | Skip compression if ratio exceeds this threshold | | `compression_chunk` | `262144` | Chunk size in bytes (256 KB) | +| `framing` | `Framing.COMPACT` | `COMPACT` writes a compact single-tensor frame (10 to 130 B of fixed overhead; raw bytes stored when no codec helps). `CONTAINER` writes a full `.ptwm` container, which readers before v1.2 and the random-access tools require. | +| `emit_payload_hash` | `True` | Store an xxh64 of the raw bytes (8 B per tensor), checked on decode | +| `emit_plane_crc` | `True` | Store a CRC32 per coded plane (4 B per plane), checked on decode | ## Full reference diff --git a/fuzz/Cargo.lock b/fuzz/Cargo.lock index 6b28fa7..b6f9168 100644 --- a/fuzz/Cargo.lock +++ b/fuzz/Cargo.lock @@ -1922,7 +1922,7 @@ dependencies = [ [[package]] name = "ptwm-core" -version = "1.0.0" +version = "1.1.0" dependencies = [ "blake3", "ciborium", diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 0419f04..2a0a9bd 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -95,3 +95,9 @@ name = "raw_binary_schema" path = "fuzz_targets/raw_binary_schema.rs" test = false doc = false + +[[bin]] +name = "compact_decode" +path = "fuzz_targets/compact_decode.rs" +test = false +doc = false diff --git a/fuzz/fuzz_targets/compact_decode.rs b/fuzz/fuzz_targets/compact_decode.rs new file mode 100644 index 0000000..4524138 --- /dev/null +++ b/fuzz/fuzz_targets/compact_decode.rs @@ -0,0 +1,42 @@ +#![no_main] + +// `Decompressor.decompress` hands any blob that starts with the compact-frame +// magic to `decode_compact`. The parser must never panic or infinite-loop on +// adversarial input. + +use libfuzzer_sys::fuzz_target; +use ptwm_core::compact::{MAGIC_COMPACT, decode_compact}; + +/// Read one LEB128 varint at `*pos`; `None` when truncated or too long. +fn varint(data: &[u8], pos: &mut usize) -> Option { + let mut v = 0u64; + for i in 0..10 { + let b = *data.get(*pos)?; + *pos += 1; + v |= u64::from(b & 0x7F) << (7 * i); + if b & 0x80 == 0 { + return Some(v); + } + } + None +} + +fuzz_target!(|data: &[u8]| { + // The frame declares its decoded size, and a decoder must be able to + // allocate it: a 100-byte frame of a constant tensor can declare 1 GB. + // Skip declared sizes above 16 MiB so libFuzzer's RSS limit flags real + // over-allocation only. Layout after the magic: flags u8, dtype varint, + // input_format u8, orig_size varint. + let mut pos = 1; + if varint(data, &mut pos).is_some() { + pos += 1; + if varint(data, &mut pos).is_some_and(|n| n > 1 << 24) { + return; + } + } + // Prefix the magic so the fuzzer spends its time past the magic check. + let mut buf = Vec::with_capacity(MAGIC_COMPACT.len() + data.len()); + buf.extend_from_slice(&MAGIC_COMPACT); + buf.extend_from_slice(data); + let _ = decode_compact(&buf); +}); diff --git a/python/ptwm/__init__.py b/python/ptwm/__init__.py index 19f343e..649b35a 100644 --- a/python/ptwm/__init__.py +++ b/python/ptwm/__init__.py @@ -7,7 +7,7 @@ Public entry points: :class:`Compressor`, :class:`Decompressor`, :class:`CompressionConfig`, :class:`DecompressionConfig`, -:class:`Method`, :class:`Format`. Third-party integrations +:class:`Method`, :class:`Format`, :class:`Framing`. Third-party integrations (HuggingFace `transformers`, `safetensors`) live in :mod:`ptwm.integrations` and are imported from there explicitly, not re-exported here. """ @@ -18,6 +18,7 @@ CompressionConfig, DecompressionConfig, Format, + Framing, Lossy, Method, ) @@ -42,6 +43,7 @@ "DecompressionConfig", "Decompressor", "Format", + "Framing", "Lossy", "Method", "Error", diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 3b824be..0c700ee 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -42,6 +42,30 @@ def _missing_(cls, value): return None +class Framing(Enum): + r"""Byte layout of a single-tensor :meth:`Compressor.compress` output. + + ``COMPACT`` writes a compact single-tensor frame (magic ``\x89PTWS``) + with a fixed cost of about 10 to 130 bytes, and stores the raw bytes when + no codec makes them smaller. It falls back to ``CONTAINER`` when the + tensor cannot use the frame (streaming, or a plane that selects shared + codec state). ``CONTAINER`` writes a full ``.ptwm`` container (magic + ``\x89PTWM``), which readers before v1.2 require, and which the + random-access, store and transcode tools open. + """ + + COMPACT = 1 + CONTAINER = 2 + + @classmethod + def _missing_(cls, value): + if isinstance(value, str): + value = value.upper() + if value in cls.__members__: + return cls.__members__[value] + return None + + class Lossy(Enum): NONE = 0 INTEGER = 1 @@ -87,6 +111,12 @@ class CompressionConfig: """ codec: str | None = None device: int | None = None + framing: Framing = Framing.COMPACT + """Output layout; see :class:`Framing`.""" + emit_payload_hash: bool = True + """Store an xxh64 of the raw bytes (8 B per tensor), checked on decode.""" + emit_plane_crc: bool = True + """Store a CRC32 of every coded plane (4 B per plane), checked on decode.""" @classmethod def from_resolved_policy( diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index c4d8820..3c7428d 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -16,7 +16,9 @@ "blake3_hash", "builtin_canonical_id", "codec_id_registry", + "compress_compact", "compress_model", + "decode_compact", "decode_model", "decode_tensor", "decode_tensor_delta_hash", @@ -216,6 +218,67 @@ def compress_model( ) +@translate_errors +def compress_compact( + dtype_code: int, + input_format: int, + raw_bytes: bytes, + shape: list[int] | None, + dtype_name: str | None, + delta_reference_blake3: bytes | None, + chains: list[bytes], + emit_payload_hash: bool, + emit_plane_crc: bool, + method_hint: int = 3, + codec_menu: list[int] | None = None, +) -> bytes | None: + """Compress one tensor into a compact single-tensor frame. + + The arguments match one record of :func:`compress_model`; the tensor name + is not stored. Returns ``None`` when the tensor cannot use the compact + frame (a plane selects shared state); write a full container then. + """ + got = _ext.compress_compact( + dtype_code, + input_format, + raw_bytes, + shape, + dtype_name, + delta_reference_blake3, + chains, + emit_payload_hash, + emit_plane_crc, + method_hint, + codec_menu, + ) + return bytes(got) if got is not None else None + + +@translate_errors +def decode_compact( + blob: bytes, +) -> tuple[bytes, list[int] | None, str | None, int, int, bytes | None, int]: + """Decode a compact single-tensor frame. + + Returns ``(raw_bytes, shape, dtype_name, dtype_code, input_format, + delta_reference_blake3, payload_len)``. ``payload_len`` counts codec + payload plus inline codec state (the raw size for a stored frame), so + ``len(blob) - payload_len`` is the frame's fixed overhead. + """ + raw, shape, dtype_name, dtype_code, input_format, delta, payload_len = ( + _ext.decode_compact(blob) + ) + return ( + bytes(raw), + shape, + dtype_name, + dtype_code, + input_format, + bytes(delta) if delta is not None else None, + payload_len, + ) + + @translate_errors def decode_tensor(blob: bytes, name: str, skip_missing: bool = False) -> bytes: """Decode a single tensor from a ``.ptwm`` container blob. diff --git a/python/ptwm/ablation/_bench.py b/python/ptwm/ablation/_bench.py index 633ad45..db7e76f 100644 --- a/python/ptwm/ablation/_bench.py +++ b/python/ptwm/ablation/_bench.py @@ -12,7 +12,7 @@ from dataclasses import dataclass, field from pathlib import Path -from ptwm import CompressionConfig, Compressor, Format +from ptwm import CompressionConfig, Compressor, Format, Framing from ptwm.ablation._resolve import resolve_active @@ -91,7 +91,9 @@ def bench_compress( # rather than silently using the unconstrained menu. config = CompressionConfig.from_resolved_policy( resolved, - base=CompressionConfig(input_format=Format.BYTE), + # Full containers: the compact frame's stored-raw fallback would + # hide which codec the policy selected. + base=CompressionConfig(input_format=Format.BYTE, framing=Framing.CONTAINER), ) for trial in range(trials): t0 = time.perf_counter_ns() diff --git a/python/ptwm/ablation/_replay.py b/python/ptwm/ablation/_replay.py index a788947..96f02c0 100644 --- a/python/ptwm/ablation/_replay.py +++ b/python/ptwm/ablation/_replay.py @@ -15,7 +15,7 @@ from dataclasses import dataclass from pathlib import Path -from ptwm import CompressionConfig, Compressor, Format +from ptwm import CompressionConfig, Compressor, Format, Framing @dataclass @@ -70,9 +70,9 @@ def bench_replay( out_file = out_dir / f"replay-{r['variant']}-trial{r['trial']}.ptwm" synthetic_input = b"\x00" * int(r["input_bytes"]) - blob = Compressor(CompressionConfig(input_format=Format.BYTE)).compress( - synthetic_input - ) + blob = Compressor( + CompressionConfig(input_format=Format.BYTE, framing=Framing.CONTAINER) + ).compress(synthetic_input) out_file.write_bytes(blob) return ReplayResult( diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index c40daf5..f8881b3 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -138,6 +138,7 @@ def compress_file( Compressor, DecompressionConfig, Decompressor, + Framing, Method, ) @@ -162,6 +163,9 @@ def compress_file( compressor = Compressor( CompressionConfig( bytearray_dtype=dtype, + # `.ptwm` files stay full containers so the random-access, store + # and inspect tools can open them. + framing=Framing.CONTAINER, is_streaming=is_streaming, streaming_chunk=chunk_size, method=Method(method), @@ -253,6 +257,7 @@ def compress_file_delta( Compressor, DecompressionConfig, Decompressor, + Framing, Method, ) @@ -285,6 +290,7 @@ def compress_file_delta( compressor = Compressor( CompressionConfig( bytearray_dtype="float32" if dtype else "bfloat16", + framing=Framing.CONTAINER, is_streaming=is_streaming, streaming_chunk=chunk_size, delta_compressed_type="file", diff --git a/python/ptwm/core/_compressor.py b/python/ptwm/core/_compressor.py index 51915cf..17d5ddd 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -4,7 +4,7 @@ import torch from .. import _rust -from .._config import CompressionConfig, Format, Method +from .._config import CompressionConfig, Format, Framing, Method from .._exceptions import CompressionMethodNotSupportedError from ..delta import encode as delta_encode from ..preprocessing._chains import ( @@ -174,7 +174,10 @@ def _compress_bin( dtype_code: int, delta_ref_blake3: bytes | None = None, ) -> bytes: - """Encode a single-tensor ``.ptwm`` via ``compress_model``.""" + """Encode one tensor as a compact frame or a single-tensor container. + + ``config.framing`` selects the layout; see :class:`ptwm.Framing`. + """ method_hint = _METHOD_HINTS.get(self.config.method.value) if method_hint is None: msg = f"Unsupported compression method: {self.config.method}" @@ -205,6 +208,28 @@ def _compress_bin( else: codec_menu_ids = [int(c) for c in self.config.codec_menu] + emit_payload_hash = self.config.emit_payload_hash + emit_plane_crc = self.config.emit_plane_crc + + # The compact frame has no chunk table, so streaming always writes a + # full container. + if self.config.framing == Framing.COMPACT and streaming_chunk is None: + blob = _rust.compress_compact( + int(dtype_code), + int(self.config.input_format.value), + ba, + tensor_shape, + dtype_name_str, + delta_ref_blake3, + chains[0], + emit_payload_hash, + emit_plane_crc, + method_hint, + codec_menu_ids, + ) + if blob is not None: + return blob + return _rust.compress_model( [ ( @@ -218,8 +243,8 @@ def _compress_bin( ) ], chains, - True, # emit_payload_hash - True, # emit_plane_crc + emit_payload_hash, + emit_plane_crc, method_hint, streaming_chunk, codec_menu_ids, diff --git a/python/ptwm/core/_decompressor.py b/python/ptwm/core/_decompressor.py index 474312f..cade22f 100644 --- a/python/ptwm/core/_decompressor.py +++ b/python/ptwm/core/_decompressor.py @@ -10,6 +10,10 @@ from ..delta import decode as delta_decode from ..utils import DType +# Full multi-tensor container and compact single-tensor frame. +_MAGIC_CONTAINER = b"\x89PTWM" +_MAGIC_COMPACT = b"\x89PTWS" + def _torch_dtype_for_code(dtype_code: int) -> torch.dtype | None: """Map a wire ``dtype_code`` to a concrete ``torch.dtype`` if supported.""" @@ -35,9 +39,11 @@ def __init__(self, config: DecompressionConfig | None = None) -> None: self._threads = self.config.threads or min(multiprocessing.cpu_count(), 16) def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Tensor: - """Decompress a ``.ptwm`` blob. + r"""Decompress a ``.ptwm`` blob. - The single-tensor wrapper uses name ``"_"``. Raw decompressed bytes + Accepts a compact single-tensor frame (magic ``\x89PTWS``) or a + full container (magic ``\x89PTWM``). The single-tensor container + uses name ``"_"``. Raw decompressed bytes come from ``decode_tensor``; shape/dtype/input_format (when present) come from the CBOR shape-metadata map. Format recovery: @@ -73,18 +79,30 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten raise ValueError(msg) mv_data = memoryview(data) - magic = b"\x89PTWM" - if len(mv_data) < len(magic): + magic = bytes(mv_data[: len(_MAGIC_CONTAINER)]) + if len(magic) < len(_MAGIC_CONTAINER): msg = "too short to be a .ptwm container" raise HeaderParseError(msg) - if bytes(mv_data[: len(magic)]) != magic: - msg = "not a .ptwm container" - raise HeaderParseError(msg) blob = bytes(mv_data) - skip = self.config.skip_missing - raw = _rust.decode_tensor(blob, "_", skip) - expected_hash = _rust.decode_tensor_delta_hash(blob, "_", skip) + info: tuple[list[int], str, int, int] | None + if magic == _MAGIC_COMPACT: + raw, shape, dtype_name, dtype_code, input_format, expected_hash, _ = ( + _rust.decode_compact(blob) + ) + info = ( + (shape, dtype_name or "", dtype_code, input_format) + if shape is not None + else None + ) + elif magic == _MAGIC_CONTAINER: + skip = self.config.skip_missing + raw = _rust.decode_tensor(blob, "_", skip) + expected_hash = _rust.decode_tensor_delta_hash(blob, "_", skip) + info = _rust.decode_tensor_info(blob, "_", skip) + else: + msg = "not a .ptwm container" + raise HeaderParseError(msg) if expected_hash is not None: if self.config.delta_second_data is None: @@ -109,7 +127,6 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten ) raise ValueError(msg) - info = _rust.decode_tensor_info(blob, "_", skip) if info is None: return raw shape, _dtype_name, dtype_code, input_format = info diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py new file mode 100644 index 0000000..5ae6a3f --- /dev/null +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -0,0 +1,72 @@ +"""Size limits for single-tensor ``Compressor.compress`` output. + +The fixed bytes per output (header, tables, framing) used to be about 2.7 KB, +so small tensors came out larger than their input. These tests hold the +default (compact) framing to the limits below and keep the full container's +Extension Table trimmed. The sweep matches ``benchmarks/container_overhead``. +""" + +from __future__ import annotations + +import pytest +import torch +from ptwm._config import Framing + +from benchmarks.container_overhead import ( + DTYPES, + SIZES, + make_tensor, + measure, + tiled_rows, +) + +# Largest allowed excess over the raw bytes, and over zstd-19 when the ptwm +# codec payload is smaller than zstd-19's output. +SLACK = 64 + + +def test_constant_64_byte_bf16_tensor_is_at_most_128_bytes() -> None: + t = torch.full((32,), 0.5, dtype=torch.bfloat16) + row = measure(t, Framing.COMPACT, "bf16", "constant") + assert row.out <= 128, f"constant 64 B bf16 tensor -> {row.out} B" + + +@pytest.mark.parametrize(("dname", "dtype"), DTYPES) +@pytest.mark.parametrize("nbytes", SIZES) +def test_random_normal_overhead_is_bounded( + dname: str, dtype: torch.dtype, nbytes: int +) -> None: + row = measure( + make_tensor(dtype, "normal", nbytes), Framing.COMPACT, dname, "normal" + ) + assert row.out <= row.raw + SLACK, row + if row.payload < row.zstd19: + assert row.out <= row.zstd19 + SLACK, row + + +@pytest.mark.parametrize(("dname", "dtype"), DTYPES) +@pytest.mark.parametrize("nbytes", SIZES) +def test_constant_overhead_is_bounded( + dname: str, dtype: torch.dtype, nbytes: int +) -> None: + row = measure( + make_tensor(dtype, "constant", nbytes), Framing.COMPACT, dname, "constant" + ) + assert row.out <= row.raw + SLACK, row + # Four-plane fp32 frames carry about 120 B of chain and plane metadata. + assert row.out <= 256, row + + +def test_tiled_rows_ratio_is_at_least_60() -> None: + row = measure(tiled_rows(64), Framing.COMPACT, "bf16", "tiled") + assert row.ratio >= 60, f"(64, 1024) tiled bf16 ratio {row.ratio:.1f}" + + +@pytest.mark.parametrize(("dname", "dtype"), DTYPES) +def test_container_framing_floor_stays_trimmed(dname: str, dtype: torch.dtype) -> None: + # v1.1 wrote about 2.5 KB of non-payload bytes per single-tensor + # container, 2.1 KB of it unreferenced codec entries. + row = measure( + make_tensor(dtype, "normal", 4096), Framing.CONTAINER, dname, "normal" + ) + assert row.non_payload < 1100, row diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py new file mode 100644 index 0000000..6f6e655 --- /dev/null +++ b/tests/ptwm/core/test_compact_frame.py @@ -0,0 +1,165 @@ +"""Compact single-tensor frame: selection, fallbacks and round trips.""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch +from ptwm import ( + CompressionConfig, + Compressor, + DecompressionConfig, + Decompressor, + Format, + Framing, + _rust, +) +from ptwm._exceptions import Error, HeaderParseError + +COMPACT = b"\x89PTWS" +CONTAINER = b"\x89PTWM" + + +def _torch_cfg(**kw) -> CompressionConfig: + return CompressionConfig(input_format=Format.TORCH, **kw) + + +def test_default_framing_is_compact() -> None: + t = torch.randn(16, 16).to(torch.bfloat16) + blob = Compressor(_torch_cfg()).compress(t) + assert blob[:5] == COMPACT + back = Decompressor().decompress(blob) + assert back.dtype == t.dtype + assert back.shape == t.shape + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + + +def test_framing_accepts_names() -> None: + assert Framing("container") is Framing.CONTAINER + assert Framing("COMPACT") is Framing.COMPACT + + +def test_container_framing_still_decodes() -> None: + t = torch.randn(8, 32, dtype=torch.float32) + blob = Compressor(_torch_cfg(framing=Framing.CONTAINER)).compress(t) + assert blob[:5] == CONTAINER + assert torch.equal(Decompressor().decompress(blob), t) + + +def test_streaming_falls_back_to_container() -> None: + t = torch.randn(4096).to(torch.bfloat16) + blob = Compressor(_torch_cfg(is_streaming=True, streaming_chunk=1024)).compress(t) + assert blob[:5] == CONTAINER + back = Decompressor().decompress(blob) + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + + +@pytest.mark.parametrize( + ("dtype", "shape"), + [ + (torch.bfloat16, (3, 5, 7)), + (torch.float16, (64,)), + (torch.float32, (2, 128)), + (torch.float8_e4m3fn, (256,)), + (torch.float8_e5m2, (16, 16)), + (torch.float32, ()), + (torch.bfloat16, (0, 16)), + ], +) +def test_torch_round_trip_keeps_shape_and_dtype(dtype: torch.dtype, shape) -> None: + g = torch.Generator().manual_seed(3) + t = torch.randn(shape, generator=g).to(dtype) + blob = Compressor(_torch_cfg()).compress(t) + assert blob[:5] == COMPACT + back = Decompressor().decompress(blob) + assert back.dtype == t.dtype + assert back.shape == t.shape + flat = back.reshape(-1).view(torch.uint8) + assert torch.equal(flat, t.reshape(-1).view(torch.uint8)) + + +def test_numpy_round_trip() -> None: + a = np.random.default_rng(0).standard_normal((12, 9)).astype(np.float32) + blob = Compressor(CompressionConfig(input_format=Format.NUMPY)).compress(a) + assert blob[:5] == COMPACT + back = Decompressor().decompress(blob) + assert isinstance(back, np.ndarray) + assert back.shape == a.shape + assert np.array_equal(back, a) + + +def test_byte_round_trip() -> None: + data = bytes(range(256)) * 9 + blob = Compressor(CompressionConfig(bytearray_dtype="float32")).compress(data) + assert blob[:5] == COMPACT + assert Decompressor().decompress(blob) == data + + +def test_incompressible_input_is_stored_with_small_overhead() -> None: + data = np.random.default_rng(1).integers(0, 256, 4000, dtype=np.uint8).tobytes() + blob = Compressor(CompressionConfig(bytearray_dtype="float32")).compress(data) + assert len(blob) <= len(data) + 32 + assert Decompressor().decompress(blob) == data + + +@pytest.mark.parametrize("framing", [Framing.COMPACT, Framing.CONTAINER]) +def test_integrity_fields_can_be_disabled(framing: Framing) -> None: + t = torch.randn(512).to(torch.bfloat16) + full = Compressor(_torch_cfg(framing=framing)).compress(t) + bare = Compressor( + _torch_cfg(framing=framing, emit_payload_hash=False, emit_plane_crc=False) + ).compress(t) + # 8 B payload hash, plus 4 B CRC per coded plane (bf16 has two). + assert len(full) - len(bare) == 8 + 2 * 4 + back = Decompressor().decompress(bare) + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + + +def test_delta_round_trip() -> None: + base = torch.randn(1024).to(torch.bfloat16) + new = base.clone() + new[::17] += 1 + ref = base.view(torch.uint8).numpy().tobytes() + cfg = CompressionConfig(bytearray_dtype="bfloat16", delta_compressed_type="file") + data = new.view(torch.uint8).numpy().tobytes() + blob = Compressor(cfg).compress(data, delta_second_data=ref) + assert blob[:5] == COMPACT + out = Decompressor(DecompressionConfig(delta_second_data=ref)).decompress(blob) + assert out == data + with pytest.raises(ValueError, match="delta_second_data is None"): + Decompressor().decompress(blob) + with pytest.raises(ValueError, match="BLAKE3"): + Decompressor(DecompressionConfig(delta_second_data=bytes(len(ref)))).decompress( + blob + ) + + +def test_corrupted_frame_is_rejected() -> None: + t = torch.randn(64, 64).to(torch.bfloat16) + blob = bytearray(Compressor(_torch_cfg()).compress(t)) + blob[len(blob) // 2] ^= 0xFF + with pytest.raises(Error): + Decompressor().decompress(bytes(blob)) + with pytest.raises(Error): + Decompressor().decompress(bytes(blob[:-3])) + + +def test_unknown_magic_is_rejected() -> None: + with pytest.raises(HeaderParseError): + Decompressor().decompress(b"\x89PTWQ" + bytes(20)) + with pytest.raises(HeaderParseError): + Decompressor().decompress(b"\x89PT") + + +def test_decode_compact_reports_metadata() -> None: + t = torch.full((4, 8), 0.25, dtype=torch.bfloat16) + blob = Compressor(_torch_cfg()).compress(t) + raw, shape, dtype_name, _code, input_format, delta, payload_len = ( + _rust.decode_compact(blob) + ) + assert raw == t.view(torch.uint8).numpy().tobytes() + assert shape == [4, 8] + assert dtype_name == "bfloat16" + assert input_format == Format.TORCH.value + assert delta is None + assert 0 < payload_len < len(blob) diff --git a/tests/ptwm/core/test_zstd.py b/tests/ptwm/core/test_zstd.py index f0d5dbc..3579beb 100644 --- a/tests/ptwm/core/test_zstd.py +++ b/tests/ptwm/core/test_zstd.py @@ -17,6 +17,7 @@ Compressor, Decompressor, Format, + Framing, Method, ) @@ -72,7 +73,9 @@ def test_zstd_numpy_roundtrip() -> None: def test_zstd_outputs_ptwm() -> None: """ZSTD produces a .ptwm container with a forced Zstd plane codec.""" arr = np.zeros(256, dtype=np.float32).tobytes() - config = CompressionConfig(method=Method.ZSTD, bytearray_dtype="float32") + config = CompressionConfig( + method=Method.ZSTD, bytearray_dtype="float32", framing=Framing.CONTAINER + ) compressed = Compressor(config).compress(arr) assert compressed[:5] == b"\x89PTWM" diff --git a/tests/ptwm/wire_format/test_extension_table.py b/tests/ptwm/wire_format/test_extension_table.py index fdde250..ece0c84 100644 --- a/tests/ptwm/wire_format/test_extension_table.py +++ b/tests/ptwm/wire_format/test_extension_table.py @@ -1,13 +1,17 @@ """Verify the v1 file carries an Extension Table immediately after the header.""" import torch -from ptwm import CompressionConfig, Compressor, Format, Method +from ptwm import CompressionConfig, Compressor, Format, Framing, Method def test_extension_table_is_present_for_huffman_chain(): tensor = torch.randn(128, 128, dtype=torch.float32) blob = Compressor( - CompressionConfig(input_format=Format.TORCH, method=Method.HUFFMAN) + CompressionConfig( + input_format=Format.TORCH, + method=Method.HUFFMAN, + framing=Framing.CONTAINER, + ) ).compress(tensor) # Header layout (little-endian): @@ -32,3 +36,16 @@ def test_extension_table_is_present_for_huffman_chain(): assert count >= 1, ( "expected at least one Extension Table entry referenced by a Huffman chain" ) + + +def test_extension_table_holds_only_referenced_codecs(): + """The writer interns the chain ops plus the codecs the planes use. The + v1.1 layout interned all 16 built-in plane codecs, about 2 KB.""" + tensor = torch.full((32,), 0.5, dtype=torch.bfloat16) + blob = Compressor( + CompressionConfig(input_format=Format.TORCH, framing=Framing.CONTAINER) + ).compress(tensor) + table_off = int.from_bytes(blob[11:19], "little") + count = int.from_bytes(blob[table_off : table_off + 4], "little") + assert count < 16, f"extension table has {count} entries" + assert len(blob) < 1024, f"single-tensor container is {len(blob)} bytes" diff --git a/tests/ptwm/wire_format/test_v1_header.py b/tests/ptwm/wire_format/test_v1_header.py index 9afc20f..3a24bd0 100644 --- a/tests/ptwm/wire_format/test_v1_header.py +++ b/tests/ptwm/wire_format/test_v1_header.py @@ -1,12 +1,14 @@ """Round-trip the v1 wire format from Python.""" import torch -from ptwm import CompressionConfig, Compressor, Decompressor, Format +from ptwm import CompressionConfig, Compressor, Decompressor, Format, Framing def test_compress_decompress_round_trip_v1_format(): tensor = torch.randn(64, 64, dtype=torch.bfloat16) - compressor = Compressor(CompressionConfig(input_format=Format.TORCH)) + compressor = Compressor( + CompressionConfig(input_format=Format.TORCH, framing=Framing.CONTAINER) + ) decompressor = Decompressor() blob = compressor.compress(tensor) From d08e50aa8b71dfbc506ee7d7014e652c2f4b3d12 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 20:42:41 +0000 Subject: [PATCH 02/34] test(core): add golden legacy .ptwm fixtures Pin the \x89PTWM wire format before the format work for issue #5. - tests/fixtures/legacy_v1_1_0.ptwm: multi-tensor container with the v1.1.0 Extension Table layout (all 16 built-in codecs, written with ContainerWriter::set_chain_registry). Holds bf16, fp16, fp32, fp8 and int8 tensors, two scale tensors that share an Order1ScaleAC prelude state, a chunked plane, an inline chain and a dependency. - legacy_v1_0_0.ptwm: the same file with built-in ids and labels re-derived for version 1.0.0 and the table hash fixed up. Its decode and transcode tests are ignored until cross-version built-ins land. - legacy_trimmed.ptwm: output of the current trimmed-table writer. - legacy_p4_byte_fp32.ptwm: fp32 from BYTE input whose Source declares four times the real size. tests/legacy_fixtures.rs checks that each fixture decodes to the expected (length, xxh64) values and that implode(explode(x)) == x. It also pins the blake3 hash of compress_model output for a fixed input, so later changes to legacy bytes are caught. An ignored test regenerates the files (see the module docs). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- .../tests/fixtures/legacy_p4_byte_fp32.ptwm | Bin 0 -> 1444 bytes .../tests/fixtures/legacy_trimmed.ptwm | Bin 0 -> 4678 bytes .../tests/fixtures/legacy_v1_0_0.ptwm | Bin 0 -> 25081 bytes .../tests/fixtures/legacy_v1_1_0.ptwm | Bin 0 -> 25081 bytes crates/ptwm-core/tests/legacy_fixtures.rs | 911 ++++++++++++++++++ 5 files changed, 911 insertions(+) create mode 100644 crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32.ptwm create mode 100644 crates/ptwm-core/tests/fixtures/legacy_trimmed.ptwm create mode 100644 crates/ptwm-core/tests/fixtures/legacy_v1_0_0.ptwm create mode 100644 crates/ptwm-core/tests/fixtures/legacy_v1_1_0.ptwm create mode 100644 crates/ptwm-core/tests/legacy_fixtures.rs diff --git a/crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32.ptwm b/crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32.ptwm new file mode 100644 index 0000000000000000000000000000000000000000..c50d05c759dc996ed392fa5967a4e58c54c420dd GIT binary patch literal 1444 zcmeAS2nqM)<&xrJV2EabfG#Eot;TX=+p=vxSj4T>uW>c*yM4WGrRcZ1vwOI9D#b6I zaznX>6{1c7LNPpr>T`nfLDJQMsy}yVi@*QpsHfy}i9;z>Mc{{egT&V)(N2Gpk8|v0 z7&7zq3QEdz^^!_6b4oJv^osLKi;`0v4D}543>X=JkVz3}&Z6v;)RNSq+|0bp;@q!! z8Hq(H<%vb9@yR)f#l^|V1*MEoOHruZHt}i4u5l?YEzeZD*PO6o@9CJ;B71A^Z+W{< z?rX;%c?}HPlQK)Pw3llRlNS=`es0^CMc^C>z^pb#4q?cA;Y;;iX^7gp(kHcBnKt{qq1C%jw zoqRPzxG_UI?^|PsV5#ZswP%Aj`_{&?zg?cLupoVN)!BlC&*`7l&;Dp#9A0cJyjP0% z`K#x~zb?OSzuf+up;vZpdKH8Fa`j)!rK)eux^RkN?q>H3ya{L3pEES@J~s$XcX%t@ z?%ST;V3p9^d)xTH^5wU9kEf@9PG8j!wK@Il0e0Wdxw+V#Hfz$iREFe1T)*@AlZ6J?xV1$OEqrvD`${H%g4RMKq1P%po0-3K&t2^>gNL?^Se)^6av|Dc2XD*71IC>dVh6 zl6RkcJQ{vV?8laXo>Muep4Epp{KM|Fd>ad0h6A25i<+wwc*0gQw%yG%6On0LzqhMX z=ZOA_318o{7Tj)JqAAF5;QKd$Ef@E&r9Ido(!OHyWYa*2y}S9=Ze*|vz3RG7XKCKHk!HYRR zYwq}F2nqEM0b5}%+FhgtL9zG{P1~U;`xRmCWeklwM+?Yurg~)>BD1r Wj5Ac!p=l=+n#ekVMYF5Fn-l<#&DWX$ literal 0 HcmV?d00001 diff --git a/crates/ptwm-core/tests/fixtures/legacy_trimmed.ptwm b/crates/ptwm-core/tests/fixtures/legacy_trimmed.ptwm new file mode 100644 index 0000000000000000000000000000000000000000..7f63ad33bbee1971b48097069c7aa47356691719 GIT binary patch literal 4678 zcmb_gc|6ox8$Ywy*X)FbFrygRE?rSXMcLa>O*7_)(JV4EN!|FRY$+rpl|5QTg%C;Q z63J4!)r~8J(!!M`j7=Xsvb#JX+V;w&YuCJulP z;Fkg+{x~Xus&e~$J=n0AIdH<>C~((sRHD&X?xCUK@~s6{B+dK<{C#NxI^uNT9W~Dlyv4}yO5C)OLrO`PwCXvnJ z2Kz#*jETm?Whj7Dv=(ws&Mqp%fr1${CXLPb#0($@Q$xtX5XqNLX0v^LgSY~oe9qr` zO8?{$6LBc>^FxA;U%0J)kz_;sPf}6a2iEUJOIP?+6XtNI&^V-Eh!sqQf=M(8f{cHd zA4#TeqW(A~<^nTRD;BR@W@9oo;p&6OmS1CiK4bQ))VB&b|L z@~i6**E$zM)9XC6^L$m~wo3|C2fA*&tJeG!Grj2Fbzy2UKa%7Yb&I9-xp$u>oTxvp zei&6wU2(r*r+usW4{4U^uY5i)nX8MxpyC%~NrKE6rt>39QKqsVNsgS^UK6JHV+Sdu z)+lS;Vd9Z(%fs>*sUcqZ+H+)K(kK-AydsmV$QUp@Yj`dt`?Ps|$;0ucktun3`4KqD zHed6inoPRh92qD)+XzA|pA!(w;`#^7k1nPS-^f8d%QtMP?XrrBV;elkejmN-ad9)g zJ>jqcS*}ELjvO>9#N^O8p?v=6|MRtyZ8w#iZ}?Cet9805w{bjjhkiTSM6{XQRFE6s zs1t9hF^7Ku*UyhZX3h_LC`t;2fB33U!LwvQTZyHJzb@lu_ufe^rrXzZIyz}oZ4Uc= zYz}pPc=tv{I`m)Dj-Yu3+bbmZ3(+p&y|$~G9ak9eklUX6JZ~;<4w=cGAKpA?$@BOa zG+Gfz15souf#&r%q{$$d0t8qz;7xcPGUKfv6^+865GpK>N)B0+Kx2^}hcpaQc~r>v z->I-@6dHv^BL$1bV9-3S3Uj#10)+Bw6#AYV3lLfq2I)~!K;-KV5nhY{!bnAmU^faS ziv9jN-vs-Y>+lXj^TOEIO9G$_7%VD8C(_BGP%xXgjJTT^0@3^fIAB(lkw@_|tXD(0 z=0NNKau5_Ds)B>4oX{XBfZ|7IkvYac2q8y_@C4(Rui_zxH*S9o&H%{4T-fkNRpg7i zX+rIm_}=K43E5mT5eLy24at?O!eQ!oqt#g_!;+E+O`EhOKeZVI8Iqs>B=Q1$2AM%^y6GAD@i8r)>m0XrB=6A1;E6~ z>b!vy55BxV;dV{#VnRi;M@OCBn%kzUjTSvCiiFLC80QF=R`=mmCJ*ASr8<@!JCvN% zG8#9y_(KOSgw)qDE)uIX-Dtcvl{WQg9$3W>G%4~%uH_`sr4?% z@N>Wa`N^o#6)=d+)9KFfJo!WwEC1|(d1Fn{*rURnY=f>XD_`3BV}!9lj1L&1oIL^! zN$+xe1_&iaGhizN18SS(rEAW$p!9Xi9eZYHXB*KPyB^3i%ha(Z#3td2Hw`1Cji&Kp z^~217NlTkeVpcQbD*I*+ovv#SEJj~QyJ&1J5hIZ?TzaQy=dic+uRhb(Vqm?frLn~d z4I|%?YWC$0N$)1C?xi?SZz06m5bOFU?00>`e(?uoqY5i!wkjg@9PIu75qg&og)e|- zu_0!<^Rvp(i_|VB^Vli;VXTf^3aavqfiq?I@zf{mq4?g0Yw3?$n*ICN zf2!dWnmLVZgS<=pom%kt#DY%ski_UyqWH;mn~!NPN?K>P^^Nm(`?v#j$38A}9Y}J# zT)hQfk#%f~5t%+vz1GBEougIrB|o9m(z|+#O3%kV9#gxMu^DY8ltf0}`&H>Dqe>-~ zj#KlC3=6F6c2qdIFH!6G>_~WEkvWkRBiC|s`4LsW@ogOh= zhrUAq{YUOviEq#&&)LmH6T}uu<17b#JR<4NCc~Ro${wnKiOQBo*C6l67i|FWKp(hw zuO`%&(xOk`E4ZsRtV+A@o4RJ_#W13SsMdvkug1K6HO>l4d!P069bbZ~&_rK+9CbhS z%7yYLZ0??*-6fY%BR@w!Dw}b-@=K0y=28c|h|wu}q5ezB@nrqsvpye_vW`~{FIc36 zYcC~acM#rQ;AE@_u+bWTn`g8q7CVn^d`a5Lt-V|D(#JLgy$q60kc%9*&uuVWyV_gQ zLQSPiuUO+;u?S-2h}*>e=~fm9+zOKw%=l(h08CpExt96#J-?~8df%-Sh>JxfRx}K&S`;y!oyMYowb{{7oNZEaw{+__cviwqm*Ib)#vpw zff}s|t71;b&(yCU8j2gJOsZ$a-uGGiB0&ERb&H5y%7N)^+06<{ZwvN11U}CUVUDCf zUB2+7j=g`H^EHE%_&EKXKx@5^Hl{a{6vj+P4}R{*S^eKdOQZI#28Yu#dJi-CdXX=%eOt*31;I$8x{de{$^XpZ2t@mt)ZLBI4b+2wX zk~L`7=K3b)m*;9I_~63`q3@hcpg?Q#tl$p7o@i(*T zvt_MGXFX(+E5dCE1`Q*nWg_ON&1kKNt9zTD?i9Cgx*qSDJJE3adAM_f*iTs8hFj5X z7ha#aOMjs+du~Lj!(L?p*~x?Q;$ZVH;s7Ho0>uCDcm#nEs(VE10E!2s1R8|>h6%@E zj|hI^crfk)qp_^z?qmRAj`bQi}&(K!xYWNy&OHd zfAwXJ7gYt!FHK8kN-&w{Wsc6=^uE12*Q+!8sZVaxv`5IEV;URtp4qwN6iw(jt}H}{ zg^r5GAH4E@W=5}jaJnXfGB~>f`yEWDk>hz~cc2|b08vBkP+NA@NzUp{r$n@ATpE}UvfuuSMkc0$z-2&Cg3P6Dv6o^WsF*%k|?c9$8 zR;pu4$ZZ5598^UCezHL2?Ld8Z^Qf=)24ID>13$gtN5pp+#4T!ye_T%cGMixTc*;W! z@M_|@V%DZ^QQO^m4imR@yJDgf2wIn3%~%Er_f&rIuc729M~i{#AqBeB@0|tbs?7!wve>lAVkwd`()%=z=WCDPtBJ4W?r1XJ`6_B|N<0vp11I*q5 zOb0Br1lFX%I(o1P3dRkB&_^)N7G7Kq>jc8`{_y%mutpa!4+IisuuvsT+6_y~z?jXj z>tk3g51^g^>FdBo4D1U8P7I)?4=X%_SBk*)pJD8E5UvD^J%e4u!Rqe=nX5oF0|#w& zfxVief!!J@=^Dj0lt+H47~)Y#C#g0;P%yfTWv^#zR};cEO@`{rk0uuRboSII`-5w8 z><1ABBpVSd=ska5@jcj7dXb*Qv#zfN?`ztX`=m%&OpO2SBS=))4lPFEelrGkGW_*H z{s83Nm74E$+gn_b;xD8KpnJ^@vz}y&+e-C*>l(HH0PlMJ$N&HU literal 0 HcmV?d00001 diff --git a/crates/ptwm-core/tests/fixtures/legacy_v1_0_0.ptwm b/crates/ptwm-core/tests/fixtures/legacy_v1_0_0.ptwm new file mode 100644 index 0000000000000000000000000000000000000000..9c1445798ad1406675fa360c7f099d605ce8e52e GIT binary patch literal 25081 zcmeHQ2|QKXzu#vt&oXD8xtWTPkg<>l$Kr`zrNzuSACt5@%4tN|vZ^4t?PXN!nx)&{l7yuR( z-LQ~#nqPlT+-DT`u~CwX)#zG$F-dCIO~B(v&Z|!0LyPRQ79W!p+(6yR0fTTyy1LsS z-4PB*BvNAc=5V5I44Cesgi#XI%09iFf=gtl>RcXTR|W)^uLTVL7U$^8)B0^JcJ168PF~JdQ}zqy zS*Xn6dL4bsELmi(Yf4cJ%_jIT-|;qjmcaj?v2X(=jj?xk#oF0#4j&T|Z$7!U8x@j3 z$TRhf8Erp?4i6GI)fJ@*+9hzm85d{Yz=wkk(govy@m{A7{fF9V%pRG0IFh1xG23zp zS!R){D)*c{wx>q7n$7$o_qqFOY#Zp?V{L6wmM)t^KDeowK(Ufgy~c;BM;jKG&8Ihk zQmuG0XlZa38mz6=$g+XFyQK?yb7;@W?PZjmL5t0VpLd8;=!Mt%o#NKCYaDIeA2d^@ z5%0K>_6an`W^-tx3A#|UdLNAollI7z(QZqiRXe}4@47HOMdrip@PR5mjt%WzrH!z0 zu*Ubqe=2vn@sim@iRwW=x@8)zr7`NIV$@mR{elgpq;v`vW9Q_VH!xu9X1zIWe@l43 z5d}mXTsCjEI(xu$vVLc1NWN4_1vX$Jq=RWWawG1TO_TY)XkS}_GQZ$mDVguLAIq^l zT12-0m3d4-paMHLjZ7mK--c$d4Z{)$w6&!(5@ESX5Cp!=vBa67^DZL}z30|G$iXcg zuShkl`Efq7kic&)JBw!n8?Q^d!_FpKD<*;QF z+8`WjR~HP@3xhy8csaP(AwZkGNs`Ixx|~}OXsfij)Myl>u%|c zaC1kl4&Scso5RPtuSo;CrU1~NkrnTWrRih`{xVvLSvzCVpv#I7N4VM|BmjU@;cMp< z06wj(_r#7IF;I0{N8@pa!<$76MxcY$33Pz!m~q2y7v+g}@d9-xC2U;K#&+ zfGBX{2hsj9E#9K~H3E)+BVe+n0Y4N0O5jIE9MA~l{7?>lWb3zRQv&Tk8*p|@1O7mO z0{C}Da^PEW60nEZL%x;kSNZ){T>C76lt7Aq=~dq^QYQs|wnzee|L6i0fl5QaKb@bS z0x|Gw#$sqLv>5tpw*AkoCIWt3Bm{nBOoAPPjX-}S2mfx`1i;T817QbYH1MCD%Ps!C z9|Eel?RJ_& zY#`RZ+mG&?jX6cl`d-+~4O(zkkZrk1T%>+?K%Xm(hUlBmb>hw-DGuU<-jQ1hx>^ zLf~gcfDn8g90rBbgOeU&@RST(1_xgp2M;Q6uNeT`=K*}0R!V|De(HQrQlEq6aohqTA^?#3tq$2!t56o`+;D9n1r=_D3Fk?Ev`OSL}agbR&U(!tIi1FSHk(ydHHnpnU}Q;}$UOE;uHF%u`!2IK99 zw70T#cD2Mv>?Q!KT5Uf*thEWOyv2MNNCJ>*QiLK!m4VOi#uGf25f}^(nWs*ZBvK;| zVIx=E=ZmwMYmvLGDbmZ`S2h3P+POOKMmulf7LkrVxs>E_tE0=cfl@4Vy}`Nj1UgW8;=O!fk%kaNuwJOEEZb^j-X&b-GgnE zHP_l8LVR~L23A&9TA*xB%~Y*a4X*Pf3%IiJ`_qWSyUvj`O}W@F$S4?+$Suz?omdG- zYq((+6|@^|)8zBGmzq*%mjuqh zd1+%Iy3vt#PBuSfDH=!dT4Gd9-T(PCHf% zSIwu4$PyXTDr2UB_a6F7$7A5h9mQ5rsPwn{E?o>NA^UC4CL>=YQ%>2UO!JU1OXs47 zV6*g<`REYZw%f9y%(in!+o&>KF63pAC|>5PW7C!AlCX*{A305~13xk=wjk{#6o928 ztiF_Cb;}hJcW-H7mdo}rXTzS5i4Jn*n>;ziXBDuA#lIGPCtmdLvHQtBiypigueXyO zIGan!DKloKALy*TbIM4OI-m?E$|w`A4r+}xbO#^+2myy)sR%ZeI0Ro1D8ueMxG%om zI!^uA_0yti#9Y}UCN1eFYP9Ko8|)q!Ixh?y{|GyBd{G}goIc<~L$bYt+N z(q+w?U8&Yrep3@55=(H-73w9QixHZ-Z1z4n<$U=R4F?0~;}XHEoq`kDnB=|o3S6VO z)@AN_9_`t~y@+Gjy1SXZW=ft=aU^9VZQz`0+WlSoRZPjHS(r-sZ?RpwMFeU&IJm;U z3@bGFVa3JO9ranM0N|sXC{`xJIp)u51^+yxtgF?oonMbAI6QXpvw2H^Mf>T}QTlQk zu~9_iAzw@z(#_Ww;(V;m$UHZzTBpMCbXq?lq=iDnF|Vg7#F4E%a$m@WZOctslapbi z<|!$hW+nq$uuPnYRS&t$h>5AN&?OQ&s*mYwNgBObMFuI}&f8b<+|q63VKQG6&EcV$PXX{_=TINq$Vu9UV%U)HCJDf3kj+;M-p zq?AYsVg%*#&+}@1c8pZ@LE$-rwE6qhJ-*uaNlw8z4^{+sWDh6Zb$%{HeQlbdQ?UgKCi6w9n532lB#B;~0j>5A-Z@dL-%)yuc*+H< z6wGDYQjtlD5@Zsl0_B1eA~48c$!Jy)n{gD4?gO9dPJ}t zPVZLMNzI?<)lkfZp7#Dge9rIY+vR2cS7RS*{H?}TEZ|@DwM4M`)o|QV*$Us+V&MQ` z6b~V_pbv(HloQ|rI04pY*fQ-DMulA?v6dZ_E#B;eYGKJwD|Hxn0s*Z^T@3b)%PD*6 z{Fh?21 zhqn{ZBP2NV5tmc1KWQ;cfe5NYG~iR%6t)Tdsz$h?S1>tIb#Ny$lJC6y(Zzy)5HKI8gq{uEr{r zF_;9nlsw6ufLzEQ;v)=ettRua2^Tt6IM6Q1VEF40a$Iz6cRT>VN%GJRac9Ibg1J6x z9n-$g25GIuP(m8;bu1_zv+^~?1%RR7YndRNU@73j2~mavOI5Z8jQq0QeXU0{Ny$1p z7XqiUY{$*j)n?z5xK+oWcsjxvow?YS)vZIbE zIM|%$iuTI^Poo3XS~|&CtFiwpDR46sivxzo;(i!useN-wW_dflO)@`ujVVw7{ zZZL1l(Y>e$p?cNPt5)rFN%%tV!+kiNT!N9tGqP`q#bKlKuOu!!C%AuQsESMDN1o`g$(0hMv`GK48loMppXjvi{t_D9nPgL5;*6 zrP$lqhUM}#s+Vd zw1C79^kHHF%It6CcN^S$;>S|OK_j^>?1D@V1v$$@b?srd*D2DSil3B~4>I>PFVZ|e ztYXIdAl-$5Ews^V*7~+ltG8@TXoX1nn+)7Jvg(8uL@*`}c_uM7U~195J(Ds6F&j~= zm-qyChqu_q{Lahs(c=dfi5{!_|52H@b<_=WfF0Zxz9xrrPsA?j|YtZbGq#kafMA*p6JZnp0@9!s&VF zM7^G~dNCSnThyTX@bP8$VCJ|B{*_~XaTYCAUbkv;I?;&)UuZ}lV~zbXKkk# zuS;Czw9!94Y_TKy(TY~Ogdd`FPZvk+k%%)*DIyhJ@ct}Pww8mQ4DW;l93$qVD#yA9 zj4Y`2jt=o2{;g|AR!9B$R9rEi};plQRw@@->ZdGX2g`y(TpSZ|+XLFk#gF5fYttlHh&|ElhdtY=(TPIVi_ zc8hsGOS>yUQXTEWzA59p9#l)WXtkd~Gx<>2bkC-Hp=?7Live(#Nb=P4zej9OjpU*a zW^5q$IVx;VI+2A96?c|kVlby;LvxGVXS)!{X#G&{=ArBpbg%ck;?6MX=y00#%#dqV zky^a_!Q_Ysmq|!<#X^yA?som&lRa1O4NH8;k;0BE$v+Y=;dsejmoGF^JvrrNwIBH86+WBbC2w(QR17w0Po!)E*#pTO8W`8N;BaBXMH6iiTJMAL*jA z&S}$bg38Ln7>-Al2TYBLBnAtKb;C-TNAtKJ~(=wkw`$`MA9x>VGCNd0p~x*o@0PWn!cA<0!;4yd@^Dnw;u zW#`m|=D5j9Wipz(@YlX(-V4zi*LimSL~vaqqi-@pe!cT#uP5}R*6ZQc@h8Seab~9B zP>c-m&Uu%na8JF?+u77&dtcL6QL>mBb!142e3B3JhNr^2=ojsJnJ#ep&_N6N0VI%6wK*$J~O&Q zFzD)^*G~SJ==?WpC%9E}0*3zh$(;ZoLyxnb1}KC8COLqr8pmmcgTesGHvp^?MdKwQ(eTt3jtpS97#7$hZIoxdIGDv&Ka;njhhwk z;z&UOWpxV58u|{1entriIFmqJv{3+e7loF(s`4u-*Lb3L#UPvRgh%C@^$j$|*Z~E! z=w|;N2n8ZIFql+f4U}Yf>pzgH}2q+TgqUvWeV!ufW?Y2_{#FW zjoW)WEWNmGnaujb-37T~xGw+lVi42YAo{tBd4}7m5VH9MwtiGVBj**SkGpt{6egi{ z7L=y~L#|!i&P`CU?QQx*k<#;^Y^tW$y4N+iE@#*U7PV?!^oK$`#>w2T(8o_+**$Z_ zhdZoTW?I~raG5(g<_^u{nDVPi+jbA-kFI_dW6Bxd1>JtBmWrL9yq|W3-F49c$Am*wm6Ln*!l2#2%XF1)FsH7I$wZ@o! zJHpTT*jV!8Np)+Ja;w-9%@>XAW&Wg#Zr)?={fvC0a1uOR0@k};rajGPEIqiMV1_4% zR%;XUwO!s(e0V*544;47A%@m_uT-VY4y`Acu1?~?un2u~Nq8kquKvjSH90KKL1(p{~pwI)mi=vuGB=J literal 0 HcmV?d00001 diff --git a/crates/ptwm-core/tests/fixtures/legacy_v1_1_0.ptwm b/crates/ptwm-core/tests/fixtures/legacy_v1_1_0.ptwm new file mode 100644 index 0000000000000000000000000000000000000000..883fd1e6a3159409fc2de04888559f96291a6710 GIT binary patch literal 25081 zcmeHQ2|QKXzu#vt&+`y6Pd8(POc4s1D;ZCQgM))O$3eFoWXKQ+iOfT!q+}==QY1>K zluF1P4P+kQ#_4vu{_pnQ=jzqFS?jaU*?XO}*6+8zzu#KF-&%X`bxs=YH`SveXD0^$ z3t%mDfafbgles}!Uu-X$SxETfmI*PXmwa1-_qJLr0Xq!!I9v79=xo;Vb+3H^G8Ahk zp{i@&OyJj1BZ1lS-&j5@sfqLMYa``F^3%=;;aMkbnNAwWjGb5I08sAY9vB>2+!l*+ z#h~28z1*>$cF0{);!@(05CDu&94`?Up}f$OZca9y_Bb0)B*M>mj z&I1csqxtox?ci*DC+Eq*p#)zC9nPz4vMgY%Jl1vnT@}}4*uvItW3gR&V>pRpLJx_h zh~>m*B+E9GJe+_NbjLb6Z44g^VsGh@y$btG zYr7OfBD}UWX1ot|d2+Q`=yCLkZ8r2JT$Q*OB+Mar4?3+Q<+ph|t3~g^VU8L%huHOhUyphU)Ki=)-#K zOcp|8C)U&U^TOC~4DBWNDdUk+2d(03d~XJhbHT@v?bBi&C4-h?BehfV;&EHowR?p& z!X9Nu=!ySS?v=ecOUo8v2RWJ#TdECOhtY-8`;@tz*N+6NPa?BQYpDguCR%$RUxqHQK!7-YZ-HuavrjbQ)5g zvjCH2#Gd@SAK5~TgU7bNX98hkZ{vYMc@xa!KUK&(DzUx7lJm0*N7Zr^`BVD&qL^43 zTL)f>GE1QdcA?5|fJ|3q!4LaNEN>Pj%kBz}uvf+!v9?(hjV* zPadF+b_DH#2g1(X9%<|D?y@nQnDWX!8*_b5hdVxEppl`Ue@QMcEJ85+ym`~hL-&-E zd_~p?$IjgigY>~5&?q02nlrL}R>A*zqJT+gboH4yykxE>^q-w;{ zT4)^}gpT{A<7^BgC1xisgsNy+81rB4lC@Ga;_~U4u@ar>H2*~X3Vn!T&w9tf4T<%% zaYcA|B3FiQch8OCW6js316@-9=+4bc#9`^$xq!cnRuXouSTE3JMMxst9S~9gz^(GN za|!@U%gg>>r%?dEa-4!3hBQHbWdvw}&A5rcCIXuXY$C9Uz$OBl2y7y-iNNZ-RKG^R8E^*7HZ|afB0vrN$cP8-0R=ylgCE)YP1@8zE6@U*+|+!9~X&%9~sZXj9`P%AIZVLn>G>fv&UfAJ{TSRXXkR0 zzwd{@ZP;yC6zu!e|MS=U?;eA`-8uwpKMVo;cD{dWjvr_Je-_GJP$n1;>}TQT|C*0K z`+~pfc@VS~`uG=}3w+<%|Ia`FUE}Zo6bg<1CD#MKtOWhqXLlgUkUPKNn&7Xp&i(Cn zT0!g~cE8+xf&Xv!1i$Cqs;9sr7{4clpRew(pOE)`5cgOS=n1Xi`uenMDn6IgkZ)c|lFKx)boi(Rh_et9R6=&`&|Uuf72ZI%p) zCTSQarRpw!y!~{OVv3GfC-3X3nTecJwZ8WpeMy_d+FmQBr@pndT)Z7D%Z`hWyp@fz zq`OC96${gzm#-9{kBf_FCgv82(66XKNWasqlnY6yYOQp_i_cbN51naTdVj|78vVuS zo6W}U^@3W}+jmQB?I;Sy%MhV+gbQ}}3$e}ya9>N%E{h6^i*1>R80P!X&W%I#w@;It zf%@IZn2jxvb1n?#xMD3#80SK2!ECPqnFrfAU z19eSK3q(ZVhIa4r@^TZD)1{H7nWo--hHMsp^X>g{#QyE4$Qs7noMz=!4v;A>PP6(g z2PM}(cD@SDPrN9lOc6$rI#znS=iBrOEyvplRJEi#pxIjUK zi8&D{CivgW-sQ9G`982n+*D-oVgs38Cws@g6}|q4!k2)_RWaDY>p9lZmkC`uvM1+- zPQdx;;~^F2x9Qm)QHU|JpK>hRr(AB#O4zs#Rqc^> zee@Pj6~Vn z=d^_z<8(T4WoKhctXi8#M&POB& z%Sc?Elxus{9g=W&es+q-;W2mp&aja-O7+Wp1tlj{u=^#y7JWNG^dGT%C_alGycusO z%>|q+q~?|%wlEHM)sr4Os7f1j6EDspAE^aujU99wAOwg2Mm@WQ8%j~3XM}EIckSDi zSZ9}@b@<|OaZOU5{6Vv(Y`+>k#^3rndizg{LT++FFFpyWOSqhW4H-d@6E1YawDLquC1(AG zjeE9a$nk_ju`f+5@>pT4^3!L<^82H>=1mw z!$6|11{-GrY7Fk49|#&p1)8G0EG#hsXc%v0Tf2F@`WQIgEORWEwZ^6x)5P8MR}tQJ zf4sDeL>6)o$`g3Sr}^1oa*YSYr%bYD?pMF?*Sk;lJDhvptPZnk({P&O1muRjQr-Q-8m+QB66KTLSj>s8+wuAJs}=$sc<|Soq#jTe~V)g2`b3x3n}P zn7neVxONQS0V`UL=$wc=y89C6%gPIGOAkaBOJHevG?9zt=G8~C%wA_av&ed|V2tyQ z;@qFzp{}1u3p?St_)ZPZP^T*R6YsZT5QnfZbN-Dm?0_J-T< zPPD1{&M}lcc)lyc3g;`wZExC?3VjIVV3m>2g`ZWk&?1$6r)%LkXy&mm$=$W;QBaT- z`e0&YD-jbyYO66KCG+BwrUU5^VfRb8xWs~EvvwRp$kFXlQ5N+;uQ7-@9SAWvL=&K+ zCG;#L+7@Xqj&j4uhdjo998zSPVg_475JCVO1X%O<>{oI?zG`{MYP$hIiCS~b{atfw z-oZfI#zp>70kyOoEvtRbm>pP&F*#S=5A=|#9TC)=xh2pQ78)(gQ}AX{-h;TeV#%?_ zHj^ck3`j~H0Dza_qwnXZ+C%USAWidW!Fo7gsPyu)i`LCYZD z)l!7n%K7@z`K8;qAkn!ebx&K9^^0NL+P&y-TE1r^ErH~dRvE8pc9Y#*`s8I$Z_CFr z<84=O)ly(^Ruw|i)7Gh3?UE#Vrt{%0ynZ3k;Jp(H?@1+LLo+>6XI>KBPwTIuG|m|~ zN3M&CZs~X9NNX6GLk~qRo>Nf6m;EkXt>q#-5?j3?zQ`c*f zd7@Q2+gq?a{s;K^1`LZ|F`}xbxY$NoibeEiBcp9JGBp-sT}{$ixQ(i!QE6u*)$9sg z2C{;ZCzv86g4DU*Deo}3`y_z9YAc-#L&O>R0xC-Ohgy0A9&gg++Etg-)%UTzZk(fg zxqr6>|ATBdX3p??K2vt(2b+BrYQk@cWxvbCpQ5Nf*MtbgBp^>D#|MqgdA44l&P7Z` zl^7;J!QbF7vA4R>bvpL#zFB2M+EPyazrH{_E3A_a}L~cYqT(sFILQs zGd(ciy)ec4TJ-b?^Hi8mKt{)%E0RX3XQ}2onn#}`KxCD&EG}@hVjaxx6znN|jyig2 zgk@QA|D@Q~^ArbG=DFYKCN7`531^^flwg*|i~n&qNhfp@s-46fYu&_Q@T~2$N`o!l z&|5#|&9v>Sp?HVu^%@T!r+9|4C7cPY91cjZZmRORT2rXVVMbSbK)69Oe_rQ2U9S0* z!x+m&sSIv=<0Avs+hQLr>t;&@Ali36-+KE{)QN_4v0KmK-Ff7kP5W@nA4G(lqh?|% zhdX)?TGJX@_Ve%m?fJI6wz|`q{8|Xr7A5iL6vd7Fikc=L`FAolpkeB7+Se6-Z)5IwWG1Sr}mu!F5!7W zbqm!N>zM!>$FvaHwpLO9^tb%pH1k*K^`1d52%z&BpN)4yIs0?xg5Yk^l$obLMQzE9 z=3xqDsi!<WHbIsb?T&8fj zijBKv=k83H9rETe3#-00dtJ0}i*e`C7a4a4q$Ubvv2WFsA4!&O?c%B}5}B+X9Sb;L z*kBWF`l7Ebz6Nl2XHM?{4?2 z#pv9;{DRu>0uKe*3oKS{g16tWDM1Y1>OVW}7h0Rl;-AV~ROdR{iGv>1eKXMf_K7J{ zl8tpB93xLEJ>%98i8E|3&!?48dc#yj&2Dk9EmuZtNjcaTo(X@>H0RjKdWQQLeH8kU zf<1k*^ULIjfkg}JXiYuwU6&H(>^=xz8}lzLRO3L-eoT4I_LxQHG}Y{r;!83s^Wv3N z5t2V_<6l-qq`s_*pq?J}X;Jp40vC@!n_$Womh5Uc98T%{fW&X`eQx}W7e!jIa6zZ* ziJ@)6AsK&OJNaXx^WUtU;8)BE7`hWjr2#;m32!$JP>BGniU3VDp4%1=g#j|}09ZSo zN)*q9!t)E_r6G9kVZiqho?8vicLUGwjAwAf>u$w!b^)@^0EG;mr~;2Tf~Tgz!w%vN zp5WQD0mxH;x)@L)1N@u;9W=l$f@kc&tCHX~=J4=hz@HgU)`8b22X=owfi+m;4A`d0 z%Rz8)WTAk%78P|3QyauMx0DQ=Ng&SJtAM+ULd)DW1l81Qa9-QvkPSB?V~Q+Z_cp{i z0@vuh8Uwc>R7l`JUtS&4d*DFnMb}se(dPxJ)s-^<_>03YEpF$8b#nnd6vDC;=+`CZ z7t5P;&;A))+RsM&3%H;(zbs6V6GkNcN+(gm>8h}aRt!7AxPwz}S(DMGF{p1n7AwK* zufYF4L8&|IVBzK ztnPuJTC$m5qn^T)`FymOOKhh&Q}bxzH5qJihsPYSlU9f>kx0-wDIOR8;7%xZkJ@F) z#QhTb+-Qu_wfOVF3d`wNyb;j;pqPP}%UjJ)GB)tOa*PyUJfv44PNyRJ@T5~yH`%MO zz1a-SVGp$(Qix%6h?q#OxX|bWWcO%%-i2mRsErI((6^5(@wA`ZdasofMzS)A{Tx;m z1ia{Sgd+Z|I44#z{Gm9j6Aub9z>5~Za06IFz#9+Xo&!EH5MM*69{@eN7f2zLrj_)a z70u^`iP6FkSJgE8wpYgTI|eBB;leTea=jJh;mU_wtA`$9d)~JSXV=NgxfSl5Wm_VD zgfsrczAU6bIuVl%(`>c#^L2APjgX0=RO(ClBHq2KD<0 zIl2otLEwR#lZB^JyszEnc_VunDNj*Ir*hC;EhiVExD0bv+I#J4x;5Nx)wd~9)Y~xFHUftum8Av|o;XCZv%_1-aCnLZmWVgL5>uC{7*}k;| z3j#qe%@zrNhsA9r``6Nk3Hcd}n49nRXvkR@ttFVROya>X2Yr_$x}2q0cWCXJl8`>& zbzFJYi(I2ZXDxk!Ac2PlZhfytw>Z%gMk;Lr-SMJbx9q5j&+M~W%U?&Rr}d~$Ait$# mPb=Z?8;%m{$w_o9c0cbrmJ0euD?x~m|0MW-51IyA?EeMEojw8p literal 0 HcmV?d00001 diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs new file mode 100644 index 0000000..1edc801 --- /dev/null +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -0,0 +1,911 @@ +//! Golden legacy `.ptwm` fixtures. +//! +//! The files in `tests/fixtures/` pin the legacy (`\x89PTWM`) wire format as +//! written by released and current writers: +//! +//! | file | content | +//! |---|---| +//! | `legacy_v1_1_0.ptwm` | Multi-tensor container with the v1.1.0 Extension Table layout (every built-in codec interned, `ContainerWriter::set_chain_registry`). Holds bf16, fp16, fp32, fp8 and int8 tensors, two scale tensors that share an `Order1ScaleAC` prelude state, a chunked bf16 tensor, an fp8 tensor with an inline chain, and a tensor with a dependency. Built-in ids and labels are derived for version `1.1.0`. | +//! | `legacy_v1_0_0.ptwm` | The same file with every built-in id and label re-derived for version `1.0.0` and the Extension Table hash fixed up, as a v1.0.0 writer would have produced it. | +//! | `legacy_trimmed.ptwm` | `compress_model` output of the current writer, whose Extension Table holds only the codecs the planes use. Same tensors as the v1.1.0 file minus the two scale tensors and the chunked tensor. | +//! | `legacy_p4_byte_fp32.ptwm` | An fp32 tensor written from `Format.BYTE` input by the v1.1 Python path: its Source node declares `[n_bytes]` fp32 elements, four times the real size. | +//! +//! The tests check that every fixture decodes to the expected bytes and that +//! `implode(explode(x)) == x`. The expected bytes are stored as +//! `(length, xxh64)` pairs in [`EXPECTED`]. +//! +//! # Regenerating +//! +//! ```text +//! PTWM_REGEN_FIXTURES=1 cargo test -p ptwm-core --test legacy_fixtures \ +//! -- --ignored regenerate_fixtures --nocapture +//! ``` +//! +//! This rewrites the four files and prints the `EXPECTED` table and the +//! golden hash for [`LEGACY_WRITER_GOLDEN_BLAKE3`]. Regenerate only when a +//! format change is intended: the point of the fixtures is that they do not +//! change. A crate version bump changes the built-in ids that the current +//! writer emits, so it changes `legacy_trimmed.ptwm` and the golden hash; +//! the v1.1.0 and v1.0.0 fixtures are pinned to their versions and do not +//! change. + +use std::collections::HashMap; +use std::io::Cursor; +use std::path::PathBuf; + +use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; +use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; +use ptwm_core::container::{ChainRegistry, ContainerReader, ContainerWriter}; +use ptwm_core::extension::{BUILTIN_PUBKEY, CanonicalId}; +use ptwm_core::tensor_record::{CHAIN_REF_INLINE_BIT, TensorRecord, parse_tensor_record}; +use ptwm_core::transcode::{explode, implode}; +use ptwm_core::transforms::op::OpId; +use ptwm_core::types::{NibbleKind, Role, ScaleFormat, ValueFormat}; +use xxhash_rust::xxh64::xxh64; + +const V1_1_0: &str = "legacy_v1_1_0.ptwm"; +const V1_0_0: &str = "legacy_v1_0_0.ptwm"; +const TRIMMED: &str = "legacy_trimmed.ptwm"; +const P4: &str = "legacy_p4_byte_fp32.ptwm"; + +/// Expected decoded bytes per fixture: `(tensor name, length, xxh64 seed 0)`. +/// Printed by `regenerate_fixtures`. +const EXPECTED: &[(&str, &[(&str, usize, u64)])] = &[ + // EXPECTED-BEGIN + ( + V1_1_0, + &[ + ("model.layers.0.q.weight", 512, 0xc82d5962a46d19ff), + ("model.layers.0.k.weight", 512, 0x73173bba2f17ef23), + ("model.norm.weight", 512, 0x946ca638752e3bfc), + ("model.layers.0.mlp.fp8.weight", 256, 0x84e10734e3fb3518), + ("model.embed.int8", 256, 0xf6193be8f075d386), + ("model.layers.0.q.scale", 2048, 0x95a622e512233480), + ("model.layers.0.k.scale", 2048, 0xda44a3268d62821d), + ("model.layers.1.q.weight", 256, 0x0d7f3dd0787bd569), + ("model.layers.2.chunked.weight", 1024, 0xf9e05338b3a99b14), + ], + ), + ( + V1_0_0, + &[ + ("model.layers.0.q.weight", 512, 0xc82d5962a46d19ff), + ("model.layers.0.k.weight", 512, 0x73173bba2f17ef23), + ("model.norm.weight", 512, 0x946ca638752e3bfc), + ("model.layers.0.mlp.fp8.weight", 256, 0x84e10734e3fb3518), + ("model.embed.int8", 256, 0xf6193be8f075d386), + ("model.layers.0.q.scale", 2048, 0x95a622e512233480), + ("model.layers.0.k.scale", 2048, 0xda44a3268d62821d), + ("model.layers.1.q.weight", 256, 0x0d7f3dd0787bd569), + ("model.layers.2.chunked.weight", 1024, 0xf9e05338b3a99b14), + ], + ), + ( + TRIMMED, + &[ + ("model.layers.0.q.weight", 512, 0xc82d5962a46d19ff), + ("model.layers.0.k.weight", 512, 0x73173bba2f17ef23), + ("model.norm.weight", 512, 0x946ca638752e3bfc), + ("model.layers.0.mlp.fp8.weight", 256, 0x84e10734e3fb3518), + ("model.embed.int8", 256, 0xf6193be8f075d386), + ("model.layers.1.q.weight", 256, 0x0d7f3dd0787bd569), + ], + ), + (P4, &[("p4.byte.fp32", 512, 0x57e3af5eb7d32ec1)]), + // EXPECTED-END +]; + +/// blake3 of `compress_model` output for [`tensors`] with the options of +/// [`main_options`] (see [`main_container`]). Printed by `regenerate_fixtures`. +const LEGACY_WRITER_GOLDEN_BLAKE3: &str = + "b74c3fd1d81d5bbf4703fb55fb6427781d9aa6d2a92a44127cf4dc91ee4b9791"; + +// ── Inputs ─────────────────────────────────────────────────────────────────── + +/// Deterministic SplitMix64 byte stream. +struct Rng(u64); + +impl Rng { + 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) + } + + /// Roughly normal value in [-1, 1] (sum of uniforms). + fn weight(&mut self) -> f32 { + let mut s = 0.0f32; + for _ in 0..4 { + s += (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32; + } + (s - 2.0) * 0.05 + } +} + +fn f32_to_bf16_bits(v: f32) -> u16 { + (v.to_bits() >> 16) as u16 +} + +fn f32_to_f16_bits(v: f32) -> u16 { + // Values are small and normal or zero; truncate the mantissa. + let b = v.to_bits(); + let sign = ((b >> 16) & 0x8000) as u16; + let exp = ((b >> 23) & 0xFF) as i32 - 127 + 15; + if exp <= 0 { + return sign; + } + sign | ((exp as u16) << 10) | ((b >> 13) & 0x3FF) as u16 +} + +struct TensorData { + name: &'static str, + dtype_code: u16, + dtype_name: &'static str, + shape: Vec, + bytes: Vec, +} + +fn bf16_tensor(name: &'static str, rng: &mut Rng, shape: Vec) -> TensorData { + let n: u64 = shape.iter().product(); + let bytes = (0..n) + .flat_map(|_| f32_to_bf16_bits(rng.weight()).to_le_bytes()) + .collect(); + TensorData { + name, + dtype_code: 6, + dtype_name: "bfloat16", + shape, + bytes, + } +} + +fn tensors() -> Vec { + let mut rng = Rng(0x5054_574D_4649_5854); + let mut out = vec![bf16_tensor( + "model.layers.0.q.weight", + &mut rng, + vec![8, 32], + )]; + + let n = 8 * 32; + out.push(TensorData { + name: "model.layers.0.k.weight", + dtype_code: 4, + dtype_name: "float16", + shape: vec![8, 32], + bytes: (0..n) + .flat_map(|_| f32_to_f16_bits(rng.weight()).to_le_bytes()) + .collect(), + }); + + let n = 4 * 32; + out.push(TensorData { + name: "model.norm.weight", + dtype_code: 1, + dtype_name: "float32", + shape: vec![4, 32], + bytes: (0..n) + .flat_map(|_| (1.0 + rng.weight()).to_le_bytes()) + .collect(), + }); + + // fp8 e4m3: small exponents, random sign and mantissa. + out.push(TensorData { + name: "model.layers.0.mlp.fp8.weight", + dtype_code: 29, + dtype_name: "float8_e4m3fn", + shape: vec![8, 32], + bytes: (0..256) + .map(|_| { + let r = rng.next_u64(); + let exp = 3 + (r % 5) as u8; + ((r >> 8) as u8 & 0x80) | (exp << 3) | ((r >> 16) as u8 & 0x07) + }) + .collect(), + }); + + out.push(TensorData { + name: "model.embed.int8", + dtype_code: 17, + dtype_name: "int8", + shape: vec![256], + bytes: (0..256) + .map(|_| ((rng.next_u64() % 9) as i8 - 4) as u8) + .collect(), + }); + + // Two auto-correlated scale planes so the shared Order1ScaleAC state wins. + for name in ["model.layers.0.q.scale", "model.layers.0.k.scale"] { + let mut bytes = Vec::with_capacity(2048); + for _ in 0..64 { + let mut prev: u8 = 0x80; + for _ in 0..32 { + let next = prev + .wrapping_add((rng.next_u64() % 3) as u8) + .wrapping_sub(1); + bytes.push(next); + prev = next; + } + } + out.push(TensorData { + name, + dtype_code: 17, + dtype_name: "int8", + shape: vec![64, 32], + bytes, + }); + } + + // Delta residual against an external safetensors reference. + let mut residual = bf16_tensor("model.layers.1.q.weight", &mut rng, vec![4, 32]); + residual.bytes.iter_mut().step_by(3).for_each(|b| *b = 0); + out.push(residual); + out +} + +/// Tensor name that carries an `ExternalSafetensors` dependency. +const DELTA_TENSOR: &str = "model.layers.1.q.weight"; +/// Tensor whose chain the v1.1.0 fixture stores inline. +const INLINE_TENSOR: &str = "model.layers.0.mlp.fp8.weight"; + +fn chunked_tensors() -> Vec { + let mut rng = Rng(0x4348_554E_4B45_4421); + vec![bf16_tensor( + "model.layers.2.chunked.weight", + &mut rng, + vec![16, 32], + )] +} + +// ── Chains (mirrors of python/ptwm/preprocessing/_chains.py) ──────────────── + +/// Source params: ndim u8, dims u32 LE, chain-internal dtype u16 LE. +fn source(shape: &[u64], internal_dtype: u16) -> ChainNode { + let mut params = vec![shape.len() as u8]; + for d in shape { + params.extend_from_slice(&(*d as u32).to_le_bytes()); + } + params.extend_from_slice(&internal_dtype.to_le_bytes()); + ChainNode { + op: OpId::Source, + params, + } +} + +fn node(op: OpId, params: Vec) -> ChainNode { + ChainNode { op, params } +} + +fn edge(src: u8, dst: u8, role_override: Option) -> ChainEdge { + ChainEdge { + src_node: src, + src_output_idx: 0, + dst_node: dst, + dst_input_idx: 0, + role_override, + vendor_bytes: vec![], + } +} + +fn term(node_idx: u8, output_idx: u8, role: Role) -> TerminalRef { + TerminalRef { + node_idx, + output_idx, + role, + } +} + +/// `_bf16_split_chain`: Source → BitReorderIeee16 → ByteSplit(2). The Python +/// builder uses the bf16 Source code for fp16 too. +fn ieee16_split_chain(shape: &[u64]) -> Chain { + Chain { + nodes: vec![ + source(shape, 0x000F), + node(OpId::BitReorderIeee16, vec![]), + node(OpId::ByteSplit, vec![2]), + ], + edges: vec![edge(0, 1, None), edge(1, 2, None)], + terminals: vec![ + term(2, 0, Role::MantissaByte { index: 0, of: 2 }), + term(2, 1, Role::ExponentByte), + ], + } +} + +/// `_fp32_split_chain`: Source → BitReorderIeee32 → ByteSplit(4). +fn fp32_split_chain(shape: &[u64]) -> Chain { + Chain { + nodes: vec![ + source(shape, 0x0003), + node(OpId::BitReorderIeee32, vec![]), + node(OpId::ByteSplit, vec![4]), + ], + edges: vec![edge(0, 1, None), edge(1, 2, None)], + terminals: vec![ + term(2, 0, Role::MantissaByte { index: 0, of: 4 }), + term(2, 1, Role::MantissaByte { index: 1, of: 4 }), + term(2, 2, Role::MantissaByte { index: 2, of: 4 }), + term(2, 3, Role::ExponentByte), + ], + } +} + +/// `_fp8e4m3_nibble_chain`: Source → BitReorderFp8E4M3 → NibbleSplit. +fn fp8_nibble_chain(shape: &[u64]) -> Chain { + Chain { + nodes: vec![ + source(shape, 0x0010), + node(OpId::BitReorderFp8E4M3, vec![]), + node(OpId::NibbleSplit, vec![]), + ], + edges: vec![edge(0, 1, None), edge(1, 2, None)], + terminals: vec![ + term( + 2, + 0, + Role::Nibble { + kind: NibbleKind::Exponent, + }, + ), + term( + 2, + 1, + Role::Nibble { + kind: NibbleKind::SignMantissa, + }, + ), + ], + } +} + +/// `_u8_byte_passthrough_value_chain`: Source(u8) → BytePassthrough. +fn byte_passthrough_chain(n_bytes: u64) -> Chain { + Chain { + nodes: vec![ + source(&[n_bytes], 0x0006), + node(OpId::BytePassthrough, vec![]), + ], + edges: vec![edge(0, 1, None)], + terminals: vec![term( + 1, + 0, + Role::Value { + format: ValueFormat::IntN { bits: 8 }, + }, + )], + } +} + +/// Byte passthrough re-tagged as an E8M0 scale plane (routes to +/// `Order1ScaleAC` together with the Rows layout from the tensor shape). +fn scale_passthrough_chain(n_bytes: u64) -> Chain { + let scale = Role::Scale { + format: ScaleFormat::E8M0, + }; + Chain { + nodes: vec![ + source(&[n_bytes], 0x0006), + node(OpId::BytePassthrough, vec![]), + ], + edges: vec![edge(0, 1, Some(scale.clone()))], + terminals: vec![term(1, 0, scale)], + } +} + +fn chain_for(t: &TensorData) -> Chain { + match t.dtype_code { + 4 | 6 => ieee16_split_chain(&t.shape), + 1 => fp32_split_chain(&t.shape), + 29 => fp8_nibble_chain(&t.shape), + 17 if t.name.ends_with(".scale") => scale_passthrough_chain(t.bytes.len() as u64), + 17 => byte_passthrough_chain(t.bytes.len() as u64), + other => panic!("no chain for dtype {other}"), + } +} + +fn inputs(data: &[TensorData]) -> Vec> { + data.iter() + .map(|t| InputTensor { + name: t.name.to_string(), + candidate_chains: vec![chain_for(t)], + dtype_code: t.dtype_code, + // Format.TORCH + input_format: 2, + raw_bytes: &t.bytes, + shape: Some(t.shape.clone()), + dtype_name: Some(t.dtype_name.to_string()), + delta_reference_blake3: (t.name == DELTA_TENSOR) + .then(|| *blake3::hash(b"reference.safetensors").as_bytes()), + }) + .collect() +} + +/// Options of the Python default (`emit_payload_hash`, `emit_plane_crc`). +fn main_options(chunk_size: Option) -> CompressorOptions { + CompressorOptions { + method_hint: 1, + emit_payload_hash: true, + emit_plane_crc: true, + forced_codec: None, + allow_codec_ids: None, + chunk_size, + } +} + +fn compress(inputs: &[InputTensor<'_>], opts: CompressorOptions) -> Vec { + let mut cursor = Cursor::new(Vec::new()); + compress_model(&mut cursor, inputs, opts).expect("compress_model"); + cursor.into_inner() +} + +/// The current legacy writer's output for the main inputs. +fn main_container() -> Vec { + let data = tensors(); + compress(&inputs(&data), main_options(None)) +} + +/// Main inputs without the scale tensors. Their shared `Order1ScaleAC` +/// state is about 15 KB; leaving it out keeps `legacy_trimmed.ptwm` small. +/// The v1.1.0 fixture and [`main_container`] cover the shared state. +fn trimmed_tensors() -> Vec { + tensors() + .into_iter() + .filter(|t| !t.name.ends_with(".scale")) + .collect() +} + +/// The current legacy writer's output for [`trimmed_tensors`]. +fn trimmed_container() -> Vec { + let data = trimmed_tensors(); + compress(&inputs(&data), main_options(None)) +} + +// ── v1.1.0 layout rewrite ──────────────────────────────────────────────────── + +/// Records of `blob` in file (offset) order. +fn records_in_file_order(reader: &ContainerReader<'_>) -> Vec { + let mut index: Vec<_> = reader.index.iter().collect(); + index.sort_by_key(|e| e.tensor_offset); + index + .iter() + .map(|e| { + let start = e.tensor_offset as usize; + let end = start + e.tensor_len as usize; + parse_tensor_record(&reader.buf[start..end]).unwrap().0 + }) + .collect() +} + +/// Rewrite trimmed-table containers into one container with the v1.1.0 +/// table layout (`set_chain_registry`, every built-in codec interned). +/// +/// Records keep their order: first all records of `parts[0]`, then of +/// `parts[1]`, and so on. Chains are merged and deduplicated; prelude +/// entries are taken from the parts unchanged (ids must not clash). The +/// record of `inline_tensor` gets its chain stored inline. +fn rewrite_as_v1_1(parts: &[&[u8]], inline_tensor: &str) -> Vec { + let readers: Vec> = parts + .iter() + .map(|b| ContainerReader::open(b).expect("open part")) + .collect(); + + let mut prelude = Vec::new(); + let mut chains: Vec = Vec::new(); + let mut records: Vec<(TensorRecord, usize)> = Vec::new(); + for (part, reader) in readers.iter().enumerate() { + for entry in &reader.prelude { + assert!( + prelude + .iter() + .all(|p: &ptwm_core::prelude::PreludeEntry| p.shared_state_id + != entry.shared_state_id), + "shared state id clash" + ); + prelude.push(entry.clone()); + } + for mut rec in records_in_file_order(reader) { + assert_eq!(rec.chain_ref & CHAIN_REF_INLINE_BIT, 0); + let chain = reader.chain_registry.lookup(rec.chain_ref).unwrap().clone(); + if rec.name == inline_tensor { + rec.chain_ref = CHAIN_REF_INLINE_BIT; + rec.inline_chain = Some(chain); + } else { + let idx = match chains.iter().position(|c| *c == chain) { + Some(i) => i, + None => { + chains.push(chain); + chains.len() - 1 + } + }; + rec.chain_ref = idx as u16; + } + records.push((rec, part)); + } + } + + let mut cursor = Cursor::new(Vec::new()); + { + let mut writer = ContainerWriter::new(&mut cursor).unwrap(); + writer.set_prelude(&prelude).unwrap(); + writer + .set_chain_registry(ChainRegistry { + chains: chains + .into_iter() + .enumerate() + .map(|(i, c)| (i as u16, c)) + .collect(), + }) + .unwrap(); + for (mut rec, part) in records { + let table = &readers[part].extension_table; + for plane in &mut rec.terminals { + let id = table.entries[plane.codec_table_idx as usize].canonical_id; + plane.codec_table_idx = writer.codec_table_idx_for(&id).unwrap(); + } + writer.append_tensor(&rec).unwrap(); + } + writer.finalize().unwrap(); + } + cursor.into_inner() +} + +/// Rewrite every built-in Extension Table entry (file-level and inline-chain +/// tables) from version `from` to version `to`: the canonical id is +/// re-derived and the `io.ptwm.builtin.@` label is changed. The +/// file-level table hash in the header is recomputed. Returns the number of +/// entries rewritten. +fn rederive_builtins(blob: &mut [u8], from: &str, to: &str) -> usize { + assert_eq!(from.len(), to.len(), "labels must keep their length"); + if from == to { + return 0; + } + let header = ContainerReader::open_partial(blob).unwrap().header; + let et_off = header.extension_table_offset as usize; + let et_len = header.extension_table_length as usize; + + let prefix = b"io.ptwm.builtin."; + let mut count = 0; + let mut p = 0; + while let Some(pos) = find(&blob[p..], prefix) { + let label_start = p + pos; + // Entry layout: canonical_id [32] | label_len u16 | label. + let label_len = u16::from_le_bytes([blob[label_start - 2], blob[label_start - 1]]) as usize; + let label = std::str::from_utf8(&blob[label_start..label_start + label_len]) + .unwrap() + .to_string(); + let (name, ver) = label[prefix.len()..].split_once('@').unwrap(); + assert_eq!(ver, from, "unexpected builtin version in {label}"); + let id_start = label_start - 34; + let old = CanonicalId::derive(&BUILTIN_PUBKEY, name, from); + assert_eq!(&blob[id_start..id_start + 32], old.as_bytes(), "{label}"); + let new = CanonicalId::derive(&BUILTIN_PUBKEY, name, to); + blob[id_start..id_start + 32].copy_from_slice(new.as_bytes()); + let ver_start = label_start + label_len - from.len(); + blob[ver_start..label_start + label_len].copy_from_slice(to.as_bytes()); + count += 1; + p = label_start + label_len; + } + + // Header: magic[9] flags u16 | et_offset u64 | et_length u64 | et_hash [32]. + let hash_off = 9 + 2 + 8 + 8; + assert_eq!(&blob[hash_off..hash_off + 32], &header.extension_table_hash); + let hash = blake3::hash(&blob[et_off..et_off + et_len]); + blob[hash_off..hash_off + 32].copy_from_slice(hash.as_bytes()); + count +} + +fn find(hay: &[u8], needle: &[u8]) -> Option { + hay.windows(needle.len()).position(|w| w == needle) +} + +/// The v1.1.0 fixture: main and chunked containers merged and rewritten with +/// the v1.1.0 table layout, built-in ids pinned to version 1.1.0. +fn v1_1_0_container() -> Vec { + let main = main_container(); + let chunked_data = chunked_tensors(); + let chunked = compress(&inputs(&chunked_data), main_options(Some(256))); + let mut blob = rewrite_as_v1_1(&[&main, &chunked], INLINE_TENSOR); + rederive_builtins(&mut blob, env!("CARGO_PKG_VERSION"), "1.1.0"); + blob +} + +fn p4_bytes() -> Vec { + let mut rng = Rng(0x0050_3442_5954_4521); + (0..128).flat_map(|_| rng.weight().to_le_bytes()).collect() +} + +/// The P4 fixture: fp32 from `Format.BYTE` input. The v1.1 Python path +/// builds the chain with `chain_shape = [len(bytes)]`, so the fp32 Source +/// declares four times the real byte count. +fn p4_container() -> Vec { + let bytes = p4_bytes(); + let input = InputTensor { + name: "p4.byte.fp32".to_string(), + candidate_chains: vec![fp32_split_chain(&[bytes.len() as u64])], + dtype_code: 1, + // Format.BYTE + input_format: 1, + raw_bytes: &bytes, + shape: None, + dtype_name: None, + delta_reference_blake3: None, + }; + compress(&[input], main_options(None)) +} + +// ── Helpers ────────────────────────────────────────────────────────────────── + +fn fixture_dir() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures") +} + +fn load(name: &str) -> Vec { + let path = fixture_dir().join(name); + std::fs::read(&path).unwrap_or_else(|e| panic!("read {}: {e}", path.display())) +} + +fn expected(fixture: &str) -> &'static [(&'static str, usize, u64)] { + EXPECTED + .iter() + .find(|(f, _)| *f == fixture) + .unwrap_or_else(|| panic!("no EXPECTED entry for {fixture}")) + .1 +} + +fn assert_decodes(fixture: &str) { + let blob = load(fixture); + let reader = ContainerReader::open(&blob).unwrap_or_else(|e| panic!("{fixture}: open: {e}")); + let want = expected(fixture); + assert_eq!(reader.index.len(), want.len(), "{fixture}: tensor count"); + for (name, len, hash) in want { + let got = reader + .decode_tensor(name) + .unwrap_or_else(|e| panic!("{fixture}: decode {name}: {e}")); + assert_eq!(got.len(), *len, "{fixture}: {name} length"); + assert_eq!(xxh64(&got, 0), *hash, "{fixture}: {name} xxh64"); + } + // decode_model must agree with per-tensor decode. + let all: HashMap> = reader.decode_model().unwrap().into_iter().collect(); + for (name, len, hash) in want { + let got = &all[*name]; + assert_eq!( + (got.len(), xxh64(got, 0)), + (*len, *hash), + "{fixture}: {name}" + ); + } +} + +fn assert_transcode_round_trip(fixture: &str) { + let blob = load(fixture); + let exploded = explode(&blob).unwrap_or_else(|e| panic!("{fixture}: explode: {e}")); + let members: HashMap = exploded + .members + .iter() + .map(|(k, v)| (k.clone(), v.as_slice())) + .collect(); + let tensor_jsons: Vec = exploded + .tensors + .iter() + .map(|t| t.manifest_json.clone()) + .collect(); + let rebuilt = implode(&exploded.registry_json, &tensor_jsons, members) + .unwrap_or_else(|e| panic!("{fixture}: implode: {e}")); + assert!( + rebuilt == blob, + "{fixture}: implode(explode(x)) differs from x ({} vs {} bytes)", + rebuilt.len(), + blob.len() + ); +} + +// ── Tests ──────────────────────────────────────────────────────────────────── + +#[test] +fn fixtures_stay_small() { + let total: u64 = [V1_1_0, V1_0_0, TRIMMED, P4] + .iter() + .map(|f| load(f).len() as u64) + .sum(); + assert!(total < 64 * 1024, "fixtures total {total} B"); +} + +#[test] +fn v1_1_0_fixture_decodes() { + assert_decodes(V1_1_0); +} + +#[test] +fn v1_1_0_fixture_transcode_round_trip() { + assert_transcode_round_trip(V1_1_0); +} + +#[test] +fn v1_1_0_fixture_has_expected_features() { + let blob = load(V1_1_0); + let reader = ContainerReader::open(&blob).unwrap(); + assert!(!reader.prelude.is_empty(), "shared prelude state"); + let records = records_in_file_order(&reader); + assert!( + records.iter().any(|r| r.inline_chain.is_some()), + "inline chain" + ); + assert!( + records + .iter() + .flat_map(|r| &r.terminals) + .any(|p| p.chunk_table.is_some()), + "chunked plane" + ); + assert!( + records + .iter() + .flat_map(|r| &r.terminals) + .any(|p| p.state_source == ptwm_core::codec::StateSource::Shared), + "plane with shared state" + ); + assert!( + records.iter().any(|r| !r.dependencies.is_empty()), + "dependency" + ); + // The v1.1.0 table interns all 16 built-in codecs. + let codecs = reader + .extension_table + .entries + .iter() + .filter(|e| e.kind == ptwm_core::extension::Kind::PlaneCodec) + .count(); + assert_eq!(codecs, 16); +} + +#[test] +#[ignore = "cross-version builtins: fixed in stage 3"] +fn v1_0_0_fixture_decodes() { + assert_decodes(V1_0_0); +} + +#[test] +#[ignore = "cross-version builtins: fixed in stage 3"] +fn v1_0_0_fixture_transcode_round_trip() { + assert_transcode_round_trip(V1_0_0); +} + +#[test] +fn v1_0_0_fixture_differs_from_v1_1_0_only_in_builtin_ids() { + let old = load(V1_0_0); + let mut cur = load(V1_1_0); + assert_eq!(old.len(), cur.len()); + assert!(rederive_builtins(&mut cur, "1.1.0", "1.0.0") > 0); + assert!(cur == old); +} + +#[test] +fn trimmed_fixture_decodes() { + assert_decodes(TRIMMED); +} + +#[test] +fn trimmed_fixture_transcode_round_trip() { + assert_transcode_round_trip(TRIMMED); +} + +#[test] +fn p4_fixture_decodes() { + assert_decodes(P4); +} + +#[test] +fn p4_fixture_transcode_round_trip() { + assert_transcode_round_trip(P4); +} + +#[test] +fn p4_fixture_source_overstates_size() { + let blob = load(P4); + let reader = ContainerReader::open(&blob).unwrap(); + let rec = &records_in_file_order(&reader)[0]; + let chain = reader.chain_registry.lookup(rec.chain_ref).unwrap(); + let p = &chain.nodes[0].params; + assert_eq!(p[0], 1, "1-D Source"); + let dim = u32::from_le_bytes(p[1..5].try_into().unwrap()) as u64; + assert_eq!(u16::from_le_bytes([p[5], p[6]]), 0x0003, "fp32 Source"); + // dim fp32 elements = 4 * dim bytes, against orig_size bytes. + assert_eq!(dim, rec.orig_size, "Source declares 4x orig_size"); +} + +/// The legacy writer must keep producing the same bytes. A failure here +/// means a change altered `\x89PTWM` output; if the change is intended +/// (for example a crate version bump, which changes built-in ids), follow +/// the regeneration steps in the module docs. +#[test] +fn legacy_writer_output_is_golden() { + let blob = main_container(); + assert_eq!( + blake3::hash(&blob).to_hex().as_str(), + LEGACY_WRITER_GOLDEN_BLAKE3, + "compress_model legacy output changed" + ); + assert!( + trimmed_container() == load(TRIMMED), + "compress_model output != {TRIMMED}" + ); +} + +/// The v1.1.0 table layout (`set_chain_registry`) must keep producing the +/// same bytes. +#[test] +fn v1_1_0_layout_writer_output_is_golden() { + assert!(v1_1_0_container() == load(V1_1_0)); +} + +/// Writes the fixtures and prints the `EXPECTED` table and the golden hash. +#[test] +#[ignore = "regenerates tests/fixtures; run with PTWM_REGEN_FIXTURES=1"] +fn regenerate_fixtures() { + assert!( + std::env::var_os("PTWM_REGEN_FIXTURES").is_some(), + "set PTWM_REGEN_FIXTURES=1 to overwrite tests/fixtures" + ); + let v1_1_0 = v1_1_0_container(); + let mut v1_0_0 = v1_1_0.clone(); + assert!(rederive_builtins(&mut v1_0_0, "1.1.0", "1.0.0") > 0); + let trimmed = trimmed_container(); + let golden = main_container(); + let p4 = p4_container(); + + let dir = fixture_dir(); + std::fs::create_dir_all(&dir).unwrap(); + let files = [ + (V1_1_0, &v1_1_0), + (V1_0_0, &v1_0_0), + (TRIMMED, &trimmed), + (P4, &p4), + ]; + for (name, bytes) in files { + std::fs::write(dir.join(name), bytes).unwrap(); + } + + // Expected decoded bytes come from the inputs, not from a decode. + let mut main: Vec<(String, Vec)> = tensors() + .into_iter() + .map(|t| (t.name.to_string(), t.bytes)) + .collect(); + let trimmed_expected: Vec<(String, Vec)> = trimmed_tensors() + .into_iter() + .map(|t| (t.name.to_string(), t.bytes)) + .collect(); + main.extend( + chunked_tensors() + .into_iter() + .map(|t| (t.name.to_string(), t.bytes)), + ); + let p4_expected = vec![("p4.byte.fp32".to_string(), p4_bytes())]; + + println!(" // EXPECTED-BEGIN"); + for (fixture, list) in [ + ("V1_1_0", &main), + ("V1_0_0", &main), + ("TRIMMED", &trimmed_expected), + ("P4", &p4_expected), + ] { + println!(" (\n {fixture},\n &["); + for (name, bytes) in list.iter() { + println!( + " ({name:?}, {}, 0x{:016x}),", + bytes.len(), + xxh64(bytes, 0) + ); + } + println!(" ],\n ),"); + } + println!(" // EXPECTED-END"); + println!( + "const LEGACY_WRITER_GOLDEN_BLAKE3: &str = \"{}\";", + blake3::hash(&golden).to_hex() + ); + for (name, bytes) in files { + println!("{name}: {} B", bytes.len()); + } +} From f07bafccba2e01f84bdc023243c0090f2fda0367 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 21:05:49 +0000 Subject: [PATCH 03/34] fix(S1-fixtures): address review findings - Add legacy_p4_byte_fp32_v1_1_0.ptwm: the P4 tensor in the v1.1.0 layout (full Extension Table). It is byte-identical to v1.1.0 compress_model output. legacy_p4_byte_fp32.ptwm is now documented as current-writer (trimmed table) output. - Add legacy_v1_0_0_native.ptwm, written once by the v1.0.0 release (a17d403). Reword the legacy_v1_0_0.ptwm docs: it is the v1.1.0 fixture with ids re-derived for 1.0.0, not v1.0.0 writer output. - rederive_builtins parses the header by hand instead of calling ContainerReader::open_partial, so it works on blobs with built-in ids of other versions. Document the equal-length version limit. - Test explode->implode and decode on in-memory main_container() output (trimmed table plus shared Order1ScaleAC state). - Fix clippy cloned_ref_to_slice_refs in transcode tests. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/transcode/tests.rs | 2 +- .../fixtures/legacy_p4_byte_fp32_v1_1_0.ptwm | Bin 0 -> 3297 bytes .../tests/fixtures/legacy_v1_0_0_native.ptwm | Bin 0 -> 9530 bytes crates/ptwm-core/tests/legacy_fixtures.rs | 246 +++++++++++++++--- 4 files changed, 209 insertions(+), 39 deletions(-) create mode 100644 crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32_v1_1_0.ptwm create mode 100644 crates/ptwm-core/tests/fixtures/legacy_v1_0_0_native.ptwm diff --git a/crates/ptwm-core/src/transcode/tests.rs b/crates/ptwm-core/src/transcode/tests.rs index b16a6e5..d53e4ba 100644 --- a/crates/ptwm-core/src/transcode/tests.rs +++ b/crates/ptwm-core/src/transcode/tests.rs @@ -345,7 +345,7 @@ fn round_trip_trimmed_codec_table() { .unwrap(); let rebuilt = implode( &exploded.registry_json, - &[b.manifest_json.clone()], + std::slice::from_ref(&b.manifest_json), members_map(&exploded), ) .unwrap(); diff --git a/crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32_v1_1_0.ptwm b/crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32_v1_1_0.ptwm new file mode 100644 index 0000000000000000000000000000000000000000..00ff06f1b96072133c171af0d403628a77d6b9ce GIT binary patch literal 3297 zcmeAS2nqM)<&xrJV2Eabfa9DHdg>h^hs4`_MGqGqGySzObnnzxfkj{6?325+dF%Gp z6wYXs5r) z$2s;g44L_Q1tsOVdP$|3IVG8Sdd2yrMaiiShI)p228;|q$fO7~XHj-aYDsEQZf0I) zaqic=jKrdp^2DOl_~e|#;^O4wf>K7Pr6|;HoA|V2*SHjymS?KnYff0P_jJr^k-fF| zx4hja_qF4XyatBtNtq?_MXC8kDXB&AnW?F%#zxd}$?q+*B6nU*Gu>j||E5A^fs767 z;oZ~zHfdRZvO0U{O?9pUrb{YIQsavYaxzP(Dhv9 z&Ly5%zBt2ErO#Lv!~TrYw6xsBJnC3)d7||F^vNrazJ54UF!%KupSi&uJGP(Kf3~%> zC7x$tlAjcY^+kz!#niDowT+ii{p_Q4BEefi7HoYUqx<^MHw(Qf@)-UG*&UyfnG8*L zl*aqs$4meHkM0PQIhPRQ_&lyduuG_7$6ID57uB673le2*B{3Y3R**~`zdvbfb$z}` zsVOtE$VFtzbJon&>XEUMCKE+o1PR4XyR`%}+)Jny=iDu750({N6#3-$B+k@i2q!Vt>@u<9Es%l?ye z;-@Y?yR9W=n7Q}zbzb~IVo04Jxm&+~(eB;s!K>Q-yttUhl9XW8cX!uvrPFq^pLEX` zzs7)BGZg@9r1YZv(t`Np{FKzB{QPX{xFlT3{QAe78-nF^ul-#X`mJ1JzM`XBt!qxy z>Dv*aW<_RRdOWa(qEgJf`tQ!l*5kUAuS%9}_16sH1Fdfu zFW&m{zT(EyrC%#}?J;8pR7J$+@aVn zAFYeSi;ac%O7T8__1yT^<=5?(+n+P^%Fa!%VsKxs{%g5Z^{rVKPBF~g?0$ha;jH>| zh6di}2EpkLZ-v`^+tVAY5}JE&8y{G{{1)%=^z_f^s~Vy}H=YFdMaT4*Hn zT196w^Ox?q?6p^_UW9A=zEV29DEr#nPIEnW=~IdCH$8TXc@(^U&iw7YicV3UUG_KS z`r^|%-tSF)`B_Es?vsy4!%vC**b>llD(BR*`tXK-*qxSdW1-7%z*A;Xb9Dkw*lNbM zyP0MpGL7r^c6I6;(O)s)>wDIM+l@;!1sM)}|0b~I;vTlN2RlUCS4^I48Yr=MH{aTg z40fScUDxR>4PLQo!TPmk>}KYdEiX=Bh%MwZ$rgT^AeFYRQFq@;jw4ws0%yE1+{tRM z!mDf?v)1V1EkV)8SDecW?Kzm)SQyyANe2{{&d>-oGWMR#ps`i|FGIvtMuwu}+}syL ztu$44$}=!{G3RH^y?$lV!W`CHtKZmVn4ga5uJ4WcdFyc1yi1uMJ}*o>-%!NF&@ri& hDS-{v;NMdE@K_$>43!jU?;{kN$U1@j8drZeDF8Q#H}n7i literal 0 HcmV?d00001 diff --git a/crates/ptwm-core/tests/fixtures/legacy_v1_0_0_native.ptwm b/crates/ptwm-core/tests/fixtures/legacy_v1_0_0_native.ptwm new file mode 100644 index 0000000000000000000000000000000000000000..24b6f480ef0ee8611f019d22d92332ab69a2c483 GIT binary patch literal 9530 zcmc&)2{@Ep`+sJ!FWI*&W9%xskjP$0wh+oN%vfiLnb2nJTb4wYLY8ECLnMW45h10b zC|Zy;q_WHZ8K(Drdwu`!`(M}Ry8icbJB1llY+R7m0yb)o}LUgO&hpi7liDBzS9hAN9mX@o5(kQ;jFx+(hk$^PM4+6)hT&`l*cnOQFElWJkSuf1 zT_g>91l2h4sTLb4GmS2!U*LTZX#F}}{9o(B2edT7&DWdg>b5g-Ov!)w>D^;pNe5xP zWoE;TU5oiVLgm?fH&@a%9rwX?k*TS!JQDe#UJ`#dZ&$GLh11YrmB@9r&Uai zL$03sTfvlpM$-m?Q>n=KC;9^0_`4CYSez4nXV^!!cTgRANNvz~foH%3cC|>%Yzn7W z_iDt+G8!6bV%)~JjlHiE-fw4kFRCBpR{P;6`{UAO_p6$N@TTDNNF&#_@$SPBKWYrq zJht;b>ql_e8Qy+WZ*Xj8Yc`eF52)6#v8S`?Usf0@mgZ!-^>jCUxL!g4I;-*O1j}j_E;{J2z%M;HrGlD!b;QxM=IX9HMiTb1bgx> z!{Lvok}JYB1(B?(ZNg&Pmc5yV$B0?QbHwF@M2M+$Tl&YeBB9dI1Itv8RadvGj{N-H991K>h}rnTGeG7IfZd`)`+^m z{&XHQ+gvMVfI6}F-wT0qa&ht@xSyeHlmD$lGPOaPuT`~Hu83YxYLmH}zQS_o2Sy;B zZ=gf02Pt{m=HIJCerHtjau45zTM#s0a)DHFPgd;C8pCnN54?X&9LioKjLMwl+`d2g zfP1tnxF7hSoV{HzXm4-Nospz%R%uSVoCOGLD)KKKv%&S2dPH&eEPir0e5EQy2rD49 z%{b29cmgJXfWo;4xZ_<>;GVrhmMMIbSy~Z}HCS(5I-e^=7x+29j&0z%QJhNrT_(9h zO7q(-L*61#US2qqG6-dMMjv-c5ha6z)U5T4i9DEx${mrFOsvm8+1d3ESWdngulJMK zrVmQs{uMYoBgm}ul=!mjinoTfhqnfkD7g zaPlZNu#o`@2cHP=hJhD(6!`t0qu@{o6at5W1c$<4;4Bab6F>t{fum^102vg8Jj<_f z5C|=7s}J2~JJ<&l?aw}x*#Ypk-^BkV$bS%@>i1DBAYTM{|A$dz{!l2WvQ1S005^aG zBi2jC%PA1!>jxtJ(=z@TcULz8u)&MVN89de=L09r#rU~7`Cvk5c@P*ELZA=E4UP5k zb|T2{rvj%Uub)DjYXT?lJoOp44q)_^srTJ|82tC6+uT8hzR*0megG zK_r*OPW4QqeS^rM`latTJ^g z;i~FGF}CF7*lucsRP3?ZT9m>U+{F2bo`vTVphRcoKVwk(f3%~ zpg~brR;a%`oTNmByD3?sJ1W7m9E_;AYFHf?nw<3fd+dbBY#+iOHQKjG6OF7&U(2(7 zDJ4$m2p>&y7&)9N5dYa%lE*otWmHK?u`Z_B?p7Aoymh7bJJ)Ss0-i7WG1oG+mlw|7 ze?hskv3%iKX>QKm_cn(HwOwJ<1*PWY>WcI2h2}#m?M#^EpIH z?1AC%#>Pe`RKT-iSNE>R-b-}Lq?*Ym^QhyB@pSE9@ovkinpSk`>x(>RH$tyHe(7-s zT9j5SufY(-koEOmQ@P_;JB@etYZ`QbF|De+%0U5H=lOcSKl&K$y5M3Zv6gmJpqb$R z-u9lqY%em)EnB0`?Z*oWz0LNn|JL-!W=r1!mUs5S)`nJirr%zDZ=@W(DiH=3WlMq7 zUf*kiKCPB$?y}%odQ{`S1!thd$dihUSI@g$$Bch85=xbf=1*ems$7kpOGqT%8shg) zj7+DMSUzGMhZIUWqH|)>42)5-9x7ggv)T>gw!Mry=8hCPce1?Bs0sp~k@7@14+^{p66cBo`8 zB}K74zpoaSBSJvBgWrXUDJmog(% zC+sc4y-XCoS|4Hwts%*9t6nkyy~Y{37mxs?0P~N!lI{20Bd($#u4|@HGXM9dq{Q;S8i+onSZgl%i zu5osjr##b4fQY*M4Oc;zGQg$>8*JTB&apoVlZ?ei!bXpiy`=ts&Sgtl%ir?XVPkQt%#YbLJHq*eT z*uG->*8anAXo%RToIQDD%G{dRXacRya`7KGhhV;YvQDt5QgK$tks;OSfDOPqg=Ej853)9J{Hv@bSQt;Cds^5~02Ssb_ z-oCRGk7EM2q9fegiDEby^5KgB|M%Cy*(Toa@zx%bKZy=od|h&5C= zq%5HOwIjV=S7p3q-qMqXfgqD7bm!rSqjixliauX|mH<>owBpB@jZlaU4bZE*J+ zdH_aE12FuV@gyli@SdT43_!?+WPkz}KfrR(R9qr5t-MDsBW+_nFoYh1=3JX9uVzv7A7k;2%{x?A)( zNv=_?pYrCOBAo4p)u*Mmp|ii&xWGlBFiJSZEnxKD*4v&jcY882CD26lwy-{tY`*|` z@xrc6_C}6Hg<*X}_<#se?L^T}SKp6|!KQ_U~riK5+WZ>2WfkT8q zBW~MpaGM*>mUggs7&QdMFckLyTq&MZ$TYGWYy_L50EB`92iQ!$$s7Gq!d;iv@7pz;&a)lLd2OH9wPxXeCQjgZe!tGK-0~$+!$YOe zi-F&1tb9H34Vp=Oc1K;WZC0KxzQ z;D1f5F}cC5sW_J6^n$rFP>@7D{AcaA=H+j60s78u5}i zH?GdRzy+>`VuWB`2(X3tT*y5jzpXrCbKL-NkXe5V{%*mVFEDU#6QrM0bN{)Kplq0K z!3U79#G7}#tv=FfpBC3&suO!36`3F@T=HpM)rWex_NQwjI+r_=4#>!wMy6wy%SXki zW4jv|F1TEha;zNgQR4CFF?*zGgu+feK3Mn4Xq4KJw(%31Mbz?jY-3q)#``9Zt20hF zeDGr~Y#@B6ER=i>0Ez3LDde=wW*q|nNJ=7XqsW33Zm`!cujA4BwIOY;82a0a6Uj*i zq|t9pJP-uK3b`yY4KXA%YA=T^%a8@qjC~B3Ah1u?`8%u zRv6baSZri^O=nlv)a>Awr-NN-#kTy;9ku7q)mAsC%ubeK%{f}7;?)Z38fptli#-Yy zre@|n&;@deo|=q!^hiUFm|y;vGISJ@D5fVk6URk1or$222TLr3{Ms`&X~7UUTl2 zFF5Zi2!ZmqIdQZp&EnGs#jum9xq~G|?$*ckF&JxWYrW*cn-6Q>4mPZ%`a7IL>ac|I z(CY@9NixEP^&}lhBp?!8t=eNAzj(TUsDRfqQK`kZ!HBU9R9$*H>+Qnn#8k`o8dSjmiG)D#(86KOj&cYN+OTG=-oYU^;?LGJi=Xe8 zFWCC`Gv&#}J9n?HN-++rUYaipNaG<*7G5s0#2V=7StN>1AfP(5+VnqISlHBPBG>z& z)Aiw~sDjVzsdFm5J$y?377g$ANek>l|=it z(awxMwzd7`tBh0G1XxLLBz*6Y@I@iL28N$?J%uk)>3{eaa+&hfPtzzF zHc*8%rW5r;=|aLexl7eAin7OPo>(j)6_|9&7H{18)O-Ky2Bs%naX-I2pUN=bo}`K} zMn2wSv&UM!#H*31ht;_{j$C z%vwu{hOTzpnMlLS#wG(E;({VqG6D~bvZy7iWhj_8Wxrd7{Kwe${1y8GA?P^qJiUSc zbL^ujvCnjXn*VRH|CA5-EN;t{DE0)lVgEqKh|VUdFJo|qOf@_ffH`0WrU{5Ptvoev zamB{PbB$VZJdD)PCNm>{si(Dd85t*(T&#~<>0>r#Yr9H{OELo5-e=d;biNzxx>i$s zZt=a-r$Kzf>c*p2$9whabX#iGTULdz`UC|f>xH+si;KICDd|RZbWAS_r3c!u&$4(2 zRldiFI~2^z2wC?EW3fo_A$b*9S(VutG(9>>m|8fDQNEoqomU!(d^b8{VV$}=k_8G0 zVPOdi;Iv9%>gi92iHI=rw;|Rlsd#@{UwpK>z7f>GJM^^6A;aEMMOo2`Ju;EEhBGCl zpDEOlg%Kk;`0ggJjU~J6x5!A0m9C!Rnbq5SLlD%w&QaQm$!^57YfIdo1qHllq_8fn zFbfM@`;v5V&Rq|Wp`wmyIIPKx%iWEZ(`1ieQjUCWw0ujV@bQ$W=qM#UtWC>#slgck zFEgKeD{AHSNUMc8V$NQAdPyOoDIdvkD^EWIVPP>>p*f$>>-kf*#V2QMFDYAn`H&ng z{l&$Nju2zvSoVxznY|`7*O%cv<7~>wVls$_L4Hxgn)^+@=KtU&ls9=D!lK-eUmF8n zhgKZeZJ;o&JAFg@1lnI5V6q@ z<##R@=e+X55YIbemajoQuMakWJ&xvzc_|G_i#%ICa*ONF#|+E#+7 zgy!VGU?S(~m7Z)oN~`H#BgLZ-?Lwv!5R#VYX*P!a`O6&Lhx2yF(cAaDv?@Q@`f5*` z`6pAp=+tvD=CAnD1NLRatLNlqeZgKQh~Et?OKy!UL=hO{P~OIDjDRq8&-Z&5I8ly! zM#OWUpZbV{KMs4{L&v`xwby1+Kup0&>LL)T+ZhGGY!(mCe!5>2M`J?c+C?{%eCN1t z2trDe>g15yYggr!)_85JPeSr*4)=Vf5&tIXgB}kh0{*kR2{rGUydHzE?6L z=kF!0KhK~4OG%3ahA){uq)jUTfGQ`+c^+Vr0(jJcT@54zngoRbN?!n2ABjntBTTAQ5FK zpku(q+{pO?Vo`994m{wYQe8B`@;tQKTTfhDr_tX}F&WeT;8J3_{m^iGvMW%*=GPIj z7o|x92Sy58IcET?d&OQ!5UO8w!p(vv0Q}<|rXb_&M#Ou9z>!id1mX6zqDocE56gd7 zF3T@0^_Ox%YyMR)q^DFNk3Dpgro3(%n~IE4DL56ZgeY51cYcKiJ|+@nxq{S0XRaQ+ zADdtGe4WAh`=e#`D!8fm`bq@POa$j*YMIsUT`0A3Ds1pBpbe4Dv!*C&tvL;Cc3?gq z9(6NyH})LJL=Y;_<@LID$~%AO*0o3?z`Vuc^~p7Ype<#d$FbF zZdzJRaVu{Y(mS&E;^VtE%s+ybO+P4i7L*<+YhsJXD~n#9n&ui?_BJ9-j9%T-b@;VI zXQ;J$ssro05UYw1EG%+=^@i$G)q`J!$C{Q}LqD#;W#;Gof@Oug|7hOV+9j@i&DNyD zL~X%tHNnp_X;6l@b! zBX{`IaW>C`40M1W3qU{sL^I$F3Gjad1SCQN&7fxi=(!_6CS^2>T;Q^F5fZkK1%|q# ztv{mpFiF%kSbfAlh9IgkTzfx8!?CAf{3-F{Opj!KtEvjV^uRLjPx@#67T@?bB-CiX zC+5QhRoa2rG1!BhN)~9XU$k$iPF)E{E4EJud;oOZ1_Hvxa zdO2<2Y+H*?xo7mQO|T|}e#Iwn!uJEW*f^YyqAm$%#rJu=<=jrpTMhOUhJO0bWrMKm zdr{;QA4n5h9{{^KS9i-tJryeQ(HSH=3iMfle#wyD$Wzzbz17}LY5(kZQDfPUUJt{# z)xJitz*AHF$h+y?Zgu28PJm=c;X7ZhRp-o9c_j|9*$lK&=IJ>X5F&a~?jswy3;@zO f$~-sITr0h=dwo4Z25TU#QTj)N|1hO*YQXnD5!pL$ literal 0 HcmV?d00001 diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs index 1edc801..c08aa91 100644 --- a/crates/ptwm-core/tests/legacy_fixtures.rs +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -5,14 +5,18 @@ //! //! | file | content | //! |---|---| -//! | `legacy_v1_1_0.ptwm` | Multi-tensor container with the v1.1.0 Extension Table layout (every built-in codec interned, `ContainerWriter::set_chain_registry`). Holds bf16, fp16, fp32, fp8 and int8 tensors, two scale tensors that share an `Order1ScaleAC` prelude state, a chunked bf16 tensor, an fp8 tensor with an inline chain, and a tensor with a dependency. Built-in ids and labels are derived for version `1.1.0`. | -//! | `legacy_v1_0_0.ptwm` | The same file with every built-in id and label re-derived for version `1.0.0` and the Extension Table hash fixed up, as a v1.0.0 writer would have produced it. | +//! | `legacy_v1_1_0.ptwm` | Multi-tensor container with the v1.1.0 Extension Table layout (every built-in codec interned, `ContainerWriter::set_chain_registry`). Holds bf16, fp16, fp32, fp8 and int8 tensors, two scale tensors that share an `Order1ScaleAC` prelude state, a chunked bf16 tensor, an fp8 tensor with an inline chain, and a tensor with a dependency. Built by [`rewrite_as_v1_1`] from current `compress_model` output; built-in ids and labels are derived for version `1.1.0`. Byte-identical to what this generator writes when built against the v1.1.0 release, whose `compress_model` emits the same layout. | +//! | `legacy_v1_0_0.ptwm` | The v1.1.0 fixture with every built-in id and label re-derived for version `1.0.0` and the Extension Table hash fixed up. It is not v1.0.0 writer output: for these inputs the v1.0.0 `compress_model` writes no shared `Order1ScaleAC` prelude state. It covers v1.0.0 ids on a file that has the v1.1.0 features. | +//! | `legacy_v1_0_0_native.ptwm` | Written once by the v1.0.0 release (commit `a17d403`): this generator, built against that tree, run as `rewrite_as_v1_1(&[main, chunked], INLINE_TENSOR)`. Real v1.0.0 writer output, with v1.0.0 built-in ids. `regenerate_fixtures` does not rewrite it. | //! | `legacy_trimmed.ptwm` | `compress_model` output of the current writer, whose Extension Table holds only the codecs the planes use. Same tensors as the v1.1.0 file minus the two scale tensors and the chunked tensor. | -//! | `legacy_p4_byte_fp32.ptwm` | An fp32 tensor written from `Format.BYTE` input by the v1.1 Python path: its Source node declares `[n_bytes]` fp32 elements, four times the real size. | +//! | `legacy_p4_byte_fp32_v1_1_0.ptwm` | An fp32 tensor from `Format.BYTE` input, with the chain built the way the v1.1 Python path builds it: the Source node declares `[n_bytes]` fp32 elements, four times the real size. v1.1.0 layout (full Extension Table), built-in ids for `1.1.0`. This is the layout of P4 files written by v1.1.x; it is byte-identical to the v1.1.0 `compress_model` output for the same input. | +//! | `legacy_p4_byte_fp32.ptwm` | The same tensor written by the current writer (trimmed Extension Table). | //! //! The tests check that every fixture decodes to the expected bytes and that //! `implode(explode(x)) == x`. The expected bytes are stored as -//! `(length, xxh64)` pairs in [`EXPECTED`]. +//! `(length, xxh64)` pairs in [`EXPECTED`]. Fixtures with built-in ids of +//! another version are checked after [`rederive_builtins`] maps the ids to +//! the crate version, until the reader accepts older built-in ids. //! //! # Regenerating //! @@ -25,9 +29,14 @@ //! golden hash for [`LEGACY_WRITER_GOLDEN_BLAKE3`]. Regenerate only when a //! format change is intended: the point of the fixtures is that they do not //! change. A crate version bump changes the built-in ids that the current -//! writer emits, so it changes `legacy_trimmed.ptwm` and the golden hash; -//! the v1.1.0 and v1.0.0 fixtures are pinned to their versions and do not -//! change. +//! writer emits, so it changes `legacy_trimmed.ptwm`, +//! `legacy_p4_byte_fp32.ptwm` and the golden hash; the fixtures pinned to +//! v1.1.0 and v1.0.0 do not change. +//! +//! [`rederive_builtins`] rewrites the version inside the built-in labels in +//! place, so it only maps between versions whose strings have the same +//! length (`1.1.0` and `1.0.0`, but not `1.10.0`). A bump to a version of +//! another length must extend it first. use std::collections::HashMap; use std::io::Cursor; @@ -46,7 +55,9 @@ use xxhash_rust::xxh64::xxh64; const V1_1_0: &str = "legacy_v1_1_0.ptwm"; const V1_0_0: &str = "legacy_v1_0_0.ptwm"; const TRIMMED: &str = "legacy_trimmed.ptwm"; +const V1_0_0_NATIVE: &str = "legacy_v1_0_0_native.ptwm"; const P4: &str = "legacy_p4_byte_fp32.ptwm"; +const P4_V1_1_0: &str = "legacy_p4_byte_fp32_v1_1_0.ptwm"; /// Expected decoded bytes per fixture: `(tensor name, length, xxh64 seed 0)`. /// Printed by `regenerate_fixtures`. @@ -80,6 +91,20 @@ const EXPECTED: &[(&str, &[(&str, usize, u64)])] = &[ ("model.layers.2.chunked.weight", 1024, 0xf9e05338b3a99b14), ], ), + ( + V1_0_0_NATIVE, + &[ + ("model.layers.0.q.weight", 512, 0xc82d5962a46d19ff), + ("model.layers.0.k.weight", 512, 0x73173bba2f17ef23), + ("model.norm.weight", 512, 0x946ca638752e3bfc), + ("model.layers.0.mlp.fp8.weight", 256, 0x84e10734e3fb3518), + ("model.embed.int8", 256, 0xf6193be8f075d386), + ("model.layers.0.q.scale", 2048, 0x95a622e512233480), + ("model.layers.0.k.scale", 2048, 0xda44a3268d62821d), + ("model.layers.1.q.weight", 256, 0x0d7f3dd0787bd569), + ("model.layers.2.chunked.weight", 1024, 0xf9e05338b3a99b14), + ], + ), ( TRIMMED, &[ @@ -92,6 +117,7 @@ const EXPECTED: &[(&str, &[(&str, usize, u64)])] = &[ ], ), (P4, &[("p4.byte.fp32", 512, 0x57e3af5eb7d32ec1)]), + (P4_V1_1_0, &[("p4.byte.fp32", 512, 0x57e3af5eb7d32ec1)]), // EXPECTED-END ]; @@ -556,14 +582,29 @@ fn rewrite_as_v1_1(parts: &[&[u8]], inline_tensor: &str) -> Vec { /// re-derived and the `io.ptwm.builtin.@` label is changed. The /// file-level table hash in the header is recomputed. Returns the number of /// entries rewritten. +/// +/// The header is parsed by hand, not with `ContainerReader`, because the +/// reader rejects built-in ids of other versions. The version is replaced +/// in place, so `from` and `to` must have the same length. fn rederive_builtins(blob: &mut [u8], from: &str, to: &str) -> usize { - assert_eq!(from.len(), to.len(), "labels must keep their length"); + assert_eq!( + from.len(), + to.len(), + "rederive_builtins only maps between versions of equal length" + ); if from == to { return 0; } - let header = ContainerReader::open_partial(blob).unwrap().header; - let et_off = header.extension_table_offset as usize; - let et_len = header.extension_table_length as usize; + // Header: magic[9] flags u16 | et_offset u64 | et_length u64 | et_hash [32]. + let u64_at = |off: usize| u64::from_le_bytes(blob[off..off + 8].try_into().unwrap()); + let et_off = u64_at(9 + 2) as usize; + let et_len = u64_at(9 + 2 + 8) as usize; + let hash_off = 9 + 2 + 8 + 8; + assert_eq!( + &blob[hash_off..hash_off + 32], + blake3::hash(&blob[et_off..et_off + et_len]).as_bytes(), + "Extension Table hash" + ); let prefix = b"io.ptwm.builtin."; let mut count = 0; @@ -588,9 +629,6 @@ fn rederive_builtins(blob: &mut [u8], from: &str, to: &str) -> usize { p = label_start + label_len; } - // Header: magic[9] flags u16 | et_offset u64 | et_length u64 | et_hash [32]. - let hash_off = 9 + 2 + 8 + 8; - assert_eq!(&blob[hash_off..hash_off + 32], &header.extension_table_hash); let hash = blake3::hash(&blob[et_off..et_off + et_len]); blob[hash_off..hash_off + 32].copy_from_slice(hash.as_bytes()); count @@ -635,6 +673,13 @@ fn p4_container() -> Vec { compress(&[input], main_options(None)) } +/// The P4 tensor in the v1.1.0 table layout, built-in ids pinned to 1.1.0. +fn p4_v1_1_0_container() -> Vec { + let mut blob = rewrite_as_v1_1(&[&p4_container()], ""); + rederive_builtins(&mut blob, env!("CARGO_PKG_VERSION"), "1.1.0"); + blob +} + // ── Helpers ────────────────────────────────────────────────────────────────── fn fixture_dir() -> PathBuf { @@ -654,10 +699,19 @@ fn expected(fixture: &str) -> &'static [(&'static str, usize, u64)] { .1 } +/// `blob` with built-in ids of version `from` re-derived for the crate +/// version, so the current reader accepts it. +fn as_current_version(mut blob: Vec, from: &str) -> Vec { + rederive_builtins(&mut blob, from, env!("CARGO_PKG_VERSION")); + blob +} + fn assert_decodes(fixture: &str) { - let blob = load(fixture); - let reader = ContainerReader::open(&blob).unwrap_or_else(|e| panic!("{fixture}: open: {e}")); - let want = expected(fixture); + assert_blob_decodes(fixture, &load(fixture), expected(fixture)); +} + +fn assert_blob_decodes(fixture: &str, blob: &[u8], want: &[(&str, usize, u64)]) { + let reader = ContainerReader::open(blob).unwrap_or_else(|e| panic!("{fixture}: open: {e}")); assert_eq!(reader.index.len(), want.len(), "{fixture}: tensor count"); for (name, len, hash) in want { let got = reader @@ -678,9 +732,21 @@ fn assert_decodes(fixture: &str) { } } +fn plane_codec_count(reader: &ContainerReader<'_>) -> usize { + reader + .extension_table + .entries + .iter() + .filter(|e| e.kind == ptwm_core::extension::Kind::PlaneCodec) + .count() +} + fn assert_transcode_round_trip(fixture: &str) { - let blob = load(fixture); - let exploded = explode(&blob).unwrap_or_else(|e| panic!("{fixture}: explode: {e}")); + assert_blob_transcode_round_trip(fixture, &load(fixture)); +} + +fn assert_blob_transcode_round_trip(fixture: &str, blob: &[u8]) { + let exploded = explode(blob).unwrap_or_else(|e| panic!("{fixture}: explode: {e}")); let members: HashMap = exploded .members .iter() @@ -705,11 +771,11 @@ fn assert_transcode_round_trip(fixture: &str) { #[test] fn fixtures_stay_small() { - let total: u64 = [V1_1_0, V1_0_0, TRIMMED, P4] + let total: u64 = [V1_1_0, V1_0_0, V1_0_0_NATIVE, TRIMMED, P4, P4_V1_1_0] .iter() .map(|f| load(f).len() as u64) .sum(); - assert!(total < 64 * 1024, "fixtures total {total} B"); + assert!(total < 96 * 1024, "fixtures total {total} B"); } #[test] @@ -751,13 +817,7 @@ fn v1_1_0_fixture_has_expected_features() { "dependency" ); // The v1.1.0 table interns all 16 built-in codecs. - let codecs = reader - .extension_table - .entries - .iter() - .filter(|e| e.kind == ptwm_core::extension::Kind::PlaneCodec) - .count(); - assert_eq!(codecs, 16); + assert_eq!(plane_codec_count(&reader), 16); } #[test] @@ -803,16 +863,121 @@ fn p4_fixture_transcode_round_trip() { #[test] fn p4_fixture_source_overstates_size() { - let blob = load(P4); + for fixture in [P4, P4_V1_1_0] { + let blob = load(fixture); + let reader = ContainerReader::open(&blob).unwrap(); + let rec = &records_in_file_order(&reader)[0]; + let chain = reader.chain_registry.lookup(rec.chain_ref).unwrap(); + let p = &chain.nodes[0].params; + assert_eq!(p[0], 1, "{fixture}: 1-D Source"); + let dim = u32::from_le_bytes(p[1..5].try_into().unwrap()) as u64; + assert_eq!(u16::from_le_bytes([p[5], p[6]]), 0x0003, "{fixture}: fp32"); + // dim fp32 elements = 4 * dim bytes, against orig_size bytes. + assert_eq!(dim, rec.orig_size, "{fixture}: Source declares 4x size"); + } +} + +#[test] +fn p4_v1_1_0_fixture_decodes() { + assert_decodes(P4_V1_1_0); +} + +#[test] +fn p4_v1_1_0_fixture_transcode_round_trip() { + assert_transcode_round_trip(P4_V1_1_0); +} + +#[test] +fn p4_v1_1_0_fixture_has_full_table() { + let blob = load(P4_V1_1_0); let reader = ContainerReader::open(&blob).unwrap(); - let rec = &records_in_file_order(&reader)[0]; - let chain = reader.chain_registry.lookup(rec.chain_ref).unwrap(); - let p = &chain.nodes[0].params; - assert_eq!(p[0], 1, "1-D Source"); - let dim = u32::from_le_bytes(p[1..5].try_into().unwrap()) as u64; - assert_eq!(u16::from_le_bytes([p[5], p[6]]), 0x0003, "fp32 Source"); - // dim fp32 elements = 4 * dim bytes, against orig_size bytes. - assert_eq!(dim, rec.orig_size, "Source declares 4x orig_size"); + assert_eq!(plane_codec_count(&reader), 16); + let trimmed = load(P4); + let reader = ContainerReader::open(&trimmed).unwrap(); + assert!(plane_codec_count(&reader) < 16); +} + +/// The v1.1.0 table layout must keep producing the same bytes for the P4 +/// input. +#[test] +fn p4_v1_1_0_layout_writer_output_is_golden() { + assert!(p4_v1_1_0_container() == load(P4_V1_1_0)); +} + +#[test] +#[ignore = "cross-version builtins: fixed in stage 3"] +fn v1_0_0_native_fixture_decodes() { + assert_decodes(V1_0_0_NATIVE); +} + +#[test] +#[ignore = "cross-version builtins: fixed in stage 3"] +fn v1_0_0_native_fixture_transcode_round_trip() { + assert_transcode_round_trip(V1_0_0_NATIVE); +} + +/// The v1.0.0 writer output decodes and transcodes byte-exactly once its +/// built-in ids are mapped to the crate version. +#[test] +fn v1_0_0_native_fixture_decodes_with_current_ids() { + let blob = load(V1_0_0_NATIVE); + let mut cur = blob.clone(); + assert!(rederive_builtins(&mut cur, "1.0.0", env!("CARGO_PKG_VERSION")) > 0); + let reader = ContainerReader::open(&cur).unwrap(); + assert!( + reader.prelude.is_empty(), + "v1.0.0 writes no shared state here" + ); + assert_eq!(plane_codec_count(&reader), 16); + assert_blob_decodes(V1_0_0_NATIVE, &cur, expected(V1_0_0_NATIVE)); + assert_blob_transcode_round_trip(V1_0_0_NATIVE, &cur); + assert!(as_current_version(blob, "1.0.0") == cur); +} + +/// `rederive_builtins` must work on blobs whose built-in ids are not those +/// of the crate version, in both directions. +#[test] +fn rederive_builtins_handles_foreign_versions() { + let orig = load(V1_0_0_NATIVE); + let mut blob = orig.clone(); + let n = rederive_builtins(&mut blob, "1.0.0", "1.2.0"); + assert!(n > 0); + assert_eq!(rederive_builtins(&mut blob, "1.2.0", "1.0.0"), n); + assert!(blob == orig); +} + +#[test] +#[should_panic(expected = "versions of equal length")] +fn rederive_builtins_rejects_length_change() { + let mut blob = load(V1_1_0); + rederive_builtins(&mut blob, "1.1.0", "1.10.0"); +} + +/// Current writer output with a trimmed table and shared `Order1ScaleAC` +/// prelude state (the default for models with scale planes) decodes to the +/// inputs and transcodes byte-exactly. No fixture file: the bytes come from +/// [`main_container`], whose hash is pinned by +/// [`legacy_writer_output_is_golden`]. +#[test] +fn current_writer_shared_state_transcode_round_trip() { + let blob = main_container(); + let reader = ContainerReader::open(&blob).unwrap(); + assert!(!reader.prelude.is_empty(), "shared prelude state"); + assert!(plane_codec_count(&reader) < 16, "trimmed table"); + assert!( + records_in_file_order(&reader) + .iter() + .flat_map(|r| &r.terminals) + .any(|p| p.state_source == ptwm_core::codec::StateSource::Shared), + "plane with shared state" + ); + let data = tensors(); + let want: Vec<(&str, usize, u64)> = data + .iter() + .map(|t| (t.name, t.bytes.len(), xxh64(&t.bytes, 0))) + .collect(); + assert_blob_decodes("main_container", &blob, &want); + assert_blob_transcode_round_trip("main_container", &blob); } /// The legacy writer must keep producing the same bytes. A failure here @@ -840,7 +1005,8 @@ fn v1_1_0_layout_writer_output_is_golden() { assert!(v1_1_0_container() == load(V1_1_0)); } -/// Writes the fixtures and prints the `EXPECTED` table and the golden hash. +/// Writes the fixtures (all but `legacy_v1_0_0_native.ptwm`, which only the +/// v1.0.0 release can write) and prints the `EXPECTED` table and the golden hash. #[test] #[ignore = "regenerates tests/fixtures; run with PTWM_REGEN_FIXTURES=1"] fn regenerate_fixtures() { @@ -854,6 +1020,7 @@ fn regenerate_fixtures() { let trimmed = trimmed_container(); let golden = main_container(); let p4 = p4_container(); + let p4_v1_1_0 = p4_v1_1_0_container(); let dir = fixture_dir(); std::fs::create_dir_all(&dir).unwrap(); @@ -862,6 +1029,7 @@ fn regenerate_fixtures() { (V1_0_0, &v1_0_0), (TRIMMED, &trimmed), (P4, &p4), + (P4_V1_1_0, &p4_v1_1_0), ]; for (name, bytes) in files { std::fs::write(dir.join(name), bytes).unwrap(); @@ -887,8 +1055,10 @@ fn regenerate_fixtures() { for (fixture, list) in [ ("V1_1_0", &main), ("V1_0_0", &main), + ("V1_0_0_NATIVE", &main), ("TRIMMED", &trimmed_expected), ("P4", &p4_expected), + ("P4_V1_1_0", &p4_expected), ] { println!(" (\n {fixture},\n &["); for (name, bytes) in list.iter() { From 85ee991c22927708406612b7b9f217632057ff2a Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 21:20:02 +0000 Subject: [PATCH 04/34] fix(core): harden the legacy container reader and writer Stage 2 of the issue #5 work (synthesis E1-E8, E11). Legacy writer output is unchanged; only new error paths are added. - metadata: the tstr bounds check uses checked arithmetic, and the shape array reservation is capped by the bytes left (E1, E2). - inspect: list_extension_table_entries checks offset + length for overflow instead of panicking (E3, bounds part only). - tensor_record / plane_record writers: lengths and counts that do not fit their wire integer return InvalidContainer before any byte is written, instead of being truncated (E4). - ChainRegistry::write rejects chain ids above 0x7FFF and more than 0x8000 chains; compress_model's registry builder fails past 0x8000 distinct chains (E5). - ContainerWriter rejects duplicate tensor names. get_tensor_record_bytes scans the whole run of equal name hashes and compares record names, so a hash collision resolves to the right tensor (E6). - The prelude and the chain registry are parsed only within shared_prelude_length (E7). - read_chain_blob and chain_blob_info take the table size from the parse position (new ExtensionTable::from_bytes_prefix), not from a re-serialization (E8). - decode_tensor_shape and list_plane_summaries take skip_missing, in the bindings and the Python wrappers (E11). Each fix has a regression test (Rust unit tests and tests/ptwm/wire_format/test_legacy_hardening.py). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/wire.rs | 55 +++- crates/ptwm-core/src/compressor.rs | 48 +++- crates/ptwm-core/src/container.rs | 257 ++++++++++++++++-- crates/ptwm-core/src/extension/table.rs | 13 + crates/ptwm-core/src/metadata.rs | 52 +++- crates/ptwm-core/src/plane_record.rs | 55 +++- crates/ptwm-core/src/tensor_record.rs | 115 +++++++- crates/ptwm-py/src/inspect.rs | 40 ++- crates/ptwm-py/src/lib.rs | 27 +- python/ptwm/_rust/__init__.py | 10 +- .../ptwm/wire_format/test_legacy_hardening.py | 117 ++++++++ 11 files changed, 710 insertions(+), 79 deletions(-) create mode 100644 tests/ptwm/wire_format/test_legacy_hardening.py diff --git a/crates/ptwm-core/src/chain/wire.rs b/crates/ptwm-core/src/chain/wire.rs index 258a454..844f981 100644 --- a/crates/ptwm-core/src/chain/wire.rs +++ b/crates/ptwm-core/src/chain/wire.rs @@ -351,19 +351,10 @@ pub fn write_chain_blob(chain: &Chain) -> Result, PtwmCoreError> { /// /// Returns `(chain, bytes_consumed)`. pub fn read_chain_blob(blob: &[u8]) -> Result<(Chain, usize), PtwmCoreError> { - // Read the local extension table. - let inline_table = ExtensionTable::from_bytes(blob) + // Read the local extension table; the chain starts where the parse + // stopped. + let (inline_table, et_consumed) = ExtensionTable::from_bytes_prefix(blob) .map_err(|e| PtwmCoreError::InvalidContainer(format!("chain blob et: {e}")))?; - // Measure how many bytes the table occupies by re-serializing it. - let et_bytes = inline_table - .to_bytes() - .map_err(|e| PtwmCoreError::InvalidContainer(format!("chain blob et re-ser: {e}")))?; - let et_consumed = et_bytes.len(); - if blob.len() < et_consumed { - return Err(PtwmCoreError::InvalidContainer( - "chain blob: buffer shorter than extension table".into(), - )); - } let (chain, chain_consumed) = read_chain(&blob[et_consumed..], &inline_table)?; Ok((chain, et_consumed + chain_consumed)) } @@ -452,6 +443,46 @@ mod tests { assert_eq!(decoded, chain); } + /// Pad the capability CBOR of the blob's first table entry with one + /// trailing byte. The CBOR parser ignores it, so the table re-serializes + /// one byte shorter than it was read. + fn pad_first_capability_blob(blob: &[u8]) -> Vec { + let label_len = u16::from_le_bytes([blob[36], blob[37]]) as usize; + let cap_len_pos = 4 + 32 + 2 + label_len + 2 + 2 + 1 + 1; + let cap_len = + u32::from_le_bytes(blob[cap_len_pos..cap_len_pos + 4].try_into().unwrap()) as usize; + let cap_end = cap_len_pos + 4 + cap_len; + let mut out = blob[..cap_len_pos].to_vec(); + out.extend_from_slice(&(cap_len as u32 + 1).to_le_bytes()); + out.extend_from_slice(&blob[cap_len_pos + 4..cap_end]); + out.push(0x00); + out.extend_from_slice(&blob[cap_end..]); + out + } + + #[test] + fn read_chain_blob_counts_consumed_bytes_from_the_parse() { + let chain = Chain { + nodes: vec![ChainNode { + op: OpId::Source, + params: make_source_params(), + }], + edges: vec![], + terminals: vec![], + }; + let blob = write_chain_blob(&chain).unwrap(); + let padded = pad_first_capability_blob(&blob); + assert_eq!(padded.len(), blob.len() + 1); + // The padded table re-serializes to the original length, so a + // re-serialization count would be one byte short. + let (table, used) = ExtensionTable::from_bytes_prefix(&padded).unwrap(); + assert_eq!(table.to_bytes().unwrap().len() + 1, used); + + let (decoded, consumed) = read_chain_blob(&padded).unwrap(); + assert_eq!(decoded, chain); + assert_eq!(consumed, padded.len()); + } + #[test] fn chain_with_node_params() { // ByteSplit{n=4} carries n in params. diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 7d602b1..b17432f 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -616,7 +616,19 @@ fn pass1_select_chain( /// Build the deduped chain registry. Returns the registry and a per-tensor /// vector of registry indices (each in 0..registry.chains.len()). -fn build_chain_registry(chains: &[Chain]) -> (ChainRegistry, Vec) { +/// +/// Fails when the tensors need more distinct chains than a registry can +/// address ([`crate::container::MAX_REGISTRY_CHAINS`]). +fn build_chain_registry(chains: &[Chain]) -> Result<(ChainRegistry, Vec), PtwmCoreError> { + build_chain_registry_with_limit(chains, crate::container::MAX_REGISTRY_CHAINS) +} + +/// [`build_chain_registry`] with an explicit chain limit, so tests can reach +/// the limit without building 32768 chains. +fn build_chain_registry_with_limit( + chains: &[Chain], + max_chains: usize, +) -> Result<(ChainRegistry, Vec), PtwmCoreError> { let mut registry: Vec<(u16, Chain)> = Vec::new(); let mut tensor_to_idx: Vec = Vec::with_capacity(chains.len()); for c in chains { @@ -624,13 +636,18 @@ fn build_chain_registry(chains: &[Chain]) -> (ChainRegistry, Vec) { match existing { Some(pos) => tensor_to_idx.push(pos as u16), None => { + if registry.len() >= max_chains { + return Err(PtwmCoreError::InvalidContainer(format!( + "compress_model: more than {max_chains} distinct chains" + ))); + } let new_id = registry.len() as u16; registry.push((new_id, c.clone())); tensor_to_idx.push(new_id); } } } - (ChainRegistry { chains: registry }, tensor_to_idx) + Ok((ChainRegistry { chains: registry }, tensor_to_idx)) } // --------------------------------------------------------------------------- @@ -900,7 +917,7 @@ pub fn compress_model( .zip(inputs.iter()) .map(|(outcome, inp)| inp.candidate_chains[outcome.chain_idx].clone()) .collect(); - let (registry, tensor_chain_indices) = build_chain_registry(&chosen_chains); + let (registry, tensor_chain_indices) = build_chain_registry(&chosen_chains)?; let mut writer = ContainerWriter::new(writer)?; writer.set_prelude(&prelude_entries)?; @@ -1730,6 +1747,31 @@ mod tests { assert!(rec.inline_chain.is_none()); } } + + #[test] + fn build_chain_registry_rejects_too_many_distinct_chains() { + let chain = |b: u8| Chain { + nodes: vec![ChainNode { + op: OpId::BytePassthrough, + params: vec![b], + }], + edges: vec![], + terminals: vec![TerminalRef { + node_idx: 0, + output_idx: 0, + role: Role::Raw, + }], + }; + // Duplicates share an id and do not count against the limit. + let (reg, idx) = + build_chain_registry_with_limit(&[chain(0), chain(1), chain(0)], 2).unwrap(); + assert_eq!(reg.chains.len(), 2); + assert_eq!(idx, vec![0, 1, 0]); + let err = build_chain_registry_with_limit(&[chain(0), chain(1), chain(2)], 2) + .unwrap_err() + .to_string(); + assert!(err.contains("more than 2 distinct chains"), "{err}"); + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 0d6a320..5a53866 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -7,7 +7,7 @@ use rayon::prelude::*; use xxhash_rust::xxh64::xxh64; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use crate::chain::Chain; use crate::chain::wire::{read_chain, write_chain}; @@ -22,6 +22,14 @@ use crate::tensor_record::{TensorRecord, write_tensor_record}; // ── Chain registry ──────────────────────────────────────────────────────────── +/// Largest registry chain id. A tensor record addresses a registry chain +/// through the low 15 bits of `chain_ref`; the high bit marks an inline chain. +pub const MAX_CHAIN_ID: u16 = 0x7FFF; + +/// Largest number of chains a registry can hold: one per id in +/// `0..=MAX_CHAIN_ID`. +pub const MAX_REGISTRY_CHAINS: usize = MAX_CHAIN_ID as usize + 1; + /// In-memory chain registry section. /// /// Wire format: @@ -52,12 +60,17 @@ impl ChainRegistry { builder: &mut ExtensionTableBuilder, out: &mut Vec, ) -> Result<(), PtwmCoreError> { - if self.chains.len() > u16::MAX as usize { + if self.chains.len() > MAX_REGISTRY_CHAINS { return Err(PtwmCoreError::InvalidContainer(format!( - "ChainRegistry: {} chains exceeds u16::MAX", + "ChainRegistry: {} chains exceeds the limit of {MAX_REGISTRY_CHAINS}", self.chains.len() ))); } + if let Some((id, _)) = self.chains.iter().find(|(id, _)| *id > MAX_CHAIN_ID) { + return Err(PtwmCoreError::InvalidContainer(format!( + "ChainRegistry: chain_id {id:#06x} does not fit in the 15-bit chain_ref" + ))); + } out.extend_from_slice(&(self.chains.len() as u16).to_le_bytes()); for (chain_id, chain) in &self.chains { out.extend_from_slice(&chain_id.to_le_bytes()); @@ -76,7 +89,8 @@ impl ChainRegistry { } let num_chains = u16::from_le_bytes(buf[..2].try_into().unwrap()) as usize; let mut pos = 2usize; - let mut chains = Vec::with_capacity(num_chains); + // Each entry takes at least its 2-byte chain_id. + let mut chains = Vec::with_capacity(num_chains.min((buf.len() - 2) / 2)); for i in 0..num_chains { if pos + 2 > buf.len() { return Err(PtwmCoreError::InvalidContainer(format!( @@ -208,6 +222,8 @@ pub struct ContainerWriter { finalized: bool, /// Total byte count of the prelude + chain-registry section. prelude_section_len: u64, + /// Names appended so far. A container holds each name at most once. + names: HashSet, } impl ContainerWriter { @@ -242,6 +258,7 @@ impl ContainerWriter { prelude_written: false, finalized: false, prelude_section_len: 0, + names: HashSet::new(), }) } @@ -417,6 +434,7 @@ impl ContainerWriter { if self.finalized { return Err(PtwmCoreError::InvalidContainer("already finalized".into())); } + self.check_new_name(&record.name)?; // On first tensor: flush prelude and record tensors_start_pos. if !self.prelude_written { @@ -442,10 +460,23 @@ impl ContainerWriter { tensor_offset, tensor_len, }); + self.names.insert(record.name.clone()); Ok(()) } + /// Fail if a tensor named `name` was already appended. Readers look + /// tensors up by name, so a second record with the same name could never + /// be reached. + fn check_new_name(&self, name: &str) -> Result<(), PtwmCoreError> { + if self.names.contains(name) { + return Err(PtwmCoreError::InvalidContainer(format!( + "duplicate tensor name '{name}'" + ))); + } + Ok(()) + } + /// Append a tensor whose record bytes were serialized ahead of time (e.g. /// in parallel via [`crate::tensor_record::write_tensor_record`]). /// Identical on-disk effect to [`Self::append_tensor`] but skips the @@ -458,6 +489,7 @@ impl ContainerWriter { if self.finalized { return Err(PtwmCoreError::InvalidContainer("already finalized".into())); } + self.check_new_name(name)?; if !self.prelude_written { self.flush_prelude()?; self.tensors_start_pos = self.header.shared_prelude_offset + self.prelude_section_len; @@ -474,6 +506,7 @@ impl ContainerWriter { tensor_offset, tensor_len: tensor_buf.len() as u64, }); + self.names.insert(name.to_owned()); Ok(()) } @@ -838,13 +871,24 @@ impl<'a> ContainerReader<'a> { )); } - let (prelude, prelude_bytes_consumed) = parse_prelude(&buf[prelude_off..])?; + // The prelude and the chain registry both live inside the + // shared-prelude section; parse them from that slice only. + let prelude_end = usize::try_from(header.shared_prelude_length) + .ok() + .and_then(|len| prelude_off.checked_add(len)) + .filter(|&end| end <= buf.len()) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer("shared prelude section exceeds buffer".into()) + })?; + let prelude_section = &buf[prelude_off..prelude_end]; + let (prelude, prelude_bytes_consumed) = parse_prelude(prelude_section)?; // The chain registry immediately follows the prelude bytes within the - // shared-prelude region. + // shared-prelude section. let cr_off = prelude_off + prelude_bytes_consumed; - let chain_registry = if cr_off < index_off && cr_off < buf.len() { - let (registry, _) = ChainRegistry::read(&buf[cr_off..], &extension_table)?; + let chain_registry = if prelude_bytes_consumed < prelude_section.len() { + let (registry, _) = + ChainRegistry::read(&prelude_section[prelude_bytes_consumed..], &extension_table)?; registry } else { ChainRegistry::new() @@ -891,22 +935,31 @@ impl<'a> ContainerReader<'a> { /// Look up a tensor by name and return a slice covering its serialized /// `TensorRecord` bytes. Returns `None` if no tensor with that name - /// exists. The on-disk index is sorted by `name_hash`, so we use - /// `binary_search_by_key` for an O(log n) lookup instead of a linear - /// scan — material on models with thousands of tensors. + /// exists. + /// + /// The on-disk index is sorted by `name_hash`, so a binary search finds + /// the run of entries with the target hash in O(log n). Two names can + /// share a hash, so every entry in the run has its record name compared + /// against `name`. Entries whose range or name cannot be read do not + /// match. pub fn get_tensor_record_bytes(&self, name: &str) -> Option<&[u8]> { - let target = name_hash(name); - let idx = self - .index - .binary_search_by_key(&target, |e| e.name_hash) - .ok()?; - let e = &self.index[idx]; - let start = e.tensor_offset as usize; - let end = start.checked_add(e.tensor_len as usize)?; - if end > self.buf.len() { - return None; - } - Some(&self.buf[start..end]) + self.record_bytes_with_hash(name, name_hash(name)) + } + + /// [`get_tensor_record_bytes`](Self::get_tensor_record_bytes) with the + /// index hash given explicitly, so tests can force a hash collision. + fn record_bytes_with_hash(&self, name: &str, target: u64) -> Option<&[u8]> { + use crate::tensor_record::peek_tensor_record_name; + + let lo = self.index.partition_point(|e| e.name_hash < target); + let hi = lo + self.index[lo..].partition_point(|e| e.name_hash == target); + self.index[lo..hi].iter().find_map(|e| { + let start = usize::try_from(e.tensor_offset).ok()?; + let len = usize::try_from(e.tensor_len).ok()?; + let end = start.checked_add(len)?; + let bytes = self.buf.get(start..end)?; + (peek_tensor_record_name(bytes).ok()? == name).then_some(bytes) + }) } /// Retrieve the shape + dtype_name embedded in this tensor's metadata, @@ -1305,6 +1358,7 @@ mod tests { use crate::plane_record::PlaneRecord; use crate::tensor_record::{CHAIN_REF_INLINE_BIT, parse_tensor_record}; use crate::types::PlaneRole; + use std::io::Cursor; fn sample_tensor_record(name: &str) -> TensorRecord { TensorRecord { @@ -1606,6 +1660,163 @@ mod tests { } } + #[test] + fn chain_registry_write_rejects_ids_beyond_15_bits() { + let mut builder = ExtensionTableBuilder::new(); + let mut out = Vec::new(); + let reg = ChainRegistry { + chains: vec![(MAX_CHAIN_ID + 1, make_minimal_chain(0))], + }; + let err = reg.write(&mut builder, &mut out).unwrap_err().to_string(); + assert!(err.contains("15-bit"), "{err}"); + + let reg = ChainRegistry { + chains: vec![(MAX_CHAIN_ID, make_minimal_chain(0))], + }; + reg.write(&mut builder, &mut out).unwrap(); + } + + #[test] + fn chain_registry_write_rejects_more_than_0x8000_chains() { + let chain = make_minimal_chain(0); + let reg = ChainRegistry { + chains: (0..=MAX_REGISTRY_CHAINS) + .map(|i| ((i & 0x7FFF) as u16, chain.clone())) + .collect(), + }; + let mut builder = ExtensionTableBuilder::new(); + let mut out = Vec::new(); + let err = reg.write(&mut builder, &mut out).unwrap_err().to_string(); + assert!(err.contains("exceeds the limit"), "{err}"); + assert!(out.is_empty()); + } + + #[test] + fn writer_rejects_duplicate_tensor_names() { + let mut buf = Cursor::new(Vec::new()); + let mut w = ContainerWriter::new(&mut buf).unwrap(); + w.append_tensor(&sample_tensor_record("a")).unwrap(); + let err = w + .append_tensor(&sample_tensor_record("a")) + .unwrap_err() + .to_string(); + assert!(err.contains("duplicate tensor name"), "{err}"); + + let mut bytes = Vec::new(); + write_tensor_record(&sample_tensor_record("b"), &mut bytes).unwrap(); + w.append_serialized_tensor("b", &bytes).unwrap(); + let err = w + .append_serialized_tensor("b", &bytes) + .unwrap_err() + .to_string(); + assert!(err.contains("duplicate tensor name"), "{err}"); + let err = w + .append_serialized_tensor("a", &bytes) + .unwrap_err() + .to_string(); + assert!(err.contains("duplicate tensor name"), "{err}"); + w.finalize().unwrap(); + } + + /// Two names with the same index hash: lookups compare the record name + /// and pick the right entry, whatever their order in the hash run. + #[test] + fn lookup_resolves_name_hash_collision() { + let names = ["first", "second", "third"]; + let mut buf = Cursor::new(Vec::new()); + { + let mut w = ContainerWriter::new(&mut buf).unwrap(); + for name in names { + w.append_tensor(&sample_tensor_record(name)).unwrap(); + } + w.finalize().unwrap(); + } + let bytes = buf.into_inner(); + let mut reader = ContainerReader::open(&bytes).unwrap(); + let expected: Vec> = names + .iter() + .map(|n| reader.get_tensor_record_bytes(n).unwrap().to_vec()) + .collect(); + + // Give "first" and "second" one hash, then try both run orders. + let forced = name_hash("first"); + for e in reader.index.iter_mut() { + if e.name_hash == name_hash("second") { + e.name_hash = forced; + } + } + reader.index.sort_by_key(|e| e.name_hash); + let run: Vec = (0..reader.index.len()) + .filter(|&i| reader.index[i].name_hash == forced) + .collect(); + assert_eq!(run.len(), 2); + for _ in 0..2 { + for (name, want) in names[..2].iter().zip(&expected) { + let got = reader.record_bytes_with_hash(name, forced).unwrap(); + assert_eq!(got, want.as_slice(), "lookup of {name}"); + } + assert_eq!( + reader.get_tensor_record_bytes("third").unwrap(), + expected[2].as_slice() + ); + // A name in the run's hash with no record of its own. + assert!(reader.record_bytes_with_hash("absent", forced).is_none()); + reader.index.swap(run[0], run[1]); + } + // The real hash of "second" no longer has an entry. + assert!(reader.get_tensor_record_bytes("second").is_none()); + assert!(reader.decode_tensor("second").is_err()); + } + + /// The prelude and the chain registry are parsed from the shared-prelude + /// section only. A shrunk shared_prelude_length cuts the registry off + /// instead of letting the parser read into the tensor records. + #[test] + fn prelude_and_registry_are_bounded_by_shared_prelude_length() { + let mut buf = Cursor::new(Vec::new()); + { + let mut w = ContainerWriter::new(&mut buf).unwrap(); + w.set_chain_registry(ChainRegistry { + chains: vec![(0, make_minimal_chain(7))], + }) + .unwrap(); + w.append_tensor(&sample_tensor_record("t")).unwrap(); + w.finalize().unwrap(); + } + let bytes = buf.into_inner(); + let reader = ContainerReader::open(&bytes).unwrap(); + assert_eq!(reader.chain_registry.chains.len(), 1); + let header = reader.header.clone(); + + let patch = |len: u64| { + let mut h = header.clone(); + h.shared_prelude_length = len; + let mut b = bytes.clone(); + b[..HEADER_LEN].copy_from_slice(&h.to_bytes().unwrap()); + b + }; + + // Section ends one byte early: the registry is truncated. + let short = patch(header.shared_prelude_length - 1); + let err = ContainerReader::open(&short).err().expect("must fail"); + assert!(err.to_string().contains("ChainRegistry"), "{err}"); + + // Section holds only the prelude: no registry is read. + let prelude_only = patch(4); + let reader = ContainerReader::open(&prelude_only).unwrap(); + assert!(reader.chain_registry.chains.is_empty()); + + // Section shorter than the prelude count. + let err = ContainerReader::open(&patch(2)).err().expect("must fail"); + assert!(err.to_string().contains("prelude"), "{err}"); + + // Section runs past the end of the buffer. + let err = ContainerReader::open(&patch(u64::MAX)) + .err() + .expect("must fail"); + assert!(err.to_string().contains("shared prelude"), "{err}"); + } + #[test] fn chain_registry_roundtrip() { use crate::extension::{ExtensionTable, ExtensionTableBuilder}; diff --git a/crates/ptwm-core/src/extension/table.rs b/crates/ptwm-core/src/extension/table.rs index ad99cf9..40fc7e3 100644 --- a/crates/ptwm-core/src/extension/table.rs +++ b/crates/ptwm-core/src/extension/table.rs @@ -114,6 +114,19 @@ impl ExtensionTable { Self::read(&mut c) } + /// Parse a table from the start of `bytes` and return it with the number + /// of bytes it occupies. The count comes from the parse position, not + /// from re-serializing: a table whose fields have more than one valid + /// encoding (for example, non-canonical capability CBOR) re-serializes + /// to a different length. + pub fn from_bytes_prefix(bytes: &[u8]) -> Result<(Self, usize), ExtensionError> { + let mut c = Cursor::new(bytes); + let table = Self::read(&mut c)?; + // The cursor reads from a slice, so its position never exceeds + // `bytes.len()`. + Ok((table, c.position() as usize)) + } + /// blake3 hash committed in the header to detect tampering / corruption. pub fn hash(bytes: &[u8]) -> [u8; 32] { let mut out = [0u8; 32]; diff --git a/crates/ptwm-core/src/metadata.rs b/crates/ptwm-core/src/metadata.rs index f09f7f9..5b031d5 100644 --- a/crates/ptwm-core/src/metadata.rs +++ b/crates/ptwm-core/src/metadata.rs @@ -78,6 +78,10 @@ impl<'a> Reader<'a> { Self { buf, pos: 0 } } + fn remaining(&self) -> usize { + self.buf.len() - self.pos + } + fn read_head(&mut self, expected_major: u8) -> Result { if self.pos >= self.buf.len() { return Err(PtwmCoreError::InvalidContainer( @@ -121,16 +125,18 @@ impl<'a> Reader<'a> { } fn read_tstr(&mut self) -> Result { - let len = self.read_head(MAJOR_TSTR)? as usize; - if self.pos + len > self.buf.len() { - return Err(PtwmCoreError::InvalidContainer( - "cbor metadata tstr truncated".into(), - )); - } - let s = std::str::from_utf8(&self.buf[self.pos..self.pos + len]) + let len = self.read_head(MAJOR_TSTR)?; + let end = usize::try_from(len) + .ok() + .and_then(|len| self.pos.checked_add(len)) + .filter(|&end| end <= self.buf.len()) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer("cbor metadata tstr truncated".into()) + })?; + let s = std::str::from_utf8(&self.buf[self.pos..end]) .map_err(|_| PtwmCoreError::InvalidContainer("cbor metadata: tstr not utf8".into()))? .to_string(); - self.pos += len; + self.pos = end; Ok(s) } } @@ -154,8 +160,10 @@ pub fn decode_shape_metadata(buf: &[u8]) -> Result<(Vec, String), PtwmCoreE let key = r.read_tstr()?; match key.as_str() { "shape" => { - let n = r.read_head(MAJOR_ARRAY)? as usize; - let mut dims = Vec::with_capacity(n); + let n = r.read_head(MAJOR_ARRAY)?; + // Each element takes at least one byte, so cap the reservation + // by the bytes left instead of trusting the declared count. + let mut dims = Vec::with_capacity(n.min(r.remaining() as u64) as usize); for _ in 0..n { dims.push(r.read_head(MAJOR_UINT)?); } @@ -253,4 +261,28 @@ mod tests { ]; assert_eq!(out, expected); } + + #[test] + fn rejects_tstr_length_that_overflows() { + // map(2), then a text-string head declaring u64::MAX bytes. The + // bounds check must not wrap around. + let mut buf = Vec::new(); + write_uint(&mut buf, MAJOR_MAP, 2); + write_uint(&mut buf, MAJOR_TSTR, u64::MAX); + buf.extend_from_slice(b"shape"); + assert!(decode_shape_metadata(&buf).is_err()); + } + + #[test] + fn huge_shape_count_fails_without_reserving() { + // {"shape": array(2^60) ...}: the declared count far exceeds the + // buffer, so decoding fails on truncation and does not try to + // reserve 2^60 elements first. + let mut buf = Vec::new(); + write_uint(&mut buf, MAJOR_MAP, 2); + write_tstr(&mut buf, "shape"); + write_uint(&mut buf, MAJOR_ARRAY, 1 << 60); + write_uint(&mut buf, MAJOR_UINT, 1); + assert!(decode_shape_metadata(&buf).is_err()); + } } diff --git a/crates/ptwm-core/src/plane_record.rs b/crates/ptwm-core/src/plane_record.rs index 0addb89..db4e351 100644 --- a/crates/ptwm-core/src/plane_record.rs +++ b/crates/ptwm-core/src/plane_record.rs @@ -3,6 +3,7 @@ use crate::codec::{CodecId, StateSource}; use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; +use crate::tensor_record::checked_len; use crate::types::PlaneRole; pub const PLANE_FLAG_CHUNK_TABLE: u8 = 0x01; @@ -49,6 +50,20 @@ pub struct PlaneRecord { } pub fn write_plane_record(rec: &PlaneRecord, out: &mut Vec) -> Result<(), PtwmCoreError> { + // Check every length field before writing, so an oversized field fails + // cleanly and leaves `out` untouched. + let inline_state_len: u32 = checked_len(rec.inline_state_bytes.len(), "inline state length")?; + let chunk_count: Option = rec + .chunk_table + .as_ref() + .map(|c| checked_len(c.len(), "chunk count")) + .transpose()?; + let ext_ref_len: Option = rec + .external_state + .as_ref() + .map(|e| checked_len(e.ref_bytes.len(), "external state ref length")) + .transpose()?; + let mut plane_flags: u8 = 0; if rec.chunk_table.is_some() { plane_flags |= PLANE_FLAG_CHUNK_TABLE; @@ -65,20 +80,19 @@ pub fn write_plane_record(rec: &PlaneRecord, out: &mut Vec) -> Result<(), Pt out.push(rec.state_source as u8); out.push(rec.state_version); out.extend_from_slice(&rec.state_info.to_le_bytes()); - let inline_state_len = rec.inline_state_bytes.len() as u32; out.extend_from_slice(&inline_state_len.to_le_bytes()); out.extend_from_slice(&rec.payload_len.to_le_bytes()); if let Some(crc) = rec.crc32 { out.extend_from_slice(&crc.to_le_bytes()); } - if let Some(chunks) = &rec.chunk_table { - out.extend_from_slice(&(chunks.len() as u32).to_le_bytes()); + if let (Some(chunks), Some(chunk_count)) = (&rec.chunk_table, chunk_count) { + out.extend_from_slice(&chunk_count.to_le_bytes()); for c in chunks { out.extend_from_slice(&c.offset_in_payload.to_le_bytes()); out.extend_from_slice(&c.decoded_size.to_le_bytes()); } } - if let Some(ext) = &rec.external_state { + if let (Some(ext), Some(ext_ref_len)) = (&rec.external_state, ext_ref_len) { let expected_hash_present = if ext.expected_hash.is_some() { 0x01u8 } else { @@ -86,7 +100,7 @@ pub fn write_plane_record(rec: &PlaneRecord, out: &mut Vec) -> Result<(), Pt }; out.push(ext.ref_kind); out.push(expected_hash_present); - out.extend_from_slice(&(ext.ref_bytes.len() as u16).to_le_bytes()); + out.extend_from_slice(&ext_ref_len.to_le_bytes()); out.extend_from_slice(&ext.ref_bytes); out.push(ext.state_format_version); if let Some(hash) = ext.expected_hash { @@ -431,4 +445,35 @@ mod tests { assert_eq!(parsed.layout, PlaneLayout::Flat); assert_eq!(consumed, buf.len()); } + + #[test] + fn write_rejects_external_ref_longer_than_u16() { + let rec = PlaneRecord { + role: PlaneRole::Value, + codec_id: CodecId::Identity, + codec_table_idx: 0, + state_source: StateSource::External, + state_version: 0, + state_info: 0, + payload_len: 0, + crc32: None, + chunk_table: None, + external_state: Some(ExternalStateRef { + ref_kind: 0x02, + expected_hash: None, + ref_bytes: vec![b'x'; u16::MAX as usize + 1], + state_format_version: 0, + }), + inline_state_bytes: Vec::new(), + payload_bytes: Vec::new(), + layout: PlaneLayout::Flat, + }; + let mut out = Vec::new(); + let err = write_plane_record(&rec, &mut out).unwrap_err(); + assert!( + err.to_string().contains("external state ref length"), + "{err}" + ); + assert!(out.is_empty()); + } } diff --git a/crates/ptwm-core/src/tensor_record.rs b/crates/ptwm-core/src/tensor_record.rs index d28e697..f77665f 100644 --- a/crates/ptwm-core/src/tensor_record.rs +++ b/crates/ptwm-core/src/tensor_record.rs @@ -79,6 +79,17 @@ fn name_padding(name_len: usize) -> usize { if rem == 0 { 0 } else { 4 - rem } } +/// Convert a length or count to its wire integer type, or fail with +/// `InvalidContainer` when it does not fit. +pub(crate) fn checked_len>(n: usize, what: &str) -> Result { + T::try_from(n).map_err(|_| { + PtwmCoreError::InvalidContainer(format!( + "{what} {n} exceeds the wire limit of {} bytes", + std::mem::size_of::() + )) + }) +} + pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), PtwmCoreError> { // Validate inline_chain vs chain_ref consistency. let has_inline = rec.chain_ref & CHAIN_REF_INLINE_BIT != 0; @@ -105,9 +116,20 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), flags |= TENSOR_FLAG_PAYLOAD_HASH; } - let num_terminals = rec.terminals.len() as u8; + // Check every length field before writing, so an oversized field fails + // cleanly and leaves `out` untouched. + let num_terminals: u8 = checked_len(rec.terminals.len(), "terminal count")?; let name_bytes = rec.name.as_bytes(); - let name_len = name_bytes.len() as u16; + let name_len: u16 = checked_len(name_bytes.len(), "tensor name length")?; + let dep_count: u16 = checked_len(rec.dependencies.len(), "dependency count")?; + for dep in &rec.dependencies { + checked_len::(dep.ref_bytes.len(), "dependency ref length")?; + } + let meta_len: Option = rec + .tensor_metadata + .as_ref() + .map(|m| checked_len(m.len(), "tensor metadata length")) + .transpose()?; // Fixed header: chain_ref(u16) + dtype_code(u16) + input_format(u8) + // num_terminals(u8) + flags(u8) + reserved(u8) + @@ -159,7 +181,6 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), // dependencies (present iff TENSOR_FLAG_DEPENDENCIES) if !rec.dependencies.is_empty() { - let dep_count = rec.dependencies.len() as u16; out.extend_from_slice(&dep_count.to_le_bytes()); for dep in &rec.dependencies { let ref_flags: u8 = if dep.expected_hash.is_some() { @@ -169,6 +190,7 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), }; out.push(dep.ref_kind as u8); out.push(ref_flags); + // Length checked above. out.extend_from_slice(&(dep.ref_bytes.len() as u16).to_le_bytes()); out.extend_from_slice(&dep.ref_bytes); if let Some(hash) = dep.expected_hash { @@ -178,8 +200,8 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), } // tensor_metadata (present iff TENSOR_FLAG_METADATA) - if let Some(meta) = &rec.tensor_metadata { - out.extend_from_slice(&(meta.len() as u32).to_le_bytes()); + if let (Some(meta), Some(meta_len)) = (&rec.tensor_metadata, meta_len) { + out.extend_from_slice(&meta_len.to_le_bytes()); out.extend_from_slice(meta); } @@ -191,6 +213,24 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), Ok(()) } +/// Read only the tensor name from serialized record bytes, without parsing +/// or copying the rest of the record. +pub fn peek_tensor_record_name(buf: &[u8]) -> Result<&str, PtwmCoreError> { + const NAME_LEN_POS: usize = 16; + const NAME_POS: usize = 18; + if buf.len() < NAME_POS { + return Err(PtwmCoreError::InvalidContainer( + "tensor record header truncated".into(), + )); + } + let name_len = u16::from_le_bytes(buf[NAME_LEN_POS..NAME_POS].try_into().unwrap()) as usize; + let name = buf.get(NAME_POS..NAME_POS + name_len).ok_or_else(|| { + PtwmCoreError::InvalidContainer("tensor record truncated reading name".into()) + })?; + std::str::from_utf8(name) + .map_err(|_| PtwmCoreError::InvalidContainer("tensor name is not valid UTF-8".into())) +} + pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCoreError> { if buf.len() < 18 { return Err(PtwmCoreError::InvalidContainer( @@ -586,4 +626,69 @@ mod tests { let mut out = Vec::new(); assert!(write_tensor_record(&rec, &mut out).is_err()); } + + fn base_record() -> TensorRecord { + TensorRecord { + chain_ref: 0x0001, + dtype_code: 0x0001, + input_format: 1, + flags: 0, + orig_size: 0, + name: "w".to_string(), + payload_hash: None, + inline_chain: None, + dependencies: Vec::new(), + tensor_metadata: None, + terminals: vec![identity_plane()], + } + } + + fn assert_write_fails_cleanly(rec: &TensorRecord, needle: &str) { + let mut out = Vec::new(); + let err = write_tensor_record(rec, &mut out).unwrap_err(); + assert!(err.to_string().contains(needle), "{err}"); + assert!(out.is_empty(), "writer emitted bytes before failing"); + } + + #[test] + fn write_rejects_more_than_255_terminals() { + let mut rec = base_record(); + rec.terminals = vec![identity_plane(); 256]; + assert_write_fails_cleanly(&rec, "terminal count"); + rec.terminals.pop(); + let mut out = Vec::new(); + write_tensor_record(&rec, &mut out).unwrap(); + let (parsed, _) = parse_tensor_record(&out).unwrap(); + assert_eq!(parsed.terminals.len(), 255); + } + + #[test] + fn write_rejects_name_longer_than_u16() { + let mut rec = base_record(); + rec.name = "n".repeat(u16::MAX as usize + 1); + assert_write_fails_cleanly(&rec, "tensor name length"); + } + + #[test] + fn write_rejects_too_many_dependencies() { + let mut rec = base_record(); + let dep = Dependency { + ref_kind: RefKind::LocalTensor, + expected_hash: None, + ref_bytes: b"x".to_vec(), + }; + rec.dependencies = vec![dep; u16::MAX as usize + 1]; + assert_write_fails_cleanly(&rec, "dependency count"); + } + + #[test] + fn write_rejects_dependency_ref_longer_than_u16() { + let mut rec = base_record(); + rec.dependencies = vec![Dependency { + ref_kind: RefKind::LocalTensor, + expected_hash: None, + ref_bytes: vec![b'x'; u16::MAX as usize + 1], + }]; + assert_write_fails_cleanly(&rec, "dependency ref length"); + } } diff --git a/crates/ptwm-py/src/inspect.rs b/crates/ptwm-py/src/inspect.rs index a064bb2..cc67933 100644 --- a/crates/ptwm-py/src/inspect.rs +++ b/crates/ptwm-py/src/inspect.rs @@ -12,6 +12,14 @@ fn to_pyerr(e: E) -> PyErr { pyo3::exceptions::PyValueError::new_err(e.to_string()) } +/// Return `buf[off..off + len]`, or `None` when the range overflows or runs +/// past the end of `buf`. +fn extension_table_region(buf: &[u8], off: u64, len: u64) -> Option<&[u8]> { + let start = usize::try_from(off).ok()?; + let end = start.checked_add(usize::try_from(len).ok()?)?; + buf.get(start..end) +} + /// List the extension-table entries recorded in a `.ptwm` file. /// /// Returns a list of ``(canonical_id_hex, human_label, kind_u16, flavor_hints)`` @@ -39,18 +47,17 @@ pub fn list_extension_table_entries(path: String) -> PyResult buf.len() { - return Err(to_pyerr(format!( - "'{}': extension table region [{}..{}] exceeds file size {}", + let off = header.extension_table_offset; + let len = header.extension_table_length; + let table_bytes = extension_table_region(&buf, off, len).ok_or_else(|| { + to_pyerr(format!( + "'{}': extension table region at offset {} with length {} exceeds file size {}", path, off, - off + len, + len, buf.len() - ))); - } - let table_bytes = &buf[off..off + len]; + )) + })?; let table = ExtensionTable::from_bytes(table_bytes).map_err(to_pyerr)?; Ok(table .entries @@ -70,3 +77,18 @@ pub fn register(m: &Bound<'_, pyo3::types::PyModule>) -> PyResult<()> { m.add_function(pyo3::wrap_pyfunction!(list_extension_table_entries, m)?)?; Ok(()) } + +#[cfg(test)] +mod tests { + use super::extension_table_region; + + #[test] + fn region_rejects_overflow_and_out_of_bounds() { + let buf = [0u8; 16]; + assert_eq!(extension_table_region(&buf, 4, 8), Some(&buf[4..12])); + assert_eq!(extension_table_region(&buf, 16, 0), Some(&buf[16..16])); + assert!(extension_table_region(&buf, 8, 9).is_none()); + assert!(extension_table_region(&buf, u64::MAX - 4, 8).is_none()); + assert!(extension_table_region(&buf, 4, u64::MAX).is_none()); + } +} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index a6f3d4c..d8c1580 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -743,17 +743,24 @@ fn list_tensor_names( /// Decode the shape-metadata map for a v3 tensor, if present. /// /// Returns ``(shape, dtype_name)`` or ``None`` when no CBOR shape metadata. +/// `skip_missing` has the same semantics as for :func:`decode_tensor`. #[pyfunction] +#[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_shape( py: Python<'_>, blob: &Bound<'_, pyo3::types::PyBytes>, name: &str, + skip_missing: bool, ) -> PyResult, String)>> { let data: Vec = blob.as_bytes().to_vec(); let name_owned = name.to_string(); let got = py .allow_threads(|| -> Result, String)>, PtwmCoreError> { - let reader = ContainerReader::open(&data)?; + let reader = if skip_missing { + ContainerReader::open_partial(&data)? + } else { + ContainerReader::open(&data)? + }; reader.tensor_shape(&name_owned) }) .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; @@ -764,10 +771,13 @@ fn decode_tensor_shape( /// /// Returns a list of (tensor_name, plane_index, role, codec_id, /// state_source, payload_len) tuples — one row per plane. +/// `skip_missing` has the same semantics as for :func:`decode_tensor`. #[pyfunction] +#[pyo3(signature = (blob, skip_missing = false))] fn list_plane_summaries( py: Python<'_>, blob: &Bound<'_, pyo3::types::PyBytes>, + skip_missing: bool, ) -> PyResult> { use ptwm_core::tensor_record::parse_tensor_record; @@ -775,7 +785,11 @@ fn list_plane_summaries( let summaries = py .allow_threads( || -> Result, PtwmCoreError> { - let reader = ContainerReader::open(&data)?; + let reader = if skip_missing { + ContainerReader::open_partial(&data)? + } else { + ContainerReader::open(&data)? + }; let names = reader.tensor_names()?; let mut out = Vec::new(); for name in names { @@ -920,13 +934,10 @@ fn py_chain_blob_info(blob: &[u8]) -> PyResult<(usize, Vec<(u16, u16)>)> { use ptwm_core::extension::dispatch::{BuiltinKind, dispatch_builtin}; use ptwm_core::extension::table::ExtensionTable; - // Parse the extension table. - let inline_table = ExtensionTable::from_bytes(blob) + // Parse the extension table; its size is the parse position, not the + // length of a re-serialization. + let (inline_table, et_size) = ExtensionTable::from_bytes_prefix(blob) .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("chain blob et: {e}")))?; - let et_bytes = inline_table - .to_bytes() - .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("chain blob et ser: {e}")))?; - let et_size = et_bytes.len(); // Build (table_idx, op_id) pairs for every entry that resolves to an op. let mut table: Vec<(u16, u16)> = Vec::with_capacity(inline_table.entries.len()); diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index 3c7428d..8045e21 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -300,9 +300,11 @@ def decode_model(blob: bytes, skip_missing: bool = False) -> list[tuple[str, byt @translate_errors -def decode_tensor_shape(blob: bytes, name: str) -> tuple[list[int], str] | None: +def decode_tensor_shape( + blob: bytes, name: str, skip_missing: bool = False +) -> tuple[list[int], str] | None: """Return ``(shape, dtype_name)`` for a tensor, or ``None``.""" - return _ext.decode_tensor_shape(blob, name) + return _ext.decode_tensor_shape(blob, name, skip_missing) @translate_errors @@ -325,14 +327,14 @@ def list_tensor_names(blob: bytes, skip_missing: bool = False) -> list[str]: @translate_errors def list_plane_summaries( - blob: bytes, + blob: bytes, skip_missing: bool = False ) -> list[tuple[str, int, int, int, int, int]]: """Return per-plane summaries for every tensor in a container. Each entry: ``(tensor_name, plane_index, role, codec_id, state_source, payload_len)``. """ - return list(_ext.list_plane_summaries(blob)) + return list(_ext.list_plane_summaries(blob, skip_missing)) @translate_errors diff --git a/tests/ptwm/wire_format/test_legacy_hardening.py b/tests/ptwm/wire_format/test_legacy_hardening.py new file mode 100644 index 0000000..def7553 --- /dev/null +++ b/tests/ptwm/wire_format/test_legacy_hardening.py @@ -0,0 +1,117 @@ +"""Regression tests for reader hardening on legacy ``\\x89PTWM`` containers.""" + +from __future__ import annotations + +import struct + +import pytest +import torch +from ptwm import CompressionConfig, Compressor, Format, Framing, Method, _rust + +HEADER_LEN = 91 +# Header field offsets (see crates/ptwm-core/src/header.rs). +ET_OFF, ET_LEN, ET_HASH = 11, 19, 27 +INDEX_OFF, PRELUDE_OFF = 59, 75 + + +def _legacy_blob() -> bytes: + tensor = torch.arange(64 * 32, dtype=torch.float32).reshape(64, 32) + return Compressor( + CompressionConfig( + input_format=Format.TORCH, + method=Method.HUFFMAN, + framing=Framing.CONTAINER, + ) + ).compress(tensor) + + +def _u64(blob: bytes | bytearray, off: int) -> int: + return struct.unpack_from(" bytes: + """One Extension Table entry for a third-party extension (wasm flavor) + that is not installed on this host.""" + label = b"io.example.missing" + return b"".join( + [ + bytes(range(1, 33)), # canonical_id + struct.pack(" bytes: + """Append a missing third-party entry to the Extension Table and fix up + every offset and the table hash. Existing table indices stay valid.""" + blake3 = pytest.importorskip("blake3") + et_off, et_len = _u64(blob, ET_OFF), _u64(blob, ET_LEN) + assert et_off == HEADER_LEN + table = bytearray(blob[et_off : et_off + et_len]) + count = struct.unpack_from(" None: + blob = _legacy_blob() + (name,) = _rust.list_tensor_names(blob) + shape = _rust.decode_tensor_shape(blob, name) + summaries = _rust.list_plane_summaries(blob) + assert shape is not None + assert list(shape[0]) == [64, 32] + + patched = _with_missing_extension(blob) + # Without skip_missing the reader refuses the file up front. + with pytest.raises(Exception, match="(?i)missing|flavor"): + _rust.decode_tensor_shape(patched, name) + with pytest.raises(Exception, match="(?i)missing|flavor"): + _rust.list_plane_summaries(patched) + # With skip_missing the metadata is still readable. + assert _rust.decode_tensor_shape(patched, name, skip_missing=True) == shape + assert _rust.list_plane_summaries(patched, skip_missing=True) == summaries + # The other skip_missing entry points agree. + assert _rust.list_tensor_names(patched, skip_missing=True) == [name] + assert _rust.decode_tensor(patched, name, skip_missing=True) == _rust.decode_tensor( + blob, name + ) + + +def test_inspect_rejects_extension_table_offset_overflow(tmp_path) -> None: + from ptwm._core import list_extension_table_entries + + blob = bytearray(_legacy_blob()) + path = tmp_path / "ok.ptwm" + path.write_bytes(bytes(blob)) + assert list_extension_table_entries(str(path)) + + # offset + length wraps around u64: must be a clean ValueError. + struct.pack_into(" Date: Tue, 27 Oct 2026 21:31:28 +0000 Subject: [PATCH 05/34] fix(S2-hardening): address review findings - write_tensor_record truncates out to its starting length on any error, so inline-chain and terminal-record failures no longer leave a partial record behind; regression test with a non-empty prefix buffer. - ContainerReader: drop the unreachable shared_prelude_length == 0 fallback for tensors_offset; test that a zero-length section is rejected and that tensors_offset is the section end. - Python test pins chain_blob_info returning the parsed Extension Table size (padded capability blob moves it by one byte). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/container.rs | 25 ++++++---- crates/ptwm-core/src/tensor_record.rs | 47 ++++++++++++++++++- .../ptwm/wire_format/test_legacy_hardening.py | 36 ++++++++++++++ 3 files changed, 96 insertions(+), 12 deletions(-) diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 5a53866..c3ebe71 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -901,16 +901,11 @@ impl<'a> ContainerReader<'a> { let index = parse_index(&buf[index_off..])?; let num_tensors = index.len() as u64; - // tensors_offset: region between prelude section end and index start. - // shared_prelude_offset + shared_prelude_length = end of prelude section. - let tensors_offset = if header.shared_prelude_length > 0 { - header - .shared_prelude_offset - .saturating_add(header.shared_prelude_length) - } else { - // Fallback: tensors sit right after the chain-registry end. - cr_off as u64 - }; + // Tensor records start at the end of the shared-prelude section. + // prelude_end is bounds-checked above, and parse_prelude has already + // rejected a section too short to hold the prelude (including + // shared_prelude_length == 0). + let tensors_offset = prelude_end as u64; // Build the codec router once from the installed extensions snapshot. // Using the same `installed` slice that the missing-flavor check and @@ -1810,6 +1805,16 @@ mod tests { let err = ContainerReader::open(&patch(2)).err().expect("must fail"); assert!(err.to_string().contains("prelude"), "{err}"); + // A zero-length section has no prelude at all and is rejected. + let err = ContainerReader::open(&patch(0)).err().expect("must fail"); + assert!(err.to_string().contains("prelude"), "{err}"); + + // An intact section puts the tensor records right after it. + assert_eq!( + ContainerReader::open(&bytes).unwrap().tensors_offset, + header.shared_prelude_offset + header.shared_prelude_length + ); + // Section runs past the end of the buffer. let err = ContainerReader::open(&patch(u64::MAX)) .err() diff --git a/crates/ptwm-core/src/tensor_record.rs b/crates/ptwm-core/src/tensor_record.rs index f77665f..5a201cf 100644 --- a/crates/ptwm-core/src/tensor_record.rs +++ b/crates/ptwm-core/src/tensor_record.rs @@ -90,7 +90,23 @@ pub(crate) fn checked_len>(n: usize, what: &str) -> Result) -> Result<(), PtwmCoreError> { + let start = out.len(); + let result = write_tensor_record_unchecked(rec, out); + if result.is_err() { + out.truncate(start); + } + result +} + +fn write_tensor_record_unchecked( + rec: &TensorRecord, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { // Validate inline_chain vs chain_ref consistency. let has_inline = rec.chain_ref & CHAIN_REF_INLINE_BIT != 0; if has_inline && rec.inline_chain.is_none() { @@ -116,8 +132,9 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), flags |= TENSOR_FLAG_PAYLOAD_HASH; } - // Check every length field before writing, so an oversized field fails - // cleanly and leaves `out` untouched. + // Check the record-level length fields before writing anything. Errors + // from the inline chain and the terminal records can still fire after + // the header is written; `write_tensor_record` truncates `out` for those. let num_terminals: u8 = checked_len(rec.terminals.len(), "terminal count")?; let name_bytes = rec.name.as_bytes(); let name_len: u16 = checked_len(name_bytes.len(), "tensor name length")?; @@ -648,6 +665,32 @@ mod tests { let err = write_tensor_record(rec, &mut out).unwrap_err(); assert!(err.to_string().contains(needle), "{err}"); assert!(out.is_empty(), "writer emitted bytes before failing"); + + // Bytes already in the buffer survive a failed write unchanged. + let mut out = b"prefix".to_vec(); + write_tensor_record(rec, &mut out).unwrap_err(); + assert_eq!(out, b"prefix"); + } + + /// A terminal-level error fires after the header, name and hash are + /// written; the partial record must not stay in `out`. + #[test] + fn terminal_error_after_header_leaves_out_unchanged() { + let mut rec = base_record(); + rec.payload_hash = Some(0x1234); + rec.dependencies = vec![Dependency { + ref_kind: RefKind::LocalTensor, + expected_hash: None, + ref_bytes: b"dep".to_vec(), + }]; + rec.terminals[0].state_source = StateSource::External; + rec.terminals[0].external_state = Some(crate::plane_record::ExternalStateRef { + ref_kind: 0x02, + expected_hash: None, + ref_bytes: vec![b'x'; u16::MAX as usize + 1], + state_format_version: 0, + }); + assert_write_fails_cleanly(&rec, "external state ref length"); } #[test] diff --git a/tests/ptwm/wire_format/test_legacy_hardening.py b/tests/ptwm/wire_format/test_legacy_hardening.py index def7553..b80f410 100644 --- a/tests/ptwm/wire_format/test_legacy_hardening.py +++ b/tests/ptwm/wire_format/test_legacy_hardening.py @@ -115,3 +115,39 @@ def test_inspect_rejects_extension_table_offset_overflow(tmp_path) -> None: path.write_bytes(bytes(blob)) with pytest.raises(ValueError, match="extension table"): list_extension_table_entries(str(path)) + + +def _pad_first_capability(blob: bytes) -> bytes: + """Append one byte to the first Extension Table entry's capability CBOR + and grow its length field to match.""" + label_len = struct.unpack_from(" None: + """chain_blob_info returns the bytes the Extension Table parse consumed, + not the length of a re-serialization. Padding a capability blob by one + byte must move the reported chain start by exactly one byte.""" + from ptwm._core import chain_blob_from_legacy, chain_blob_info + + nodes = struct.pack(" Date: Tue, 27 Oct 2026 21:59:33 +0000 Subject: [PATCH 06/34] feat(core): read built-in ids of other releases Stage 3 of the issue #5 work (synthesis X1, E9, E12; legacy manifest parts of section 6). Legacy writer output is unchanged. - extension::builtin_registry: a OnceLock map over every built-in, so dispatch_builtin, is_builtin and builtin_canonical_id are O(1). OpId::ALL and OpId::builtin_name replace the per-call id derivation. - normalize_builtin_ids maps entries labelled io.ptwm.builtin.@ to the current ids, but only when has a built-in entry and the id equals derive(BUILTIN_PUBKEY, name, ver). Forged entries stay non-built-in. It runs after the blake3 check in ContainerReader, on legacy inline-chain tables and in read_chain_blob / chain_blob_info, so older chain-blob caches load. ContainerReader exposes the table's builtin_version. - E9: when every Extension Table entry is a built-in, the reader skips discovery, the keyring and apply_fresh_install. An integration test with temporary HOME/XDG directories checks that no file is created. - Writers take a built-in version: build_extension_table(.., version), ExtensionTableBuilder::with_builtin_version / intern_builtin, ContainerWriter::set_builtin_version (file-level and inline tables), write_tensor_record_with_builtin_version. - New ContainerWriter::set_raw_extension_table (validated by encode_registry_for_table: canonical bytes, no trailing bytes, no embedded WASM, every chain op present), written verbatim. - Transcode: optional container_flags.builtin_version and registry.extension_table_b64 (omitted when unused, so existing manifests are unchanged). Explode records the version, carries the raw table only when no builder reproduces it, refuses inline chains of two different versions, and checks that the shared-prelude section re-encodes byte-exactly. Implode tries the trimmed table, then all built-ins, at that version, then the raw bytes, and errors when none matches the recorded hash instead of falling back silently. - E12: implode checks (ptwm_format, manifest_version) before parsing; every manifest struct denies unknown fields. - The v1.0.0 fixture tests are no longer ignored; the v1.0.0 fixtures decode and explode->implode byte-exactly. New tests cover mixed-version tables, forged labels, single-tensor implode of v1.0.0 files, manifest rejection, and old-id chain blobs from Python. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/wire.rs | 78 ++- crates/ptwm-core/src/container.rs | 343 +++++++++++-- crates/ptwm-core/src/extension/builder.rs | 112 ++++- crates/ptwm-core/src/extension/builtins.rs | 471 +++++++++++++++--- crates/ptwm-core/src/extension/dispatch.rs | 139 +----- crates/ptwm-core/src/extension/mod.rs | 10 +- crates/ptwm-core/src/tensor_record.rs | 36 +- crates/ptwm-core/src/transcode/manifest.rs | 27 + crates/ptwm-core/src/transcode/mod.rs | 215 ++++++-- crates/ptwm-core/src/transcode/tests.rs | 193 +++++++ crates/ptwm-core/src/transforms/op.rs | 39 +- crates/ptwm-core/tests/legacy_fixtures.rs | 275 +++++++++- crates/ptwm-core/tests/trust_fast_path.rs | 96 ++++ crates/ptwm-py/src/lib.rs | 4 +- .../test_cross_version_builtins.py | 145 ++++++ 15 files changed, 1878 insertions(+), 305 deletions(-) create mode 100644 crates/ptwm-core/tests/trust_fast_path.rs create mode 100644 tests/ptwm/wire_format/test_cross_version_builtins.py diff --git a/crates/ptwm-core/src/chain/wire.rs b/crates/ptwm-core/src/chain/wire.rs index 844f981..c954a79 100644 --- a/crates/ptwm-core/src/chain/wire.rs +++ b/crates/ptwm-core/src/chain/wire.rs @@ -26,7 +26,7 @@ use crate::chain::graph::{Chain, ChainEdge, ChainNode, TerminalRef}; use crate::error::PtwmCoreError; use crate::extension::dispatch::{BuiltinKind, dispatch_builtin}; use crate::extension::table::ExtensionTable; -use crate::extension::{CanonicalId, ExtensionTableBuilder, ExtensionTableEntry}; +use crate::extension::{ExtensionTableBuilder, normalize_builtin_ids}; use crate::transforms::op::OpId; use crate::types::role::Role; @@ -69,9 +69,9 @@ pub fn write_chain( ))); } // Intern the op into the extension table and write the table index. - let canonical_id = node.op.canonical_id(); - let entry = builtin_table_entry(canonical_id); - let table_idx = builder.intern(entry); + let table_idx = builder + .intern_builtin(node.op.builtin_name()) + .expect("missing curated builtin entry for op (invariant violated)"); nodes_buf.extend_from_slice(&table_idx.to_le_bytes()); nodes_buf.push(node.params.len() as u8); nodes_buf.extend_from_slice(&node.params); @@ -309,17 +309,6 @@ fn resolve_op_from_table(table_idx: u16, table: &ExtensionTable) -> Result ExtensionTableEntry { - crate::extension::builtins::builtin_entries() - .into_iter() - .find(|e| e.canonical_id == id) - .expect("missing curated builtin entry for canonical id (invariant violated)") -} - /// Serialize `chain` into a **self-contained blob** suitable for use as a /// Python chain blob or for any context where the caller does not have access /// to a shared file-level Extension Table. @@ -353,8 +342,10 @@ pub fn write_chain_blob(chain: &Chain) -> Result, PtwmCoreError> { pub fn read_chain_blob(blob: &[u8]) -> Result<(Chain, usize), PtwmCoreError> { // Read the local extension table; the chain starts where the parse // stopped. - let (inline_table, et_consumed) = ExtensionTable::from_bytes_prefix(blob) + let (mut inline_table, et_consumed) = ExtensionTable::from_bytes_prefix(blob) .map_err(|e| PtwmCoreError::InvalidContainer(format!("chain blob et: {e}")))?; + // Blobs cached by another release carry that release's built-in ids. + normalize_builtin_ids(&mut inline_table); let (chain, chain_consumed) = read_chain(&blob[et_consumed..], &inline_table)?; Ok((chain, et_consumed + chain_consumed)) } @@ -483,6 +474,61 @@ mod tests { assert_eq!(consumed, padded.len()); } + /// Chain blobs cached by another release carry that release's built-in + /// ids; `read_chain_blob` maps them to the current ones. + #[test] + fn read_chain_blob_accepts_builtin_ids_of_other_releases() { + let chain = Chain { + nodes: vec![ + ChainNode { + op: OpId::Source, + params: make_source_params(), + }, + ChainNode { + op: OpId::ByteSplit, + params: vec![2u8], + }, + ], + edges: vec![ChainEdge { + src_node: 0, + src_output_idx: 0, + dst_node: 1, + dst_input_idx: 0, + role_override: None, + vendor_bytes: vec![], + }], + terminals: vec![], + }; + let mut builder = ExtensionTableBuilder::with_builtin_version("1.0.0"); + let mut chain_bytes = Vec::new(); + write_chain(&chain, &mut builder, &mut chain_bytes).unwrap(); + let table = ExtensionTable { + entries: builder.finish(), + }; + assert!( + table + .entries + .iter() + .all(|e| e.human_label.ends_with("@1.0.0") + && dispatch_builtin(&e.canonical_id).is_none()) + ); + // Without normalisation the ops do not resolve. + assert!(read_chain(&chain_bytes, &table).is_err()); + + let mut blob = table.to_bytes().unwrap(); + blob.extend_from_slice(&chain_bytes); + let (decoded, consumed) = read_chain_blob(&blob).unwrap(); + assert_eq!(decoded, chain); + assert_eq!(consumed, blob.len()); + + // A forged entry (built-in label, wrong id) stays unresolved. + let mut forged = table.clone(); + forged.entries[1].canonical_id = crate::extension::CanonicalId::from_bytes([7; 32]); + let mut blob = forged.to_bytes().unwrap(); + blob.extend_from_slice(&chain_bytes); + assert!(read_chain_blob(&blob).is_err()); + } + #[test] fn chain_with_node_params() { // ByteSplit{n=4} carries n in params. diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index c3ebe71..a4f2de5 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -18,7 +18,7 @@ use crate::index::{ EOF_SENTINEL, IndexEntry, name_hash, parse_index, verify_eof_sentinel, write_index, }; use crate::prelude::{PreludeEntry, parse_prelude, write_prelude}; -use crate::tensor_record::{TensorRecord, write_tensor_record}; +use crate::tensor_record::{TensorRecord, write_tensor_record_with_builtin_version}; // ── Chain registry ──────────────────────────────────────────────────────────── @@ -129,6 +129,7 @@ pub enum CodecSet<'a> { } /// Output of [`build_extension_table`]. +#[derive(Debug)] pub struct BuiltExtensionTable { pub table: ExtensionTable, /// Serialized chain registry whose op indices refer to `table`. @@ -142,16 +143,22 @@ pub struct BuiltExtensionTable { /// plane codecs selected by `codecs`, in the order of /// [`crate::extension::builtin_entries`]. /// -/// The result is a pure function of `registry` and the codec set, so a -/// reader-side tool (the transcoder) can rebuild the same bytes. +/// Built-in entries carry the ids and labels of `builtin_version` +/// ([`crate::extension::CURRENT_BUILTIN_VERSION`] for new files). The +/// returned `codec_idx_map` is keyed by the current canonical ids whatever +/// the version. +/// +/// The result is a pure function of `registry`, the codec set and the +/// version, so a reader-side tool (the transcoder) can rebuild the same +/// bytes. pub fn build_extension_table( registry: &ChainRegistry, codecs: CodecSet<'_>, + builtin_version: &str, ) -> Result { - use crate::extension::Kind; - use crate::extension::builtin_entries; + use crate::extension::{Kind, builtin_entries, builtin_registry}; - let mut builder = ExtensionTableBuilder::new(); + let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version); let mut registry_bytes = Vec::new(); registry.write(&mut builder, &mut registry_bytes)?; @@ -160,28 +167,35 @@ pub fn build_extension_table( // excludes `huffman_nibble`, which `dispatch_builtin` keeps for legacy // wire compat but which has no `builtin_entry`, so `is_builtin` and // `collect_missing_flavors` do not recognise it. - let codec_entries: Vec<_> = builtin_entries() + let codec_ids: Vec = builtin_entries() .into_iter() .filter(|e| e.kind == Kind::PlaneCodec) + .map(|e| e.canonical_id) .collect(); if let CodecSet::Only(wanted) = codecs { for id in wanted { - if !codec_entries.iter().any(|e| e.canonical_id == *id) { + if !codec_ids.contains(id) { return Err(PtwmCoreError::InvalidContainer(format!( "codec {id:?} has no built-in Extension Table entry" ))); } } } + let reg = builtin_registry(); let mut codec_idx_map = HashMap::new(); - for entry in codec_entries { + for canonical_id in codec_ids { let keep = match codecs { CodecSet::AllBuiltins => true, - CodecSet::Only(wanted) => wanted.contains(&entry.canonical_id), + CodecSet::Only(wanted) => wanted.contains(&canonical_id), }; if keep { - let canonical_id = entry.canonical_id; - let idx = builder.intern(entry); + let name = reg + .get(&canonical_id) + .map(|r| r.name) + .expect("built-in codec entry is in the registry"); + let idx = builder + .intern_builtin(name) + .expect("built-in codec has an Extension Table entry"); codec_idx_map.insert(canonical_id, idx); } } @@ -195,6 +209,65 @@ pub fn build_extension_table( }) } +/// Validate serialized Extension Table bytes and encode `registry` against +/// them, for [`ContainerWriter::set_raw_extension_table`]. +/// +/// Checks: +/// - `table_bytes` parse as one Extension Table with no trailing bytes and +/// re-serialize to the same bytes; +/// - no entry references an embedded WASM module; +/// - every op of every registry chain is already in the table (built-in +/// ids of other releases count once normalised). +/// +/// Registry ops resolve to the first entry with their id; `codec_idx_map` +/// maps each built-in codec's current id to its first entry. The returned +/// `table` is the table as parsed, without normalisation. +pub fn encode_registry_for_table( + table_bytes: &[u8], + registry: &ChainRegistry, +) -> Result { + use crate::extension::{BuiltinKind, dispatch_builtin, normalize_builtin_ids}; + + let bad = |msg: String| PtwmCoreError::InvalidContainer(format!("raw extension table: {msg}")); + let (table, used) = + ExtensionTable::from_bytes_prefix(table_bytes).map_err(|e| bad(e.to_string()))?; + if used != table_bytes.len() { + return Err(bad(format!("{} trailing bytes", table_bytes.len() - used))); + } + if table.to_bytes().map_err(|e| bad(e.to_string()))? != table_bytes { + return Err(bad("bytes are not in canonical form".into())); + } + if table + .entries + .iter() + .any(|e| e.embedded_wasm_offset.is_some() || e.embedded_wasm_length.is_some()) + { + return Err(bad("embedded WASM entries are not supported".into())); + } + + let mut normalised = table.clone(); + normalize_builtin_ids(&mut normalised); + let n_entries = normalised.entries.len(); + let mut codec_idx_map = HashMap::new(); + for (i, e) in normalised.entries.iter().enumerate() { + if let Some(BuiltinKind::Codec(_)) = dispatch_builtin(&e.canonical_id) { + let idx = u16::try_from(i).map_err(|_| bad("too many entries".into()))?; + codec_idx_map.entry(e.canonical_id).or_insert(idx); + } + } + let mut builder = ExtensionTableBuilder::from_entries(normalised.entries); + let mut registry_bytes = Vec::new(); + registry.write(&mut builder, &mut registry_bytes)?; + if builder.len() != n_entries { + return Err(bad("a chain op is missing from the table".into())); + } + Ok(BuiltExtensionTable { + table, + registry_bytes, + codec_idx_map, + }) +} + /// Streaming writer for a `.ptwm` container. pub struct ContainerWriter { writer: W, @@ -224,6 +297,14 @@ pub struct ContainerWriter { prelude_section_len: u64, /// Names appended so far. A container holds each name at most once. names: HashSet, + /// Version whose ids and labels built-in Extension Table entries are + /// written with (file-level and inline-chain tables). + builtin_version: String, + /// Extension Table bytes to write verbatim instead of serializing + /// `extension_table` (see `set_raw_extension_table`). + raw_extension_table: Option>, + /// Whether an Extension Table or chain registry was set. + table_set: bool, } impl ContainerWriter { @@ -259,6 +340,9 @@ impl ContainerWriter { finalized: false, prelude_section_len: 0, names: HashSet::new(), + builtin_version: crate::extension::CURRENT_BUILTIN_VERSION.to_owned(), + raw_extension_table: None, + table_set: false, }) } @@ -327,8 +411,10 @@ impl ContainerWriter { "chain registry must be set before appending tensors".into(), )); } - let built = build_extension_table(®istry, codecs)?; + let built = build_extension_table(®istry, codecs, &self.builtin_version)?; self.extension_table = built.table; + self.raw_extension_table = None; + self.table_set = true; self.chain_registry_bytes = built.registry_bytes; self.chain_registry = registry; self.codec_idx_map = built.codec_idx_map; @@ -349,6 +435,62 @@ impl ContainerWriter { /// with a real table; the default is `ExtensionTable::empty()`. pub fn set_extension_table(&mut self, table: ExtensionTable) { self.extension_table = table; + self.raw_extension_table = None; + self.table_set = true; + } + + /// Write built-in Extension Table entries with the ids and labels of + /// `version` instead of the crate version. This applies to the table + /// built by [`set_chain_registry`](Self::set_chain_registry) and + /// [`set_chain_registry_with_codecs`](Self::set_chain_registry_with_codecs) + /// and to inline-chain tables. The transcoder uses it to reproduce files + /// written by other releases; new files keep the default. + /// + /// Must be called before the Extension Table or chain registry is set + /// and before any tensor is appended. + pub fn set_builtin_version(&mut self, version: &str) -> Result<(), PtwmCoreError> { + if self.table_set || self.prelude_written || !self.index_entries.is_empty() { + return Err(PtwmCoreError::InvalidContainer( + "builtin version must be set before the extension table and tensors".into(), + )); + } + if version.is_empty() || version.len() > 64 || version.contains('@') { + return Err(PtwmCoreError::InvalidContainer(format!( + "invalid builtin version {version:?}" + ))); + } + self.builtin_version = version.to_owned(); + Ok(()) + } + + /// Use `table_bytes` verbatim as the file's Extension Table and encode + /// `registry` against it. + /// + /// The transcoder uses this for tables that no table builder reproduces + /// (for example a table that mixes built-in versions). The bytes are + /// validated by [`encode_registry_for_table`] first. Plane codec indices + /// for [`codec_table_idx_for`](Self::codec_table_idx_for) come from the + /// first table entry of each codec. + /// + /// Must be called before any tensor is appended. + pub fn set_raw_extension_table( + &mut self, + table_bytes: &[u8], + registry: ChainRegistry, + ) -> Result<(), PtwmCoreError> { + if self.prelude_written || !self.index_entries.is_empty() { + return Err(PtwmCoreError::InvalidContainer( + "extension table must be set before appending tensors".into(), + )); + } + let built = encode_registry_for_table(table_bytes, ®istry)?; + self.extension_table = built.table; + self.raw_extension_table = Some(table_bytes.to_vec()); + self.chain_registry_bytes = built.registry_bytes; + self.chain_registry = registry; + self.codec_idx_map = built.codec_idx_map; + self.table_set = true; + Ok(()) } /// Write the extension table, prelude section, and chain registry section, @@ -362,10 +504,13 @@ impl ContainerWriter { // Extension table sits immediately after the header at HEADER_LEN. // Serialize it now (set_extension_table must be called before this point). - let et_bytes = self - .extension_table - .to_bytes() - .map_err(|e| PtwmCoreError::InvalidContainer(format!("extension table: {e}")))?; + let et_bytes = match &self.raw_extension_table { + Some(raw) => raw.clone(), + None => self + .extension_table + .to_bytes() + .map_err(|e| PtwmCoreError::InvalidContainer(format!("extension table: {e}")))?, + }; let et_len = et_bytes.len() as u64; let et_hash = if et_len > 0 { ExtensionTable::hash(&et_bytes) @@ -448,7 +593,7 @@ impl ContainerWriter { .map_err(|e| PtwmCoreError::Io(e.to_string()))?; let mut tensor_buf = Vec::new(); - write_tensor_record(record, &mut tensor_buf)?; + write_tensor_record_with_builtin_version(record, &mut tensor_buf, &self.builtin_version)?; let tensor_len = tensor_buf.len() as u64; self.writer @@ -699,6 +844,11 @@ pub struct ContainerReader<'a> { /// request). Populated by `open_partial`; always empty when `open` returns /// `Ok` (because `open` returns `Err` instead). pub missing_extensions: Vec, + /// Built-in version of the file's Extension Table, as returned by + /// [`crate::extension::normalize_builtin_ids`]: the release whose + /// built-in ids the writer used. `None` when the table holds no built-in + /// entry. `extension_table` holds the ids mapped to the crate version. + pub builtin_version: Option, } impl<'a> ContainerReader<'a> { @@ -740,7 +890,7 @@ impl<'a> ContainerReader<'a> { // Parse the Extension Table (sits at extension_table_offset for // extension_table_length bytes, immediately after the header). - let extension_table = if header.extension_table_length == 0 { + let mut extension_table = if header.extension_table_length == 0 { ExtensionTable::empty() } else { let et_start = usize::try_from(header.extension_table_offset).map_err(|_| { @@ -767,10 +917,25 @@ impl<'a> ContainerReader<'a> { ExtensionTable::from_bytes(table_bytes) .map_err(|e| PtwmCoreError::InvalidContainer(format!("extension table: {e}")))? }; + // Map built-in ids written by other releases to the current ones. + // This runs after the hash check, which covers the on-disk bytes. + let builtin_version = crate::extension::normalize_builtin_ids(&mut extension_table); + + // A table of built-ins only needs no discovery, keyring or trust + // pass: built-ins are trusted by construction and always available. + // Skipping it keeps opening such a file free of filesystem access. + let all_builtin = extension_table + .entries + .iter() + .all(|e| crate::extension::is_builtin(&e.canonical_id)); // Scan installed extensions once; used by both the missing-flavor check // and the trust verifier pass below. - let installed = crate::discovery::scan_all_cached().unwrap_or_default(); + let installed = if all_builtin { + Vec::new() + } else { + crate::discovery::scan_all_cached().unwrap_or_default() + }; // Check for missing non-builtin flavors immediately after parsing the // Extension Table — before touching any tensor payload. @@ -785,7 +950,7 @@ impl<'a> ContainerReader<'a> { // (i.e. not flagged as missing above), verify that its author is // trusted by the active keyring. Entries with no matching install // (already covered by MissingFlavor above) are skipped here. - { + if !all_builtin { // Build the active keyring: start from the bundled keys, merge the // user keyring on top. let mut keyring = match crate::trust::evaluate_bundled() { @@ -925,6 +1090,7 @@ impl<'a> ContainerReader<'a> { tensors_offset, chain_registry_offset, missing_extensions: needs, + builtin_version, }) } @@ -1351,6 +1517,7 @@ mod tests { use crate::codec::{CodecId, StateSource}; use crate::layout::PlaneLayout; use crate::plane_record::PlaneRecord; + use crate::tensor_record::write_tensor_record; use crate::tensor_record::{CHAIN_REF_INLINE_BIT, parse_tensor_record}; use crate::types::PlaneRole; use std::io::Cursor; @@ -1686,6 +1853,126 @@ mod tests { assert!(out.is_empty()); } + #[test] + fn set_builtin_version_must_come_first() { + let mut buf = Cursor::new(Vec::new()); + let mut w = ContainerWriter::new(&mut buf).unwrap(); + for bad in ["", "1@2", &"9".repeat(65)] { + assert!(w.set_builtin_version(bad).is_err(), "{bad:?}"); + } + w.set_builtin_version("1.0.0").unwrap(); + w.set_chain_registry(ChainRegistry { + chains: vec![(0, make_minimal_chain(1))], + }) + .unwrap(); + let err = w.set_builtin_version("1.0.1").unwrap_err().to_string(); + assert!(err.contains("must be set before"), "{err}"); + assert!( + w.extension_table + .entries + .iter() + .all(|e| e.human_label.ends_with("@1.0.0")) + ); + // Codec indices are keyed by the current ids whatever the version. + assert!( + w.codec_table_idx_for(&crate::extension::builtin_canonical_id("zstd")) + .is_some() + ); + } + + /// Serialized Extension Table of a container written with `registry`. + fn table_bytes_for(registry: &ChainRegistry, codecs: CodecSet<'_>) -> Vec { + build_extension_table(registry, codecs, crate::extension::CURRENT_BUILTIN_VERSION) + .unwrap() + .table + .to_bytes() + .unwrap() + } + + #[test] + fn raw_extension_table_is_validated() { + let registry = ChainRegistry { + chains: vec![(0, v3_passthrough_chain(4))], + }; + let good = table_bytes_for(®istry, CodecSet::AllBuiltins); + let built = encode_registry_for_table(&good, ®istry).unwrap(); + let reference = build_extension_table( + ®istry, + CodecSet::AllBuiltins, + crate::extension::CURRENT_BUILTIN_VERSION, + ) + .unwrap(); + assert_eq!(built.registry_bytes, reference.registry_bytes); + assert_eq!(built.codec_idx_map, reference.codec_idx_map); + + // Trailing bytes. + let mut trailing = good.clone(); + trailing.push(0); + let e = encode_registry_for_table(&trailing, ®istry).unwrap_err(); + assert!(e.to_string().contains("trailing bytes"), "{e}"); + + // A table that lacks a chain op. + let other = ChainRegistry { + chains: vec![(0, make_minimal_chain(0))], + }; + let only_passthrough = table_bytes_for(&other, CodecSet::Only(&[])); + let e = encode_registry_for_table(&only_passthrough, ®istry).unwrap_err(); + assert!(e.to_string().contains("missing from the table"), "{e}"); + + // Embedded WASM entries. + let mut table = ExtensionTable::from_bytes(&good).unwrap(); + table.entries[0].embedded_wasm_offset = Some(0); + table.entries[0].embedded_wasm_length = Some(1); + let wasm = table.to_bytes().unwrap(); + let e = encode_registry_for_table(&wasm, ®istry).unwrap_err(); + assert!(e.to_string().contains("embedded WASM"), "{e}"); + + // Garbage. + assert!(encode_registry_for_table(b"\xff\xff\xff\xff", ®istry).is_err()); + } + + /// A raw table is written verbatim; the registry indices resolve + /// against it, here a table whose ops come after its codecs. + #[test] + fn raw_extension_table_is_written_verbatim() { + use crate::extension::{Kind, builtin_entries}; + let registry = ChainRegistry { + chains: vec![(0, v3_passthrough_chain(4))], + }; + let mut entries: Vec<_> = builtin_entries() + .into_iter() + .filter(|e| e.kind == Kind::PlaneCodec) + .collect(); + entries + .push(crate::extension::builtin_entry_at_version("byte_passthrough", "1.0.0").unwrap()); + entries.push(crate::extension::builtin_entry_at_version("source", "1.0.0").unwrap()); + let raw = ExtensionTable { entries }.to_bytes().unwrap(); + + let mut buf = Cursor::new(Vec::new()); + { + let mut w = ContainerWriter::new(&mut buf).unwrap(); + w.set_raw_extension_table(&raw, registry.clone()).unwrap(); + let identity = crate::extension::builtin_canonical_id("identity"); + assert_eq!(w.codec_table_idx_for(&identity), Some(0)); + let mut rec = sample_tensor_record("t"); + rec.chain_ref = 0; + rec.orig_size = 4; + rec.terminals[0].codec_id = CodecId::Identity; + rec.terminals[0].codec_table_idx = 0; + rec.terminals[0].payload_bytes = vec![1, 2, 3, 4]; + rec.terminals[0].payload_len = 4; + rec.terminals[0].crc32 = None; + w.append_tensor(&rec).unwrap(); + w.finalize().unwrap(); + } + let blob = buf.into_inner(); + let reader = ContainerReader::open(&blob).unwrap(); + let et = &blob[reader.header.extension_table_offset as usize..][..raw.len()]; + assert_eq!(et, raw.as_slice()); + assert_eq!(reader.chain_registry, registry); + assert_eq!(reader.decode_tensor("t").unwrap(), vec![1, 2, 3, 4]); + } + #[test] fn writer_rejects_duplicate_tensor_names() { let mut buf = Cursor::new(Vec::new()); @@ -1781,10 +2068,10 @@ mod tests { let bytes = buf.into_inner(); let reader = ContainerReader::open(&bytes).unwrap(); assert_eq!(reader.chain_registry.chains.len(), 1); - let header = reader.header.clone(); + let header = reader.header; let patch = |len: u64| { - let mut h = header.clone(); + let mut h = header; h.shared_prelude_length = len; let mut b = bytes.clone(); b[..HEADER_LEN].copy_from_slice(&h.to_bytes().unwrap()); @@ -1793,7 +2080,7 @@ mod tests { // Section ends one byte early: the registry is truncated. let short = patch(header.shared_prelude_length - 1); - let err = ContainerReader::open(&short).err().expect("must fail"); + let err = ContainerReader::open(&short).expect_err("must fail"); assert!(err.to_string().contains("ChainRegistry"), "{err}"); // Section holds only the prelude: no registry is read. @@ -1802,11 +2089,11 @@ mod tests { assert!(reader.chain_registry.chains.is_empty()); // Section shorter than the prelude count. - let err = ContainerReader::open(&patch(2)).err().expect("must fail"); + let err = ContainerReader::open(&patch(2)).expect_err("must fail"); assert!(err.to_string().contains("prelude"), "{err}"); // A zero-length section has no prelude at all and is rejected. - let err = ContainerReader::open(&patch(0)).err().expect("must fail"); + let err = ContainerReader::open(&patch(0)).expect_err("must fail"); assert!(err.to_string().contains("prelude"), "{err}"); // An intact section puts the tensor records right after it. @@ -1816,9 +2103,7 @@ mod tests { ); // Section runs past the end of the buffer. - let err = ContainerReader::open(&patch(u64::MAX)) - .err() - .expect("must fail"); + let err = ContainerReader::open(&patch(u64::MAX)).expect_err("must fail"); assert!(err.to_string().contains("shared prelude"), "{err}"); } diff --git a/crates/ptwm-core/src/extension/builder.rs b/crates/ptwm-core/src/extension/builder.rs index 8fef2f3..3f5c7ad 100644 --- a/crates/ptwm-core/src/extension/builder.rs +++ b/crates/ptwm-core/src/extension/builder.rs @@ -2,14 +2,24 @@ use std::collections::HashMap; -use super::{CanonicalId, ExtensionTableEntry}; +use super::{ + CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionTableEntry, builtin_entry_at_version, + builtin_registry, +}; /// Intern `ExtensionTableEntry` values by canonical id, returning a /// stable `u16` index for each first insertion. -#[derive(Default)] pub struct ExtensionTableBuilder { by_id: HashMap, entries: Vec, + /// Version whose ids and labels built-in entries are interned with. + builtin_version: String, +} + +impl Default for ExtensionTableBuilder { + fn default() -> Self { + Self::with_builtin_version(CURRENT_BUILTIN_VERSION) + } } impl ExtensionTableBuilder { @@ -17,6 +27,67 @@ impl ExtensionTableBuilder { Self::default() } + /// A builder that interns built-ins with the ids and labels of + /// `version`, as a writer of that release does. The transcoder uses it + /// to rebuild tables written by other releases byte-exactly. + pub fn with_builtin_version(version: &str) -> Self { + Self { + by_id: HashMap::new(), + entries: Vec::new(), + builtin_version: version.to_owned(), + } + } + + /// A builder whose table starts as `entries`. Interning an id already in + /// `entries` returns its first index. Ids are matched as given, so + /// callers pass entries whose built-in ids were normalised to the crate + /// version. + pub fn from_entries(entries: Vec) -> Self { + let mut by_id = HashMap::with_capacity(entries.len()); + for (i, e) in entries.iter().enumerate() { + let idx = u16::try_from(i).expect("extension table index overflow (>= 65536 entries)"); + by_id.entry(e.canonical_id).or_insert(idx); + } + Self { + by_id, + entries, + builtin_version: CURRENT_BUILTIN_VERSION.to_owned(), + } + } + + /// The version built-in entries are interned with. + pub fn builtin_version(&self) -> &str { + &self.builtin_version + } + + /// Number of entries interned so far. + pub fn len(&self) -> usize { + self.entries.len() + } + + /// Whether no entry has been interned. + pub fn is_empty(&self) -> bool { + self.entries.is_empty() + } + + /// Intern the built-in `name` at the builder's built-in version. + /// + /// Entries are keyed by their current canonical id, so the same built-in + /// interns once whatever version the builder writes. Returns `None` if + /// `name` has no built-in entry. + pub fn intern_builtin(&mut self, name: &str) -> Option { + let current = builtin_registry().id_for_entry_name(name)?; + if let Some(&idx) = self.by_id.get(¤t) { + return Some(idx); + } + let entry = builtin_entry_at_version(name, &self.builtin_version)?; + let idx = u16::try_from(self.entries.len()) + .expect("extension table index overflow (>= 65536 entries)"); + self.by_id.insert(current, idx); + self.entries.push(entry); + Some(idx) + } + /// Insert `entry` if its `canonical_id` is not yet known, returning the /// index. Otherwise return the existing index. /// @@ -65,6 +136,43 @@ mod tests { assert_eq!(b.finish().len(), 2); } + #[test] + fn intern_builtin_uses_the_builder_version() { + let mut b = ExtensionTableBuilder::with_builtin_version("1.0.0"); + assert_eq!(b.builtin_version(), "1.0.0"); + assert_eq!(b.intern_builtin("source"), Some(0)); + assert_eq!(b.intern_builtin("zstd"), Some(1)); + assert_eq!(b.intern_builtin("source"), Some(0)); + assert_eq!(b.intern_builtin("not_a_builtin"), None); + let entries = b.finish(); + assert_eq!(entries.len(), 2); + assert_eq!(entries[0].human_label, "io.ptwm.builtin.source@1.0.0"); + assert_eq!( + entries[0].canonical_id, + CanonicalId::derive(&crate::extension::BUILTIN_PUBKEY, "source", "1.0.0") + ); + + // The default version produces the current entries. + let mut b = ExtensionTableBuilder::new(); + b.intern_builtin("source").unwrap(); + assert_eq!(b.finish(), vec![make_entry("source")]); + } + + #[test] + fn from_entries_returns_existing_indices() { + let entries = vec![ + make_entry("terminal"), + make_entry("source"), + make_entry("source"), + ]; + let mut b = ExtensionTableBuilder::from_entries(entries); + assert_eq!(b.len(), 3); + assert_eq!(b.intern(make_entry("source")), 1); + assert_eq!(b.intern_builtin("terminal"), Some(0)); + assert_eq!(b.intern_builtin("byte_split"), Some(3)); + assert_eq!(b.len(), 4); + } + #[test] fn same_id_interned_twice_returns_same_index() { let mut b = ExtensionTableBuilder::new(); diff --git a/crates/ptwm-core/src/extension/builtins.rs b/crates/ptwm-core/src/extension/builtins.rs index 908fad9..10d5b97 100644 --- a/crates/ptwm-core/src/extension/builtins.rs +++ b/crates/ptwm-core/src/extension/builtins.rs @@ -1,24 +1,191 @@ -//! Materialize ExtensionTableEntry values for in-tree built-ins. +//! Materialize ExtensionTableEntry values for in-tree built-ins, and the +//! static registry that maps built-in canonical ids to their names and +//! dispatch targets. + +use std::collections::HashMap; +use std::sync::OnceLock; use super::capability::CapabilityValue; -use super::table::FLAVOR_NATIVE; +use super::dispatch::BuiltinKind; +use super::table::{FLAVOR_HOST, FLAVOR_NATIVE}; use super::{ - Attestation, CanonicalId, CapabilityMap, ExtensionTableEntry, Kind, Lifecycle, - builtin_canonical_id, + Attestation, BUILTIN_PUBKEY, CanonicalId, CapabilityMap, ExtensionTable, ExtensionTableEntry, + Kind, Lifecycle, }; +use crate::codec::CodecId; +use crate::transforms::op::OpId; -pub fn builtin_entry(name: &str, kind: Kind, deterministic: bool) -> ExtensionTableEntry { +/// Version that the built-in ids and labels of this build derive from. +pub const CURRENT_BUILTIN_VERSION: &str = env!("CARGO_PKG_VERSION"); + +/// Prefix of every built-in label: `io.ptwm.builtin.@`. +pub const BUILTIN_LABEL_PREFIX: &str = "io.ptwm.builtin."; + +/// Name, kind and flavor of every built-in that has an Extension Table +/// entry, in the order of [`builtin_entries`]. Transform names must match +/// [`OpId::builtin_name`]. +const BUILTIN_SPECS: &[(&str, Kind, u8)] = &[ + // Plane codecs + ("identity", Kind::PlaneCodec, FLAVOR_NATIVE), + ("huffman", Kind::PlaneCodec, FLAVOR_NATIVE), + ("rans", Kind::PlaneCodec, FLAVOR_NATIVE), + ("zstd", Kind::PlaneCodec, FLAVOR_NATIVE), + ("zstd_dict", Kind::PlaneCodec, FLAVOR_NATIVE), + ("fpc", Kind::PlaneCodec, FLAVOR_NATIVE), + ("tans", Kind::PlaneCodec, FLAVOR_NATIVE), + ("order1_scale_ac", Kind::PlaneCodec, FLAVOR_NATIVE), + ("arithmetic_o0", Kind::PlaneCodec, FLAVOR_NATIVE), + ("arithmetic_o0_adaptive", Kind::PlaneCodec, FLAVOR_NATIVE), + ("arithmetic_o1", Kind::PlaneCodec, FLAVOR_NATIVE), + ("per_group_codebook", Kind::PlaneCodec, FLAVOR_NATIVE), + ("context_mixing_lite", Kind::PlaneCodec, FLAVOR_NATIVE), + ("huff_llm_5bit", Kind::PlaneCodec, FLAVOR_NATIVE), + ("order1_arithmetic", Kind::PlaneCodec, FLAVOR_NATIVE), + ("neural_predictor", Kind::PlaneCodec, FLAVOR_NATIVE), + // Transforms + ("bit_reorder_ieee16", Kind::Transform, FLAVOR_NATIVE), + ("bit_reorder_ieee32", Kind::Transform, FLAVOR_NATIVE), + ("bit_reorder_fp8_e4m3", Kind::Transform, FLAVOR_NATIVE), + ("bit_reorder_fp8_e5m2", Kind::Transform, FLAVOR_NATIVE), + ("byte_split", Kind::Transform, FLAVOR_NATIVE), + ("nibble_split", Kind::Transform, FLAVOR_NATIVE), + ("byte_passthrough", Kind::Transform, FLAVOR_NATIVE), + ("xor_delta", Kind::Transform, FLAVOR_NATIVE), + ("int_delta", Kind::Transform, FLAVOR_NATIVE), + ("predictor_xor", Kind::Transform, FLAVOR_NATIVE), + ("burrows_wheeler", Kind::Transform, FLAVOR_NATIVE), + ("move_to_front", Kind::Transform, FLAVOR_NATIVE), + ("mxfp4_deinterleave", Kind::Transform, FLAVOR_NATIVE), + ("block_microscaling_repack", Kind::Transform, FLAVOR_NATIVE), + ("concat", Kind::Transform, FLAVOR_NATIVE), + ("reshape", Kind::Transform, FLAVOR_NATIVE), + ("index_bitwidth_pack", Kind::Transform, FLAVOR_NATIVE), + ("mantissa_zero_strip", Kind::Transform, FLAVOR_NATIVE), + ("source", Kind::Transform, FLAVOR_NATIVE), + ("terminal", Kind::Transform, FLAVOR_NATIVE), + ("spherical_normalize", Kind::Transform, FLAVOR_NATIVE), + ("alpha_stable_normalize", Kind::Transform, FLAVOR_NATIVE), + ("float_delta", Kind::Transform, FLAVOR_NATIVE), + ("entropy_estimate", Kind::Transform, FLAVOR_NATIVE), + // Encode-only Python-resident contributions (chain builder + classifiers). + ("production_chains", Kind::ChainBuilder, FLAVOR_HOST), + ("classifier_heuristic", Kind::Classifier, FLAVOR_HOST), + ("classifier_hf_quant_config", Kind::Classifier, FLAVOR_HOST), + ("classifier_ptwm_config", Kind::Classifier, FLAVOR_HOST), + ("classifier_explicit_flags", Kind::Classifier, FLAVOR_HOST), +]; + +/// One built-in known to the registry. +#[derive(Debug, Clone, Copy)] +pub struct BuiltinRef { + /// Built-in name (`io.ptwm.builtin.`). + pub name: &'static str, + /// The in-tree op or codec this id dispatches to, if any. + pub dispatch: Option, + /// Position in [`builtin_entries`], or `None` for ids that dispatch but + /// have no Extension Table entry (`huffman_nibble`). + pub entry_index: Option, +} + +/// Lookup tables over every built-in at [`CURRENT_BUILTIN_VERSION`]. +#[derive(Debug)] +pub struct BuiltinRegistry { + by_id: HashMap, + by_name: HashMap<&'static str, CanonicalId>, + entries: Vec, +} + +impl BuiltinRegistry { + fn build() -> Self { + let mut by_id: HashMap = HashMap::new(); + let mut by_name: HashMap<&'static str, CanonicalId> = HashMap::new(); + let mut entries = Vec::with_capacity(BUILTIN_SPECS.len()); + for (i, &(name, kind, flavor)) in BUILTIN_SPECS.iter().enumerate() { + let id = CanonicalId::derive(&BUILTIN_PUBKEY, name, CURRENT_BUILTIN_VERSION); + entries.push(make_entry(name, kind, flavor, CURRENT_BUILTIN_VERSION, id)); + by_name.insert(name, id); + by_id.insert( + id, + BuiltinRef { + name, + dispatch: None, + entry_index: Some(i), + }, + ); + } + let mut add_dispatch = |name: &'static str, kind: BuiltinKind| { + let id = CanonicalId::derive(&BUILTIN_PUBKEY, name, CURRENT_BUILTIN_VERSION); + by_id + .entry(id) + .or_insert(BuiltinRef { + name, + dispatch: None, + entry_index: None, + }) + .dispatch = Some(kind); + }; + for op in OpId::ALL { + add_dispatch(op.builtin_name(), BuiltinKind::Op(op)); + } + for &(name, raw) in crate::codec::REGISTRY { + if let Some(codec) = CodecId::from_u16(raw) { + add_dispatch(name, BuiltinKind::Codec(codec)); + } + } + Self { + by_id, + by_name, + entries, + } + } + + /// The built-in with this canonical id, if any. + pub fn get(&self, id: &CanonicalId) -> Option<&BuiltinRef> { + self.by_id.get(id) + } + + /// Current canonical id of the built-in named `name` that has an + /// Extension Table entry. + pub fn id_for_entry_name(&self, name: &str) -> Option { + self.by_name.get(name).copied() + } +} + +/// The process-wide built-in registry, built on first use. +pub fn builtin_registry() -> &'static BuiltinRegistry { + static REGISTRY: OnceLock = OnceLock::new(); + REGISTRY.get_or_init(BuiltinRegistry::build) +} + +/// Canonical id of the built-in `name` at the crate version. Uses the +/// registry for known names and derives the id otherwise. +pub(super) fn current_builtin_id(name: &str) -> CanonicalId { + builtin_registry() + .by_name + .get(name) + .copied() + // Dispatch-only built-ins (`huffman_nibble`) are not in `by_name`. + .unwrap_or_else(|| CanonicalId::derive(&BUILTIN_PUBKEY, name, CURRENT_BUILTIN_VERSION)) +} + +fn make_entry( + name: &str, + kind: Kind, + flavor: u8, + version: &str, + canonical_id: CanonicalId, +) -> ExtensionTableEntry { let mut caps = CapabilityMap::new(); - caps.set("determinism", CapabilityValue::Bool(deterministic)); + caps.set("determinism", CapabilityValue::Bool(true)); caps.set("hardware_class", CapabilityValue::Text("cpu".into())); ExtensionTableEntry { - canonical_id: builtin_canonical_id(name), - human_label: format!("io.ptwm.builtin.{name}@{}", env!("CARGO_PKG_VERSION")), + canonical_id, + human_label: format!("{BUILTIN_LABEL_PREFIX}{name}@{version}"), kind, abi_version: 1, lifecycle: Lifecycle::Thread, - flavor_hints: FLAVOR_NATIVE, + flavor_hints: flavor, capabilities: caps, // Built-ins are not attested by an external author key. A zero-byte // signature marks "no external attestation needed"; the verifier @@ -30,86 +197,230 @@ pub fn builtin_entry(name: &str, kind: Kind, deterministic: bool) -> ExtensionTa } } +pub fn builtin_entry(name: &str, kind: Kind, deterministic: bool) -> ExtensionTableEntry { + let mut e = make_entry( + name, + kind, + FLAVOR_NATIVE, + CURRENT_BUILTIN_VERSION, + current_builtin_id(name), + ); + e.capabilities + .set("determinism", CapabilityValue::Bool(deterministic)); + e +} + +/// Every built-in Extension Table entry at the crate version. pub fn builtin_entries() -> Vec { - vec![ - // Plane codecs - builtin_entry("identity", Kind::PlaneCodec, true), - builtin_entry("huffman", Kind::PlaneCodec, true), - builtin_entry("rans", Kind::PlaneCodec, true), - builtin_entry("zstd", Kind::PlaneCodec, true), - builtin_entry("zstd_dict", Kind::PlaneCodec, true), - builtin_entry("fpc", Kind::PlaneCodec, true), - builtin_entry("tans", Kind::PlaneCodec, true), - builtin_entry("order1_scale_ac", Kind::PlaneCodec, true), - builtin_entry("arithmetic_o0", Kind::PlaneCodec, true), - builtin_entry("arithmetic_o0_adaptive", Kind::PlaneCodec, true), - builtin_entry("arithmetic_o1", Kind::PlaneCodec, true), - builtin_entry("per_group_codebook", Kind::PlaneCodec, true), - builtin_entry("context_mixing_lite", Kind::PlaneCodec, true), - builtin_entry("huff_llm_5bit", Kind::PlaneCodec, true), - builtin_entry("order1_arithmetic", Kind::PlaneCodec, true), - builtin_entry("neural_predictor", Kind::PlaneCodec, true), - // Transforms — names MUST match OpId::canonical_id() in transforms/op.rs - builtin_entry("bit_reorder_ieee16", Kind::Transform, true), - builtin_entry("bit_reorder_ieee32", Kind::Transform, true), - builtin_entry("bit_reorder_fp8_e4m3", Kind::Transform, true), - builtin_entry("bit_reorder_fp8_e5m2", Kind::Transform, true), - builtin_entry("byte_split", Kind::Transform, true), - builtin_entry("nibble_split", Kind::Transform, true), - builtin_entry("byte_passthrough", Kind::Transform, true), - builtin_entry("xor_delta", Kind::Transform, true), - builtin_entry("int_delta", Kind::Transform, true), - builtin_entry("predictor_xor", Kind::Transform, true), - builtin_entry("burrows_wheeler", Kind::Transform, true), - builtin_entry("move_to_front", Kind::Transform, true), - builtin_entry("mxfp4_deinterleave", Kind::Transform, true), - builtin_entry("block_microscaling_repack", Kind::Transform, true), - builtin_entry("concat", Kind::Transform, true), - builtin_entry("reshape", Kind::Transform, true), - builtin_entry("index_bitwidth_pack", Kind::Transform, true), - builtin_entry("mantissa_zero_strip", Kind::Transform, true), - builtin_entry("source", Kind::Transform, true), - builtin_entry("terminal", Kind::Transform, true), - builtin_entry("spherical_normalize", Kind::Transform, true), - builtin_entry("alpha_stable_normalize", Kind::Transform, true), - builtin_entry("float_delta", Kind::Transform, true), - builtin_entry("entropy_estimate", Kind::Transform, true), - // Encode-only Python-resident contributions (chain builder + classifiers). - { - let mut e = builtin_entry("production_chains", Kind::ChainBuilder, true); - e.flavor_hints = super::table::FLAVOR_HOST; - e - }, - { - let mut e = builtin_entry("classifier_heuristic", Kind::Classifier, true); - e.flavor_hints = super::table::FLAVOR_HOST; - e - }, - { - let mut e = builtin_entry("classifier_hf_quant_config", Kind::Classifier, true); - e.flavor_hints = super::table::FLAVOR_HOST; - e - }, - { - let mut e = builtin_entry("classifier_ptwm_config", Kind::Classifier, true); - e.flavor_hints = super::table::FLAVOR_HOST; - e - }, - { - let mut e = builtin_entry("classifier_explicit_flags", Kind::Classifier, true); - e.flavor_hints = super::table::FLAVOR_HOST; - e - }, - ] + builtin_registry().entries.clone() } +/// The Extension Table entry of the built-in `name` with its id and label +/// derived for `version`, as a writer of that version emits it. Returns +/// `None` if `name` has no built-in entry. +pub fn builtin_entry_at_version(name: &str, version: &str) -> Option { + let reg = builtin_registry(); + let idx = reg.by_id.get(reg.by_name.get(name)?)?.entry_index?; + if version == CURRENT_BUILTIN_VERSION { + return Some(reg.entries[idx].clone()); + } + let (name, kind, flavor) = BUILTIN_SPECS[idx]; + Some(make_entry( + name, + kind, + flavor, + version, + CanonicalId::derive(&BUILTIN_PUBKEY, name, version), + )) +} + +/// Whether `id` is a built-in with an Extension Table entry at the crate +/// version. Ids of other versions are built-in only after +/// [`normalize_builtin_ids`] has mapped them. pub fn is_builtin(id: &CanonicalId) -> bool { - builtin_entries().iter().any(|e| &e.canonical_id == id) + builtin_registry() + .by_id + .get(id) + .is_some_and(|r| r.entry_index.is_some()) +} + +/// Map built-in entries written by another release to the ids and labels of +/// the crate version. +/// +/// An entry is changed only when all of these hold: +/// - its label is `io.ptwm.builtin.@`; +/// - `` has a built-in entry ([`builtin_entries`]); +/// - its id equals `CanonicalId::derive(BUILTIN_PUBKEY, name, ver)`. +/// +/// Its id and label are then replaced in memory with the current ones; all +/// other fields stay as read. Entries that fail any check (third-party +/// entries, or an entry with a built-in label and some other id) are left +/// untouched, so they are still treated as non-built-in. +/// +/// Returns the built-in version the table was written with: the first +/// version other than the current one, or the current version if every +/// built-in entry already has it. Returns `None` if no entry is a built-in. +/// Writers use one version per table, so a table that mixes versions is +/// not writer output; for such a table the result names one of them. +pub fn normalize_builtin_ids(table: &mut ExtensionTable) -> Option { + let reg = builtin_registry(); + let mut found_current = false; + let mut other: Option = None; + for entry in &mut table.entries { + let Some(rest) = entry.human_label.strip_prefix(BUILTIN_LABEL_PREFIX) else { + continue; + }; + let Some((name, ver)) = rest.split_once('@') else { + continue; + }; + let Some(¤t_id) = reg.by_name.get(name) else { + continue; + }; + if ver == CURRENT_BUILTIN_VERSION { + found_current |= entry.canonical_id == current_id; + continue; + } + if entry.canonical_id != CanonicalId::derive(&BUILTIN_PUBKEY, name, ver) { + continue; + } + if other.is_none() { + other = Some(ver.to_owned()); + } + entry.canonical_id = current_id; + entry.human_label = format!("{BUILTIN_LABEL_PREFIX}{name}@{CURRENT_BUILTIN_VERSION}"); + } + other.or_else(|| found_current.then(|| CURRENT_BUILTIN_VERSION.to_owned())) } #[cfg(test)] mod tests { use super::*; + use crate::extension::builtin_canonical_id; + + fn table_at(version: &str) -> ExtensionTable { + ExtensionTable { + entries: ["source", "byte_split", "huffman", "zstd"] + .iter() + .map(|n| builtin_entry_at_version(n, version).unwrap()) + .collect(), + } + } + + #[test] + fn registry_covers_every_entry_and_dispatch_target() { + let reg = builtin_registry(); + for (i, e) in builtin_entries().iter().enumerate() { + let r = reg.get(&e.canonical_id).expect("entry in registry"); + assert_eq!(r.entry_index, Some(i)); + assert_eq!( + e.human_label, + format!("{BUILTIN_LABEL_PREFIX}{}@{CURRENT_BUILTIN_VERSION}", r.name) + ); + assert_eq!(builtin_canonical_id(r.name), e.canonical_id); + } + for op in OpId::ALL { + assert_eq!( + reg.get(&op.canonical_id()).and_then(|r| r.dispatch), + Some(BuiltinKind::Op(op)) + ); + } + // `huffman_nibble` dispatches but has no entry, so it is not a + // built-in for `is_builtin`. + let nibble = builtin_canonical_id("huffman_nibble"); + assert_eq!( + reg.get(&nibble).and_then(|r| r.dispatch), + Some(BuiltinKind::Codec(CodecId::HuffmanNibble)) + ); + assert!(!is_builtin(&nibble)); + assert!(builtin_entry_at_version("huffman_nibble", "1.0.0").is_none()); + } + + #[test] + fn builtin_entry_at_version_derives_id_and_label() { + let e = builtin_entry_at_version("huffman", "1.0.0").unwrap(); + assert_eq!(e.human_label, "io.ptwm.builtin.huffman@1.0.0"); + assert_eq!( + e.canonical_id, + CanonicalId::derive(&BUILTIN_PUBKEY, "huffman", "1.0.0") + ); + let cur = builtin_entry_at_version("huffman", CURRENT_BUILTIN_VERSION).unwrap(); + assert_eq!(cur, builtin_entries()[1]); + // Only id and label depend on the version. + let mut e2 = e.clone(); + e2.canonical_id = cur.canonical_id; + e2.human_label = cur.human_label.clone(); + assert_eq!(e2, cur); + // Host-flavored built-ins keep their flavor. + let host = builtin_entry_at_version("production_chains", "1.0.0").unwrap(); + assert_eq!(host.flavor_hints, FLAVOR_HOST); + assert!(builtin_entry_at_version("not_a_builtin", "1.0.0").is_none()); + } + + #[test] + fn ids_of_other_versions_are_builtin_only_after_normalisation() { + let old = "0.9.0"; + let mut table = table_at(old); + for e in &table.entries { + assert!(!is_builtin(&e.canonical_id)); + assert_eq!(crate::extension::dispatch_builtin(&e.canonical_id), None); + } + assert_eq!(normalize_builtin_ids(&mut table).as_deref(), Some(old)); + assert_eq!(table, table_at(CURRENT_BUILTIN_VERSION)); + for e in &table.entries { + assert!(is_builtin(&e.canonical_id)); + } + assert_eq!( + crate::extension::dispatch_builtin(&table.entries[2].canonical_id), + Some(BuiltinKind::Codec(CodecId::Huffman)) + ); + // Normalising again changes nothing and reports the current version. + assert_eq!( + normalize_builtin_ids(&mut table).as_deref(), + Some(CURRENT_BUILTIN_VERSION) + ); + assert_eq!(table, table_at(CURRENT_BUILTIN_VERSION)); + } + + #[test] + fn forged_builtin_labels_are_not_normalised() { + let mut table = table_at("1.0.0"); + // Built-in label, but an id that does not re-derive from it. + table.entries[0].canonical_id = CanonicalId::from_bytes([0x11; 32]); + // Built-in label of the current version with a wrong id. + table.entries[1].human_label = + format!("io.ptwm.builtin.byte_split@{CURRENT_BUILTIN_VERSION}"); + // Id of one name under the label of another. + table.entries[2].human_label = "io.ptwm.builtin.zstd@1.0.0".into(); + let before = table.clone(); + let version = normalize_builtin_ids(&mut table); + assert_eq!(table.entries[..3], before.entries[..3]); + for e in &table.entries[..3] { + assert!(!is_builtin(&e.canonical_id)); + } + // The genuine entry is still mapped. + assert!(is_builtin(&table.entries[3].canonical_id)); + assert_eq!(version.as_deref(), Some("1.0.0")); + } + + #[test] + fn normalisation_ignores_third_party_and_unknown_names() { + let id = CanonicalId::derive(&[0x22; 32], "huffman", "1.0.0"); + let mut entry = builtin_entry_at_version("huffman", "1.0.0").unwrap(); + entry.canonical_id = id; + entry.human_label = "com.example.huffman@1.0.0".into(); + let mut unknown = entry.clone(); + unknown.canonical_id = CanonicalId::derive(&BUILTIN_PUBKEY, "nope", "1.0.0"); + unknown.human_label = "io.ptwm.builtin.nope@1.0.0".into(); + let mut no_version = entry.clone(); + no_version.human_label = "io.ptwm.builtin.huffman".into(); + let mut table = ExtensionTable { + entries: vec![entry, unknown, no_version], + }; + let before = table.clone(); + assert_eq!(normalize_builtin_ids(&mut table), None); + assert_eq!(table, before); + assert_eq!(normalize_builtin_ids(&mut ExtensionTable::empty()), None); + } #[test] fn every_builtin_has_unique_id() { diff --git a/crates/ptwm-core/src/extension/dispatch.rs b/crates/ptwm-core/src/extension/dispatch.rs index 599febb..2d63228 100644 --- a/crates/ptwm-core/src/extension/dispatch.rs +++ b/crates/ptwm-core/src/extension/dispatch.rs @@ -4,7 +4,8 @@ //! [`CodecId`] for files that reference contributions via the Extension Table. use crate::codec::CodecId; -use crate::extension::{CanonicalId, builtin_canonical_id}; +use crate::extension::CanonicalId; +use crate::extension::builtins::builtin_registry; use crate::transforms::op::OpId; /// The kind of a resolved built-in contribution. @@ -18,142 +19,18 @@ pub enum BuiltinKind { /// /// Returns `Some(BuiltinKind)` for every built-in op and codec that ships /// with ptwm-core, or `None` if the id is not recognized (e.g. a third-party -/// extension handled via dynamic loading instead). +/// extension handled via dynamic loading instead). Only ids of the crate +/// version match; tables from other releases go through +/// [`crate::extension::normalize_builtin_ids`] first. The lookup is a single +/// hash-map probe. pub fn dispatch_builtin(id: &CanonicalId) -> Option { - // Ops — 16 variants. - if *id == OpId::Source.canonical_id() { - return Some(BuiltinKind::Op(OpId::Source)); - } - if *id == OpId::Terminal.canonical_id() { - return Some(BuiltinKind::Op(OpId::Terminal)); - } - if *id == OpId::BitReorderIeee16.canonical_id() { - return Some(BuiltinKind::Op(OpId::BitReorderIeee16)); - } - if *id == OpId::BitReorderIeee32.canonical_id() { - return Some(BuiltinKind::Op(OpId::BitReorderIeee32)); - } - if *id == OpId::BitReorderFp8E4M3.canonical_id() { - return Some(BuiltinKind::Op(OpId::BitReorderFp8E4M3)); - } - if *id == OpId::BitReorderFp8E5M2.canonical_id() { - return Some(BuiltinKind::Op(OpId::BitReorderFp8E5M2)); - } - if *id == OpId::ByteSplit.canonical_id() { - return Some(BuiltinKind::Op(OpId::ByteSplit)); - } - if *id == OpId::NibbleSplit.canonical_id() { - return Some(BuiltinKind::Op(OpId::NibbleSplit)); - } - if *id == OpId::BytePassthrough.canonical_id() { - return Some(BuiltinKind::Op(OpId::BytePassthrough)); - } - if *id == OpId::MxFp4Deinterleave.canonical_id() { - return Some(BuiltinKind::Op(OpId::MxFp4Deinterleave)); - } - if *id == OpId::BlockMicroscalingRepack.canonical_id() { - return Some(BuiltinKind::Op(OpId::BlockMicroscalingRepack)); - } - if *id == OpId::XorDelta.canonical_id() { - return Some(BuiltinKind::Op(OpId::XorDelta)); - } - if *id == OpId::FloatDelta.canonical_id() { - return Some(BuiltinKind::Op(OpId::FloatDelta)); - } - if *id == OpId::IntDelta.canonical_id() { - return Some(BuiltinKind::Op(OpId::IntDelta)); - } - if *id == OpId::PredictorXor.canonical_id() { - return Some(BuiltinKind::Op(OpId::PredictorXor)); - } - if *id == OpId::BurrowsWheeler.canonical_id() { - return Some(BuiltinKind::Op(OpId::BurrowsWheeler)); - } - if *id == OpId::MoveToFront.canonical_id() { - return Some(BuiltinKind::Op(OpId::MoveToFront)); - } - if *id == OpId::IndexBitwidthPack.canonical_id() { - return Some(BuiltinKind::Op(OpId::IndexBitwidthPack)); - } - if *id == OpId::EntropyEstimate.canonical_id() { - return Some(BuiltinKind::Op(OpId::EntropyEstimate)); - } - if *id == OpId::Concat.canonical_id() { - return Some(BuiltinKind::Op(OpId::Concat)); - } - if *id == OpId::Reshape.canonical_id() { - return Some(BuiltinKind::Op(OpId::Reshape)); - } - if *id == OpId::MantissaZeroStrip.canonical_id() { - return Some(BuiltinKind::Op(OpId::MantissaZeroStrip)); - } - if *id == OpId::SphericalNormalize.canonical_id() { - return Some(BuiltinKind::Op(OpId::SphericalNormalize)); - } - if *id == OpId::AlphaStableNormalize.canonical_id() { - return Some(BuiltinKind::Op(OpId::AlphaStableNormalize)); - } - - // Codecs — 7 variants (including HuffmanNibble which has no struct, - // represented only as CodecId for wire compatibility). - if *id == builtin_canonical_id("identity") { - return Some(BuiltinKind::Codec(CodecId::Identity)); - } - if *id == builtin_canonical_id("huffman") { - return Some(BuiltinKind::Codec(CodecId::Huffman)); - } - if *id == builtin_canonical_id("huffman_nibble") { - return Some(BuiltinKind::Codec(CodecId::HuffmanNibble)); - } - if *id == builtin_canonical_id("rans") { - return Some(BuiltinKind::Codec(CodecId::Rans)); - } - if *id == builtin_canonical_id("zstd") { - return Some(BuiltinKind::Codec(CodecId::Zstd)); - } - if *id == builtin_canonical_id("zstd_dict") { - return Some(BuiltinKind::Codec(CodecId::ZstdDict)); - } - if *id == builtin_canonical_id("fpc") { - return Some(BuiltinKind::Codec(CodecId::Fpc)); - } - if *id == builtin_canonical_id("tans") { - return Some(BuiltinKind::Codec(CodecId::Tans)); - } - if *id == builtin_canonical_id("per_group_codebook") { - return Some(BuiltinKind::Codec(CodecId::PerGroupCodebook)); - } - if *id == builtin_canonical_id("order1_scale_ac") { - return Some(BuiltinKind::Codec(CodecId::Order1ScaleAC)); - } - if *id == builtin_canonical_id("arithmetic_o0") { - return Some(BuiltinKind::Codec(CodecId::ArithmeticO0)); - } - if *id == builtin_canonical_id("arithmetic_o0_adaptive") { - return Some(BuiltinKind::Codec(CodecId::ArithmeticO0Adaptive)); - } - if *id == builtin_canonical_id("arithmetic_o1") { - return Some(BuiltinKind::Codec(CodecId::ArithmeticO1)); - } - if *id == builtin_canonical_id("context_mixing_lite") { - return Some(BuiltinKind::Codec(CodecId::ContextMixingLite)); - } - if *id == builtin_canonical_id("huff_llm_5bit") { - return Some(BuiltinKind::Codec(CodecId::HuffLlm5Bit)); - } - if *id == builtin_canonical_id("order1_arithmetic") { - return Some(BuiltinKind::Codec(CodecId::Order1Arithmetic)); - } - if *id == builtin_canonical_id("neural_predictor") { - return Some(BuiltinKind::Codec(CodecId::NeuralPredictor)); - } - - None + builtin_registry().get(id).and_then(|r| r.dispatch) } #[cfg(test)] mod tests { use super::*; + use crate::extension::builtin_canonical_id; #[test] fn every_op_id_resolves() { diff --git a/crates/ptwm-core/src/extension/mod.rs b/crates/ptwm-core/src/extension/mod.rs index d6a2be5..76f9f19 100644 --- a/crates/ptwm-core/src/extension/mod.rs +++ b/crates/ptwm-core/src/extension/mod.rs @@ -15,7 +15,10 @@ pub mod resolve; pub mod table; pub use builder::ExtensionTableBuilder; -pub use builtins::{builtin_entries, builtin_entry, is_builtin}; +pub use builtins::{ + BUILTIN_LABEL_PREFIX, BuiltinRef, BuiltinRegistry, CURRENT_BUILTIN_VERSION, builtin_entries, + builtin_entry, builtin_entry_at_version, builtin_registry, is_builtin, normalize_builtin_ids, +}; pub use capability::{Capability, CapabilityMap}; pub use dispatch::{BuiltinKind, dispatch_builtin}; pub use error::ExtensionError; @@ -32,9 +35,10 @@ pub const BUILTIN_PUBKEY: [u8; 32] = [0xBB; 32]; /// Compute the canonical id of an in-tree built-in by name. /// The version is the PTWM crate version (`CARGO_PKG_VERSION`), so -/// built-in IDs are stable per release. +/// built-in IDs change with each release. Files from other releases are +/// mapped to these ids by [`normalize_builtin_ids`]. pub fn builtin_canonical_id(name: &str) -> CanonicalId { - CanonicalId::derive(&BUILTIN_PUBKEY, name, env!("CARGO_PKG_VERSION")) + builtins::current_builtin_id(name) } #[cfg(test)] diff --git a/crates/ptwm-core/src/tensor_record.rs b/crates/ptwm-core/src/tensor_record.rs index 5a201cf..41354da 100644 --- a/crates/ptwm-core/src/tensor_record.rs +++ b/crates/ptwm-core/src/tensor_record.rs @@ -8,8 +8,8 @@ use crate::chain::Chain; use crate::chain::wire::{read_chain, write_chain}; use crate::error::PtwmCoreError; -use crate::extension::ExtensionTableBuilder; use crate::extension::table::ExtensionTable; +use crate::extension::{CURRENT_BUILTIN_VERSION, ExtensionTableBuilder, normalize_builtin_ids}; use crate::plane_record::{PlaneRecord, parse_plane_record, write_plane_record}; pub const TENSOR_FLAG_DEPENDENCIES: u8 = 0x01; @@ -95,8 +95,19 @@ pub(crate) fn checked_len>(n: usize, what: &str) -> Result) -> Result<(), PtwmCoreError> { + write_tensor_record_with_builtin_version(rec, out, CURRENT_BUILTIN_VERSION) +} + +/// [`write_tensor_record`] with the inline chain's Extension Table written +/// with the built-in ids and labels of `builtin_version`. The transcoder +/// uses it to reproduce records written by other releases. +pub fn write_tensor_record_with_builtin_version( + rec: &TensorRecord, + out: &mut Vec, + builtin_version: &str, +) -> Result<(), PtwmCoreError> { let start = out.len(); - let result = write_tensor_record_unchecked(rec, out); + let result = write_tensor_record_unchecked(rec, out, builtin_version); if result.is_err() { out.truncate(start); } @@ -106,6 +117,7 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), fn write_tensor_record_unchecked( rec: &TensorRecord, out: &mut Vec, + builtin_version: &str, ) -> Result<(), PtwmCoreError> { // Validate inline_chain vs chain_ref consistency. let has_inline = rec.chain_ref & CHAIN_REF_INLINE_BIT != 0; @@ -183,7 +195,7 @@ fn write_tensor_record_unchecked( // inline_et_entries: [ExtensionTableEntry; count] // chain_bytes: (num_nodes, ..., nodes_section, edges_section, terminals_section) if let Some(chain) = &rec.inline_chain { - let mut builder = ExtensionTableBuilder::new(); + let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version); let mut chain_bytes = Vec::new(); write_chain(chain, &mut builder, &mut chain_bytes)?; let inline_table = ExtensionTable { @@ -249,6 +261,19 @@ pub fn peek_tensor_record_name(buf: &[u8]) -> Result<&str, PtwmCoreError> { } pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCoreError> { + parse_tensor_record_with_builtin_version(buf).map(|(rec, consumed, _)| (rec, consumed)) +} + +/// [`parse_tensor_record`] that also returns the built-in version of the +/// inline chain's Extension Table (see +/// [`crate::extension::normalize_builtin_ids`]), or `None` when the record +/// has no inline chain or its table holds no built-in entry. +/// +/// Inline tables written by another release are mapped to the current +/// built-in ids before the chain is resolved. +pub fn parse_tensor_record_with_builtin_version( + buf: &[u8], +) -> Result<(TensorRecord, usize, Option), PtwmCoreError> { if buf.len() < 18 { return Err(PtwmCoreError::InvalidContainer( "tensor record header truncated".into(), @@ -314,6 +339,7 @@ pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCore // // Wire layout: a local extension table precedes the chain // bytes. Read it first, then pass it to `read_chain`. + let mut inline_version = None; let inline_chain = if chain_ref & CHAIN_REF_INLINE_BIT != 0 { // Read the inline extension table (count: u32 + entries) directly // from the buffer slice via a Cursor. The cursor's position tells @@ -322,10 +348,11 @@ pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCore // tables — count is already bounded by MAX_TABLE_ENTRIES, but // the re-serialize was redundant work either way). let mut cursor = std::io::Cursor::new(&buf[pos..]); - let inline_table = ExtensionTable::read(&mut cursor).map_err(|e| { + let mut inline_table = ExtensionTable::read(&mut cursor).map_err(|e| { PtwmCoreError::InvalidContainer(format!("inline chain extension table: {e}")) })?; pos += cursor.position() as usize; + inline_version = normalize_builtin_ids(&mut inline_table); let (chain, consumed) = read_chain(&buf[pos..], &inline_table).map_err(|e| { PtwmCoreError::InvalidContainer(format!("inline chain parse error: {e}")) })?; @@ -411,6 +438,7 @@ pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCore terminals, }, pos, + inline_version, )) } diff --git a/crates/ptwm-core/src/transcode/manifest.rs b/crates/ptwm-core/src/transcode/manifest.rs index 6122099..5c02725 100644 --- a/crates/ptwm-core/src/transcode/manifest.rs +++ b/crates/ptwm-core/src/transcode/manifest.rs @@ -18,6 +18,7 @@ pub const MANIFEST_VERSION: u32 = 1; /// `shared_prelude` to archive scope so codecs/ops/codebooks/chains are defined /// once and referenced by id. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct RegistryManifest { pub ptwm_format: String, pub manifest_version: u32, @@ -28,19 +29,33 @@ pub struct RegistryManifest { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct ContainerFlags { pub embedded_wasm: bool, + /// Release whose built-in Extension Table ids and labels the container + /// was written with, when it is not the release reading it. `implode` + /// writes built-in entries (file-level and inline-chain tables) with + /// this version. Absent for the crate version. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub builtin_version: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct Registry { pub ops: Vec, pub codecs: Vec, pub shared_state: Vec, pub chains: Vec, + /// The container's Extension Table bytes (base64), present only when no + /// table builder reproduces them from `chains`, the codec list and + /// `builtin_version`. `implode` then writes them verbatim. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub extension_table_b64: Option, } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct OpEntry { /// Stable wire op id (`OpId::as_u16`). pub id: u16, @@ -49,6 +64,7 @@ pub struct OpEntry { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct CodecEntry { /// Wire-stable `CodecId` value. pub id: u16, @@ -60,6 +76,7 @@ pub struct CodecEntry { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct SharedStateEntry { pub id: u16, pub codec: u16, @@ -73,6 +90,7 @@ pub struct SharedStateEntry { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct ChainEntry { pub id: u16, pub nodes: Vec, @@ -81,6 +99,7 @@ pub struct ChainEntry { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct NodeJson { /// Stable wire op id (`OpId::as_u16`). pub op: u16, @@ -88,6 +107,7 @@ pub struct NodeJson { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct EdgeJson { pub src: u8, pub src_out: u8, @@ -98,6 +118,7 @@ pub struct EdgeJson { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct TerminalJson { pub node: u8, pub out: u8, @@ -109,6 +130,7 @@ pub struct TerminalJson { /// Per-tensor manifest — a direct projection of `TensorRecord`. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct TensorManifest { pub name: String, /// Canonical dtype name from the CBOR metadata; `null` when absent. @@ -132,6 +154,7 @@ pub struct TensorManifest { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct ChainSpec { /// Registry chain index when shared; `null` when the chain is inlined. #[serde(rename = "ref")] @@ -143,6 +166,7 @@ pub struct ChainSpec { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct DependencyJson { pub ref_kind: String, pub ref_b64: String, @@ -150,6 +174,7 @@ pub struct DependencyJson { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct PlaneJson { pub index: u32, /// Human-readable plane role label. @@ -184,6 +209,7 @@ pub struct PlaneJson { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct ExternalStateJson { pub ref_kind: u8, pub ref_b64: String, @@ -192,6 +218,7 @@ pub struct ExternalStateJson { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct ChunkJson { pub offset_in_payload: u32, pub decoded_size: u32, diff --git a/crates/ptwm-core/src/transcode/mod.rs b/crates/ptwm-core/src/transcode/mod.rs index 988eebf..e3b6659 100644 --- a/crates/ptwm-core/src/transcode/mod.rs +++ b/crates/ptwm-core/src/transcode/mod.rs @@ -24,17 +24,19 @@ use std::io::Cursor; use crate::chain::graph::{Chain, ChainEdge, ChainNode, TerminalRef}; use crate::codec::{CodecId, StateSource}; use crate::container::{ - ChainRegistry, CodecSet, ContainerReader, ContainerWriter, build_extension_table, + BuiltExtensionTable, ChainRegistry, CodecSet, ContainerReader, ContainerWriter, + build_extension_table, encode_registry_for_table, }; use crate::error::PtwmCoreError; -use crate::extension::ExtensionTable; +use crate::extension::{CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionTable}; use crate::header::FLAG_PTWX; use crate::layout::PlaneLayout; use crate::metadata::{decode_shape_metadata, encode_shape_metadata}; use crate::plane_record::{ChunkEntry, ExternalStateRef, PlaneRecord}; -use crate::prelude::PreludeEntry; +use crate::prelude::{PreludeEntry, write_prelude}; use crate::tensor_record::{ - CHAIN_REF_INLINE_BIT, Dependency, RefKind, TensorRecord, parse_tensor_record, + CHAIN_REF_INLINE_BIT, Dependency, RefKind, TensorRecord, + parse_tensor_record_with_builtin_version, }; use crate::transforms::op::OpId; use crate::types::PlaneRole; @@ -232,6 +234,73 @@ fn shared_state_key(shared_state_id: u16) -> String { format!("shared/s{shared_state_id:05}.bin") } +// ── Extension Table reconstruction ──────────────────────────────────────────── + +/// Which table builder reproduced a container's Extension Table. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum TableLayout { + /// Chain ops plus the codecs the planes use (current writers). + Trimmed, + /// Chain ops plus every built-in codec (v1.1 and earlier). + AllBuiltins, +} + +/// Current canonical ids of the manifest's codec list. The list is +/// archive-wide, so a one-tensor implode still rebuilds the archive's table. +fn used_codec_ids(codecs: &[CodecEntry]) -> Vec { + codecs + .iter() + .filter_map(|c| CodecId::from_u16(c.id)) + .map(crate::compressor::codec_canonical_id) + .collect() +} + +/// Header-style hash of serialized table bytes. +fn table_hash(bytes: &[u8]) -> [u8; 32] { + if bytes.is_empty() { + [0u8; 32] + } else { + ExtensionTable::hash(bytes) + } +} + +/// Rebuild the Extension Table from the chain registry, trying the trimmed +/// layout and then the all-built-ins layout, both with built-in entries of +/// `version`. Returns the first layout whose table bytes hash to +/// `expected_hash`, or `None`. +fn rebuild_extension_table( + registry: &ChainRegistry, + used_codecs: &[CanonicalId], + version: &str, + expected_hash: &[u8; 32], +) -> Option<(TableLayout, BuiltExtensionTable)> { + [ + (TableLayout::Trimmed, CodecSet::Only(used_codecs)), + (TableLayout::AllBuiltins, CodecSet::AllBuiltins), + ] + .into_iter() + .find_map(|(layout, codecs)| { + let built = build_extension_table(registry, codecs, version).ok()?; + let bytes = built.table.to_bytes().ok()?; + (table_hash(&bytes) == *expected_hash).then_some((layout, built)) + }) +} + +/// The container's on-disk Extension Table bytes. +fn extension_table_bytes<'a>( + reader: &ContainerReader<'_>, + blob: &'a [u8], +) -> Result<&'a [u8], PtwmCoreError> { + let start = usize::try_from(reader.header.extension_table_offset) + .map_err(|_| err("extension_table_offset overflows usize"))?; + let len = usize::try_from(reader.header.extension_table_length) + .map_err(|_| err("extension_table_length overflows usize"))?; + start + .checked_add(len) + .and_then(|end| blob.get(start..end)) + .ok_or_else(|| err("extension table region exceeds buffer")) +} + // ── Explode ─────────────────────────────────────────────────────────────────── /// Parse a `.ptwm` container into a registry manifest, per-tensor manifests, @@ -285,6 +354,7 @@ pub fn explode(blob: &[u8]) -> Result { let mut codec_ids: std::collections::BTreeSet = std::collections::BTreeSet::new(); let mut tensors: Vec = Vec::with_capacity(index.len()); + let mut inline_builtin_version: Option = None; for (file_idx, entry) in index.iter().enumerate() { // `usize` is ≥ 32 bits on every supported platform; widening these @@ -294,7 +364,24 @@ pub fn explode(blob: &[u8]) -> Result { .checked_add(entry.tensor_len as usize) .filter(|e| *e <= blob.len()) .ok_or_else(|| err("tensor record slice out of bounds"))?; - let (record, _) = parse_tensor_record(&blob[start..end])?; + let (record, _, inline_version) = + parse_tensor_record_with_builtin_version(&blob[start..end])?; + // `implode` writes every inline-chain table with one built-in + // version, so they must agree. + if let Some(v) = inline_version { + match &inline_builtin_version { + None => inline_builtin_version = Some(v), + Some(b) if *b != v => { + return Err(err(format!( + "tensor {:?}: inline chain built-in version {v} differs from \ + {b} in an earlier inline chain; the container cannot be \ + transcoded byte-exactly", + record.name + ))); + } + Some(_) => {} + } + } for plane in &record.terminals { codec_ids.insert(plane.codec_id.as_u16()); } @@ -331,16 +418,59 @@ pub fn explode(blob: &[u8]) -> Result { }) .collect(); + // `implode` writes built-in entries with one version. Inline-chain + // tables are always re-encoded, so their version wins; the file-level + // table is rebuilt with it when that reproduces the bytes. + let builtin_version = inline_builtin_version.or_else(|| reader.builtin_version.clone()); + let version = builtin_version + .as_deref() + .unwrap_or(CURRENT_BUILTIN_VERSION); + + // Check that implode can rebuild the Extension Table and the + // shared-prelude section byte-exactly. When no table builder reproduces + // the table, carry its bytes in the manifest. + let used_codecs = used_codec_ids(&codecs); + let et_bytes = extension_table_bytes(&reader, blob)?; + let (built, extension_table_b64) = match rebuild_extension_table( + &reader.chain_registry, + &used_codecs, + version, + &reader.header.extension_table_hash, + ) { + Some((_, built)) => (built, None), + None => ( + encode_registry_for_table(et_bytes, &reader.chain_registry)?, + Some(base64::encode(et_bytes)), + ), + }; + let mut section = Vec::new(); + write_prelude(&reader.prelude, &mut section)?; + section.extend_from_slice(&built.registry_bytes); + let prelude_off = usize::try_from(reader.header.shared_prelude_offset) + .map_err(|_| err("shared_prelude_offset overflows usize"))?; + let prelude_end = usize::try_from(reader.tensors_offset) + .map_err(|_| err("tensors_offset overflows usize"))?; + if blob.get(prelude_off..prelude_end) != Some(section.as_slice()) { + return Err(err( + "the shared prelude or chain registry differs from its re-encoding; \ + the container cannot be transcoded byte-exactly", + )); + } + let registry = RegistryManifest { ptwm_format: PTWM_MANIFEST_FORMAT.to_string(), manifest_version: MANIFEST_VERSION, - container_flags: ContainerFlags { embedded_wasm }, + container_flags: ContainerFlags { + embedded_wasm, + builtin_version: builtin_version.filter(|v| v != CURRENT_BUILTIN_VERSION), + }, registry_hash_blake3: hex::encode(reader.header.extension_table_hash), registry: Registry { ops, codecs, shared_state, chains, + extension_table_b64, }, }; let registry_value = @@ -502,19 +632,32 @@ pub fn implode( tensor_jsons: &[String], mut members: HashMap, ) -> Result, PtwmCoreError> { - let registry: RegistryManifest = serde_json::from_str(registry_json) + // Check the format and version before the strict parse, so a manifest + // of another format or version gets a clear error. + let value: serde_json::Value = serde_json::from_str(registry_json) .map_err(|e| err(format!("parse registry manifest: {e}")))?; - if registry.ptwm_format != PTWM_MANIFEST_FORMAT { + let format = value.get("ptwm_format").and_then(|v| v.as_str()); + let version = value.get("manifest_version").and_then(|v| v.as_u64()); + if format != Some(PTWM_MANIFEST_FORMAT) || version != Some(u64::from(MANIFEST_VERSION)) { return Err(err(format!( - "unexpected ptwm_format {:?}", - registry.ptwm_format + "unsupported manifest: ptwm_format {format:?}, manifest_version {version:?} \ + (this reader handles {PTWM_MANIFEST_FORMAT:?} version {MANIFEST_VERSION})" ))); } + let registry: RegistryManifest = + serde_json::from_value(value).map_err(|e| err(format!("parse registry manifest: {e}")))?; if registry.container_flags.embedded_wasm { return Err(err( "embedded-WASM (.ptwx) containers are not supported by the transcoder", )); } + let expected_hash = parse_hash32(®istry.registry_hash_blake3) + .map_err(|_| err("registry_hash_blake3 must be 32 bytes of hex"))?; + let builtin_version = registry + .container_flags + .builtin_version + .as_deref() + .unwrap_or(CURRENT_BUILTIN_VERSION); // Rebuild the shared-state prelude (order preserved). let mut prelude_entries: Vec = @@ -560,33 +703,47 @@ pub fn implode( { let mut writer = ContainerWriter::new(&mut cursor)?; writer.set_prelude(&prelude_entries)?; + // Built-in entries (file-level and inline-chain tables) carry the + // ids and labels of the release that wrote the container. + writer.set_builtin_version(builtin_version)?; // Containers written by v1.1 and earlier intern every built-in codec; // later writers intern only the codecs the planes reference. Rebuild // whichever layout reproduces the recorded table hash, so the planes' - // `codec_table_idx` values resolve to the same entries. + // `codec_table_idx` values resolve to the same entries. A table that + // neither layout reproduces travels as `extension_table_b64`. // // Take the codec set from the archive-wide registry, not from // `records`: a caller may implode a subset of the tensors (the // exploded store layouts rebuild one tensor at a time). - let used_codecs: Vec = registry - .registry - .codecs - .iter() - .filter_map(|c| CodecId::from_u16(c.id)) - .map(crate::compressor::codec_canonical_id) - .collect(); - let trimmed = build_extension_table(&chain_registry, CodecSet::Only(&used_codecs)).ok(); - let trimmed_matches = trimmed.is_some_and(|built| { - built - .table - .to_bytes() - .map(|b| hex::encode(ExtensionTable::hash(&b)) == registry.registry_hash_blake3) - .unwrap_or(false) - }); - if trimmed_matches { - writer.set_chain_registry_with_codecs(chain_registry, &used_codecs)?; + if let Some(b64) = ®istry.registry.extension_table_b64 { + let bytes = + base64::decode(b64).ok_or_else(|| err("extension_table_b64 not valid base64"))?; + if table_hash(&bytes) != expected_hash { + return Err(err( + "extension_table_b64 does not match registry_hash_blake3", + )); + } + writer.set_raw_extension_table(&bytes, chain_registry)?; } else { - writer.set_chain_registry(chain_registry)?; + let used_codecs = used_codec_ids(®istry.registry.codecs); + match rebuild_extension_table( + &chain_registry, + &used_codecs, + builtin_version, + &expected_hash, + ) { + Some((TableLayout::Trimmed, _)) => { + writer.set_chain_registry_with_codecs(chain_registry, &used_codecs)?; + } + Some((TableLayout::AllBuiltins, _)) => writer.set_chain_registry(chain_registry)?, + None => { + return Err(err(format!( + "no Extension Table layout reproduces registry_hash_blake3 {} \ + (built-in version {builtin_version})", + registry.registry_hash_blake3 + ))); + } + } } for record in &records { writer.append_tensor(record)?; diff --git a/crates/ptwm-core/src/transcode/tests.rs b/crates/ptwm-core/src/transcode/tests.rs index d53e4ba..ef5de2b 100644 --- a/crates/ptwm-core/src/transcode/tests.rs +++ b/crates/ptwm-core/src/transcode/tests.rs @@ -358,3 +358,196 @@ fn round_trip_trimmed_codec_table() { "imploded subset must resolve b.weight's plane to rANS" ); } + +fn minimal_container() -> Vec { + let registry = ChainRegistry { + chains: vec![(0u16, passthrough_chain())], + }; + let rec = TensorRecord { + chain_ref: 0, + dtype_code: 6, + input_format: 2, + flags: 0, + orig_size: 8, + name: "w".to_string(), + payload_hash: None, + inline_chain: None, + dependencies: Vec::new(), + tensor_metadata: None, + terminals: vec![flat_plane( + PlaneRole::Value, + CodecId::Identity, + vec![1, 2, 3, 4, 5, 6, 7, 8], + )], + }; + build_container(&[], registry, &[rec]) +} + +/// Implode `registry` (a JSON value) with the tensors and members of +/// `exploded`. +fn implode_value( + exploded: &Exploded, + registry: &serde_json::Value, +) -> Result, PtwmCoreError> { + implode( + ®istry.to_string(), + &tensor_jsons(exploded), + members_map(exploded), + ) +} + +#[test] +fn implode_rejects_other_manifest_format_or_version() { + let blob = minimal_container(); + let exploded = explode(&blob).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + assert_eq!(implode_value(&exploded, ®).unwrap(), blob); + + for (field, value) in [ + ("ptwm_format", serde_json::json!("ptwm-manifest-c")), + ("manifest_version", serde_json::json!(2)), + ("manifest_version", serde_json::json!(0)), + ] { + let mut bad = reg.clone(); + bad[field] = value.clone(); + let e = implode_value(&exploded, &bad).unwrap_err(); + assert!( + e.to_string().contains("unsupported manifest"), + "{field}={value}: {e}" + ); + } + let mut missing = reg.clone(); + missing.as_object_mut().unwrap().remove("manifest_version"); + assert!(implode_value(&exploded, &missing).is_err()); +} + +#[test] +fn implode_rejects_unknown_fields() { + let blob = minimal_container(); + let exploded = explode(&blob).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + + let mut cases = Vec::new(); + let mut top = reg.clone(); + top["layout"] = serde_json::json!("compact"); + cases.push(top); + let mut flags = reg.clone(); + flags["container_flags"]["layout"] = serde_json::json!("compact"); + cases.push(flags); + let mut registry = reg.clone(); + registry["registry"]["extension_table"] = serde_json::json!([]); + cases.push(registry); + let mut chain = reg.clone(); + chain["registry"]["chains"][0]["extra"] = serde_json::json!(1); + cases.push(chain); + for bad in cases { + let e = implode_value(&exploded, &bad).unwrap_err(); + assert!(e.to_string().contains("unknown field"), "{e}"); + } + + // Unknown fields in a tensor manifest are rejected too. + let mut tm: serde_json::Value = + serde_json::from_str(&exploded.tensors[0].manifest_json).unwrap(); + tm["planes"][0]["extra"] = serde_json::json!(true); + let e = implode( + &exploded.registry_json, + &[tm.to_string()], + members_map(&exploded), + ) + .unwrap_err(); + assert!(e.to_string().contains("unknown field"), "{e}"); +} + +/// When no table layout reproduces the recorded hash, implode errors +/// instead of writing a container with another table. +#[test] +fn implode_rejects_unreproducible_table_hash() { + let blob = minimal_container(); + let exploded = explode(&blob).unwrap(); + let mut reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + reg["registry_hash_blake3"] = serde_json::json!("00".repeat(32)); + let e = implode_value(&exploded, ®).unwrap_err(); + assert!(e.to_string().contains("no Extension Table layout"), "{e}"); + + // The wrong built-in version does not reproduce the hash either. + let mut reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + reg["container_flags"]["builtin_version"] = serde_json::json!("0.9.0"); + let e = implode_value(&exploded, ®).unwrap_err(); + assert!(e.to_string().contains("no Extension Table layout"), "{e}"); +} + +#[test] +fn implode_writes_raw_extension_table() { + let blob = minimal_container(); + let exploded = explode(&blob).unwrap(); + let reader = ContainerReader::open(&blob).unwrap(); + let et = extension_table_bytes(&reader, &blob).unwrap(); + + // A raw table equal to the on-disk one reproduces the container. + let mut reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + reg["registry"]["extension_table_b64"] = serde_json::json!(base64::encode(et)); + assert_eq!(implode_value(&exploded, ®).unwrap(), blob); + + // Raw bytes that do not match the recorded hash are rejected. + let mut other = et.to_vec(); + other.push(0); + reg["registry"]["extension_table_b64"] = serde_json::json!(base64::encode(&other)); + let e = implode_value(&exploded, ®).unwrap_err(); + assert!(e.to_string().contains("does not match"), "{e}"); + + // So are bytes that match a forged hash but do not parse cleanly. + reg["registry_hash_blake3"] = serde_json::json!(hex::encode(ExtensionTable::hash(&other))); + let e = implode_value(&exploded, ®).unwrap_err(); + assert!(e.to_string().contains("trailing bytes"), "{e}"); +} + +/// A container written with the built-ins of another release explodes with +/// that version and implodes to the same bytes. +#[test] +fn round_trip_other_builtin_version() { + let registry = ChainRegistry { + chains: vec![(0u16, passthrough_chain())], + }; + let rec = |name: &str, inline: bool| TensorRecord { + chain_ref: if inline { CHAIN_REF_INLINE_BIT } else { 0 }, + dtype_code: 6, + input_format: 2, + flags: 0, + orig_size: 8, + name: name.to_string(), + payload_hash: None, + inline_chain: inline.then(passthrough_chain), + dependencies: Vec::new(), + tensor_metadata: None, + terminals: vec![flat_plane(PlaneRole::Value, CodecId::Identity, vec![9; 8])], + }; + let mut cursor = Cursor::new(Vec::::new()); + { + let mut writer = ContainerWriter::new(&mut cursor).unwrap(); + writer.set_builtin_version("0.9.0").unwrap(); + writer.set_chain_registry(registry).unwrap(); + let idx = writer + .codec_table_idx_for(&crate::compressor::codec_canonical_id(CodecId::Identity)) + .unwrap(); + for (name, inline) in [("a", false), ("b", true)] { + let mut r = rec(name, inline); + r.terminals[0].codec_table_idx = idx; + writer.append_tensor(&r).unwrap(); + } + writer.finalize().unwrap(); + } + let blob = cursor.into_inner(); + assert!( + blob.windows(b"@0.9.0".len()).any(|w| w == b"@0.9.0"), + "labels carry the requested version" + ); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(reader.builtin_version.as_deref(), Some("0.9.0")); + assert_eq!(reader.decode_tensor("b").unwrap(), vec![9; 8]); + + let exploded = explode(&blob).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + assert_eq!(reg["container_flags"]["builtin_version"], "0.9.0"); + assert!(reg["registry"].get("extension_table_b64").is_none()); + assert_round_trip(&blob); +} diff --git a/crates/ptwm-core/src/transforms/op.rs b/crates/ptwm-core/src/transforms/op.rs index d49dc9b..e0547d8 100644 --- a/crates/ptwm-core/src/transforms/op.rs +++ b/crates/ptwm-core/src/transforms/op.rs @@ -36,8 +36,42 @@ pub enum OpId { } impl OpId { + /// Every op, in wire-id order. + pub const ALL: [OpId; 24] = [ + OpId::Source, + OpId::Terminal, + OpId::BitReorderIeee16, + OpId::BitReorderIeee32, + OpId::BitReorderFp8E4M3, + OpId::BitReorderFp8E5M2, + OpId::ByteSplit, + OpId::NibbleSplit, + OpId::BytePassthrough, + OpId::MxFp4Deinterleave, + OpId::BlockMicroscalingRepack, + OpId::XorDelta, + OpId::FloatDelta, + OpId::IntDelta, + OpId::PredictorXor, + OpId::IndexBitwidthPack, + OpId::EntropyEstimate, + OpId::Concat, + OpId::Reshape, + OpId::MantissaZeroStrip, + OpId::BurrowsWheeler, + OpId::MoveToFront, + OpId::SphericalNormalize, + OpId::AlphaStableNormalize, + ]; + + /// Canonical id of the op's built-in contribution at the crate version. pub fn canonical_id(self) -> crate::extension::CanonicalId { - let name = match self { + crate::extension::builtin_canonical_id(self.builtin_name()) + } + + /// Name of the op's built-in contribution (`io.ptwm.builtin.`). + pub fn builtin_name(self) -> &'static str { + match self { OpId::Source => "source", OpId::Terminal => "terminal", OpId::BitReorderIeee16 => "bit_reorder_ieee16", @@ -62,8 +96,7 @@ impl OpId { OpId::MoveToFront => "move_to_front", OpId::SphericalNormalize => "spherical_normalize", OpId::AlphaStableNormalize => "alpha_stable_normalize", - }; - crate::extension::builtin_canonical_id(name) + } } pub fn as_u16(self) -> u16 { diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs index c08aa91..0ffea62 100644 --- a/crates/ptwm-core/tests/legacy_fixtures.rs +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -14,9 +14,10 @@ //! //! The tests check that every fixture decodes to the expected bytes and that //! `implode(explode(x)) == x`. The expected bytes are stored as -//! `(length, xxh64)` pairs in [`EXPECTED`]. Fixtures with built-in ids of -//! another version are checked after [`rederive_builtins`] maps the ids to -//! the crate version, until the reader accepts older built-in ids. +//! `(length, xxh64)` pairs in [`EXPECTED`]. The reader maps built-in ids +//! of other releases to the crate version +//! (`ptwm_core::extension::normalize_builtin_ids`), so fixtures with +//! v1.0.0 ids decode and transcode as they are. //! //! # Regenerating //! @@ -821,13 +822,11 @@ fn v1_1_0_fixture_has_expected_features() { } #[test] -#[ignore = "cross-version builtins: fixed in stage 3"] fn v1_0_0_fixture_decodes() { assert_decodes(V1_0_0); } #[test] -#[ignore = "cross-version builtins: fixed in stage 3"] fn v1_0_0_fixture_transcode_round_trip() { assert_transcode_round_trip(V1_0_0); } @@ -905,13 +904,11 @@ fn p4_v1_1_0_layout_writer_output_is_golden() { } #[test] -#[ignore = "cross-version builtins: fixed in stage 3"] fn v1_0_0_native_fixture_decodes() { assert_decodes(V1_0_0_NATIVE); } #[test] -#[ignore = "cross-version builtins: fixed in stage 3"] fn v1_0_0_native_fixture_transcode_round_trip() { assert_transcode_round_trip(V1_0_0_NATIVE); } @@ -1079,3 +1076,267 @@ fn regenerate_fixtures() { println!("{name}: {} B", bytes.len()); } } + +// ── Cross-version built-ins ────────────────────────────────────────────────── + +fn explode_implode(blob: &[u8]) -> (String, Vec) { + let exploded = explode(blob).expect("explode"); + let members: HashMap = exploded + .members + .iter() + .map(|(k, v)| (k.clone(), v.as_slice())) + .collect(); + let tensor_jsons: Vec = exploded + .tensors + .iter() + .map(|t| t.manifest_json.clone()) + .collect(); + let rebuilt = implode(&exploded.registry_json, &tensor_jsons, members).expect("implode"); + (exploded.registry_json, rebuilt) +} + +/// Offset of the Extension Table entry whose label is `label`, as +/// `(id_start, label_start)`. +fn find_label(blob: &[u8], label: &str) -> (usize, usize) { + let label_start = find(blob, label.as_bytes()).unwrap_or_else(|| panic!("no {label}")); + (label_start - 34, label_start) +} + +/// Recompute the header's Extension Table hash after an in-place edit. +fn fix_table_hash(blob: &mut [u8]) { + let u64_at = |b: &[u8], off: usize| u64::from_le_bytes(b[off..off + 8].try_into().unwrap()); + let et_off = u64_at(blob, 11) as usize; + let et_len = u64_at(blob, 19) as usize; + let hash = blake3::hash(&blob[et_off..et_off + et_len]); + blob[27..59].copy_from_slice(hash.as_bytes()); +} + +#[test] +fn v1_0_0_fixture_reader_reports_builtin_version() { + for fixture in [V1_0_0, V1_0_0_NATIVE] { + let blob = load(fixture); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!( + reader.builtin_version.as_deref(), + Some("1.0.0"), + "{fixture}" + ); + // The in-memory table holds current built-in ids. + assert!( + reader + .extension_table + .entries + .iter() + .all(|e| ptwm_core::extension::is_builtin(&e.canonical_id)), + "{fixture}" + ); + } + let blob = load(V1_1_0); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(reader.builtin_version.as_deref(), Some("1.1.0")); +} + +#[test] +fn v1_0_0_fixture_manifest_records_builtin_version() { + let blob = load(V1_0_0); + let (registry_json, rebuilt) = explode_implode(&blob); + assert!(rebuilt == blob); + let reg: serde_json::Value = serde_json::from_str(®istry_json).unwrap(); + let want = if env!("CARGO_PKG_VERSION") == "1.0.0" { + serde_json::Value::Null + } else { + serde_json::Value::String("1.0.0".into()) + }; + assert_eq!(reg["container_flags"]["builtin_version"], want); + // The table is rebuilt from the chains, so its bytes are not carried. + assert!(reg["registry"].get("extension_table_b64").is_none()); + + // A crate-version file carries neither field. + let (registry_json, _) = explode_implode(&main_container()); + let reg: serde_json::Value = serde_json::from_str(®istry_json).unwrap(); + assert!(reg["container_flags"].get("builtin_version").is_none()); + assert!(reg["registry"].get("extension_table_b64").is_none()); +} + +/// The exploded stores implode one tensor at a time with the archive-wide +/// registry; that must work for files of other releases too. +#[test] +fn v1_0_0_fixture_single_tensor_implode_decodes() { + let blob = load(V1_0_0); + let exploded = explode(&blob).unwrap(); + let members: HashMap = exploded + .members + .iter() + .map(|(k, v)| (k.clone(), v.as_slice())) + .collect(); + let want: HashMap<&str, (usize, u64)> = expected(V1_0_0) + .iter() + .map(|(n, l, h)| (*n, (*l, *h))) + .collect(); + for t in &exploded.tensors { + let one = implode( + &exploded.registry_json, + std::slice::from_ref(&t.manifest_json), + members.clone(), + ) + .unwrap_or_else(|e| panic!("implode {}: {e}", t.name)); + let reader = ContainerReader::open(&one).unwrap(); + assert_eq!(reader.builtin_version.as_deref(), Some("1.0.0")); + let got = reader.decode_tensor(&t.name).unwrap(); + assert_eq!( + (got.len(), xxh64(&got, 0)), + want[t.name.as_str()], + "{}", + t.name + ); + } +} + +/// A table that mixes built-in versions is not rebuilt by any table +/// builder; its bytes travel in the manifest and implode stays byte-exact. +#[test] +fn mixed_version_table_round_trips_through_raw_bytes() { + let mut blob = load(V1_1_0); + // Move the file-level `identity` codec entry (and only it) to 1.0.0. + let (id_start, label_start) = find_label(&blob, "io.ptwm.builtin.identity@1.1.0"); + let et_end = { + let u64_at = |off: usize| u64::from_le_bytes(blob[off..off + 8].try_into().unwrap()); + (u64_at(11) + u64_at(19)) as usize + }; + assert!(label_start < et_end, "entry is in the file-level table"); + let new = CanonicalId::derive(&BUILTIN_PUBKEY, "identity", "1.0.0"); + blob[id_start..id_start + 32].copy_from_slice(new.as_bytes()); + let ver = label_start + "io.ptwm.builtin.identity@".len(); + blob[ver..ver + 5].copy_from_slice(b"1.0.0"); + fix_table_hash(&mut blob); + + assert_blob_decodes("mixed", &blob, expected(V1_1_0)); + let (registry_json, rebuilt) = explode_implode(&blob); + let reg: serde_json::Value = serde_json::from_str(®istry_json).unwrap(); + assert!(reg["registry"]["extension_table_b64"].is_string()); + assert!(rebuilt == blob, "implode(explode(x)) != x"); +} + +/// An entry with a built-in label whose id does not re-derive from that +/// label is not a built-in: the reader treats it as a missing third-party +/// contribution. +#[test] +fn forged_builtin_label_stays_non_builtin() { + let mut blob = load(V1_0_0); + let (id_start, _) = find_label(&blob, "io.ptwm.builtin.huffman@1.0.0"); + blob[id_start..id_start + 32].copy_from_slice(&[0x5A; 32]); + fix_table_hash(&mut blob); + + match ContainerReader::open(&blob) { + Err(ptwm_core::error::PtwmCoreError::MissingFlavor { needs }) => { + assert_eq!(needs.len(), 1); + assert_eq!(needs[0].id, CanonicalId::from_bytes([0x5A; 32])); + } + other => panic!("expected MissingFlavor, got {other:?}"), + } + let reader = ContainerReader::open_partial(&blob).unwrap(); + assert_eq!(reader.missing_extensions.len(), 1); + assert!( + !reader.extension_table.entries.iter().any(|e| e.canonical_id + == CanonicalId::derive(&BUILTIN_PUBKEY, "huffman", env!("CARGO_PKG_VERSION"))), + "the forged entry is not mapped to the current huffman id" + ); +} + +/// Rewrite every built-in label (and id) at or after `from` to `to`. +fn rederive_labels_after(blob: &mut [u8], from: usize, to: &str) -> usize { + let prefix = b"io.ptwm.builtin."; + let mut p = from; + let mut n = 0; + while let Some(pos) = find(&blob[p..], prefix) { + let label_start = p + pos; + let label_len = u16::from_le_bytes([blob[label_start - 2], blob[label_start - 1]]) as usize; + let label = std::str::from_utf8(&blob[label_start..label_start + label_len]) + .unwrap() + .to_string(); + let (name, ver) = label[prefix.len()..].split_once('@').unwrap(); + assert_eq!(ver.len(), to.len()); + let id = CanonicalId::derive(&BUILTIN_PUBKEY, name, to); + blob[label_start - 34..label_start - 2].copy_from_slice(id.as_bytes()); + blob[label_start + label_len - to.len()..label_start + label_len] + .copy_from_slice(to.as_bytes()); + p = label_start + label_len; + n += 1; + } + n +} + +fn table_end(blob: &[u8]) -> usize { + let u64_at = |off: usize| u64::from_le_bytes(blob[off..off + 8].try_into().unwrap()); + (u64_at(11) + u64_at(19)) as usize +} + +/// An inline chain whose table version differs from the file-level table: +/// the inline version is used for implode and the file-level table travels +/// as raw bytes. +#[test] +fn inline_chain_of_another_version_round_trips() { + let mut blob = load(V1_1_0); + let end = table_end(&blob); + assert!( + rederive_labels_after(&mut blob, end, "1.0.0") > 0, + "fixture has an inline chain" + ); + assert_blob_decodes("inline-1.0.0", &blob, expected(V1_1_0)); + let (registry_json, rebuilt) = explode_implode(&blob); + let reg: serde_json::Value = serde_json::from_str(®istry_json).unwrap(); + assert_eq!(reg["container_flags"]["builtin_version"], "1.0.0"); + assert!(reg["registry"]["extension_table_b64"].is_string()); + assert!(rebuilt == blob, "implode(explode(x)) != x"); +} + +/// Inline chains of two different built-in versions cannot be rebuilt with +/// one version, so explode refuses the container instead of producing +/// different bytes. +#[test] +fn explode_rejects_inline_chains_of_two_versions() { + let blob = load(V1_1_0); + let reader = ContainerReader::open(&blob).unwrap(); + let mut records = records_in_file_order(&reader); + // Store a second registry chain inline, on the last record. + let last = records.last_mut().unwrap(); + assert!(last.inline_chain.is_none()); + last.inline_chain = Some( + reader + .chain_registry + .lookup(last.chain_ref) + .unwrap() + .clone(), + ); + last.chain_ref = CHAIN_REF_INLINE_BIT; + let mut cursor = Cursor::new(Vec::new()); + { + let mut w = ContainerWriter::new(&mut cursor).unwrap(); + w.set_prelude(&reader.prelude).unwrap(); + w.set_chain_registry(reader.chain_registry.clone()).unwrap(); + for r in &records { + w.append_tensor(r).unwrap(); + } + w.finalize().unwrap(); + } + let mut blob = cursor.into_inner(); + // The last record's inline table is the last run of built-in labels. + let last_off = ContainerReader::open(&blob) + .unwrap() + .index + .iter() + .map(|e| e.tensor_offset) + .max() + .unwrap() as usize; + let from = env!("CARGO_PKG_VERSION"); + let to = if from == "1.0.0" { "1.0.1" } else { "1.0.0" }; + if to.len() != from.len() { + return; // `rederive_labels_after` needs equal-length versions. + } + assert!(rederive_labels_after(&mut blob, last_off, to) > 0); + let e = explode(&blob).err().expect("explode must fail"); + assert!( + e.to_string().contains("inline chain built-in version"), + "{e}" + ); +} diff --git a/crates/ptwm-core/tests/trust_fast_path.rs b/crates/ptwm-core/tests/trust_fast_path.rs new file mode 100644 index 0000000..aecf455 --- /dev/null +++ b/crates/ptwm-core/tests/trust_fast_path.rs @@ -0,0 +1,96 @@ +//! Opening a container whose Extension Table holds only built-ins must not +//! touch the filesystem: no extension discovery, no keyring load and no +//! first-run trust lock (synthesis E9). +//! +//! The test points `HOME` and the XDG directories at a temporary directory, +//! so it runs in its own test binary: environment changes are process-wide. + +use std::io::Cursor; +use std::path::{Path, PathBuf}; + +use ptwm_core::container::ContainerReader; +use ptwm_core::container::ContainerWriter; +use ptwm_core::extension::{CanonicalId, ExtensionTable, Kind, builtin_entry}; + +fn fixture(name: &str) -> Vec { + let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures") + .join(name); + std::fs::read(&path).unwrap_or_else(|e| panic!("read {}: {e}", path.display())) +} + +/// Every file below `dir`, recursively. +fn files_under(dir: &Path) -> Vec { + let mut out = Vec::new(); + let mut stack = vec![dir.to_path_buf()]; + while let Some(d) = stack.pop() { + for entry in std::fs::read_dir(&d).unwrap() { + let path = entry.unwrap().path(); + if path.is_dir() { + stack.push(path); + } else { + out.push(path); + } + } + } + out +} + +/// A container whose table holds one third-party codec entry. +fn third_party_container() -> Vec { + let mut entry = builtin_entry("identity", Kind::PlaneCodec, true); + entry.canonical_id = CanonicalId::derive(&[0x42; 32], "demo", "0.1.0"); + entry.human_label = "com.example.demo@0.1.0".into(); + let mut cursor = Cursor::new(Vec::new()); + { + let mut w = ContainerWriter::new(&mut cursor).unwrap(); + w.set_extension_table(ExtensionTable { + entries: vec![entry], + }); + w.finalize().unwrap(); + } + cursor.into_inner() +} + +#[test] +fn builtin_only_files_open_without_filesystem_access() { + let tmp = tempfile::tempdir().unwrap(); + let root = tmp.path(); + // SAFETY: this binary has a single test, so no other thread reads the + // environment while it is changed. + unsafe { + std::env::set_var("HOME", root.join("home")); + std::env::set_var("XDG_CONFIG_HOME", root.join("config")); + std::env::set_var("XDG_DATA_HOME", root.join("data")); + std::env::set_var("XDG_CACHE_HOME", root.join("cache")); + std::env::remove_var("PTWM_EXTENSION_PATH"); + } + + for name in [ + "legacy_v1_1_0.ptwm", + "legacy_v1_0_0.ptwm", + "legacy_v1_0_0_native.ptwm", + "legacy_trimmed.ptwm", + ] { + let blob = fixture(name); + let reader = ContainerReader::open(&blob).unwrap_or_else(|e| panic!("{name}: {e}")); + reader.decode_model().unwrap(); + assert!(reader.missing_extensions.is_empty()); + } + assert_eq!( + files_under(root), + Vec::::new(), + "opening built-in-only files created files" + ); + + // Control: a table with a third-party entry runs the discovery and + // trust pass, which records the first-run trust lock. This shows the + // assertion above would catch the pass running. + let blob = third_party_container(); + let reader = ContainerReader::open_partial(&blob).unwrap(); + assert_eq!(reader.missing_extensions.len(), 1); + assert!( + !files_under(root).is_empty(), + "the trust pass was expected to write its lock file" + ); +} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index d8c1580..88c1d25 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -936,8 +936,10 @@ fn py_chain_blob_info(blob: &[u8]) -> PyResult<(usize, Vec<(u16, u16)>)> { // Parse the extension table; its size is the parse position, not the // length of a re-serialization. - let (inline_table, et_size) = ExtensionTable::from_bytes_prefix(blob) + let (mut inline_table, et_size) = ExtensionTable::from_bytes_prefix(blob) .map_err(|e| pyo3::exceptions::PyValueError::new_err(format!("chain blob et: {e}")))?; + // Blobs cached by another release carry that release's built-in ids. + ptwm_core::extension::normalize_builtin_ids(&mut inline_table); // Build (table_idx, op_id) pairs for every entry that resolves to an op. let mut table: Vec<(u16, u16)> = Vec::with_capacity(inline_table.entries.len()); diff --git a/tests/ptwm/wire_format/test_cross_version_builtins.py b/tests/ptwm/wire_format/test_cross_version_builtins.py new file mode 100644 index 0000000..0aff70f --- /dev/null +++ b/tests/ptwm/wire_format/test_cross_version_builtins.py @@ -0,0 +1,145 @@ +"""Built-in Extension Table ids of other releases. + +Built-in canonical ids are derived from the release version, so files and +chain blobs written by another release carry other ids. The reader maps them +to the current ids when the label and id re-derive consistently. +""" + +from __future__ import annotations + +import json +import re +import struct +from pathlib import Path + +import pytest +from ptwm import _rust + +FIXTURES = Path(__file__).resolve().parents[3] / "crates/ptwm-core/tests/fixtures" +BUILTIN_PUBKEY = b"\xbb" * 32 +_LABEL = re.compile(rb"io\.ptwm\.builtin\.([a-z0-9_]+)@([0-9.]+)") + + +def _derive(name: str, version: str) -> bytes: + blake3 = pytest.importorskip("blake3") + return blake3.blake3( + BUILTIN_PUBKEY + b"\x00" + name.encode() + b"\x00" + version.encode() + ).digest() + + +def _rederive(blob: bytes, to: str) -> bytes: + """Rewrite every built-in entry's id and label to version ``to`` (which + must have the length of the current version). No hash is fixed up, so + use it only on blobs without a header hash (chain blobs).""" + out = bytearray(blob) + count = 0 + for m in _LABEL.finditer(blob): + name, ver = m.group(1).decode(), m.group(2).decode() + assert len(ver) == len(to) + id_start = m.start() - 34 + assert bytes(out[id_start : id_start + 32]) == _derive(name, ver) + out[id_start : id_start + 32] = _derive(name, to) + out[m.start(2) : m.end(2)] = to.encode() + count += 1 + assert count > 0 + return bytes(out) + + +def _fp32_chain_blob(shape: list[int]) -> bytes: + from ptwm.preprocessing import _chains + + return _chains.CHAIN_FP32_SPLIT(shape).to_bytes() + + +def _other_version() -> str: + import ptwm + + return "1.0.0" if ptwm.__version__ != "1.0.0" else "1.0.1" + + +def test_chain_blob_info_maps_old_builtin_ids() -> None: + from ptwm._core import chain_blob_info + + blob = _fp32_chain_blob([16]) + old = _rederive(blob, _other_version()) + assert old != blob + assert chain_blob_info(old) == chain_blob_info(blob) + + +def test_compress_model_accepts_chain_blobs_of_other_releases() -> None: + """A chain blob cached by another release compresses to the same + container as the current blob.""" + import torch + + shape = [8, 16] + raw = torch.linspace(-1, 1, 128, dtype=torch.float32).numpy().tobytes() + blob = _fp32_chain_blob(shape) + old = _rederive(blob, _other_version()) + + def compress(chain: bytes) -> bytes: + return _rust.compress_model( + records=[("w", 3, 0, raw, shape, "float32", None)], + chains_per_tensor=[[chain]], + emit_payload_hash=True, + emit_plane_crc=True, + ) + + current = compress(blob) + assert compress(old) == current + assert _rust.decode_tensor(current, "w") == raw + + +def test_forged_chain_blob_entry_is_rejected() -> None: + """A built-in label whose id does not re-derive is not a built-in.""" + from ptwm._core import chain_blob_info + + old = bytearray(_rederive(_fp32_chain_blob([16]), _other_version())) + first = _LABEL.search(old) + assert first is not None + old[first.start() - 34 : first.start() - 2] = b"\x5a" * 32 + _, table = chain_blob_info(bytes(old)) + _, full = chain_blob_info(_fp32_chain_blob([16])) + assert len(table) == len(full) - 1 + + +def test_v1_0_0_fixture_decodes_and_transcodes() -> None: + old = (FIXTURES / "legacy_v1_0_0.ptwm").read_bytes() + cur = (FIXTURES / "legacy_v1_1_0.ptwm").read_bytes() + assert dict(_rust.decode_model(old)) == dict(_rust.decode_model(cur)) + + registry_json, tensors, members = _rust.explode_ptwm(old) + rebuilt = _rust.implode_ptwm(registry_json, [t[2] for t in tensors], members) + assert rebuilt == old + reg = json.loads(registry_json) + import ptwm + + if ptwm.__version__ != "1.0.0": + assert reg["container_flags"]["builtin_version"] == "1.0.0" + assert "extension_table_b64" not in reg["registry"] + + # One tensor at a time, as the exploded stores do. + name, _key, tjson = tensors[0] + one = _rust.implode_ptwm(registry_json, [tjson], members) + assert _rust.decode_tensor(one, name) == _rust.decode_tensor(old, name) + + +def test_implode_rejects_unknown_manifest_fields() -> None: + blob = (FIXTURES / "legacy_trimmed.ptwm").read_bytes() + registry_json, tensors, members = _rust.explode_ptwm(blob) + reg = json.loads(registry_json) + reg["container_flags"]["layout"] = "compact" + with pytest.raises(Exception, match="unknown field"): + _rust.implode_ptwm(json.dumps(reg), [t[2] for t in tensors], members) + reg = json.loads(registry_json) + reg["manifest_version"] = 2 + with pytest.raises(Exception, match="unsupported manifest"): + _rust.implode_ptwm(json.dumps(reg), [t[2] for t in tensors], members) + + +def test_label_layout_assumption() -> None: + """``_rederive`` assumes ``canonical_id[32] | label_len u16 | label``.""" + blob = _fp32_chain_blob([16]) + m = _LABEL.search(blob) + assert m is not None + (label_len,) = struct.unpack_from(" Date: Tue, 27 Oct 2026 22:17:43 +0000 Subject: [PATCH 07/34] fix(S3-cross-version): address review findings - explode re-encodes every inline-chain record with the built-in version implode will use and refuses the container when the bytes differ. An inline Extension Table that mixes built-in versions was accepted before and imploded to other bytes. - One built-in version rule, extension::is_valid_builtin_version (non-empty, at most 64 bytes, no '@'), now applies in normalize_builtin_ids, builtin_entry_at_version, ExtensionTableBuilder::with_builtin_version (now fallible, new ExtensionError::InvalidBuiltinVersion), build_extension_table and ContainerWriter::set_builtin_version. Entries with an invalid version stay non-built-in, so explode never records a version implode rejects. - ptwm-py list_extension_table_entries normalises built-in ids, so ptwm.ablation compare/attribute match built-ins across releases. Tests: explode_rejects_inline_table_with_mixed_versions, invalid_builtin_versions_are_not_recorded, invalid_versions_are_not_normalised, with_builtin_version_rejects_invalid_versions, a build_extension_table check in set_builtin_version_must_come_first, and the Python test_list_extension_table_entries_reports_current_builtin_ids. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/wire.rs | 2 +- crates/ptwm-core/src/container.rs | 12 +- crates/ptwm-core/src/extension/builder.rs | 36 +++-- crates/ptwm-core/src/extension/builtins.rs | 55 ++++++- crates/ptwm-core/src/extension/error.rs | 3 + crates/ptwm-core/src/extension/mod.rs | 3 +- crates/ptwm-core/src/tensor_record.rs | 2 +- crates/ptwm-core/src/transcode/mod.rs | 23 ++- crates/ptwm-core/src/transcode/tests.rs | 149 +++++++++++++++++- crates/ptwm-py/src/inspect.rs | 9 +- .../test_cross_version_builtins.py | 23 +++ 11 files changed, 291 insertions(+), 26 deletions(-) diff --git a/crates/ptwm-core/src/chain/wire.rs b/crates/ptwm-core/src/chain/wire.rs index c954a79..583f041 100644 --- a/crates/ptwm-core/src/chain/wire.rs +++ b/crates/ptwm-core/src/chain/wire.rs @@ -499,7 +499,7 @@ mod tests { }], terminals: vec![], }; - let mut builder = ExtensionTableBuilder::with_builtin_version("1.0.0"); + let mut builder = ExtensionTableBuilder::with_builtin_version("1.0.0").unwrap(); let mut chain_bytes = Vec::new(); write_chain(&chain, &mut builder, &mut chain_bytes).unwrap(); let table = ExtensionTable { diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index a4f2de5..9d473a9 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -158,7 +158,7 @@ pub fn build_extension_table( ) -> Result { use crate::extension::{Kind, builtin_entries, builtin_registry}; - let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version); + let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version)?; let mut registry_bytes = Vec::new(); registry.write(&mut builder, &mut registry_bytes)?; @@ -454,7 +454,7 @@ impl ContainerWriter { "builtin version must be set before the extension table and tensors".into(), )); } - if version.is_empty() || version.len() > 64 || version.contains('@') { + if !crate::extension::is_valid_builtin_version(version) { return Err(PtwmCoreError::InvalidContainer(format!( "invalid builtin version {version:?}" ))); @@ -1859,6 +1859,14 @@ mod tests { let mut w = ContainerWriter::new(&mut buf).unwrap(); for bad in ["", "1@2", &"9".repeat(65)] { assert!(w.set_builtin_version(bad).is_err(), "{bad:?}"); + // The table builder applies the same rule. + let registry = ChainRegistry { + chains: vec![(0, make_minimal_chain(1))], + }; + assert!( + build_extension_table(®istry, CodecSet::AllBuiltins, bad).is_err(), + "{bad:?}" + ); } w.set_builtin_version("1.0.0").unwrap(); w.set_chain_registry(ChainRegistry { diff --git a/crates/ptwm-core/src/extension/builder.rs b/crates/ptwm-core/src/extension/builder.rs index 3f5c7ad..7ac7e77 100644 --- a/crates/ptwm-core/src/extension/builder.rs +++ b/crates/ptwm-core/src/extension/builder.rs @@ -3,8 +3,8 @@ use std::collections::HashMap; use super::{ - CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionTableEntry, builtin_entry_at_version, - builtin_registry, + CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionError, ExtensionTableEntry, + builtin_entry_at_version, builtin_registry, is_valid_builtin_version, }; /// Intern `ExtensionTableEntry` values by canonical id, returning a @@ -18,7 +18,11 @@ pub struct ExtensionTableBuilder { impl Default for ExtensionTableBuilder { fn default() -> Self { - Self::with_builtin_version(CURRENT_BUILTIN_VERSION) + Self { + by_id: HashMap::new(), + entries: Vec::new(), + builtin_version: CURRENT_BUILTIN_VERSION.to_owned(), + } } } @@ -30,12 +34,16 @@ impl ExtensionTableBuilder { /// A builder that interns built-ins with the ids and labels of /// `version`, as a writer of that release does. The transcoder uses it /// to rebuild tables written by other releases byte-exactly. - pub fn with_builtin_version(version: &str) -> Self { - Self { - by_id: HashMap::new(), - entries: Vec::new(), - builtin_version: version.to_owned(), + /// + /// Fails when `version` does not pass [`is_valid_builtin_version`]. + pub fn with_builtin_version(version: &str) -> Result { + if !is_valid_builtin_version(version) { + return Err(ExtensionError::InvalidBuiltinVersion(version.to_owned())); } + Ok(Self { + builtin_version: version.to_owned(), + ..Self::default() + }) } /// A builder whose table starts as `entries`. Interning an id already in @@ -138,7 +146,7 @@ mod tests { #[test] fn intern_builtin_uses_the_builder_version() { - let mut b = ExtensionTableBuilder::with_builtin_version("1.0.0"); + let mut b = ExtensionTableBuilder::with_builtin_version("1.0.0").unwrap(); assert_eq!(b.builtin_version(), "1.0.0"); assert_eq!(b.intern_builtin("source"), Some(0)); assert_eq!(b.intern_builtin("zstd"), Some(1)); @@ -158,6 +166,16 @@ mod tests { assert_eq!(b.finish(), vec![make_entry("source")]); } + #[test] + fn with_builtin_version_rejects_invalid_versions() { + for bad in ["", "9.9@x", &"7".repeat(65)] { + assert!( + ExtensionTableBuilder::with_builtin_version(bad).is_err(), + "{bad:?}" + ); + } + } + #[test] fn from_entries_returns_existing_indices() { let entries = vec![ diff --git a/crates/ptwm-core/src/extension/builtins.rs b/crates/ptwm-core/src/extension/builtins.rs index 10d5b97..2629307 100644 --- a/crates/ptwm-core/src/extension/builtins.rs +++ b/crates/ptwm-core/src/extension/builtins.rs @@ -217,8 +217,12 @@ pub fn builtin_entries() -> Vec { /// The Extension Table entry of the built-in `name` with its id and label /// derived for `version`, as a writer of that version emits it. Returns -/// `None` if `name` has no built-in entry. +/// `None` if `name` has no built-in entry or `version` fails +/// [`is_valid_builtin_version`]. pub fn builtin_entry_at_version(name: &str, version: &str) -> Option { + if !is_valid_builtin_version(version) { + return None; + } let reg = builtin_registry(); let idx = reg.by_id.get(reg.by_name.get(name)?)?.entry_index?; if version == CURRENT_BUILTIN_VERSION { @@ -244,11 +248,23 @@ pub fn is_builtin(id: &CanonicalId) -> bool { .is_some_and(|r| r.entry_index.is_some()) } +/// Whether `version` is a valid built-in version: non-empty, at most 64 +/// bytes, and free of `@`. Every writer and +/// [`normalize_builtin_ids`] apply this rule, so a version the reader accepts +/// is one a writer can reproduce. +pub fn is_valid_builtin_version(version: &str) -> bool { + !version.is_empty() && version.len() <= MAX_BUILTIN_VERSION_LEN && !version.contains('@') +} + +/// Longest built-in version string, in bytes. +pub const MAX_BUILTIN_VERSION_LEN: usize = 64; + /// Map built-in entries written by another release to the ids and labels of /// the crate version. /// /// An entry is changed only when all of these hold: -/// - its label is `io.ptwm.builtin.@`; +/// - its label is `io.ptwm.builtin.@`, and `` passes +/// [`is_valid_builtin_version`]; /// - `` has a built-in entry ([`builtin_entries`]); /// - its id equals `CanonicalId::derive(BUILTIN_PUBKEY, name, ver)`. /// @@ -273,6 +289,9 @@ pub fn normalize_builtin_ids(table: &mut ExtensionTable) -> Option { let Some((name, ver)) = rest.split_once('@') else { continue; }; + if !is_valid_builtin_version(ver) { + continue; + } let Some(¤t_id) = reg.by_name.get(name) else { continue; }; @@ -402,6 +421,38 @@ mod tests { assert_eq!(version.as_deref(), Some("1.0.0")); } + /// Versions that no writer accepts are not normalised, so a reader never + /// reports a version that the transcoder cannot write back. + #[test] + fn invalid_versions_are_not_normalised() { + let long = "7".repeat(MAX_BUILTIN_VERSION_LEN + 1); + for bad in ["", "9.9@x", long.as_str()] { + assert!(!is_valid_builtin_version(bad), "{bad:?}"); + assert!( + builtin_entry_at_version("huffman", bad).is_none(), + "{bad:?}" + ); + let mut entry = builtin_entry_at_version("huffman", "1.0.0").unwrap(); + entry.canonical_id = CanonicalId::derive(&BUILTIN_PUBKEY, "huffman", bad); + entry.human_label = format!("{BUILTIN_LABEL_PREFIX}huffman@{bad}"); + let mut table = ExtensionTable { + entries: vec![entry], + }; + let before = table.clone(); + assert_eq!(normalize_builtin_ids(&mut table), None, "{bad:?}"); + assert_eq!(table, before); + assert!(!is_builtin(&table.entries[0].canonical_id)); + } + let max = "7".repeat(MAX_BUILTIN_VERSION_LEN); + for good in ["1.0.0", "0.9.0-rc.1", max.as_str()] { + assert!(is_valid_builtin_version(good), "{good:?}"); + assert_eq!( + normalize_builtin_ids(&mut table_at(good)).as_deref(), + Some(good) + ); + } + } + #[test] fn normalisation_ignores_third_party_and_unknown_names() { let id = CanonicalId::derive(&[0x22; 32], "huffman", "1.0.0"); diff --git a/crates/ptwm-core/src/extension/error.rs b/crates/ptwm-core/src/extension/error.rs index 93f8a64..799ef8b 100644 --- a/crates/ptwm-core/src/extension/error.rs +++ b/crates/ptwm-core/src/extension/error.rs @@ -33,4 +33,7 @@ pub enum ExtensionError { #[error("attestation parse error: {0}")] AttestationParse(String), + + #[error("invalid builtin version {0:?}")] + InvalidBuiltinVersion(String), } diff --git a/crates/ptwm-core/src/extension/mod.rs b/crates/ptwm-core/src/extension/mod.rs index 76f9f19..ff2065b 100644 --- a/crates/ptwm-core/src/extension/mod.rs +++ b/crates/ptwm-core/src/extension/mod.rs @@ -17,7 +17,8 @@ pub mod table; pub use builder::ExtensionTableBuilder; pub use builtins::{ BUILTIN_LABEL_PREFIX, BuiltinRef, BuiltinRegistry, CURRENT_BUILTIN_VERSION, builtin_entries, - builtin_entry, builtin_entry_at_version, builtin_registry, is_builtin, normalize_builtin_ids, + builtin_entry, builtin_entry_at_version, builtin_registry, is_builtin, + is_valid_builtin_version, normalize_builtin_ids, }; pub use capability::{Capability, CapabilityMap}; pub use dispatch::{BuiltinKind, dispatch_builtin}; diff --git a/crates/ptwm-core/src/tensor_record.rs b/crates/ptwm-core/src/tensor_record.rs index 41354da..1ac5658 100644 --- a/crates/ptwm-core/src/tensor_record.rs +++ b/crates/ptwm-core/src/tensor_record.rs @@ -195,7 +195,7 @@ fn write_tensor_record_unchecked( // inline_et_entries: [ExtensionTableEntry; count] // chain_bytes: (num_nodes, ..., nodes_section, edges_section, terminals_section) if let Some(chain) = &rec.inline_chain { - let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version); + let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version)?; let mut chain_bytes = Vec::new(); write_chain(chain, &mut builder, &mut chain_bytes)?; let inline_table = ExtensionTable { diff --git a/crates/ptwm-core/src/transcode/mod.rs b/crates/ptwm-core/src/transcode/mod.rs index e3b6659..9507d4d 100644 --- a/crates/ptwm-core/src/transcode/mod.rs +++ b/crates/ptwm-core/src/transcode/mod.rs @@ -36,7 +36,7 @@ use crate::plane_record::{ChunkEntry, ExternalStateRef, PlaneRecord}; use crate::prelude::{PreludeEntry, write_prelude}; use crate::tensor_record::{ CHAIN_REF_INLINE_BIT, Dependency, RefKind, TensorRecord, - parse_tensor_record_with_builtin_version, + parse_tensor_record_with_builtin_version, write_tensor_record_with_builtin_version, }; use crate::transforms::op::OpId; use crate::types::PlaneRole; @@ -364,7 +364,7 @@ pub fn explode(blob: &[u8]) -> Result { .checked_add(entry.tensor_len as usize) .filter(|e| *e <= blob.len()) .ok_or_else(|| err("tensor record slice out of bounds"))?; - let (record, _, inline_version) = + let (record, consumed, inline_version) = parse_tensor_record_with_builtin_version(&blob[start..end])?; // `implode` writes every inline-chain table with one built-in // version, so they must agree. @@ -382,6 +382,25 @@ pub fn explode(blob: &[u8]) -> Result { Some(_) => {} } } + // `implode` re-encodes inline-chain tables with a single built-in + // version. A table that mixes versions (or is otherwise not what + // that writer emits) would come back with different bytes, so + // refuse it here instead of producing a silent mismatch. + if record.inline_chain.is_some() { + let version = inline_builtin_version + .as_deref() + .unwrap_or(CURRENT_BUILTIN_VERSION); + let mut reencoded = Vec::new(); + write_tensor_record_with_builtin_version(&record, &mut reencoded, version)?; + if blob.get(start..start + consumed) != Some(reencoded.as_slice()) { + return Err(err(format!( + "tensor {:?}: inline chain record differs from its re-encoding \ + (for example, its Extension Table mixes built-in versions); \ + the container cannot be transcoded byte-exactly", + record.name + ))); + } + } for plane in &record.terminals { codec_ids.insert(plane.codec_id.as_u16()); } diff --git a/crates/ptwm-core/src/transcode/tests.rs b/crates/ptwm-core/src/transcode/tests.rs index ef5de2b..3a24661 100644 --- a/crates/ptwm-core/src/transcode/tests.rs +++ b/crates/ptwm-core/src/transcode/tests.rs @@ -501,10 +501,9 @@ fn implode_writes_raw_extension_table() { assert!(e.to_string().contains("trailing bytes"), "{e}"); } -/// A container written with the built-ins of another release explodes with -/// that version and implodes to the same bytes. -#[test] -fn round_trip_other_builtin_version() { +/// A container with built-ins of `version`: tensor "a" uses the chain +/// registry, tensor "b" an inline chain. +fn container_at_builtin_version(version: &str) -> Vec { let registry = ChainRegistry { chains: vec![(0u16, passthrough_chain())], }; @@ -524,7 +523,7 @@ fn round_trip_other_builtin_version() { let mut cursor = Cursor::new(Vec::::new()); { let mut writer = ContainerWriter::new(&mut cursor).unwrap(); - writer.set_builtin_version("0.9.0").unwrap(); + writer.set_builtin_version(version).unwrap(); writer.set_chain_registry(registry).unwrap(); let idx = writer .codec_table_idx_for(&crate::compressor::codec_canonical_id(CodecId::Identity)) @@ -536,7 +535,14 @@ fn round_trip_other_builtin_version() { } writer.finalize().unwrap(); } - let blob = cursor.into_inner(); + cursor.into_inner() +} + +/// A container written with the built-ins of another release explodes with +/// that version and implodes to the same bytes. +#[test] +fn round_trip_other_builtin_version() { + let blob = container_at_builtin_version("0.9.0"); assert!( blob.windows(b"@0.9.0".len()).any(|w| w == b"@0.9.0"), "labels carry the requested version" @@ -551,3 +557,134 @@ fn round_trip_other_builtin_version() { assert!(reg["registry"].get("extension_table_b64").is_none()); assert_round_trip(&blob); } + +fn find_last(haystack: &[u8], needle: &[u8]) -> usize { + haystack + .windows(needle.len()) + .rposition(|w| w == needle) + .expect("needle present") +} + +/// Rewrite the built-in entry labelled `old_label` (its last occurrence) to +/// `@` with the id derived for that version. Both +/// versions must have the same length. An entry's id directly precedes its +/// u16 label length. +fn relabel_builtin(blob: &mut [u8], old_label: &str, new_version: &str) { + let (name, old_version) = old_label + .strip_prefix(crate::extension::BUILTIN_LABEL_PREFIX) + .unwrap() + .split_once('@') + .unwrap(); + assert_eq!(old_version.len(), new_version.len()); + let pos = find_last(blob, old_label.as_bytes()); + let id = CanonicalId::derive(&crate::extension::BUILTIN_PUBKEY, name, new_version); + blob[pos - 34..pos - 2].copy_from_slice(id.as_bytes()); + let ver_pos = pos + old_label.len() - old_version.len(); + blob[ver_pos..pos + old_label.len()].copy_from_slice(new_version.as_bytes()); +} + +/// One inline-chain table that mixes built-in versions cannot be rebuilt by +/// implode, which writes each inline table with a single version. Explode +/// must refuse it rather than produce a manifest that implodes to other +/// bytes. +#[test] +fn explode_rejects_inline_table_with_mixed_versions() { + let mut blob = container_at_builtin_version("0.9.0"); + // The last `source@0.9.0` label is the one in the inline table of "b". + relabel_builtin(&mut blob, "io.ptwm.builtin.source@0.9.0", "0.8.0"); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(reader.decode_tensor("b").unwrap(), vec![9; 8]); + let Err(e) = explode(&blob) else { + panic!("explode accepted an inline table with mixed versions"); + }; + let e = e.to_string(); + assert!(e.contains("differs from its re-encoding"), "{e}"); + + // The control: relabelling every entry of that table to one version + // round-trips. + let mut blob = container_at_builtin_version("0.9.0"); + relabel_builtin(&mut blob, "io.ptwm.builtin.source@0.9.0", "0.8.0"); + relabel_builtin(&mut blob, "io.ptwm.builtin.byte_passthrough@0.9.0", "0.8.0"); + let exploded = explode(&blob).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + assert_eq!(reg["container_flags"]["builtin_version"], "0.8.0"); + assert_round_trip(&blob); +} + +/// Built-in labels with a version that no writer accepts are not +/// normalised: the writer refuses such a table, and explode never records +/// the version in a manifest that implode would then reject. +#[test] +fn invalid_builtin_versions_are_not_recorded() { + // Writer side: a raw table whose built-ins carry an invalid version + // does not cover the chain ops, so it is refused. + let long = "7".repeat(70); + for bad in ["9.9@x", long.as_str(), ""] { + let registry = ChainRegistry { + chains: vec![(0u16, passthrough_chain())], + }; + let mut built = + crate::container::build_extension_table(®istry, CodecSet::AllBuiltins, "1.0.0") + .unwrap(); + for e in &mut built.table.entries { + let name = e + .human_label + .strip_prefix(crate::extension::BUILTIN_LABEL_PREFIX) + .unwrap() + .split_once('@') + .unwrap() + .0 + .to_owned(); + e.canonical_id = CanonicalId::derive(&crate::extension::BUILTIN_PUBKEY, &name, bad); + e.human_label = format!("{}{name}@{bad}", crate::extension::BUILTIN_LABEL_PREFIX); + } + let et = built.table.to_bytes().unwrap(); + let mut writer = ContainerWriter::new(Cursor::new(Vec::::new())).unwrap(); + let e = writer.set_raw_extension_table(&et, registry).unwrap_err(); + assert!( + e.to_string().contains("missing from the table"), + "{bad:?}: {e}" + ); + } + + // Reader side: patch every built-in entry of a 1.0.0 file to the + // same-length invalid version "9.9@x" and fix the table hash. + let mut blob = container_at_builtin_version("1.0.0"); + let names: Vec = ContainerReader::open(&blob) + .unwrap() + .extension_table + .entries + .iter() + .map(|e| { + e.human_label + .strip_prefix(crate::extension::BUILTIN_LABEL_PREFIX) + .unwrap() + .split_once('@') + .unwrap() + .0 + .to_owned() + }) + .collect(); + for name in &names { + let label = format!("{}{name}@1.0.0", crate::extension::BUILTIN_LABEL_PREFIX); + while blob.windows(label.len()).any(|w| w == label.as_bytes()) { + relabel_builtin(&mut blob, &label, "9.9@x"); + } + } + let mut header = crate::header::Header::from_bytes(&blob).unwrap(); + let start = header.extension_table_offset as usize; + let end = start + header.extension_table_length as usize; + header.extension_table_hash = ExtensionTable::hash(&blob[start..end]); + let header_bytes = header.to_bytes().unwrap(); + blob[..header_bytes.len()].copy_from_slice(&header_bytes); + + // The entries are not built-ins, so the chain ops do not resolve and + // the file is rejected on read instead of exploding to a manifest with + // builtin_version "9.9@x". + assert!(ContainerReader::open(&blob).is_err()); + let Err(e) = explode(&blob) else { + panic!("explode accepted built-ins with an invalid version"); + }; + let msg = e.to_string(); + assert!(msg.contains("unknown canonical id"), "{msg}"); +} diff --git a/crates/ptwm-py/src/inspect.rs b/crates/ptwm-py/src/inspect.rs index cc67933..3051de7 100644 --- a/crates/ptwm-py/src/inspect.rs +++ b/crates/ptwm-py/src/inspect.rs @@ -5,7 +5,7 @@ use std::io::Read; use pyo3::prelude::*; -use ptwm_core::extension::ExtensionTable; +use ptwm_core::extension::{ExtensionTable, normalize_builtin_ids}; use ptwm_core::header::{HEADER_LEN, Header}; fn to_pyerr(e: E) -> PyErr { @@ -26,6 +26,10 @@ fn extension_table_region(buf: &[u8], off: u64, len: u64) -> Option<&[u8]> { /// tuples -- one per entry in the file's extension table. Returns an empty /// list when the file carries no extension table (``extension_table_length == 0``). /// +/// Built-in entries written by another release are reported with the ids and +/// labels of the current release (see ``normalize_builtin_ids``), so the ids +/// of files written by different releases compare equal. +/// /// Raises ``ValueError`` for any parse error (bad magic, truncated header, /// out-of-bounds extension table region, ...). #[pyfunction] @@ -58,7 +62,8 @@ pub fn list_extension_table_entries(path: String) -> PyResult None: assert m is not None (label_len,) = struct.unpack_from(" None: + """Built-in entries of an older file are listed with the current ids, so + ``ptwm.ablation.bench_compare`` sees the same contributions in both files.""" + from ptwm._core import list_extension_table_entries + from ptwm.ablation import bench_compare + + old = FIXTURES / "legacy_v1_0_0.ptwm" + cur = FIXTURES / "legacy_v1_1_0.ptwm" + old_entries = list_extension_table_entries(str(old)) + cur_entries = list_extension_table_entries(str(cur)) + assert old_entries + assert {cid for cid, *_ in old_entries} <= {cid for cid, *_ in cur_entries} + + import ptwm + + for _cid, label, _kind, _flavor in old_entries: + assert label.endswith(f"@{ptwm.__version__}"), label + + result = bench_compare(old, cur) + assert result.only_in_a == [] + assert result.common From ce3e9449e051f7163ec0d3c7dbb35ee44e7f0cd5 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 22:40:32 +0000 Subject: [PATCH 08/34] fix(core): size chains symbolically and enforce exact Source sizes Stage 4 of the issue #5 work (synthesis section 3, P4 and P5). P5, decode side: - chain::runtime::propagate_terminal_descriptors(_bounded) walks propagate_descriptors in topological order with no plane buffers. It mirrors forward_chain: inputs sorted by dst_input_idx, edge role overrides, terminal role re-tagging (fused kernels emit the same descriptors). The bounded variant rejects the Source, and any node output, above max_plane_len(orig_size) = 16 x orig_size + 4096. chain_terminal_descriptors now delegates to it; chain_source_descriptor is shared by all callers. - Source::output_descriptor uses checked arithmetic (no overflow panic or wrap on 255 x u32::MAX dims). - validate_chain looks edge sources up with .get() (no index panic). - Legacy container decode: Source <= 8 x orig_size (P4 files stay readable), chunk-table decoded sizes and their sum bounded before decode, decoded planes bounded, inverse output must equal orig_size. - Compact frame decode: Source must equal orig_size; each plane must decode to its descriptor length. P4, encode side: - Python build_chain_blobs/source_shape_for in preprocessing/_chains.py give each builder a Source shape in units of its own Source dtype that covers exactly the raw byte count, skipping builders that cannot, and fall back to a byte passthrough over all bytes. Used by core/_compressor.py, sharding/_writer.py and integrations/_compress_safetensors.py. - Compressor reads dtype and shape before delta encoding, fixing the AttributeError for delta + TORCH/NUMPY input. - compress_model skips candidate chains whose Source size differs from the raw length and errors when none is left. - The forward pass takes the source element width and nibble packing from the chain's Source (chain-internal dtype codes) instead of the record's public dtype code, so encode and decode descriptors agree. Legacy golden output and measured sizes are unchanged. - New binding _core.chain_source_byte_len. Tests: symbolic == forward reference on every production/explorer chain blob (new fixture production_chains.bin, kept current by test_production_chain_fixture.py) and on 3000 random chains; bounded walk, delta-op, missing-plane and overflow cases; container 4 TiB/4 GiB Source fail fast, 8x tolerance, understated Source, orig_size mismatch, chunk-table bound; compact exact-Source checks; writer skip of mis-sized candidates; Python test_source_shape.py. The P4 legacy fixture generator now writes the exact chain and patches the stored Source dim, which reproduces the pinned v1.1 bytes (fixture files unchanged). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/runtime.rs | 589 ++++++++++++++++-- crates/ptwm-core/src/chain/validate.rs | 34 +- crates/ptwm-core/src/codecs/huff_llm.rs | 17 +- crates/ptwm-core/src/compact.rs | 73 ++- crates/ptwm-core/src/compressor.rs | 131 +++- crates/ptwm-core/src/container.rs | 247 +++++++- crates/ptwm-core/src/transforms/source.rs | 65 +- .../tests/fixtures/production_chains.bin | Bin 0 -> 249646 bytes crates/ptwm-core/tests/legacy_fixtures.rs | 24 +- crates/ptwm-py/src/lib.rs | 14 + docs/architecture/ppg.md | 14 +- python/ptwm/core/_compressor.py | 82 +-- .../integrations/_compress_safetensors.py | 10 +- python/ptwm/preprocessing/_chains.py | 130 +++- python/ptwm/sharding/_writer.py | 25 +- .../test_production_chain_fixture.py | 93 +++ tests/ptwm/preprocessing/test_source_shape.py | 209 +++++++ 17 files changed, 1567 insertions(+), 190 deletions(-) create mode 100644 crates/ptwm-core/tests/fixtures/production_chains.bin create mode 100644 tests/ptwm/preprocessing/test_production_chain_fixture.py create mode 100644 tests/ptwm/preprocessing/test_source_shape.py diff --git a/crates/ptwm-core/src/chain/runtime.rs b/crates/ptwm-core/src/chain/runtime.rs index c8b9fca..b4c11bd 100644 --- a/crates/ptwm-core/src/chain/runtime.rs +++ b/crates/ptwm-core/src/chain/runtime.rs @@ -270,65 +270,167 @@ pub fn forward_chain(chain: &Chain, ctx: &ForwardContext) -> Result Result, PtwmCoreError> { - if chain.nodes.is_empty() { +/// Checks that the chain is non-empty and that node 0 is `Source`. The +/// length comes from the Source params (shape × dtype) with checked +/// arithmetic, so a corrupt shape returns an error instead of wrapping. +pub fn chain_source_descriptor(chain: &Chain) -> Result { + let first = chain.nodes.first().ok_or_else(|| { + PtwmCoreError::InvalidContainer("chain descriptors: chain has no nodes".into()) + })?; + if first.op != OpId::Source { return Err(PtwmCoreError::InvalidContainer( - "chain_terminal_descriptors: chain has no nodes".into(), + "chain descriptors: node[0] must be Source".into(), )); } - if chain.nodes[0].op != OpId::Source { - return Err(PtwmCoreError::InvalidContainer( - "chain_terminal_descriptors: node[0] must be Source".into(), - )); + let mut descs = op_from_id(OpId::Source, &first.params)?.propagate_descriptors(&[])?; + if descs.len() != 1 { + return Err(PtwmCoreError::InvalidContainer(format!( + "chain descriptors: Source op produced {} descriptors (expected 1)", + descs.len() + ))); } + Ok(descs.remove(0)) +} + +/// Largest plane, at any node of a chain, that a decoder accepts for a +/// tensor of `orig_size` raw bytes. +/// +/// The factor 16 covers a Source that overstates the raw size by up to 8× +/// (files written by the Python byte-input path before the P4 fix) followed +/// by ops that at most double a plane (MxFp4 values, spherical and +/// alpha-stable radius planes for 16-bit rows). The constant covers block +/// padding such as the BurrowsWheeler index (4 bytes per 256 KiB block). +/// Decoders use this one helper so the bound cannot drift between layouts. +pub fn max_plane_len(orig_size: u64) -> u64 { + orig_size.saturating_mul(16).saturating_add(4096) +} - // Re-instantiate the Source op to call its `propagate_descriptors` on - // no inputs, producing the canonical raw-tensor descriptor (with - // element_width / layout / length_bytes derived from shape + - // dtype_code). - let source_op = op_from_id(OpId::Source, &chain.nodes[0].params)?; - let source_descs = source_op.propagate_descriptors(&[])?; - if source_descs.len() != 1 { +/// Recover the [`PlaneDescriptor`] of every chain terminal in +/// `chain.terminals` order, without the tensor bytes and without allocating +/// any plane buffer. +/// +/// Same as [`propagate_terminal_descriptors_bounded`] with no bound. +pub fn propagate_terminal_descriptors( + chain: &Chain, +) -> Result, PtwmCoreError> { + propagate_terminal_descriptors_bounded(chain, u64::MAX) +} + +/// Recover terminal descriptors symbolically and reject any node output +/// longer than `max_plane_len` bytes. +/// +/// The walk calls `propagate_descriptors` on every node in topological +/// order. It mirrors [`forward_chain`]: +/// +/// - inputs are gathered from the incoming edges sorted by `dst_input_idx`, +/// and an edge's `role_override` replaces the input's role; +/// - fused pairs (see [`try_fuse_at`]) need no special case: fusion applies +/// only to edges without a role override, and every fused kernel emits the +/// same descriptors as the unfused producer and consumer; +/// - each terminal's descriptor takes the terminal's role, as +/// `forward_chain` re-tags terminal planes. +/// +/// No op's output size, layout or role depends on plane contents, so the +/// result equals the descriptors `forward_chain` produces on a tensor of +/// the Source's declared size. The one data-dependent field is the estimate +/// `EntropyEstimate` writes into `vendor_bytes`: this walk reports zeros +/// (its `propagate_descriptors` value). The earlier decoder computed it +/// from placeholder planes, which was not the encode-time value either, and +/// no codec or decoder reads it. +/// +/// Cross-tensor delta ops (`XorDelta`, `FloatDelta`) need dependency planes +/// and return an error here. +/// +/// The Source descriptor is checked against the bound first, so a Source +/// that declares gigabytes fails before any other node is visited. +pub fn propagate_terminal_descriptors_bounded( + chain: &Chain, + max_plane_len: u64, +) -> Result, PtwmCoreError> { + let source = chain_source_descriptor(chain)?; + if source.length_bytes > max_plane_len { return Err(PtwmCoreError::InvalidContainer(format!( - "chain_terminal_descriptors: Source op produced {} descriptors (expected 1)", - source_descs.len() + "chain descriptors: Source declares {} bytes, bound is {max_plane_len}", + source.length_bytes ))); } - let source_descriptor = source_descs.into_iter().next().ok_or_else(|| { - PtwmCoreError::InvalidContainer( - "chain_terminal_descriptors: Source op produced no descriptors".into(), - ) - })?; - let placeholder_bytes = vec![0u8; source_descriptor.length_bytes as usize]; - - let ctx = ForwardContext { - source_bytes: &placeholder_bytes, - source_descriptor, - deps: &[], - }; - let result = forward_chain(chain, &ctx)?; - Ok(result - .terminal_planes - .into_iter() - .map(|(plane, _role)| plane.descriptor) - .collect()) + let topo = topo_sort(chain)?; + let mut outs: Vec>> = vec![None; chain.nodes.len()]; + for &node_idx in &topo { + if node_idx == 0 { + outs[0] = Some(vec![source.clone()]); + continue; + } + let node = &chain.nodes[node_idx]; + if is_delta_op(node.op) { + return Err(PtwmCoreError::InvalidContainer(format!( + "chain descriptors: node[{node_idx}] ({:?}) is a cross-tensor delta op \ + and needs dependency planes", + node.op + ))); + } + let mut input_edges: Vec<_> = chain + .edges + .iter() + .filter(|e| e.dst_node as usize == node_idx) + .collect(); + input_edges.sort_by_key(|e| e.dst_input_idx); + let mut inputs = Vec::with_capacity(input_edges.len()); + for edge in input_edges { + let mut desc = outs + .get(edge.src_node as usize) + .and_then(|o| o.as_ref()) + .and_then(|o| o.get(edge.src_output_idx as usize)) + .cloned() + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "chain descriptors: missing plane for node={}, output={}", + edge.src_node, edge.src_output_idx + )) + })?; + if let Some(ref role) = edge.role_override { + desc.role = role.clone(); + } + inputs.push(desc); + } + let produced = op_from_id(node.op, &node.params)?.propagate_descriptors(&inputs)?; + if let Some(d) = produced.iter().find(|d| d.length_bytes > max_plane_len) { + return Err(PtwmCoreError::InvalidContainer(format!( + "chain descriptors: node[{node_idx}] ({:?}) declares a {}-byte plane, \ + bound is {max_plane_len}", + node.op, d.length_bytes + ))); + } + outs[node_idx] = Some(produced); + } + chain + .terminals + .iter() + .map(|term| { + let mut desc = outs + .get(term.node_idx as usize) + .and_then(|o| o.as_ref()) + .and_then(|o| o.get(term.output_idx as usize)) + .cloned() + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "chain descriptors: terminal references missing plane \ + (node={}, output={})", + term.node_idx, term.output_idx + )) + })?; + desc.role = term.role.clone(); + Ok(desc) + }) + .collect() +} + +/// Recover terminal descriptors. Kept for callers of the earlier API; it is +/// [`propagate_terminal_descriptors`]. +pub fn chain_terminal_descriptors(chain: &Chain) -> Result, PtwmCoreError> { + propagate_terminal_descriptors(chain) } // --------------------------------------------------------------------------- @@ -1362,4 +1464,391 @@ mod tests { plane.descriptor.role, ); } + + // ── Symbolic descriptor propagation (P5) ──────────────────────────────── + + /// The earlier decoder path: run `forward_chain` over a zero buffer of + /// the Source's declared size and read the terminal descriptors back. + /// Kept here as the reference the symbolic walk must match. + fn forward_descriptors(chain: &Chain) -> Result, PtwmCoreError> { + let source_descriptor = chain_source_descriptor(chain)?; + let placeholder = vec![0u8; source_descriptor.length_bytes as usize]; + let ctx = ForwardContext { + source_bytes: &placeholder, + source_descriptor, + deps: &[], + }; + Ok(forward_chain(chain, &ctx)? + .terminal_planes + .into_iter() + .map(|(plane, _)| plane.descriptor) + .collect()) + } + + /// Compare the symbolic walk with the forward reference on one chain. + /// Returns true when both ran and agreed; false when the forward pass + /// rejects the chain (it checks some lengths `propagate_descriptors` + /// does not, so the symbolic walk may accept those). + fn assert_symbolic_matches_forward(label: &str, chain: &Chain) -> bool { + let Ok(expected) = forward_descriptors(chain) else { + return false; + }; + let mut got = propagate_terminal_descriptors(chain) + .unwrap_or_else(|e| panic!("{label}: forward succeeded but symbolic failed: {e}")); + let mut expected = expected; + // EntropyEstimate writes an estimate of the bytes it sees into + // `vendor_bytes`. The forward reference computes it from placeholder + // planes (zeros at the Source, arbitrary after ops such as + // AlphaStableNormalize), the symbolic walk always reports zeros. + // Neither is the encode-time value and no decoder reads it, so it + // is left out of the comparison. + if chain.nodes.iter().any(|n| n.op == OpId::EntropyEstimate) { + for d in got.iter_mut().chain(expected.iter_mut()) { + d.vendor_bytes.clear(); + } + } + assert_eq!( + got, expected, + "{label}: symbolic descriptors differ from forward" + ); + true + } + + #[test] + fn symbolic_descriptors_match_forward_on_production_chains() { + // Written by tests/ptwm/preprocessing/test_production_chain_fixture.py + // from every Python chain builder and the explorer's first candidates. + let data: &[u8] = include_bytes!("../../tests/fixtures/production_chains.bin"); + let mut pos = 0usize; + let (mut total, mut valid, mut compared) = (0usize, 0usize, 0usize); + let mut fused = 0usize; + while pos < data.len() { + let name_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize; + pos += 2; + let label = std::str::from_utf8(&data[pos..pos + name_len]).unwrap(); + pos += name_len; + let blob_len = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap()) as usize; + pos += 4; + let blob = &data[pos..pos + blob_len]; + pos += blob_len; + total += 1; + let (chain, _) = crate::chain::wire::read_chain_blob(blob) + .unwrap_or_else(|e| panic!("{label}: fixture blob does not parse: {e}")); + // Some explorer chains do not fit every fixture shape (a 0-d + // fp16 tensor cannot be split into 4 bytes); the decoder + // rejects those through `validate_chain` before descriptors. + if crate::chain::validate::validate_chain(&chain, 0).is_err() { + assert!(label.starts_with("explore"), "{label}: invalid chain"); + continue; + } + valid += 1; + if (1..chain.nodes.len()).any(|i| try_fuse_at(&chain, i as u8).is_some()) { + fused += 1; + } + if assert_symbolic_matches_forward(label, &chain) { + compared += 1; + } + } + assert!(total > 500, "fixture holds only {total} chains"); + // Forward rejects a few builder/shape pairs (e.g. microscale repack + // on lengths that are not whole blocks); nearly all must compare. + assert!( + compared * 10 >= valid * 9, + "only {compared} of {valid} valid chains ran forward ({total} in fixture)" + ); + assert!(fused > 50, "only {fused} chains exercise fusion"); + } + + /// Small deterministic PRNG (xorshift64*), so the test needs no crate. + struct Rng(u64); + impl Rng { + fn next(&mut self) -> u64 { + self.0 ^= self.0 >> 12; + self.0 ^= self.0 << 25; + self.0 ^= self.0 >> 27; + self.0.wrapping_mul(0x2545_F491_4F6C_DD1D) + } + fn below(&mut self, n: u64) -> u64 { + self.next() % n + } + fn pick(&mut self, items: &[T]) -> T { + items[self.below(items.len() as u64) as usize].clone() + } + } + + fn random_role(rng: &mut Rng) -> Role { + use crate::types::role::{NibbleKind, ScaleFormat, ValueFormat}; + rng.pick(&[ + Role::Raw, + Role::ExponentByte, + Role::Index, + Role::MantissaByte { index: 0, of: 2 }, + Role::Nibble { + kind: NibbleKind::Exponent, + }, + Role::Value { + format: ValueFormat::IntN { bits: 8 }, + }, + Role::Scale { + format: ScaleFormat::E8M0, + }, + Role::GlobalScale { + format: ScaleFormat::F32, + }, + ]) + } + + /// A random op and its params, or a Concat of two open outputs. + fn random_node(rng: &mut Rng) -> (OpId, Vec) { + match rng.below(19) { + 0 => (OpId::BitReorderIeee16, vec![]), + 1 => (OpId::BitReorderIeee32, vec![]), + 2 => (OpId::BitReorderFp8E4M3, vec![]), + 3 => (OpId::BitReorderFp8E5M2, vec![]), + 4 => (OpId::ByteSplit, vec![rng.pick(&[2u8, 4, 8])]), + 5 => (OpId::NibbleSplit, vec![]), + 6 => (OpId::BytePassthrough, vec![]), + 7 => (OpId::MxFp4Deinterleave, vec![32]), + 8 => (OpId::BlockMicroscalingRepack, vec![32]), + 9 => (OpId::IntDelta, vec![]), + 10 => (OpId::PredictorXor, vec![]), + 11 => (OpId::IndexBitwidthPack, vec![1 + rng.below(8) as u8]), + 12 => (OpId::EntropyEstimate, vec![]), + 13 => (OpId::MantissaZeroStrip, vec![rng.below(5) as u8]), + 14 => (OpId::BurrowsWheeler, vec![]), + 15 => (OpId::MoveToFront, vec![]), + 16 => { + let mut p = Vec::new(); + if rng.below(2) == 0 { + Layout::Flat.write(&mut p); + } else { + Layout::Rows { + row_len: 1 + rng.below(16) as u32, + } + .write(&mut p); + } + (OpId::Reshape, p) + } + 17 => ( + OpId::SphericalNormalize, + vec![rng.below(2) as u8, 0, rng.below(3) as u8], + ), + _ => (OpId::AlphaStableNormalize, vec![0, rng.below(3) as u8]), + } + } + + /// Build a random well-formed chain: each step consumes one or two open + /// outputs; whatever stays open becomes a terminal. Op choices whose + /// descriptors do not propagate are dropped. + fn random_chain(rng: &mut Rng) -> Chain { + let dtype = rng.pick(&[ + 0x0002u16, 0x000F, 0x0003, 0x0004, 0x0005, 0x0006, 0x0009, 0x0010, 0x0011, 0x001F, + ]); + let shape: Vec = (0..1 + rng.below(2)) + .map(|_| rng.below(48) as u32 + u32::from(rng.below(8) != 0)) + .collect(); + let mut src = vec![shape.len() as u8]; + for d in &shape { + src.extend_from_slice(&d.to_le_bytes()); + } + src.extend_from_slice(&dtype.to_le_bytes()); + let mut chain = Chain { + nodes: vec![ChainNode { + op: OpId::Source, + params: src, + }], + edges: vec![], + terminals: vec![], + }; + let Ok(source) = chain_source_descriptor(&chain) else { + return chain; + }; + // Open outputs: (node, output, descriptor as seen by a consumer). + let mut open: Vec<(u8, u8, PlaneDescriptor)> = vec![(0, 0, source)]; + for _ in 0..1 + rng.below(5) { + if open.is_empty() { + break; + } + let concat = open.len() >= 2 && rng.below(6) == 0; + let n_in = if concat { 2 } else { 1 }; + let mut picked = Vec::new(); + for _ in 0..n_in { + let i = rng.below(open.len() as u64) as usize; + picked.push(open.swap_remove(i)); + } + let overrides: Vec> = picked + .iter() + .map(|_| (rng.below(4) == 0).then(|| random_role(rng))) + .collect(); + let inputs: Vec = picked + .iter() + .zip(&overrides) + .map(|((_, _, d), o)| { + let mut d = d.clone(); + if let Some(r) = o { + d.role = r.clone(); + } + d + }) + .collect(); + let (op, params) = if concat { + let mut p = vec![2u8]; + for d in &inputs { + p.extend_from_slice(&(d.length_bytes as u32).to_le_bytes()); + } + (OpId::Concat, p) + } else { + random_node(rng) + }; + let produced = op_from_id(op, ¶ms).and_then(|o| o.propagate_descriptors(&inputs)); + let Ok(produced) = produced else { + open.extend(picked); + continue; + }; + let idx = chain.nodes.len() as u8; + chain.nodes.push(ChainNode { op, params }); + for (slot, ((src_node, src_out, _), role_override)) in + picked.into_iter().zip(overrides).enumerate() + { + chain.edges.push(ChainEdge { + src_node, + src_output_idx: src_out, + dst_node: idx, + dst_input_idx: slot as u8, + role_override, + vendor_bytes: vec![], + }); + } + for (o, d) in produced.into_iter().enumerate() { + open.push((idx, o as u8, d)); + } + } + for (node_idx, output_idx, _) in open { + chain.terminals.push(TerminalRef { + node_idx, + output_idx, + role: random_role(rng), + }); + } + chain + } + + #[test] + fn symbolic_descriptors_match_forward_on_random_chains() { + let mut rng = Rng(0x9E37_79B9_7F4A_7C15); + let (mut valid, mut compared, mut multi_node) = (0usize, 0usize, 0usize); + for i in 0..3000 { + let chain = random_chain(&mut rng); + if crate::chain::validate::validate_chain(&chain, 0).is_err() { + continue; + } + valid += 1; + if assert_symbolic_matches_forward(&format!("random chain {i}: {chain:?}"), &chain) { + compared += 1; + if chain.nodes.len() > 2 { + multi_node += 1; + } + } + } + assert!(valid > 1000, "only {valid} random chains were valid"); + assert!(compared > 800, "only {compared} random chains ran forward"); + assert!( + multi_node > 300, + "only {multi_node} compared chains have 2+ ops" + ); + } + + #[test] + fn symbolic_applies_terminal_role() { + use crate::types::role::ScaleFormat; + let mut chain = passthrough_terminal_chain(); + chain.terminals[0].role = Role::Scale { + format: ScaleFormat::E8M0, + }; + let descs = propagate_terminal_descriptors(&chain).unwrap(); + assert_eq!( + descs[0].role, + Role::Scale { + format: ScaleFormat::E8M0 + } + ); + assert_eq!(descs, forward_descriptors(&chain).unwrap()); + } + + #[test] + fn bounded_walk_rejects_huge_source_without_allocating() { + // [u32::MAX, 64] fp64 declares about 2.2 TB. The bound from a 16-byte + // tensor rejects it before any node runs; a forward pass would have + // tried to allocate the whole buffer. + let mut params = vec![2u8]; + params.extend_from_slice(&u32::MAX.to_le_bytes()); + params.extend_from_slice(&64u32.to_le_bytes()); + params.extend_from_slice(&0x0004u16.to_le_bytes()); + let mut chain = passthrough_terminal_chain(); + chain.nodes[0].params = params; + let t0 = std::time::Instant::now(); + let err = propagate_terminal_descriptors_bounded(&chain, max_plane_len(16)) + .unwrap_err() + .to_string(); + assert!(err.contains("Source declares"), "{err}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(1)); + } + + #[test] + fn bounded_walk_rejects_oversized_intermediate_plane() { + // Source [256, 1] fp16 (512 bytes) → SphericalNormalize(mode 0, + // fp16): the fp32 radius of 256 one-value rows is 1024 bytes, larger + // than the source. A bound of 512 passes node 0 and stops at node 1. + let mut params = vec![2u8]; + params.extend_from_slice(&256u32.to_le_bytes()); + params.extend_from_slice(&1u32.to_le_bytes()); + params.extend_from_slice(&0x0002u16.to_le_bytes()); + let mut chain = passthrough_terminal_chain(); + chain.nodes[0].params = params; + chain.nodes[1] = ChainNode { + op: OpId::SphericalNormalize, + params: vec![0, 0, 2], + }; + chain.terminals = (0..3) + .map(|o| TerminalRef { + node_idx: 1, + output_idx: o, + role: Role::Raw, + }) + .collect(); + let descs = propagate_terminal_descriptors(&chain).unwrap(); + let largest = descs.iter().map(|d| d.length_bytes).max().unwrap(); + assert!(largest > 512, "planes {descs:?}"); + let err = propagate_terminal_descriptors_bounded(&chain, 512) + .unwrap_err() + .to_string(); + assert!(err.contains("node[1]"), "{err}"); + assert!(propagate_terminal_descriptors_bounded(&chain, largest).is_ok()); + } + + #[test] + fn symbolic_walk_rejects_delta_ops() { + let mut chain = passthrough_terminal_chain(); + chain.nodes[1] = ChainNode { + op: OpId::XorDelta, + params: vec![0], + }; + let err = propagate_terminal_descriptors(&chain) + .unwrap_err() + .to_string(); + assert!(err.contains("delta"), "{err}"); + } + + #[test] + fn symbolic_walk_reports_missing_planes_instead_of_panicking() { + let mut chain = passthrough_terminal_chain(); + chain.edges[0].src_output_idx = 5; + assert!(propagate_terminal_descriptors(&chain).is_err()); + let mut chain = passthrough_terminal_chain(); + chain.terminals[0].output_idx = 3; + assert!(propagate_terminal_descriptors(&chain).is_err()); + let mut chain = passthrough_terminal_chain(); + chain.terminals[0].node_idx = 9; + assert!(propagate_terminal_descriptors(&chain).is_err()); + } } diff --git a/crates/ptwm-core/src/chain/validate.rs b/crates/ptwm-core/src/chain/validate.rs index 6b2b617..6fa9599 100644 --- a/crates/ptwm-core/src/chain/validate.rs +++ b/crates/ptwm-core/src/chain/validate.rs @@ -141,15 +141,26 @@ pub fn validate_chain(chain: &Chain, n_dependencies: u8) -> Result<(), PtwmCoreE let inputs: Vec = input_edges .iter() .map(|e| { - // Apply role_override if present. - let src_desc = &node_outputs[e.src_node as usize][e.src_output_idx as usize]; + // A wire-derived edge can name an output the source node + // does not produce; that is an error, not a panic. + let src_desc = node_outputs + .get(e.src_node as usize) + .and_then(|o| o.get(e.src_output_idx as usize)) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "validate_chain: edge into node[{node_idx}] references missing \ + output (node={}, output={})", + e.src_node, e.src_output_idx + )) + })?; let mut desc = src_desc.clone(); + // Apply role_override if present. if let Some(ref role) = e.role_override { desc.role = role.clone(); } - desc + Ok(desc) }) - .collect(); + .collect::>()?; let outputs = op.propagate_descriptors(&inputs).map_err(|e| { PtwmCoreError::InvalidContainer(format!( @@ -509,4 +520,19 @@ mod tests { // n_dependencies = 2, dep_idx = 10 → should fail assert!(validate_chain(&chain, 2).is_err()); } + + #[test] + fn edge_to_missing_output_is_an_error_not_a_panic() { + // Source produces one output; an edge from output 5 used to index + // past the end of node_outputs[0]. + let mut chain = passthrough_terminal_chain(); + chain.edges[0].src_output_idx = 5; + chain.terminals.push(TerminalRef { + node_idx: 0, + output_idx: 0, + role: Role::Raw, + }); + let err = validate_chain(&chain, 0).unwrap_err().to_string(); + assert!(err.contains("missing output"), "{err}"); + } } diff --git a/crates/ptwm-core/src/codecs/huff_llm.rs b/crates/ptwm-core/src/codecs/huff_llm.rs index c331d54..6c45c4a 100644 --- a/crates/ptwm-core/src/codecs/huff_llm.rs +++ b/crates/ptwm-core/src/codecs/huff_llm.rs @@ -36,14 +36,15 @@ impl PlaneCodec for HuffLlm5Bit { } fn accepts(&self, descriptor: &PlaneDescriptor) -> bool { - // Gate on the *role*, not `element_width`. `source_descriptor_for` - // derives `element_width` from a chain-internal dtype namespace, but - // production passes *external* dtype codes, so a raw fp16/bf16 source - // plane reaches a codec mislabeled (bf16 -> Byte, fp16 -> Word8 via a - // namespace collision). `Role::Raw` is stamped reliably by - // `source_descriptor_for` and preserved through the terminal re-tag, - // and the codec reads raw 16-bit LE words directly (it never relied on - // `element_width` for correctness). Only terminal planes are + // Gate on the *role*, not `element_width`. The compressor now takes + // the source plane's width from the chain's Source node, so encode + // and decode agree on it, but files written before that carried + // planes whose width came from the public dtype code (bf16 -> Byte, + // fp16 -> Word8), and the codec must keep accepting the same planes. + // `Role::Raw` is stamped reliably by the Source and preserved + // through the terminal re-tag, and the codec reads raw 16-bit LE + // words directly (it never relied on `element_width` for + // correctness). Only terminal planes are // trial-encoded. The huff_llm raw chain is the only one terminating on // the *un-transformed* source plane, but the spherical_normalize and // alpha_stable_normalize chains also expose `Role::Raw` terminals diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index 862f68a..603f6ad 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -51,8 +51,10 @@ //! unknown flag bits, and an incompatible layout gets a new magic. use crate::chain::Chain; -use crate::chain::dispatch::op_from_id; -use crate::chain::runtime::{InverseContext, chain_terminal_descriptors, inverse_chain}; +use crate::chain::runtime::{ + InverseContext, chain_source_descriptor, inverse_chain, max_plane_len, + propagate_terminal_descriptors_bounded, +}; use crate::chain::validate::validate_chain; use crate::chain::wire::{read_chain, write_chain}; use crate::codec::{CodecId, StateSource, codec_for}; @@ -378,31 +380,18 @@ fn decode_body( let (chain, consumed) = read_chain(r.rest(), &table)?; r.pos += consumed; validate_chain(&chain, 0)?; - // `chain_terminal_descriptors` runs the chain forward on a zero buffer - // of the size the Source node declares, and codecs pre-allocate the - // plane lengths that follow from it. Bound the declared size by the - // frame's `orig_size` first, so a corrupt chain cannot request gigabytes - // for a small tensor. The factor covers a source declared in elements of - // up to 8 bytes for byte input; the constant covers block padding. - let max_len = orig_size.saturating_mul(8).saturating_add(4096); - let first = chain - .nodes - .first() - .ok_or_else(|| bad("chain has no nodes"))?; - if first.op != OpId::Source { - return Err(bad("chain node 0 is not Source")); - } - let source = op_from_id(OpId::Source, &first.params)?.propagate_descriptors(&[])?; - if source.iter().any(|d| d.length_bytes > max_len) { + // Writers check that the chain's Source declares exactly the raw byte + // count, so a frame whose Source disagrees with `orig_size` is corrupt. + // Checking before the descriptor walk keeps a corrupt chain from sizing + // planes (and codec allocations) beyond the tensor. + let source = chain_source_descriptor(&chain)?; + if source.length_bytes != orig_size { return Err(bad(&format!( - "chain declares a source larger than {max_len} bytes for a \ - {orig_size}-byte tensor" + "chain Source declares {} bytes, header says {orig_size}", + source.length_bytes ))); } - let descriptors = chain_terminal_descriptors(&chain)?; - if descriptors.iter().any(|d| d.length_bytes > max_len) { - return Err(bad("chain declares a plane larger than its bound")); - } + let descriptors = propagate_terminal_descriptors_bounded(&chain, max_plane_len(orig_size))?; let mut terminal_planes: Vec = Vec::with_capacity(descriptors.len()); for descriptor in descriptors { @@ -432,11 +421,15 @@ fn decode_body( let layout = legacy_plane_layout(descriptor.layout); let decoded_len = usize::try_from(descriptor.length_bytes) .map_err(|_| bad("plane length overflows usize"))?; - // Self-delimiting codecs may return a length other than the chain - // descriptor's (the descriptor can overstate it for byte input with a - // multi-byte dtype). The container reader accepts this too; the - // `orig_size` check and the payload hash below cover the result. let bytes = codec.decode(state_version, state, payload, &layout, decoded_len)?; + // The Source is exact, so every plane decodes to its descriptor's + // length. + if bytes.len() != decoded_len { + return Err(bad(&format!( + "plane decoded to {} bytes, chain says {decoded_len}", + bytes.len() + ))); + } if let Some(expected) = crc && crc32fast::hash(&bytes) != expected { @@ -732,10 +725,32 @@ mod tests { b[at + 1..at + 5].copy_from_slice(&0xFFFF_FFF0u32.to_le_bytes()); let t0 = std::time::Instant::now(); let err = decode_compact(&b).unwrap_err().to_string(); - assert!(err.contains("source larger than"), "{err}"); + assert!(err.contains("Source declares"), "{err}"); assert!(t0.elapsed() < std::time::Duration::from_secs(1)); } + #[test] + fn source_must_declare_orig_size_exactly() { + // Unlike legacy containers, compact frames carry no files written + // before the writer check, so any Source size other than + // `orig_size` is corrupt, overstated or understated. + let raw: Vec = (0..3000u32).map(|i| (i % 5) as u8).collect(); + let blob = compress_compact(&input(&raw, None, None), &opts(false, false)) + .unwrap() + .unwrap(); + let needle: Vec = [&[1u8][..], &3000u32.to_le_bytes()].concat(); + let at = blob + .windows(needle.len()) + .position(|w| w == needle.as_slice()) + .unwrap(); + for dim in [3001u32, 1500, 6000] { + let mut b = blob.clone(); + b[at + 1..at + 5].copy_from_slice(&dim.to_le_bytes()); + let err = decode_compact(&b).unwrap_err().to_string(); + assert!(err.contains("Source declares"), "{dim}: {err}"); + } + } + #[test] fn varint_round_trips() { for v in [0u64, 1, 127, 128, 300, u32::MAX as u64, u64::MAX] { diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index b17432f..e2c6c66 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -12,7 +12,7 @@ use rayon::prelude::*; use xxhash_rust::xxh64::xxh64; use crate::chain::Chain; -use crate::chain::runtime::{ForwardContext, forward_chain}; +use crate::chain::runtime::{ForwardContext, chain_source_descriptor, forward_chain}; use crate::chain::validate::validate_chain; use crate::codec::{CodecId, StateSource, codec_for}; use crate::codecs::order1_scale_ac; @@ -518,9 +518,43 @@ fn pass1_select_chain( let mut n_failed: usize = 0; for (idx, chain) in candidates.iter().enumerate() { + // Decoders rebuild every plane descriptor from the chain's Source + // node, so a chain whose Source does not declare exactly the raw + // byte count would write planes the decoder sizes wrongly. Skip it + // like any other failing candidate. + let chain_src = match chain_source_descriptor(chain) { + Ok(d) => d, + Err(e) => { + last_err = Some(e); + n_failed += 1; + continue; + } + }; + let declared = chain_src.length_bytes; + if declared != raw_bytes.len() as u64 { + let e = PtwmCoreError::InvalidContainer(format!( + "candidate chain {idx} Source declares {declared} bytes, tensor has {}", + raw_bytes.len() + )); + tracing::warn!("pass1_select_chain: {e}; skipping and trying remaining candidates"); + last_err = Some(e); + n_failed += 1; + continue; + } + // `source_descriptor` derives the element width from the record's + // dtype code, which is in the public `Dtype` code space, while the + // decoder derives it from the chain's Source node, which uses the + // chain-internal code space (fp32 is 1 in one and 0x0003 in the + // other). Take the width and packing from the chain so encode-time + // and decode-time descriptors agree. The layout stays as given: the + // record's shape decides it, and changing it would change the + // legacy writer's output. + let mut forward_source = source_descriptor.clone(); + forward_source.element_width = chain_src.element_width; + forward_source.is_nibble_packed = chain_src.is_nibble_packed; let ctx = ForwardContext { source_bytes: raw_bytes, - source_descriptor: source_descriptor.clone(), + source_descriptor: forward_source, deps: &[], }; let fwd = match forward_chain(chain, &ctx) { @@ -1401,6 +1435,99 @@ mod tests { } } + #[test] + fn candidate_with_mis_sized_source_is_skipped() { + let raw: Vec = (0..256u32).map(|i| (i % 7) as u8).collect(); + let opts = || CompressorOptions { + method_hint: 1, + emit_payload_hash: true, + emit_plane_crc: true, + forced_codec: None, + allow_codec_ids: None, + chunk_size: None, + }; + let input = |chains: Vec| InputTensor { + name: "t".into(), + candidate_chains: chains, + dtype_code: 0x0005, + input_format: 0, + raw_bytes: &raw, + shape: None, + dtype_name: None, + delta_reference_blake3: None, + }; + + // The first candidate declares 4× the raw size (the old Python + // byte-input bug); the writer must pick the exact one. + let mut buf: Vec = Vec::new(); + let picked = compress_model( + Cursor::new(&mut buf), + &[input(vec![ + passthrough_value_chain(raw.len() as u32 * 4), + passthrough_value_chain(raw.len() as u32), + ])], + opts(), + ) + .unwrap(); + assert_eq!(picked, vec![0]); + let reader = ContainerReader::open(&buf).unwrap(); + assert_eq!(reader.chain_registry.chains.len(), 1); + assert_eq!( + reader.chain_registry.chains[0].1, + passthrough_value_chain(raw.len() as u32) + ); + assert_eq!(reader.decode_tensor("t").unwrap(), raw); + + // With only mis-sized candidates (over- and understated), compression + // fails and names the mismatch. + let mut buf: Vec = Vec::new(); + let err = compress_model( + Cursor::new(&mut buf), + &[input(vec![ + passthrough_value_chain(raw.len() as u32 * 4), + passthrough_value_chain(raw.len() as u32 / 2), + ])], + opts(), + ) + .unwrap_err() + .to_string(); + assert!(err.contains("all 2 candidate chains failed"), "{err}"); + assert!( + err.contains("Source declares 128 bytes, tensor has 256"), + "{err}" + ); + } + + #[test] + fn encode_descriptors_take_element_width_from_the_chain() { + // The record's dtype code (fp16 is 4 in the public code space) maps + // to Word8 in the chain-internal space; the chain's Source says + // Word2. Encode-time planes must carry the chain's width, as the + // decoder rebuilds them from the chain. + let n = 64u32; + let raw: Vec = (0..n) + .flat_map(|i| (0x3800u16 | i as u16).to_le_bytes()) + .collect(); + let chain = fp16_split_chain(n); + let outcome = pass1_select_chain( + std::slice::from_ref(&chain), + &raw, + &source_descriptor_for(4, raw.len() as u64, None), + None, + None, + false, + None, + ) + .unwrap(); + let decoded = crate::chain::runtime::propagate_terminal_descriptors(&chain).unwrap(); + assert_eq!(outcome.terminals.len(), decoded.len()); + for ((plane, _), want) in outcome.terminals.iter().zip(&decoded) { + assert_eq!(plane.descriptor.element_width, want.element_width); + assert_eq!(plane.descriptor.is_nibble_packed, want.is_nibble_packed); + assert_eq!(plane.descriptor.length_bytes, want.length_bytes); + } + } + #[test] fn compress_forced_chunked_skip_roundtrip() { // Exercises the deferred-encode fast path: a forced codec + single diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 9d473a9..ab593cf 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -1219,7 +1219,10 @@ impl<'a> ContainerReader<'a> { &self, rec: &crate::tensor_record::TensorRecord, ) -> Result, PtwmCoreError> { - use crate::chain::runtime::{InverseContext, chain_terminal_descriptors, inverse_chain}; + use crate::chain::runtime::{ + InverseContext, chain_source_descriptor, inverse_chain, max_plane_len, + propagate_terminal_descriptors_bounded, + }; use crate::chain::validate::validate_chain; use crate::codec::codec_for; use crate::tensor_record::{CHAIN_REF_INLINE_BIT, TENSOR_FLAG_PAYLOAD_HASH}; @@ -1253,10 +1256,27 @@ impl<'a> ContainerReader<'a> { })?; validate_chain(chain, n_dependencies)?; - // Discover the descriptor of every terminal by walking the chain - // forward symbolically (no source bytes needed). This is a pure - // function of the chain structure — see `chain_terminal_descriptors`. - let term_descriptors = chain_terminal_descriptors(chain)?; + // Bound the chain by the record's `orig_size` before any plane is + // sized from it. The Source may overstate the raw size by up to 8×: + // the Python byte-input path wrote such chains before the P4 fix, + // and those files must keep decoding. A Source declaring more is + // corrupt; rejecting it here keeps codecs from pre-allocating + // gigabytes for a small tensor. + let source_desc = chain_source_descriptor(chain)?; + let max_source = rec.orig_size.saturating_mul(8); + if source_desc.length_bytes > max_source { + return Err(PtwmCoreError::InvalidContainer(format!( + "decode_tensor: chain Source declares {} bytes for a {}-byte tensor \ + (at most {max_source} accepted)", + source_desc.length_bytes, rec.orig_size + ))); + } + let plane_bound = max_plane_len(rec.orig_size); + + // Discover the descriptor of every terminal by walking the chain's + // descriptors symbolically (no source bytes, no plane buffers). Every + // node's output is bounded by `plane_bound`. + let term_descriptors = propagate_terminal_descriptors_bounded(chain, plane_bound)?; if term_descriptors.len() != rec.terminals.len() { return Err(PtwmCoreError::InvalidContainer(format!( "decode_tensor: chain has {} terminals but record has {} planes", @@ -1349,7 +1369,19 @@ impl<'a> ContainerReader<'a> { // own state, so decode them across the rayon pool and // concatenate in order. (The encoder frames chunks sequentially, // so the offsets are contiguous.) - let total: usize = chunks.iter().map(|e| e.decoded_size as usize).sum(); + // The chunk table is file data: bound every decoded size and + // their sum before any chunk is decoded or allocated. + let total = chunks + .iter() + .try_fold(0u64, |acc, e| acc.checked_add(u64::from(e.decoded_size))) + .filter(|&t| t <= plane_bound) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "chunk table decodes to more than {plane_bound} bytes for a \ + {}-byte tensor", + rec.orig_size + )) + })? as usize; let decoded_chunks: Vec> = chunks .par_iter() .enumerate() @@ -1461,8 +1493,11 @@ impl<'a> ContainerReader<'a> { )); } }; - // Decoded length comes from the propagated terminal descriptor. - let decoded_len = descriptor.length_bytes as usize; + // Decoded length comes from the propagated terminal descriptor, + // which `plane_bound` already limits. + let decoded_len = usize::try_from(descriptor.length_bytes).map_err(|_| { + PtwmCoreError::InvalidContainer("plane length overflows usize".into()) + })?; codec.decode( state_version, &state_bytes, @@ -1471,6 +1506,19 @@ impl<'a> ContainerReader<'a> { decoded_len, )? }; + // Self-delimiting codecs may return a length other than the + // descriptor's: older files carry Sources that overstate (Python + // byte input) or understate (byte passthrough over an element + // shape) the raw size. Bound the result by `orig_size` instead; + // the length check after `inverse_chain` and the payload hash + // cover the rest. + if plane_bytes.len() as u64 > plane_bound { + return Err(PtwmCoreError::InvalidContainer(format!( + "decoded plane is {} bytes, bound is {plane_bound} for a {}-byte tensor", + plane_bytes.len(), + rec.orig_size + ))); + } // Plane CRC enforcement (when emitted by writer). if let Some(expected_crc) = plane_rec.crc32 { @@ -1494,6 +1542,13 @@ impl<'a> ContainerReader<'a> { deps: &[], }; let source_plane = inverse_chain(chain, &inv_ctx)?; + if source_plane.bytes.len() as u64 != rec.orig_size { + return Err(PtwmCoreError::InvalidContainer(format!( + "decode_tensor: chain decoded {} bytes, record says {}", + source_plane.bytes.len(), + rec.orig_size + ))); + } let raw: Vec = source_plane.bytes.as_ref().to_vec(); // Tensor-level payload hash validation. @@ -3104,4 +3159,180 @@ label = "test-author" "round_trip_with_builtin_via_router: bit-exact roundtrip failed" ); } + + // ── Decode bounds (P5) ─────────────────────────────────────────────── + + /// Compress `raw` through a passthrough chain with the Identity codec + /// (so plane bytes equal tensor bytes), then rewrite the container with + /// the record's chain inline and `edit` applied to the record. + fn identity_container_with( + raw: &[u8], + chunk_size: Option, + edit: impl FnOnce(&mut TensorRecord), + ) -> Vec { + use crate::compressor::{CompressorOptions, InputTensor, compress_model}; + use crate::tensor_record::{CHAIN_REF_INLINE_BIT, parse_tensor_record}; + use std::io::Cursor; + + let chain = v3_passthrough_chain(raw.len() as u32); + let inputs = vec![InputTensor { + name: "t".into(), + candidate_chains: vec![chain.clone()], + dtype_code: 0x0005, + input_format: 0, + raw_bytes: raw, + shape: None, + dtype_name: None, + delta_reference_blake3: None, + }]; + let mut buf: Vec = Vec::new(); + compress_model( + Cursor::new(&mut buf), + &inputs, + CompressorOptions { + method_hint: 1, + emit_payload_hash: false, + emit_plane_crc: false, + forced_codec: Some(CodecId::Identity), + allow_codec_ids: None, + chunk_size, + }, + ) + .unwrap(); + let reader = ContainerReader::open(&buf).unwrap(); + assert_eq!(reader.decode_tensor("t").unwrap(), raw); + let (mut rec, _) = + parse_tensor_record(reader.get_tensor_record_bytes("t").unwrap()).unwrap(); + rec.chain_ref = CHAIN_REF_INLINE_BIT; + rec.inline_chain = Some(chain); + edit(&mut rec); + + let mut out: Vec = Vec::new(); + { + let mut w = ContainerWriter::new(std::io::Cursor::new(&mut out)).unwrap(); + w.set_extension_table(reader.extension_table.clone()); + w.append_tensor(&rec).unwrap(); + w.finalize().unwrap(); + } + out + } + + fn set_source_shape(rec: &mut TensorRecord, dims: &[u32], dtype_code: u16) { + let mut params = vec![dims.len() as u8]; + for d in dims { + params.extend_from_slice(&d.to_le_bytes()); + } + params.extend_from_slice(&dtype_code.to_le_bytes()); + rec.inline_chain.as_mut().unwrap().nodes[0].params = params; + } + + #[test] + fn rewritten_identity_container_still_decodes() { + let raw: Vec = (0..64u8).collect(); + let buf = identity_container_with(&raw, None, |_| {}); + assert_eq!( + ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap(), + raw + ); + } + + #[test] + fn huge_source_fails_fast_without_allocating() { + // A 16-byte tensor whose chain declares [1<<20, 1<<20] fp32 (4 TiB), + // and one that declares 4 GiB: both fail before any plane buffer is + // sized from the chain. + let raw = vec![7u8; 16]; + for (dims, dtype) in [ + (vec![1u32 << 20, 1 << 20], 0x0003u16), + (vec![u32::MAX], 0x0005), + ] { + let buf = + identity_container_with(&raw, None, |rec| set_source_shape(rec, &dims, dtype)); + let reader = ContainerReader::open(&buf).unwrap(); + let t0 = std::time::Instant::now(); + let err = reader.decode_tensor("t").unwrap_err().to_string(); + assert!(err.contains("Source declares"), "{err}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(1)); + } + } + + #[test] + fn overstated_source_within_8x_still_decodes() { + // Files written by the Python byte-input path before the P4 fix + // declare fp32 elements for a byte count: 4× the raw size. + let raw: Vec = (0..64u8).collect(); + let buf = identity_container_with(&raw, None, |rec| set_source_shape(rec, &[64], 0x0003)); + assert_eq!( + ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap(), + raw + ); + // Above 8× is rejected. + let buf = identity_container_with(&raw, None, |rec| set_source_shape(rec, &[129], 0x0003)); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert!(err.contains("Source declares"), "{err}"); + } + + #[test] + fn understated_source_still_decodes() { + // Byte passthrough over an element shape (Python fallback for dtypes + // without a production chain) declares fewer bytes than the tensor. + let raw: Vec = (0..64u8).collect(); + let buf = identity_container_with(&raw, None, |rec| set_source_shape(rec, &[16], 0x0005)); + assert_eq!( + ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap(), + raw + ); + } + + #[test] + fn decoded_length_must_equal_orig_size() { + let raw: Vec = (0..64u8).collect(); + let buf = identity_container_with(&raw, None, |rec| rec.orig_size = 63); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert!(err.contains("record says 63"), "{err}"); + } + + #[test] + fn chunk_table_beyond_bound_is_rejected_before_decode() { + let raw: Vec = (0..2048u32).map(|i| (i % 251) as u8).collect(); + let clean = identity_container_with(&raw, Some(256), |rec| { + assert!(rec.terminals[0].chunk_table.is_some()); + }); + assert_eq!( + ContainerReader::open(&clean) + .unwrap() + .decode_tensor("t") + .unwrap(), + raw + ); + // One chunk claims u32::MAX decoded bytes for a 2 KiB tensor. + let buf = identity_container_with(&raw, Some(256), |rec| { + rec.terminals[0].chunk_table.as_mut().unwrap()[0].decoded_size = u32::MAX; + }); + let t0 = std::time::Instant::now(); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert!(err.contains("chunk table decodes to more than"), "{err}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(1)); + } } diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index 2abde55..b510cb5 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -48,11 +48,19 @@ pub fn is_nibble_packed_dtype(dtype_code: u16) -> bool { /// counts. Element width describes the *value* width for such a dtype, so /// deriving the byte count from it would halve the length and leave the /// decompressor rebuilding every packed plane at half its size. -fn storage_bytes(dtype_code: u16, n_elements: u64) -> u64 { +/// +/// Returns `None` when the byte count does not fit in `u64`. +pub fn storage_bytes(dtype_code: u16, n_elements: u64) -> Option { + n_elements.checked_mul(storage_bytes_per_element(dtype_code)) +} + +/// Storage bytes of one element of a chain-internal dtype. Packed FP4 is +/// one byte per element (two values). +pub fn storage_bytes_per_element(dtype_code: u16) -> u64 { if is_nibble_packed_dtype(dtype_code) { - return n_elements; + return 1; } - n_elements * (element_width_for(dtype_code).bits_per_element() as u64 / 8) + element_width_for(dtype_code).bits_per_element() as u64 / 8 } /// Mandatory entry node of every PPG chain. Takes no inputs (the runtime @@ -89,17 +97,34 @@ impl Source { } /// Build the `PlaneDescriptor` for the raw tensor plane this source emits. - fn output_descriptor(&self) -> PlaneDescriptor { - // product() on an empty iterator yields 1 (multiplicative - // identity), so an empty shape represents a scalar (1 element). - let n_elements: u64 = self.shape.iter().map(|&d| d as u64).product(); - let length_bytes = storage_bytes(self.dtype_code, n_elements); + /// + /// The element count and byte length use checked arithmetic: a shape + /// read from a corrupt file can hold up to 255 dimensions of `u32::MAX`, + /// which overflows `u64`. + fn output_descriptor(&self) -> Result { + // An empty shape represents a scalar (1 element). + let n_elements = self + .shape + .iter() + .try_fold(1u64, |acc, &d| acc.checked_mul(d as u64)) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "Source: element count of shape {:?} overflows u64", + self.shape + )) + })?; + let length_bytes = storage_bytes(self.dtype_code, n_elements).ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "Source: byte length of {n_elements} elements of dtype {:#06x} overflows u64", + self.dtype_code + )) + })?; let layout = if let Some(&last) = self.shape.last() { Layout::Rows { row_len: last } } else { Layout::Flat }; - PlaneDescriptor { + Ok(PlaneDescriptor { role: Role::Raw, element_width: element_width_for(self.dtype_code), length_bytes, @@ -108,7 +133,7 @@ impl Source { residual_of: None, is_nibble_packed: is_nibble_packed_dtype(self.dtype_code), vendor_bytes: vec![], - } + }) } } @@ -129,7 +154,7 @@ impl Op for Source { self.shape.len() ))); } - Ok(vec![self.output_descriptor()]) + Ok(vec![self.output_descriptor()?]) } fn forward(&self, inputs: &[Plane]) -> Result, PtwmCoreError> { @@ -362,4 +387,22 @@ mod tests { assert!(!d.is_nibble_packed, "dtype 0x{dtype_code:04X}"); } } + + #[test] + fn overflowing_shape_is_an_error_not_a_panic() { + // 3 × u32::MAX elements overflow u64 in the product. + let src = Source { + shape: vec![u32::MAX; 3], + dtype_code: 0x0004, + }; + let err = src.propagate_descriptors(&[]).unwrap_err().to_string(); + assert!(err.contains("overflows u64"), "{err}"); + // The element count fits but the fp64 byte count does not. + let src = Source { + shape: vec![u32::MAX, u32::MAX, 2], + dtype_code: 0x0004, + }; + let err = src.propagate_descriptors(&[]).unwrap_err().to_string(); + assert!(err.contains("overflows u64"), "{err}"); + } } diff --git a/crates/ptwm-core/tests/fixtures/production_chains.bin b/crates/ptwm-core/tests/fixtures/production_chains.bin new file mode 100644 index 0000000000000000000000000000000000000000..e432e61e3e1446f291736193bb4a33fbcea7b267 GIT binary patch literal 249646 zcmeHw4RBp$oqq03lu-F7&{$BUp+#COgfvM>K}Hgw*p4avNRf|ly(HeK~Y+5hvt=e*~A zfA4wEx#zao$UxuET;O~ae7tzLP}x=X6VBil9&S4Vda zUpca|dhOO7TZTtA-?VY;l7S1m`}-=Fm)hFOZEbDu+y3TPzP@((ORo<1tQPu&-)LANbt|W5p3j?Re44m4o)|%Jg-HErB5n*A==BsvW z9ldH8=lvZUw+&zQj9 zr(S;Xx%WKura$}spIn?e-A_!-Ox{|Zu1_{*<|c8f-}~-rqc%G}dF#UXkBNVO<&M53 zwflbi?3(j-e`xrE2TpqGWA8utBb685aiDz4BfB48lRlvOJdR+dv8&do)@!xex=R-3 zNPhRw-8Vk`-0t-c4SxI=2hRKa*_XfOJCAH-#qdBFIMe7@XpaYXV)+L>dC+UMCIil|DgLnE=Y|h z=4SWQ8gBTg;Ohd{(gmQEa5|FLF^(S1Qu^_uEZV2Slsgy@H z?;NW<0v^z3uUy<9%I?G!ySj^=w8vvAWO|-{on3BSR|m5_WlQUDfLSjW>LRq#r2!+ z>vKqcW2%O$o3Gt@^~RB%!()|=Eapkl23f6rMQV~!-#cI3Rhyb0U#NW+<)xS(sjWLsAEx(IU`sM;iqzYTl?8LD(kIm7MUrqlGoy|CLCn8ES8zq8y=6;21RO#$f6G-wYHA-vXvU? zSDQ(VY#h@D$(}oP3o8p7m7A!Is*Uu}&D3^o&N(!0tGCFBln74rru24Zd?F3>_02R# zwwHS=T-54|#gh*y@}qk?mEXxr*1V-dP|K)j#+bqqB){C4j0mSt#MX!-i)M)=E&X6WS!F^5?Sh~ z?iSz5ah07nENLRUz-%r)y~eYW8{`~IlqbS-;J`Un-94htvB9ZH6Y-Jz*~yw04NXLQ z4%YZ6+H>$A+6#QncJsFg2sj?zSdS%Z{58Ke-X z@q;XU4kD0Vc#wki>KFzoec1Fb-Gzg%=#7E_z1gN#{gH}r(_7{8V`FlLcYR|rNTDFh z8b|OTn)3%K`RiQPAf<0i1}OyN=m$BqI-n5Q`Gb^v*<}q<`b-w9G_wxS_$b12@PP30 z2PygD%NnHMJgd`bMSBj`qCE-*(OzhflCQn2K}z443{t{rY_jq>0)+aS4^r|Km^Dc0 zmqCgHwEig4a}a^_!h;lSfjWjkO1}(J9B`wLqBjQ;=q-DYa;%z@L5hR0{wSJr5FwiL z2PyfQ%NnHg%OJ%8IQb~DOQ1jsy}Q);gOq&PWermLWsoB6MAK&tL3sItlzj1J4O04L zkfI&2Q{wugEP4(iM0=q@O1}281}XhANXdAxn5d7culXP)Ux8VJlmQu}I6%w7=YxTT zFFZ)W!-$SykTM{H6dz%^TIM1`x(oFO?v9jW)tn4cT!__YIlVfF5Y736lzh!)4N?YV zkm3NGd=%LwP&65&|RQhBvnFTKkFrd1+NDGIq96+f|>OpJ`MN%rq7(B;^=lV*`w2a%n(B?W&WXQ}fZ;)myCRjKUas1%7?)IcX+M0iGcHs^B) zzv{-j@OjpIwqaU)(H%abJ%H-GM0*+xEP7UZfh{TcOF0YQ>U5MgUHHr$KG)#}4Rplv z9LQLKqCTU(;|!Zk#>hAB+*sWy7Z5*qtfqJ36MHb#&UDGdK z^tp?^^yx2rcKxBRJ^NqJ{OWgSoN>mlZu{i$m(Tg?8&19Z{M6`R;yv@V>g@R3-29%# z%)WQ;S(xcTxl}q456yiI0BLV~t*LDMghv*aI+N=`VJv$q7YdkKw$T@Wd2$S9lpoqE zpo0Ra1E+e{kprkY7o`AAGR?yZ)!m(l)CHAbs+yV$;!vkta^wky^w4WvuKYaTayO@CkY z?IT;S93H8TUNt-dH?>+2rHCpwrTQ8?nvb}t6WYyI=@f{)m9>H*MN!!(g<@o*3L`lv zNIEzs{wYFIEuKD-=IN0sO={yLvUJgmjZ*}sGEP32rt?gJ89SznrOF*V;}g4r(>f*} zQnL>>1!-iJ1W^Z2W0fLAjaBjyH7$J#QR9*NMWqg?u}4yvQjg^0+7(;-Yc^~g*wSCU zcJq~2Du47C(8N-)$kV3#A2@R2!JlnE^0_siI&|mRA1*z->+&a_y5;I0uD|5qhP!_{ z_S?6lr%m5w z;hsBgzVw5CIsKXaH-33#>R9&HCnlz9)w$WJ`uxHiLwmV=Qv68o_O_+Ur>Tjc@oV~X z8UCYO>L|TF$ZK!qTSBaY%y}^(c4 zWwe|^la7{KMpQ=3Ify2v+|j7Hnw+Uo88uTx@l$Hl?90^bFU|y2M$RcvHLU7dG95Yl z(3(rhnb6AUISnlt<`d87tORCnXeP48@W}|;ML2T=?JM^J(-3c4cS&`4WZS0U>gdj4 zjG?R7Zr!nEcx3ZU8%>VG4eq0TN_s%q?$EWJ=U7(}E!>7)`>Tn{kx6Nr8L*cTT4g)at zr4zd%OXbcr1)z@B2|?P`a`c%XZB~m6q?0M067{C*uDPZfd5?+op{>He-8U`m}!;i_`wvhrf5no*z$Nzxw8h=YO$k_knkg-Z{H| z*;h~g^(QJX|M&;p|8YU;v_Chyr`D)Xj!*qBob#6j&-qb-(vyDD8DyaxNnSvPP~8_! z_ZDPZeI<3ygnTK;*2H43shHTtbbW7qZu;onJ>!jC`1$X1EKvs6!*Rx%pGf$kyvd2E~ z!bAW0_&G=a{J&SMN!f}CXOmmKR;%^*EzDeF@$cnrFL#`bi*u1>R8(U9FJbPf`HxQq zYOK*wE@6{f8Ly6VJH{VZZ5rNuZS|UM3AeqKLqe`oD;%;>K<;TVxz7C&$4Ky?5KQQ0 zN3gOMg)2DOt!gV+&p@PzDakBWXj{a|-6TbdIJp_mavTqfo{^=R(9DA6u6U`ogR{cb zygYqAjx4MYO?YOISB7)4CED!<3mGqaks~o)X8ib%QqzQM7BziVT+?(ETez;1oy*Az zUfq1{&Z{?$>>M7eY{jHT?NE?rD&l?d3IE1hdp`92@4V@ce{;^*gP)wc`{^@RU9{%o zlh<#1^kyT?3qt>nyZvBRUw0wXuFBp-c2u<+nwIYi8mHr zZ{=O0JW*cO<}s{(U0DN`1yJ9hsE=rm!+eF;x@XnrO4hb00p?ZLI?lh*b7PH1NeKL+ApyLNw3>`!u!tgi- zE7x(2WBSHq9OK}uKZ^1kM4&v|5UbxOTlw(e_|+%l7zb#jJ5mzl>uX$VG%c z7!YB29D|Kn$2E@WmvM}bwU|K{5z>#SKXAt}$FDva$GBjt&vNp05FzUG$1z3Kmo<** zmvM{(b@EXJm_X5F98*+)S>u>~8OKNi(x_D@r^e^zCUf?u+D3Q?+U&;rSkA&( zV;rcHk0QVXiYDWjq5{ks#|+3gM#eG6KaMFX!>n=4fQ(~g9CO^`n4&7o8pjOCI7Y@X z&7W~pXdH7KTHLl{N0pRBrJPuD4>Yic(h z>K@P+E4t?ORG|Wyt9qkuKMKsOrv^Wrmo_0I-Ys^2*C82m)`W@SY}H&x%Ds854!zli!Y7DRpCJf^7n&PzSmqDt33 zeb16Ce^z|zlL%1bK?LaCt%Sc&X!-LV$cU%dc>;_hWF)v!L!NvQ2^tC7?9?IrsvGZP zI177EtW3)=x@SpL2zZ^Bs8EA}h0v-nussHUsnBA$y4l&-+_O~0Ehp%!<bBw0(Vd%iY`yO7o5m`up=hXZktKLn*Yt}QeeR+!efkTZ zU4Q6n&;Hjlzxv%7XPoh?+det`<#WFJhEwl8KQ(}uc+Y&TI*Tor^LrXI``*20VP+8J zQt3oIzFyU2d)o?A?f9vLf9xfvAe6n8wSuMAZ}bINo*ZKtJ@ckTJO>3&2TfF|i>Ep} z`FPsk_@F*@3n&tGF(r_y&c}z;_SHd1q;`KCZR_3gT*vXTK=cffmSi#fqhmPVmvSPlMjDabQb-Y^cU6j{#1#Npx`^ zF?s(3M@~HWv+YMdx8_ra?mYX$rH6N2{=`$aT>ZoKmmJ)1_fN-u`ExQ&P69H z;eQ=G@;ARJ&@Qc??pv5vg$c>FT{^9ByW|kdrCkzc=^BDYn6^vD#4eY1N!ZO{SDl3O z*>&3`r*;{&OG0oC!Kz)#C)jP59D-T3OG0sq;&i)|PqW)DIW%)>mxSdUmQA~q&$8Pt zIV|&PmxSjO&*63{pKG&lU$>z;deyLO)XB?J!zDYrQRlVhnht2d{GiPgO%O=grbERF zBnCEiEV+^QSYrfkunXyPX#Pi206K+UeFAfU2ACk3d6)NUjh!e8jd z+lSTM&lZGLw&wsb?**Y+RE@dwQ8nA1f~Xo(C!2FzXn|{ZbB>Sgscp3bvr{vT+Ss}Z z64=HnYv2WBYqP-;I4Z~V8vNXjf0R!$UrMFUr5<>_l?w$u10BA{ry+%2wBw{Hf`fvf zgP(_>S|3vgp5bXwyCIRGiz1PsaziN)yF3tM2h=Oo?BEC()wx<#v)!GLf^wM$W#oDU zNe4)?DLF#YxSkY}ulJBNPG>+=>3|7UDyO5cq*vKfcoM4vx5LEhEdK3#WGZ2l*34lnf z4nSmOby&c7dR)!afb39KXPLD+Hlfj2of0f%XQ|ycJ%F*da*5!my*F}Y=$M{XupAUDHIiz^ z<6xQCpFAvy5rQtY5fW)ecoH}hBa{axu|$9~wnUdnt(#g`mPpmQXNmGaCFTgAM&?K` zO+;!jCFUp(Q(}{Vsj*3K6QvrZDx0K~8n8)u$P&ZU55u&Y1*9DAaKLm(T>^`42&Y4} zjkIr0G{eh3O=v?GWehQ<;3 z7!nf%3|VagVuGT&{H02n+8`fE;(UN4r+q){fONf&4 zMvSU^Nfo1x4J^j$ps;EQ)t^|hnt0PNtHh$BR<%Wya;40&BsMY3VPc7I2C=bk{;sUV zgj+_`s(VJGtye=h#BRLxR3l6;G4r4|GV{VNP;*`^*(I$-m|fCgfL+sJ9FppF*llZL z#IL$XG-I=0m}AoRte3Wj5UYMz{HyYaS~(fC+6{=f1I5|Zpn1h_N}aee?`=f{q0;vF zl(Jy-w3t!S@DyiN+8*YtIt@=YtEBBgtxDVDv&v-ww&m(HJlVvOwg<#Y+XKs(ez91I zyW*ADyKIqnjwNjm=#{o7%`Fq>WZCr^p6uc!Z4cO$wr3yT>OCv2-oOQ(yb305&n0Mk zo|K-J)G&voxkI%*b~CbSd+aVHuk9(TZOP@-wq#J-TX|Zz<;SgZRCI&%G4eYo{1SF~ zoOnAdzium3pngeHg!2o%Td(f4JYCM>fcSsCc`=JtYX zHEEPUbkryb&xB?cBjygvv)?EcC}`3$foIb)9hQnFRI{jP)iQ-|0w`G3qzOC8)l{ zbB?G|NpFNo)xD7vDT7v*x}ywUNkar)QiXfM0JTBW_34&_TP9bP*bEH9Sn|UNOvIJ2X))a`= zqedXc9#yhy-Fk)=(<#6yD6Jukpo|QmAeo7zv0^DCt^10QH111MfITe3DZnW_iPb^W zprioDIFb}#1)H;f3@*n@0fra>9v7>#_*Wf(9!&u*h#(iM!wf-VasnD8RtFknWObN< zXiQE3L}GOSA}g!I0>*R-a7INFs{<%GSsi9bx+W(eN@8_@BrmJO0?Tj;aE3C8)j?dK zqyQ&Y$4LR!;MpA@#_EJpfMI$7qoe>E=S`wk)h}DdQc{3}SSI!-PnX09L6_PHNeXZT zD3NAjg!14dmI!dhmPit76F_ZZ?GS#~jd#;}C*}yCM&?K`1$GL)QisKqn4>(Q5}O1} zjZKmi-~`oh3UE`k62pYRK}i9Q5Yt;U?(ZS33jM+NrxTyb_PWpZw2Ut_X1w614e53XCNap<}D-}&OLyK3{b#&mseeQx^b-aX@uUHiuy zwd&;5_}tv&c%9%JHAQ?Fy0WyNo z1Jd{lQuz?wC zOMEkkjeRqZoLHgOP-?z|+H?wV&T34|Jm`(gys&HVjJ`8+PL^HLT7-+2bQoaQbQnBR zfTen=%5btQm%=H)ITcLW9z>o=+vB0S_}d#u3UExNvrEKwJlKNlrU0X2L8!Dn9-%Bk zH!v)Va-E|T;FejHwud>ZPQ#O3tEBBgtxDVDvx=WtEWoy0orWiySkm@@SZRCoc+O{+ z1+_LOb_mDUjd#=4n6y2hSK1z+mU=N`*K2sPiajRJgTgi1xCR603fF9Q<-g0TzBU_;p*M z0`*IpBGj*Lill@sxat#&<(Z}cw@|{QJp#pXd&DChR_fSF4ePW`haJN2yYX&;XwoQw z=%`T=o`ItCVo}VNXTMP@ux68%2|SyYiAM^sRJ0WvDmt73+(K28<_W=POj3aP(9B_5 zGf4rC$%PE`9L~!YB|GK7CzljpA=K8by1HAo!J_I9S|gQ93b2rxNvhK*;qQbhmGnla zRNWi#NC6gJGkJ9`8D;QF8Y1u-H$*&wVue|cSsN5Pgg@!VJAWxAT@cufx*#DJz&S4_ zrpc1)^+7n)(2`7AACPNWA08>dQoD}Al!XeX0O!}Q?RU`HN(!)xsQQd;XEp`+R8u7Q z1KbWm7K6uvKt2MUYE09!mt(qMbs}+U7`u+BQoHuu_|O5g97OngY>! z)Ck1bqe^zIflz%1jOi5M6qMEwMo>nEP>{?-(pa$+lGc4iNE-JgDZn0<;S}H$p2X@P zYEV*uV;qrhumn*hOYmd6UK^eLVqyQTfi3p7&@-ZYP2pGmDNWyC!6w@idp>ih92S`TF zM}Q0h=?;c`NfNsQAdTIT6krF+a0+lpm&D`X{$Yv3GBC#9xorAC*4asjP*M-ij;_aH z4<=r8hn~Ov2!u);R!n5_jS-_X1z6W-F;*oGE1@&~z>?L(n}%5>78SLsE$TN>(vkvf z>19c5Vwl6k65p&+IR-B5?(eHePOPO>ePTqd&50etfpz1Zvl2BkBs;YrvT?vFll=bc_wWS<(5U>3Z|q{W^B^%IJ-o2?cz3; zQYW?`dv{%cP-%NSLU}N1V@1{B#RXcNRcU*ev+6WF*{qVb2em3~kIyQ8Ui|%KxOvc) ztJCmg6HD425G!qu9?$vgvYpCT$PsmA1!cR{hM3U9aKEE?(01fL&>O{@_x8 zn^rJsd$2J>Qh>b*=1{y&3a}y6ZdCNzo=^%fc*TDt1=ypO#lq~R03Qdxk^;<}U$+%1 zP`{)pLjCHdNK$|)$t;#qL?kuexp=i%_c1qcs4DQqyRfK<0A{E!zsWmR5fXy5PT*nz?9`SCV_Tpn_MRj zExBrh&w37*nF6eqXw@-Euqb7e386FvIIh@&q)G}frL>Jpm0D7rMk#|-(i@>tb#Ela zN||L>tj;B)3|>h?1YYBYNP=QnGHZikhj4t|c;_$0qzeMOQ5PiS0yyWzq--s@ULS-* z4Vfp>`hZ;1`bY|}!!15Cp*NfYoL|2--EAH2vyvEKN~{+H{902bx?Zp$?Xvjr4dk&6OsI9Xk0V}jY4w zjt3~Ray%?>OeX?o6g71`z><^WVFsn^bpo=ajt5Zkay%@+3}*soXp=e~qz6hUaN>BJ zP+*Oo-4J3Nj|~N;(;na`sldj9ljv3T%hth^Sl}R-i3`e8CiO#5ruIXU3mhRz#F_Y^ zJUpp00-mumk`UVj(R46y(~_qi38+RMNk9#NN<31YRH<76sKzZxG;jiII2yRAUa4=w z$AwVSQs6raLTzH9IO4*T`IN7{PrG6Hbs{O1Kt4l1iVoe-%9T+K47-D z!sA&eLajrsjU`mSJRW)ng0fuQ` z1cu|j=o#@=4#_MywsDg~IJ9oOTUe-RhXkIZ4oPSRjLC}`&3=bepqy#11e#5+^p8@{ z3C%nvlO5qu(H3f&c1;LUsyr+#J1jHFPqLY!F}FBG@T9)PZBqyKJ$IC7s{}fGD~}1E z{P?u#lj_)7-E0d+y((1HD8MM+)2(-qX z(6?j_c9>;Ctc{!;!ohXpoxcv#_6Ouf?T>H^7m}?NqsJv%ixGV4rauMW z33`EP4k4FpEtU|fWjh6-bC6bCm>7ecmO3Y$HM~5aGGd_HmT!5;24vr2o6X( zX`&iEyAi}Vo=}>o@;<34AeHx72J7f}pO(Pn5*5J=%)|xdfl2)kl&SrYq=`n2 zO)5>aX{A$V1UzGBB*B*nqUkizrl?Yn1XLrBB%lUBB_1iS^r>3{sKzbvh>BpbnugOv zo9dPNCPbD=nrNg{e8sTCrim&f@ru9rX`;$30leOd#5fxGNQ|TUWg!G5P1N8QND~bq zNNo``sBMwR(7;G!Xa+_2LsH1p3jxH~3rPm117kW(G^9srfBdv3P1kTZAvmXwX)lCOH0gy>z%Kq+pQHrGY&x4s zY#ZdI1gmK@IF()~#VL;zLz)!aI4kz*DHbcSEn24w%I22# zLf}?=4$LZoHc=%d{5)!_o_dhsw+d_PU_#5~jTn2$o(b>$(98io@x_ zP3xHULI^!$l0nL;*X6*klOSw3wVTGgUMQ3xthyqw+gst0S}645B%8bt5KU8rkB8zx zp*R87d|*+*mSVR*Dp19=Q$iK%PDu*cuxq|EW7+ML3b0IjCa@g$Ogz#@Wx+NoQ{S;* zrxS%+ShQ*P1g@j*NvL)~R-ajvw58hbo(hyT?V~`o>7#xm*OWZAS&$u07H*-oX-9=f zH72R0oK|oqX?ZqXIOZ3E)HB%3bYay-fzsXzkMu*~lrAB-W@uwe)gQIU>!9#T0mV;D zYSr@UbWb?EP`%QI3H7QQCMj8;TXxhsEmH=yv`YfDahJp+uT<7!Dlx3b>9pbeHJP?X zAUJA`gkOp*n|0ap>orFi1x))P@N4>^r=%b9C}y%ehf|01tJpR{{fHcv)M1ZVm9lA5 zhfgymMwbjp9gflClEKGiV=oDNRoFI%kV^(1O9<7{or2K1=val?MN8_i(wlh^Au9Kp zLeU!57{%DIN>r}FP<;oE>JDoaBmTPb5WaTfox;+3!x+oR8w!*FD|rzpjVDV%X>C{x zrLkd>I-DSh4*)ZqI-CNOIv#`?O6qV7Btj@`>afDi*>q+y`0TCBN{&E}i{n}RYZ8Eu zrVbZ{kc;DC2BGmf0Toim0~IoIJj_ruUMGMebv!_kmE&Q7V>)#>qo}Fl0hXK`4>KrT zuM?0Zbv%HQm*ZgpW;k^?Lz~p`Aec~6hZD!+qz-HJ>_!pectWYe%KHFENgXy8oCLe7 zU$zdWqz(tcOk7Z&GN~VeGPNI))ZqwGBF@AQ<>5)45%7$ikp$!>h^AAAo0dHFNI*66 zNCIjARN|5Hq)OcqKs9bjQil^*!>Pkf^-6scq75Z=ID!;kG3>CZ!wN~f;xB&curf;k zPf~{y@DzILVVs@#P4tk4LI_Iga0Y_LBl#Mnwg?*3wn$Qk4U9yF#u)h^QZED$V=p9G zxeko!)ZtJ?Qv(DjBLgHz29b0vkS|K=dVr*HJ(4=?U>Qyw4k?pb9R#W|N!t3i>jz#P zTG8?J6*oNd;EyT?NALRByYBty)|a+^_iKm$b@0gF{H8$a@cQY#|HG-nn#^!W9=%it zPPrs)38|(=?e|pba4W=eN!k)(EurQUAvQ6p;fkex6&0)fDv#7*;npx}zY}f~{~YF) z+Guba+h`u4w8E}s)qD!34(BY()ZK&O$lVLU22k^z%g$OtFzGhJB}|(R5Nw(b9;w4p z!$x5CJ17pP4(HS{?S&AiCcTgaH2$>sn;=N)aLlH&nM7A`xeXWbu`{&PNqZr1D!ovO zQx+K>WgB3s4zP7%r59pPtkVT$S1j#?P_fbr>FaAgyDU~>TeMCWl+7*eg}|-!LVsAP z!%dfE+6#eU>4pBFQiq!oOnV^^EWOYlVCryFnrSbDP&6i~!<>3u4rgE|b=YufH;s9{ zP$+d+?a=|dTvCUHUf1TOQiqR+VlJt}LUAU=ZhusuifN~WD%PD6kJMpdIg@3#Q!2nR z?U}%G+%xfrq7|M!o^2HE5dNea?-mwq+C72msCyEs0V8>_uwIsG?;296Kv~m13RIgu zibv|O)U~TO7_o5ba0|6fJ1WGgF-aZf#59vUHJdsd^9wD+_c+8YRR?zRcOVd_c3^gD zrcoPHZ4@Zwk~%D$y49R*x5z7()M4Q@hgYY2%BWY`Fri*`!^9(XSg6gR)@hkC zsHI&JsExZM9;viKZi-x+N;`x<>&82OO{T382##7K;TOn}7ZVR<`SqG39Bjx-(S8X0 zntq5!>abKW5t(vS;nd;$Dt4M6(&+Z;<`8np;A07)TDnsZIxad=p?1-dI_xfxsADTc3G8k%g7rFl$lT(PnLqxv0)LE#)e7ium@&1bvOkmbv(e7iQ`%PD<}w|u&Kie zH)r3d$>4J)tbk4{lE;aL1$;DhxG02N91k-Hjn@gNkUAbE#NJ*m)|(lM#_I%7q>cwD zvT{5ua7?ETXB0JcJiwBZ<6#D+>vaOMq>cwr@^U;ZzznAjXK0f;9t0Cg>Tqlwam`pc zo@My4U3Y}eMiJw9f~mt$C2*9~VPnBbu&er|mZIHtQvgg!9S(w-xS%{RsULzewI7ny z;RsP8&cqMp;Ypnl@Qj_21mq@&rc;NTmOS-HKsEA60%`zM;*s*CO5GAbHEu~#hZ9)C zsl!e6N_`We4JCCrf)rm#>aao*ulP$Uby!(=WtJrF)xbyMUezxPAtP}?Grp@EUe&=?~hMCyeAV(f(^E7yTBojM$E$ut$Yn&t z2*{-|kj{o6*FNTe$z?>u63oN|g~6m=2+CwMB4UYW;)lX`QdX3jcrx6iLP>D+lgG$X3Ks9EG$B2j#tl@;x9KBNCgz%-xBj5nr3)|ke{mrj@ zeeLp>ULEdPIdZsjW$nDBFR#As%%c-sw{IQzHiR+ixB z*OiCxwHxoIOELBLU^nvjLN35dUM!YlOD^d!@=BNX86em689W}mA+>81hT2u>IjVL~ zZCA~T>d)09P@C2-?ROB8RQax)aRHopaVT9Ufz~PGCVYeG9z%cELz;N6b z@kp1IMcN1ri?ofQ9KvtA@or(ErX3P^jyfcv889X<7Nu-y_B*5k)H0_!YrBr!X&ddg2d}or7Wm86DZgFVfNqvdirc}x(S|!lg%Op9LOVu*r zQ*AU-Qw?XMk>;@wr-Q<&l78}H$EnjH;UGht(xwP?s+*$6Wa&*vWwZ3!Qgzy*3|eU? z1X|-xh(~O!tie`Ym^0hd$szo{8}IycD{X&3Zq)t=w?Hm=F Vec { /// The P4 fixture: fp32 from `Format.BYTE` input. The v1.1 Python path /// builds the chain with `chain_shape = [len(bytes)]`, so the fp32 Source /// declares four times the real byte count. +/// +/// The writer now skips chains whose Source does not declare the raw byte +/// count, so this compresses with the exact chain (`[len / 4]`) and then +/// rewrites the stored Source dimension to `len(bytes)`. The encoder never +/// read the Source size (it sizes the source plane from the raw bytes), and +/// the dimension is a fixed-width u32, so the result is byte-identical to +/// what v1.1 `compress_model` wrote for the overstated chain. fn p4_container() -> Vec { let bytes = p4_bytes(); + let n = bytes.len() as u64; let input = InputTensor { name: "p4.byte.fp32".to_string(), - candidate_chains: vec![fp32_split_chain(&[bytes.len() as u64])], + candidate_chains: vec![fp32_split_chain(&[n / 4])], dtype_code: 1, // Format.BYTE input_format: 1, @@ -671,7 +679,19 @@ fn p4_container() -> Vec { dtype_name: None, delta_reference_blake3: None, }; - compress(&[input], main_options(None)) + let mut blob = compress(&[input], main_options(None)); + let exact = fp32_split_chain(&[n / 4]).nodes[0].params.clone(); + let overstated = fp32_split_chain(&[n]).nodes[0].params.clone(); + assert_eq!(exact.len(), overstated.len()); + let at: Vec = blob + .windows(exact.len()) + .enumerate() + .filter(|(_, w)| *w == exact.as_slice()) + .map(|(i, _)| i) + .collect(); + assert_eq!(at.len(), 1, "Source params must occur once in the P4 blob"); + blob[at[0]..at[0] + exact.len()].copy_from_slice(&overstated); + blob } /// The P4 tensor in the v1.1.0 table layout, built-in ids pinned to 1.1.0. diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 88c1d25..5876a2c 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -60,6 +60,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_function(wrap_pyfunction!(py_builtin_canonical_id, m)?)?; m.add_function(wrap_pyfunction!(py_write_chain_blob, m)?)?; m.add_function(wrap_pyfunction!(py_chain_blob_info, m)?)?; + m.add_function(wrap_pyfunction!(py_chain_source_byte_len, m)?)?; trust::register(m)?; policy::register(m)?; host_flavor::register(m)?; @@ -952,6 +953,19 @@ fn py_chain_blob_info(blob: &[u8]) -> PyResult<(usize, Vec<(u16, u16)>)> { Ok((et_size, table)) } +/// Byte count the Source node of a chain blob declares (shape × dtype +/// storage size). Writers require it to equal the tensor's raw byte count; +/// tests use it to check the Python chain builders against the Rust rule. +#[pyfunction] +#[pyo3(name = "chain_source_byte_len")] +fn py_chain_source_byte_len(blob: &[u8]) -> PyResult { + let (chain, _) = ptwm_core::chain::wire::read_chain_blob(blob) + .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?; + let desc = ptwm_core::chain::runtime::chain_source_descriptor(&chain) + .map_err(|e| pyo3::exceptions::PyValueError::new_err(e.to_string()))?; + Ok(desc.length_bytes) +} + /// Parse a chain from the **legacy wire format**, where node op-ids were /// stored as raw `OpId` discriminant u16 values rather than extension table /// indices. Used only by the [`py_write_chain_blob`] migration helper. diff --git a/docs/architecture/ppg.md b/docs/architecture/ppg.md index 06fc0d2..9a24f69 100644 --- a/docs/architecture/ppg.md +++ b/docs/architecture/ppg.md @@ -68,9 +68,17 @@ values. Adding an eighteenth is a wire-format change. The current catalogue covers: The control ops `Source` and `Terminal` mark chain boundaries. The -source carries the input dtype and tensor shape. Each terminal -carries the role bytes that the codec dispatcher uses to score -candidate codecs. +source carries a chain-internal dtype code and a shape that counts +elements of that dtype. The shape's element count times the dtype's +storage size must equal the tensor's raw byte count: the decoder +derives every plane descriptor from the Source alone, by walking +`propagate_descriptors` through the chain without touching tensor +bytes. Writers skip candidate chains whose Source declares another +size. Readers of the legacy container still accept files whose Source +overstates the size by up to 8× (byte input written by v1.1 and +earlier) or understates it (byte passthrough over an element shape). +Each terminal carries the role bytes that the codec dispatcher uses +to score candidate codecs. The bit reorder family (`BitReorderIeee16`, `BitReorderIeee32`, `BitReorderFp8E4M3`, `BitReorderFp8E5M2`) does the per-element bit diff --git a/python/ptwm/core/_compressor.py b/python/ptwm/core/_compressor.py index 17d5ddd..bdcca53 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -8,8 +8,8 @@ from .._exceptions import CompressionMethodNotSupportedError from ..delta import encode as delta_encode from ..preprocessing._chains import ( - CHAIN_BYTE_PASSTHROUGH_VALUE, ClassifierRole, + build_chain_blobs, chains_for, ) from ..utils import DType @@ -25,15 +25,6 @@ Method.IDENTITY.value: 5, } -# Dtypes that are packed into uint8 on the Rust side (single byte per element). -_UINT8_VIEW_DTYPE_CODES = frozenset( - { - DType.FLOAT8_E4M3FN.code, - DType.FLOAT8_E5M2.code, - DType.FLOAT4_E2M1FN_X2.code, - } -) - def _to_raw_bytes( data: bytes | bytearray | memoryview | np.ndarray | torch.Tensor, @@ -115,42 +106,6 @@ def compress( msg = "delta compression required but delta_second_data is None" raise ValueError(msg) - delta_ref_blake3: bytes | None = None - if delta_second_data is not None: - ref_bytes = _to_raw_bytes(delta_second_data) - delta_ref_blake3 = _rust.blake3_hash(ref_bytes) - data = delta_encode(_to_raw_bytes(data), reference=ref_bytes) - - return self._compress_torch_numpy_byte(data, delta_ref_blake3=delta_ref_blake3) - - def _compress_torch_numpy_byte( - self, - data: bytes | bytearray | np.ndarray | torch.Tensor, - *, - delta_ref_blake3: bytes | None = None, - ) -> bytes: - if self.config.input_format == Format.BYTE: - dtype_enum = DType.from_dtype(self.config.bytearray_dtype).code - shape: tuple[int, ...] | None = None - else: - dtype_enum = DType.from_dtype(data.dtype).code - shape = tuple(data.shape) - - # BF16 in TORCH: numpy cannot hold bfloat16, so reinterpret as uint16 - # before going through numpy/memoryview. - if ( - dtype_enum == DType.BFLOAT16.code - and self.config.input_format == Format.TORCH - ): - data = data.view(torch.uint16) - # FP8 / packed FP4 on TORCH / NUMPY input: reinterpret as uint8 to - # hand a flat byte buffer to the Rust compressor. - elif ( - dtype_enum in _UINT8_VIEW_DTYPE_CODES - and self.config.input_format != Format.BYTE - ): - data = data.view(torch.uint8) - if self.config.input_format not in ( Format.TORCH, Format.NUMPY, @@ -158,12 +113,28 @@ def _compress_torch_numpy_byte( ): msg = f"Unsupported input_format: {self.config.input_format}" raise ValueError(msg) + + # Read dtype and shape from the input before any delta encoding: + # the delta residual is a flat byte buffer, but the record still + # describes the original tensor. + if self.config.input_format == Format.BYTE: + dtype_code = DType.from_dtype(self.config.bytearray_dtype).code + shape: tuple[int, ...] | None = None + else: + dtype_code = DType.from_dtype(data.dtype).code + shape = tuple(data.shape) ba_bytes = _to_raw_bytes(data) + delta_ref_blake3: bytes | None = None + if delta_second_data is not None: + ref_bytes = _to_raw_bytes(delta_second_data) + delta_ref_blake3 = _rust.blake3_hash(ref_bytes) + ba_bytes = delta_encode(ba_bytes, reference=ref_bytes) + return self._compress_bin( ba=ba_bytes, shape=shape, - dtype_code=dtype_enum, + dtype_code=dtype_code, delta_ref_blake3=delta_ref_blake3, ) @@ -194,13 +165,16 @@ def _compress_bin( ) tensor_shape = list(shape) if shape is not None else None - builders = list(chains_for(int(dtype_code), ClassifierRole.STANDARD)) - if not builders: - builders = [CHAIN_BYTE_PASSTHROUGH_VALUE] - chain_shape = tensor_shape if tensor_shape is not None else [len(ba)] - # Builders may decline a shape by returning None (e.g. the - # spherical-normalize chain only accepts 2D tensors); skip those. - chains = [[c.to_bytes() for b in builders if (c := b(chain_shape)) is not None]] + # Each chain's Source shape counts elements of that chain's own + # Source dtype and covers exactly len(ba) bytes; byte input has no + # shape, so the element count comes from the byte count. + chains = [ + build_chain_blobs( + chains_for(int(dtype_code), ClassifierRole.STANDARD), + shape=tensor_shape, + nbytes=len(ba), + ) + ] codec_menu_ids: list[int] | None if self.config.codec_menu is None: diff --git a/python/ptwm/integrations/_compress_safetensors.py b/python/ptwm/integrations/_compress_safetensors.py index 6454649..d5c71bd 100644 --- a/python/ptwm/integrations/_compress_safetensors.py +++ b/python/ptwm/integrations/_compress_safetensors.py @@ -26,6 +26,7 @@ from ..preprocessing._chains import ( CHAIN_BYTE_PASSTHROUGH_VALUE, ClassifierRole, + build_chain_blobs, chains_for, ) from ..preprocessing._explorer import ExploreOptions, explore_chains @@ -108,7 +109,7 @@ def _chain_bytes_for( shape: tuple[int, ...], explore_options: ExploreOptions | None = None, *, - raw: bytes = b"", + raw: bytes, builder_cache: dict[tuple[int, TensorRole, int], list] | None = None, warned_passthrough: set[tuple[int, TensorRole]] | None = None, use_user_cache: bool = True, @@ -171,9 +172,10 @@ def _chain_bytes_for( if builder_cache is not None: builder_cache[cache_key] = builders - # Builders may decline a shape by returning None (e.g. spherical-normalize - # only accepts 2D tensors); skip those. - return [c.to_bytes() for b in builders if (c := b(list(shape))) is not None] + # Each chain's Source covers exactly len(raw) bytes in units of its own + # Source dtype; builders that decline the shape (e.g. spherical-normalize + # only accepts 2D tensors) or cannot cover the bytes are skipped. + return build_chain_blobs(builders, shape=list(shape), nbytes=len(raw)) def compress_safetensors_file( diff --git a/python/ptwm/preprocessing/_chains.py b/python/ptwm/preprocessing/_chains.py index 0bdb34d..9b24415 100644 --- a/python/ptwm/preprocessing/_chains.py +++ b/python/ptwm/preprocessing/_chains.py @@ -16,16 +16,20 @@ ``list[ChainBuilder]`` table * :func:`chains_for` — thin wrapper over :data:`PRODUCTION_CHAINS`. -Chain builders are callables ``(shape: list[int]) -> Chain``. The compressor -driver calls them per-tensor so each chain instance carries the tensor's -concrete shape in its Source node params, which the decoder's placeholder- -forward descriptor recovery path requires. +Chain builders are callables ``(shape: list[int]) -> Chain``. The shape +counts elements of the builder's *Source dtype*, and the shape's element +count times that dtype's storage size must equal the raw byte count of the +tensor: the decoder rebuilds every plane descriptor from the Source node +alone, and the Rust writer skips any candidate chain whose Source declares +another size. Call builders through :func:`build_chain_blobs`, which derives +that shape from the tensor shape and byte count. """ from __future__ import annotations +import math import struct -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass, field from enum import IntEnum @@ -96,6 +100,31 @@ class _Op(IntEnum): } +# Storage bytes per element of a chain-internal dtype. Mirrors +# ``storage_bytes_per_element`` in crates/ptwm-core/src/transforms/source.rs: +# 16-bit, 32-bit and 64-bit codes; everything else, packed FP4 (0x001F, one +# byte holds two values) included, is one byte. +_CHAIN_STORAGE_BYTES: dict[int, int] = { + 0x0002: 2, + 0x000F: 2, + 0x0007: 2, + 0x0008: 2, + 0x0003: 4, + 0x0009: 4, + 0x000A: 4, + 0x0004: 8, + 0x000B: 8, + 0x000C: 8, +} + +_MAX_SOURCE_DIM = 0xFFFF_FFFF + + +def chain_storage_bytes(internal_dtype: int) -> int: + """Storage bytes of one element of chain-internal dtype ``internal_dtype``.""" + return _CHAIN_STORAGE_BYTES.get(internal_dtype, 1) + + # --------------------------------------------------------------------------- # Role wire helpers (from types/role.rs) # --------------------------------------------------------------------------- @@ -962,6 +991,94 @@ def _builder(shape: list[int]) -> Chain | None: """FP16 raw-word chain feeding the huff_llm_5bit codec.""" +# --------------------------------------------------------------------------- +# Source shape derivation +# --------------------------------------------------------------------------- + + +def _source_dtype(chain: Chain) -> int: + """Chain-internal dtype code of the chain's Source node (its last 2 bytes).""" + params = chain.nodes[0].params + return struct.unpack(" list[int] | None: + """Shape, in elements of ``storage`` bytes each, that covers ``nbytes``. + + - ``shape is None`` (flat byte input): ``[nbytes // storage]``. + - The shape already covers ``nbytes``: the shape, unchanged. + - Otherwise the last dimension is rescaled so that the rank stays the + same and the leading dimensions still count rows; a scalar becomes + ``[nbytes // storage]``. + + Returns ``None`` when no such shape exists (the byte count does not + divide into whole elements or rows) or a dimension exceeds ``u32``. + """ + dims = None if shape is None else [int(d) for d in shape] + if dims is not None and math.prod(dims) * storage == nbytes: + out = dims + elif not dims: + if nbytes % storage: + return None + out = [nbytes // storage] + else: + rows = math.prod(dims[:-1]) + if rows == 0 or nbytes % rows or (nbytes // rows) % storage: + return None + out = [*dims[:-1], nbytes // rows // storage] + if any(d > _MAX_SOURCE_DIM for d in out): + return None + return out + + +def build_chain_blobs( + builders: Sequence[ChainBuilder], + *, + shape: Sequence[int] | None, + nbytes: int, +) -> list[bytes]: + """Build every candidate chain for a tensor and serialise it. + + ``shape`` is the tensor shape in elements of the tensor's own dtype, or + ``None`` for flat byte input; ``nbytes`` is the raw byte count. Each + builder's Source shape is derived from these so that it counts elements + of the builder's own Source dtype and covers exactly ``nbytes`` (see + :func:`source_shape_for`). A builder is skipped when it declines the + shape (returns ``None``) or when no Source shape fits the byte count, + for example 4094 bytes for an fp32 chain. + + When no builder remains, the result is the byte passthrough chain over + all ``nbytes``, which fits any buffer. + """ + probe_shape = list(shape) if shape is not None else [nbytes] + blobs: list[bytes] = [] + for builder in builders: + chain = builder(probe_shape) + if chain is None: + continue + storage = chain_storage_bytes(_source_dtype(chain)) + src_shape = source_shape_for(shape, nbytes, storage) + if src_shape is None: + continue + if src_shape != probe_shape: + chain = builder(src_shape) + if chain is None: + continue + blobs.append(chain.to_bytes()) + if not blobs: + src_shape = source_shape_for(shape, nbytes, 1) + if src_shape is None: + src_shape = [nbytes] + chain = CHAIN_BYTE_PASSTHROUGH_VALUE(src_shape) + if chain is None: # pragma: no cover - the passthrough never declines + msg = "byte passthrough chain declined its shape" + raise RuntimeError(msg) + blobs.append(chain.to_bytes()) + return blobs + + # --------------------------------------------------------------------------- # Production chains table # --------------------------------------------------------------------------- @@ -1073,5 +1190,8 @@ def chains_for(dtype_code: int, role: ClassifierRole) -> list[ChainBuilder]: "ClassifierRole", "PRODUCTION_CHAINS", "TerminalRef", + "build_chain_blobs", + "chain_storage_bytes", "chains_for", + "source_shape_for", ] diff --git a/python/ptwm/sharding/_writer.py b/python/ptwm/sharding/_writer.py index 1d2e030..405a415 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -8,8 +8,8 @@ from .. import _rust from ..preprocessing._chains import ( - CHAIN_BYTE_PASSTHROUGH_VALUE, ClassifierRole, + build_chain_blobs, chains_for, ) from ..utils import DType @@ -17,14 +17,19 @@ __all__ = ["write_ptwm_shard", "compress_ptwm_blob"] -def _default_chain_bytes(dtype_code: int, shape: tuple[int, ...]) -> list[bytes]: - """Return chain bytes for ``(dtype_code, STANDARD)``; fallback to passthrough.""" - builders = chains_for(dtype_code, ClassifierRole.STANDARD) - if not builders: - builders = [CHAIN_BYTE_PASSTHROUGH_VALUE] - # Builders may decline a shape by returning None (e.g. spherical-normalize - # only accepts 2D tensors); skip those. - return [c.to_bytes() for b in builders if (c := b(list(shape))) is not None] +def _default_chain_bytes( + dtype_code: int, shape: tuple[int, ...], nbytes: int +) -> list[bytes]: + """Return chain bytes for ``(dtype_code, STANDARD)``; fallback to passthrough. + + Each chain's Source covers exactly ``nbytes`` (see + :func:`~ptwm.preprocessing._chains.build_chain_blobs`). + """ + return build_chain_blobs( + chains_for(dtype_code, ClassifierRole.STANDARD), + shape=list(shape), + nbytes=nbytes, + ) def compress_ptwm_blob( @@ -59,7 +64,7 @@ def compress_ptwm_blob( None, # delta_ref_blake3 ) ) - auto_chains.append(_default_chain_bytes(int(dt.code), shape)) + auto_chains.append(_default_chain_bytes(int(dt.code), shape, len(raw))) effective_chains = ( chains_per_tensor if chains_per_tensor is not None else auto_chains diff --git a/tests/ptwm/preprocessing/test_production_chain_fixture.py b/tests/ptwm/preprocessing/test_production_chain_fixture.py new file mode 100644 index 0000000..023932c --- /dev/null +++ b/tests/ptwm/preprocessing/test_production_chain_fixture.py @@ -0,0 +1,93 @@ +"""Pin the chain-blob fixture the Rust descriptor tests read. + +``crates/ptwm-core/tests/fixtures/production_chains.bin`` holds every +production chain builder's output (and the explorer's first candidates) for +a set of shapes. The Rust test +``chain::runtime::tests::symbolic_descriptors_match_forward_on_production_chains`` +checks that the symbolic descriptor walk equals the forward pass on each of +them. This test keeps the fixture in step with the Python builders. + +Regenerate after changing a builder with:: + + PTWM_REGEN_FIXTURES=1 uv run --no-sync pytest \ + tests/ptwm/preprocessing/test_production_chain_fixture.py + +Format: repeated records of ``u16 name_len | name (UTF-8) | u32 blob_len | +blob`` (little-endian), where ``blob`` is ``Chain.to_bytes()``. +""" + +from __future__ import annotations + +import os +import struct +from pathlib import Path + +from ptwm.preprocessing import _chains +from ptwm.preprocessing._chains import PRODUCTION_CHAINS, ClassifierRole +from ptwm.preprocessing._explorer import ExploreOptions, explore_chains + +FIXTURE = ( + Path(__file__).resolve().parents[3] + / "crates" + / "ptwm-core" + / "tests" + / "fixtures" + / "production_chains.bin" +) + +# Shapes in elements of each builder's Source dtype. Small enough for the +# Rust forward pass to stay fast; they cover 0-d, empty, odd, 1-D, 2-D +# (spherical / alpha-stable only accept 2-D) and 3-D tensors. +SHAPES: list[list[int]] = [ + [], + [0], + [7], + [1000], + [64, 32], + [7, 33], + [3, 5, 8], +] + + +def _builders() -> list[tuple[str, _chains.ChainBuilder]]: + named = [ + (name, getattr(_chains, name)) + for name in sorted(_chains.__all__) + if name.startswith("CHAIN_") + ] + seen = {id(b) for _, b in named} + for (dtype_code, role), builders in sorted(PRODUCTION_CHAINS.items()): + for i, b in enumerate(builders): + if id(b) not in seen: + seen.add(id(b)) + named.append((f"production[{dtype_code},{role.name}][{i}]", b)) + opts = ExploreOptions(max_candidates_per_dtype=8, time_budget_ms=2000) + for dtype_code in (1, 4, 6, 13, 29, 30, 31): + for i, b in enumerate( + explore_chains(dtype_code, ClassifierRole.STANDARD, opts) + ): + named.append((f"explore[{dtype_code}][{i}]", b)) + return named + + +def _fixture_bytes() -> bytes: + out = bytearray() + for name, builder in _builders(): + for shape in SHAPES: + chain = builder(list(shape)) + if chain is None: + continue + label = f"{name}{shape}".encode() + blob = chain.to_bytes() + out += struct.pack(" None: + data = _fixture_bytes() + if os.environ.get("PTWM_REGEN_FIXTURES"): + FIXTURE.write_bytes(data) + assert FIXTURE.read_bytes() == data, ( + "production chain fixture is stale; regenerate it (see module docstring)" + ) diff --git a/tests/ptwm/preprocessing/test_source_shape.py b/tests/ptwm/preprocessing/test_source_shape.py new file mode 100644 index 0000000..b663d3a --- /dev/null +++ b/tests/ptwm/preprocessing/test_source_shape.py @@ -0,0 +1,209 @@ +"""Source shapes of candidate chains cover the raw byte count exactly (P4). + +The decoder sizes every plane from the chain's Source node, and the Rust +writer skips candidates whose Source declares another byte count. These +tests pin :func:`build_chain_blobs` and its three call sites. +""" + +from __future__ import annotations + +import struct + +import numpy as np +import pytest +import torch +from ptwm import ( + CompressionConfig, + Compressor, + DecompressionConfig, + Decompressor, + Format, + Framing, + _rust, +) +from ptwm._core import chain_blob_info, chain_source_byte_len +from ptwm.preprocessing._chains import ( + _CHAIN_DTYPE, + CHAIN_BYTE_PASSTHROUGH_VALUE, + CHAIN_FP32_SPLIT, + ClassifierRole, + _Op, + build_chain_blobs, + chain_storage_bytes, + chains_for, + source_shape_for, +) +from ptwm.sharding import compress_ptwm_blob +from ptwm.utils import DType + + +def _source(blob: bytes) -> tuple[list[int], int]: + """(dims, chain-internal dtype) of a chain blob's Source node.""" + et_size, table = chain_blob_info(blob) + ops = dict(table) + wire = blob[et_size:] + # Legacy chain header: n_nodes, n_edges, n_terms (u8), 3 × u16 lengths. + pos = 9 + op_idx, plen = struct.unpack_from(" list: + return chains_for(dtype.code, ClassifierRole.STANDARD) + + +@pytest.mark.parametrize( + "dtype", + [ + DType.FLOAT32, + DType.BFLOAT16, + DType.FLOAT16, + DType.FLOAT8_E4M3FN, + DType.FLOAT8_E5M2, + DType.FLOAT4_E2M1FN_X2, + DType.UINT8, + DType.FLOAT64, + ], +) +@pytest.mark.parametrize("nbytes", [0, 8, 4096]) +def test_byte_input_source_covers_nbytes(dtype: DType, nbytes: int) -> None: + blobs = build_chain_blobs(_builders(dtype), shape=None, nbytes=nbytes) + assert blobs + for blob in blobs: + assert chain_source_byte_len(blob) == nbytes + dims, src_dtype = _source(blob) + assert dims == [nbytes // chain_storage_bytes(src_dtype)] + + +@pytest.mark.parametrize( + ("dtype", "nbytes"), [(DType.FLOAT32, 4094), (DType.BFLOAT16, 1001)] +) +def test_unaligned_byte_input_falls_back_to_passthrough( + dtype: DType, nbytes: int +) -> None: + blobs = build_chain_blobs(_builders(dtype), shape=None, nbytes=nbytes) + expected = CHAIN_BYTE_PASSTHROUGH_VALUE([nbytes]).to_bytes() + assert blobs == [expected] + + +@pytest.mark.parametrize( + ("shape", "nbytes", "dims"), + [ + ((64, 64), 64 * 64 * 4, [64, 256]), # int32 tensor, element shape + ((), 8, [8]), # int64 scalar + ((0, 16), 0, [0, 16]), # empty + ], +) +def test_passthrough_fallback_counts_bytes( + shape: tuple[int, ...], nbytes: int, dims: list[int] +) -> None: + (blob,) = build_chain_blobs([], shape=list(shape), nbytes=nbytes) + assert _source(blob) == (dims, _CHAIN_DTYPE[DType.UINT8.code]) + assert chain_source_byte_len(blob) == nbytes + + +def test_typed_shapes_are_unchanged() -> None: + blobs = build_chain_blobs( + _builders(DType.BFLOAT16), shape=[64, 32], nbytes=64 * 32 * 2 + ) + assert len(blobs) == len(_builders(DType.BFLOAT16)) + for blob in blobs: + assert _source(blob)[0] == [64, 32] + + +def test_source_shape_for_rules() -> None: + assert source_shape_for(None, 4096, 4) == [1024] + assert source_shape_for(None, 4094, 4) is None + assert source_shape_for([], 4, 4) == [] + assert source_shape_for([], 8, 1) == [8] + assert source_shape_for([3, 5], 60, 4) == [3, 5] + assert source_shape_for([3, 5], 60, 1) == [3, 20] + assert source_shape_for([3, 5], 61, 1) is None + assert source_shape_for([0, 5], 1, 1) is None + assert source_shape_for(None, 1 << 33, 1) is None # dim > u32 + + +def test_compressor_byte_input_declares_raw_size(monkeypatch) -> None: + seen: list[list[bytes]] = [] + real = _rust.compress_compact + + def spy(*args, **kwargs): + seen.append(list(args[6])) + return real(*args, **kwargs) + + monkeypatch.setattr(_rust, "compress_compact", spy) + raw = np.random.default_rng(0).standard_normal(1024).astype(np.float32).tobytes() + cfg = CompressionConfig(bytearray_dtype="float32", framing=Framing.COMPACT) + Compressor(cfg).compress(raw) + (chains,) = seen + assert chains + for blob in chains: + assert chain_source_byte_len(blob) == 4096 + assert CHAIN_FP32_SPLIT([1024]).to_bytes() in chains + + +@pytest.mark.parametrize("framing", [Framing.COMPACT, Framing.CONTAINER]) +@pytest.mark.parametrize( + ("dtype", "nbytes"), + [("float32", 4094), ("float32", 4096), ("float64", 1024), ("int32", 1024)], +) +def test_byte_input_round_trips(framing: Framing, dtype: str, nbytes: int) -> None: + raw = bytes(np.random.default_rng(1).integers(0, 256, nbytes, dtype=np.uint8)) + cfg = CompressionConfig(bytearray_dtype=dtype, framing=framing) + blob = Compressor(cfg).compress(raw) + assert Decompressor().decompress(blob) == raw + + +@pytest.mark.parametrize("framing", [Framing.COMPACT, Framing.CONTAINER]) +@pytest.mark.parametrize("fmt", [Format.TORCH, Format.NUMPY]) +def test_delta_with_tensor_input_round_trips(framing: Framing, fmt: Format) -> None: + # Delta encoding used to replace the tensor with its byte residual and + # then read `.dtype` from the bytes (AttributeError). + rng = np.random.default_rng(2) + base = rng.standard_normal((16, 8)).astype(np.float32) + ref = rng.standard_normal((16, 8)).astype(np.float32) + data = torch.from_numpy(base) if fmt == Format.TORCH else base + ref_in = torch.from_numpy(ref) if fmt == Format.TORCH else ref + cfg = CompressionConfig( + input_format=fmt, framing=framing, delta_compressed_type="byte" + ) + blob = Compressor(cfg).compress(data, delta_second_data=ref_in) + out = Decompressor(DecompressionConfig(delta_second_data=ref.tobytes())).decompress( + blob + ) + if fmt == Format.TORCH: + assert torch.equal(out, data) + else: + np.testing.assert_array_equal(out, base) + + +@pytest.mark.parametrize( + ("dtype", "shape"), + [ + (torch.int64, (32, 4)), + (torch.float64, (8, 8)), + (torch.int64, ()), + (torch.int32, (64, 64)), + (torch.float32, (16, 16)), + ], +) +def test_sharding_writer_chains_cover_raw_bytes(monkeypatch, dtype, shape) -> None: + t = (torch.arange(int(np.prod(shape))) % 13).reshape(shape).to(dtype) + raw = t.contiguous().reshape(-1).view(torch.uint8).numpy().tobytes() + seen: list[list[list[bytes]]] = [] + real = _rust.compress_model + + def spy(records, chains, *args, **kwargs): + seen.append(chains) + return real(records, chains, *args, **kwargs) + + monkeypatch.setattr(_rust, "compress_model", spy) + blob = compress_ptwm_blob([("t", raw, dtype, tuple(shape))]) + for chain in seen[0][0]: + assert chain_source_byte_len(chain) == len(raw) + assert bytes(_rust.decode_tensor(blob, "t", False)) == raw From 5aa2ccef791dc054204cd4bfa3ace6b79e293cdc Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 23:03:34 +0000 Subject: [PATCH 09/34] wip(S4-p5-p4): snapshot of an interrupted review-fix Saved so a container restart cannot lose it. The resumed Stage 4 review and fix validate these changes and follow up in the next commit. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/compact.rs | 38 +++++++++- crates/ptwm-core/src/container.rs | 74 ++++++++++++++----- .../src/transforms/alpha_stable_normalize.rs | 40 ++++++++-- .../src/transforms/burrows_wheeler.rs | 29 +++++++- crates/ptwm-core/src/transforms/concat.rs | 4 +- crates/ptwm-core/src/transforms/index_pack.rs | 30 +++++++- .../src/transforms/mxfp4_deinterleave.rs | 25 ++++++- .../src/transforms/spherical_normalize.rs | 33 ++++++++- python/ptwm/preprocessing/_chains.py | 50 ++++++++++++- tests/ptwm/core/test_compact_frame.py | 2 +- tests/ptwm/preprocessing/test_chains.py | 1 - .../test_production_chain_fixture.py | 5 +- 12 files changed, 291 insertions(+), 40 deletions(-) diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index 603f6ad..190949e 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -379,10 +379,10 @@ fn decode_body( } let (chain, consumed) = read_chain(r.rest(), &table)?; r.pos += consumed; - validate_chain(&chain, 0)?; // Writers check that the chain's Source declares exactly the raw byte // count, so a frame whose Source disagrees with `orig_size` is corrupt. - // Checking before the descriptor walk keeps a corrupt chain from sizing + // Checking before `validate_chain` and the descriptor walk (both call + // every op's `propagate_descriptors`) keeps a corrupt chain from sizing // planes (and codec allocations) beyond the tensor. let source = chain_source_descriptor(&chain)?; if source.length_bytes != orig_size { @@ -391,6 +391,7 @@ fn decode_body( source.length_bytes ))); } + validate_chain(&chain, 0)?; let descriptors = propagate_terminal_descriptors_bounded(&chain, max_plane_len(orig_size))?; let mut terminal_planes: Vec = Vec::with_capacity(descriptors.len()); @@ -729,6 +730,39 @@ mod tests { assert!(t0.elapsed() < std::time::Duration::from_secs(1)); } + #[test] + fn huge_source_is_rejected_before_op_descriptor_arithmetic() { + // Source [2, 1500] int8 → BurrowsWheeler, then the dims are patched + // in place to [u32::MAX, u32::MAX] (about 2^64 bytes). The Source + // check must run before `validate_chain`, which calls + // `BurrowsWheeler::propagate_descriptors`; before the fix, a build + // with overflow checks panicked there. + let raw: Vec = (0..3000u32).map(|i| (i % 5) as u8).collect(); + let mut chain = passthrough_chain(raw.len() as u32); + chain.nodes[0].params = [&[2u8][..], &2u32.to_le_bytes(), &1500u32.to_le_bytes()] + .concat() + .into_iter() + .chain(17u16.to_le_bytes()) + .collect(); + chain.nodes[1].op = OpId::BurrowsWheeler; + let mut inp = input(&raw, None, None); + inp.candidate_chains = vec![chain]; + let blob = compress_compact(&inp, &opts(false, false)) + .unwrap() + .unwrap(); + assert_eq!(blob[FLAGS_AT] & FLAG_STORED, 0); + assert_eq!(decode_compact(&blob).unwrap().raw, raw); + let needle: Vec = [&[2u8][..], &2u32.to_le_bytes(), &1500u32.to_le_bytes()].concat(); + let at = blob + .windows(needle.len()) + .position(|w| w == needle.as_slice()) + .unwrap(); + let mut b = blob.clone(); + b[at + 1..at + 9].copy_from_slice(&[0xFF; 8]); + let err = decode_compact(&b).unwrap_err().to_string(); + assert!(err.contains("Source declares"), "{err}"); + } + #[test] fn source_must_declare_orig_size_exactly() { // Unlike legacy containers, compact frames carry no files written diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index ab593cf..12f5f40 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -1244,6 +1244,25 @@ impl<'a> ContainerReader<'a> { })? }; + // Bound the chain by the record's `orig_size` before any plane is + // sized from it, and before `validate_chain`, which calls every op's + // `propagate_descriptors` on lengths derived from the Source. The + // Source may overstate the raw size by up to 8×: the Python + // byte-input path wrote such chains before the P4 fix, and those + // files must keep decoding. A Source declaring more is corrupt; + // rejecting it here keeps codecs from pre-allocating gigabytes for a + // small tensor. `chain_source_descriptor` checks node 0 itself and + // needs no prior validation. + let source_desc = chain_source_descriptor(chain)?; + let max_source = rec.orig_size.saturating_mul(8); + if source_desc.length_bytes > max_source { + return Err(PtwmCoreError::InvalidContainer(format!( + "decode_tensor: chain Source declares {} bytes for a {}-byte tensor \ + (at most {max_source} accepted)", + source_desc.length_bytes, rec.orig_size + ))); + } + // Validate the chain structurally before handing it to the runtime. // A malformed wire-derived chain (out-of-range edge nodes, missing // Source, cycles, …) would otherwise reach `topo_sort` / @@ -1255,22 +1274,6 @@ impl<'a> ContainerReader<'a> { )) })?; validate_chain(chain, n_dependencies)?; - - // Bound the chain by the record's `orig_size` before any plane is - // sized from it. The Source may overstate the raw size by up to 8×: - // the Python byte-input path wrote such chains before the P4 fix, - // and those files must keep decoding. A Source declaring more is - // corrupt; rejecting it here keeps codecs from pre-allocating - // gigabytes for a small tensor. - let source_desc = chain_source_descriptor(chain)?; - let max_source = rec.orig_size.saturating_mul(8); - if source_desc.length_bytes > max_source { - return Err(PtwmCoreError::InvalidContainer(format!( - "decode_tensor: chain Source declares {} bytes for a {}-byte tensor \ - (at most {max_source} accepted)", - source_desc.length_bytes, rec.orig_size - ))); - } let plane_bound = max_plane_len(rec.orig_size); // Discover the descriptor of every terminal by walking the chain's @@ -3169,12 +3172,26 @@ label = "test-author" raw: &[u8], chunk_size: Option, edit: impl FnOnce(&mut TensorRecord), + ) -> Vec { + identity_container_with_chain( + raw, + v3_passthrough_chain(raw.len() as u32), + chunk_size, + edit, + ) + } + + /// [`identity_container_with`] for a given single-terminal chain. + fn identity_container_with_chain( + raw: &[u8], + chain: Chain, + chunk_size: Option, + edit: impl FnOnce(&mut TensorRecord), ) -> Vec { use crate::compressor::{CompressorOptions, InputTensor, compress_model}; use crate::tensor_record::{CHAIN_REF_INLINE_BIT, parse_tensor_record}; use std::io::Cursor; - let chain = v3_passthrough_chain(raw.len() as u32); let inputs = vec![InputTensor { name: "t".into(), candidate_chains: vec![chain.clone()], @@ -3259,6 +3276,29 @@ label = "test-author" } } + #[test] + fn huge_source_is_rejected_before_op_descriptor_arithmetic() { + // Source [u32::MAX, u32::MAX] int8 (about 2^64 bytes) feeding + // BurrowsWheeler, whose output length adds a block header per + // 256 KiB. The Source bound must reject the chain before + // `validate_chain` walks the op descriptors; before the fix, a build + // with overflow checks (debug, cargo-fuzz) panicked inside + // `BurrowsWheeler::propagate_descriptors`. + use crate::transforms::op::OpId; + let raw: Vec = (0..64u8).collect(); + let mut chain = v3_passthrough_chain(raw.len() as u32); + chain.nodes[1].op = OpId::BurrowsWheeler; + let buf = identity_container_with_chain(&raw, chain, None, |rec| { + set_source_shape(rec, &[u32::MAX, u32::MAX], 0x0005) + }); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert!(err.contains("Source declares"), "{err}"); + } + #[test] fn overstated_source_within_8x_still_decodes() { // Files written by the Python byte-input path before the P4 fix diff --git a/crates/ptwm-core/src/transforms/alpha_stable_normalize.rs b/crates/ptwm-core/src/transforms/alpha_stable_normalize.rs index bb92367..9d32396 100644 --- a/crates/ptwm-core/src/transforms/alpha_stable_normalize.rs +++ b/crates/ptwm-core/src/transforms/alpha_stable_normalize.rs @@ -248,19 +248,28 @@ impl Op for AlphaStableNormalize { )); } let rows = d.length_bytes / (row_len * w); + // Checked: a corrupt chain can declare a length near `u64::MAX`. + let overflow = || { + PtwmCoreError::InvalidContainer(format!( + "AlphaStableNormalize: output length overflows for {} input bytes", + d.length_bytes + )) + }; + let scale_len = rows + .checked_mul(SCALE_BYTES_PER_ROW as u64) + .ok_or_else(overflow)?; + let normalized_len = rows + .checked_mul(row_len) + .and_then(|n| n.checked_mul(2)) + .ok_or_else(overflow)?; Ok(vec![ // scale: (γ, δ) f32 pair per row = SCALE_BYTES_PER_ROW bytes/row - Self::raw_descriptor( - Role::Raw, - ElementWidth::Word4, - rows * SCALE_BYTES_PER_ROW as u64, - Layout::Flat, - ), + Self::raw_descriptor(Role::Raw, ElementWidth::Word4, scale_len, Layout::Flat), // normalized: bf16 per element Self::raw_descriptor( Role::Raw, ElementWidth::Word2, - rows * row_len * 2, + normalized_len, Layout::Rows { row_len: row_len as u32, }, @@ -551,4 +560,21 @@ mod tests { inp.descriptor.layout = Layout::Flat; assert!(op.forward(&[inp]).is_err()); } + + #[test] + fn propagate_descriptors_overflow_is_an_error() { + // 16-bit rows of one element: the scale plane is four times the input. + let op = AlphaStableNormalize::new(0, 2).unwrap(); + let d = PlaneDescriptor { + role: Role::Raw, + element_width: ElementWidth::Word2, + length_bytes: u64::MAX - 1, + layout: Layout::Rows { row_len: 1 }, + derives_from_tensor: None, + residual_of: None, + is_nibble_packed: false, + vendor_bytes: vec![], + }; + assert!(op.propagate_descriptors(&[d]).is_err()); + } } diff --git a/crates/ptwm-core/src/transforms/burrows_wheeler.rs b/crates/ptwm-core/src/transforms/burrows_wheeler.rs index bc81600..1852dd2 100644 --- a/crates/ptwm-core/src/transforms/burrows_wheeler.rs +++ b/crates/ptwm-core/src/transforms/burrows_wheeler.rs @@ -83,9 +83,17 @@ impl Op for BurrowsWheeler { } let d = &inputs[0]; require_byte_plane(d, "propagate_descriptors")?; - let n = d.length_bytes as usize; + // Checked: a corrupt chain can declare a length near `u64::MAX`. + let n = d.length_bytes; + let headers = n + .div_ceil(BWT_BLOCK_SIZE as u64) + .checked_mul(HEADER_BYTES as u64); let mut out = d.clone(); - out.length_bytes = (n + HEADER_BYTES * block_count(n)) as u64; + out.length_bytes = headers.and_then(|h| n.checked_add(h)).ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "BurrowsWheeler.propagate_descriptors: output length overflows for {n} input bytes" + )) + })?; Ok(vec![out]) } @@ -499,4 +507,21 @@ mod tests { .unwrap(); assert_eq!(d1[0].length_bytes, 100 + HEADER_BYTES as u64); } + + #[test] + fn descriptor_length_overflow_is_an_error() { + // A corrupt chain can declare a length near u64::MAX; the header + // bytes must not overflow (a panic under overflow checks). + let d = PlaneDescriptor { + role: Role::Raw, + element_width: ElementWidth::Byte, + length_bytes: u64::MAX, + layout: Layout::Flat, + derives_from_tensor: None, + residual_of: None, + is_nibble_packed: false, + vendor_bytes: vec![], + }; + assert!(BurrowsWheeler::new().propagate_descriptors(&[d]).is_err()); + } } diff --git a/crates/ptwm-core/src/transforms/concat.rs b/crates/ptwm-core/src/transforms/concat.rs index dccb7da..0b1f9f2 100644 --- a/crates/ptwm-core/src/transforms/concat.rs +++ b/crates/ptwm-core/src/transforms/concat.rs @@ -54,7 +54,9 @@ impl Op for Concat { i, input.length_bytes, i, expected_len ))); } - total_bytes += input.length_bytes; + total_bytes = total_bytes.checked_add(input.length_bytes).ok_or_else(|| { + PtwmCoreError::InvalidContainer("Concat: total length overflows u64".into()) + })?; } Ok(vec![PlaneDescriptor { diff --git a/crates/ptwm-core/src/transforms/index_pack.rs b/crates/ptwm-core/src/transforms/index_pack.rs index 476d6d0..325c9ad 100644 --- a/crates/ptwm-core/src/transforms/index_pack.rs +++ b/crates/ptwm-core/src/transforms/index_pack.rs @@ -48,6 +48,12 @@ fn packed_len(n_indices: u64, bits: u8) -> u64 { (n_indices * bits as u64).div_ceil(8) } +/// [`packed_len`] for a length read from a chain: `None` instead of an +/// overflow when `n_indices * bits` does not fit in `u64`. +fn checked_packed_len(n_indices: u64, bits: u8) -> Option { + n_indices.checked_mul(bits as u64).map(|b| b.div_ceil(8)) +} + /// `IndexBitwidthPack` packs a stream of u8 indices (each ∈ `0..2^bits`) into /// a tight bit stream, MSB-first within each output byte. Used downstream of /// codebook ops to recover the bit budget the codebook earned. @@ -93,7 +99,12 @@ impl Op for IndexBitwidthPack { "IndexBitwidthPack: input must not be nibble-packed".into(), )); } - let out_len = packed_len(input.length_bytes, self.bits); + let out_len = checked_packed_len(input.length_bytes, self.bits).ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "IndexBitwidthPack.propagate_descriptors: {} indices of {} bits overflow u64", + input.length_bytes, self.bits + )) + })?; Ok(vec![PlaneDescriptor { role: Role::Index, element_width: ElementWidth::Byte, @@ -264,4 +275,21 @@ mod tests { assert_eq!(decoded.bits, bits); } } + + #[test] + fn propagate_descriptors_overflow_is_an_error() { + // `length_bytes * bits` overflows u64 for a corrupt chain. + let op = IndexBitwidthPack { bits: 8 }; + let d = PlaneDescriptor { + role: Role::Raw, + element_width: ElementWidth::Byte, + length_bytes: u64::MAX, + layout: Layout::Flat, + derives_from_tensor: None, + residual_of: None, + is_nibble_packed: false, + vendor_bytes: vec![], + }; + assert!(op.propagate_descriptors(&[d]).is_err()); + } } diff --git a/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs b/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs index 0867a0c..9ca3b4d 100644 --- a/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs +++ b/crates/ptwm-core/src/transforms/mxfp4_deinterleave.rs @@ -66,7 +66,13 @@ impl Op for MxFp4Deinterleave { ))); } let block_count = inp.length_bytes / BLOCK_BYTES; - let value_len = block_count * BLOCK_VALUES; // 32 nibble-byte slots per block + // 32 nibble-byte slots per block. Checked: a corrupt chain can + // declare a length near `u64::MAX`. + let value_len = block_count.checked_mul(BLOCK_VALUES).ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "MxFp4Deinterleave.propagate_descriptors: {block_count} blocks overflow u64" + )) + })?; let scale_len = block_count; // 1 scale byte per block // Determine row layout for each plane based on input layout. @@ -345,4 +351,21 @@ mod tests { op.write_params(&mut buf); assert_eq!(buf, vec![32u8]); } + + #[test] + fn propagate_descriptors_overflow_is_an_error() { + // 17-byte blocks expand to 32 value slots: overflows u64 near the top. + let op = MxFp4Deinterleave::new(32).unwrap(); + let d = PlaneDescriptor { + role: Role::Raw, + element_width: ElementWidth::Byte, + length_bytes: u64::MAX / 17 * 17, + layout: Layout::Flat, + derives_from_tensor: None, + residual_of: None, + is_nibble_packed: false, + vendor_bytes: vec![], + }; + assert!(op.propagate_descriptors(&[d]).is_err()); + } } diff --git a/crates/ptwm-core/src/transforms/spherical_normalize.rs b/crates/ptwm-core/src/transforms/spherical_normalize.rs index 3c47f7d..aa8cb0a 100644 --- a/crates/ptwm-core/src/transforms/spherical_normalize.rs +++ b/crates/ptwm-core/src/transforms/spherical_normalize.rs @@ -267,14 +267,26 @@ impl Op for SphericalNormalize { )); } let rows = d.length_bytes / (row_len * w); + // Checked: a corrupt chain can declare a length near `u64::MAX`. + let overflow = || { + PtwmCoreError::InvalidContainer(format!( + "SphericalNormalize: output length overflows for {} input bytes", + d.length_bytes + )) + }; + let radius_len = rows.checked_mul(4).ok_or_else(overflow)?; + let direction_len = rows + .checked_mul(row_len) + .and_then(|n| n.checked_mul(2)) + .ok_or_else(overflow)?; Ok(vec![ // radius: one f32 per row - Self::raw_descriptor(Role::Raw, ElementWidth::Word4, rows * 4, Layout::Flat), + Self::raw_descriptor(Role::Raw, ElementWidth::Word4, radius_len, Layout::Flat), // direction: bf16 per component Self::raw_descriptor( Role::Raw, ElementWidth::Word2, - rows * row_len * 2, + direction_len, Layout::Rows { row_len: row_len as u32, }, @@ -740,4 +752,21 @@ mod op_tests { inp.descriptor.layout = Layout::Flat; assert!(op.forward(&[inp]).is_err()); } + + #[test] + fn propagate_descriptors_overflow_is_an_error() { + // 16-bit rows of one element: the radius plane is twice the input. + let op = SphericalNormalize::new(0, 0, 2).unwrap(); + let d = PlaneDescriptor { + role: Role::Raw, + element_width: ElementWidth::Word2, + length_bytes: u64::MAX - 1, + layout: Layout::Rows { row_len: 1 }, + derives_from_tensor: None, + residual_of: None, + is_nibble_packed: false, + vendor_bytes: vec![], + }; + assert!(op.propagate_descriptors(&[d]).is_err()); + } } diff --git a/python/ptwm/preprocessing/_chains.py b/python/ptwm/preprocessing/_chains.py index 9b24415..6bc7948 100644 --- a/python/ptwm/preprocessing/_chains.py +++ b/python/ptwm/preprocessing/_chains.py @@ -118,6 +118,8 @@ class _Op(IntEnum): } _MAX_SOURCE_DIM = 0xFFFF_FFFF +# Largest number of row counts _byte_passthrough_shape tries. +_PASSTHROUGH_ROW_SEARCH = 1 << 16 def chain_storage_bytes(internal_dtype: int) -> int: @@ -1033,6 +1035,29 @@ def source_shape_for( return out +def _byte_passthrough_shape(nbytes: int) -> list[int]: + """Source shape, in bytes, for a passthrough over ``nbytes`` raw bytes. + + ``[nbytes]`` when it fits in ``u32``; otherwise ``[rows, cols]`` with + ``rows * cols == nbytes`` and both dimensions in ``u32``, using the + smallest such ``rows``. The search tries at most + ``_PASSTHROUGH_ROW_SEARCH`` values of ``rows``; real tensor byte counts + have small factors. Raises :class:`ValueError` when no split is found + (for example ``nbytes`` is a prime above ``u32``). + """ + if nbytes <= _MAX_SOURCE_DIM: + return [nbytes] + first = -(-nbytes // _MAX_SOURCE_DIM) + for rows in range(first, min(first + _PASSTHROUGH_ROW_SEARCH, _MAX_SOURCE_DIM + 1)): + if nbytes % rows == 0: + return [rows, nbytes // rows] + msg = ( + f"cannot express {nbytes} bytes as a chain Source shape: no split into " + "two dimensions of at most 2**32 - 1" + ) + raise ValueError(msg) + + def build_chain_blobs( builders: Sequence[ChainBuilder], *, @@ -1050,9 +1075,19 @@ def build_chain_blobs( for example 4094 bytes for an fp32 chain. When no builder remains, the result is the byte passthrough chain over - all ``nbytes``, which fits any buffer. + all ``nbytes``. Its Source is 1-D when ``nbytes`` fits in one ``u32`` + dimension and 2-D otherwise; a byte count with no such 2-D split + raises :class:`ValueError`. + + Each builder is first called once with a probe shape to learn its Source + dtype. The probe clamps every dimension to ``u32`` so that it can always + be serialised: for byte input of 4 GiB or more, ``[nbytes]`` does not + fit, while the real Source shape of a multi-byte dtype often does. """ - probe_shape = list(shape) if shape is not None else [nbytes] + if shape is None: + probe_shape = [min(nbytes, _MAX_SOURCE_DIM)] + else: + probe_shape = [min(int(d), _MAX_SOURCE_DIM) for d in shape] blobs: list[bytes] = [] for builder in builders: chain = builder(probe_shape) @@ -1070,7 +1105,7 @@ def build_chain_blobs( if not blobs: src_shape = source_shape_for(shape, nbytes, 1) if src_shape is None: - src_shape = [nbytes] + src_shape = _byte_passthrough_shape(nbytes) chain = CHAIN_BYTE_PASSTHROUGH_VALUE(src_shape) if chain is None: # pragma: no cover - the passthrough never declines msg = "byte passthrough chain declined its shape" @@ -1101,11 +1136,18 @@ def build_chain_blobs( # PredictorXor variants compete with the plain split chain via # trial-encode; the dispatcher only commits to a variant when it # beats the baseline on byte count. + # + # CHAIN_BF16_MICROSCALE_REPACK is not a candidate. Until the encoder + # took the element width from the chain's Source, the chain failed on + # every bf16 tensor (the public bf16 code mapped to byte elements and + # BlockMicroscalingRepack needs 16-bit ones), so files never used it. + # Now that it works it often wins trial-encode, which would change the + # legacy writer's output; that output stays byte-identical. Callers can + # still pass the builder explicitly. (_DT.BFLOAT16.code, ClassifierRole.STANDARD): [ CHAIN_BF16_SPLIT, CHAIN_BF16_SPLIT_INTDELTA, CHAIN_BF16_SPLIT_PREDICTOR_XOR, - CHAIN_BF16_MICROSCALE_REPACK, CHAIN_BF16_SPLIT_MZS_K1, CHAIN_BF16_SPLIT_MZS_K2, CHAIN_BF16_SPLIT_MZS_K3, diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index 6f6e655..2da047f 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -104,7 +104,7 @@ def test_incompressible_input_is_stored_with_small_overhead() -> None: @pytest.mark.parametrize("framing", [Framing.COMPACT, Framing.CONTAINER]) def test_integrity_fields_can_be_disabled(framing: Framing) -> None: - t = torch.randn(512).to(torch.bfloat16) + t = torch.randn(512, generator=torch.Generator().manual_seed(0)).to(torch.bfloat16) full = Compressor(_torch_cfg(framing=framing)).compress(t) bare = Compressor( _torch_cfg(framing=framing, emit_payload_hash=False, emit_plane_crc=False) diff --git a/tests/ptwm/preprocessing/test_chains.py b/tests/ptwm/preprocessing/test_chains.py index d2a0b3d..15ec03a 100644 --- a/tests/ptwm/preprocessing/test_chains.py +++ b/tests/ptwm/preprocessing/test_chains.py @@ -193,7 +193,6 @@ def test_bfloat16_standard(self): CHAIN_BF16_SPLIT, CHAIN_BF16_SPLIT_INTDELTA, CHAIN_BF16_SPLIT_PREDICTOR_XOR, - CHAIN_BF16_MICROSCALE_REPACK, CHAIN_BF16_SPLIT_MZS_K1, CHAIN_BF16_SPLIT_MZS_K2, CHAIN_BF16_SPLIT_MZS_K3, diff --git a/tests/ptwm/preprocessing/test_production_chain_fixture.py b/tests/ptwm/preprocessing/test_production_chain_fixture.py index 023932c..8e289af 100644 --- a/tests/ptwm/preprocessing/test_production_chain_fixture.py +++ b/tests/ptwm/preprocessing/test_production_chain_fixture.py @@ -61,7 +61,10 @@ def _builders() -> list[tuple[str, _chains.ChainBuilder]]: if id(b) not in seen: seen.add(id(b)) named.append((f"production[{dtype_code},{role.name}][{i}]", b)) - opts = ExploreOptions(max_candidates_per_dtype=8, time_budget_ms=2000) + # The candidate cap bounds the work. The time budget is effectively + # unlimited so that a slow machine still returns all 8 candidates and the + # bytes stay equal to the checked-in fixture. + opts = ExploreOptions(max_candidates_per_dtype=8, time_budget_ms=10**9) for dtype_code in (1, 4, 6, 13, 29, 30, 31): for i, b in enumerate( explore_chains(dtype_code, ClassifierRole.STANDARD, opts) From f64364cd4b2da043b8e3dc33c8e406b5523ae761 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 23:32:44 +0000 Subject: [PATCH 10/34] fix(S4-p5-p4): address review findings Bound codec output while decoding (spec section 3, P5: no buffer beyond orig_size): - PlaneCodec::decode_bounded takes a hard max_len. The default decodes and checks the length. Zstd and ZstdDict decode through a streaming reader that stops one byte past the bound (ZstdDict with a dictionary caps its bulk capacity). rANS, tANS, the arithmetic coders, Order1Arithmetic, ContextMixingLite and NeuralPredictor reject a length prefix above the bound before the decode loop. - Compact frames pass the exact plane length, legacy whole planes pass max_plane_len(orig_size), chunks pass the table's decoded size. A 98 KB frame of 3 GiB zeros now fails at 4 KiB instead of growing the process by 3 GiB. - New tests: per-codec bound tests, a hand-built Zstd RLE bomb, and tests/decode_bomb.rs, which plants a 512 MiB bomb in compact, legacy and chunked files and checks the largest allocation stays under 64 MiB. Other findings: - The spherical and alpha-stable builders decline shapes with an empty dimension, so empty 2-D float tensors such as (5, 0) compress again. - max_plane_len documents why 16x holds: the 8x overstated Source and the 4x alpha-stable scale plane never occur in the same chain. - test_chains.py uses the CHAIN_BF16_MICROSCALE_REPACK import (lint). Production table change carried by the previous commit (5aa2cce): CHAIN_BF16_MICROSCALE_REPACK is no longer a bf16 production candidate. With the encoder taking the element width from the chain's Source, the chain stopped failing and would win trial-encode on some bf16 tensors, changing the legacy writer's output. Removing it keeps that output byte-identical; the builder still works when passed explicitly. The stage report of ce3e944 did not mention this. Legacy and compact output for the reviewer's typed tensor set is unchanged by this commit. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/runtime.rs | 23 +- crates/ptwm-core/src/codec.rs | 43 +++ crates/ptwm-core/src/codec_tagged.rs | 47 ++++ crates/ptwm-core/src/codecs/arithmetic.rs | 57 ++++ crates/ptwm-core/src/codecs/context_mixing.rs | 20 ++ crates/ptwm-core/src/codecs/length_prefix.rs | 23 ++ crates/ptwm-core/src/codecs/mod.rs | 145 +++++++++++ .../ptwm-core/src/codecs/neural_predictor.rs | 20 ++ .../ptwm-core/src/codecs/order1_arithmetic.rs | 20 ++ crates/ptwm-core/src/codecs/rans.rs | 21 ++ crates/ptwm-core/src/codecs/tans.rs | 21 ++ crates/ptwm-core/src/codecs/zstd.rs | 22 ++ crates/ptwm-core/src/codecs/zstd_dict.rs | 52 +++- crates/ptwm-core/src/compact.rs | 12 +- crates/ptwm-core/src/container.rs | 24 +- crates/ptwm-core/tests/decode_bomb.rs | 244 ++++++++++++++++++ python/ptwm/preprocessing/_chains.py | 8 +- tests/ptwm/core/test_compact_frame.py | 13 + tests/ptwm/preprocessing/test_chains.py | 28 ++ 19 files changed, 815 insertions(+), 28 deletions(-) create mode 100644 crates/ptwm-core/tests/decode_bomb.rs diff --git a/crates/ptwm-core/src/chain/runtime.rs b/crates/ptwm-core/src/chain/runtime.rs index b4c11bd..98d8d5f 100644 --- a/crates/ptwm-core/src/chain/runtime.rs +++ b/crates/ptwm-core/src/chain/runtime.rs @@ -297,12 +297,23 @@ pub fn chain_source_descriptor(chain: &Chain) -> Result u64 { orig_size.saturating_mul(16).saturating_add(4096) } diff --git a/crates/ptwm-core/src/codec.rs b/crates/ptwm-core/src/codec.rs index eb6eb4f..aef514e 100644 --- a/crates/ptwm-core/src/codec.rs +++ b/crates/ptwm-core/src/codec.rs @@ -170,6 +170,49 @@ pub trait PlaneCodec: Send + Sync { layout: &PlaneLayout, decoded_len: usize, ) -> Result, PtwmCoreError>; + + /// [`Self::decode`] that produces at most `max_len` bytes. + /// + /// Readers pass a `max_len` derived from the tensor's `orig_size`, so a + /// payload that decodes to more is corrupt. Codecs whose output can be + /// far larger than their payload (Zstd, and the coders that read the + /// decoded length from a payload prefix) override this method and fail + /// before they allocate or decode past `max_len`. The default decodes + /// and then checks the length; it fits codecs whose output is bounded + /// by a small multiple of the payload length or by `decoded_len`. + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + let out = self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + )?; + check_decoded_len(self.id(), out.len(), max_len)?; + Ok(out) + } +} + +/// Error unless a plane of `len` decoded bytes fits in `max_len`. +pub(crate) fn check_decoded_len( + id: CodecId, + len: usize, + max_len: usize, +) -> Result<(), PtwmCoreError> { + if len > max_len { + return Err(PtwmCoreError::InvalidContainer(format!( + "{id:?} plane decodes to {len} bytes, bound is {max_len}" + ))); + } + Ok(()) } /// Look up a codec by id. Returns `None` for unregistered ids; the diff --git a/crates/ptwm-core/src/codec_tagged.rs b/crates/ptwm-core/src/codec_tagged.rs index 481b88a..2978411 100644 --- a/crates/ptwm-core/src/codec_tagged.rs +++ b/crates/ptwm-core/src/codec_tagged.rs @@ -214,6 +214,53 @@ pub fn zstd_decode(blob: &[u8]) -> Result, PtwmCoreError> { }) } +/// Decode a Zstd blob (one or more frames) to at most `max_len` bytes. +/// +/// The output buffer starts at `size_hint.min(max_len)` bytes and the +/// decoder stops reading one byte past `max_len`, so a small frame that +/// expands to gigabytes fails after at most `max_len + 1` decoded bytes. +/// `size_hint` is the expected length; a plane that turns out longer is +/// still read, up to the bound. +#[cfg(feature = "codec-zstd")] +pub fn zstd_decode_bounded( + blob: &[u8], + size_hint: usize, + max_len: usize, +) -> Result, PtwmCoreError> { + use std::io::Read; + if blob.is_empty() { + return Ok(Vec::new()); + } + let err = |e: std::io::Error| PtwmCoreError::CodecDecode { + codec: "zstd", + msg: e.to_string(), + }; + let mut decoder = zstd::stream::read::Decoder::with_buffer(blob).map_err(err)?; + let first = size_hint.min(max_len); + let mut out = Vec::with_capacity(first); + // Read the expected length into the exact-size buffer first, then probe + // for more. In the usual case the probe reads nothing and the buffer is + // never grown. + (&mut decoder) + .take(first as u64) + .read_to_end(&mut out) + .map_err(err)?; + if out.len() == first { + let rest = ((max_len - first) as u64).saturating_add(1); + (&mut decoder) + .take(rest) + .read_to_end(&mut out) + .map_err(err)?; + } + if out.len() > max_len { + return Err(PtwmCoreError::CodecDecode { + codec: "zstd", + msg: format!("frame decodes to more than {max_len} bytes"), + }); + } + Ok(out) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/ptwm-core/src/codecs/arithmetic.rs b/crates/ptwm-core/src/codecs/arithmetic.rs index 03c2a1d..76305bc 100644 --- a/crates/ptwm-core/src/codecs/arithmetic.rs +++ b/crates/ptwm-core/src/codecs/arithmetic.rs @@ -138,6 +138,25 @@ impl PlaneCodec for ArithmeticO0 { let mut model = Order0Static::deserialize(state_bytes)?; Ok(decode_bytes(&mut model, body, n)) } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + length_prefix::check_within(payload, max_len, "arithmetic_o0")?; + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } /// Order-0 adaptive arithmetic coder. No inline state. @@ -199,6 +218,25 @@ impl PlaneCodec for ArithmeticO0Adaptive { let (n, body) = unframe_payload(payload, "arithmetic_o0_adaptive")?; Ok(decode_bytes(&mut Order0Adaptive::new(), body, n)) } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + length_prefix::check_within(payload, max_len, "arithmetic_o0_adaptive")?; + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } /// Order-1 adaptive arithmetic coder (prev-byte context). No inline state. @@ -273,6 +311,25 @@ impl PlaneCodec for ArithmeticO1 { let (n, body) = unframe_payload(payload, "arithmetic_o1")?; Ok(decode_bytes(&mut Order1Adaptive::new(), body, n)) } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + length_prefix::check_within(payload, max_len, "arithmetic_o1")?; + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/context_mixing.rs b/crates/ptwm-core/src/codecs/context_mixing.rs index 5836a4f..e617bc7 100644 --- a/crates/ptwm-core/src/codecs/context_mixing.rs +++ b/crates/ptwm-core/src/codecs/context_mixing.rs @@ -4,6 +4,7 @@ //! named in `codec_menu`. use crate::codec::{CodecId, Encoded, PlaneCodec}; +use crate::codecs::length_prefix; use crate::entropy::context_mixing::{decode_plane, encode_plane}; use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; @@ -76,6 +77,25 @@ impl PlaneCodec for ContextMixingLite { ) -> Result, PtwmCoreError> { decode_plane(payload) } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + length_prefix::check_within(payload, max_len, "context_mixing_lite")?; + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/length_prefix.rs b/crates/ptwm-core/src/codecs/length_prefix.rs index f2e1f59..5bbccbd 100644 --- a/crates/ptwm-core/src/codecs/length_prefix.rs +++ b/crates/ptwm-core/src/codecs/length_prefix.rs @@ -32,6 +32,29 @@ pub(crate) fn accepted(decoded_len: usize, body_len: usize) -> bool { && (decoded_len <= RATIO_FREE_LEN || decoded_len <= body_len.saturating_mul(MAX_RATIO)) } +/// Error when the `u32` length prefix at the start of `prefixed` declares +/// more than `max_len` decoded bytes. +/// +/// Used by the bounded decode paths before the decode loop runs. A payload +/// shorter than the prefix passes; the codec's own decoder reports it. +pub(crate) fn check_within( + prefixed: &[u8], + max_len: usize, + codec: &'static str, +) -> Result<(), crate::error::PtwmCoreError> { + let Some(prefix) = prefixed.first_chunk::<4>() else { + return Ok(()); + }; + let declared = u32::from_le_bytes(*prefix) as usize; + if declared > max_len { + return Err(crate::error::PtwmCoreError::CodecDecode { + codec, + msg: format!("payload declares {declared} decoded bytes, bound is {max_len}"), + }); + } + Ok(()) +} + /// Bytes from a skewed distribution (about 2 bits per byte), long enough to /// exercise lengths above [`RATIO_FREE_LEN`] while staying far below /// [`MAX_RATIO`]. diff --git a/crates/ptwm-core/src/codecs/mod.rs b/crates/ptwm-core/src/codecs/mod.rs index 39d3049..1bc1f8e 100644 --- a/crates/ptwm-core/src/codecs/mod.rs +++ b/crates/ptwm-core/src/codecs/mod.rs @@ -13,3 +13,148 @@ pub mod rans; pub mod tans; pub mod zstd; pub mod zstd_dict; + +/// A Zstd frame that decodes to `total` zero bytes, built from RLE blocks +/// of 128 KiB (4 bytes each), so it is about `total / 32768` bytes long and +/// is built without compressing anything. +#[cfg(test)] +pub(crate) fn zstd_zero_bomb(total: u64) -> Vec { + const BLOCK: u64 = 128 * 1024; + // Magic, then a frame header with no content size, no checksum and a + // 128 KiB window (Window_Descriptor exponent 7, mantissa 0). + let mut out = vec![0x28, 0xB5, 0x2F, 0xFD, 0x00, 0x38]; + let mut left = total; + while left > 0 { + let n = left.min(BLOCK); + left -= n; + let last = u32::from(left == 0); + // Block header: Last_Block bit, Block_Type 1 (RLE), Block_Size. + let header = ((n as u32) << 3) | (1 << 1) | last; + out.extend_from_slice(&header.to_le_bytes()[..3]); + out.push(0); + } + out +} + +#[cfg(test)] +mod bounded_decode_tests { + use crate::codec::{CodecId, PlaneCodec, codec_for}; + use crate::layout::PlaneLayout; + + /// Codecs that encode any byte plane with a flat layout. + const FLAT_CODECS: &[CodecId] = &[ + CodecId::Identity, + CodecId::Huffman, + CodecId::Rans, + CodecId::Tans, + CodecId::Zstd, + CodecId::ZstdDict, + CodecId::ArithmeticO0, + CodecId::ArithmeticO0Adaptive, + CodecId::ArithmeticO1, + CodecId::Order1Arithmetic, + CodecId::ContextMixingLite, + CodecId::NeuralPredictor, + ]; + + fn plane() -> Vec { + (0..20_000u32).map(|i| ((i * 7) % 13) as u8).collect() + } + + #[test] + fn decode_bounded_accepts_the_exact_length_and_rejects_one_less() { + let data = plane(); + for &id in FLAT_CODECS { + let c = codec_for(id).unwrap(); + let enc = c.encode(&data, None, &PlaneLayout::Flat).unwrap(); + let dec = |max_len| { + c.decode_bounded( + enc.state_format_version, + &enc.state_bytes, + &enc.payload, + &PlaneLayout::Flat, + data.len(), + max_len, + ) + }; + assert_eq!(dec(data.len()).unwrap(), data, "{id:?}"); + assert_eq!(dec(usize::MAX).unwrap(), data, "{id:?}"); + assert!(dec(data.len() - 1).is_err(), "{id:?} exceeded its bound"); + } + } + + /// Codecs whose payload starts with a `u32` decoded length, and the + /// bytes before it. + fn prefixed_codecs() -> Vec<(Box, Vec, u8)> { + vec![ + (codec_for(CodecId::Rans).unwrap(), vec![1], 0), + (codec_for(CodecId::Tans).unwrap(), vec![1], 0), + (codec_for(CodecId::ArithmeticO0).unwrap(), vec![], 1), + (codec_for(CodecId::ArithmeticO0Adaptive).unwrap(), vec![], 0), + (codec_for(CodecId::Order1Arithmetic).unwrap(), vec![], 1), + (codec_for(CodecId::ArithmeticO1).unwrap(), vec![], 0), + (codec_for(CodecId::ContextMixingLite).unwrap(), vec![], 0), + (codec_for(CodecId::NeuralPredictor).unwrap(), vec![], 0), + ] + } + + #[test] + fn a_length_prefix_above_the_bound_fails_before_decoding() { + // 1 GiB declared from a 64-byte body: within `length_prefix`'s + // ratio-free length, so before the bound the decode loop ran for + // 2^30 symbols. + let t0 = std::time::Instant::now(); + for (c, head, version) in prefixed_codecs() { + let mut payload = head; + payload.extend_from_slice(&(1u32 << 30).to_le_bytes()); + payload.extend_from_slice(&[0u8; 64]); + let err = c + .decode_bounded(version, &[], &payload, &PlaneLayout::Flat, 4096, 4096) + .unwrap_err() + .to_string(); + assert!(err.contains("bound is 4096"), "{:?}: {err}", c.id()); + } + assert!(t0.elapsed() < std::time::Duration::from_secs(2)); + } + + #[cfg(feature = "codec-zstd")] + #[test] + fn zstd_bomb_stops_at_the_bound() { + let bomb = super::zstd_zero_bomb(3 << 30); + assert!(bomb.len() < 100_000); + let t0 = std::time::Instant::now(); + for id in [CodecId::Zstd, CodecId::ZstdDict] { + let c = codec_for(id).unwrap(); + let err = c + .decode_bounded(0, &[], &bomb, &PlaneLayout::Flat, 4096, 4096) + .unwrap_err() + .to_string(); + assert!(err.contains("more than 4096 bytes"), "{id:?}: {err}"); + } + assert!(t0.elapsed() < std::time::Duration::from_secs(1)); + } + + #[cfg(feature = "codec-zstd")] + #[test] + fn zstd_reads_past_a_short_hint_up_to_the_bound() { + // Legacy readers pass the descriptor length as the hint, which can + // understate the plane; the bound is what limits it. + let data = plane(); + let c = codec_for(CodecId::Zstd).unwrap(); + let enc = c.encode(&data, None, &PlaneLayout::Flat).unwrap(); + for hint in [0, 1, data.len() / 2, data.len(), data.len() * 3] { + let got = c + .decode_bounded(0, &[], &enc.payload, &PlaneLayout::Flat, hint, data.len()) + .unwrap(); + assert_eq!(got, data, "hint {hint}"); + } + // A bomb frame followed by a skippable frame still decodes when it + // fits: skippable frames are part of the format. + let mut two = super::zstd_zero_bomb(1000); + two.extend_from_slice(&[0x50, 0x2A, 0x4D, 0x18, 3, 0, 0, 0, 9, 9, 9]); + let got = c + .decode_bounded(0, &[], &two, &PlaneLayout::Flat, 10, 1000) + .unwrap(); + assert_eq!(got, vec![0u8; 1000]); + } +} diff --git a/crates/ptwm-core/src/codecs/neural_predictor.rs b/crates/ptwm-core/src/codecs/neural_predictor.rs index e7f5f32..323c975 100644 --- a/crates/ptwm-core/src/codecs/neural_predictor.rs +++ b/crates/ptwm-core/src/codecs/neural_predictor.rs @@ -3,6 +3,7 @@ //! in the default dispatch menu — selected only when named in `codec_menu`. use crate::codec::{CodecId, Encoded, PlaneCodec}; +use crate::codecs::length_prefix; use crate::entropy::neural_predictor::{decode_plane, encode_plane}; use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; @@ -74,6 +75,25 @@ impl PlaneCodec for NeuralPredictor { ) -> Result, PtwmCoreError> { decode_plane(payload) } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + length_prefix::check_within(payload, max_len, "neural_predictor")?; + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/order1_arithmetic.rs b/crates/ptwm-core/src/codecs/order1_arithmetic.rs index c1f16c9..f27a38d 100644 --- a/crates/ptwm-core/src/codecs/order1_arithmetic.rs +++ b/crates/ptwm-core/src/codecs/order1_arithmetic.rs @@ -7,6 +7,7 @@ use crate::codec::{CodecId, Encoded, PlaneCodec}; use crate::codecs::arithmetic::{frame_payload, unframe_payload}; +use crate::codecs::length_prefix; use crate::entropy::arithmetic::{Order1Static, decode_bytes, encode_bytes}; use crate::error::PtwmCoreError; use crate::layout::PlaneLayout; @@ -107,6 +108,25 @@ impl PlaneCodec for Order1Arithmetic { let (n, body) = unframe_payload(payload, "order1_arithmetic")?; Ok(decode_bytes(&mut model, body, n)) } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + length_prefix::check_within(payload, max_len, "order1_arithmetic")?; + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/rans.rs b/crates/ptwm-core/src/codecs/rans.rs index 03205ab..ab02b1e 100644 --- a/crates/ptwm-core/src/codecs/rans.rs +++ b/crates/ptwm-core/src/codecs/rans.rs @@ -142,6 +142,27 @@ impl PlaneCodec for Rans { }), } } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + if let Some((&TAG_RANS, rest)) = payload.split_first() { + length_prefix::check_within(rest, max_len, "rans")?; + } + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/tans.rs b/crates/ptwm-core/src/codecs/tans.rs index 0302582..2c2246d 100644 --- a/crates/ptwm-core/src/codecs/tans.rs +++ b/crates/ptwm-core/src/codecs/tans.rs @@ -135,6 +135,27 @@ impl PlaneCodec for Tans { }), } } + + fn decode_bounded( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + if let Some((&TAG_TANS, rest)) = payload.split_first() { + length_prefix::check_within(rest, max_len, "tans")?; + } + self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + ) + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/zstd.rs b/crates/ptwm-core/src/codecs/zstd.rs index 7d17b51..b0b599d 100644 --- a/crates/ptwm-core/src/codecs/zstd.rs +++ b/crates/ptwm-core/src/codecs/zstd.rs @@ -82,6 +82,28 @@ impl PlaneCodec for Zstd { )) } } + + fn decode_bounded( + &self, + _state_format_version: u8, + _state_bytes: &[u8], + payload: &[u8], + _layout: &PlaneLayout, + decoded_len: usize, + max_len: usize, + ) -> Result, PtwmCoreError> { + #[cfg(feature = "codec-zstd")] + { + crate::codec_tagged::zstd_decode_bounded(payload, decoded_len, max_len) + } + #[cfg(not(feature = "codec-zstd"))] + { + let _ = (payload, decoded_len, max_len); + Err(PtwmCoreError::InvalidContainer( + "codec-zstd feature is disabled".into(), + )) + } + } } #[cfg(test)] diff --git a/crates/ptwm-core/src/codecs/zstd_dict.rs b/crates/ptwm-core/src/codecs/zstd_dict.rs index c436639..f963139 100644 --- a/crates/ptwm-core/src/codecs/zstd_dict.rs +++ b/crates/ptwm-core/src/codecs/zstd_dict.rs @@ -105,12 +105,31 @@ impl PlaneCodec for ZstdDict { } fn decode( + &self, + state_format_version: u8, + state_bytes: &[u8], + payload: &[u8], + layout: &PlaneLayout, + decoded_len: usize, + ) -> Result, PtwmCoreError> { + self.decode_bounded( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + usize::MAX, + ) + } + + fn decode_bounded( &self, state_format_version: u8, state_bytes: &[u8], payload: &[u8], _layout: &PlaneLayout, decoded_len: usize, + max_len: usize, ) -> Result, PtwmCoreError> { #[cfg(feature = "codec-zstd")] { @@ -127,18 +146,24 @@ impl PlaneCodec for ZstdDict { return Ok(Vec::new()); } if state_bytes.is_empty() { - zstd::decode_all(payload) - } else { - decode_with_dict(payload, state_bytes, decoded_len) + return crate::codec_tagged::zstd_decode_bounded(payload, decoded_len, max_len); } - .map_err(|e| PtwmCoreError::CodecDecode { - codec: "zstd_dict", - msg: e.to_string(), + decode_with_dict(payload, state_bytes, decoded_len, max_len).map_err(|e| { + PtwmCoreError::CodecDecode { + codec: "zstd_dict", + msg: e.to_string(), + } }) } #[cfg(not(feature = "codec-zstd"))] { - let _ = (state_format_version, state_bytes, payload, decoded_len); + let _ = ( + state_format_version, + state_bytes, + payload, + decoded_len, + max_len, + ); Err(PtwmCoreError::InvalidContainer( "codec-zstd feature is disabled".into(), )) @@ -161,20 +186,27 @@ fn encode_with_dict(plane: &[u8], dict: &[u8], level: i32) -> std::io::Result std::io::Result> { +fn decode_with_dict( + payload: &[u8], + dict: &[u8], + decoded_len: usize, + max_len: usize, +) -> std::io::Result> { let mut decompressor = zstd::bulk::Decompressor::with_dictionary(dict)?; // The bulk Decompressor writes into a buffer of exactly the // capacity we pass in. The dispatcher knows the decoded length // from the propagated terminal descriptor and forwards it via the // `decoded_len` argument. When that is zero (older callers that // don't propagate it), fall back to the frame's pledged source - // size if present, then to an 8x heuristic. + // size if present, then to an 8x heuristic. The capacity never + // exceeds `max_len`: a frame that needs more fails instead. let capacity = if decoded_len > 0 { decoded_len } else { zstd::bulk::Decompressor::upper_bound(payload) .unwrap_or_else(|| payload.len().saturating_mul(8)) - }; + } + .min(max_len); decompressor.decompress(payload, capacity) } diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index 190949e..69690ad 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -422,9 +422,17 @@ fn decode_body( let layout = legacy_plane_layout(descriptor.layout); let decoded_len = usize::try_from(descriptor.length_bytes) .map_err(|_| bad("plane length overflows usize"))?; - let bytes = codec.decode(state_version, state, payload, &layout, decoded_len)?; // The Source is exact, so every plane decodes to its descriptor's - // length. + // length. The codec stops at that length, so a corrupt payload cannot + // allocate past it. + let bytes = codec.decode_bounded( + state_version, + state, + payload, + &layout, + decoded_len, + decoded_len, + )?; if bytes.len() != decoded_len { return Err(bad(&format!( "plane decoded to {} bytes, chain says {decoded_len}", diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 12f5f40..43e4ebe 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -1418,12 +1418,15 @@ impl<'a> ContainerReader<'a> { } let state_bytes = &frame[5..5 + state_len]; let chunk_payload = &frame[5 + state_len..]; - let decoded = codec.decode( + // The table's decoded size is bounded above; the codec + // stops at it, so a corrupt chunk cannot decode past it. + let decoded = codec.decode_bounded( state_version, state_bytes, chunk_payload, &plane_rec.layout, entry.decoded_size as usize, + entry.decoded_size as usize, )?; if decoded.len() != entry.decoded_size as usize { return Err(PtwmCoreError::InvalidContainer( @@ -1501,20 +1504,25 @@ impl<'a> ContainerReader<'a> { let decoded_len = usize::try_from(descriptor.length_bytes).map_err(|_| { PtwmCoreError::InvalidContainer("plane length overflows usize".into()) })?; - codec.decode( + // Self-delimiting codecs may return a length other than the + // descriptor's: older files carry Sources that overstate + // (Python byte input) or understate (byte passthrough over an + // element shape) the raw size. So the codec's hard limit is + // `plane_bound`, derived from `orig_size`, not `decoded_len`; + // the length check after `inverse_chain` and the payload hash + // cover the rest. The codec stops at the limit, so a corrupt + // payload cannot allocate past it. + codec.decode_bounded( state_version, &state_bytes, &plane_rec.payload_bytes, &plane_rec.layout, decoded_len, + usize::try_from(plane_bound).unwrap_or(usize::MAX), )? }; - // Self-delimiting codecs may return a length other than the - // descriptor's: older files carry Sources that overstate (Python - // byte input) or understate (byte passthrough over an element - // shape) the raw size. Bound the result by `orig_size` instead; - // the length check after `inverse_chain` and the payload hash - // cover the rest. + // The codecs enforce the bound; this check keeps it for any + // codec that does not. if plane_bytes.len() as u64 > plane_bound { return Err(PtwmCoreError::InvalidContainer(format!( "decoded plane is {} bytes, bound is {plane_bound} for a {}-byte tensor", diff --git a/crates/ptwm-core/tests/decode_bomb.rs b/crates/ptwm-core/tests/decode_bomb.rs new file mode 100644 index 0000000..8dbc087 --- /dev/null +++ b/crates/ptwm-core/tests/decode_bomb.rs @@ -0,0 +1,244 @@ +//! Decompression bombs must fail without allocating past the tensor. +//! +//! Each test writes a small tensor with the Zstd plane codec, then replaces +//! one Zstd payload in place with a frame of the same length that decodes +//! to 512 MiB of zeros (RLE blocks, padded with a Zstd skippable frame). +//! Every length and offset in the file stays valid, so the reader reaches +//! the codec. The global allocator records the largest single allocation; +//! the tests assert it stays far below the bomb's size. Before the codecs +//! took a bound, the reader decoded the whole frame and only then compared +//! its length with the chain. +//! +//! This binary holds only these tests: they share the allocator's counter. + +use std::alloc::{GlobalAlloc, Layout, System}; +use std::io::Cursor; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; +use ptwm_core::codec::{CodecId, codec_for}; +use ptwm_core::compact::{compress_compact, decode_compact}; +use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; +use ptwm_core::container::ContainerReader; +use ptwm_core::layout::PlaneLayout; +use ptwm_core::transforms::op::OpId; +use ptwm_core::types::{Role, ValueFormat}; + +struct PeakAlloc; + +static LARGEST: AtomicUsize = AtomicUsize::new(0); + +unsafe impl GlobalAlloc for PeakAlloc { + unsafe fn alloc(&self, layout: Layout) -> *mut u8 { + LARGEST.fetch_max(layout.size(), Ordering::Relaxed); + unsafe { System.alloc(layout) } + } + unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 { + LARGEST.fetch_max(layout.size(), Ordering::Relaxed); + unsafe { System.alloc_zeroed(layout) } + } + unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) { + unsafe { System.dealloc(ptr, layout) } + } + unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 { + LARGEST.fetch_max(new_size, Ordering::Relaxed); + unsafe { System.realloc(ptr, layout, new_size) } + } +} + +#[global_allocator] +static GLOBAL: PeakAlloc = PeakAlloc; + +const BOMB_LEN: u64 = 512 << 20; +/// Largest single allocation a bomb test may cause. +const ALLOC_LIMIT: usize = 64 << 20; +const RAW_LEN: usize = 128 * 1024; + +/// Random bytes in `0..16`: Zstd codes them at about half their size. +fn raw() -> Vec { + let mut x: u64 = 0x9E37_79B9_7F4A_7C15; + (0..RAW_LEN) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + (x & 0x0F) as u8 + }) + .collect() +} + +/// Source(int8) → BytePassthrough → one Value terminal. +fn passthrough_chain(n_bytes: u32) -> Chain { + let role = Role::Value { + format: ValueFormat::IntN { bits: 8 }, + }; + let mut source = vec![1u8]; + source.extend_from_slice(&n_bytes.to_le_bytes()); + source.extend_from_slice(&17u16.to_le_bytes()); + Chain { + nodes: vec![ + ChainNode { + op: OpId::Source, + params: source, + }, + ChainNode { + op: OpId::BytePassthrough, + params: Vec::new(), + }, + ], + edges: vec![ChainEdge { + src_node: 0, + src_output_idx: 0, + dst_node: 1, + dst_input_idx: 0, + role_override: Some(role.clone()), + vendor_bytes: Vec::new(), + }], + terminals: vec![TerminalRef { + node_idx: 1, + output_idx: 0, + role, + }], + } +} + +fn input(raw: &[u8]) -> InputTensor<'_> { + InputTensor { + name: "t".into(), + candidate_chains: vec![passthrough_chain(raw.len() as u32)], + dtype_code: 0, + input_format: 1, + raw_bytes: raw, + shape: None, + dtype_name: None, + delta_reference_blake3: None, + } +} + +fn opts(chunk_size: Option) -> CompressorOptions { + CompressorOptions { + method_hint: 1, + emit_payload_hash: false, + emit_plane_crc: false, + forced_codec: Some(CodecId::Zstd), + allow_codec_ids: None, + chunk_size, + } +} + +/// A Zstd frame of `BOMB_LEN` zeros from 128 KiB RLE blocks, followed by a +/// skippable frame that pads it to exactly `len` bytes. +fn bomb_of_len(len: usize) -> Vec { + const BLOCK: u64 = 128 * 1024; + let mut out = vec![0x28, 0xB5, 0x2F, 0xFD, 0x00, 0x38]; + let mut left = BOMB_LEN; + while left > 0 { + let n = left.min(BLOCK); + left -= n; + let header = ((n as u32) << 3) | (1 << 1) | u32::from(left == 0); + out.extend_from_slice(&header.to_le_bytes()[..3]); + out.push(0); + } + let pad = len + .checked_sub(out.len() + 8) + .expect("payload too short for the bomb"); + out.extend_from_slice(&0x184D_2A50u32.to_le_bytes()); + out.extend_from_slice(&(pad as u32).to_le_bytes()); + out.resize(len, 0); + out +} + +/// Replace the one occurrence of `payload` in `file` with a bomb. +fn plant_bomb(file: &mut [u8], payload: &[u8]) { + let at = file + .windows(payload.len()) + .position(|w| w == payload) + .expect("Zstd payload in file"); + assert!( + file[at + 1..].windows(payload.len()).all(|w| w != payload), + "payload must be unique" + ); + file[at..at + payload.len()].copy_from_slice(&bomb_of_len(payload.len())); +} + +fn zstd_payload(plane: &[u8]) -> Vec { + codec_for(CodecId::Zstd) + .unwrap() + .encode(plane, None, &PlaneLayout::Flat) + .unwrap() + .payload +} + +fn assert_bounded(err: String, t0: Instant) { + assert!(err.contains("more than"), "{err}"); + assert!(t0.elapsed() < Duration::from_secs(1), "{:?}", t0.elapsed()); + let largest = LARGEST.load(Ordering::Relaxed); + assert!( + largest < ALLOC_LIMIT, + "largest allocation was {largest} bytes" + ); +} + +#[test] +fn compact_frame_bomb_is_bounded() { + let raw = raw(); + let mut blob = compress_compact(&input(&raw), &opts(None)) + .unwrap() + .expect("compact frame"); + assert_eq!(decode_compact(&blob).unwrap().raw, raw); + plant_bomb(&mut blob, &zstd_payload(&raw)); + let t0 = Instant::now(); + let err = decode_compact(&blob).unwrap_err().to_string(); + assert_bounded(err, t0); +} + +#[test] +fn legacy_plane_bomb_is_bounded() { + let raw = raw(); + let mut buf = Vec::new(); + compress_model(Cursor::new(&mut buf), &[input(&raw)], opts(None)).unwrap(); + assert_eq!( + ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap(), + raw + ); + plant_bomb(&mut buf, &zstd_payload(&raw)); + let t0 = Instant::now(); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert_bounded(err, t0); +} + +#[test] +fn legacy_chunk_bomb_is_bounded() { + const CHUNK: usize = 64 * 1024; + let raw = raw(); + let mut buf = Vec::new(); + compress_model( + Cursor::new(&mut buf), + &[input(&raw)], + opts(Some(CHUNK as u32)), + ) + .unwrap(); + assert_eq!( + ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap(), + raw + ); + plant_bomb(&mut buf, &zstd_payload(&raw[..CHUNK])); + let t0 = Instant::now(); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert_bounded(err, t0); +} diff --git a/python/ptwm/preprocessing/_chains.py b/python/ptwm/preprocessing/_chains.py index 6bc7948..58a1ead 100644 --- a/python/ptwm/preprocessing/_chains.py +++ b/python/ptwm/preprocessing/_chains.py @@ -762,7 +762,9 @@ def _builder(shape: list[int]) -> Chain | None: # meaningful (rows, embedding-dim) factorization to normalize. # Returning None drops this candidate so the op never runs on 1D # (bias/norm) or higher-rank tensors, bounding trial-encode cost. - if len(shape) != 2: + # An empty dimension leaves no row to normalize, and the op rejects + # it, so the candidate is dropped as well. + if len(shape) != 2 or 0 in shape: return None return Chain( nodes=[ @@ -818,7 +820,9 @@ def _builder(shape: list[int]) -> Chain | None: # meaningful (rows, embedding-dim) factorization to normalize. # Returning None drops this candidate so the op never runs on 1D # (bias/norm) or higher-rank tensors, bounding trial-encode cost. - if len(shape) != 2: + # An empty dimension leaves no row to normalize, and the op rejects + # it, so the candidate is dropped as well. + if len(shape) != 2 or 0 in shape: return None return Chain( nodes=[ diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index 2da047f..5be10d0 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -163,3 +163,16 @@ def test_decode_compact_reports_metadata() -> None: assert input_format == Format.TORCH.value assert delta is None assert 0 < payload_len < len(blob) + + +@pytest.mark.parametrize("framing", [Framing.COMPACT, Framing.CONTAINER]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("shape", [(5, 0), (0, 5)]) +def test_empty_matrix_round_trips(framing: Framing, dtype: torch.dtype, shape) -> None: + # The row-normalizing chains reject a matrix with an empty dimension; + # the builders drop them so the tensor still compresses. + t = torch.empty(shape, dtype=dtype) + blob = Compressor(_torch_cfg(framing=framing)).compress(t) + back = Decompressor().decompress(blob) + assert back.dtype == dtype + assert back.shape == t.shape diff --git a/tests/ptwm/preprocessing/test_chains.py b/tests/ptwm/preprocessing/test_chains.py index 15ec03a..51a8c77 100644 --- a/tests/ptwm/preprocessing/test_chains.py +++ b/tests/ptwm/preprocessing/test_chains.py @@ -202,6 +202,33 @@ def test_bfloat16_standard(self): CHAIN_BF16_HUFF_LLM_RAW, ] + def test_microscale_repack_is_not_a_production_candidate(self): + # Kept out of every table so the legacy writer's bf16 output stays + # byte-identical; see the comment above PRODUCTION_CHAINS. The + # builder itself still works when passed explicitly + # (TestCanonicalChains). + for code in (DT.BFLOAT16.code, DT.FLOAT16.code): + for role in ClassifierRole: + assert CHAIN_BF16_MICROSCALE_REPACK not in chains_for(code, role) + + @pytest.mark.parametrize( + "builder", + [ + CHAIN_BF16_SPHERICAL_NORMALIZE, + CHAIN_FP16_SPHERICAL_NORMALIZE, + CHAIN_FP32_SPHERICAL_NORMALIZE, + CHAIN_BF16_ALPHA_STABLE_NORMALIZE, + CHAIN_FP16_ALPHA_STABLE_NORMALIZE, + CHAIN_FP32_ALPHA_STABLE_NORMALIZE, + ], + ) + def test_row_normalizers_decline_empty_dimensions(self, builder): + # The ops reject a matrix with no rows or no columns, so the builders + # must drop the candidate rather than fail the whole tensor. + assert builder([5, 0]) is None + assert builder([0, 5]) is None + assert builder([5, 3]) is not None + def test_float16_standard_same_chain_as_bf16(self): result = chains_for(DT.FLOAT16.code, ClassifierRole.STANDARD) assert result == [ @@ -268,6 +295,7 @@ class TestCanonicalChains: (CHAIN_F32_GLOBAL_SCALE, [1]), (CHAIN_FP4_NIBBLE_PASSTHROUGH, [64]), (CHAIN_BYTE_PASSTHROUGH_VALUE, [256]), + (CHAIN_BF16_MICROSCALE_REPACK, [4, 64]), ], ) def test_builder_produces_chain(self, builder, shape): From 3c5487864cd58e560946aaee8e2ad377d02d5d6b Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 23:47:27 +0000 Subject: [PATCH 11/34] feat(core): compact chain wire (CW2) and a smaller PTWS frame Stage 5 of the issue #5 work (synthesis section 2 CW2 and PTWS, P3). - crate::varint: canonical unsigned LEB128 (minimal encoding, at most 10 bytes, fits u64) and a bounds-checked Cursor whose len_varint never exceeds the remaining bytes. The compact frame's private helpers moved here. - chain::wire_compact (CW2): varint counts, op_ref resolved by the caller (OpId wire id for PTWS, Extension Table index for PTWC later), spine-edge compression (leading (i,0)->(i+1,0) edges are a count), one-byte terminal refs, one-byte compact roles with a 0xFF escape, and a bound Source (param_len 0 = params of the tensor's own shape and dtype). Every field has one encoding; the reader rejects the others (overlong varints, needless escapes, explicit spine edges, explicit Source params equal to the binding, unknown etag bits or role bytes). - transforms::source: chain_dtype_for_public (mirror of _CHAIN_DTYPE) and bound_source_params. - PTWS frame, changed in place (unreleased): CW2 chain replaces the op table and registry chain wire; plane head is one varint codec_id << 2 | state with state 2 = inline, version 0, empty (state 1 with an empty version-0 state is rejected); the last plane's payload_len is omitted; orig_size is stored only when the shape does not give it (new FLAG_SIZE, rejected when redundant); a stored dtype name equal to the canonical one is rejected. Shapes above 64 dims fall back to the full container instead of writing an unreadable frame. - Fixed overhead with hash and CRCs on: fp32 random 125 -> 68 B, bf16 random 98 -> 48 B, fp32 constant 124 -> 62 B. - Fuzz: compact_decode follows the new header order; new target chain_wire_compact asserts accepted chains re-encode to the same bytes. - tests/ptwm/benchmarks/test_container_overhead.py bounds compact non-payload bytes per dtype (fp32 <= 70 B, bf16 <= 55 B) for random and constant tensors at every size. - wire-format.md and the Framing docstring describe the new frame. Tests: varint (round trip, non-minimal, overflow, cursor bounds); CW2 round trip on all 500+ production/explorer chains in the fixture (explicit and bound), exact bytes of the fp32 split chain (19 B), edge cases (non-spine edges, vendor bytes, escaped refs and roles, large op refs, 0-d Source, 255 nodes), unencodable chains, canonicality rejections, and 60k random mutations that must be rejected or re-encode identically; PTWS bound Source, shape-independent chain bytes, bound Source without a frame shape, derived vs stored orig_size, header and plane-state canonicality, last-plane truncation. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/mod.rs | 1 + crates/ptwm-core/src/chain/wire_compact.rs | 892 ++++++++++++++++++ crates/ptwm-core/src/compact.rs | 693 +++++++++----- crates/ptwm-core/src/lib.rs | 1 + crates/ptwm-core/src/transforms/source.rs | 72 ++ crates/ptwm-core/src/types/role.rs | 8 +- crates/ptwm-core/src/varint.rs | 263 ++++++ docs/architecture/wire-format.md | 71 +- fuzz/Cargo.toml | 6 + fuzz/fuzz_targets/chain_wire_compact.rs | 31 + fuzz/fuzz_targets/compact_decode.rs | 46 +- python/ptwm/_config.py | 2 +- .../benchmarks/test_container_overhead.py | 8 +- 13 files changed, 1839 insertions(+), 255 deletions(-) create mode 100644 crates/ptwm-core/src/chain/wire_compact.rs create mode 100644 crates/ptwm-core/src/varint.rs create mode 100644 fuzz/fuzz_targets/chain_wire_compact.rs diff --git a/crates/ptwm-core/src/chain/mod.rs b/crates/ptwm-core/src/chain/mod.rs index 8ae6dc3..0549967 100644 --- a/crates/ptwm-core/src/chain/mod.rs +++ b/crates/ptwm-core/src/chain/mod.rs @@ -5,5 +5,6 @@ pub mod graph; pub mod runtime; pub mod validate; pub mod wire; +pub mod wire_compact; pub use graph::{Chain, ChainEdge, ChainNode, TerminalRef}; diff --git a/crates/ptwm-core/src/chain/wire_compact.rs b/crates/ptwm-core/src/chain/wire_compact.rs new file mode 100644 index 0000000..4ed7021 --- /dev/null +++ b/crates/ptwm-core/src/chain/wire_compact.rs @@ -0,0 +1,892 @@ +//! CW2, the compact chain wire. +//! +//! [`crate::chain::wire`] stores a chain with fixed-width fields and section +//! lengths, which costs 50-70 bytes for a typical production chain. CW2 +//! stores the same [`Chain`] in about 15-20 bytes. The `\x89PTWS` compact +//! frame uses it, and the compact container will too. The legacy wire stays +//! unchanged for `\x89PTWM` containers and Python chain blobs. +//! +//! ```text +//! n_nodes varint 1..=255 +//! n_edges varint 0..=255 +//! n_terms varint 0..=255 +//! node × n_nodes: +//! op_ref varint resolved by the caller (see below) +//! param_len varint 0..=255 +//! params bytes +//! [n_spine] varint 0..=n_edges; present only when n_edges > 0 +//! edge × (n_edges − n_spine): +//! src, src_out, dst, dst_in varint each, 0..=255 +//! etag u8 bit0 ROLE, bit1 VENDOR; other bits must be 0 +//! [role] compact role (ROLE) +//! [vendor] varint len 1..=65535, then bytes (VENDOR) +//! terminal × n_terms: +//! ref u8 node << 4 | out, for node <= 14 and out <= 15; +//! 0xFF = escape, then node u8, out u8 +//! role compact role +//! ``` +//! +//! All varints are canonical unsigned LEB128 ([`crate::varint`]). +//! +//! **op_ref.** The caller maps ops to integers: the `\x89PTWS` frame uses +//! the `OpId` wire id, a container uses an Extension Table index. +//! +//! **Spine edges.** Most chains start with edges `(i, 0) → (i + 1, 0)` with +//! no role override and no vendor bytes. The first `n_spine` edges are of +//! that form for `i = 0..n_spine` and are not stored. The writer always +//! emits the longest such run, so the reader rejects an explicit first edge +//! that would extend it. +//! +//! **Bound Source.** A `Source` node with `param_len = 0` takes its params +//! from the tensor that the chain encodes: the caller passes +//! [`crate::transforms::source::bound_source_params`] of the tensor's shape +//! and dtype. The writer binds a Source exactly when its params equal those +//! bytes, and the reader rejects explicit params that equal them. Real +//! Source params are never empty (ndim and dtype take 3 bytes), so the +//! marker cannot collide with an explicit Source. Chains that differ only +//! in the tensor shape then encode to the same bytes. +//! +//! **Compact roles.** One byte for every role production chains use: +//! +//! | byte | role | +//! |---|---| +//! | `0x00..=0x04` | `Scale { format }` (`ScaleFormat` tag) | +//! | `0x08..=0x0C` | `GlobalScale { format }` | +//! | `0x10..=0x12` | `Nibble { kind }` (`NibbleKind` tag) | +//! | `0x18` | `Value { Fp4E2m1 }` | +//! | `0x19` | `ExponentByte` | +//! | `0x1A` | `Raw` | +//! | `0x1B` | `Index` | +//! | `0x1C` | `Residual { Xor }` | +//! | `0x40 + 8·(of − 1) + index` | `MantissaByte { index, of }`, `index < of <= 8` | +//! | `0x80 + 8·(of − 1) + index` | `IntegerByte { index, of }`, `index < of <= 8` | +//! | `0xC0 + bits − 1` | `Value { IntN { bits } }`, `1 <= bits <= 32` | +//! | `0xFF` | escape, followed by the [`Role::write`] bytes | +//! +//! Other bytes are rejected, and so is an escaped role that has a one-byte +//! form. + +use crate::chain::graph::{Chain, ChainEdge, ChainNode, TerminalRef}; +use crate::error::PtwmCoreError; +use crate::transforms::op::OpId; +use crate::types::role::{NibbleKind, ResidualFormat, Role, ScaleFormat, ValueFormat}; +use crate::varint::{Cursor, write_varint}; + +/// Largest node, edge and terminal count (node indices are `u8`). +const MAX_COUNT: u64 = 255; +/// Largest params length per node (matches the legacy chain wire). +const MAX_PARAMS: u64 = 255; +/// Largest edge vendor-bytes length (matches the legacy chain wire). +const MAX_VENDOR: u64 = 65535; + +const ETAG_ROLE: u8 = 0x01; +const ETAG_VENDOR: u8 = 0x02; + +const TERM_ESCAPE: u8 = 0xFF; +const ROLE_ESCAPE: u8 = 0xFF; + +/// `op_ref` resolver for the `\x89PTWS` frame: the `OpId` wire id. +pub fn op_wire_ref(op: OpId) -> Result { + Ok(u64::from(op.as_u16())) +} + +/// Inverse of [`op_wire_ref`]; unknown ids are rejected. +pub fn op_from_wire_ref(op_ref: u64) -> Result { + let wire = u16::try_from(op_ref).map_err(|_| { + PtwmCoreError::InvalidContainer(format!("chain: op id {op_ref} exceeds u16")) + })?; + OpId::from_u16(wire) +} + +/// Encode `chain` as CW2 into `out`. +/// +/// `bound_source` is the tensor's [`bound_source_params`] (`None` when the +/// tensor cannot bind one); a Source whose params equal it is written as a +/// bound Source. `op_ref` maps each node's op to its stored integer. +/// +/// [`bound_source_params`]: crate::transforms::source::bound_source_params +pub fn write_chain_compact( + chain: &Chain, + bound_source: Option<&[u8]>, + mut op_ref: impl FnMut(OpId) -> Result, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + let err = |msg: String| PtwmCoreError::InvalidContainer(format!("chain: {msg}")); + for (what, n) in [ + ("nodes", chain.nodes.len()), + ("edges", chain.edges.len()), + ("terminals", chain.terminals.len()), + ] { + if n as u64 > MAX_COUNT { + return Err(err(format!("{n} {what} (max {MAX_COUNT})"))); + } + } + if chain.nodes.is_empty() { + return Err(err("no nodes".into())); + } + write_varint(out, chain.nodes.len() as u64); + write_varint(out, chain.edges.len() as u64); + write_varint(out, chain.terminals.len() as u64); + + for node in &chain.nodes { + if node.params.len() as u64 > MAX_PARAMS { + return Err(err(format!( + "node params length {} (max {MAX_PARAMS})", + node.params.len() + ))); + } + write_varint(out, op_ref(node.op)?); + if node.op == OpId::Source { + if node.params.is_empty() { + // An empty Source would read back as a bound Source. + return Err(err("Source node with empty params".into())); + } + if bound_source == Some(node.params.as_slice()) { + write_varint(out, 0); + continue; + } + } + write_varint(out, node.params.len() as u64); + out.extend_from_slice(&node.params); + } + + if !chain.edges.is_empty() { + let n_spine = spine_len(&chain.edges); + write_varint(out, n_spine as u64); + for edge in &chain.edges[n_spine..] { + for v in [ + edge.src_node, + edge.src_output_idx, + edge.dst_node, + edge.dst_input_idx, + ] { + write_varint(out, u64::from(v)); + } + let mut etag = 0u8; + if edge.role_override.is_some() { + etag |= ETAG_ROLE; + } + if !edge.vendor_bytes.is_empty() { + etag |= ETAG_VENDOR; + } + out.push(etag); + if let Some(role) = &edge.role_override { + write_role_compact(role, out)?; + } + if !edge.vendor_bytes.is_empty() { + if edge.vendor_bytes.len() as u64 > MAX_VENDOR { + return Err(err(format!( + "edge vendor bytes length {} (max {MAX_VENDOR})", + edge.vendor_bytes.len() + ))); + } + write_varint(out, edge.vendor_bytes.len() as u64); + out.extend_from_slice(&edge.vendor_bytes); + } + } + } + + for term in &chain.terminals { + if term.node_idx <= 14 && term.output_idx <= 15 { + out.push((term.node_idx << 4) | term.output_idx); + } else { + out.extend_from_slice(&[TERM_ESCAPE, term.node_idx, term.output_idx]); + } + write_role_compact(&term.role, out)?; + } + Ok(()) +} + +/// Decode one CW2 chain at the cursor. The cursor advances past the chain; +/// bytes after it are left for the caller. +/// +/// `bound_source` supplies the params of a bound Source (see +/// [`write_chain_compact`]); a bound Source without it is an error. The +/// returned chain is not validated: callers run `validate_chain`. +pub fn read_chain_compact( + r: &mut Cursor<'_>, + bound_source: Option<&[u8]>, + mut resolve: impl FnMut(u64) -> Result, +) -> Result { + let n_nodes = r.varint_max(MAX_COUNT, "chain node count")? as usize; + if n_nodes == 0 { + return Err(r.err("chain has no nodes")); + } + let n_edges = r.varint_max(MAX_COUNT, "chain edge count")? as usize; + let n_terms = r.varint_max(MAX_COUNT, "chain terminal count")? as usize; + // Each node takes at least 2 bytes and each terminal 2, so a short + // buffer fails here instead of after the allocations below. + if r.remaining() < 2 * n_nodes + 2 * n_terms { + return Err(r.err("chain truncated")); + } + + let mut nodes = Vec::with_capacity(n_nodes); + for _ in 0..n_nodes { + let op_ref = r.varint()?; + let op = resolve(op_ref)?; + let param_len = r.varint_max(MAX_PARAMS, "chain node params length")? as usize; + let params = if op == OpId::Source && param_len == 0 { + bound_source + .ok_or_else(|| { + r.err("chain has a bound Source but the tensor has no bindable shape and dtype") + })? + .to_vec() + } else { + let params = r.take(param_len)?; + if op == OpId::Source && bound_source == Some(params) { + return Err(r.err("chain Source params equal the bound Source (non-canonical)")); + } + params.to_vec() + }; + nodes.push(ChainNode { op, params }); + } + + let mut edges = Vec::with_capacity(n_edges); + if n_edges > 0 { + let n_spine = r.varint_max(n_edges as u64, "chain spine edge count")? as usize; + for i in 0..n_spine { + edges.push(spine_edge(i as u8)); + } + for i in n_spine..n_edges { + let mut idx = [0u8; 4]; + for v in &mut idx { + *v = r.varint_max(MAX_COUNT, "chain edge index")? as u8; + } + let etag = r.u8()?; + if etag & !(ETAG_ROLE | ETAG_VENDOR) != 0 { + return Err(r.err(format_args!("chain edge tag {etag:#04x} has unknown bits"))); + } + let role_override = if etag & ETAG_ROLE != 0 { + Some(read_role_compact(r)?) + } else { + None + }; + let vendor_bytes = if etag & ETAG_VENDOR != 0 { + let len = r.varint_max(MAX_VENDOR, "chain edge vendor length")?; + if len == 0 { + return Err(r.err("chain edge VENDOR flag with no vendor bytes")); + } + r.take(len as usize)?.to_vec() + } else { + Vec::new() + }; + let edge = ChainEdge { + src_node: idx[0], + src_output_idx: idx[1], + dst_node: idx[2], + dst_input_idx: idx[3], + role_override, + vendor_bytes, + }; + if i == n_spine && edge == spine_edge(i as u8) { + return Err(r.err("chain spine edge stored explicitly (non-canonical)")); + } + edges.push(edge); + } + } + + let mut terminals = Vec::with_capacity(n_terms); + for _ in 0..n_terms { + let b = r.u8()?; + let (node_idx, output_idx) = if b == TERM_ESCAPE { + let node = r.u8()?; + let out = r.u8()?; + if node <= 14 && out <= 15 { + return Err(r.err("chain terminal ref escaped needlessly (non-canonical)")); + } + (node, out) + } else if b >> 4 == 15 { + return Err(r.err(format_args!("chain terminal ref byte {b:#04x} is reserved"))); + } else { + (b >> 4, b & 0x0F) + }; + let role = read_role_compact(r)?; + terminals.push(TerminalRef { + node_idx, + output_idx, + role, + }); + } + + Ok(Chain { + nodes, + edges, + terminals, + }) +} + +/// Decode a CW2 chain that must fill `buf` exactly. +pub fn read_chain_compact_exact( + buf: &[u8], + bound_source: Option<&[u8]>, + resolve: impl FnMut(u64) -> Result, +) -> Result { + let mut r = Cursor::new(buf, "compact chain"); + let chain = read_chain_compact(&mut r, bound_source, resolve)?; + r.finish("chain")?; + Ok(chain) +} + +fn spine_edge(i: u8) -> ChainEdge { + ChainEdge { + src_node: i, + src_output_idx: 0, + dst_node: i + 1, + dst_input_idx: 0, + role_override: None, + vendor_bytes: Vec::new(), + } +} + +/// Length of the leading run of spine edges. +fn spine_len(edges: &[ChainEdge]) -> usize { + edges + .iter() + .enumerate() + .take_while(|&(i, e)| i < 255 && *e == spine_edge(i as u8)) + .count() +} + +/// The one-byte form of `role`, if it has one. +fn role_code(role: &Role) -> Option { + let small = + |index: u8, of: u8| (index < of && (1..=8).contains(&of)).then(|| (of - 1) * 8 + index); + match role { + Role::Scale { format } => Some(format.as_u8()), + Role::GlobalScale { format } => Some(0x08 + format.as_u8()), + Role::Nibble { kind } => Some(0x10 + kind.as_u8()), + Role::Value { + format: ValueFormat::Fp4E2m1, + } => Some(0x18), + Role::ExponentByte => Some(0x19), + Role::Raw => Some(0x1A), + Role::Index => Some(0x1B), + Role::Residual { + format: ResidualFormat::Xor, + } => Some(0x1C), + Role::MantissaByte { index, of } => small(*index, *of).map(|c| 0x40 + c), + Role::IntegerByte { index, of } => small(*index, *of).map(|c| 0x80 + c), + Role::Value { + format: ValueFormat::IntN { bits }, + } if (1..=32).contains(bits) => Some(0xC0 + bits - 1), + _ => None, + } +} + +/// Append the compact encoding of `role`. +pub fn write_role_compact(role: &Role, out: &mut Vec) -> Result<(), PtwmCoreError> { + if let Some(code) = role_code(role) { + out.push(code); + return Ok(()); + } + if let Role::Vendor { bytes, .. } = role + && bytes.len() > u16::MAX as usize + { + return Err(PtwmCoreError::InvalidContainer(format!( + "chain: vendor role of {} bytes (max {})", + bytes.len(), + u16::MAX + ))); + } + out.push(ROLE_ESCAPE); + role.write(out); + Ok(()) +} + +/// Read one compact role at the cursor. +pub fn read_role_compact(r: &mut Cursor<'_>) -> Result { + let b = r.u8()?; + let byte_role = |code: u8| ((code & 7), (code >> 3) + 1); + let role = match b { + 0x00..=0x04 => Role::Scale { + format: ScaleFormat::from_u8(b)?, + }, + 0x08..=0x0C => Role::GlobalScale { + format: ScaleFormat::from_u8(b - 0x08)?, + }, + 0x10..=0x12 => Role::Nibble { + kind: NibbleKind::from_u8(b - 0x10)?, + }, + 0x18 => Role::Value { + format: ValueFormat::Fp4E2m1, + }, + 0x19 => Role::ExponentByte, + 0x1A => Role::Raw, + 0x1B => Role::Index, + 0x1C => Role::Residual { + format: ResidualFormat::Xor, + }, + 0x40..=0x7F => { + let (index, of) = byte_role(b - 0x40); + if index >= of { + return Err(r.err(format_args!("role byte {b:#04x} has index >= of"))); + } + Role::MantissaByte { index, of } + } + 0x80..=0xBF => { + let (index, of) = byte_role(b - 0x80); + if index >= of { + return Err(r.err(format_args!("role byte {b:#04x} has index >= of"))); + } + Role::IntegerByte { index, of } + } + 0xC0..=0xDF => Role::Value { + format: ValueFormat::IntN { bits: b - 0xC0 + 1 }, + }, + ROLE_ESCAPE => { + let (role, n) = Role::read(r.rest())?; + r.take(n)?; + if role_code(&role).is_some() { + return Err(r.err("role escaped needlessly (non-canonical)")); + } + role + } + _ => return Err(r.err(format_args!("unknown role byte {b:#04x}"))), + }; + Ok(role) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::transforms::source::bound_source_params; + + fn write(chain: &Chain, bound: Option<&[u8]>) -> Vec { + let mut out = Vec::new(); + write_chain_compact(chain, bound, op_wire_ref, &mut out).unwrap(); + out + } + + fn read(buf: &[u8], bound: Option<&[u8]>) -> Result { + read_chain_compact_exact(buf, bound, op_from_wire_ref) + } + + fn assert_round_trip(label: &str, chain: &Chain, bound: Option<&[u8]>) -> Vec { + let bytes = write(chain, bound); + let back = read(&bytes, bound).unwrap_or_else(|e| panic!("{label}: {e}")); + assert_eq!(&back, chain, "{label}"); + // Canonical: re-encoding gives the same bytes. + assert_eq!(write(&back, bound), bytes, "{label}"); + bytes + } + + fn node(op: OpId, params: &[u8]) -> ChainNode { + ChainNode { + op, + params: params.to_vec(), + } + } + + fn edge(src: u8, out: u8, dst: u8, inp: u8) -> ChainEdge { + ChainEdge { + src_node: src, + src_output_idx: out, + dst_node: dst, + dst_input_idx: inp, + role_override: None, + vendor_bytes: Vec::new(), + } + } + + fn term(node_idx: u8, output_idx: u8, role: Role) -> TerminalRef { + TerminalRef { + node_idx, + output_idx, + role, + } + } + + /// fp32 [64, 64]: Source → BitReorderIeee32 → ByteSplit(4), 4 terminals. + fn fp32_split() -> (Chain, Vec) { + let source = bound_source_params(&[64, 64], 1).unwrap(); + let chain = Chain { + nodes: vec![ + node(OpId::Source, &source), + node(OpId::BitReorderIeee32, &[]), + node(OpId::ByteSplit, &[4]), + ], + edges: vec![edge(0, 0, 1, 0), edge(1, 0, 2, 0)], + terminals: (0..4) + .map(|i| term(2, i, Role::MantissaByte { index: i, of: 4 })) + .collect(), + }; + (chain, source) + } + + #[test] + fn production_chains_round_trip() { + // Same fixture as the symbolic-descriptor test in `runtime.rs`. + let data: &[u8] = include_bytes!("../../tests/fixtures/production_chains.bin"); + let mut pos = 0usize; + let (mut total, mut bound_shorter) = (0usize, 0usize); + let mut legacy_bytes = 0usize; + let mut compact_bytes = 0usize; + while pos < data.len() { + let name_len = u16::from_le_bytes([data[pos], data[pos + 1]]) as usize; + pos += 2; + let label = std::str::from_utf8(&data[pos..pos + name_len]).unwrap(); + pos += name_len; + let blob_len = u32::from_le_bytes(data[pos..pos + 4].try_into().unwrap()) as usize; + pos += 4; + let blob = &data[pos..pos + blob_len]; + pos += blob_len; + total += 1; + let (chain, _) = crate::chain::wire::read_chain_blob(blob).unwrap(); + let explicit = assert_round_trip(label, &chain, None); + // Bind the chain's own Source params, as a writer does when + // they match the tensor. + let own = chain.nodes[0].params.clone(); + let bound = assert_round_trip(label, &chain, Some(&own)); + assert!(bound.len() < explicit.len(), "{label}"); + bound_shorter += 1; + // A different binding leaves the Source explicit. + assert_eq!( + assert_round_trip(label, &chain, Some(&[0, 1, 0][..])), + explicit + ); + + let mut b = crate::extension::ExtensionTableBuilder::new(); + let mut legacy = Vec::new(); + crate::chain::wire::write_chain(&chain, &mut b, &mut legacy).unwrap(); + legacy_bytes += legacy.len(); + compact_bytes += bound.len(); + } + assert!(total > 500, "fixture holds only {total} chains"); + assert_eq!(bound_shorter, total); + // CW2 with a bound Source is far smaller than the legacy wire. + assert!( + compact_bytes * 3 < legacy_bytes, + "{compact_bytes} vs {legacy_bytes}" + ); + } + + #[test] + fn fp32_split_chain_is_small() { + let (chain, source) = fp32_split(); + let bytes = assert_round_trip("fp32", &chain, Some(&source)); + // 3 counts, nodes 2 + 2 + 3, n_spine 1, 4 × (ref + role). + assert_eq!(bytes.len(), 19, "{bytes:02x?}"); + assert_eq!( + bytes, + [ + 3, 2, 4, // counts + 0x00, 0, // bound Source + 0x11, 0, // BitReorderIeee32 + 0x20, 1, 4, // ByteSplit(4) + 2, // both edges are spine edges + 0x20, 0x58, 0x21, 0x59, 0x22, 0x5A, 0x23, 0x5B, + ] + ); + } + + #[test] + fn edge_cases_round_trip() { + let (base, source) = fp32_split(); + + // Non-spine edges, role overrides and vendor bytes. + let mut c = base.clone(); + c.edges[1].role_override = Some(Role::ExponentByte); + c.edges.push(ChainEdge { + vendor_bytes: vec![9; 300], + ..edge(1, 1, 2, 1) + }); + let bytes = assert_round_trip("non-spine", &c, Some(&source)); + assert_eq!(bytes[10], 1, "only the first edge is a spine edge"); + // An edge that matches the spine pattern after a break stays explicit. + let mut c = base.clone(); + c.edges = vec![edge(0, 0, 1, 1), edge(1, 0, 2, 0)]; + assert_round_trip("late spine", &c, None); + + // Escaped terminal refs. + let mut c = base.clone(); + c.terminals.push(term(15, 0, Role::Raw)); + c.terminals.push(term(1, 16, Role::Raw)); + c.terminals.push(term(255, 255, Role::Raw)); + assert_round_trip("terminal escape", &c, None); + + // Every one-byte role and escaped roles. + let mut roles = vec![ + Role::ExponentByte, + Role::Raw, + Role::Index, + Role::Residual { + format: ResidualFormat::Xor, + }, + Role::Residual { + format: ResidualFormat::FloatDelta { dtype: 0x0203 }, + }, + Role::Value { + format: ValueFormat::Fp4E2m1, + }, + Role::Vendor { + tag: 7, + bytes: vec![1, 2, 3], + }, + Role::MantissaByte { index: 9, of: 9 }, + Role::MantissaByte { index: 3, of: 2 }, + Role::IntegerByte { index: 0, of: 0 }, + Role::IntegerByte { index: 0, of: 12 }, + ]; + for f in [ + ScaleFormat::E8M0, + ScaleFormat::E4M3, + ScaleFormat::E5M2, + ScaleFormat::F32, + ScaleFormat::F16, + ] { + roles.push(Role::Scale { format: f }); + roles.push(Role::GlobalScale { format: f }); + } + for kind in [ + NibbleKind::Exponent, + NibbleKind::SignMantissa, + NibbleKind::Value, + ] { + roles.push(Role::Nibble { kind }); + } + for bits in [0u8, 1, 4, 8, 32, 33, 255] { + roles.push(Role::Value { + format: ValueFormat::IntN { bits }, + }); + } + for of in 1..=8u8 { + for index in 0..of { + roles.push(Role::MantissaByte { index, of }); + roles.push(Role::IntegerByte { index, of }); + } + } + let mut c = base.clone(); + c.terminals = roles + .iter() + .enumerate() + .map(|(i, role)| term(2, (i % 200) as u8, role.clone())) + .collect(); + c.edges[0].role_override = Some(Role::Vendor { + tag: 1, + bytes: vec![], + }); + assert_round_trip("roles", &c, None); + for role in &roles { + let mut out = Vec::new(); + write_role_compact(role, &mut out).unwrap(); + let escaped = out[0] == ROLE_ESCAPE; + assert_eq!(escaped, role_code(role).is_none(), "{role:?}"); + if !escaped { + assert_eq!(out.len(), 1, "{role:?}"); + } + } + + // Large op refs (multi-byte varints) through a custom resolver. + let mut out = Vec::new(); + write_chain_compact( + &base, + None, + |op| Ok(u64::from(op.as_u16()) + 100_000), + &mut out, + ) + .unwrap(); + let back = read_chain_compact_exact(&out, None, |r| op_from_wire_ref(r - 100_000)).unwrap(); + assert_eq!(back, base); + + // 0-d tensor Source and a single-node chain. + let scalar = bound_source_params(&[], 1).unwrap(); + let c = Chain { + nodes: vec![node(OpId::Source, &scalar)], + edges: vec![], + terminals: vec![term(0, 0, Role::Raw)], + }; + assert_eq!( + assert_round_trip("scalar", &c, Some(&scalar)), + [1, 0, 1, 0, 0, 0x00, 0x1A] + ); + + // 255 nodes and edges, all spine. + let mut c = base.clone(); + c.nodes = std::iter::once(node(OpId::Source, &source)) + .chain((1..255).map(|_| node(OpId::BytePassthrough, &[]))) + .collect(); + c.edges = (0..254).map(|i| edge(i, 0, i + 1, 0)).collect(); + c.edges.push(edge(254, 1, 0, 0)); + assert_round_trip("max counts", &c, Some(&source)); + } + + #[test] + fn unencodable_chains_are_rejected() { + let (base, _) = fp32_split(); + let mut out = Vec::new(); + let mut c = base.clone(); + c.nodes[0].params.clear(); + let err = write_chain_compact(&c, None, op_wire_ref, &mut out).unwrap_err(); + assert!(err.to_string().contains("empty params"), "{err}"); + let mut c = base.clone(); + c.nodes[1].params = vec![0; 256]; + assert!(write_chain_compact(&c, None, op_wire_ref, &mut out).is_err()); + let mut c = base.clone(); + c.terminals[0].role = Role::Vendor { + tag: 0, + bytes: vec![0; 65536], + }; + assert!(write_chain_compact(&c, None, op_wire_ref, &mut out).is_err()); + let mut c = base.clone(); + c.edges[1].vendor_bytes = vec![0; 65536]; + assert!(write_chain_compact(&c, None, op_wire_ref, &mut out).is_err()); + let c = Chain { + nodes: vec![], + edges: vec![], + terminals: vec![], + }; + assert!(write_chain_compact(&c, None, op_wire_ref, &mut out).is_err()); + } + + fn rejects(label: &str, buf: &[u8], bound: Option<&[u8]>, needle: &str) { + let err = read(buf, bound) + .err() + .unwrap_or_else(|| panic!("{label}: accepted {buf:02x?}")) + .to_string(); + assert!(err.contains(needle), "{label}: {err}"); + } + + #[test] + fn non_canonical_and_corrupt_chains_are_rejected() { + let (chain, source) = fp32_split(); + let good = write(&chain, Some(&source)); + let src = Some(source.as_slice()); + + // Every truncation, and trailing bytes. + for cut in 0..good.len() { + assert!(read(&good[..cut], src).is_err(), "cut at {cut}"); + } + let mut longer = good.clone(); + longer.push(0); + rejects("trailing", &longer, src, "trailing"); + + // Bound Source with no binding available. + rejects("unbound", &good, None, "bound Source"); + // Explicit params equal to the binding. + let explicit = write(&chain, None); + rejects("explicit bound", &explicit, src, "non-canonical"); + + // Non-minimal varint (n_nodes = 3 as 0x83 0x00). + let mut b = vec![0x83, 0x00]; + b.extend_from_slice(&good[1..]); + rejects("overlong", &b, src, "non-minimal"); + + // Counts. + rejects("no nodes", &[0, 0, 0], src, "no nodes"); + rejects("256 nodes", &[0x80, 0x02, 0, 0], src, "exceeds 255"); + + // n_spine > n_edges. + let mut b = good.clone(); + b[10] = 3; + rejects("n_spine", &b, src, "spine edge count"); + + // Unknown op. + let mut b = good.clone(); + b[5] = 0x7F; + assert!(read(&b, src).is_err()); + + // Params longer than 255. + let mut b = good[..6].to_vec(); + b.extend_from_slice(&[0x80, 0x02]); + b.extend_from_slice(&[0; 300]); + rejects("params len", &b, src, "params length"); + + // Reserved terminal ref byte and needless escape. + let mut b = good.clone(); + b[11] = 0xF0; + rejects("ref 0xF0", &b, src, "reserved"); + let mut b = good[..11].to_vec(); + b.extend_from_slice(&[0xFF, 2, 0]); + b.extend_from_slice(&good[12..]); + rejects("ref escape", &b, src, "non-canonical"); + + // Unknown and needlessly escaped roles. + for code in [0x05u8, 0x0D, 0x13, 0x1D, 0x3F, 0xE0, 0xFE] { + let mut b = good.clone(); + b[12] = code; + rejects(&format!("role {code:#04x}"), &b, src, "role byte"); + } + // MantissaByte index >= of: 0x40 + 8*0 + 1 = {index 1, of 1}. + let mut b = good.clone(); + b[12] = 0x41; + rejects("index >= of", &b, src, "index >= of"); + let mut b = good[..12].to_vec(); + b.extend_from_slice(&[0xFF, 0x08]); // Raw through the escape + b.extend_from_slice(&good[13..]); + rejects("role escape", &b, src, "non-canonical"); + + // Explicit edges: unknown etag bits, VENDOR with no bytes, and an + // explicit edge that extends the spine. + let mut c = chain.clone(); + c.edges[0].dst_input_idx = 1; + let ok = write(&c, Some(&source)); + // counts(3) nodes(7) n_spine(1) = 11; edge 0 indices at 11..15, + // etag at 15. + assert_eq!(ok[10], 0); + assert_eq!(&ok[11..16], &[0, 0, 1, 1, 0]); + for etag in [0x04u8, 0x80] { + let mut b = ok.clone(); + b[15] = etag; + rejects("etag", &b, src, "unknown bits"); + } + let mut b = ok[..15].to_vec(); + b.extend_from_slice(&[ETAG_VENDOR, 0]); + b.extend_from_slice(&ok[16..]); + rejects("empty vendor", &b, src, "no vendor bytes"); + let mut b = ok.clone(); + b[14] = 0; // dst_in 1 → 0: now edge 0 is a spine edge + rejects("explicit spine", &b, src, "spine edge stored explicitly"); + } + + #[test] + fn mutated_chains_are_rejected_or_canonical() { + // The invariant of the `chain_wire_compact` fuzz target, on random + // mutations of a few chains: every accepted input re-encodes to + // itself. + let (base, source) = fp32_split(); + let mut with_extras = base.clone(); + with_extras.edges.push(ChainEdge { + role_override: Some(Role::Vendor { + tag: 3, + bytes: vec![1, 2], + }), + vendor_bytes: vec![7], + ..edge(1, 2, 2, 1) + }); + with_extras + .terminals + .push(term(200, 3, Role::IntegerByte { index: 0, of: 9 })); + let seeds = [ + write(&base, Some(&source)), + write(&base, None), + write(&with_extras, Some(&source)), + ]; + let mut x: u64 = 0x2545_F491_4F6C_DD1D; + let mut next = move || { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x + }; + let mut accepted = 0; + for i in 0..60_000 { + let mut b = seeds[i % seeds.len()].clone(); + for _ in 0..1 + next() % 3 { + let at = (next() as usize) % b.len(); + match next() % 3 { + 0 => b[at] = next() as u8, + 1 => b[at] ^= 1 << (next() % 8), + _ => { + b.insert(at, next() as u8); + } + } + } + let bound = (next() % 2 == 0).then_some(source.as_slice()); + if let Ok(chain) = read(&b, bound) { + assert_eq!(write(&chain, bound), b, "{b:02x?}"); + accepted += 1; + } + } + assert!(accepted > 1000, "only {accepted} mutations parsed"); + } +} diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index 69690ad..17040f1 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -8,37 +8,53 @@ //! ```text //! magic: [u8; 5] \x89PTWS //! flags: u8 see FLAG_*; unknown bits are rejected -//! dtype_code: varint +//! dtype_code: varint at most u16::MAX //! input_format: u8 -//! orig_size: varint raw byte count -//! [shape] varint ndim, then varint per dimension (FLAG_SHAPE) -//! [dtype_name] varint len, then UTF-8 bytes (FLAG_DTYPE_NAME) +//! [orig_size] varint raw byte count (FLAG_SIZE) +//! [shape] varint ndim (<= 64), then varint per dimension +//! (FLAG_SHAPE) +//! [dtype_name] varint len, then UTF-8 bytes (FLAG_DTYPE_NAME) //! [delta_ref] [u8; 32] BLAKE3 of the delta reference (FLAG_DELTA_REF) -//! [payload_hash] u64 xxh64 of the raw bytes (FLAG_PAYLOAD_HASH) +//! [payload_hash] u64 xxh64 of the raw bytes (FLAG_PAYLOAD_HASH) //! body //! ``` //! +//! `orig_size` is stored only when the shape does not give it: without +//! `FLAG_SHAPE`, for a dtype code without a known element size, or when +//! `product(shape) × element_size` differs from the raw byte count (for +//! example a byte buffer with a shape in other units). A reader rejects +//! `FLAG_SIZE` when the shape gives the same value. `FLAG_DTYPE_NAME` +//! requires `FLAG_SHAPE` and a name that differs from the canonical name of +//! `dtype_code`; `FLAG_SHAPE` without it requires a known dtype code. +//! //! When `FLAG_STORED` is set, the body is the `orig_size` raw bytes. Else the //! body is a coded tensor: //! //! ```text -//! n_ops: u8 -//! op_ids: [u16; n_ops] OpId wire ids; the chain's u16 op -//! indices refer to this list -//! chain: chain::wire::write_chain bytes +//! chain: CW2 chain (chain::wire_compact); op_ref is the OpId +//! wire id, and a bound Source takes the frame's shape +//! and dtype_code //! per terminal plane, in chain terminal order: -//! codec_id: varint CodecId wire id -//! state_source: u8 0 = none, 1 = inline -//! [state_version] u8 (inline only) -//! [state_len] varint (inline only) -//! [state] bytes (inline only) -//! payload_len: varint +//! head: varint CodecId wire id << 2 | state +//! state 0 = none, 1 = inline, +//! 2 = inline, version 0, empty; +//! 3 is rejected +//! [state_version] u8 (state 1) +//! [state_len] varint (state 1) +//! [state] bytes (state 1) +//! [payload_len] varint every plane but the last; the last +//! payload runs to the end of the frame //! [crc32] u32 CRC32 of the raw plane (FLAG_PLANE_CRC) //! payload: bytes //! ``` //! -//! All integers are little-endian; varints are unsigned LEB128. The frame -//! ends exactly after the body; trailing bytes are rejected. +//! State 1 with version 0 and an empty state is rejected: state 2 is its +//! only encoding. All integers are little-endian; varints are canonical +//! unsigned LEB128 ([`crate::varint`]). Every header and chain field has +//! exactly one encoding. A stored body must be exactly `orig_size` bytes, +//! and a coded body ends with the last plane's payload, so bytes appended to +//! a coded frame reach the last plane's codec, which rejects them or ignores +//! them; the plane CRC and payload hash still guard the decoded bytes. //! //! The plane layout, the decoded plane length and the plane count all come //! from the chain, so the frame does not store them. Planes that need shared @@ -56,14 +72,16 @@ use crate::chain::runtime::{ propagate_terminal_descriptors_bounded, }; use crate::chain::validate::validate_chain; -use crate::chain::wire::{read_chain, write_chain}; +use crate::chain::wire_compact::{ + op_from_wire_ref, op_wire_ref, read_chain_compact, write_chain_compact, +}; use crate::codec::{CodecId, StateSource, codec_for}; use crate::compressor::{CompressorOptions, InputTensor, legacy_plane_layout, plan_model}; use crate::dtype::Dtype; use crate::error::PtwmCoreError; -use crate::extension::dispatch::{BuiltinKind, dispatch_builtin}; -use crate::extension::{ExtensionTable, ExtensionTableBuilder, builtin_entries}; -use crate::transforms::op::{OpId, Plane}; +use crate::transforms::op::Plane; +use crate::transforms::source::bound_source_params; +use crate::varint::{Cursor, write_varint}; use std::sync::Arc; use xxhash_rust::xxh64::xxh64; @@ -84,17 +102,25 @@ pub const FLAG_DTYPE_NAME: u8 = 0x08; pub const FLAG_DELTA_REF: u8 = 0x10; /// The body is the raw bytes, not a coded tensor. pub const FLAG_STORED: u8 = 0x20; +/// The frame stores `orig_size` because the shape does not give it. +pub const FLAG_SIZE: u8 = 0x40; const KNOWN_FLAGS: u8 = FLAG_PAYLOAD_HASH | FLAG_PLANE_CRC | FLAG_SHAPE | FLAG_DTYPE_NAME | FLAG_DELTA_REF - | FLAG_STORED; + | FLAG_STORED + | FLAG_SIZE; /// Upper bound on the dimension count, so a corrupt `ndim` cannot drive a /// large allocation. const MAX_NDIM: u64 = 64; +/// Plane state codes in the plane head. +const STATE_NONE: u64 = 0; +const STATE_INLINE: u64 = 1; +const STATE_INLINE_EMPTY: u64 = 2; + /// A tensor decoded from a compact frame. #[derive(Clone, Debug, PartialEq, Eq)] pub struct CompactTensor { @@ -103,8 +129,7 @@ pub struct CompactTensor { /// Shape, when the writer supplied one. pub shape: Option>, /// Dtype name paired with `shape`. When the frame does not store a name, - /// this is the canonical name of `dtype_code` (or `None` if the code is - /// unknown). + /// this is the canonical name of `dtype_code`. pub dtype_name: Option, /// BLAKE3 hash of the delta reference, when the tensor is delta-coded. pub delta_reference_blake3: Option<[u8; 32]>, @@ -128,6 +153,18 @@ fn canonical_dtype_name(dtype_code: u16) -> Option<&'static str> { .map(Dtype::canonical_name) } +/// Raw byte count implied by a shape and a public dtype code: +/// `product(shape) × element_size`. `None` for an unknown dtype code or on +/// overflow. +pub fn size_from_shape(shape: &[u64], dtype_code: u16) -> Option { + let dtype = Dtype::from_code(u8::try_from(dtype_code).ok()?).ok()?; + shape + .iter() + .try_fold(u64::from(dtype.element_size()), |acc, &d| { + acc.checked_mul(d) + }) +} + /// Encode one tensor as a compact frame. /// /// Runs the same chain and codec search as @@ -143,6 +180,7 @@ pub fn compress_compact( return Ok(None); } let raw = input.raw_bytes; + let orig_size = raw.len() as u64; let mut out = Vec::with_capacity(64); out.extend_from_slice(&MAGIC_COMPACT); @@ -151,10 +189,22 @@ pub fn compress_compact( let mut flags = 0u8; write_varint(&mut out, u64::from(input.dtype_code)); out.push(input.input_format); - write_varint(&mut out, raw.len() as u64); // Shape metadata is recorded only when both shape and dtype name are // known, matching the full container's tensor metadata. - if let (Some(shape), Some(name)) = (&input.shape, &input.dtype_name) { + let shape = match (&input.shape, &input.dtype_name) { + (Some(shape), Some(_)) => Some(shape.as_slice()), + _ => None, + }; + if shape.and_then(|s| size_from_shape(s, input.dtype_code)) != Some(orig_size) { + flags |= FLAG_SIZE; + write_varint(&mut out, orig_size); + } + let mut bound_source = None; + if let (Some(shape), Some(name)) = (shape, &input.dtype_name) { + if shape.len() as u64 > MAX_NDIM { + // The full container has no such limit. + return Ok(None); + } flags |= FLAG_SHAPE; write_varint(&mut out, shape.len() as u64); for &d in shape { @@ -165,6 +215,7 @@ pub fn compress_compact( write_varint(&mut out, name.len() as u64); out.extend_from_slice(name.as_bytes()); } + bound_source = bound_source_params(shape, input.dtype_code); } if let Some(h) = &input.delta_reference_blake3 { flags |= FLAG_DELTA_REF; @@ -175,7 +226,7 @@ pub fn compress_compact( out.extend_from_slice(&xxh64(raw, 0).to_le_bytes()); } - let coded = encode_body(input, opts)?; + let coded = encode_body(input, opts, bound_source.as_deref())?; match coded { None => return Ok(None), Some(body) if body.len() < raw.len() => { @@ -193,11 +244,12 @@ pub fn compress_compact( Ok(Some(out)) } -/// Encode the coded body (ops, chain, planes). `Ok(None)` when a plane needs -/// shared state. +/// Encode the coded body (chain, planes). `Ok(None)` when a plane needs +/// shared or external state. fn encode_body( input: &InputTensor<'_>, opts: &CompressorOptions, + bound_source: Option<&[u8]>, ) -> Result>, PtwmCoreError> { let plan = plan_model(std::slice::from_ref(input), opts)?; if !plan.prelude_entries.is_empty() { @@ -206,51 +258,34 @@ fn encode_body( let chain: &Chain = &input.candidate_chains[plan.pass1[0].chain_idx]; let planes = &plan.chosen_per_tensor[0]; - let mut builder = ExtensionTableBuilder::new(); - let mut chain_bytes = Vec::new(); - write_chain(chain, &mut builder, &mut chain_bytes)?; - let entries = builder.finish(); - let n_ops = u8::try_from(entries.len()).map_err(|_| { - PtwmCoreError::InvalidContainer(format!( - "compact frame: chain uses {} ops (max 255)", - entries.len() - )) - })?; - let mut body = Vec::with_capacity( - 16 + chain_bytes.len() - + planes - .iter() - .map(|p| p.payload_bytes.len() + 16) - .sum::(), + 64 + planes + .iter() + .map(|p| p.payload_bytes.len() + p.encoded_state_bytes.len() + 16) + .sum::(), ); - body.push(n_ops); - for entry in &entries { - let op = match dispatch_builtin(&entry.canonical_id) { - Some(BuiltinKind::Op(op)) => op, - _ => { - return Err(PtwmCoreError::InvalidContainer( - "compact frame: chain op is not a built-in transform".into(), - )); - } - }; - body.extend_from_slice(&op.as_u16().to_le_bytes()); - } - body.extend_from_slice(&chain_bytes); + write_chain_compact(chain, bound_source, op_wire_ref, &mut body)?; - for cp in planes { - write_varint(&mut body, u64::from(cp.codec_id.as_u16())); + for (i, cp) in planes.iter().enumerate() { + let codec = u64::from(cp.codec_id.as_u16()) << 2; match cp.state_source { - StateSource::None => body.push(0), + StateSource::None => write_varint(&mut body, codec | STATE_NONE), + StateSource::Inline + if cp.encoded_state_version == 0 && cp.encoded_state_bytes.is_empty() => + { + write_varint(&mut body, codec | STATE_INLINE_EMPTY) + } StateSource::Inline => { - body.push(1); + write_varint(&mut body, codec | STATE_INLINE); body.push(cp.encoded_state_version); write_varint(&mut body, cp.encoded_state_bytes.len() as u64); body.extend_from_slice(&cp.encoded_state_bytes); } StateSource::Shared | StateSource::External => return Ok(None), } - write_varint(&mut body, cp.payload_bytes.len() as u64); + if i + 1 < planes.len() { + write_varint(&mut body, cp.payload_bytes.len() as u64); + } if opts.emit_plane_crc { body.extend_from_slice(&crc32fast::hash(&cp.raw_plane_for_crc).to_le_bytes()); } @@ -259,32 +294,42 @@ fn encode_body( Ok(Some(body)) } -/// Decode a compact frame. Verifies the payload hash and plane CRCs when the -/// frame carries them. -pub fn decode_compact(buf: &[u8]) -> Result { - let mut r = Reader { buf, pos: 0 }; - if !is_compact(buf) { - return Err(bad("bad magic")); +/// The frame fields before the body. +struct FrameHeader { + flags: u8, + dtype_code: u16, + input_format: u8, + orig_size: u64, + shape: Option>, + dtype_name: Option, + delta_reference_blake3: Option<[u8; 32]>, + payload_hash: Option, +} + +fn read_header(r: &mut Cursor<'_>) -> Result { + if !is_compact(r.rest()) { + return Err(r.err("bad magic")); } - r.pos = MAGIC_COMPACT.len(); + r.take(MAGIC_COMPACT.len())?; let flags = r.u8()?; if flags & !KNOWN_FLAGS != 0 { - return Err(bad(&format!( + return Err(r.err(format_args!( "unknown flag bits {:#04x}", flags & !KNOWN_FLAGS ))); } - let dtype_code = u16::try_from(r.varint()?).map_err(|_| bad("dtype_code exceeds u16"))?; + let dtype_code = r.varint_max(u64::from(u16::MAX), "dtype_code")? as u16; let input_format = r.u8()?; - let orig_size = r.varint()?; + let stored_size = if flags & FLAG_SIZE != 0 { + Some(r.varint()?) + } else { + None + }; let mut shape = None; let mut dtype_name = None; if flags & FLAG_SHAPE != 0 { - let ndim = r.varint()?; - if ndim > MAX_NDIM { - return Err(bad(&format!("ndim {ndim} exceeds {MAX_NDIM}"))); - } + let ndim = r.varint_max(MAX_NDIM, "ndim")?; let mut dims = Vec::with_capacity(ndim as usize); for _ in 0..ndim { dims.push(r.varint()?); @@ -294,35 +339,69 @@ pub fn decode_compact(buf: &[u8]) -> Result { } if flags & FLAG_DTYPE_NAME != 0 { if flags & FLAG_SHAPE == 0 { - return Err(bad("dtype name without shape")); + return Err(r.err("dtype name without shape")); + } + let bytes = r.len_bytes()?; + let name = std::str::from_utf8(bytes).map_err(|_| r.err("dtype name is not UTF-8"))?; + if dtype_name.as_deref() == Some(name) { + return Err(r.err("stored dtype name equals the canonical name")); } - let len = r.len_varint()?; - let bytes = r.take(len)?; - let name = std::str::from_utf8(bytes).map_err(|_| bad("dtype name is not UTF-8"))?; dtype_name = Some(name.to_owned()); + } else if shape.is_some() && dtype_name.is_none() { + return Err(r.err(format_args!( + "shape without a dtype name for unknown dtype_code {dtype_code}" + ))); } + + let derived_size = shape + .as_deref() + .and_then(|s| size_from_shape(s, dtype_code)); + let orig_size = match (stored_size, derived_size) { + (Some(s), Some(d)) if s == d => { + return Err(r.err("stored orig_size equals the size the shape gives")); + } + (Some(s), _) => s, + (None, Some(d)) => d, + (None, None) => return Err(r.err("no orig_size and no shape that gives it")), + }; + let delta_reference_blake3 = if flags & FLAG_DELTA_REF != 0 { - let mut h = [0u8; 32]; - h.copy_from_slice(r.take(32)?); - Some(h) + Some(r.array::<32>()?) } else { None }; let payload_hash = if flags & FLAG_PAYLOAD_HASH != 0 { - Some(u64::from_le_bytes(r.take(8)?.try_into().unwrap())) + Some(r.u64_le()?) } else { None }; + Ok(FrameHeader { + flags, + dtype_code, + input_format, + orig_size, + shape, + dtype_name, + delta_reference_blake3, + payload_hash, + }) +} + +/// Decode a compact frame. Verifies the payload hash and plane CRCs when the +/// frame carries them. +pub fn decode_compact(buf: &[u8]) -> Result { + let mut r = Cursor::new(buf, "compact frame"); + let h = read_header(&mut r)?; + let orig_size = h.orig_size; let mut payload_len = 0u64; - let raw = if flags & FLAG_STORED != 0 { - if flags & FLAG_PLANE_CRC != 0 { - return Err(bad("stored frame with plane CRC flag")); + let raw = if h.flags & FLAG_STORED != 0 { + if h.flags & FLAG_PLANE_CRC != 0 { + return Err(r.err("stored frame with plane CRC flag")); } - let rest = r.rest(); - r.pos = buf.len(); + let rest = r.take(r.remaining())?; if rest.len() as u64 != orig_size { - return Err(bad(&format!( + return Err(r.err(format_args!( "stored body is {} bytes, header says {orig_size}", rest.len() ))); @@ -330,55 +409,45 @@ pub fn decode_compact(buf: &[u8]) -> Result { payload_len = orig_size; rest.to_vec() } else { + let bound = h + .shape + .as_deref() + .and_then(|s| bound_source_params(s, h.dtype_code)); decode_body( &mut r, - flags & FLAG_PLANE_CRC != 0, + h.flags & FLAG_PLANE_CRC != 0, orig_size, + bound.as_deref(), &mut payload_len, )? }; - if r.pos != buf.len() { - return Err(bad("trailing bytes after body")); - } + r.finish("body")?; - if let Some(expected) = payload_hash + if let Some(expected) = h.payload_hash && xxh64(&raw, 0) != expected { - return Err(bad("payload hash mismatch")); + return Err(r.err("payload hash mismatch")); } Ok(CompactTensor { - dtype_code, - input_format, - shape, - dtype_name, - delta_reference_blake3, + dtype_code: h.dtype_code, + input_format: h.input_format, + shape: h.shape, + dtype_name: h.dtype_name, + delta_reference_blake3: h.delta_reference_blake3, raw, payload_len, }) } fn decode_body( - r: &mut Reader<'_>, + r: &mut Cursor<'_>, has_crc: bool, orig_size: u64, + bound_source: Option<&[u8]>, payload_len_out: &mut u64, ) -> Result, PtwmCoreError> { - let n_ops = r.u8()? as usize; - let all = builtin_entries(); - let mut table = ExtensionTable::empty(); - for _ in 0..n_ops { - let wire = u16::from_le_bytes(r.take(2)?.try_into().unwrap()); - let op = OpId::from_u16(wire)?; - let id = op.canonical_id(); - let entry = all - .iter() - .find(|e| e.canonical_id == id) - .ok_or_else(|| bad(&format!("op {wire:#06x} has no built-in entry")))?; - table.entries.push(entry.clone()); - } - let (chain, consumed) = read_chain(r.rest(), &table)?; - r.pos += consumed; + let chain = read_chain_compact(r, bound_source, op_from_wire_ref)?; // Writers check that the chain's Source declares exactly the raw byte // count, so a frame whose Source disagrees with `orig_size` is corrupt. // Checking before `validate_chain` and the descriptor walk (both call @@ -386,7 +455,7 @@ fn decode_body( // planes (and codec allocations) beyond the tensor. let source = chain_source_descriptor(&chain)?; if source.length_bytes != orig_size { - return Err(bad(&format!( + return Err(r.err(format_args!( "chain Source declares {} bytes, header says {orig_size}", source.length_bytes ))); @@ -394,34 +463,36 @@ fn decode_body( validate_chain(&chain, 0)?; let descriptors = propagate_terminal_descriptors_bounded(&chain, max_plane_len(orig_size))?; - let mut terminal_planes: Vec = Vec::with_capacity(descriptors.len()); - for descriptor in descriptors { - let wire = u16::try_from(r.varint()?).map_err(|_| bad("codec id exceeds u16"))?; - let codec_id = - CodecId::from_u16(wire).ok_or_else(|| bad(&format!("unknown codec {wire:#06x}")))?; - let (state_version, state) = match r.u8()? { - 0 => (0u8, &[][..]), - 1 => { + let n_planes = descriptors.len(); + let mut terminal_planes: Vec = Vec::with_capacity(n_planes); + for (i, descriptor) in descriptors.into_iter().enumerate() { + let head = r.varint()?; + let wire = u16::try_from(head >> 2).map_err(|_| r.err("codec id exceeds u16"))?; + let codec_id = CodecId::from_u16(wire) + .ok_or_else(|| r.err(format_args!("unknown codec {wire:#06x}")))?; + let (state_version, state) = match head & 3 { + STATE_NONE | STATE_INLINE_EMPTY => (0u8, &[][..]), + STATE_INLINE => { let v = r.u8()?; - let len = r.len_varint()?; - (v, r.take(len)?) + let state = r.len_bytes()?; + if v == 0 && state.is_empty() { + return Err(r.err("empty version-0 inline state must use state code 2")); + } + (v, state) } - s => return Err(bad(&format!("unknown state source {s}"))), - }; - let payload_len = r.len_varint()?; - let crc = if has_crc { - Some(u32::from_le_bytes(r.take(4)?.try_into().unwrap())) - } else { - None + s => return Err(r.err(format_args!("unknown state code {s}"))), }; - let payload = r.take(payload_len)?; + let last = i + 1 == n_planes; + let payload_len = if last { None } else { Some(r.len_varint()?) }; + let crc = if has_crc { Some(r.u32_le()?) } else { None }; + let payload = r.take(payload_len.unwrap_or(r.remaining()))?; *payload_len_out += (state.len() + payload.len()) as u64; let codec = codec_for(codec_id) - .ok_or_else(|| bad(&format!("no codec implementation for {codec_id:?}")))?; + .ok_or_else(|| r.err(format_args!("no codec implementation for {codec_id:?}")))?; let layout = legacy_plane_layout(descriptor.layout); let decoded_len = usize::try_from(descriptor.length_bytes) - .map_err(|_| bad("plane length overflows usize"))?; + .map_err(|_| r.err("plane length overflows usize"))?; // The Source is exact, so every plane decodes to its descriptor's // length. The codec stops at that length, so a corrupt payload cannot // allocate past it. @@ -434,7 +505,7 @@ fn decode_body( decoded_len, )?; if bytes.len() != decoded_len { - return Err(bad(&format!( + return Err(r.err(format_args!( "plane decoded to {} bytes, chain says {decoded_len}", bytes.len() ))); @@ -442,7 +513,7 @@ fn decode_body( if let Some(expected) = crc && crc32fast::hash(&bytes) != expected { - return Err(bad("plane CRC mismatch")); + return Err(r.err("plane CRC mismatch")); } terminal_planes.push(Plane { bytes: Arc::from(bytes.into_boxed_slice()), @@ -456,7 +527,7 @@ fn decode_body( }; let source = inverse_chain(&chain, &ctx)?; if source.bytes.len() as u64 != orig_size { - return Err(bad(&format!( + return Err(r.err(format_args!( "decoded {} bytes, header says {orig_size}", source.bytes.len() ))); @@ -464,74 +535,11 @@ fn decode_body( Ok(source.bytes.as_ref().to_vec()) } -fn bad(msg: &str) -> PtwmCoreError { - PtwmCoreError::InvalidContainer(format!("compact frame: {msg}")) -} - -fn write_varint(out: &mut Vec, mut v: u64) { - while v >= 0x80 { - out.push((v as u8) | 0x80); - v >>= 7; - } - out.push(v as u8); -} - -struct Reader<'a> { - buf: &'a [u8], - pos: usize, -} - -impl<'a> Reader<'a> { - fn take(&mut self, n: usize) -> Result<&'a [u8], PtwmCoreError> { - let end = self - .pos - .checked_add(n) - .filter(|&e| e <= self.buf.len()) - .ok_or_else(|| bad("truncated"))?; - let s = &self.buf[self.pos..end]; - self.pos = end; - Ok(s) - } - - fn u8(&mut self) -> Result { - Ok(self.take(1)?[0]) - } - - fn rest(&self) -> &'a [u8] { - &self.buf[self.pos..] - } - - fn varint(&mut self) -> Result { - let mut v: u64 = 0; - for i in 0..10 { - let b = self.u8()?; - let bits = u64::from(b & 0x7F); - if i == 9 && b > 1 { - return Err(bad("varint overflows u64")); - } - v |= bits << (7 * i); - if b & 0x80 == 0 { - return Ok(v); - } - } - Err(bad("varint overflows u64")) - } - - /// A varint used as a byte length; must fit in the remaining buffer. - fn len_varint(&mut self) -> Result { - let v = self.varint()?; - let remaining = (self.buf.len() - self.pos) as u64; - if v > remaining { - return Err(bad("length exceeds remaining bytes")); - } - Ok(v as usize) - } -} - #[cfg(test)] mod tests { use super::*; use crate::chain::graph::{ChainEdge, ChainNode, TerminalRef}; + use crate::transforms::op::OpId; use crate::types::role::Role; const FLAGS_AT: usize = MAGIC_COMPACT.len(); @@ -793,14 +801,273 @@ mod tests { } } + /// Source(fp32, bound) → BitReorderIeee32 → ByteSplit(4), the + /// production fp32 split chain. + fn fp32_chain(shape: &[u64]) -> Chain { + let mut terminals: Vec = (0..3) + .map(|i| TerminalRef { + node_idx: 2, + output_idx: i, + role: Role::MantissaByte { index: i, of: 4 }, + }) + .collect(); + terminals.push(TerminalRef { + node_idx: 2, + output_idx: 3, + role: Role::ExponentByte, + }); + let spine = |i: u8| ChainEdge { + src_node: i, + src_output_idx: 0, + dst_node: i + 1, + dst_input_idx: 0, + role_override: None, + vendor_bytes: Vec::new(), + }; + Chain { + nodes: vec![ + ChainNode { + op: OpId::Source, + params: bound_source_params(shape, 1).unwrap(), + }, + ChainNode { + op: OpId::BitReorderIeee32, + params: Vec::new(), + }, + ChainNode { + op: OpId::ByteSplit, + params: vec![4], + }, + ], + edges: vec![spine(0), spine(1)], + terminals, + } + } + + fn fp32_raw(n: usize) -> Vec { + let mut x: u64 = 0x9E37_79B9_7F4A_7C15; + (0..n) + .flat_map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + // Normal-ish magnitudes: a narrow exponent range. + let v = ((x >> 40) as f32 / (1u64 << 24) as f32 - 0.5) * 4.0; + v.to_le_bytes() + }) + .collect() + } + + fn fp32_input<'a>(raw: &'a [u8], shape: &[u64]) -> InputTensor<'a> { + InputTensor { + name: "_".into(), + candidate_chains: vec![fp32_chain(shape)], + dtype_code: 1, + input_format: 1, + raw_bytes: raw, + shape: Some(shape.to_vec()), + dtype_name: Some("float32".into()), + delta_reference_blake3: None, + } + } + + /// Byte offset where the frame body starts. + fn body_start(blob: &[u8]) -> usize { + let mut r = Cursor::new(blob, "test"); + read_header(&mut r).unwrap(); + r.pos() + } + + /// Byte offset of the first plane head (after the chain). + fn planes_start(blob: &[u8]) -> usize { + let mut r = Cursor::new(blob, "test"); + let h = read_header(&mut r).unwrap(); + let bound = bound_source_params(h.shape.as_deref().unwrap_or(&[]), h.dtype_code); + let bound = h.shape.as_ref().and(bound); + read_chain_compact(&mut r, bound.as_deref(), op_from_wire_ref).unwrap(); + r.pos() + } + + #[test] + fn fp32_frame_binds_source_and_stays_small() { + let raw = fp32_raw(4096); + for (h, c) in [(true, true), (false, false)] { + let blob = compress_compact(&fp32_input(&raw, &[64, 64]), &opts(h, c)) + .unwrap() + .unwrap(); + let flags = blob[FLAGS_AT]; + assert_eq!(flags & (FLAG_STORED | FLAG_SIZE), 0, "{flags:#04x}"); + let t = decode_compact(&blob).unwrap(); + assert_eq!(t.raw, raw); + assert_eq!(t.shape, Some(vec![64, 64])); + assert_eq!(t.dtype_name.as_deref(), Some("float32")); + let overhead = blob.len() as u64 - t.payload_len; + // magic 5, flags 1, dtype 1, format 1, shape 3, chain 19, + // 4 plane heads, 3 payload lengths; plus hash 8 and 4 CRCs. + let limit = if h { 33 + 8 + 4 + 16 + 8 } else { 33 + 4 + 8 }; + assert!( + overhead <= limit, + "overhead {overhead} > {limit} (hash {h})" + ); + // The Source dims are not stored. + let dims: Vec = [&64u32.to_le_bytes()[..], &64u32.to_le_bytes()].concat(); + assert!(!blob.windows(8).any(|w| w == dims.as_slice())); + } + // Chains for other shapes encode to the same bytes. + let raw2 = fp32_raw(1000); + let a = compress_compact(&fp32_input(&raw, &[64, 64]), &opts(false, false)) + .unwrap() + .unwrap(); + let b = compress_compact(&fp32_input(&raw2, &[10, 100]), &opts(false, false)) + .unwrap() + .unwrap(); + let chain = |blob: &[u8]| blob[body_start(blob)..planes_start(blob)].to_vec(); + assert_eq!(chain(&a), chain(&b)); + } + + #[test] + fn bound_source_needs_the_frame_shape() { + let raw = fp32_raw(256); + let blob = compress_compact(&fp32_input(&raw, &[256]), &opts(false, false)) + .unwrap() + .unwrap(); + // magic, flags, dtype 1, format 1, shape (1, 256 = 0x80 0x02), body. + assert_eq!( + &blob[FLAGS_AT..FLAGS_AT + 7], + &[FLAG_SHAPE, 1, 1, 1, 0x80, 0x02, 3] + ); + // The same body in a frame without a shape: the size is explicit and + // the bound Source has nothing to bind to. + let mut b = blob[..FLAGS_AT].to_vec(); + b.extend_from_slice(&[FLAG_SIZE, 1, 1]); + write_varint(&mut b, raw.len() as u64); + b.extend_from_slice(&blob[FLAGS_AT + 6..]); + let err = decode_compact(&b).unwrap_err().to_string(); + assert!(err.contains("bound Source"), "{err}"); + } + + #[test] + fn orig_size_is_stored_only_when_the_shape_does_not_give_it() { + // A byte count the shape does not give: stored. + let raw = vec![3u8; 64]; + let mut inp = input(&raw, Some(vec![10]), Some("float32")); + inp.dtype_code = 1; + let blob = compress_compact(&inp, &opts(true, true)).unwrap().unwrap(); + assert_ne!(blob[FLAGS_AT] & FLAG_SIZE, 0); + let t = decode_compact(&blob).unwrap(); + assert_eq!((t.raw, t.shape), (raw.clone(), Some(vec![10]))); + + // Packed fp4 and 0-d tensors derive it. + for (shape, code, name, len) in [ + (vec![16u64], 31u16, "float4_e2m1fn_x2", 16usize), + (vec![], 1, "float32", 4), + (vec![0, 7], 6, "bfloat16", 0), + ] { + let raw = vec![0x11u8; len]; + let mut inp = input(&raw, Some(shape.clone()), Some(name)); + inp.dtype_code = code; + let blob = compress_compact(&inp, &opts(true, false)).unwrap().unwrap(); + assert_eq!(blob[FLAGS_AT] & FLAG_SIZE, 0, "{shape:?}"); + let t = decode_compact(&blob).unwrap(); + assert_eq!((t.raw, t.shape), (raw, Some(shape))); + } + } + + #[test] + fn non_canonical_headers_are_rejected() { + let raw = fp32_raw(256); + let blob = compress_compact(&fp32_input(&raw, &[256]), &opts(false, false)) + .unwrap() + .unwrap(); + let rest = &blob[FLAGS_AT + 6..]; + let frame = |head: &[u8]| [&MAGIC_COMPACT[..], head, rest].concat(); + let check = |head: &[u8], needle: &str| { + let err = decode_compact(&frame(head)).unwrap_err().to_string(); + assert!(err.contains(needle), "{head:02x?}: {err}"); + }; + // The original decodes. + decode_compact(&frame(&[FLAG_SHAPE, 1, 1, 1, 0x80, 0x02])).unwrap(); + // orig_size stored although the shape gives it. + check( + &[FLAG_SHAPE | FLAG_SIZE, 1, 1, 0x80, 0x08, 1, 0x80, 0x02], + "equals the size", + ); + // Neither orig_size nor shape. + check(&[0, 1, 1], "no orig_size"); + // Dtype name equal to the canonical name. + let mut h = vec![FLAG_SHAPE | FLAG_DTYPE_NAME, 1, 1, 1, 0x80, 0x02, 7]; + h.extend_from_slice(b"float32"); + check(&h, "canonical name"); + // Shape without a name for a dtype code with no canonical name. + check( + &[ + FLAG_SHAPE | FLAG_SIZE, + 0xC8, + 0x01, + 1, + 0x80, + 0x08, + 1, + 0x80, + 0x02, + ], + "dtype name", + ); + // Non-minimal varints. + check(&[FLAG_SHAPE, 0x81, 0x00, 1, 1, 0x80, 0x02], "non-minimal"); + check(&[FLAG_SHAPE, 1, 1, 0x81, 0x00, 0x80, 0x02], "non-minimal"); + // dtype_code above u16, unknown flag bit. + check(&[FLAG_SIZE, 0x80, 0x80, 0x04, 1, 0x80, 0x08], "dtype_code"); + check(&[FLAG_SHAPE | 0x80, 1, 1, 1, 0x80, 0x02], "unknown flag"); + // ndim above 64. + check(&[FLAG_SHAPE, 1, 1, 65], "ndim"); + } + + #[test] + fn plane_state_codes_are_canonical() { + let raw: Vec = (0..4096u32).map(|i| (i % 7) as u8).collect(); + let mut o = opts(false, false); + o.forced_codec = Some(CodecId::Zstd); + let blob = compress_compact(&input(&raw, None, None), &o) + .unwrap() + .unwrap(); + assert_eq!(decode_compact(&blob).unwrap().raw, raw); + let at = planes_start(&blob); + // Zstd carries an empty version-0 inline state: state code 2. + assert_eq!(blob[at], (4 << 2) | 2); + let with_head = |head: &[u8]| [&blob[..at], head, &blob[at + 1..]].concat(); + // The same plane with state code 0 decodes (decode treats both alike). + assert_eq!(decode_compact(&with_head(&[4 << 2])).unwrap().raw, raw); + // State code 1 with an empty version-0 state is non-canonical. + let err = decode_compact(&with_head(&[(4 << 2) | 1, 0, 0])) + .unwrap_err() + .to_string(); + assert!(err.contains("state code 2"), "{err}"); + let err = decode_compact(&with_head(&[(4 << 2) | 3])) + .unwrap_err() + .to_string(); + assert!(err.contains("unknown state code"), "{err}"); + } + #[test] - fn varint_round_trips() { - for v in [0u64, 1, 127, 128, 300, u32::MAX as u64, u64::MAX] { - let mut out = Vec::new(); - write_varint(&mut out, v); - let mut r = Reader { buf: &out, pos: 0 }; - assert_eq!(r.varint().unwrap(), v); - assert_eq!(r.pos, out.len()); + fn last_plane_runs_to_the_end_of_the_frame() { + let raw = fp32_raw(1024); + let blob = compress_compact(&fp32_input(&raw, &[1024]), &opts(true, true)) + .unwrap() + .unwrap(); + assert_eq!(decode_compact(&blob).unwrap().raw, raw); + // Every truncation fails: a short last payload fails its codec, CRC + // or the payload hash. + for cut in 0..blob.len() { + assert!(decode_compact(&blob[..cut]).is_err(), "cut at {cut}"); + } + // A byte appended to the frame becomes part of the last payload. + // Whether that fails depends on the codec; it must never return + // other data. + let mut longer = blob.clone(); + longer.push(0); + if let Ok(t) = decode_compact(&longer) { + assert_eq!(t.raw, raw); } } } diff --git a/crates/ptwm-core/src/lib.rs b/crates/ptwm-core/src/lib.rs index aa232c3..b13fe6d 100644 --- a/crates/ptwm-core/src/lib.rs +++ b/crates/ptwm-core/src/lib.rs @@ -48,6 +48,7 @@ pub mod transcode; pub mod transforms; pub mod trust; pub mod types; +pub mod varint; pub use dtype::{Dtype, PreprocessingModes}; pub use error::PtwmCoreError; diff --git a/crates/ptwm-core/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index b510cb5..6aa5a5d 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -63,6 +63,47 @@ pub fn storage_bytes_per_element(dtype_code: u16) -> u64 { element_width_for(dtype_code).bits_per_element() as u64 / 8 } +/// Chain-internal Source dtype for a public dtype code (`Dtype` registry), +/// mirroring `_CHAIN_DTYPE` in `python/ptwm/preprocessing/_chains.py`. +/// `None` for codes without a chain dtype (for example `bool`). +pub fn chain_dtype_for_public(dtype_code: u16) -> Option { + Some(match dtype_code { + 1 | 2 => 0x0003, // float32, float + 4 | 5 => 0x0002, // float16, half + 6 => 0x000F, // bfloat16 + 13 => 0x0006, // uint8 + 15 => 0x000A, // uint32 + 17 => 0x0005, // int8 + 18 | 19 => 0x0007, // int16, short + 20 | 21 => 0x0009, // int32, int + 22 | 23 => 0x000B, // int64, long + 29 => 0x0010, // float8_e4m3fn + 30 => 0x0011, // float8_e5m2 + 31 => DTYPE_FP4_E2M1FN_X2, + _ => return None, + }) +} + +/// Source params for a tensor's own shape and public dtype code: the params +/// a *bound* Source stands for in the compact chain wire +/// (`chain::wire_compact`). Same layout as [`Source::write_params`]: +/// `u8` ndim, `u32` LE dims, `u16` LE chain dtype. +/// +/// `None` when the tensor cannot bind a Source: more than +/// [`MAX_SOURCE_DIMS`] dimensions, a dimension above `u32::MAX`, or a dtype +/// without a chain dtype. +pub fn bound_source_params(shape: &[u64], dtype_code: u16) -> Option> { + let chain_dtype = chain_dtype_for_public(dtype_code)?; + let ndim = u8::try_from(shape.len()).ok()?; + let mut out = Vec::with_capacity(3 + 4 * shape.len()); + out.push(ndim); + for &d in shape { + out.extend_from_slice(&u32::try_from(d).ok()?.to_le_bytes()); + } + out.extend_from_slice(&chain_dtype.to_le_bytes()); + Some(out) +} + /// Mandatory entry node of every PPG chain. Takes no inputs (the runtime /// injects raw tensor bytes externally) and emits one `Role::Raw` plane /// whose descriptor follows from the tensor `shape` and `dtype_code`. @@ -235,6 +276,37 @@ mod tests { } } + #[test] + fn bound_source_params_match_write_params() { + let params = bound_source_params(&[4, 8, 16], 4).unwrap(); + let mut expected = Vec::new(); + make_source().write_params(&mut expected); + assert_eq!(params, expected); + // 0-d tensor: no dims. + assert_eq!( + bound_source_params(&[], 1).unwrap(), + [0u8, 0x03, 0x00].to_vec() + ); + // Unbindable: no chain dtype, a dim above u32, too many dims. + assert_eq!(bound_source_params(&[4], 24), None); + assert_eq!(bound_source_params(&[1 << 32], 1), None); + assert_eq!(bound_source_params(&[1; 256], 1), None); + // Every public dtype with a chain dtype maps to a code whose element + // width matches the public element size. + for code in 0..=255u8 { + let Ok(d) = crate::dtype::Dtype::from_code(code) else { + continue; + }; + if let Some(c) = chain_dtype_for_public(u16::from(code)) { + assert_eq!( + storage_bytes_per_element(c), + u64::from(d.element_size()), + "{d:?}" + ); + } + } + } + #[test] fn propagate_descriptors_with_empty_inputs_returns_one_plane() { let src = make_source(); diff --git a/crates/ptwm-core/src/types/role.rs b/crates/ptwm-core/src/types/role.rs index 8c957ed..21bf9ba 100644 --- a/crates/ptwm-core/src/types/role.rs +++ b/crates/ptwm-core/src/types/role.rs @@ -13,7 +13,7 @@ pub enum ScaleFormat { } impl ScaleFormat { - fn as_u8(self) -> u8 { + pub(crate) fn as_u8(self) -> u8 { match self { ScaleFormat::E8M0 => 0, ScaleFormat::E4M3 => 1, @@ -22,7 +22,7 @@ impl ScaleFormat { ScaleFormat::F16 => 4, } } - fn from_u8(v: u8) -> Result { + pub(crate) fn from_u8(v: u8) -> Result { match v { 0 => Ok(ScaleFormat::E8M0), 1 => Ok(ScaleFormat::E4M3), @@ -83,14 +83,14 @@ pub enum NibbleKind { } impl NibbleKind { - fn as_u8(self) -> u8 { + pub(crate) fn as_u8(self) -> u8 { match self { NibbleKind::Exponent => 0, NibbleKind::SignMantissa => 1, NibbleKind::Value => 2, } } - fn from_u8(v: u8) -> Result { + pub(crate) fn from_u8(v: u8) -> Result { match v { 0 => Ok(NibbleKind::Exponent), 1 => Ok(NibbleKind::SignMantissa), diff --git a/crates/ptwm-core/src/varint.rs b/crates/ptwm-core/src/varint.rs new file mode 100644 index 0000000..7bd23a6 --- /dev/null +++ b/crates/ptwm-core/src/varint.rs @@ -0,0 +1,263 @@ +//! Unsigned LEB128 varints and a bounds-checked byte cursor for the compact +//! wire formats (the `\x89PTWS` frame and the compact chain wire). +//! +//! Parsing is canonical: a varint must use the fewest bytes that hold its +//! value (no trailing `0x80`/`0x00` groups), is at most 10 bytes long, and +//! must fit in `u64`. Together with the other fail-closed checks of the +//! compact formats, this makes every accepted byte string the only encoding +//! of its value, so `write(parse(b)) == b`. + +use crate::error::PtwmCoreError; + +/// Longest LEB128 encoding of a `u64`. +pub const MAX_VARINT_LEN: usize = 10; + +/// Append `v` as an unsigned LEB128 varint. +pub fn write_varint(out: &mut Vec, mut v: u64) { + while v >= 0x80 { + out.push((v as u8) | 0x80); + v >>= 7; + } + out.push(v as u8); +} + +/// Number of bytes [`write_varint`] uses for `v`. +pub fn varint_len(v: u64) -> usize { + let bits = 64 - v.leading_zeros() as usize; + bits.div_ceil(7).max(1) +} + +/// Why a varint did not parse. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VarintError { + /// The buffer ended inside the varint. + Truncated, + /// The value does not fit in `u64` (more than 10 bytes, or a 10th byte + /// above 1). + Overflow, + /// The varint has a redundant trailing zero group. + NonMinimal, +} + +impl VarintError { + fn message(self) -> &'static str { + match self { + VarintError::Truncated => "truncated varint", + VarintError::Overflow => "varint overflows u64", + VarintError::NonMinimal => "non-minimal varint", + } + } +} + +/// Read one canonical varint from `buf` at `*pos` and advance `*pos`. +pub fn read_varint(buf: &[u8], pos: &mut usize) -> Result { + let mut v: u64 = 0; + for i in 0..MAX_VARINT_LEN { + let b = *buf.get(*pos + i).ok_or(VarintError::Truncated)?; + if i == MAX_VARINT_LEN - 1 && b > 1 { + return Err(VarintError::Overflow); + } + v |= u64::from(b & 0x7F) << (7 * i); + if b & 0x80 == 0 { + // A zero final group after the first byte adds nothing: the + // value has a shorter encoding. + if i > 0 && b == 0 { + return Err(VarintError::NonMinimal); + } + *pos += i + 1; + return Ok(v); + } + } + Err(VarintError::Overflow) +} + +/// A bounds-checked reader over one wire section. Every error is an +/// [`PtwmCoreError::InvalidContainer`] prefixed with the cursor's context +/// (for example `"compact frame"`). +pub struct Cursor<'a> { + buf: &'a [u8], + pos: usize, + context: &'static str, +} + +impl<'a> Cursor<'a> { + pub fn new(buf: &'a [u8], context: &'static str) -> Self { + Cursor { + buf, + pos: 0, + context, + } + } + + /// Error carrying this cursor's context. + pub fn err(&self, msg: impl std::fmt::Display) -> PtwmCoreError { + PtwmCoreError::InvalidContainer(format!("{}: {msg}", self.context)) + } + + pub fn pos(&self) -> usize { + self.pos + } + + /// Move to an absolute position at or before the end of the buffer. + pub fn set_pos(&mut self, pos: usize) { + assert!(pos <= self.buf.len(), "cursor position past the end"); + self.pos = pos; + } + + pub fn remaining(&self) -> usize { + self.buf.len() - self.pos + } + + pub fn is_empty(&self) -> bool { + self.pos == self.buf.len() + } + + /// The unread bytes, without consuming them. + pub fn rest(&self) -> &'a [u8] { + &self.buf[self.pos..] + } + + pub fn take(&mut self, n: usize) -> Result<&'a [u8], PtwmCoreError> { + if n > self.remaining() { + return Err(self.err("truncated")); + } + let s = &self.buf[self.pos..self.pos + n]; + self.pos += n; + Ok(s) + } + + pub fn array(&mut self) -> Result<[u8; N], PtwmCoreError> { + Ok(self.take(N)?.try_into().expect("take returned N bytes")) + } + + pub fn u8(&mut self) -> Result { + Ok(self.take(1)?[0]) + } + + pub fn u32_le(&mut self) -> Result { + Ok(u32::from_le_bytes(self.array()?)) + } + + pub fn u64_le(&mut self) -> Result { + Ok(u64::from_le_bytes(self.array()?)) + } + + /// One canonical varint. + pub fn varint(&mut self) -> Result { + read_varint(self.buf, &mut self.pos).map_err(|e| self.err(e.message())) + } + + /// A varint that must not exceed `max`; `what` names the field. + pub fn varint_max(&mut self, max: u64, what: &str) -> Result { + let v = self.varint()?; + if v > max { + return Err(self.err(format_args!("{what} {v} exceeds {max}"))); + } + Ok(v) + } + + /// A varint used as a byte length: it must not exceed the bytes left. + pub fn len_varint(&mut self) -> Result { + let v = self.varint()?; + if v > self.remaining() as u64 { + return Err(self.err("length exceeds remaining bytes")); + } + Ok(v as usize) + } + + /// A length-prefixed byte string (`len_varint` + bytes). + pub fn len_bytes(&mut self) -> Result<&'a [u8], PtwmCoreError> { + let n = self.len_varint()?; + self.take(n) + } + + /// Fail unless every byte was consumed. + pub fn finish(&self, what: &str) -> Result<(), PtwmCoreError> { + if self.is_empty() { + Ok(()) + } else { + Err(self.err(format_args!( + "{} trailing bytes after {what}", + self.remaining() + ))) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn varint_round_trips() { + for v in [ + 0u64, + 1, + 127, + 128, + 300, + 16383, + 16384, + u32::MAX as u64, + 1 << 63, + u64::MAX, + ] { + let mut out = Vec::new(); + write_varint(&mut out, v); + assert_eq!(out.len(), varint_len(v), "{v}"); + let mut pos = 0; + assert_eq!(read_varint(&out, &mut pos), Ok(v)); + assert_eq!(pos, out.len()); + } + assert_eq!(varint_len(u64::MAX), MAX_VARINT_LEN); + } + + #[test] + fn non_minimal_varints_are_rejected() { + for bytes in [ + &[0x80, 0x00][..], + &[0x81, 0x00], + &[0xFF, 0x80, 0x00], + &[0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x80, 0x00], + ] { + let mut pos = 0; + assert_eq!( + read_varint(bytes, &mut pos), + Err(VarintError::NonMinimal), + "{bytes:02x?}" + ); + assert_eq!(pos, 0); + } + } + + #[test] + fn overflowing_and_truncated_varints_are_rejected() { + let mut pos = 0; + let eleven = [0xFFu8; 11]; + assert_eq!(read_varint(&eleven, &mut pos), Err(VarintError::Overflow)); + let mut ten = [0xFFu8; 10]; + ten[9] = 0x02; + assert_eq!(read_varint(&ten, &mut pos), Err(VarintError::Overflow)); + ten[9] = 0x01; + assert_eq!(read_varint(&ten, &mut pos), Ok(u64::MAX)); + let mut pos = 0; + assert_eq!(read_varint(&[0x80], &mut pos), Err(VarintError::Truncated)); + assert_eq!(read_varint(&[], &mut pos), Err(VarintError::Truncated)); + } + + #[test] + fn cursor_bounds_lengths_by_remaining_bytes() { + let mut c = Cursor::new(&[3, 1, 2, 3, 9], "test"); + assert_eq!(c.len_bytes().unwrap(), &[1, 2, 3]); + assert!(c.finish("x").is_err()); + assert_eq!(c.u8().unwrap(), 9); + c.finish("x").unwrap(); + + let mut c = Cursor::new(&[4, 1, 2, 3], "test"); + let err = c.len_varint().unwrap_err().to_string(); + assert!(err.contains("test: length exceeds"), "{err}"); + + let mut c = Cursor::new(&[0x80, 0x02], "test"); + assert!(c.varint_max(255, "count").is_err()); + } +} diff --git a/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index b4c473a..4313f8d 100644 --- a/docs/architecture/wire-format.md +++ b/docs/architecture/wire-format.md @@ -199,10 +199,11 @@ a compact frame instead, with its own magic `\x89PTWS`: ``` magic [u8;5] \x89PTWS flags u8 payload hash, plane CRC, shape, dtype name, - delta reference, stored; unknown bits rejected + delta reference, stored, size; unknown bits + rejected dtype_code varint input_format u8 -orig_size varint +orig_size varint (if flagged) shape varint ndim + varint dims (if flagged) dtype_name varint length + utf8 (if flagged) delta_ref [u8;32] blake3 of the reference (if flagged) @@ -210,21 +211,59 @@ payload_hash u64 xxh64 of the raw bytes (if flagged) body ``` +`orig_size` is stored only when the shape does not give it as +`product(shape) × element_size(dtype_code)`. The dtype name is stored +only when it differs from the canonical name of `dtype_code`. Varints +are minimal unsigned LEB128, and a reader rejects any field written +in a form the writer would not produce, so each tensor has one frame +encoding. + The body is either the raw bytes (the stored flag, written when no -codec makes the tensor smaller) or a coded tensor: a list of op wire -ids, the chain in the wire form the chain registry uses (its op -indices refer to that list), and one record per terminal plane -(codec id, inline state, payload length, optional CRC32, payload). -The frame does not store plane layouts, plane lengths or the plane -count, because the decoder derives them from the chain. The dtype -name is stored only when it differs from the canonical name of -`dtype_code`. - -The fixed cost is 10 to 130 bytes, depending on the chain and on -which integrity fields are on. A tensor whose planes need shared -state (a prelude) or chunked encoding cannot use the frame; the -compressor writes a full container for it. The frame holds no -extension table, so it can reference built-in ops and codecs only. +codec makes the tensor smaller) or a coded tensor: the chain in the +compact chain wire (CW2, below), then one record per terminal plane: +a varint `codec_id << 2 | state` (state 0 none, 1 inline with a +version byte and length-prefixed bytes, 2 inline with version 0 and +no bytes), the payload length (omitted on the last plane, whose +payload runs to the end of the frame), an optional CRC32 and the +payload. The frame does not store plane layouts, plane lengths or +the plane count, because the decoder derives them from the chain. + +The fixed cost is 10 to 70 bytes, depending on the chain and on +which integrity fields are on: about 48 bytes for a bf16 tensor and +68 bytes for an fp32 tensor with the payload hash and plane CRCs on. +A tensor whose planes need shared state (a prelude) or chunked +encoding cannot use the frame; the compressor writes a full +container for it. The frame holds no extension table, so it can +reference built-in ops and codecs only. + +### Compact chain wire (CW2) + +`chain::wire_compact` stores a chain in about 15 to 20 bytes, against +50 to 70 for the registry wire form: + +``` +n_nodes, n_edges, n_terms varint each (n_nodes 1..=255, others <= 255) +node: op_ref varint, param_len varint (<= 255), params +n_spine varint (<= n_edges; only when n_edges > 0) +edge: src, src_out, dst, dst_in varint; etag u8 (bit0 role, + bit1 vendor); [compact role]; [varint length + vendor bytes] +terminal: ref u8 (node << 4 | out, or 0xFF + node u8 + out u8); + compact role +``` + +- `op_ref` is the `OpId` wire id in the compact frame. +- The first `n_spine` edges are `(i, 0) → (i + 1, 0)` with no role + override and no vendor bytes, and are not stored. +- A Source node with `param_len = 0` is a bound Source: its params + are those of the frame's shape and dtype (`u8` ndim, `u32` dims, + `u16` chain dtype). The writer binds a Source exactly when its + params equal these, so chains that differ only in shape encode to + the same bytes. +- A compact role is one byte for every role production chains use + (scale formats, nibble kinds, mantissa and integer bytes up to 8, + `IntN` up to 32 bits, raw, index, exponent byte, fp4 value, XOR + residual), and `0xFF` followed by the registry role encoding for + the rest. The frame follows the same versioning rule as the container: no numeric version field, and a new magic for an incompatible layout. diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index 2a0a9bd..fe20860 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -101,3 +101,9 @@ name = "compact_decode" path = "fuzz_targets/compact_decode.rs" test = false doc = false + +[[bin]] +name = "chain_wire_compact" +path = "fuzz_targets/chain_wire_compact.rs" +test = false +doc = false diff --git a/fuzz/fuzz_targets/chain_wire_compact.rs b/fuzz/fuzz_targets/chain_wire_compact.rs new file mode 100644 index 0000000..11f37e3 --- /dev/null +++ b/fuzz/fuzz_targets/chain_wire_compact.rs @@ -0,0 +1,31 @@ +#![no_main] + +// CW2, the compact chain wire, is read from untrusted `\x89PTWS` frames (and +// compact containers). The parser must never panic, infinite-loop or +// allocate unboundedly, and every chain it accepts must re-encode to the +// same bytes (the format is canonical). +// +// Input layout: byte 0 selects the bound Source (bit 0 set: none; else +// params for a [3, 5] fp32 tensor); the rest is the chain. + +use libfuzzer_sys::fuzz_target; +use ptwm_core::chain::validate::validate_chain; +use ptwm_core::chain::wire_compact::{ + op_from_wire_ref, op_wire_ref, read_chain_compact_exact, write_chain_compact, +}; +use ptwm_core::transforms::source::bound_source_params; + +fuzz_target!(|data: &[u8]| { + let Some((&sel, chain_bytes)) = data.split_first() else { + return; + }; + let params = bound_source_params(&[3, 5], 1).expect("fp32 binds"); + let bound = (sel & 1 == 0).then_some(params.as_slice()); + let Ok(chain) = read_chain_compact_exact(chain_bytes, bound, op_from_wire_ref) else { + return; + }; + let mut out = Vec::new(); + write_chain_compact(&chain, bound, op_wire_ref, &mut out).expect("accepted chain re-encodes"); + assert_eq!(out, chain_bytes, "CW2 is not canonical"); + let _ = validate_chain(&chain, 0); +}); diff --git a/fuzz/fuzz_targets/compact_decode.rs b/fuzz/fuzz_targets/compact_decode.rs index 4524138..9766828 100644 --- a/fuzz/fuzz_targets/compact_decode.rs +++ b/fuzz/fuzz_targets/compact_decode.rs @@ -5,34 +5,40 @@ // adversarial input. use libfuzzer_sys::fuzz_target; -use ptwm_core::compact::{MAGIC_COMPACT, decode_compact}; +use ptwm_core::compact::{FLAG_SHAPE, FLAG_SIZE, MAGIC_COMPACT, decode_compact, size_from_shape}; +use ptwm_core::varint::read_varint; -/// Read one LEB128 varint at `*pos`; `None` when truncated or too long. -fn varint(data: &[u8], pos: &mut usize) -> Option { - let mut v = 0u64; - for i in 0..10 { - let b = *data.get(*pos)?; - *pos += 1; - v |= u64::from(b & 0x7F) << (7 * i); - if b & 0x80 == 0 { - return Some(v); - } +/// Raw byte count the frame declares, from `orig_size` or the shape. Layout +/// after the magic: flags u8, dtype_code varint, input_format u8, +/// [orig_size varint] (FLAG_SIZE), [ndim varint, dims varint...] +/// (FLAG_SHAPE). `None` when the header does not parse that far. +fn declared_size(data: &[u8]) -> Option { + let flags = *data.first()?; + let mut pos = 1; + let dtype_code = u16::try_from(read_varint(data, &mut pos).ok()?).ok()?; + pos += 1; + if flags & FLAG_SIZE != 0 { + return read_varint(data, &mut pos).ok(); + } + if flags & FLAG_SHAPE == 0 { + return None; } - None + let ndim = read_varint(data, &mut pos).ok()?.min(64); + let mut shape = Vec::new(); + for _ in 0..ndim { + shape.push(read_varint(data, &mut pos).ok()?); + } + // An unknown dtype or an overflowing product fails in the decoder. + Some(size_from_shape(&shape, dtype_code).unwrap_or(0)) } fuzz_target!(|data: &[u8]| { // The frame declares its decoded size, and a decoder must be able to // allocate it: a 100-byte frame of a constant tensor can declare 1 GB. // Skip declared sizes above 16 MiB so libFuzzer's RSS limit flags real - // over-allocation only. Layout after the magic: flags u8, dtype varint, - // input_format u8, orig_size varint. - let mut pos = 1; - if varint(data, &mut pos).is_some() { - pos += 1; - if varint(data, &mut pos).is_some_and(|n| n > 1 << 24) { - return; - } + // over-allocation only. + if declared_size(data).is_some_and(|n| n > 1 << 24) { + return; } // Prefix the magic so the fuzzer spends its time past the magic check. let mut buf = Vec::with_capacity(MAGIC_COMPACT.len() + data.len()); diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 0c700ee..8c930b1 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -46,7 +46,7 @@ class Framing(Enum): r"""Byte layout of a single-tensor :meth:`Compressor.compress` output. ``COMPACT`` writes a compact single-tensor frame (magic ``\x89PTWS``) - with a fixed cost of about 10 to 130 bytes, and stores the raw bytes when + with a fixed cost of about 10 to 70 bytes, and stores the raw bytes when no codec makes them smaller. It falls back to ``CONTAINER`` when the tensor cannot use the frame (streaming, or a plane that selects shared codec state). ``CONTAINER`` writes a full ``.ptwm`` container (magic diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py index 5ae6a3f..4124e04 100644 --- a/tests/ptwm/benchmarks/test_container_overhead.py +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -24,6 +24,11 @@ # codec payload is smaller than zstd-19's output. SLACK = 64 +# Largest compact-frame non-payload bytes (header, compact chain, plane heads, +# payload hash and plane CRCs) per dtype. fp32 frames carry four planes, bf16 +# frames two. Measured maxima over SIZES: fp32 68 B, bf16 48 B. +COMPACT_OVERHEAD = {"bf16": 55, "fp32": 70} + def test_constant_64_byte_bf16_tensor_is_at_most_128_bytes() -> None: t = torch.full((32,), 0.5, dtype=torch.bfloat16) @@ -42,6 +47,7 @@ def test_random_normal_overhead_is_bounded( assert row.out <= row.raw + SLACK, row if row.payload < row.zstd19: assert row.out <= row.zstd19 + SLACK, row + assert row.non_payload <= COMPACT_OVERHEAD[dname], row @pytest.mark.parametrize(("dname", "dtype"), DTYPES) @@ -53,7 +59,7 @@ def test_constant_overhead_is_bounded( make_tensor(dtype, "constant", nbytes), Framing.COMPACT, dname, "constant" ) assert row.out <= row.raw + SLACK, row - # Four-plane fp32 frames carry about 120 B of chain and plane metadata. + assert row.non_payload <= COMPACT_OVERHEAD[dname], row assert row.out <= 256, row From 1ca8aa519c99e6a1d2e6d983e693a69dc888460f Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 00:01:42 +0000 Subject: [PATCH 12/34] fix(S5-cw2-ptws): address review findings - PTWS frames are self-delimiting again: every plane, including the last, stores its payload length, and the reader rejects bytes appended to a coded or stored frame. - The compact writer stores plane CRCs only without the payload hash, which already covers every plane of the one tensor; a reader rejects a frame with both flags. This pays for the restored length: default fp32 overhead is now at most 55 B (split chain) and 61 B (IntDelta chain). - Docs (compact.rs, varint.rs, wire-format.md, _config.py) now claim one encoding per field, not per tensor, and name the remaining writer choices (stored vs coded body, plane state 0 vs 2). - Overhead tests: tighter random/constant bounds and a smooth-data case (IntDelta chains) held to 70 B. - Fuzz: seed corpora for compact_decode and chain_wire_compact; compact_decode asserts that no accepted frame survives an appended byte. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/compact.rs | 178 +++++++++++++----- crates/ptwm-core/src/varint.rs | 7 +- docs/architecture/wire-format.md | 28 ++- fuzz/fuzz_targets/chain_wire_compact.rs | 2 + fuzz/fuzz_targets/compact_decode.rs | 10 +- python/ptwm/_config.py | 8 +- .../benchmarks/test_container_overhead.py | 40 +++- tests/ptwm/core/test_compact_frame.py | 47 ++++- 8 files changed, 252 insertions(+), 68 deletions(-) diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index 17040f1..e71b282 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -7,7 +7,9 @@ //! //! ```text //! magic: [u8; 5] \x89PTWS -//! flags: u8 see FLAG_*; unknown bits are rejected +//! flags: u8 see FLAG_*; unknown bits are rejected, +//! and FLAG_PLANE_CRC with FLAG_PAYLOAD_HASH +//! is rejected //! dtype_code: varint at most u16::MAX //! input_format: u8 //! [orig_size] varint raw byte count (FLAG_SIZE) @@ -42,19 +44,29 @@ //! [state_version] u8 (state 1) //! [state_len] varint (state 1) //! [state] bytes (state 1) -//! [payload_len] varint every plane but the last; the last -//! payload runs to the end of the frame +//! payload_len: varint //! [crc32] u32 CRC32 of the raw plane (FLAG_PLANE_CRC) //! payload: bytes //! ``` //! //! State 1 with version 0 and an empty state is rejected: state 2 is its -//! only encoding. All integers are little-endian; varints are canonical -//! unsigned LEB128 ([`crate::varint`]). Every header and chain field has -//! exactly one encoding. A stored body must be exactly `orig_size` bytes, -//! and a coded body ends with the last plane's payload, so bytes appended to -//! a coded frame reach the last plane's codec, which rejects them or ignores -//! them; the plane CRC and payload hash still guard the decoded bytes. +//! only encoding. States 0 and 2 decode alike; the frame keeps the +//! distinction the codec reported, and the reader accepts either. All +//! integers are little-endian; varints are canonical unsigned LEB128 +//! ([`crate::varint`]). A stored body must be exactly `orig_size` bytes, and +//! a coded body must end exactly with the last plane's payload, so the reader +//! rejects bytes appended to either. +//! +//! The payload hash covers the decoded bytes of the one tensor, and so every +//! plane, so the writer stores plane CRCs only when it does not store the +//! hash. Both checks verify decoded bytes; they do not cover the coded +//! payload. +//! +//! Each header, chain and plane field has one encoding for the value it +//! holds. The frame as a whole is not unique for a tensor: the writer +//! chooses between a stored and a coded body, between plane state codes 0 +//! and 2, and between codecs and chains, and the reader accepts any valid +//! choice. //! //! The plane layout, the decoded plane length and the plane count all come //! from the chain, so the frame does not store them. Planes that need shared @@ -91,7 +103,8 @@ pub const MAGIC_COMPACT: [u8; 5] = [0x89, b'P', b'T', b'W', b'S']; /// The frame carries an xxh64 of the raw bytes. pub const FLAG_PAYLOAD_HASH: u8 = 0x01; -/// Every coded plane carries a CRC32 of its raw bytes. +/// Every coded plane carries a CRC32 of its raw bytes. Not combined with +/// [`FLAG_PAYLOAD_HASH`], which already covers every plane. pub const FLAG_PLANE_CRC: u8 = 0x02; /// The frame carries the tensor shape. pub const FLAG_SHAPE: u8 = 0x04; @@ -226,11 +239,14 @@ pub fn compress_compact( out.extend_from_slice(&xxh64(raw, 0).to_le_bytes()); } - let coded = encode_body(input, opts, bound_source.as_deref())?; + // The payload hash covers the decoded bytes of every plane, so plane + // CRCs would only repeat it. + let plane_crc = opts.emit_plane_crc && !opts.emit_payload_hash; + let coded = encode_body(input, opts, bound_source.as_deref(), plane_crc)?; match coded { None => return Ok(None), Some(body) if body.len() < raw.len() => { - if opts.emit_plane_crc { + if plane_crc { flags |= FLAG_PLANE_CRC; } out.extend_from_slice(&body); @@ -250,6 +266,7 @@ fn encode_body( input: &InputTensor<'_>, opts: &CompressorOptions, bound_source: Option<&[u8]>, + plane_crc: bool, ) -> Result>, PtwmCoreError> { let plan = plan_model(std::slice::from_ref(input), opts)?; if !plan.prelude_entries.is_empty() { @@ -266,7 +283,7 @@ fn encode_body( ); write_chain_compact(chain, bound_source, op_wire_ref, &mut body)?; - for (i, cp) in planes.iter().enumerate() { + for cp in planes { let codec = u64::from(cp.codec_id.as_u16()) << 2; match cp.state_source { StateSource::None => write_varint(&mut body, codec | STATE_NONE), @@ -283,10 +300,8 @@ fn encode_body( } StateSource::Shared | StateSource::External => return Ok(None), } - if i + 1 < planes.len() { - write_varint(&mut body, cp.payload_bytes.len() as u64); - } - if opts.emit_plane_crc { + write_varint(&mut body, cp.payload_bytes.len() as u64); + if plane_crc { body.extend_from_slice(&crc32fast::hash(&cp.raw_plane_for_crc).to_le_bytes()); } body.extend_from_slice(&cp.payload_bytes); @@ -318,6 +333,9 @@ fn read_header(r: &mut Cursor<'_>) -> Result { flags & !KNOWN_FLAGS ))); } + if flags & FLAG_PLANE_CRC != 0 && flags & FLAG_PAYLOAD_HASH != 0 { + return Err(r.err("plane CRC flag with payload hash")); + } let dtype_code = r.varint_max(u64::from(u16::MAX), "dtype_code")? as u16; let input_format = r.u8()?; let stored_size = if flags & FLAG_SIZE != 0 { @@ -465,7 +483,7 @@ fn decode_body( let n_planes = descriptors.len(); let mut terminal_planes: Vec = Vec::with_capacity(n_planes); - for (i, descriptor) in descriptors.into_iter().enumerate() { + for descriptor in descriptors { let head = r.varint()?; let wire = u16::try_from(head >> 2).map_err(|_| r.err("codec id exceeds u16"))?; let codec_id = CodecId::from_u16(wire) @@ -482,10 +500,9 @@ fn decode_body( } s => return Err(r.err(format_args!("unknown state code {s}"))), }; - let last = i + 1 == n_planes; - let payload_len = if last { None } else { Some(r.len_varint()?) }; + let payload_len = r.len_varint()?; let crc = if has_crc { Some(r.u32_le()?) } else { None }; - let payload = r.take(payload_len.unwrap_or(r.remaining()))?; + let payload = r.take(payload_len)?; *payload_len_out += (state.len() + payload.len()) as u64; let codec = codec_for(codec_id) @@ -891,7 +908,7 @@ mod tests { #[test] fn fp32_frame_binds_source_and_stays_small() { let raw = fp32_raw(4096); - for (h, c) in [(true, true), (false, false)] { + for (h, c) in [(true, true), (false, false), (false, true)] { let blob = compress_compact(&fp32_input(&raw, &[64, 64]), &opts(h, c)) .unwrap() .unwrap(); @@ -903,11 +920,27 @@ mod tests { assert_eq!(t.dtype_name.as_deref(), Some("float32")); let overhead = blob.len() as u64 - t.payload_len; // magic 5, flags 1, dtype 1, format 1, shape 3, chain 19, - // 4 plane heads, 3 payload lengths; plus hash 8 and 4 CRCs. - let limit = if h { 33 + 8 + 4 + 16 + 8 } else { 33 + 4 + 8 }; + // 4 plane heads, 4 payload lengths (up to 2 bytes each); plus + // the hash (8) or 4 CRCs (16), never both. + let integrity = match (h, c) { + (true, _) => 8, + (false, true) => 16, + (false, false) => 0, + }; + let limit = 31 + 4 + 8 + integrity; assert!( overhead <= limit, - "overhead {overhead} > {limit} (hash {h})" + "overhead {overhead} > {limit} (hash {h}, crc {c})" + ); + assert_eq!( + flags & (FLAG_PAYLOAD_HASH | FLAG_PLANE_CRC), + if h { + FLAG_PAYLOAD_HASH + } else if c { + FLAG_PLANE_CRC + } else { + 0 + } ); // The Source dims are not stored. let dims: Vec = [&64u32.to_le_bytes()[..], &64u32.to_le_bytes()].concat(); @@ -1024,7 +1057,7 @@ mod tests { } #[test] - fn plane_state_codes_are_canonical() { + fn plane_state_codes_are_checked() { let raw: Vec = (0..4096u32).map(|i| (i % 7) as u8).collect(); let mut o = opts(false, false); o.forced_codec = Some(CodecId::Zstd); @@ -1036,7 +1069,9 @@ mod tests { // Zstd carries an empty version-0 inline state: state code 2. assert_eq!(blob[at], (4 << 2) | 2); let with_head = |head: &[u8]| [&blob[..at], head, &blob[at + 1..]].concat(); - // The same plane with state code 0 decodes (decode treats both alike). + // The same plane with state code 0 decodes: decode treats both + // alike, and the frame keeps the state kind the codec reported, so + // these two frames are both valid for one tensor. assert_eq!(decode_compact(&with_head(&[4 << 2])).unwrap().raw, raw); // State code 1 with an empty version-0 state is non-canonical. let err = decode_compact(&with_head(&[(4 << 2) | 1, 0, 0])) @@ -1050,24 +1085,79 @@ mod tests { } #[test] - fn last_plane_runs_to_the_end_of_the_frame() { - let raw = fp32_raw(1024); - let blob = compress_compact(&fp32_input(&raw, &[1024]), &opts(true, true)) - .unwrap() - .unwrap(); - assert_eq!(decode_compact(&blob).unwrap().raw, raw); - // Every truncation fails: a short last payload fails its codec, CRC - // or the payload hash. - for cut in 0..blob.len() { - assert!(decode_compact(&blob[..cut]).is_err(), "cut at {cut}"); + fn appended_bytes_are_rejected() { + let fp32 = fp32_raw(1024); + let bytes: Vec = (0..4096u32).map(|i| (i % 7) as u8).collect(); + let mut blobs = Vec::new(); + for (h, c) in [(true, true), (false, true), (false, false)] { + blobs.push( + compress_compact(&fp32_input(&fp32, &[1024]), &opts(h, c)) + .unwrap() + .unwrap(), + ); + blobs.push( + compress_compact(&input(&bytes, None, None), &opts(h, c)) + .unwrap() + .unwrap(), + ); } - // A byte appended to the frame becomes part of the last payload. - // Whether that fails depends on the codec; it must never return - // other data. - let mut longer = blob.clone(); - longer.push(0); - if let Ok(t) = decode_compact(&longer) { - assert_eq!(t.raw, raw); + let mut zstd = opts(true, false); + zstd.forced_codec = Some(CodecId::Zstd); + blobs.push( + compress_compact(&input(&bytes, None, None), &zstd) + .unwrap() + .unwrap(), + ); + for blob in &blobs { + assert_eq!(blob[FLAGS_AT] & FLAG_STORED, 0); + let t = decode_compact(blob).unwrap(); + // Every payload byte is accounted for: the coded body ends with + // the last payload. + assert!(t.payload_len > 0); + for cut in 0..blob.len() { + assert!(decode_compact(&blob[..cut]).is_err(), "cut at {cut}"); + } + for tail in [&[0u8][..], &[0u8; 7], &[0xFFu8; 100]] { + let longer = [blob.as_slice(), tail].concat(); + let err = decode_compact(&longer).unwrap_err().to_string(); + assert!(err.contains("trailing"), "{} appended: {err}", tail.len()); + } } + // A stored body must be exactly orig_size bytes. + let mut x: u64 = 0x9E37_79B9_7F4A_7C15; + let noise: Vec = (0..300) + .map(|_| { + x ^= x << 13; + x ^= x >> 7; + x ^= x << 17; + x as u8 + }) + .collect(); + let stored = compress_compact(&input(&noise, None, None), &opts(true, true)) + .unwrap() + .unwrap(); + assert_ne!(stored[FLAGS_AT] & FLAG_STORED, 0); + assert!(decode_compact(&[stored.as_slice(), &[0]].concat()).is_err()); + } + + #[test] + fn plane_crc_is_not_combined_with_the_payload_hash() { + let raw = fp32_raw(256); + let with_crc = compress_compact(&fp32_input(&raw, &[256]), &opts(false, true)) + .unwrap() + .unwrap(); + assert_eq!(with_crc[FLAGS_AT], FLAG_SHAPE | FLAG_PLANE_CRC); + decode_compact(&with_crc).unwrap(); + // Add the payload hash to the CRC frame: the writer never does. + let at = body_start(&with_crc); + let mut b = [ + &with_crc[..at], + &xxh64(&raw, 0).to_le_bytes()[..], + &with_crc[at..], + ] + .concat(); + b[FLAGS_AT] |= FLAG_PAYLOAD_HASH; + let err = decode_compact(&b).unwrap_err().to_string(); + assert!(err.contains("plane CRC flag with payload hash"), "{err}"); } } diff --git a/crates/ptwm-core/src/varint.rs b/crates/ptwm-core/src/varint.rs index 7bd23a6..accac41 100644 --- a/crates/ptwm-core/src/varint.rs +++ b/crates/ptwm-core/src/varint.rs @@ -3,9 +3,10 @@ //! //! Parsing is canonical: a varint must use the fewest bytes that hold its //! value (no trailing `0x80`/`0x00` groups), is at most 10 bytes long, and -//! must fit in `u64`. Together with the other fail-closed checks of the -//! compact formats, this makes every accepted byte string the only encoding -//! of its value, so `write(parse(b)) == b`. +//! must fit in `u64`, so each value has one varint encoding. The compact +//! formats build on this: each of their fields has one encoding for the +//! value it holds. (A compact frame as a whole still leaves the writer some +//! choices; see [`crate::compact`].) use crate::error::PtwmCoreError; diff --git a/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index 4313f8d..af4b709 100644 --- a/docs/architecture/wire-format.md +++ b/docs/architecture/wire-format.md @@ -200,7 +200,7 @@ a compact frame instead, with its own magic `\x89PTWS`: magic [u8;5] \x89PTWS flags u8 payload hash, plane CRC, shape, dtype name, delta reference, stored, size; unknown bits - rejected + and plane CRC with payload hash rejected dtype_code varint input_format u8 orig_size varint (if flagged) @@ -215,22 +215,32 @@ body `product(shape) × element_size(dtype_code)`. The dtype name is stored only when it differs from the canonical name of `dtype_code`. Varints are minimal unsigned LEB128, and a reader rejects any field written -in a form the writer would not produce, so each tensor has one frame -encoding. +in a form the writer would not produce, so each field has one +encoding for the value it holds. The frame as a whole is not unique +for a tensor: the writer chooses the codecs, the chain, a stored or a +coded body, and plane state 0 or 2 (below), and the reader accepts +any valid choice. A reader rejects bytes after the end of the body. The body is either the raw bytes (the stored flag, written when no codec makes the tensor smaller) or a coded tensor: the chain in the compact chain wire (CW2, below), then one record per terminal plane: a varint `codec_id << 2 | state` (state 0 none, 1 inline with a version byte and length-prefixed bytes, 2 inline with version 0 and -no bytes), the payload length (omitted on the last plane, whose -payload runs to the end of the frame), an optional CRC32 and the -payload. The frame does not store plane layouts, plane lengths or +no bytes; states 0 and 2 decode alike), the payload length, an +optional CRC32 and the payload. The coded body ends with the last +payload. The payload hash and the plane CRCs check decoded bytes, not +the coded payload. The payload hash already covers every plane of the +one tensor, so the writer stores plane CRCs only when it does not +store the hash, and a reader rejects a frame with both. The frame does not store plane layouts, plane lengths or the plane count, because the decoder derives them from the chain. -The fixed cost is 10 to 70 bytes, depending on the chain and on -which integrity fields are on: about 48 bytes for a bf16 tensor and -68 bytes for an fp32 tensor with the payload hash and plane CRCs on. +The fixed cost is about 10 to 70 bytes, depending on the chain and +on which integrity fields are on. With the default settings (payload +hash on) it is at most about 50 bytes for a bf16 tensor, 55 bytes for +an fp32 tensor with the split chain, and 61 bytes for an fp32 tensor +with the IntDelta or PredictorXor chain that smooth data selects. +Plane CRCs instead of the hash add up to 8 bytes more for four +planes. A tensor whose planes need shared state (a prelude) or chunked encoding cannot use the frame; the compressor writes a full container for it. The frame holds no extension table, so it can diff --git a/fuzz/fuzz_targets/chain_wire_compact.rs b/fuzz/fuzz_targets/chain_wire_compact.rs index 11f37e3..ebd0b6d 100644 --- a/fuzz/fuzz_targets/chain_wire_compact.rs +++ b/fuzz/fuzz_targets/chain_wire_compact.rs @@ -7,6 +7,8 @@ // // Input layout: byte 0 selects the bound Source (bit 0 set: none; else // params for a [3, 5] fp32 tensor); the rest is the chain. +// Seeds in `corpus/chain_wire_compact` are chains from production frames, +// bound and explicit, with the selector byte. use libfuzzer_sys::fuzz_target; use ptwm_core::chain::validate::validate_chain; diff --git a/fuzz/fuzz_targets/compact_decode.rs b/fuzz/fuzz_targets/compact_decode.rs index 9766828..74a5ee1 100644 --- a/fuzz/fuzz_targets/compact_decode.rs +++ b/fuzz/fuzz_targets/compact_decode.rs @@ -2,7 +2,8 @@ // `Decompressor.decompress` hands any blob that starts with the compact-frame // magic to `decode_compact`. The parser must never panic or infinite-loop on -// adversarial input. +// adversarial input. Seeds in `corpus/compact_decode` are production frames +// without the magic. use libfuzzer_sys::fuzz_target; use ptwm_core::compact::{FLAG_SHAPE, FLAG_SIZE, MAGIC_COMPACT, decode_compact, size_from_shape}; @@ -44,5 +45,10 @@ fuzz_target!(|data: &[u8]| { let mut buf = Vec::with_capacity(MAGIC_COMPACT.len() + data.len()); buf.extend_from_slice(&MAGIC_COMPACT); buf.extend_from_slice(data); - let _ = decode_compact(&buf); + if decode_compact(&buf).is_ok() { + // A frame is self-delimiting: no accepted frame stays valid with a + // byte appended. + buf.push(0); + assert!(decode_compact(&buf).is_err(), "appended byte accepted"); + } }); diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 8c930b1..c3d0e1c 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -46,7 +46,7 @@ class Framing(Enum): r"""Byte layout of a single-tensor :meth:`Compressor.compress` output. ``COMPACT`` writes a compact single-tensor frame (magic ``\x89PTWS``) - with a fixed cost of about 10 to 70 bytes, and stores the raw bytes when + with a fixed cost of at most about 70 bytes, and stores the raw bytes when no codec makes them smaller. It falls back to ``CONTAINER`` when the tensor cannot use the frame (streaming, or a plane that selects shared codec state). ``CONTAINER`` writes a full ``.ptwm`` container (magic @@ -116,7 +116,11 @@ class CompressionConfig: emit_payload_hash: bool = True """Store an xxh64 of the raw bytes (8 B per tensor), checked on decode.""" emit_plane_crc: bool = True - """Store a CRC32 of every coded plane (4 B per plane), checked on decode.""" + """Store a CRC32 of every coded plane (4 B per plane), checked on decode. + + A compact frame stores plane CRCs only when ``emit_payload_hash`` is off: + the payload hash already covers every plane of its one tensor. + """ @classmethod def from_resolved_policy( diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py index 4124e04..0f033a5 100644 --- a/tests/ptwm/benchmarks/test_container_overhead.py +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -10,6 +10,7 @@ import pytest import torch +from ptwm import CompressionConfig, Compressor, Decompressor, Format from ptwm._config import Framing from benchmarks.container_overhead import ( @@ -17,6 +18,7 @@ SIZES, make_tensor, measure, + non_payload_bytes, tiled_rows, ) @@ -24,10 +26,17 @@ # codec payload is smaller than zstd-19's output. SLACK = 64 -# Largest compact-frame non-payload bytes (header, compact chain, plane heads, -# payload hash and plane CRCs) per dtype. fp32 frames carry four planes, bf16 -# frames two. Measured maxima over SIZES: fp32 68 B, bf16 48 B. -COMPACT_OVERHEAD = {"bf16": 55, "fp32": 70} +# Largest compact-frame non-payload bytes (header, compact chain, plane heads +# and payload hash; the frame omits plane CRCs when it stores the hash) per +# dtype. fp32 frames carry four planes, bf16 frames two. Measured maxima over +# SIZES: fp32 55 B, bf16 43 B. +COMPACT_OVERHEAD = {"bf16": 50, "fp32": 62} + +# Bound for any fp32 or bf16 chain the compressor picks, including the +# IntDelta and PredictorXor chains for smooth data, which carry an extra node +# and a non-spine edge. Measured maxima for cumulative-sum data: fp32 61 B +# with the default config, 69 B with plane CRCs and no payload hash. +CHAIN_OVERHEAD_LIMIT = 70 def test_constant_64_byte_bf16_tensor_is_at_most_128_bytes() -> None: @@ -63,6 +72,29 @@ def test_constant_overhead_is_bounded( assert row.out <= 256, row +@pytest.mark.parametrize( + ("hash_", "crc"), [(True, True), (False, True), (False, False)] +) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +@pytest.mark.parametrize("n", [64, 1024, 65536, 262144]) +def test_smooth_data_overhead_is_bounded( + hash_: bool, crc: bool, dtype: torch.dtype, n: int +) -> None: + g = torch.Generator().manual_seed(n) + t = torch.cumsum(torch.randn(n, generator=g), 0).to(dtype) + blob = Compressor( + CompressionConfig( + input_format=Format.TORCH, + emit_payload_hash=hash_, + emit_plane_crc=crc, + ) + ).compress(t) + assert blob[:5] == b"\x89PTWS" + back = Decompressor().decompress(blob) + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + assert non_payload_bytes(blob) <= CHAIN_OVERHEAD_LIMIT, (dtype, n) + + def test_tiled_rows_ratio_is_at_least_60() -> None: row = measure(tiled_rows(64), Framing.COMPACT, "bf16", "tiled") assert row.ratio >= 60, f"(64, 1024) tiled bf16 ratio {row.ratio:.1f}" diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index 5be10d0..a922705 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -109,10 +109,49 @@ def test_integrity_fields_can_be_disabled(framing: Framing) -> None: bare = Compressor( _torch_cfg(framing=framing, emit_payload_hash=False, emit_plane_crc=False) ).compress(t) - # 8 B payload hash, plus 4 B CRC per coded plane (bf16 has two). - assert len(full) - len(bare) == 8 + 2 * 4 - back = Decompressor().decompress(bare) - assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + crc_only = Compressor( + _torch_cfg(framing=framing, emit_payload_hash=False, emit_plane_crc=True) + ).compress(t) + # 8 B payload hash; 4 B CRC per coded plane (bf16 has two). The compact + # frame omits plane CRCs when it stores the payload hash, which already + # covers every plane of its one tensor. + if framing is Framing.COMPACT: + assert len(full) - len(bare) == 8 + else: + assert len(full) - len(bare) == 8 + 2 * 4 + assert len(crc_only) - len(bare) == 2 * 4 + for blob in (bare, crc_only): + back = Decompressor().decompress(blob) + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + + +@pytest.mark.parametrize("tail", [b"\x00", b"\x00" * 7, b"\xff" * 100]) +@pytest.mark.parametrize( + "t", + [ + torch.randn(64, 64, generator=torch.Generator().manual_seed(0)).to( + torch.bfloat16 + ), + torch.randn(64, 64, generator=torch.Generator().manual_seed(0)), + torch.randint( + -128, 127, (4096,), generator=torch.Generator().manual_seed(0) + ).to(torch.int8), + torch.cumsum(torch.randn(4096, generator=torch.Generator().manual_seed(0)), 0), + ], + ids=["bf16", "fp32", "int8", "fp32-smooth"], +) +@pytest.mark.parametrize( + ("hash_", "crc"), [(True, True), (False, True), (False, False)] +) +def test_appended_bytes_are_rejected( + t: torch.Tensor, tail: bytes, hash_: bool, crc: bool +) -> None: + blob = Compressor(_torch_cfg(emit_payload_hash=hash_, emit_plane_crc=crc)).compress( + t + ) + assert blob[:5] == b"\x89PTWS" + with pytest.raises(Error, match="trailing|stored body is"): + Decompressor().decompress(blob + tail) def test_delta_round_trip() -> None: From f6e0e48a8087576c0120eb879f4d8ecaa40143d6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 00:02:05 +0000 Subject: [PATCH 13/34] fix(S5-cw2-ptws): add the compact fuzz seed corpora fuzz/corpus/ is gitignored, so the seeds the previous commit's fuzz targets point to were left out. Force-add them, as for container_open. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- fuzz/corpus/chain_wire_compact/bf16_const | Bin 0 -> 8 bytes fuzz/corpus/chain_wire_compact/bf16_randn | Bin 0 -> 16 bytes fuzz/corpus/chain_wire_compact/bf16_smooth | Bin 0 -> 23 bytes fuzz/corpus/chain_wire_compact/byte_bf16 | Bin 0 -> 30 bytes fuzz/corpus/chain_wire_compact/fp32_const | Bin 0 -> 20 bytes fuzz/corpus/chain_wire_compact/fp32_smooth | Bin 0 -> 27 bytes fuzz/corpus/compact_decode/bf16_const_bare | Bin 0 -> 31 bytes fuzz/corpus/compact_decode/bf16_const_crc | Bin 0 -> 35 bytes fuzz/corpus/compact_decode/bf16_const_hash | Bin 0 -> 39 bytes fuzz/corpus/compact_decode/bf16_randn_bare | Bin 0 -> 397 bytes fuzz/corpus/compact_decode/bf16_randn_crc | Bin 0 -> 405 bytes fuzz/corpus/compact_decode/bf16_randn_hash | Bin 0 -> 405 bytes fuzz/corpus/compact_decode/bf16_smooth_bare | Bin 0 -> 380 bytes fuzz/corpus/compact_decode/bf16_smooth_crc | Bin 0 -> 388 bytes fuzz/corpus/compact_decode/bf16_smooth_hash | Bin 0 -> 388 bytes fuzz/corpus/compact_decode/byte_bf16_hash | Bin 0 -> 674 bytes fuzz/corpus/compact_decode/fp16_randn_bare | Bin 0 -> 133 bytes fuzz/corpus/compact_decode/fp16_randn_crc | Bin 0 -> 133 bytes fuzz/corpus/compact_decode/fp16_randn_hash | Bin 0 -> 141 bytes fuzz/corpus/compact_decode/fp32_const_bare | Bin 0 -> 100 bytes fuzz/corpus/compact_decode/fp32_const_crc | Bin 0 -> 116 bytes fuzz/corpus/compact_decode/fp32_const_hash | Bin 0 -> 108 bytes fuzz/corpus/compact_decode/fp32_randn_bare | Bin 0 -> 489 bytes fuzz/corpus/compact_decode/fp32_randn_crc | Bin 0 -> 505 bytes fuzz/corpus/compact_decode/fp32_randn_hash | Bin 0 -> 497 bytes fuzz/corpus/compact_decode/fp32_smooth_bare | Bin 0 -> 902 bytes fuzz/corpus/compact_decode/fp32_smooth_crc | Bin 0 -> 918 bytes fuzz/corpus/compact_decode/fp32_smooth_hash | Bin 0 -> 910 bytes fuzz/corpus/compact_decode/fp8_e4m3_bare | Bin 0 -> 134 bytes fuzz/corpus/compact_decode/fp8_e4m3_crc | Bin 0 -> 134 bytes fuzz/corpus/compact_decode/fp8_e4m3_hash | Bin 0 -> 142 bytes fuzz/corpus/compact_decode/int8_stored_bare | 1 + fuzz/corpus/compact_decode/int8_stored_crc | 1 + fuzz/corpus/compact_decode/int8_stored_hash | 1 + 34 files changed, 3 insertions(+) create mode 100644 fuzz/corpus/chain_wire_compact/bf16_const create mode 100644 fuzz/corpus/chain_wire_compact/bf16_randn create mode 100644 fuzz/corpus/chain_wire_compact/bf16_smooth create mode 100644 fuzz/corpus/chain_wire_compact/byte_bf16 create mode 100644 fuzz/corpus/chain_wire_compact/fp32_const create mode 100644 fuzz/corpus/chain_wire_compact/fp32_smooth create mode 100644 fuzz/corpus/compact_decode/bf16_const_bare create mode 100644 fuzz/corpus/compact_decode/bf16_const_crc create mode 100644 fuzz/corpus/compact_decode/bf16_const_hash create mode 100644 fuzz/corpus/compact_decode/bf16_randn_bare create mode 100644 fuzz/corpus/compact_decode/bf16_randn_crc create mode 100644 fuzz/corpus/compact_decode/bf16_randn_hash create mode 100644 fuzz/corpus/compact_decode/bf16_smooth_bare create mode 100644 fuzz/corpus/compact_decode/bf16_smooth_crc create mode 100644 fuzz/corpus/compact_decode/bf16_smooth_hash create mode 100644 fuzz/corpus/compact_decode/byte_bf16_hash create mode 100644 fuzz/corpus/compact_decode/fp16_randn_bare create mode 100644 fuzz/corpus/compact_decode/fp16_randn_crc create mode 100644 fuzz/corpus/compact_decode/fp16_randn_hash create mode 100644 fuzz/corpus/compact_decode/fp32_const_bare create mode 100644 fuzz/corpus/compact_decode/fp32_const_crc create mode 100644 fuzz/corpus/compact_decode/fp32_const_hash create mode 100644 fuzz/corpus/compact_decode/fp32_randn_bare create mode 100644 fuzz/corpus/compact_decode/fp32_randn_crc create mode 100644 fuzz/corpus/compact_decode/fp32_randn_hash create mode 100644 fuzz/corpus/compact_decode/fp32_smooth_bare create mode 100644 fuzz/corpus/compact_decode/fp32_smooth_crc create mode 100644 fuzz/corpus/compact_decode/fp32_smooth_hash create mode 100644 fuzz/corpus/compact_decode/fp8_e4m3_bare create mode 100644 fuzz/corpus/compact_decode/fp8_e4m3_crc create mode 100644 fuzz/corpus/compact_decode/fp8_e4m3_hash create mode 100644 fuzz/corpus/compact_decode/int8_stored_bare create mode 100644 fuzz/corpus/compact_decode/int8_stored_crc create mode 100644 fuzz/corpus/compact_decode/int8_stored_hash diff --git a/fuzz/corpus/chain_wire_compact/bf16_const b/fuzz/corpus/chain_wire_compact/bf16_const new file mode 100644 index 0000000000000000000000000000000000000000..2e2d8112c9576ee6b880fcc1755f14b3a1b648ea GIT binary patch literal 8 PcmZQzWME`qV2}a;04@L> literal 0 HcmV?d00001 diff --git a/fuzz/corpus/chain_wire_compact/bf16_randn b/fuzz/corpus/chain_wire_compact/bf16_randn new file mode 100644 index 0000000000000000000000000000000000000000..618ad3359a7992239b5e8d8481922d00c958d632 GIT binary patch literal 16 XcmZQzW@2Js5MWSXWMWeAP?Q7!1LFYS literal 0 HcmV?d00001 diff --git a/fuzz/corpus/chain_wire_compact/bf16_smooth b/fuzz/corpus/chain_wire_compact/bf16_smooth new file mode 100644 index 0000000000000000000000000000000000000000..d35a4612a40b06232201cf3c57c55d47452aae1b GIT binary patch literal 23 ecmZQzVP;}r5MWSXWC~(nVq#=wU{LTdkOTk;+5t!a literal 0 HcmV?d00001 diff --git a/fuzz/corpus/chain_wire_compact/byte_bf16 b/fuzz/corpus/chain_wire_compact/byte_bf16 new file mode 100644 index 0000000000000000000000000000000000000000..b5a99467020db6c83ec1d941a07c0c817b3ca668 GIT binary patch literal 30 lcmZQ%VP;}rXJlYvVBlvEU{GLW3SwYlVq|7uQ1CF21ON<90b&3E literal 0 HcmV?d00001 diff --git a/fuzz/corpus/chain_wire_compact/fp32_const b/fuzz/corpus/chain_wire_compact/fp32_const new file mode 100644 index 0000000000000000000000000000000000000000..a247021627c9b1499ed4bedb4fdeeef9354c438c GIT binary patch literal 20 bcmZQzW@2Gr5M)qbWMNW>P>fWHQkDb&3rGR% literal 0 HcmV?d00001 diff --git a/fuzz/corpus/chain_wire_compact/fp32_smooth b/fuzz/corpus/chain_wire_compact/fp32_smooth new file mode 100644 index 0000000000000000000000000000000000000000..ced20073b14c7723d99577c41a86b583360778b6 GIT binary patch literal 27 icmZQzVP;`q5M)qbWC>zmVq#`yU{Huqj8uv;kOTk{-U3ws literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/bf16_const_bare b/fuzz/corpus/compact_decode/bf16_const_bare new file mode 100644 index 0000000000000000000000000000000000000000..8b6cd35e98fb7d9b2ba2e113cecd6e7c8c0f9d73 GIT binary patch literal 31 mcmZQ!V`5ZbWME`qV2~0L)Yz*3mm$KHfkA-5o{_=3C;|X8i3E56 literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/bf16_const_crc b/fuzz/corpus/compact_decode/bf16_const_crc new file mode 100644 index 0000000000000000000000000000000000000000..02a8160deac9f56c8fbd45f96fd722a262833c4f GIT binary patch literal 35 qcmZQ$V`5ZbWME`qV2~0L-0Ix4L}RP|Uxo-*1_l8JdqxKDq6h$9#s-)G literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/bf16_const_hash b/fuzz/corpus/compact_decode/bf16_const_hash new file mode 100644 index 0000000000000000000000000000000000000000..a07363a133fce90d69570cb08062a4c6071319f4 GIT binary patch literal 39 ucmZQ&V`5YY)G*1-^O@+y$iT?Jz#t_gsIgW5FGGYY1A_pAJtKp6Q3L>m=LaYN literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/bf16_randn_bare b/fuzz/corpus/compact_decode/bf16_randn_bare new file mode 100644 index 0000000000000000000000000000000000000000..888b28f0b90df72508f9a91a428920f33b23687a GIT binary patch literal 397 zcmV;80doEX1_A;Q5CZ}N000mGAOQjbAV?t@0DuCJlz`q@F`|bhqB9o@tGeYz`&8B=RiL zjrM?v50Fm3FxP5W%>_676YoRy1qGwkM{?RL(7 z6>Ku%*3J|3e-jTJ|}iig&Yz zx5%~%2T4eZxjYRk|CpJN2~bKfV3Pxwj7b79VsQ5rmCu{ z0R$cMi1D!|7l(AuHP+9RzP5Yr;;bEM$!sEZqI7J;r{*4=!3+^Zts>>X-r(UNoq&dU r^9uf8%aZj2Lqo@tGeYz`&8 zB=RiLjrM?v50Fm3FxP5W%>_676YoRy1qGwkM{ z?RL(76>Ku%*3J|3e-jTJ|}i zig&Yzx5%~%2T4eZxjYRk|CpJN2~bKfV3Pxwj7b7 z9VsQ5rmCu{0R$cMi1D!|7l(AuHP+9RzP5Yr;;bEM$!sEZqI7J;r{*4=!3+^Zts>>X z-r(UNoq&dU^9uf8%aZj2Lqo@tGe zYz`&8B=RiLjrM?v50Fm3FxP5W%>_676YoRy1q zGwkM{?RL(76>Ku%*3J|3e-j zTJ|}iig&Yzx5%~%2T4eZxjYRk|CpJN2~bKfV3Pxwj7b7 z9VsQ5rmCu{0R$cMi1D!|7l(AuHP+9RzP5Yr;;bEM$!sEZqI7J;r{*4=!3+^Zts>>X z-r(UNoq&dU^9uf8%aZj2Lq2rV{3JUBikQd%uROJosZc5fF4AT&Ni zK}p87dxeHL<@l;$69L50vg a!KBE7?Z|$x#%g^df~SVX3Ti^_3v@0zIBJss literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/bf16_smooth_crc b/fuzz/corpus/compact_decode/bf16_smooth_crc new file mode 100644 index 0000000000000000000000000000000000000000..3966e3014524699ee9dbd123fc36ca5b245d05d9 GIT binary patch literal 388 zcmV-~0ek)i1_A+q0t5pB000mGAOQkW00II50{{RZNH7@yfCB9RofnXjqP~Z=-uL2} zj)Z}-)%4Myjpo3Tox_v3!??X(2UyO&-JZ086as5oHavd`H6IZ;AS)aLCJkXA0d74N zPd+S3Z3|}{6&e62rV{3JUBikQd%uROJosZc5fF4 zAT&NiK}p87dxeHL<@ il;$69L50vg!KBE7?Z|$x#%g^df~SVX3Ti^_3v@16Ol{Es literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/bf16_smooth_hash b/fuzz/corpus/compact_decode/bf16_smooth_hash new file mode 100644 index 0000000000000000000000000000000000000000..d91b847d3de8b27b422ce9dfef1ea501f87a88cf GIT binary patch literal 388 zcmV-~0ek)h1_A+q0*U76kV5$Pw*&(M000mGAOQkW00II50{{RZNH7@yfC7+`qP~Z= z-uL2}j)Z}-)%4Myjpo3Tox_v3!??X(2UyO&-JZ086as5oHavd`H6IZ;AS)aLCJkXA z0d74NPd+S3Z3|}{6&e62rV{3JUBikQd%uROJosZ zc5fF4AT&NiK}p87dxeHL<@ il;$69L50vg!KBE7?Z|$x#%g^df~SVX3Ti^_3v@2lg>Yj4 literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/byte_bf16_hash b/fuzz/corpus/compact_decode/byte_bf16_hash new file mode 100644 index 0000000000000000000000000000000000000000..1c61d57e5f83316d5065a86abff4f746af0251d9 GIT binary patch literal 674 zcmV;T0$u$<1_6KwaE)EtN1yGu1Ooy92LS*A000jF5C9+n0#X110s#X603b*(832F; zkdmUlhqm7L;+c+wfwI-~(Vvaxz>=NAleojUyJUEm2HnV|iX)d29y_2`vu~d|h636@CQ`Mi2-Q z85}1yL?}Tc932=IRBA6Q1WaK^PYo(pJ1r9>IYiXnqI{SZP>1J7H~SLO512LtR%tJslqqYY=xOR6t)+J~t;oL?Iq} zE=+qE2?tm=Cm>WZ7dKZ)B1=5!fd>1G@K3HFL0RjjB za&2dCcLD_w6%h~?6Audke*^*m0|^rw9}^WDCrD95K0{4NLPSL~J~~K0O+!gjO-){H zY;9`9h;9`9h;9`9h;P>fWHQkE1F)Yz*3mm$KHfkA+QfstV{gBa;TwY6Z4 E0QpG~IsgCw literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_const_crc b/fuzz/corpus/compact_decode/fp32_const_crc new file mode 100644 index 0000000000000000000000000000000000000000..d6683c997d359719f813c2cd4f444cddd00139a0 GIT binary patch literal 116 zcmZQ$WMXt+W@2Gr5M)qbWMNW>P>fWHQkE1FG)wL+)!3^4mm$KHfkA+QfstV{gBYp8 P>slYDAq=Ul1sehY-qje1 literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_const_hash b/fuzz/corpus/compact_decode/fp32_const_hash new file mode 100644 index 0000000000000000000000000000000000000000..d5408c4577dfff4a1bddbed2e61fa62c010d53dc GIT binary patch literal 108 zcmZQ&WMXvq`u_j4VtF5sHyYQOc4+f*M=(|1w0lGB5}*FfcMq NW)LG?sJ0fY5dhhp75xAJ literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_randn_bare b/fuzz/corpus/compact_decode/fp32_randn_bare new file mode 100644 index 0000000000000000000000000000000000000000..70356b33e9cd61eeb6029cd4995d90dcfeec4031 GIT binary patch literal 489 zcmV5fgKfb$^#8kela74C zMg(1pTgpyrE9<1d0Gu<(WGHaZ*7!WRoKqJsRkxm>-t?*v2wfMHBz3LXBf|VXKRKAr zi^5OA0M6E{Of?-^MS`pgpW9931BF7JUW z{`Yg0pf%(GfB_qTOC?1o%wGYtFvD+a@-`zB5FHJCKwuv?3g3s_ zPZHW!1Q<_o6&&4^6;gJ1uSNa%>@d26FB&vWJk5=~tk#mY3`H&Co5_aJ4PWP6+pfXg zSjejyY8F5yxK`SB52y5q140$pR5fSD%H%#mTEm@ajge~( zCA}o_E6|Mf|A-EdPQNeKX;;hzH~SOt%Z2m_HBzoH1*h5R4}2XhXldO@?ckf0i*_>W z=AZ3#&V3bZF%m^6wJ-euSmgl#LPrb=;AspgCHWWtR2X}71~bHnT7@K!y}`pHoj^nR f@|u6POl18rWMm`-(OZ4_sGGDt0sufY+#kvq#H`hR literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_randn_crc b/fuzz/corpus/compact_decode/fp32_randn_crc new file mode 100644 index 0000000000000000000000000000000000000000..fb7510fdc8251b8ad79196c26bc1f5003d17adcd GIT binary patch literal 505 zcmV5fgKfb$^#8ke zla74CMg(1pTgpyrE9<1d0Gu<(WGHaZ*7!WRoKqJsRkxm>-t?*v2wfMHBz3LXBf|VX zKRKAri^5OA0M6E{Of?-^MS`pgpW9931B zF7JUW{`Yg0pf%(GfB_-Ry*wL$OC?1o%wGYtFvD+a@-`zB5FHJC zKwuv?3g3s_PZHW!1Q<_o6&&4^6;gJ1uSNa%>@d26FB&vWJk5=~tk#mY3`H&Co5_aJ z4PWP6+pfXgSjejyY8F5yxK`SB52y5q140$pR5fSD z%H%#mTEm@ajge~(CA}o_E6|Mf|A-EdPQNeKX;;hzH~SOt%Z2m_HBzoH1*h5R4}2Xh zXldO@?ckf0i*_>W=AZ3#&V3bZF%m_HM^_RkwJ-euSmgl#LPrb=;AspgCHWWtR2X}7 v1~bHnT7@K!y}`pHoj^nR@|u6POl18rWMm`-(OZ4_sGGDt0sufY+#kvqTl?Jg literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_randn_hash b/fuzz/corpus/compact_decode/fp32_randn_hash new file mode 100644 index 0000000000000000000000000000000000000000..65c1b1540700656852b304e4583ebff78ebf92cf GIT binary patch literal 497 zcmV5fgKfb$ z^#8kela74CMg(1pTgpyrE9<1d0Gu<(WGHaZ*7!WRoKqJsRkxm>-t?*v2wfMHBz3LX zBf|VXKRKAri^5OA0M6E{Of?-^MS`pgpW z9931BF7JUW{`Yg0pf%(GfB_qTOC?1o%wGYtFvD+a@-`zB5FHJC zKwuv?3g3s_PZHW!1Q<_o6&&4^6;gJ1uSNa%>@d26FB&vWJk5=~tk#mY3`H&Co5_aJ z4PWP6+pfXgSjejyY8F5yxK`SB52y5q140$pR5fSD%H%#m zTEm@ajge~(CA}o_E6|Mf|A-EdPQNeKX;;hzH~SOt%Z2m_HBzoH1*h5R4}2XhXldO@ z?ckf0i*_>W=AZ3#&V3bZF%m^6wJ-euSmgl#LPrb=;AspgCHWWtR2X}71~bHnT7@K! ny}`pHoj^nR@|u6POl18rWMm`-(OZ4_sGGDt0sufY+#kvq(d^nj literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_smooth_bare b/fuzz/corpus/compact_decode/fp32_smooth_bare new file mode 100644 index 0000000000000000000000000000000000000000..5ef04dcc209c58b38ad659f106f966b9fed54875 GIT binary patch literal 902 zcmV;119|)e0RjPl0t5pD000pHAOQqY00II70{{RZSRq*=S}+*^fC7aGf-dP0D^@U7 zMF~sI4BXFocpA5|AD4VbH5+!J*z}gE>CKW)4pz0O)MR`~H_-#|ay@M8LjPGdcu&~S z0O6Sd&$>Z_od{4MQtmp?ZYp)fOPOU*OG>N9VY2N-_x|{94sVD9L8W+8MY{XkRug2b z2{anur1GZ^gEcvY%Y^4D27sGUjWWJSPYku;dpCAWDw39%du;1yj+1QAiBbe2qgx4( zmb7F9k8Hk)e29k+K&l;>HJu*bv8gZdJ>iCdCvLZEOHm zDZ?8giLV3zJuBO1%!eRfcmVM(Sl*bH$b)d-P5h6R+7@*NKxVPSAPq4@?&@Z<)m~`B zjG1kf9ubW}C+z$fvx@Z8)uVqSTjgI~}pz&R=pP zIc{8GxVr--rdEjwg!$njvr#AwvS@TvnR{naer4_Qy4`^W&JQYtp zD@kk$W*ih60B&DTEl^BmV|iX)d29v@2rUl|dtF|16n+E@Mh^%P85<@wLnuKb932=I zQ)(|O159B@PYfznIxQ0=IY$FzVP#S(CkA)~cYg;gHbOi&J|$6FEkR3T5Mp(17Y83R zJw!lELLVn}a9M0VBrFIaeoz-6H5FeHT@WioPXlxt7;6F*Dgze literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_smooth_crc b/fuzz/corpus/compact_decode/fp32_smooth_crc new file mode 100644 index 0000000000000000000000000000000000000000..684f7fa7d7a61b73ac7b2e291093f33eb362b634 GIT binary patch literal 918 zcmV;H18MvQ0RjPl0t5pD000pHAOQqY00II70{{RZSRq*=S}+*^fC7}fI&y^wf-dP0 zD^@U7MF~sI4BXFocpA5|AD4VbH5+!J*z}gE>CKW)4pz0O)MR`~H_-#|ay@M8LjPGd zcu&~S0O6Sd&$>Z_od{4MQtmp?ZYp)fOPOU*OG>N9VY2N-_x|{94sVD9L8W+8MY{Xk zRug2b2{anur1GZ^gEcvY%Y^4D27sGUjWWJSPYku;dpCAWDw39%du;1yj+1QAiBbe2 zqgx4(mb7F9k8Hk)e29k+K&l;>HJu*bv8gZdJ>i zCdCvLZEOHmDZ?8giLV3zJuBO1%!eRfcmVM(Sl*bH$b)d-P5h6R+7@*NKxVPSAPq4@ z?&@Z<)m~`BjG1kf9ubW}C+z$fvx@Z8)uVqSTjg zI~}pz&R=pPIc{8GxVr--rdEjwg!$njvr#AwvS@TvnR{naer4_Qy4`^W&JQYtpD@kk$W*ih60B&DTEl^BmV|iX)d29v@2rUl|dtF|16n+E@Mh^%P z85<@wLnuKb932=IQ)(|O159B@PYfznIxQ0=IY$FzVP#S(CkA)~cYg;gHbOi&J|$6F zEkR3T5Mp(17Y83RJw!lELLVn}a9M0VBrFIaeoz-6H5FeHT@WioPXlxt7;6F*Dgze< zD=tlFKwnQLMQ&MZbaPQnLUnrs0SgCXUQIn%JXd@}6gWL_R84YcV-hbQ1PMk`9$Tec z#fkv{0RR9C3IG6@1py!d0tf&D2mk~K0DlMo{0IR32mt;F0RIpHXaN8K00000!ETgy s42YikIcR+ra-7cTD}{HbO85pEK;v^pGOWgGeItUWhQ$hMLhTE5E&?HJMF0Q* literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp32_smooth_hash b/fuzz/corpus/compact_decode/fp32_smooth_hash new file mode 100644 index 0000000000000000000000000000000000000000..ead5f2d9be1996b2f82445c48c8551f1f444d52e GIT binary patch literal 910 zcmV;919AKX0RjPl0-^T2kr`Y*1q1^G000pHAOQqY00II70{{RZSRq*=S}+*^fC7aG zf-dP0D^@U7MF~sI4BXFocpA5|AD4VbH5+!J*z}gE>CKW)4pz0O)MR`~H_-#|ay@M8 zLjPGdcu&~S0O6Sd&$>Z_od{4MQtmp?ZYp)fOPOU*OG>N9VY2N-_x|{94sVD9L8W+8 zMY{XkRug2b2{anur1GZ^gEcvY%Y^4D27sGUjWWJSPYku;dpCAWDw39%du;1yj+1QA ziBbe2qgx4(mb7F9k8Hk)e29k+K&l;>HJu*bv8gZdJ>i zCdCvLZEOHmDZ?8giLV3zJuBO1%!eRfcmVM(Sl*bH$b)d-P5h6R+7@*NKxVPSAPq4@ z?&@Z<)m~`BjG1kf9ubW}C+z$fvx@Z8)uVqSTjg zI~}pz&R=pPIc{8GxVr--rdEjwg!$njvr#AwvS@TvnR{naer4_Qy4`^W&JQYtpD@kk$W*ih60B&DTEl^BmV|iX)d29v@2rUl|dtF|16n+E@Mh^%P85<@w zLnuKb932=IQ)(|O159B@PYfznIxQ0=IY$FzVP#S(CkA)~cYg;gHbOi&J|$6FEkR3T z5Mp(17Y83RJw!lELLVn}a9M0VBrFIaeoz-6H5FeHT@WioPXlxt7;6F*Dgze0lB%Tt~N8Oz$&^cDjzyIsJp%~H8Q-bKQytUBs)DfwXzVfC@vl{ zE}b>06)-`(y|XqrBd4t~t*$gYM!dQ|v9UKfI=n)!E4{I=u(>&*HJmfG05z>RE2u%g omAJPNGd;U2Fv7m4DJU8%FE6$=x2D3GkvA*1ma?xutspr!J1#mmg#Z8m literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp8_e4m3_crc b/fuzz/corpus/compact_decode/fp8_e4m3_crc new file mode 100644 index 0000000000000000000000000000000000000000..54f284ede5819e04f5095192402f4f86a555d3e7 GIT binary patch literal 134 zcmV;10D1o;9RdM>0lB%Tt~N8Oz$&^cDjzyIsJp%~H8Q-bKQytUBs)DfwXzVfC@vl{ zE}b>06)-`(y|XqrBd4t~t*$gYM!dQ|v9UKfI=n)!E4{I=u(>&*HJmfG05z>RE2u%g omAJPNGd;U2Fv7m4DJU8%FE6$=x2D3GkvA*1ma?xutspr!J1#mmg#Z8m literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/fp8_e4m3_hash b/fuzz/corpus/compact_decode/fp8_e4m3_hash new file mode 100644 index 0000000000000000000000000000000000000000..0dc87ce623d0b02b5ae48ad9c5b0a93a32efa08d GIT binary patch literal 142 zcmV;90CE2%9RdM>0l2(tg06)-`(y|XqrBd4t~t*$gYM!dQ|v9UKfI=n)!E4{I=u(>&*HJmfG w05z>RE2u%gmAJPNGd;U2Fv7m4DJU8%FE6$=x2D3GkvA*1ma?xutspr!JG!epX8-^I literal 0 HcmV?d00001 diff --git a/fuzz/corpus/compact_decode/int8_stored_bare b/fuzz/corpus/compact_decode/int8_stored_bare new file mode 100644 index 0000000..e360453 --- /dev/null +++ b/fuzz/corpus/compact_decode/int8_stored_bare @@ -0,0 +1 @@ +$@Bi_L#sщ2~F]lQ3Ny@-abcn-9DN s \ No newline at end of file diff --git a/fuzz/corpus/compact_decode/int8_stored_crc b/fuzz/corpus/compact_decode/int8_stored_crc new file mode 100644 index 0000000..e360453 --- /dev/null +++ b/fuzz/corpus/compact_decode/int8_stored_crc @@ -0,0 +1 @@ +$@Bi_L#sщ2~F]lQ3Ny@-abcn-9DN s \ No newline at end of file diff --git a/fuzz/corpus/compact_decode/int8_stored_hash b/fuzz/corpus/compact_decode/int8_stored_hash new file mode 100644 index 0000000..aad562a --- /dev/null +++ b/fuzz/corpus/compact_decode/int8_stored_hash @@ -0,0 +1 @@ +%@3׽hNBi_L#sщ2~F]lQ3Ny@-abcn-9DN s \ No newline at end of file From b7c3dcef5d5064f4fd3746a82c1509169bc91957 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 00:16:25 +0000 Subject: [PATCH 14/34] wip(S6-ptwc): snapshot of an interrupted implementation Partial work on the compact multi-tensor container, saved so a container restart cannot lose it. May not compile yet; the resumed Stage 6 implementer completes it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/wire_compact.rs | 155 +++- crates/ptwm-core/src/compact.rs | 9 +- crates/ptwm-core/src/container.rs | 811 ++++++++++++++---- crates/ptwm-core/src/dtype.rs | 10 + crates/ptwm-core/src/extension/builtins.rs | 13 + .../ptwm-core/src/extension/compact_table.rs | 409 +++++++++ crates/ptwm-core/src/extension/mod.rs | 1 + crates/ptwm-core/src/extension/table.rs | 22 + crates/ptwm-core/src/header.rs | 46 + crates/ptwm-core/src/index.rs | 8 + crates/ptwm-core/src/lib.rs | 2 + crates/ptwm-core/src/ptwc.rs | 294 +++++++ crates/ptwm-core/src/tensor_record_compact.rs | 798 +++++++++++++++++ 13 files changed, 2410 insertions(+), 168 deletions(-) create mode 100644 crates/ptwm-core/src/extension/compact_table.rs create mode 100644 crates/ptwm-core/src/ptwc.rs create mode 100644 crates/ptwm-core/src/tensor_record_compact.rs diff --git a/crates/ptwm-core/src/chain/wire_compact.rs b/crates/ptwm-core/src/chain/wire_compact.rs index 4ed7021..27d5dcd 100644 --- a/crates/ptwm-core/src/chain/wire_compact.rs +++ b/crates/ptwm-core/src/chain/wire_compact.rs @@ -108,6 +108,41 @@ pub fn op_from_wire_ref(op_ref: u64) -> Result { pub fn write_chain_compact( chain: &Chain, bound_source: Option<&[u8]>, + op_ref: impl FnMut(OpId) -> Result, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + write_chain_with(chain, SourceBinding::Tensor(bound_source), op_ref, out) +} + +/// Encode a shared chain (a container registry entry) as CW2 into `out`. +/// +/// A registry chain serves many tensors, so it cannot be bound to one. In +/// memory, a Source node with empty params stands for a bound Source: it is +/// written as one, and each tensor record that uses the chain supplies the +/// params ([`SourceBinding::Deferred`]). Other Source nodes are written with +/// their params. +pub fn write_chain_compact_deferred( + chain: &Chain, + op_ref: impl FnMut(OpId) -> Result, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + write_chain_with(chain, SourceBinding::Deferred, op_ref, out) +} + +/// How a chain's Source params relate to the tensor it encodes. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum SourceBinding<'a> { + /// The chain belongs to one tensor, whose bound Source params are given + /// (`None` when the tensor cannot bind a Source). + Tensor(Option<&'a [u8]>), + /// The chain is shared; a bound Source stays unbound (empty params) in + /// memory, and the tensor record that uses the chain binds it. + Deferred, +} + +fn write_chain_with( + chain: &Chain, + binding: SourceBinding<'_>, mut op_ref: impl FnMut(OpId) -> Result, out: &mut Vec, ) -> Result<(), PtwmCoreError> { @@ -137,11 +172,17 @@ pub fn write_chain_compact( } write_varint(out, op_ref(node.op)?); if node.op == OpId::Source { - if node.params.is_empty() { - // An empty Source would read back as a bound Source. - return Err(err("Source node with empty params".into())); - } - if bound_source == Some(node.params.as_slice()) { + let bind = match binding { + SourceBinding::Tensor(bound_source) => { + if node.params.is_empty() { + // An empty Source would read back as a bound Source. + return Err(err("Source node with empty params".into())); + } + bound_source == Some(node.params.as_slice()) + } + SourceBinding::Deferred => node.params.is_empty(), + }; + if bind { write_varint(out, 0); continue; } @@ -206,6 +247,16 @@ pub fn write_chain_compact( pub fn read_chain_compact( r: &mut Cursor<'_>, bound_source: Option<&[u8]>, + resolve: impl FnMut(u64) -> Result, +) -> Result { + read_chain_with(r, SourceBinding::Tensor(bound_source), resolve) +} + +/// Decode one CW2 chain with an explicit [`SourceBinding`]. With +/// [`SourceBinding::Deferred`] a bound Source is returned with empty params. +pub fn read_chain_with( + r: &mut Cursor<'_>, + binding: SourceBinding<'_>, mut resolve: impl FnMut(u64) -> Result, ) -> Result { let n_nodes = r.varint_max(MAX_COUNT, "chain node count")? as usize; @@ -226,14 +277,20 @@ pub fn read_chain_compact( let op = resolve(op_ref)?; let param_len = r.varint_max(MAX_PARAMS, "chain node params length")? as usize; let params = if op == OpId::Source && param_len == 0 { - bound_source - .ok_or_else(|| { - r.err("chain has a bound Source but the tensor has no bindable shape and dtype") - })? - .to_vec() + match binding { + SourceBinding::Tensor(bound_source) => bound_source + .ok_or_else(|| { + r.err( + "chain has a bound Source but the tensor has no bindable shape \ + and dtype", + ) + })? + .to_vec(), + SourceBinding::Deferred => Vec::new(), + } } else { let params = r.take(param_len)?; - if op == OpId::Source && bound_source == Some(params) { + if op == OpId::Source && binding == SourceBinding::Tensor(Some(params)) { return Err(r.err("chain Source params equal the bound Source (non-canonical)")); } params.to_vec() @@ -315,6 +372,37 @@ pub fn read_chain_compact( }) } +/// Decode a shared (registry) CW2 chain that must fill `buf` exactly. A +/// bound Source comes back with empty params; see +/// [`write_chain_compact_deferred`]. +pub fn read_chain_compact_deferred_exact( + buf: &[u8], + resolve: impl FnMut(u64) -> Result, +) -> Result { + let mut r = Cursor::new(buf, "compact chain"); + let chain = read_chain_with(&mut r, SourceBinding::Deferred, resolve)?; + r.finish("chain")?; + Ok(chain) +} + +/// Whether `chain` has a bound Source that a tensor record must supply (a +/// Source node with empty params; see [`write_chain_compact_deferred`]). +pub fn has_deferred_source(chain: &Chain) -> bool { + chain + .nodes + .iter() + .any(|n| n.op == OpId::Source && n.params.is_empty()) +} + +/// Replace the params of every bound Source of `chain` with `params`. +pub fn bind_deferred_source(chain: &mut Chain, params: &[u8]) { + for n in &mut chain.nodes { + if n.op == OpId::Source && n.params.is_empty() { + n.params = params.to_vec(); + } + } +} + /// Decode a CW2 chain that must fill `buf` exactly. pub fn read_chain_compact_exact( buf: &[u8], @@ -451,6 +539,51 @@ mod tests { use super::*; use crate::transforms::source::bound_source_params; + #[test] + fn deferred_source_round_trips_unbound() { + let params = bound_source_params(&[4, 8], 1).unwrap(); + let mut chain = Chain { + nodes: vec![ + ChainNode { + op: OpId::Source, + params: Vec::new(), + }, + ChainNode { + op: OpId::BytePassthrough, + params: Vec::new(), + }, + ], + edges: vec![spine_edge(0)], + terminals: vec![TerminalRef { + node_idx: 1, + output_idx: 0, + role: Role::Raw, + }], + }; + let mut deferred = Vec::new(); + write_chain_compact_deferred(&chain, op_wire_ref, &mut deferred).unwrap(); + // The tensor-bound writer refuses an unbound Source. + let mut out = Vec::new(); + assert!(write_chain_compact(&chain, Some(¶ms), op_wire_ref, &mut out).is_err()); + let back = read_chain_compact_deferred_exact(&deferred, op_from_wire_ref).unwrap(); + assert_eq!(back, chain); + assert!(has_deferred_source(&back)); + // The same bytes read with a tensor binding give the bound chain. + let bound = read(&deferred, Some(¶ms)).unwrap(); + bind_deferred_source(&mut chain, ¶ms); + assert_eq!(bound, chain); + assert!(!has_deferred_source(&bound)); + // An explicit Source stays explicit in a shared chain, even when its + // params equal some tensor's binding. + let mut explicit = Vec::new(); + write_chain_compact_deferred(&chain, op_wire_ref, &mut explicit).unwrap(); + assert_eq!( + read_chain_compact_deferred_exact(&explicit, op_from_wire_ref).unwrap(), + chain + ); + assert!(explicit.len() > deferred.len()); + } + fn write(chain: &Chain, bound: Option<&[u8]>) -> Vec { let mut out = Vec::new(); write_chain_compact(chain, bound, op_wire_ref, &mut out).unwrap(); diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index e71b282..fce88ea 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -89,7 +89,7 @@ use crate::chain::wire_compact::{ }; use crate::codec::{CodecId, StateSource, codec_for}; use crate::compressor::{CompressorOptions, InputTensor, legacy_plane_layout, plan_model}; -use crate::dtype::Dtype; +use crate::dtype::{Dtype, canonical_dtype_name}; use crate::error::PtwmCoreError; use crate::transforms::op::Plane; use crate::transforms::source::bound_source_params; @@ -159,13 +159,6 @@ pub fn is_compact(buf: &[u8]) -> bool { buf.len() >= MAGIC_COMPACT.len() && buf[..MAGIC_COMPACT.len()] == MAGIC_COMPACT } -fn canonical_dtype_name(dtype_code: u16) -> Option<&'static str> { - u8::try_from(dtype_code) - .ok() - .and_then(|c| Dtype::from_code(c).ok()) - .map(Dtype::canonical_name) -} - /// Raw byte count implied by a shape and a public dtype code: /// `product(shape) × element_size`. `None` for an unknown dtype code or on /// overflow. diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 43e4ebe..4b1e920 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -11,9 +11,14 @@ use std::collections::{HashMap, HashSet}; use crate::chain::Chain; use crate::chain::wire::{read_chain, write_chain}; +use crate::chain::wire_compact::{bind_deferred_source, has_deferred_source}; +use crate::codec::CodecId; use crate::error::PtwmCoreError; use crate::extension::{CanonicalId, ExtensionTable, ExtensionTableBuilder}; -use crate::header::{HEADER_LEN, Header, MAGIC_PTWM}; +use crate::header::{ContainerMagic, HEADER_LEN, Header, MAGIC_PTWC, MAGIC_PTWM, MAGIC_PTWX, sniff}; +use crate::tensor_record_compact::{ + RecordContext, parse_compact_record, plane_codec_ids, record_binding, write_compact_record, +}; use crate::index::{ EOF_SENTINEL, IndexEntry, name_hash, parse_index, verify_eof_sentinel, write_index, }; @@ -112,6 +117,12 @@ impl ChainRegistry { /// Return a reference to the `Chain` with the given `chain_id`, or `None`. pub fn lookup(&self, chain_id: u16) -> Option<&Chain> { + // Writers assign ids by position; try that first. + if let Some((id, chain)) = self.chains.get(usize::from(chain_id)) + && *id == chain_id + { + return Some(chain); + } self.chains .iter() .find(|(id, _)| *id == chain_id) @@ -268,9 +279,112 @@ pub fn encode_registry_for_table( }) } +/// Encode `registry` as a compact registry whose op references index +/// `table` (the entry with each op's built-in id). +fn compact_registry_bytes( + registry: &ChainRegistry, + table: &ExtensionTable, +) -> Result, PtwmCoreError> { + let reg = crate::extension::builtin_registry(); + let mut index: HashMap = HashMap::new(); + for (i, e) in table.entries.iter().enumerate() { + index.entry(e.canonical_id).or_insert(i as u64); + } + crate::ptwc::write_registry(registry, |op| { + reg.id_for_entry_name(op.builtin_name()) + .and_then(|id| index.get(&id).copied()) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "op {op:?} is not in the extension table" + )) + }) + }) +} + +/// On-disk layout of a container. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] +pub enum ContainerLayout { + /// `\x89PTWM`: fixed-width header, full Extension Table, hash-sorted + /// index. Readable by every release. + #[default] + Legacy, + /// `\x89PTWC`: varint fields, built-ins stored by id, names in the index + /// ([`crate::ptwc`]). Readable by v1.2 and later. + Compact, +} + +impl ContainerLayout { + /// The layout's name as the Python bindings spell it. + pub fn as_str(self) -> &'static str { + match self { + Self::Legacy => "legacy", + Self::Compact => "compact", + } + } + + /// Parse a layout name (`"legacy"` or `"compact"`). + pub fn parse(name: &str) -> Option { + match name { + "legacy" => Some(Self::Legacy), + "compact" => Some(Self::Compact), + _ => None, + } + } +} + +/// Serializes tensor records for one container writer's layout and tables. +/// +/// [`ContainerWriter::record_encoder`] returns it so that callers can +/// serialize records in parallel and append the bytes with +/// [`ContainerWriter::append_serialized_tensor`]. +#[derive(Debug, Clone)] +pub struct RecordEncoder { + layout: ContainerLayout, + builtin_version: String, + codec_ids: Vec>, + registry_bound: Vec, +} + +impl RecordEncoder { + /// Serialize `rec` in the writer's layout. A compact record leaves out + /// the name, which the index stores. + pub fn encode(&self, rec: &TensorRecord) -> Result, PtwmCoreError> { + let est: usize = rec + .terminals + .iter() + .map(|t| t.payload_bytes.len() + t.inline_state_bytes.len()) + .sum(); + let mut buf = Vec::with_capacity(est + 256); + self.encode_into(rec, &mut buf)?; + Ok(buf) + } + + /// [`encode`](Self::encode) into `out`; on error `out` is unchanged. + pub fn encode_into(&self, rec: &TensorRecord, out: &mut Vec) -> Result<(), PtwmCoreError> { + match self.layout { + ContainerLayout::Legacy => { + write_tensor_record_with_builtin_version(rec, out, &self.builtin_version) + } + ContainerLayout::Compact => { + let ctx = RecordContext { + codec_ids: &self.codec_ids, + registry_bound: &self.registry_bound, + }; + write_compact_record(rec, &ctx, out) + .map_err(|e| PtwmCoreError::InvalidContainer(format!("tensor '{}': {e}", rec.name))) + } + } + } +} + /// Streaming writer for a `.ptwm` container. pub struct ContainerWriter { writer: W, + layout: ContainerLayout, + /// Record names and lengths in body order (compact layout). + compact_index: Vec<(String, u64)>, + /// Record encoder state for the current tables. + encoder: RecordEncoder, header: Header, header_position: u64, tensors_start_pos: u64, @@ -308,24 +422,42 @@ pub struct ContainerWriter { } impl ContainerWriter { - /// Create a new container writer. Writes a 91-byte header placeholder - /// immediately. - pub fn new(mut writer: W) -> Result { + /// Create a new legacy-layout container writer. Writes a 91-byte header + /// placeholder immediately. + pub fn new(writer: W) -> Result { + Self::with_layout(writer, ContainerLayout::Legacy) + } + + /// Create a container writer for `layout`. The legacy layout writes a + /// header placeholder now and patches it in [`finalize`](Self::finalize); + /// the compact layout writes its header with the first tensor and never + /// seeks. + pub fn with_layout(mut writer: W, layout: ContainerLayout) -> Result { let header_position = writer .stream_position() .map_err(|e| PtwmCoreError::Io(e.to_string()))?; - // Write placeholder (magic bytes, rest zeros). - let mut placeholder = [0u8; HEADER_LEN]; - placeholder[..9].copy_from_slice(&MAGIC_PTWM); - writer - .write_all(&placeholder) - .map_err(|e| PtwmCoreError::Io(e.to_string()))?; + if layout == ContainerLayout::Legacy { + // Write placeholder (magic bytes, rest zeros). + let mut placeholder = [0u8; HEADER_LEN]; + placeholder[..9].copy_from_slice(&MAGIC_PTWM); + writer + .write_all(&placeholder) + .map_err(|e| PtwmCoreError::Io(e.to_string()))?; + } let header = Header::new_ptwm(); Ok(Self { writer, + layout, + compact_index: Vec::new(), + encoder: RecordEncoder { + layout, + builtin_version: crate::extension::CURRENT_BUILTIN_VERSION.to_owned(), + codec_ids: Vec::new(), + registry_bound: Vec::new(), + }, header, header_position, tensors_start_pos: 0, @@ -411,13 +543,103 @@ impl ContainerWriter { "chain registry must be set before appending tensors".into(), )); } + self.check_registry_for_layout(®istry)?; let built = build_extension_table(®istry, codecs, &self.builtin_version)?; + self.chain_registry_bytes = match self.layout { + ContainerLayout::Legacy => built.registry_bytes, + ContainerLayout::Compact => compact_registry_bytes(®istry, &built.table)?, + }; self.extension_table = built.table; self.raw_extension_table = None; self.table_set = true; - self.chain_registry_bytes = built.registry_bytes; self.chain_registry = registry; self.codec_idx_map = built.codec_idx_map; + self.refresh_encoder(); + Ok(()) + } + + /// Fail on registry chains the layout cannot store: a bound Source (a + /// Source node with empty params) needs the compact layout. + fn check_registry_for_layout(&self, registry: &ChainRegistry) -> Result<(), PtwmCoreError> { + if self.layout == ContainerLayout::Legacy + && let Some((id, _)) = registry.chains.iter().find(|(_, c)| has_deferred_source(c)) + { + return Err(PtwmCoreError::InvalidContainer(format!( + "chain {id} has a bound Source (empty Source params), which needs the \ + compact layout" + ))); + } + Ok(()) + } + + /// Recompute the record encoder's view of the tables. + fn refresh_encoder(&mut self) { + self.encoder.builtin_version = self.builtin_version.clone(); + self.encoder.codec_ids = plane_codec_ids(&self.extension_table); + self.encoder.registry_bound = self + .chain_registry + .chains + .iter() + .map(|(_, c)| has_deferred_source(c)) + .collect(); + } + + /// The writer's layout. + pub fn layout(&self) -> ContainerLayout { + self.layout + } + + /// A record encoder for the current layout and tables. Take it after the + /// Extension Table and chain registry are set. + pub fn record_encoder(&self) -> RecordEncoder { + self.encoder.clone() + } + + /// Serialize `rec` the way [`append_tensor`](Self::append_tensor) would, + /// for [`append_serialized_tensor`](Self::append_serialized_tensor). + pub fn encode_record(&self, rec: &TensorRecord) -> Result, PtwmCoreError> { + self.encoder.encode(rec) + } + + /// Use `table` as the Extension Table and encode `registry` against it + /// (compact layout only). Registry ops resolve to the table entry with + /// their built-in id, and [`codec_table_idx_for`](Self::codec_table_idx_for) + /// to the entry of each built-in codec. The transcoder uses this to write + /// a table and registry exactly as they were read. + /// + /// Must be called before any tensor is appended. + pub fn set_compact_tables( + &mut self, + table: ExtensionTable, + registry: ChainRegistry, + ) -> Result<(), PtwmCoreError> { + if self.layout != ContainerLayout::Compact { + return Err(PtwmCoreError::InvalidContainer( + "set_compact_tables needs the compact layout".into(), + )); + } + if self.prelude_written || !self.index_entries.is_empty() { + return Err(PtwmCoreError::InvalidContainer( + "extension table must be set before appending tensors".into(), + )); + } + // Fails early on entries the compact table cannot store. + crate::extension::compact_table::compact_table_bytes(&table)?; + self.chain_registry_bytes = compact_registry_bytes(®istry, &table)?; + let mut codec_idx_map = HashMap::new(); + for (i, e) in table.entries.iter().enumerate() { + if let Some(crate::extension::BuiltinKind::Codec(_)) = + crate::extension::dispatch_builtin(&e.canonical_id) + { + codec_idx_map.entry(e.canonical_id).or_insert(i as u16); + } + } + self.extension_table = table; + self.raw_extension_table = None; + self.chain_registry = registry; + self.codec_idx_map = codec_idx_map; + self.table_set = true; + self.refresh_encoder(); Ok(()) } @@ -437,6 +659,7 @@ impl ContainerWriter { self.extension_table = table; self.raw_extension_table = None; self.table_set = true; + self.refresh_encoder(); } /// Write built-in Extension Table entries with the ids and labels of @@ -454,12 +677,18 @@ impl ContainerWriter { "builtin version must be set before the extension table and tensors".into(), )); } + if self.layout == ContainerLayout::Compact { + return Err(PtwmCoreError::InvalidContainer( + "the compact layout stores built-ins without a version".into(), + )); + } if !crate::extension::is_valid_builtin_version(version) { return Err(PtwmCoreError::InvalidContainer(format!( "invalid builtin version {version:?}" ))); } self.builtin_version = version.to_owned(); + self.refresh_encoder(); Ok(()) } @@ -483,6 +712,12 @@ impl ContainerWriter { "extension table must be set before appending tensors".into(), )); } + if self.layout != ContainerLayout::Legacy { + return Err(PtwmCoreError::InvalidContainer( + "set_raw_extension_table needs the legacy layout".into(), + )); + } + self.check_registry_for_layout(®istry)?; let built = encode_registry_for_table(table_bytes, ®istry)?; self.extension_table = built.table; self.raw_extension_table = Some(table_bytes.to_vec()); @@ -490,6 +725,7 @@ impl ContainerWriter { self.chain_registry = registry; self.codec_idx_map = built.codec_idx_map; self.table_set = true; + self.refresh_encoder(); Ok(()) } @@ -501,6 +737,9 @@ impl ContainerWriter { if self.prelude_written { return Ok(()); } + if self.layout == ContainerLayout::Compact { + return self.flush_compact_prelude(); + } // Extension table sits immediately after the header at HEADER_LEN. // Serialize it now (set_extension_table must be called before this point). @@ -574,12 +813,40 @@ impl ContainerWriter { Ok(()) } + /// Write the compact header, Extension Table, prelude and registry. + fn flush_compact_prelude(&mut self) -> Result<(), PtwmCoreError> { + let et = crate::extension::compact_table::compact_table_bytes(&self.extension_table)?; + let mut prelude = Vec::new(); + write_prelude(&self.prelude_entries, &mut prelude)?; + let registry = if self.chain_registry_bytes.is_empty() { + crate::ptwc::write_registry(&self.chain_registry, |op| { + Err(PtwmCoreError::InvalidContainer(format!( + "op {op:?} is not in the extension table" + ))) + })? + } else { + self.chain_registry_bytes.clone() + }; + let header = crate::ptwc::header_bytes(&et, &prelude, ®istry); + for part in [&header, &et, &prelude, ®istry] { + self.writer + .write_all(part) + .map_err(|e| PtwmCoreError::Io(e.to_string()))?; + } + self.prelude_written = true; + Ok(()) + } + /// Append a tensor record to the container. pub fn append_tensor(&mut self, record: &TensorRecord) -> Result<(), PtwmCoreError> { if self.finalized { return Err(PtwmCoreError::InvalidContainer("already finalized".into())); } self.check_new_name(&record.name)?; + if self.layout == ContainerLayout::Compact { + let bytes = self.encoder.encode(record)?; + return self.append_serialized_tensor(&record.name, &bytes); + } // On first tensor: flush prelude and record tensors_start_pos. if !self.prelude_written { @@ -604,6 +871,7 @@ impl ContainerWriter { name_hash: name_hash(&record.name), tensor_offset, tensor_len, + name: None, }); self.names.insert(record.name.clone()); @@ -622,10 +890,12 @@ impl ContainerWriter { Ok(()) } - /// Append a tensor whose record bytes were serialized ahead of time (e.g. - /// in parallel via [`crate::tensor_record::write_tensor_record`]). - /// Identical on-disk effect to [`Self::append_tensor`] but skips the - /// (payload-copy-heavy) serialization, which the caller has already done. + /// Append a tensor whose record bytes were serialized ahead of time by + /// [`encode_record`](Self::encode_record) or a [`RecordEncoder`] (for + /// example in parallel). Identical on-disk effect to + /// [`Self::append_tensor`] but skips the (payload-copy-heavy) + /// serialization, which the caller has already done. The bytes must be + /// in this writer's layout. pub fn append_serialized_tensor( &mut self, name: &str, @@ -635,6 +905,29 @@ impl ContainerWriter { return Err(PtwmCoreError::InvalidContainer("already finalized".into())); } self.check_new_name(name)?; + if self.layout == ContainerLayout::Compact { + if name.is_empty() || name.len() > crate::ptwc::MAX_NAME_LEN { + return Err(PtwmCoreError::InvalidContainer(format!( + "tensor name of {} bytes (the compact layout needs 1..={})", + name.len(), + crate::ptwc::MAX_NAME_LEN + ))); + } + if (tensor_buf.len() as u64) < crate::tensor_record_compact::MIN_RECORD_LEN { + return Err(PtwmCoreError::InvalidContainer(format!( + "tensor '{name}': {}-byte record is not a compact record", + tensor_buf.len() + ))); + } + self.flush_prelude()?; + self.writer + .write_all(tensor_buf) + .map_err(|e| PtwmCoreError::Io(e.to_string()))?; + self.compact_index + .push((name.to_owned(), tensor_buf.len() as u64)); + self.names.insert(name.to_owned()); + return Ok(()); + } if !self.prelude_written { self.flush_prelude()?; self.tensors_start_pos = self.header.shared_prelude_offset + self.prelude_section_len; @@ -650,6 +943,7 @@ impl ContainerWriter { name_hash: name_hash(name), tensor_offset, tensor_len: tensor_buf.len() as u64, + name: None, }); self.names.insert(name.to_owned()); Ok(()) @@ -666,6 +960,17 @@ impl ContainerWriter { if !self.prelude_written { self.flush_prelude()?; } + if self.layout == ContainerLayout::Compact { + let index = crate::ptwc::write_index(&self.compact_index)?; + self.writer + .write_all(&index) + .map_err(|e| PtwmCoreError::Io(e.to_string()))?; + self.writer + .write_all(&crate::ptwc::footer_bytes(index.len() as u64)) + .map_err(|e| PtwmCoreError::Io(e.to_string()))?; + self.finalized = true; + return Ok(()); + } // Record tensor-index offset and write sorted index. let index_offset = self @@ -809,6 +1114,120 @@ fn decode_pubkey_to_canonical( crate::extension::CanonicalId::derive(&key, name, version) } +/// Discovery, missing-flavor and trust pass over an Extension Table (spec +/// §6.5: trusted ⇒ allowed). Returns the installed-extension snapshot and +/// the missing entries. With `skip_missing`, missing or untrusted entries +/// do not fail the call. +#[allow(clippy::type_complexity)] +fn check_extensions( + extension_table: &ExtensionTable, + skip_missing: bool, +) -> Result< + ( + Vec, + Vec, + ), + PtwmCoreError, +> { + // A table of built-ins only needs no discovery, keyring or trust + // pass: built-ins are trusted by construction and always available. + // Skipping it keeps opening such a file free of filesystem access. + let all_builtin = extension_table + .entries + .iter() + .all(|e| crate::extension::is_builtin(&e.canonical_id)); + + // Scan installed extensions once; used by both the missing-flavor check + // and the trust verifier pass below. + let installed = if all_builtin { + Vec::new() + } else { + crate::discovery::scan_all_cached().unwrap_or_default() + }; + + // Check for missing non-builtin flavors immediately after parsing the + // Extension Table — before touching any tensor payload. + let needs = collect_missing_flavors(extension_table, &installed)?; + if !needs.is_empty() && !skip_missing { + return Err(PtwmCoreError::MissingFlavor { needs }); + } + + // ── Trust Verifier pass (spec §6.5: trusted ⇒ allowed) ─────────────── + // + // For every non-builtin entry in the Extension Table that IS installed + // (i.e. not flagged as missing above), verify that its author is + // trusted by the active keyring. Entries with no matching install + // (already covered by MissingFlavor above) are skipped here. + if !all_builtin { + // Build the active keyring: start from the bundled keys, merge the + // user keyring on top. + let mut keyring = match crate::trust::evaluate_bundled() { + Ok(crate::trust::BundleStatus::Match { keyring }) => keyring, + Ok(crate::trust::BundleStatus::FreshInstall { keyring, hash }) => { + // First-run convenience: auto-trust the bundled keyring and + // persist the lock so subsequent opens see Match. + let _ = crate::trust::apply_fresh_install(&hash); + keyring + } + // Bundled keyring changed since last acceptance — do not + // activate it until the user runs `ptwm trust update --bundled`. + Ok(crate::trust::BundleStatus::Mismatch { .. }) => crate::trust::Keyring::default(), + Err(_) => crate::trust::Keyring::default(), + }; + + // Merge user keyring (keys.toml lives beside bundled.lock). + let user_keys_path = crate::trust::lock_path().with_file_name("keys.toml"); + if let Ok(user) = crate::trust::Keyring::load_from_path(&user_keys_path) { + keyring.entries.extend(user.entries); + } + + let verifier = crate::trust::Verifier::new(keyring); + + let mut untrusted: Vec<(crate::extension::CanonicalId, String)> = Vec::new(); + for entry in &extension_table.entries { + if crate::extension::is_builtin(&entry.canonical_id) { + // Built-ins are trusted by construction; verifier would + // short-circuit anyway, but skip the lookup entirely. + continue; + } + let Some((pubkey, contrib_idx)) = + lookup_author_pubkey(&entry.canonical_id, &installed) + else { + // No matching installed manifest — MissingFlavor already + // accounts for this entry. If skip_missing is true we keep + // going; treat as unverifiable (but not actively untrusted). + continue; + }; + // Recompute the signature's signed material from the on-disk + // bundle (manifest + binaries, in the order + // `ext_tooling/sign.py::sign_bundle` writes them). Tampering + // with either the manifest TOML or any binary after signing + // changes the digest and trips Ed25519 verification. + let signed_material = match installed[contrib_idx].signed_digest() { + Ok(d) => d, + Err(e) => { + untrusted.push((entry.canonical_id, e.to_string())); + continue; + } + }; + match verifier.verify(entry, &signed_material, &pubkey) { + Ok(crate::trust::TrustVerdict::Trusted) => {} + Ok(crate::trust::TrustVerdict::Untrusted { reason }) => { + untrusted.push((entry.canonical_id, reason)); + } + Err(e) => { + untrusted.push((entry.canonical_id, e.to_string())); + } + } + } + if !untrusted.is_empty() && !skip_missing { + let (id, reason) = untrusted.remove(0); + return Err(PtwmCoreError::ContributionUntrusted { id, reason }); + } + } + Ok((installed, needs)) +} + /// Zero-copy reader for a `.ptwm` container backed by an in-memory byte slice. #[derive(Debug)] pub struct ContainerReader<'a> { @@ -848,7 +1267,16 @@ pub struct ContainerReader<'a> { /// [`crate::extension::normalize_builtin_ids`]: the release whose /// built-in ids the writer used. `None` when the table holds no built-in /// entry. `extension_table` holds the ids mapped to the crate version. + /// Always `None` for the compact layout, which stores built-ins without + /// a version. pub builtin_version: Option, + /// The file's layout, from its magic. + pub layout: ContainerLayout, + /// Compact layout: the codec id of each Extension Table entry (see + /// [`plane_codec_ids`]). Empty for the legacy layout. + codec_ids: Vec>, + /// Compact layout: whether each registry chain has a bound Source. + registry_bound: Vec, } impl<'a> ContainerReader<'a> { @@ -881,12 +1309,20 @@ impl<'a> ContainerReader<'a> { /// [`open_partial`](Self::open_partial). fn open_with(buf: &'a [u8], skip_missing: bool) -> Result { verify_eof_sentinel(buf)?; + if sniff(buf) == Some(ContainerMagic::Ptwc) { + return Self::open_compact(buf, skip_missing); + } if buf.len() < HEADER_LEN { return Err(PtwmCoreError::InvalidContainer( "file shorter than header".into(), )); } let header = Header::from_bytes(&buf[..HEADER_LEN])?; + if header.is_ptwx() != (header.magic == MAGIC_PTWX) { + return Err(PtwmCoreError::InvalidContainer( + "the PTWX header flag disagrees with the magic".into(), + )); + } // Parse the Extension Table (sits at extension_table_offset for // extension_table_length bytes, immediately after the header). @@ -921,102 +1357,7 @@ impl<'a> ContainerReader<'a> { // This runs after the hash check, which covers the on-disk bytes. let builtin_version = crate::extension::normalize_builtin_ids(&mut extension_table); - // A table of built-ins only needs no discovery, keyring or trust - // pass: built-ins are trusted by construction and always available. - // Skipping it keeps opening such a file free of filesystem access. - let all_builtin = extension_table - .entries - .iter() - .all(|e| crate::extension::is_builtin(&e.canonical_id)); - - // Scan installed extensions once; used by both the missing-flavor check - // and the trust verifier pass below. - let installed = if all_builtin { - Vec::new() - } else { - crate::discovery::scan_all_cached().unwrap_or_default() - }; - - // Check for missing non-builtin flavors immediately after parsing the - // Extension Table — before touching any tensor payload. - let needs = collect_missing_flavors(&extension_table, &installed)?; - if !needs.is_empty() && !skip_missing { - return Err(PtwmCoreError::MissingFlavor { needs }); - } - - // ── Trust Verifier pass (spec §6.5: trusted ⇒ allowed) ─────────────── - // - // For every non-builtin entry in the Extension Table that IS installed - // (i.e. not flagged as missing above), verify that its author is - // trusted by the active keyring. Entries with no matching install - // (already covered by MissingFlavor above) are skipped here. - if !all_builtin { - // Build the active keyring: start from the bundled keys, merge the - // user keyring on top. - let mut keyring = match crate::trust::evaluate_bundled() { - Ok(crate::trust::BundleStatus::Match { keyring }) => keyring, - Ok(crate::trust::BundleStatus::FreshInstall { keyring, hash }) => { - // First-run convenience: auto-trust the bundled keyring and - // persist the lock so subsequent opens see Match. - let _ = crate::trust::apply_fresh_install(&hash); - keyring - } - // Bundled keyring changed since last acceptance — do not - // activate it until the user runs `ptwm trust update --bundled`. - Ok(crate::trust::BundleStatus::Mismatch { .. }) => crate::trust::Keyring::default(), - Err(_) => crate::trust::Keyring::default(), - }; - - // Merge user keyring (keys.toml lives beside bundled.lock). - let user_keys_path = crate::trust::lock_path().with_file_name("keys.toml"); - if let Ok(user) = crate::trust::Keyring::load_from_path(&user_keys_path) { - keyring.entries.extend(user.entries); - } - - let verifier = crate::trust::Verifier::new(keyring); - - let mut untrusted: Vec<(crate::extension::CanonicalId, String)> = Vec::new(); - for entry in &extension_table.entries { - if crate::extension::is_builtin(&entry.canonical_id) { - // Built-ins are trusted by construction; verifier would - // short-circuit anyway, but skip the lookup entirely. - continue; - } - let Some((pubkey, contrib_idx)) = - lookup_author_pubkey(&entry.canonical_id, &installed) - else { - // No matching installed manifest — MissingFlavor already - // accounts for this entry. If skip_missing is true we keep - // going; treat as unverifiable (but not actively untrusted). - continue; - }; - // Recompute the signature's signed material from the on-disk - // bundle (manifest + binaries, in the order - // `ext_tooling/sign.py::sign_bundle` writes them). Tampering - // with either the manifest TOML or any binary after signing - // changes the digest and trips Ed25519 verification. - let signed_material = match installed[contrib_idx].signed_digest() { - Ok(d) => d, - Err(e) => { - untrusted.push((entry.canonical_id, e.to_string())); - continue; - } - }; - match verifier.verify(entry, &signed_material, &pubkey) { - Ok(crate::trust::TrustVerdict::Trusted) => {} - Ok(crate::trust::TrustVerdict::Untrusted { reason }) => { - untrusted.push((entry.canonical_id, reason)); - } - Err(e) => { - untrusted.push((entry.canonical_id, e.to_string())); - } - } - } - if !untrusted.is_empty() && !skip_missing { - let (id, reason) = untrusted.remove(0); - return Err(PtwmCoreError::ContributionUntrusted { id, reason }); - } - } + let (installed, needs) = check_extensions(&extension_table, skip_missing)?; let prelude_off = usize::try_from(header.shared_prelude_offset).map_err(|_| { PtwmCoreError::InvalidContainer("shared_prelude_offset overflows usize".into()) @@ -1091,9 +1432,149 @@ impl<'a> ContainerReader<'a> { chain_registry_offset, missing_extensions: needs, builtin_version, + layout: ContainerLayout::Legacy, + codec_ids: Vec::new(), + registry_bound: Vec::new(), + }) + } + + /// [`open_with`](Self::open_with) for the compact layout. + fn open_compact(buf: &'a [u8], skip_missing: bool) -> Result { + use crate::extension::compact_table::read_compact_table; + + let sec = crate::ptwc::locate_sections(buf)?; + let extension_table = read_compact_table(&buf[sec.extension_table.clone()])?; + let (installed, needs) = check_extensions(&extension_table, skip_missing)?; + + let prelude_bytes = &buf[sec.prelude.clone()]; + let (prelude, used) = parse_prelude(prelude_bytes)?; + if used != prelude_bytes.len() { + return Err(PtwmCoreError::InvalidContainer(format!( + "compact container: {} trailing bytes after the prelude", + prelude_bytes.len() - used + ))); + } + // Reserved fields and padding must be zero: one encoding per prelude. + let mut rewritten = Vec::with_capacity(prelude_bytes.len()); + write_prelude(&prelude, &mut rewritten)?; + if rewritten != prelude_bytes { + return Err(PtwmCoreError::InvalidContainer( + "compact container: prelude is not in canonical form".into(), + )); + } + + let chain_registry = + crate::ptwc::read_registry(&buf[sec.registry.clone()], &extension_table)?; + let index = crate::ptwc::read_index(&buf[sec.index.clone()], sec.body.clone())?; + let registry_bound = chain_registry + .chains + .iter() + .map(|(_, c)| has_deferred_source(c)) + .collect(); + let codec_ids = plane_codec_ids(&extension_table); + let router = crate::flavor::PlaneCodecRouter::new(installed); + + // The compact header has no fixed fields; this view gives code that + // reads offsets from `header` the section positions. + let header = Header { + magic: MAGIC_PTWC, + flags: 0, + extension_table_offset: sec.extension_table.start as u64, + extension_table_length: sec.extension_table.len() as u64, + extension_table_hash: [0u8; 32], + tensor_index_offset: sec.index.start as u64, + tensor_index_length: sec.index.len() as u64, + shared_prelude_offset: sec.prelude.start as u64, + shared_prelude_length: (sec.registry.end - sec.prelude.start) as u64, + }; + Ok(Self { + buf, + header, + extension_table, + prelude, + router, + num_tensors: index.len() as u64, + index, + tensors_offset: sec.body.start as u64, + chain_registry_offset: if chain_registry.chains.is_empty() { + 0 + } else { + sec.registry.start as u64 + }, + chain_registry, + missing_extensions: needs, + builtin_version: None, + layout: ContainerLayout::Compact, + codec_ids, + registry_bound, }) } + /// The context compact records are parsed against. + fn record_context(&self) -> RecordContext<'_> { + RecordContext { + codec_ids: &self.codec_ids, + registry_bound: &self.registry_bound, + } + } + + /// Parse serialized record bytes of this container's layout. A compact + /// record does not hold its name, so the returned `name` is empty; use + /// [`tensor_record`](Self::tensor_record) for a record with its name. + pub fn parse_record(&self, bytes: &[u8]) -> Result { + match self.layout { + ContainerLayout::Legacy => { + crate::tensor_record::parse_tensor_record(bytes).map(|(rec, _)| rec) + } + ContainerLayout::Compact => parse_compact_record(bytes, &self.record_context()), + } + } + + /// The record bytes an index entry points at. + fn entry_bytes(&self, e: &IndexEntry) -> Result<&'a [u8], PtwmCoreError> { + let buf: &'a [u8] = self.buf; + usize::try_from(e.tensor_offset) + .ok() + .zip(usize::try_from(e.tensor_len).ok()) + .and_then(|(start, len)| buf.get(start..start.checked_add(len)?)) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "tensor index entry ({}+{}) exceeds buffer ({} bytes)", + e.tensor_offset, + e.tensor_len, + buf.len() + )) + }) + } + + /// Parse the record of an index entry, with its name. + fn entry_record(&self, e: &IndexEntry) -> Result { + let mut rec = self.parse_record(self.entry_bytes(e)?)?; + if let Some(name) = &e.name { + rec.name = name.to_string(); + } + Ok(rec) + } + + /// Look up a tensor by name and parse its record. + pub fn tensor_record(&self, name: &str) -> Result { + let bytes = self.get_tensor_record_bytes(name).ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!("tensor '{}' not found", name)) + })?; + let mut rec = self.parse_record(bytes)?; + if self.layout == ContainerLayout::Compact { + rec.name = name.to_owned(); + } + Ok(rec) + } + + /// Index entries in body (file) order. + fn entries_in_body_order(&self) -> Vec<&IndexEntry> { + let mut entries: Vec<&IndexEntry> = self.index.iter().collect(); + entries.sort_by_key(|e| e.tensor_offset); + entries + } + /// Look up a tensor by name and return a slice covering its serialized /// `TensorRecord` bytes. Returns `None` if no tensor with that name /// exists. @@ -1115,11 +1596,12 @@ impl<'a> ContainerReader<'a> { let lo = self.index.partition_point(|e| e.name_hash < target); let hi = lo + self.index[lo..].partition_point(|e| e.name_hash == target); self.index[lo..hi].iter().find_map(|e| { - let start = usize::try_from(e.tensor_offset).ok()?; - let len = usize::try_from(e.tensor_len).ok()?; - let end = start.checked_add(len)?; - let bytes = self.buf.get(start..end)?; - (peek_tensor_record_name(bytes).ok()? == name).then_some(bytes) + let bytes = self.entry_bytes(e).ok()?; + let found = match &e.name { + Some(n) => &**n, + None => peek_tensor_record_name(bytes).ok()?, + }; + (found == name).then_some(bytes) }) } @@ -1128,24 +1610,29 @@ impl<'a> ContainerReader<'a> { /// malformed or the metadata isn't the expected shape-convention map. pub fn tensor_shape(&self, name: &str) -> Result, String)>, PtwmCoreError> { use crate::metadata::decode_shape_metadata; - use crate::tensor_record::parse_tensor_record; - let record_bytes = self.get_tensor_record_bytes(name).ok_or_else(|| { - PtwmCoreError::InvalidContainer(format!("tensor '{}' not found", name)) - })?; - let (rec, _) = parse_tensor_record(record_bytes)?; + let rec = self.tensor_record(name)?; match rec.tensor_metadata { Some(bytes) => Ok(Some(decode_shape_metadata(&bytes)?)), None => Ok(None), } } - /// Return every tensor name in index order (i.e. sorted by name hash, - /// which is how the on-disk index is stored). Requires parsing each - /// tensor record header; use sparingly on very large containers. + /// Return every tensor name. For the legacy layout the names come in + /// index order (sorted by name hash, which is how the on-disk index is + /// stored) and each record header is parsed; use sparingly on very large + /// containers. For the compact layout they come from the index alone, in + /// body (write) order. pub fn tensor_names(&self) -> Result, PtwmCoreError> { use crate::tensor_record::parse_tensor_record; + if self.layout == ContainerLayout::Compact { + return Ok(self + .entries_in_body_order() + .into_iter() + .map(|e| e.name.as_deref().unwrap_or_default().to_owned()) + .collect()); + } let mut names = Vec::with_capacity(self.index.len()); for e in &self.index { let start = e.tensor_offset as usize; @@ -1173,25 +1660,12 @@ impl<'a> ContainerReader<'a> { if self.index.is_empty() { return Ok(Vec::new()); } - use crate::tensor_record::parse_tensor_record; // Parse each record header and decode it in the same parallel pass — // no sequential tensor_names() pre-pass. self.index .par_iter() .map(|e| { - let start = e.tensor_offset as usize; - let end = start.checked_add(e.tensor_len as usize).ok_or_else(|| { - PtwmCoreError::InvalidContainer( - "tensor_offset + tensor_len overflows usize".into(), - ) - })?; - if end > self.buf.len() { - return Err(PtwmCoreError::InvalidContainer(format!( - "tensor index entry ({start}..{end}) exceeds buffer ({} bytes)", - self.buf.len() - ))); - } - let (rec, _) = parse_tensor_record(&self.buf[start..end])?; + let rec = self.entry_record(e)?; let bytes = self.decode_tensor_record(&rec)?; Ok((rec.name, bytes)) }) @@ -1203,11 +1677,7 @@ impl<'a> ContainerReader<'a> { /// terminal planes to reconstruct the source bytes. Validates plane CRCs /// (when present) and the optional tensor-level `payload_hash`. pub fn decode_tensor(&self, name: &str) -> Result, PtwmCoreError> { - use crate::tensor_record::parse_tensor_record; - let record_bytes = self.get_tensor_record_bytes(name).ok_or_else(|| { - PtwmCoreError::InvalidContainer(format!("tensor '{}' not found", name)) - })?; - let (rec, _consumed) = parse_tensor_record(record_bytes)?; + let rec = self.tensor_record(name)?; self.decode_tensor_record(&rec) } @@ -1228,7 +1698,13 @@ impl<'a> ContainerReader<'a> { use crate::tensor_record::{CHAIN_REF_INLINE_BIT, TENSOR_FLAG_PAYLOAD_HASH}; use crate::transforms::op::Plane; + // The compact layout requires exact sizes throughout: the Source + // equals the raw size and each plane decodes to its descriptor + // length. The legacy layout tolerates older writers' slack. + let strict = self.layout == ContainerLayout::Compact; + // Resolve the chain via the inline bit or the chain registry. + let bound_chain: Chain; let chain: &Chain = if rec.chain_ref & CHAIN_REF_INLINE_BIT != 0 { rec.inline_chain.as_ref().ok_or_else(|| { PtwmCoreError::InvalidContainer( @@ -1237,11 +1713,27 @@ impl<'a> ContainerReader<'a> { })? } else { let chain_id = rec.chain_ref & !CHAIN_REF_INLINE_BIT; - self.chain_registry.lookup(chain_id).ok_or_else(|| { + let shared = self.chain_registry.lookup(chain_id).ok_or_else(|| { PtwmCoreError::InvalidContainer(format!( "decode_tensor: chain_id {chain_id} not in registry" )) - })? + })?; + if strict && has_deferred_source(shared) { + // A bound Source takes its params from the record's shape. + let params = record_binding(rec.tensor_metadata.as_deref(), rec.dtype_code) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "decode_tensor: chain {chain_id} has a bound Source but the \ + record has no shape that binds it" + )) + })?; + let mut c = shared.clone(); + bind_deferred_source(&mut c, ¶ms); + bound_chain = c; + &bound_chain + } else { + shared + } }; // Bound the chain by the record's `orig_size` before any plane is @@ -1254,6 +1746,12 @@ impl<'a> ContainerReader<'a> { // small tensor. `chain_source_descriptor` checks node 0 itself and // needs no prior validation. let source_desc = chain_source_descriptor(chain)?; + if strict && source_desc.length_bytes != rec.orig_size { + return Err(PtwmCoreError::InvalidContainer(format!( + "decode_tensor: chain Source declares {} bytes for a {}-byte tensor", + source_desc.length_bytes, rec.orig_size + ))); + } let max_source = rec.orig_size.saturating_mul(8); if source_desc.length_bytes > max_source { return Err(PtwmCoreError::InvalidContainer(format!( @@ -1378,10 +1876,13 @@ impl<'a> ContainerReader<'a> { .iter() .try_fold(0u64, |acc, e| acc.checked_add(u64::from(e.decoded_size))) .filter(|&t| t <= plane_bound) + .filter(|&t| !strict || t == descriptor.length_bytes) .ok_or_else(|| { PtwmCoreError::InvalidContainer(format!( - "chunk table decodes to more than {plane_bound} bytes for a \ + "chunk table decodes to more than {plane_bound} bytes (or, in \ + a compact container, not to the plane's {} bytes) for a \ {}-byte tensor", + descriptor.length_bytes, rec.orig_size )) })? as usize; @@ -1512,14 +2013,26 @@ impl<'a> ContainerReader<'a> { // the length check after `inverse_chain` and the payload hash // cover the rest. The codec stops at the limit, so a corrupt // payload cannot allocate past it. - codec.decode_bounded( + let limit = if strict { + decoded_len + } else { + usize::try_from(plane_bound).unwrap_or(usize::MAX) + }; + let plane = codec.decode_bounded( state_version, &state_bytes, &plane_rec.payload_bytes, &plane_rec.layout, decoded_len, - usize::try_from(plane_bound).unwrap_or(usize::MAX), - )? + limit, + )?; + if strict && plane.len() != decoded_len { + return Err(PtwmCoreError::InvalidContainer(format!( + "decoded plane is {} bytes, its descriptor says {decoded_len}", + plane.len() + ))); + } + plane }; // The codecs enforce the bound; this check keeps it for any // codec that does not. diff --git a/crates/ptwm-core/src/dtype.rs b/crates/ptwm-core/src/dtype.rs index 9a112c5..adf96c6 100644 --- a/crates/ptwm-core/src/dtype.rs +++ b/crates/ptwm-core/src/dtype.rs @@ -57,6 +57,16 @@ pub struct PreprocessingModes { pub byte_reorder: i32, } +/// Canonical dtype name of a public `dtype_code`, or `None` for a code +/// outside the [`Dtype`] set. Compact formats store a tensor's dtype name +/// only when it differs from this name. +pub fn canonical_dtype_name(dtype_code: u16) -> Option<&'static str> { + u8::try_from(dtype_code) + .ok() + .and_then(|c| Dtype::from_code(c).ok()) + .map(Dtype::canonical_name) +} + impl Dtype { /// Wire `dtype_code` for this dtype (matches Python /// `DType.code`). diff --git a/crates/ptwm-core/src/extension/builtins.rs b/crates/ptwm-core/src/extension/builtins.rs index 2629307..e071006 100644 --- a/crates/ptwm-core/src/extension/builtins.rs +++ b/crates/ptwm-core/src/extension/builtins.rs @@ -144,6 +144,19 @@ impl BuiltinRegistry { self.by_id.get(id) } + /// The Extension Table entry at position `index` of + /// [`builtin_entries`], without cloning the list. + pub fn entry_at(&self, index: usize) -> Option<&ExtensionTableEntry> { + self.entries.get(index) + } + + /// The Extension Table entry of the built-in named `name` at the crate + /// version, if it has one. + pub fn entry_by_name(&self, name: &str) -> Option<&ExtensionTableEntry> { + let id = self.by_name.get(name)?; + self.entry_at(self.by_id.get(id)?.entry_index?) + } + /// Current canonical id of the built-in named `name` that has an /// Extension Table entry. pub fn id_for_entry_name(&self, name: &str) -> Option { diff --git a/crates/ptwm-core/src/extension/compact_table.rs b/crates/ptwm-core/src/extension/compact_table.rs new file mode 100644 index 0000000..7fc210b --- /dev/null +++ b/crates/ptwm-core/src/extension/compact_table.rs @@ -0,0 +1,409 @@ +//! Compact Extension Table, used by the compact container (`\x89PTWC`). +//! +//! The legacy table stores every entry in full: id, label, kind, ABI, +//! capabilities and attestation, about 125 bytes per built-in. A built-in +//! entry is fully determined by its name and the crate version, so the +//! compact table stores a built-in as a tag and a wire id or name, and +//! stores third-party entries in full: +//! +//! ```text +//! count varint at most 65536 and at most (table bytes) / 2 +//! entry × count: +//! tag u8 +//! 0x01 OpId wire id (varint) built-in transform +//! 0x02 CodecId wire id (varint) built-in plane codec +//! 0x03 name_len varint (1..=255), then the name +//! built-in without a wire id +//! 0x7F len varint, then one legacy table entry of exactly len +//! bytes third-party entry +//! ``` +//! +//! The reader expands tags 0x01-0x03 to the built-in entry of the reader's +//! own version ([`builtin_registry`]), so dispatch, trust and the +//! missing-flavor check see the same entries as for a legacy file, and a +//! compact file does not depend on the version that wrote it. +//! +//! Each entry has one encoding: +//! - a built-in with an op or codec wire id must use tag 0x01 or 0x02; tag +//! 0x03 is only for built-ins without one; +//! - a codec id without a built-in entry (`huffman_nibble`) is rejected; +//! - a 0x7F entry must not carry a built-in id, of this version or (by its +//! label) of another one, must re-serialize to the same bytes, and must +//! not reference embedded WASM; +//! - unknown tags, ids and names are rejected, and so are duplicate ids. +//! +//! The writer refuses entries it cannot store this way: a built-in id whose +//! other fields differ from the built-in entry, or a built-in of another +//! release. + +use std::collections::HashSet; + +use super::{ + BuiltinKind, ExtensionTable, ExtensionTableEntry, builtin_registry, normalize_builtin_ids, +}; +use crate::codec::CodecId; +use crate::error::PtwmCoreError; +use crate::transforms::op::OpId; +use crate::varint::{Cursor, write_varint}; + +/// Built-in transform, stored by `OpId` wire id. +pub const TAG_OP: u8 = 0x01; +/// Built-in plane codec, stored by `CodecId` wire id. +pub const TAG_CODEC: u8 = 0x02; +/// Built-in without a wire id, stored by name. +pub const TAG_NAMED: u8 = 0x03; +/// Third-party entry, stored in the legacy entry layout. +pub const TAG_FULL: u8 = 0x7F; + +/// Largest entry count: table indices are `u16`. +pub const MAX_COMPACT_ENTRIES: u64 = 1 << 16; +/// Longest built-in name stored under [`TAG_NAMED`]. +const MAX_NAME_LEN: u64 = 255; + +/// How [`write_compact_table`] stores one entry. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CompactEntry<'a> { + Op(OpId), + Codec(CodecId), + Named(&'static str), + Full(&'a ExtensionTableEntry), +} + +fn err(msg: impl std::fmt::Display) -> PtwmCoreError { + PtwmCoreError::InvalidContainer(format!("compact extension table: {msg}")) +} + +/// Whether `entry` is a built-in of another release: its label names a +/// built-in and its id re-derives from that label (see +/// [`normalize_builtin_ids`]). +fn is_other_version_builtin(entry: &ExtensionTableEntry) -> bool { + let mut one = ExtensionTable { + entries: vec![entry.clone()], + }; + normalize_builtin_ids(&mut one); + one.entries[0].canonical_id != entry.canonical_id +} + +/// Decide how `entry` is stored, or fail when the compact table cannot +/// store it. +pub fn classify(entry: &ExtensionTableEntry) -> Result, PtwmCoreError> { + let reg = builtin_registry(); + if let Some(r) = reg.get(&entry.canonical_id) { + let canonical = r.entry_index.and_then(|i| reg.entry_at(i)); + if canonical != Some(entry) { + return Err(err(format_args!( + "entry with the id of built-in {:?} differs from that built-in's entry", + r.name + ))); + } + return Ok(match r.dispatch { + Some(BuiltinKind::Op(op)) => CompactEntry::Op(op), + Some(BuiltinKind::Codec(c)) => CompactEntry::Codec(c), + None => CompactEntry::Named(r.name), + }); + } + if is_other_version_builtin(entry) { + return Err(err(format_args!( + "{:?} is a built-in of another release; the compact table stores \ + built-ins of the reader's version only", + entry.human_label + ))); + } + if entry.embedded_wasm_offset.is_some() || entry.embedded_wasm_length.is_some() { + return Err(err(format_args!( + "{:?} references embedded WASM, which needs the legacy layout", + entry.human_label + ))); + } + Ok(CompactEntry::Full(entry)) +} + +/// Append the compact encoding of `table` to `out`. On error `out` is left +/// unchanged. +pub fn write_compact_table( + table: &ExtensionTable, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + let n = table.entries.len() as u64; + if n > MAX_COMPACT_ENTRIES { + return Err(err(format_args!( + "{n} entries (max {MAX_COMPACT_ENTRIES})" + ))); + } + let mut buf = Vec::with_capacity(1 + 3 * table.entries.len()); + write_varint(&mut buf, n); + let mut seen = HashSet::with_capacity(table.entries.len()); + for entry in &table.entries { + if !seen.insert(entry.canonical_id) { + return Err(err(format_args!( + "duplicate entry {:?}", + entry.human_label + ))); + } + match classify(entry)? { + CompactEntry::Op(op) => { + buf.push(TAG_OP); + write_varint(&mut buf, u64::from(op.as_u16())); + } + CompactEntry::Codec(c) => { + buf.push(TAG_CODEC); + write_varint(&mut buf, u64::from(c.as_u16())); + } + CompactEntry::Named(name) => { + buf.push(TAG_NAMED); + write_varint(&mut buf, name.len() as u64); + buf.extend_from_slice(name.as_bytes()); + } + CompactEntry::Full(e) => { + let bytes = e.to_bytes().map_err(err)?; + buf.push(TAG_FULL); + write_varint(&mut buf, bytes.len() as u64); + buf.extend_from_slice(&bytes); + } + } + } + out.extend_from_slice(&buf); + Ok(()) +} + +/// Serialize `table` in the compact layout. +pub fn compact_table_bytes(table: &ExtensionTable) -> Result, PtwmCoreError> { + let mut out = Vec::new(); + write_compact_table(table, &mut out)?; + Ok(out) +} + +/// The built-in entry for a wire id or name, or an error naming `what`. +fn builtin_by_name(name: &str, what: impl std::fmt::Display) -> Result { + builtin_registry() + .entry_by_name(name) + .cloned() + .ok_or_else(|| err(format_args!("{what} has no built-in entry"))) +} + +/// Parse one entry at the cursor. +fn read_entry(r: &mut Cursor<'_>) -> Result { + let tag = r.u8()?; + match tag { + TAG_OP => { + let id = r.varint_max(u64::from(u16::MAX), "op id")? as u16; + let op = OpId::from_u16(id).map_err(|_| r.err(format_args!("unknown op id {id:#06x}")))?; + builtin_by_name(op.builtin_name(), format_args!("op {op:?}")) + } + TAG_CODEC => { + let id = r.varint_max(u64::from(u16::MAX), "codec id")? as u16; + let name = crate::codec::REGISTRY + .iter() + .find(|&&(_, raw)| raw == id) + .map(|&(name, _)| name) + .ok_or_else(|| r.err(format_args!("unknown codec id {id:#06x}")))?; + builtin_by_name(name, format_args!("codec {name:?}")) + } + TAG_NAMED => { + let len = r.varint_max(MAX_NAME_LEN, "built-in name length")?; + if len == 0 { + return Err(r.err("empty built-in name")); + } + let bytes = r.take(len as usize)?; + let name = std::str::from_utf8(bytes).map_err(|_| r.err("built-in name is not UTF-8"))?; + let reg = builtin_registry(); + let id = reg + .id_for_entry_name(name) + .ok_or_else(|| r.err(format_args!("unknown built-in name {name:?}")))?; + if reg.get(&id).is_some_and(|b| b.dispatch.is_some()) { + return Err(r.err(format_args!( + "built-in {name:?} has a wire id and must use its tag (non-canonical)" + ))); + } + builtin_by_name(name, format_args!("{name:?}")) + } + TAG_FULL => { + let bytes = r.len_bytes()?; + let entry = ExtensionTableEntry::from_bytes_exact(bytes) + .map_err(|e| r.err(format_args!("third-party entry: {e}")))?; + if entry.to_bytes().map_err(err)? != bytes { + return Err(r.err("third-party entry is not in canonical form")); + } + if builtin_registry().get(&entry.canonical_id).is_some() + || is_other_version_builtin(&entry) + { + return Err(r.err(format_args!( + "full entry {:?} carries a built-in id (non-canonical)", + entry.human_label + ))); + } + if entry.embedded_wasm_offset.is_some() || entry.embedded_wasm_length.is_some() { + return Err(r.err("embedded WASM entries are not allowed in a compact table")); + } + Ok(entry) + } + other => Err(r.err(format_args!("unknown entry tag {other:#04x}"))), + } +} + +/// Parse a compact table that fills `buf` exactly. +pub fn read_compact_table(buf: &[u8]) -> Result { + let mut r = Cursor::new(buf, "compact extension table"); + let count = r.varint_max(MAX_COMPACT_ENTRIES, "entry count")?; + // Every entry takes at least a tag and one more byte. + if count > (buf.len() / 2) as u64 { + return Err(r.err(format_args!("entry count {count} exceeds the table bytes"))); + } + let count = count as usize; + let mut entries = Vec::with_capacity(count); + let mut seen = HashSet::with_capacity(count); + for _ in 0..count { + let entry = read_entry(&mut r)?; + if !seen.insert(entry.canonical_id) { + return Err(r.err(format_args!("duplicate entry {:?}", entry.human_label))); + } + entries.push(entry); + } + r.finish("the last entry")?; + Ok(ExtensionTable { entries }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::extension::{ + Attestation, CanonicalId, CapabilityMap, Kind, Lifecycle, builtin_entries, + builtin_entry_at_version, + }; + + fn third_party() -> ExtensionTableEntry { + ExtensionTableEntry { + canonical_id: CanonicalId::derive(&[7; 32], "acme_codec", "0.1.0"), + human_label: "com.acme.codec@0.1.0".into(), + kind: Kind::PlaneCodec, + abi_version: 1, + lifecycle: Lifecycle::Thread, + flavor_hints: crate::extension::table::FLAVOR_NATIVE, + capabilities: CapabilityMap::new(), + attestation: Attestation::PgpSignature(vec![1, 2, 3]), + install_hint: Some("pip install acme".into()), + embedded_wasm_offset: None, + embedded_wasm_length: None, + } + } + + #[test] + fn every_builtin_round_trips_in_a_few_bytes() { + let table = ExtensionTable { + entries: builtin_entries(), + }; + let bytes = compact_table_bytes(&table).unwrap(); + // Built-ins cost a tag and a short id or name. + assert!( + bytes.len() <= 4 * table.entries.len() + 30 * 5, + "{} bytes for {} entries", + bytes.len(), + table.entries.len() + ); + let back = read_compact_table(&bytes).unwrap(); + assert_eq!(back, table); + } + + #[test] + fn ops_and_codecs_cost_at_most_three_bytes() { + let table = ExtensionTable { + entries: builtin_entries() + .into_iter() + .filter(|e| matches!(e.kind, Kind::PlaneCodec | Kind::Transform)) + .collect(), + }; + let bytes = compact_table_bytes(&table).unwrap(); + assert!(bytes.len() <= 1 + 3 * table.entries.len()); + assert_eq!(read_compact_table(&bytes).unwrap(), table); + } + + #[test] + fn third_party_entry_round_trips() { + let mut entries = vec![builtin_registry().entry_by_name("zstd").unwrap().clone()]; + entries.push(third_party()); + let table = ExtensionTable { entries }; + let bytes = compact_table_bytes(&table).unwrap(); + assert_eq!(bytes[1], TAG_CODEC); + assert_eq!(read_compact_table(&bytes).unwrap(), table); + } + + #[test] + fn writer_refuses_entries_it_cannot_store() { + // A built-in id with other fields changed. + let mut e = builtin_registry().entry_by_name("zstd").unwrap().clone(); + e.flavor_hints = 0; + let t = ExtensionTable { entries: vec![e] }; + assert!(compact_table_bytes(&t).is_err()); + // A built-in of another release. + let t = ExtensionTable { + entries: vec![builtin_entry_at_version("zstd", "1.0.0").unwrap()], + }; + let msg = compact_table_bytes(&t).unwrap_err().to_string(); + assert!(msg.contains("another release"), "{msg}"); + // Embedded WASM. + let mut e = third_party(); + e.embedded_wasm_offset = Some(10); + e.embedded_wasm_length = Some(10); + assert!(compact_table_bytes(&ExtensionTable { entries: vec![e] }).is_err()); + // Duplicates. + let z = builtin_registry().entry_by_name("zstd").unwrap().clone(); + let t = ExtensionTable { + entries: vec![z.clone(), z], + }; + assert!(compact_table_bytes(&t).is_err()); + } + + fn rejects(bytes: &[u8], needle: &str) { + let msg = read_compact_table(bytes).unwrap_err().to_string(); + assert!(msg.contains(needle), "{needle:?} not in {msg:?}"); + } + + #[test] + fn reader_rejects_non_canonical_and_unknown_entries() { + let zstd = CodecId::Zstd.as_u16() as u8; + // Overlong varint for the count and for an id. + rejects(&[0x81, 0x00, TAG_CODEC, zstd], "non-minimal"); + rejects(&[1, TAG_CODEC, 0x80 | zstd, 0x00], "non-minimal"); + // Unknown tag, op id, codec id and name. + rejects(&[1, 0x04, 0], "unknown entry tag"); + rejects(&[1, TAG_OP, 0x7F], "unknown op id"); + rejects(&[1, TAG_CODEC, 0x7F], "unknown codec id"); + rejects(&[1, TAG_NAMED, 3, b'f', b'o', b'o'], "unknown built-in name"); + // huffman_nibble has a codec id but no built-in entry. + rejects(&[1, TAG_CODEC, CodecId::HuffmanNibble.as_u16() as u8], "no built-in entry"); + // A named entry for a built-in that has a wire id. + rejects(&[1, TAG_NAMED, 4, b'z', b's', b't', b'd'], "non-canonical"); + // Duplicates. + rejects(&[2, TAG_CODEC, zstd, TAG_CODEC, zstd], "duplicate"); + // Trailing bytes and a count that does not fit. + rejects(&[1, TAG_CODEC, zstd, 0], "trailing"); + rejects(&[3, TAG_CODEC, zstd], "exceeds the table bytes"); + rejects(&[0x81, 0x80, 0x04], "exceeds 65536"); + } + + fn full_entry_bytes(e: &ExtensionTableEntry) -> Vec { + let body = e.to_bytes().unwrap(); + let mut out = vec![1, TAG_FULL]; + write_varint(&mut out, body.len() as u64); + out.extend_from_slice(&body); + out + } + + #[test] + fn reader_rejects_builtins_stored_in_full() { + let z = builtin_registry().entry_by_name("zstd").unwrap().clone(); + rejects(&full_entry_bytes(&z), "built-in id"); + let old = builtin_entry_at_version("zstd", "1.0.0").unwrap(); + rejects(&full_entry_bytes(&old), "built-in id"); + let mut wasm = third_party(); + wasm.embedded_wasm_offset = Some(1); + wasm.embedded_wasm_length = Some(1); + rejects(&full_entry_bytes(&wasm), "embedded WASM"); + // A full entry with bytes after it inside its length. + let mut bytes = full_entry_bytes(&third_party()); + let body_len = bytes.len() - 3; + bytes[2] = (body_len + 1) as u8; + bytes.push(0); + rejects(&bytes, "trailing"); + } +} diff --git a/crates/ptwm-core/src/extension/mod.rs b/crates/ptwm-core/src/extension/mod.rs index ff2065b..f7f30d1 100644 --- a/crates/ptwm-core/src/extension/mod.rs +++ b/crates/ptwm-core/src/extension/mod.rs @@ -5,6 +5,7 @@ pub mod builder; pub mod builtins; pub mod capability; +pub mod compact_table; pub mod dispatch; pub mod error; pub mod id; diff --git a/crates/ptwm-core/src/extension/table.rs b/crates/ptwm-core/src/extension/table.rs index 40fc7e3..8c22e87 100644 --- a/crates/ptwm-core/src/extension/table.rs +++ b/crates/ptwm-core/src/extension/table.rs @@ -136,6 +136,28 @@ impl ExtensionTable { } impl ExtensionTableEntry { + /// This entry in the legacy table layout (the bytes one entry takes in + /// [`ExtensionTable::to_bytes`]). + pub fn to_bytes(&self) -> Result, ExtensionError> { + let mut buf = Vec::new(); + self.write(&mut buf)?; + Ok(buf) + } + + /// Parse one entry in the legacy table layout that fills `bytes` + /// exactly. + pub fn from_bytes_exact(bytes: &[u8]) -> Result { + let mut c = Cursor::new(bytes); + let entry = Self::read(&mut c)?; + if c.position() as usize != bytes.len() { + return Err(ExtensionError::TableParse(format!( + "{} trailing bytes after the entry", + bytes.len() - c.position() as usize + ))); + } + Ok(entry) + } + fn write(&self, w: &mut W) -> Result<(), ExtensionError> { w.write_all(self.canonical_id.as_bytes()) .map_err(io_err("write canonical_id"))?; diff --git a/crates/ptwm-core/src/header.rs b/crates/ptwm-core/src/header.rs index 82ff0b1..652f100 100644 --- a/crates/ptwm-core/src/header.rs +++ b/crates/ptwm-core/src/header.rs @@ -23,8 +23,35 @@ use crate::error::PtwmCoreError; pub const MAGIC_PTWM: [u8; 9] = [0x89, b'P', b'T', b'W', b'M', b'\r', b'\n', 0x1A, b'\n']; pub const MAGIC_PTWX: [u8; 9] = [0x89, b'P', b'T', b'W', b'X', b'\r', b'\n', 0x1A, b'\n']; +/// Magic of the compact container (`ContainerLayout::Compact`). Readers that +/// predate it fail [`Header::read`] with "bad magic". +pub const MAGIC_PTWC: [u8; 9] = [0x89, b'P', b'T', b'W', b'C', b'\r', b'\n', 0x1A, b'\n']; + pub const FLAG_PTWX: u16 = 0x0001; +/// Container kind named by the first 9 bytes of a file. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum ContainerMagic { + /// `\x89PTWM`: legacy container. + Ptwm, + /// `\x89PTWX`: legacy container with embedded WASM. + Ptwx, + /// `\x89PTWC`: compact container. + Ptwc, +} + +/// Identify the container magic at the start of `buf`, or `None` when `buf` +/// does not start with one of the three 9-byte container magics. +pub fn sniff(buf: &[u8]) -> Option { + let magic: [u8; 9] = buf.get(..9)?.try_into().ok()?; + match magic { + MAGIC_PTWM => Some(ContainerMagic::Ptwm), + MAGIC_PTWX => Some(ContainerMagic::Ptwx), + MAGIC_PTWC => Some(ContainerMagic::Ptwc), + _ => None, + } +} + pub const HEADER_LEN: usize = 91; #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -178,6 +205,25 @@ mod tests { assert!(res.is_err()); } + #[test] + fn ptwc_magic_is_bad_magic_for_the_legacy_header() { + // A reader that only knows the legacy header must fail on the magic, + // not on a later field. + let mut bytes = Header::new_ptwm().to_bytes().unwrap(); + bytes[..9].copy_from_slice(&MAGIC_PTWC); + let err = Header::from_bytes(&bytes).unwrap_err().to_string(); + assert!(err.contains("bad magic"), "{err}"); + } + + #[test] + fn sniff_names_each_magic() { + assert_eq!(sniff(&MAGIC_PTWM), Some(ContainerMagic::Ptwm)); + assert_eq!(sniff(&MAGIC_PTWX), Some(ContainerMagic::Ptwx)); + assert_eq!(sniff(&MAGIC_PTWC), Some(ContainerMagic::Ptwc)); + assert_eq!(sniff(&MAGIC_PTWC[..8]), None); + assert_eq!(sniff(b"\x89PTWS\r\n\x1a\n"), None); + } + #[test] fn rejects_truncated() { let bytes = Header::new_ptwm().to_bytes().unwrap(); diff --git a/crates/ptwm-core/src/index.rs b/crates/ptwm-core/src/index.rs index 68beed1..664636a 100644 --- a/crates/ptwm-core/src/index.rs +++ b/crates/ptwm-core/src/index.rs @@ -10,6 +10,10 @@ pub struct IndexEntry { pub name_hash: u64, pub tensor_offset: u64, pub tensor_len: u64, + /// Tensor name. The compact container stores names in its index, so + /// lookups compare them directly; the legacy index stores only the hash + /// and leaves this `None` (the name is in the record). + pub name: Option>, } pub fn write_index(entries: &mut [IndexEntry], out: &mut Vec) { @@ -53,6 +57,7 @@ pub fn parse_index(buf: &[u8]) -> Result, PtwmCoreError> { name_hash: u64::from_le_bytes(buf[pos..pos + 8].try_into().unwrap()), tensor_offset: u64::from_le_bytes(buf[pos + 8..pos + 16].try_into().unwrap()), tensor_len: u64::from_le_bytes(buf[pos + 16..pos + 24].try_into().unwrap()), + name: None, }); pos += 24; } @@ -86,16 +91,19 @@ mod tests { name_hash: 5, tensor_offset: 100, tensor_len: 10, + name: None, }, IndexEntry { name_hash: 1, tensor_offset: 0, tensor_len: 20, + name: None, }, IndexEntry { name_hash: 3, tensor_offset: 50, tensor_len: 30, + name: None, }, ]; let mut out = Vec::new(); diff --git a/crates/ptwm-core/src/lib.rs b/crates/ptwm-core/src/lib.rs index b13fe6d..cb2661e 100644 --- a/crates/ptwm-core/src/lib.rs +++ b/crates/ptwm-core/src/lib.rs @@ -39,11 +39,13 @@ pub mod metadata; pub mod plane_record; pub mod policy; pub mod prelude; +pub mod ptwc; pub mod quantize; pub mod range_coder; pub mod select; pub mod split; pub mod tensor_record; +pub mod tensor_record_compact; pub mod transcode; pub mod transforms; pub mod trust; diff --git a/crates/ptwm-core/src/ptwc.rs b/crates/ptwm-core/src/ptwc.rs new file mode 100644 index 0000000..17f21f7 --- /dev/null +++ b/crates/ptwm-core/src/ptwc.rs @@ -0,0 +1,294 @@ +//! Sections of the compact container (`\x89PTWC`): header, chain registry, +//! index and footer. [`crate::container`] writes and reads whole files; +//! [`crate::extension::compact_table`] and [`crate::tensor_record_compact`] +//! hold the Extension Table and record codecs. +//! +//! ```text +//! FILE = HEADER | EXTENSION TABLE | PRELUDE | REGISTRY | BODY | INDEX | FOOTER +//! +//! HEADER +//! magic [u8; 9] \x89PTWC\r\n\x1a\n +//! hflags u16 must be 0 (every bit is reserved) +//! et_len varint compact Extension Table bytes +//! prelude_len varint shared-state prelude bytes (legacy prelude layout) +//! registry_len varint chain registry bytes +//! meta_xxh64 u64 xxh64 (seed 0) of the header bytes before this +//! field, then the Extension Table, prelude and +//! registry bytes +//! REGISTRY +//! count varint at most 32768 +//! chain × count varint chain_len, then a CW2 chain (chain::wire_compact) +//! whose op_ref is an Extension Table index of a +//! built-in transform; chain id = position +//! BODY tensor records (tensor_record_compact), back to back +//! INDEX +//! count varint at most (index bytes) / 3 +//! entry × count varint name_len (1..=65535), UTF-8 name, varint rec_len +//! (>= 3), in body order +//! FOOTER +//! index_len u64 +//! sentinel [u8; 8] \x89TWMEOF\x1a (the legacy EOF sentinel) +//! ``` +//! +//! The body starts right after the registry and the index right after the +//! body: the record lengths must add up to exactly the bytes between them, +//! and names must be unique. Each section parser gets exactly its slice and +//! must consume all of it. Integers are little-endian and varints are +//! canonical unsigned LEB128 ([`crate::varint`]), so each file has one +//! encoding for what it holds. +//! +//! There is no version field. A future layout change sets an `hflags` bit, +//! which current readers reject. + +use std::collections::HashSet; +use std::ops::Range; + +use xxhash_rust::xxh64::Xxh64; + +use crate::chain::wire_compact::{read_chain_compact_deferred_exact, write_chain_compact_deferred}; +use crate::container::{ChainRegistry, MAX_REGISTRY_CHAINS}; +use crate::error::PtwmCoreError; +use crate::extension::{BuiltinKind, ExtensionTable, dispatch_builtin}; +use crate::header::MAGIC_PTWC; +use crate::index::{EOF_SENTINEL, IndexEntry, name_hash}; +use crate::tensor_record_compact::MIN_RECORD_LEN; +use crate::transforms::op::OpId; +use crate::varint::{Cursor, write_varint}; + +/// Footer size: `index_len` and the EOF sentinel. +pub const FOOTER_LEN: usize = 16; +/// Longest tensor name. +pub const MAX_NAME_LEN: usize = u16::MAX as usize; + +fn err(msg: impl std::fmt::Display) -> PtwmCoreError { + PtwmCoreError::InvalidContainer(format!("compact container: {msg}")) +} + +/// Serialize the header for sections of the given bytes. +pub fn header_bytes(et: &[u8], prelude: &[u8], registry: &[u8]) -> Vec { + let mut h = Vec::with_capacity(9 + 2 + 3 * 10 + 8); + h.extend_from_slice(&MAGIC_PTWC); + h.extend_from_slice(&0u16.to_le_bytes()); + write_varint(&mut h, et.len() as u64); + write_varint(&mut h, prelude.len() as u64); + write_varint(&mut h, registry.len() as u64); + let mut hasher = Xxh64::new(0); + hasher.update(&h); + hasher.update(et); + hasher.update(prelude); + hasher.update(registry); + h.extend_from_slice(&hasher.digest().to_le_bytes()); + h +} + +/// Byte ranges of each section of a compact container. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Sections { + pub header: Range, + pub extension_table: Range, + pub prelude: Range, + pub registry: Range, + pub body: Range, + pub index: Range, +} + +/// Locate and check the sections of a compact container: magic, `hflags`, +/// section lengths, footer, and `meta_xxh64`. The caller checks the EOF +/// sentinel first. +pub fn locate_sections(buf: &[u8]) -> Result { + if buf.len() < MAGIC_PTWC.len() + 2 + FOOTER_LEN || buf[..MAGIC_PTWC.len()] != MAGIC_PTWC { + return Err(err("bad magic or file too short")); + } + let footer_start = buf.len() - FOOTER_LEN; + let mut r = Cursor::new(&buf[..footer_start], "compact container header"); + r.take(MAGIC_PTWC.len())?; + let hflags = u16::from_le_bytes(r.array()?); + if hflags != 0 { + return Err(r.err(format_args!("unknown hflags bits {hflags:#06x}"))); + } + let et_len = r.varint()?; + let prelude_len = r.varint()?; + let registry_len = r.varint()?; + let hash_pos = r.pos(); + let stored_hash = r.u64_le()?; + let header_end = r.pos(); + + let section = |start: usize, len: u64, what: &str| -> Result, PtwmCoreError> { + usize::try_from(len) + .ok() + .and_then(|len| start.checked_add(len)) + .filter(|&end| end <= footer_start) + .map(|end| start..end) + .ok_or_else(|| err(format_args!("{what} section exceeds the file"))) + }; + let extension_table = section(header_end, et_len, "extension table")?; + let prelude = section(extension_table.end, prelude_len, "prelude")?; + let registry = section(prelude.end, registry_len, "registry")?; + + let index_len = u64::from_le_bytes(buf[footer_start..footer_start + 8].try_into().unwrap()); + let index_start = usize::try_from(index_len) + .ok() + .and_then(|len| footer_start.checked_sub(len)) + .ok_or_else(|| err("index_len exceeds the file"))?; + if registry.end > index_start { + return Err(err("index overlaps the registry")); + } + + let mut hasher = Xxh64::new(0); + hasher.update(&buf[..hash_pos]); + hasher.update(&buf[header_end..registry.end]); + if hasher.digest() != stored_hash { + return Err(err("meta_xxh64 mismatch")); + } + Ok(Sections { + header: 0..header_end, + extension_table, + prelude, + body: registry.end..index_start, + registry, + index: index_start..footer_start, + }) +} + +/// Encode the chain registry. Chain ids must equal their positions; a +/// Source node with empty params is written as a bound Source. `op_index` +/// maps an op to its Extension Table index. +pub fn write_registry( + registry: &ChainRegistry, + mut op_index: impl FnMut(OpId) -> Result, +) -> Result, PtwmCoreError> { + if registry.chains.len() > MAX_REGISTRY_CHAINS { + return Err(err(format_args!( + "{} chains exceed the limit of {MAX_REGISTRY_CHAINS}", + registry.chains.len() + ))); + } + let mut out = Vec::new(); + write_varint(&mut out, registry.chains.len() as u64); + let mut chain_bytes = Vec::new(); + for (pos, (id, chain)) in registry.chains.iter().enumerate() { + if usize::from(*id) != pos { + return Err(err(format_args!( + "registry chain at position {pos} has id {id}; compact chain ids are positions" + ))); + } + chain_bytes.clear(); + write_chain_compact_deferred(chain, &mut op_index, &mut chain_bytes) + .map_err(|e| err(format_args!("registry chain {id}: {e}")))?; + write_varint(&mut out, chain_bytes.len() as u64); + out.extend_from_slice(&chain_bytes); + } + Ok(out) +} + +/// Parse a chain registry that fills `buf` exactly. Op references resolve +/// through `table`, whose entry must be a built-in transform. +pub fn read_registry(buf: &[u8], table: &ExtensionTable) -> Result { + let mut r = Cursor::new(buf, "compact chain registry"); + let count = r.varint_max(MAX_REGISTRY_CHAINS as u64, "chain count")? as usize; + // Each chain takes at least its length byte and 5 chain bytes. + if count > buf.len() / 6 { + return Err(r.err("chain count exceeds the registry bytes")); + } + let resolve = |op_ref: u64| -> Result { + let entry = usize::try_from(op_ref) + .ok() + .and_then(|i| table.entries.get(i)) + .ok_or_else(|| err(format_args!("op_ref {op_ref} is outside the Extension Table")))?; + match dispatch_builtin(&entry.canonical_id) { + Some(BuiltinKind::Op(op)) => Ok(op), + _ => Err(err(format_args!( + "op_ref {op_ref} names {:?}, which is not a built-in transform", + entry.human_label + ))), + } + }; + let mut chains = Vec::with_capacity(count); + for id in 0..count { + let bytes = r.len_bytes()?; + let chain = read_chain_compact_deferred_exact(bytes, resolve) + .map_err(|e| r.err(format_args!("chain {id}: {e}")))?; + chains.push((id as u16, chain)); + } + r.finish("the last chain")?; + Ok(ChainRegistry { chains }) +} + +/// Encode the index for `(name, record length)` pairs in body order. +pub fn write_index(entries: &[(String, u64)]) -> Result, PtwmCoreError> { + let mut out = Vec::new(); + write_varint(&mut out, entries.len() as u64); + for (name, len) in entries { + if name.is_empty() || name.len() > MAX_NAME_LEN { + return Err(err(format_args!( + "tensor name of {} bytes (must be 1..={MAX_NAME_LEN})", + name.len() + ))); + } + if *len < MIN_RECORD_LEN { + return Err(err(format_args!("record of {len} bytes for {name:?}"))); + } + write_varint(&mut out, name.len() as u64); + out.extend_from_slice(name.as_bytes()); + write_varint(&mut out, *len); + } + Ok(out) +} + +/// Parse the index of a body at `body` (absolute range). Returns entries +/// sorted by `(name_hash, name)` with absolute offsets; their names are set. +pub fn read_index(buf: &[u8], body: Range) -> Result, PtwmCoreError> { + let mut r = Cursor::new(buf, "compact container index"); + let count = r.varint()?; + // Each entry takes at least 3 bytes. + if count > (buf.len() / 3) as u64 { + return Err(r.err("entry count exceeds the index bytes")); + } + let mut entries = Vec::with_capacity(count as usize); + let mut names = HashSet::with_capacity(count as usize); + let mut offset = body.start as u64; + for _ in 0..count { + let name_len = r.varint_max(MAX_NAME_LEN as u64, "name length")? as usize; + if name_len == 0 { + return Err(r.err("empty tensor name")); + } + let name = std::str::from_utf8(r.take(name_len)?) + .map_err(|_| r.err("tensor name is not UTF-8"))?; + if !names.insert(name) { + return Err(r.err(format_args!("duplicate tensor name {name:?}"))); + } + let len = r.varint()?; + if len < MIN_RECORD_LEN { + return Err(r.err(format_args!("record of {len} bytes for {name:?}"))); + } + let end = offset + .checked_add(len) + .filter(|&end| end <= body.end as u64) + .ok_or_else(|| r.err(format_args!("record of {name:?} runs past the body")))?; + entries.push(IndexEntry { + name_hash: name_hash(name), + tensor_offset: offset, + tensor_len: len, + name: Some(name.into()), + }); + offset = end; + } + r.finish("the last index entry")?; + if offset != body.end as u64 { + return Err(err(format_args!( + "records cover {} of {} body bytes", + offset - body.start as u64, + body.len() + ))); + } + entries.sort_by(|a, b| (a.name_hash, &a.name).cmp(&(b.name_hash, &b.name))); + Ok(entries) +} + +/// Footer bytes for an index of `index_len` bytes. +pub fn footer_bytes(index_len: u64) -> [u8; FOOTER_LEN] { + let mut f = [0u8; FOOTER_LEN]; + f[..8].copy_from_slice(&index_len.to_le_bytes()); + f[8..].copy_from_slice(&EOF_SENTINEL); + f +} diff --git a/crates/ptwm-core/src/tensor_record_compact.rs b/crates/ptwm-core/src/tensor_record_compact.rs new file mode 100644 index 0000000..530ed87 --- /dev/null +++ b/crates/ptwm-core/src/tensor_record_compact.rs @@ -0,0 +1,798 @@ +//! Tensor and plane records of the compact container (`\x89PTWC`). +//! +//! A compact record holds the same [`TensorRecord`] as a legacy record, with +//! varints instead of fixed-width fields, without the name (the container +//! index stores it), and without fields the reader can derive. All varints +//! are canonical unsigned LEB128 ([`crate::varint`]); integers are +//! little-endian. +//! +//! ```text +//! TENSOR RECORD +//! tflags u8 TF_* bits; bits 6-7 must be 0 +//! [chain_id] varint registry index, < registry count (no TF_INLINE_CHAIN) +//! dtype_code varint at most u16::MAX +//! input_format u8 +//! orig_size varint +//! [shape] varint ndim (<= 64), then varint per dimension (TF_SHAPE) +//! [dtype_name] varint len, then UTF-8 bytes (TF_DTYPE_NAME) +//! [raw_meta] varint len, then bytes (TF_RAW_META) +//! [payload_hash] u64 (TF_PAYLOAD_HASH) +//! [inline chain] varint et_len, compact Extension Table of the chain's +//! ops in first-use node order, varint chain_len, CW2 +//! chain (TF_INLINE_CHAIN) +//! [deps] varint count >= 1, then per dependency: ref_kind u8, +//! ref_flags u8 (bit 0: hash), varint ref_len, bytes, +//! [hash 32 B] (TF_DEPS) +//! n_planes varint at most 255 +//! plane × n_planes +//! +//! PLANE +//! pflags u8 PF_* bits, state source in bits 2-3; bits 6-7 +//! must be 0 +//! role u8 PlaneRole +//! codec_ref varint Extension Table index of a plane codec +//! [state_version] u8 (PF_STATE_EXTRA) +//! [state_info] varint at most u16::MAX (PF_STATE_EXTRA) +//! [inline_len] varint inline state length (state source Inline) +//! payload_len varint +//! [row_len] varint 1..=u32::MAX (PF_LAYOUT_ROWS) +//! [crc32] u32 (PF_CRC32) +//! [chunks] varint n (<= payload_len), then n × (varint offset, +//! varint decoded_size), each at most u32::MAX (PF_CHUNKS) +//! [external] ref_kind u8, ref_flags u8 (bit 0: hash), varint +//! ref_len, bytes, state_format_version u8, [hash 32 B] +//! (state source External) +//! [inline state] bytes +//! payload bytes +//! ``` +//! +//! **Metadata.** `TF_SHAPE` stands for the CBOR shape map of +//! [`encode_shape_metadata`] with the canonical dtype name of `dtype_code`, +//! or with the stored name under `TF_DTYPE_NAME` (which must then differ +//! from the canonical name). Other metadata is stored as `TF_RAW_META`. The +//! writer uses `TF_SHAPE` exactly when it reproduces the record's metadata +//! bytes, so the reader rejects `TF_RAW_META` bytes that `TF_SHAPE` could +//! store. +//! +//! **Chains.** A registry chain may have a bound Source (see +//! [`crate::chain::wire_compact::write_chain_compact_deferred`]); the record +//! then needs `TF_SHAPE` and a shape and dtype that bind +//! ([`bound_source_params`]). An inline chain binds its Source to the +//! record the same way as the `\x89PTWS` frame does. +//! +//! **Derived fields.** The in-memory `flags`, `chain_ref` (`0x8000` for an +//! inline chain), `payload_len`, and each plane's `codec_id` (from the +//! Extension Table entry its `codec_ref` names) are not stored. The writer +//! rejects a record whose stored values disagree with them, so each record +//! has one encoding and parsing then writing gives the same bytes. + +use crate::chain::Chain; +use crate::chain::wire_compact::{read_chain_compact_exact, write_chain_compact}; +use crate::codec::{CodecId, StateSource}; +use crate::dtype::canonical_dtype_name; +use crate::error::PtwmCoreError; +use crate::extension::compact_table::{read_compact_table, write_compact_table}; +use crate::extension::{ + BuiltinKind, ExtensionTable, Kind, builtin_registry, dispatch_builtin, +}; +use crate::layout::PlaneLayout; +use crate::metadata::{decode_shape_metadata, encode_shape_metadata}; +use crate::plane_record::{ChunkEntry, ExternalStateRef, PlaneRecord}; +use crate::tensor_record::{ + CHAIN_REF_INLINE_BIT, Dependency, RefKind, TENSOR_FLAG_DEPENDENCIES, TENSOR_FLAG_METADATA, + TENSOR_FLAG_PAYLOAD_HASH, TensorRecord, +}; +use crate::transforms::op::OpId; +use crate::transforms::source::bound_source_params; +use crate::types::PlaneRole; +use crate::varint::{Cursor, write_varint}; + +pub const TF_PAYLOAD_HASH: u8 = 0x01; +pub const TF_INLINE_CHAIN: u8 = 0x02; +pub const TF_DEPS: u8 = 0x04; +pub const TF_SHAPE: u8 = 0x08; +pub const TF_DTYPE_NAME: u8 = 0x10; +pub const TF_RAW_META: u8 = 0x20; +const TF_KNOWN: u8 = + TF_PAYLOAD_HASH | TF_INLINE_CHAIN | TF_DEPS | TF_SHAPE | TF_DTYPE_NAME | TF_RAW_META; + +pub const PF_CRC32: u8 = 0x01; +pub const PF_CHUNKS: u8 = 0x02; +/// Plane state source (`StateSource` value) in bits 2-3. +pub const PF_STATE_SHIFT: u8 = 2; +pub const PF_STATE_MASK: u8 = 0x0C; +pub const PF_LAYOUT_ROWS: u8 = 0x10; +pub const PF_STATE_EXTRA: u8 = 0x20; +const PF_KNOWN: u8 = PF_CRC32 | PF_CHUNKS | PF_STATE_MASK | PF_LAYOUT_ROWS | PF_STATE_EXTRA; + +/// Largest dimension count a compact record stores. Shape metadata with +/// more dimensions is stored as `TF_RAW_META`. +pub const MAX_NDIM: usize = 64; +/// Largest plane count (the legacy record stores it as `u8`). +const MAX_PLANES: u64 = 255; +/// Smallest compact record: tflags, dtype_code, input_format, orig_size, +/// n_planes would be 5; the index requires at least 3. +pub const MIN_RECORD_LEN: u64 = 3; + +const DEP_FLAG_HASH: u8 = 0x01; + +/// What a compact record refers to outside itself. +#[derive(Debug, Clone, Copy)] +pub struct RecordContext<'a> { + /// Per Extension Table index: the codec id of a plane that references + /// the entry, or `None` when the entry is not a plane codec. See + /// [`plane_codec_ids`]. + pub codec_ids: &'a [Option], + /// Per registry chain: whether it has a bound Source that the record + /// must supply. + pub registry_bound: &'a [bool], +} + +/// The codec id each Extension Table entry gives a plane that references +/// it: the built-in codec for a built-in codec entry, `Identity` for a +/// third-party plane codec (which has no wire id), and `None` for entries +/// that are not plane codecs. +pub fn plane_codec_ids(table: &ExtensionTable) -> Vec> { + table + .entries + .iter() + .map(|e| match dispatch_builtin(&e.canonical_id) { + Some(BuiltinKind::Codec(c)) => Some(c), + Some(BuiltinKind::Op(_)) => None, + None if e.kind == Kind::PlaneCodec => Some(CodecId::Identity), + None => None, + }) + .collect() +} + +fn err(msg: impl std::fmt::Display) -> PtwmCoreError { + PtwmCoreError::InvalidContainer(format!("compact tensor record: {msg}")) +} + +/// How a record's metadata is stored. +enum Meta<'a> { + None, + Shape { + shape: Vec, + /// Stored name; `None` when it is the canonical name. + stored_name: Option, + }, + Raw(&'a [u8]), +} + +fn classify_meta(meta: Option<&[u8]>, dtype_code: u16) -> Meta<'_> { + let Some(m) = meta else { + return Meta::None; + }; + match decode_shape_metadata(m) { + Ok((shape, name)) + if shape.len() <= MAX_NDIM && encode_shape_metadata(&shape, &name) == m => + { + let stored_name = (canonical_dtype_name(dtype_code) != Some(name.as_str())).then_some(name); + Meta::Shape { shape, stored_name } + } + _ => Meta::Raw(m), + } +} + +/// Bound Source params for a record's shape metadata as the compact +/// container stores it: `Some` only when the metadata is stored as +/// `TF_SHAPE` and the shape and dtype bind. +pub fn record_binding(meta: Option<&[u8]>, dtype_code: u16) -> Option> { + match classify_meta(meta, dtype_code) { + Meta::Shape { shape, .. } => bound_source_params(&shape, dtype_code), + _ => None, + } +} + +/// Ops of `chain` in first-use node order, without repeats. +fn ops_in_first_use_order(chain: &Chain) -> Vec { + let mut ops: Vec = Vec::new(); + for n in &chain.nodes { + if !ops.contains(&n.op) { + ops.push(n.op); + } + } + ops +} + +/// Append the compact encoding of `rec` to `out`. `rec.name` is not +/// written. On error `out` is left unchanged. +pub fn write_compact_record( + rec: &TensorRecord, + ctx: &RecordContext<'_>, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + let start = out.len(); + let result = write_unchecked(rec, ctx, out); + if result.is_err() { + out.truncate(start); + } + result +} + +fn write_unchecked( + rec: &TensorRecord, + ctx: &RecordContext<'_>, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + let inline = rec.chain_ref & CHAIN_REF_INLINE_BIT != 0; + if inline != rec.inline_chain.is_some() { + return Err(err("chain_ref inline bit and inline_chain disagree")); + } + if inline && rec.chain_ref != CHAIN_REF_INLINE_BIT { + return Err(err(format_args!( + "inline chain_ref {:#06x} has low bits set; the compact record stores none", + rec.chain_ref + ))); + } + if rec.terminals.len() as u64 > MAX_PLANES { + return Err(err(format_args!("{} planes (max {MAX_PLANES})", rec.terminals.len()))); + } + if rec.dependencies.len() > u16::MAX as usize { + return Err(err(format_args!( + "{} dependencies (max {})", + rec.dependencies.len(), + u16::MAX + ))); + } + if let Some(m) = &rec.tensor_metadata + && m.len() as u64 > u64::from(u32::MAX) + { + return Err(err("tensor metadata longer than u32::MAX bytes")); + } + let meta = classify_meta(rec.tensor_metadata.as_deref(), rec.dtype_code); + let binding = match &meta { + Meta::Shape { shape, .. } => bound_source_params(shape, rec.dtype_code), + _ => None, + }; + + let mut tflags = 0u8; + if rec.payload_hash.is_some() { + tflags |= TF_PAYLOAD_HASH; + } + if inline { + tflags |= TF_INLINE_CHAIN; + } + if !rec.dependencies.is_empty() { + tflags |= TF_DEPS; + } + match &meta { + Meta::None => {} + Meta::Shape { stored_name, .. } => { + tflags |= TF_SHAPE; + if stored_name.is_some() { + tflags |= TF_DTYPE_NAME; + } + } + Meta::Raw(_) => tflags |= TF_RAW_META, + } + out.push(tflags); + + if !inline { + let id = usize::from(rec.chain_ref); + let bound = ctx.registry_bound.get(id).ok_or_else(|| { + err(format_args!( + "chain_ref {id} is not in the registry ({} chains)", + ctx.registry_bound.len() + )) + })?; + if *bound && binding.is_none() { + return Err(err(format_args!( + "registry chain {id} has a bound Source, but the record has no shape \ + metadata that binds it" + ))); + } + write_varint(out, id as u64); + } + write_varint(out, u64::from(rec.dtype_code)); + out.push(rec.input_format); + write_varint(out, rec.orig_size); + match &meta { + Meta::None => {} + Meta::Shape { shape, stored_name } => { + write_varint(out, shape.len() as u64); + for &d in shape { + write_varint(out, d); + } + if let Some(name) = stored_name { + write_varint(out, name.len() as u64); + out.extend_from_slice(name.as_bytes()); + } + } + Meta::Raw(m) => { + write_varint(out, m.len() as u64); + out.extend_from_slice(m); + } + } + if let Some(h) = rec.payload_hash { + out.extend_from_slice(&h.to_le_bytes()); + } + if let Some(chain) = &rec.inline_chain { + let ops = ops_in_first_use_order(chain); + let reg = builtin_registry(); + let table = ExtensionTable { + entries: ops + .iter() + .map(|op| { + reg.entry_by_name(op.builtin_name()) + .cloned() + .ok_or_else(|| err(format_args!("op {op:?} has no built-in entry"))) + }) + .collect::>()?, + }; + let mut et = Vec::new(); + write_compact_table(&table, &mut et)?; + let mut chain_bytes = Vec::new(); + write_chain_compact( + chain, + binding.as_deref(), + |op| { + Ok(ops + .iter() + .position(|&o| o == op) + .expect("every op of the chain is listed") as u64) + }, + &mut chain_bytes, + )?; + write_varint(out, et.len() as u64); + out.extend_from_slice(&et); + write_varint(out, chain_bytes.len() as u64); + out.extend_from_slice(&chain_bytes); + } + if !rec.dependencies.is_empty() { + write_varint(out, rec.dependencies.len() as u64); + for dep in &rec.dependencies { + if dep.ref_bytes.len() > u16::MAX as usize { + return Err(err(format_args!( + "dependency ref of {} bytes (max {})", + dep.ref_bytes.len(), + u16::MAX + ))); + } + out.push(dep.ref_kind as u8); + out.push(if dep.expected_hash.is_some() { + DEP_FLAG_HASH + } else { + 0 + }); + write_varint(out, dep.ref_bytes.len() as u64); + out.extend_from_slice(&dep.ref_bytes); + if let Some(h) = &dep.expected_hash { + out.extend_from_slice(h); + } + } + } + write_varint(out, rec.terminals.len() as u64); + for (i, plane) in rec.terminals.iter().enumerate() { + write_plane(plane, ctx, out).map_err(|e| err(format_args!("plane {i}: {e}")))?; + } + Ok(()) +} + +fn write_plane( + p: &PlaneRecord, + ctx: &RecordContext<'_>, + out: &mut Vec, +) -> Result<(), PtwmCoreError> { + let idx = usize::from(p.codec_table_idx); + match ctx.codec_ids.get(idx) { + Some(Some(c)) if *c == p.codec_id => {} + Some(Some(c)) => { + return Err(err(format_args!( + "codec_id {:?} differs from {c:?}, the codec of Extension Table entry {idx}", + p.codec_id + ))); + } + Some(None) => { + return Err(err(format_args!( + "Extension Table entry {idx} is not a plane codec" + ))); + } + None => { + return Err(err(format_args!( + "codec_table_idx {idx} is outside the Extension Table ({} entries)", + ctx.codec_ids.len() + ))); + } + } + if p.payload_len != p.payload_bytes.len() as u64 { + return Err(err(format_args!( + "payload_len {} differs from the {} payload bytes", + p.payload_len, + p.payload_bytes.len() + ))); + } + let is_inline = p.state_source == StateSource::Inline; + if !is_inline && !p.inline_state_bytes.is_empty() { + return Err(err("inline state bytes without an Inline state source")); + } + let is_external = p.state_source == StateSource::External; + if is_external != p.external_state.is_some() { + return Err(err("external state ref and External state source disagree")); + } + let mut pflags = (p.state_source as u8) << PF_STATE_SHIFT; + if p.crc32.is_some() { + pflags |= PF_CRC32; + } + if p.chunk_table.is_some() { + pflags |= PF_CHUNKS; + } + if p.layout.row_len().is_some() { + pflags |= PF_LAYOUT_ROWS; + } + let extra = p.state_version != 0 || p.state_info != 0; + if extra { + pflags |= PF_STATE_EXTRA; + } + out.push(pflags); + out.push(p.role as u8); + write_varint(out, idx as u64); + if extra { + out.push(p.state_version); + write_varint(out, u64::from(p.state_info)); + } + if is_inline { + write_varint(out, p.inline_state_bytes.len() as u64); + } + write_varint(out, p.payload_bytes.len() as u64); + if let Some(row_len) = p.layout.row_len() { + write_varint(out, u64::from(row_len)); + } + if let Some(crc) = p.crc32 { + out.extend_from_slice(&crc.to_le_bytes()); + } + if let Some(chunks) = &p.chunk_table { + if chunks.len() > p.payload_bytes.len() { + return Err(err(format_args!( + "{} chunks for a {}-byte payload", + chunks.len(), + p.payload_bytes.len() + ))); + } + write_varint(out, chunks.len() as u64); + for c in chunks { + write_varint(out, u64::from(c.offset_in_payload)); + write_varint(out, u64::from(c.decoded_size)); + } + } + if let Some(ext) = &p.external_state { + if ext.ref_bytes.len() > u16::MAX as usize { + return Err(err(format_args!( + "external state ref of {} bytes (max {})", + ext.ref_bytes.len(), + u16::MAX + ))); + } + out.push(ext.ref_kind); + out.push(if ext.expected_hash.is_some() { + DEP_FLAG_HASH + } else { + 0 + }); + write_varint(out, ext.ref_bytes.len() as u64); + out.extend_from_slice(&ext.ref_bytes); + out.push(ext.state_format_version); + if let Some(h) = &ext.expected_hash { + out.extend_from_slice(h); + } + } + out.extend_from_slice(&p.inline_state_bytes); + out.extend_from_slice(&p.payload_bytes); + Ok(()) +} + +/// Parse a compact record that fills `buf` exactly. The returned record has +/// an empty `name`; the container index supplies it. +pub fn parse_compact_record( + buf: &[u8], + ctx: &RecordContext<'_>, +) -> Result { + let mut r = Cursor::new(buf, "compact tensor record"); + let tflags = r.u8()?; + if tflags & !TF_KNOWN != 0 { + return Err(r.err(format_args!( + "unknown tflags bits {:#04x}", + tflags & !TF_KNOWN + ))); + } + let inline = tflags & TF_INLINE_CHAIN != 0; + if tflags & TF_DTYPE_NAME != 0 && tflags & TF_SHAPE == 0 { + return Err(r.err("TF_DTYPE_NAME without TF_SHAPE")); + } + if tflags & TF_RAW_META != 0 && tflags & TF_SHAPE != 0 { + return Err(r.err("TF_RAW_META with TF_SHAPE")); + } + + let chain_id = if inline { + None + } else { + let n = ctx.registry_bound.len() as u64; + let id = r.varint()?; + if id >= n { + return Err(r.err(format_args!("chain id {id} is not in the registry ({n} chains)"))); + } + Some(id as usize) + }; + let dtype_code = r.varint_max(u64::from(u16::MAX), "dtype_code")? as u16; + let input_format = r.u8()?; + let orig_size = r.varint()?; + + let mut tensor_metadata = None; + let mut binding = None; + if tflags & TF_SHAPE != 0 { + let ndim = r.varint_max(MAX_NDIM as u64, "ndim")? as usize; + if ndim > r.remaining() { + return Err(r.err("truncated shape")); + } + let mut shape = Vec::with_capacity(ndim); + for _ in 0..ndim { + shape.push(r.varint()?); + } + let canonical = canonical_dtype_name(dtype_code); + let name: String = if tflags & TF_DTYPE_NAME != 0 { + let bytes = r.len_bytes()?; + let name = + std::str::from_utf8(bytes).map_err(|_| r.err("dtype name is not UTF-8"))?; + if canonical == Some(name) { + return Err(r.err("stored dtype name equals the canonical name (non-canonical)")); + } + name.to_owned() + } else { + canonical + .ok_or_else(|| { + r.err(format_args!( + "shape without a dtype name for dtype_code {dtype_code}, which has no \ + canonical name" + )) + })? + .to_owned() + }; + binding = bound_source_params(&shape, dtype_code); + tensor_metadata = Some(encode_shape_metadata(&shape, &name)); + } else if tflags & TF_RAW_META != 0 { + let bytes = r.len_bytes()?; + if bytes.len() as u64 > u64::from(u32::MAX) { + return Err(r.err("raw metadata longer than u32::MAX bytes")); + } + if !matches!(classify_meta(Some(bytes), dtype_code), Meta::Raw(_)) { + return Err(r.err("raw metadata that TF_SHAPE stores (non-canonical)")); + } + tensor_metadata = Some(bytes.to_vec()); + } + + if let Some(id) = chain_id + && ctx.registry_bound[id] + && binding.is_none() + { + return Err(r.err(format_args!( + "registry chain {id} has a bound Source, but the record has no shape that binds it" + ))); + } + + let payload_hash = if tflags & TF_PAYLOAD_HASH != 0 { + Some(r.u64_le()?) + } else { + None + }; + + let inline_chain = if inline { + let et_bytes = r.len_bytes()?; + let table = read_compact_table(et_bytes) + .map_err(|e| r.err(format_args!("inline chain: {e}")))?; + let ops: Vec = table + .entries + .iter() + .map(|e| match dispatch_builtin(&e.canonical_id) { + Some(BuiltinKind::Op(op)) => Ok(op), + _ => Err(r.err(format_args!( + "inline chain table entry {:?} is not a transform", + e.human_label + ))), + }) + .collect::>()?; + let chain_bytes = r.len_bytes()?; + let chain = read_chain_compact_exact(chain_bytes, binding.as_deref(), |i| { + usize::try_from(i) + .ok() + .and_then(|i| ops.get(i).copied()) + .ok_or_else(|| err(format_args!("inline chain op_ref {i} outside its table"))) + })?; + if ops_in_first_use_order(&chain) != ops { + return Err(r.err( + "inline chain table does not list the ops in first-use order (non-canonical)", + )); + } + Some(chain) + } else { + None + }; + + let mut dependencies = Vec::new(); + if tflags & TF_DEPS != 0 { + let n = r.varint_max(u64::from(u16::MAX), "dependency count")?; + if n == 0 { + return Err(r.err("TF_DEPS with no dependencies")); + } + // Each dependency takes at least 3 bytes. + if n > (r.remaining() / 3) as u64 { + return Err(r.err("dependency count exceeds the record")); + } + dependencies.reserve_exact(n as usize); + for _ in 0..n { + let kind = r.u8()?; + let ref_kind = RefKind::from_u8(kind) + .ok_or_else(|| r.err(format_args!("unknown ref_kind {kind:#04x}")))?; + let flags = r.u8()?; + if flags & !DEP_FLAG_HASH != 0 { + return Err(r.err(format_args!("unknown dependency flag bits {flags:#04x}"))); + } + let len = r.varint_max(u64::from(u16::MAX), "dependency ref length")? as usize; + let ref_bytes = r.take(len)?.to_vec(); + let expected_hash = if flags & DEP_FLAG_HASH != 0 { + Some(r.array::<32>()?) + } else { + None + }; + dependencies.push(Dependency { + ref_kind, + expected_hash, + ref_bytes, + }); + } + } + + let n_planes = r.varint_max(MAX_PLANES, "plane count")?; + // Each plane takes at least 4 bytes. + if n_planes > (r.remaining() / 4) as u64 { + return Err(r.err("plane count exceeds the record")); + } + let mut terminals = Vec::with_capacity(n_planes as usize); + for _ in 0..n_planes { + terminals.push(read_plane(&mut r, ctx)?); + } + r.finish("the last plane")?; + + let mut flags = 0u8; + if !dependencies.is_empty() { + flags |= TENSOR_FLAG_DEPENDENCIES; + } + if tensor_metadata.is_some() { + flags |= TENSOR_FLAG_METADATA; + } + if payload_hash.is_some() { + flags |= TENSOR_FLAG_PAYLOAD_HASH; + } + Ok(TensorRecord { + chain_ref: match chain_id { + Some(id) => id as u16, + None => CHAIN_REF_INLINE_BIT, + }, + dtype_code, + input_format, + flags, + orig_size, + name: String::new(), + payload_hash, + inline_chain, + dependencies, + tensor_metadata, + terminals, + }) +} + +fn read_plane(r: &mut Cursor<'_>, ctx: &RecordContext<'_>) -> Result { + let pflags = r.u8()?; + if pflags & !PF_KNOWN != 0 { + return Err(r.err(format_args!( + "unknown pflags bits {:#04x}", + pflags & !PF_KNOWN + ))); + } + let state_source = StateSource::from_u8((pflags & PF_STATE_MASK) >> PF_STATE_SHIFT) + .expect("two bits always name a state source"); + let role_byte = r.u8()?; + let role = PlaneRole::from_u8(role_byte) + .ok_or_else(|| r.err(format_args!("unknown plane role {role_byte:#04x}")))?; + let codec_ref = r.varint_max(u64::from(u16::MAX), "codec_ref")? as u16; + let codec_id = match ctx.codec_ids.get(usize::from(codec_ref)) { + Some(Some(c)) => *c, + Some(None) => { + return Err(r.err(format_args!( + "codec_ref {codec_ref} names an entry that is not a plane codec" + ))); + } + None => { + return Err(r.err(format_args!( + "codec_ref {codec_ref} is outside the Extension Table ({} entries)", + ctx.codec_ids.len() + ))); + } + }; + let (state_version, state_info) = if pflags & PF_STATE_EXTRA != 0 { + let v = r.u8()?; + let info = r.varint_max(u64::from(u16::MAX), "state_info")? as u16; + if v == 0 && info == 0 { + return Err(r.err("PF_STATE_EXTRA with zero state_version and state_info")); + } + (v, info) + } else { + (0, 0) + }; + let inline_len = if state_source == StateSource::Inline { + r.len_varint()? + } else { + 0 + }; + let payload_len = r.len_varint()?; + let layout = if pflags & PF_LAYOUT_ROWS != 0 { + let row_len = r.varint_max(u64::from(u32::MAX), "row_len")? as u32; + PlaneLayout::rows(row_len).ok_or_else(|| r.err("row_len 0"))? + } else { + PlaneLayout::Flat + }; + let crc32 = if pflags & PF_CRC32 != 0 { + Some(r.u32_le()?) + } else { + None + }; + let chunk_table = if pflags & PF_CHUNKS != 0 { + let n = r.varint_max(payload_len as u64, "chunk count")?; + // Each chunk entry takes at least 2 bytes. + if n > (r.remaining() / 2) as u64 { + return Err(r.err("chunk count exceeds the record")); + } + let mut chunks = Vec::with_capacity(n as usize); + for _ in 0..n { + let offset = r.varint_max(u64::from(u32::MAX), "chunk offset")? as u32; + let size = r.varint_max(u64::from(u32::MAX), "chunk decoded size")? as u32; + chunks.push(ChunkEntry { + offset_in_payload: offset, + decoded_size: size, + }); + } + Some(chunks) + } else { + None + }; + let external_state = if state_source == StateSource::External { + let ref_kind = r.u8()?; + let flags = r.u8()?; + if flags & !DEP_FLAG_HASH != 0 { + return Err(r.err(format_args!("unknown external state flag bits {flags:#04x}"))); + } + let len = r.varint_max(u64::from(u16::MAX), "external state ref length")? as usize; + let ref_bytes = r.take(len)?.to_vec(); + let state_format_version = r.u8()?; + let expected_hash = if flags & DEP_FLAG_HASH != 0 { + Some(r.array::<32>()?) + } else { + None + }; + Some(ExternalStateRef { + ref_kind, + expected_hash, + ref_bytes, + state_format_version, + }) + } else { + None + }; + let inline_state_bytes = r.take(inline_len)?.to_vec(); + let payload_bytes = r.take(payload_len)?.to_vec(); + Ok(PlaneRecord { + role, + codec_id, + codec_table_idx: codec_ref, + state_source, + state_version, + state_info, + payload_len: payload_len as u64, + crc32, + chunk_table, + external_state, + inline_state_bytes, + payload_bytes, + layout, + }) +} From a67f73c9ef79ed8cc054650958de729d3e75b9f9 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 00:37:02 +0000 Subject: [PATCH 15/34] feat(core): compact multi-tensor container (PTWC) Stage 6 of the issue #5 work (synthesis sections 1, 2, 4, 5). Completes the interrupted snapshot b7c3dce. - Magic \x89PTWC\r\n\x1a\n. Header: hflags u16 (all bits reserved), varint lengths of the Extension Table, prelude and registry, and meta_xxh64 over the header and those sections. Footer: index_len u64 and the unchanged EOF sentinel. Each section parser gets exactly its slice and must consume it (crate::ptwc). - Compact Extension Table (extension::compact_table): tags 0x01 op id, 0x02 codec id, 0x03 built-in name, 0x7F full third-party entry. Built-ins expand to the reader's own entries, so dispatch, trust and the missing-flavor check are shared with the legacy reader. - Registry: CW2 chains whose op_ref is a table index. A Source with empty params in memory is a bound Source; the record's shape supplies its params. - Records (tensor_record_compact): no name, tflags/pflags with every reserved bit rejected, shape stored as varints (dtype name only when not canonical, otherwise raw metadata), inline chains with a first-use-order table, plane codec_id derived from the table entry. The writer rejects records whose stored fields disagree with the derived ones, so parse-then-write is the identity. - Index: names in body order; lookup compares names; duplicates, gaps and overlaps are rejected. - ContainerLayout {Legacy, Compact}; ContainerWriter::with_layout (new() stays Legacy); RecordEncoder / encode_record for parallel record serialization; set_compact_tables for verbatim tables. - CompressorOptions now derives Default and has container_layout (default Legacy until the stage that flips defaults). In Compact, compress_model binds Sources, so tensors that differ only in shape share one registry chain. Registry dedup now uses a HashMap for both layouts (Chain and Role derive Hash); legacy output is unchanged (the golden writer hash test passes). - Reader sniffs the magic and stores the layout; parse_record, tensor_record and tensor_records are used by tensor_shape, tensor_names (index only for PTWC), decode_tensor, decode_model and the ptwm-py info, delta-hash and plane-summary functions. The compact reader requires exact sizes (Source == orig_size, plane output == descriptor length). - Transcode explode rejects PTWC explicitly until the compact manifest lands. Tests: legacy fixtures rewritten as PTWC (registry and inline chains, shared prelude state, chunked planes, dependency, Rows layout, bf16/fp16/fp32/fp8/int8) decode identically and rewrite byte-exactly; compress_model matrix over both layouts x hash/CRC x chunked, with 0-d, non-canonical dtype name and dependency tensors; raw metadata, external state and state_extra records; canonicality rejection for header, table, record, plane, registry and index; bounded counts; a P4 file with Source slack fails in PTWC; Header::from_bytes on PTWC gives "bad magic"; bound Source gives one chain for 64 shapes; duplicate names; 30k record mutations re-encode identically or fail. Overhead for 64 fp32 (64,) tensors: 81 B/tensor with hash and CRC (legacy 229), 57 B without; built-in-only table at most 4 B/entry. Fuzz: new targets parse_wire_compact and extension_table_compact (re-encode identity), container_open also runs decode_model; PTWC seeds added. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/chain/graph.rs | 8 +- crates/ptwm-core/src/compact.rs | 1 + crates/ptwm-core/src/compressor.rs | 102 ++- crates/ptwm-core/src/container.rs | 49 +- .../ptwm-core/src/extension/compact_table.rs | 35 +- crates/ptwm-core/src/ptwc.rs | 149 +++- crates/ptwm-core/src/tensor_record_compact.rs | 467 ++++++++++- crates/ptwm-core/src/transcode/mod.rs | 5 + crates/ptwm-core/src/transforms/op.rs | 2 +- crates/ptwm-core/src/types/role.rs | 10 +- crates/ptwm-core/tests/decode_bomb.rs | 1 + crates/ptwm-core/tests/legacy_fixtures.rs | 1 + crates/ptwm-core/tests/ptwc.rs | 725 ++++++++++++++++++ crates/ptwm-py/src/lib.rs | 23 +- fuzz/Cargo.toml | 12 + fuzz/corpus/container_open/ptwc_bound_source | Bin 0 -> 365 bytes fuzz/corpus/container_open/ptwc_trimmed | Bin 0 -> 2457 bytes .../extension_table_compact/bound_source | Bin 0 -> 9 bytes fuzz/corpus/extension_table_compact/trimmed | Bin 0 -> 23 bytes fuzz/corpus/parse_wire_compact/record_0 | 2 + fuzz/corpus/parse_wire_compact/record_1 | 1 + fuzz/corpus/parse_wire_compact/record_2 | 1 + fuzz/corpus/parse_wire_compact/record_3 | Bin 0 -> 81 bytes fuzz/fuzz_targets/container_open.rs | 17 +- fuzz/fuzz_targets/extension_table_compact.rs | 14 + fuzz/fuzz_targets/parse_wire_compact.rs | 25 + 26 files changed, 1537 insertions(+), 113 deletions(-) create mode 100644 crates/ptwm-core/tests/ptwc.rs create mode 100644 fuzz/corpus/container_open/ptwc_bound_source create mode 100644 fuzz/corpus/container_open/ptwc_trimmed create mode 100644 fuzz/corpus/extension_table_compact/bound_source create mode 100644 fuzz/corpus/extension_table_compact/trimmed create mode 100644 fuzz/corpus/parse_wire_compact/record_0 create mode 100644 fuzz/corpus/parse_wire_compact/record_1 create mode 100644 fuzz/corpus/parse_wire_compact/record_2 create mode 100644 fuzz/corpus/parse_wire_compact/record_3 create mode 100644 fuzz/fuzz_targets/extension_table_compact.rs create mode 100644 fuzz/fuzz_targets/parse_wire_compact.rs diff --git a/crates/ptwm-core/src/chain/graph.rs b/crates/ptwm-core/src/chain/graph.rs index 9290984..934d337 100644 --- a/crates/ptwm-core/src/chain/graph.rs +++ b/crates/ptwm-core/src/chain/graph.rs @@ -3,13 +3,13 @@ use crate::transforms::op::OpId; use crate::types::role::Role; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ChainNode { pub op: OpId, pub params: Vec, // op-specific TLV payload (encoded by Op::write_params) } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct ChainEdge { pub src_node: u8, pub src_output_idx: u8, @@ -19,14 +19,14 @@ pub struct ChainEdge { pub vendor_bytes: Vec, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct TerminalRef { pub node_idx: u8, pub output_idx: u8, pub role: Role, } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Chain { pub nodes: Vec, pub edges: Vec, diff --git a/crates/ptwm-core/src/compact.rs b/crates/ptwm-core/src/compact.rs index fce88ea..a582008 100644 --- a/crates/ptwm-core/src/compact.rs +++ b/crates/ptwm-core/src/compact.rs @@ -602,6 +602,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() } } diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index e2c6c66..9050817 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -5,6 +5,7 @@ //! capability-keyed [`dispatch`] menu, fits Pass-2 shared state on the //! chosen terminals, and writes a chain-registry-backed `.ptwm` container. +use std::collections::HashMap; use std::io::{Seek, Write}; use std::sync::Arc; @@ -16,7 +17,7 @@ use crate::chain::runtime::{ForwardContext, chain_source_descriptor, forward_cha use crate::chain::validate::validate_chain; use crate::codec::{CodecId, StateSource, codec_for}; use crate::codecs::order1_scale_ac; -use crate::container::{ChainRegistry, ContainerWriter}; +use crate::container::{ChainRegistry, ContainerLayout, ContainerWriter}; use crate::dispatch::dispatch; use crate::error::PtwmCoreError; use crate::extension::CanonicalId; @@ -27,7 +28,8 @@ use crate::plane_record::PlaneRecord; use crate::prelude::PreludeEntry; use crate::select::{MenuItem, TrialResult, trial_encode_plane}; use crate::tensor_record::{Dependency, RefKind, TENSOR_FLAG_PAYLOAD_HASH, TensorRecord}; -use crate::transforms::op::Plane; +use crate::tensor_record_compact::record_binding; +use crate::transforms::op::{OpId, Plane}; use crate::types::PlaneRole; use crate::types::descriptor::{ElementWidth, Layout, PlaneDescriptor}; use crate::types::role::{NibbleKind, Role}; @@ -87,7 +89,9 @@ pub struct InputTensor<'a> { pub delta_reference_blake3: Option<[u8; 32]>, } -/// Compressor options. +/// Compressor options. `Default` gives no hint, no hash or CRC, dispatcher +/// codec selection, whole planes and the legacy container layout. +#[derive(Default)] pub struct CompressorOptions { /// Header hint to readers about the dominant codec. pub method_hint: u16, @@ -105,6 +109,8 @@ pub struct CompressorOptions { pub allow_codec_ids: Option>, /// When `Some`, encode each plane in fixed-size chunks. pub chunk_size: Option, + /// On-disk layout of the written container. + pub container_layout: ContainerLayout, } // --------------------------------------------------------------------------- @@ -664,26 +670,40 @@ fn build_chain_registry_with_limit( max_chains: usize, ) -> Result<(ChainRegistry, Vec), PtwmCoreError> { let mut registry: Vec<(u16, Chain)> = Vec::new(); + let mut ids: HashMap<&Chain, u16> = HashMap::with_capacity(chains.len()); let mut tensor_to_idx: Vec = Vec::with_capacity(chains.len()); for c in chains { - let existing = registry.iter().position(|(_, rc)| rc == c); - match existing { - Some(pos) => tensor_to_idx.push(pos as u16), - None => { - if registry.len() >= max_chains { - return Err(PtwmCoreError::InvalidContainer(format!( - "compress_model: more than {max_chains} distinct chains" - ))); - } - let new_id = registry.len() as u16; - registry.push((new_id, c.clone())); - tensor_to_idx.push(new_id); - } + if let Some(&id) = ids.get(c) { + tensor_to_idx.push(id); + continue; + } + if registry.len() >= max_chains { + return Err(PtwmCoreError::InvalidContainer(format!( + "compress_model: more than {max_chains} distinct chains" + ))); } + let new_id = registry.len() as u16; + ids.insert(c, new_id); + registry.push((new_id, c.clone())); + tensor_to_idx.push(new_id); } Ok((ChainRegistry { chains: registry }, tensor_to_idx)) } +/// Replace the Source params of `chain` with empty params (a bound Source, +/// see [`crate::chain::wire_compact::write_chain_compact_deferred`]) when +/// they equal `binding`, the params the tensor record's shape gives. The +/// compact layout then stores one registry chain for tensors that differ +/// only in shape. +fn bind_source(chain: &mut Chain, binding: Option<&[u8]>) { + if let (Some(node), Some(params)) = (chain.nodes.first_mut(), binding) + && node.op == OpId::Source + && node.params == params + { + node.params.clear(); + } +} + // --------------------------------------------------------------------------- // compress_model // --------------------------------------------------------------------------- @@ -946,14 +966,30 @@ pub fn compress_model( } // ---- Pass 4: write container. ----------------------------------------- // Build a deduped chain registry from the chosen chain per tensor. + let tensor_metadata: Vec>> = inputs + .iter() + .map(|inp| match (&inp.shape, &inp.dtype_name) { + (Some(shape), Some(name)) => Some(encode_shape_metadata(shape, name)), + _ => None, + }) + .collect(); + let compact = opts.container_layout == ContainerLayout::Compact; let chosen_chains: Vec = pass1 .iter() .zip(inputs.iter()) - .map(|(outcome, inp)| inp.candidate_chains[outcome.chain_idx].clone()) + .zip(tensor_metadata.iter()) + .map(|((outcome, inp), meta)| { + let mut chain = inp.candidate_chains[outcome.chain_idx].clone(); + if compact { + let binding = record_binding(meta.as_deref(), inp.dtype_code); + bind_source(&mut chain, binding.as_deref()); + } + chain + }) .collect(); let (registry, tensor_chain_indices) = build_chain_registry(&chosen_chains)?; - let mut writer = ContainerWriter::new(writer)?; + let mut writer = ContainerWriter::with_layout(writer, opts.container_layout)?; writer.set_prelude(&prelude_entries)?; // Intern only the codecs the planes reference: interning every built-in // codec costs about 125 B per unused entry, which dominates small files. @@ -1007,6 +1043,7 @@ pub fn compress_model( prof!(prof_t, "pass4_chunked_encode"); + let mut tensor_metadata = tensor_metadata; let mut tensor_records: Vec = Vec::with_capacity(inputs.len()); for (idx, inp) in inputs.iter().enumerate() { let chosen_planes = &chosen_per_tensor[idx]; @@ -1080,11 +1117,6 @@ pub fn compress_model( None }; - let tensor_metadata = match (&inp.shape, &inp.dtype_name) { - (Some(shape), Some(name)) => Some(encode_shape_metadata(shape, name)), - _ => None, - }; - let dependencies: Vec = match inp.delta_reference_blake3 { Some(hash) => vec![Dependency { ref_kind: RefKind::ExternalSafetensors, @@ -1105,26 +1137,20 @@ pub fn compress_model( payload_hash, inline_chain: None, dependencies, - tensor_metadata, + tensor_metadata: tensor_metadata[idx].take(), terminals: plane_records, }; tensor_records.push(rec); } // Serialize tensor records (the payload-copy-heavy step) in parallel — - // write_tensor_record is pure per record — then append the finished byte - // blocks sequentially, which is all the ordered writer needs to do. + // encoding is pure per record — then append the finished byte blocks + // sequentially, which is all the ordered writer needs to do. The encoder + // writes the writer's layout. + let encoder = writer.record_encoder(); let serialized: Vec> = tensor_records .par_iter() - .map(|rec| { - // Pre-size to the plane payloads (the dominant term) plus headroom - // for record/plane metadata, so the large payload copies don't - // repeatedly reallocate. - let est: usize = rec.terminals.iter().map(|t| t.payload_bytes.len()).sum(); - let mut buf = Vec::with_capacity(est + 1024); - crate::tensor_record::write_tensor_record(rec, &mut buf)?; - Ok::, PtwmCoreError>(buf) - }) + .map(|rec| encoder.encode(rec)) .collect::, _>>()?; for (rec, buf) in tensor_records.iter().zip(serialized.iter()) { writer.append_serialized_tensor(&rec.name, buf)?; @@ -1371,6 +1397,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -1445,6 +1472,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }; let input = |chains: Vec| InputTensor { name: "t".into(), @@ -1565,6 +1593,7 @@ mod tests { forced_codec: Some(CodecId::Rans), allow_codec_ids: None, chunk_size: Some(512), // small → multiple chunks per plane + ..Default::default() }, ) .expect("compress succeeds"); @@ -1684,6 +1713,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .unwrap(); @@ -1771,6 +1801,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .unwrap(); @@ -1853,6 +1884,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .unwrap(); diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 4b1e920..6b43609 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -15,15 +15,17 @@ use crate::chain::wire_compact::{bind_deferred_source, has_deferred_source}; use crate::codec::CodecId; use crate::error::PtwmCoreError; use crate::extension::{CanonicalId, ExtensionTable, ExtensionTableBuilder}; -use crate::header::{ContainerMagic, HEADER_LEN, Header, MAGIC_PTWC, MAGIC_PTWM, MAGIC_PTWX, sniff}; -use crate::tensor_record_compact::{ - RecordContext, parse_compact_record, plane_codec_ids, record_binding, write_compact_record, +use crate::header::{ + ContainerMagic, HEADER_LEN, Header, MAGIC_PTWC, MAGIC_PTWM, MAGIC_PTWX, sniff, }; use crate::index::{ EOF_SENTINEL, IndexEntry, name_hash, parse_index, verify_eof_sentinel, write_index, }; use crate::prelude::{PreludeEntry, parse_prelude, write_prelude}; use crate::tensor_record::{TensorRecord, write_tensor_record_with_builtin_version}; +use crate::tensor_record_compact::{ + RecordContext, parse_compact_record, plane_codec_ids, record_binding, write_compact_record, +}; // ── Chain registry ──────────────────────────────────────────────────────────── @@ -294,9 +296,7 @@ fn compact_registry_bytes( reg.id_for_entry_name(op.builtin_name()) .and_then(|id| index.get(&id).copied()) .ok_or_else(|| { - PtwmCoreError::InvalidContainer(format!( - "op {op:?} is not in the extension table" - )) + PtwmCoreError::InvalidContainer(format!("op {op:?} is not in the extension table")) }) }) } @@ -370,8 +370,9 @@ impl RecordEncoder { codec_ids: &self.codec_ids, registry_bound: &self.registry_bound, }; - write_compact_record(rec, &ctx, out) - .map_err(|e| PtwmCoreError::InvalidContainer(format!("tensor '{}': {e}", rec.name))) + write_compact_record(rec, &ctx, out).map_err(|e| { + PtwmCoreError::InvalidContainer(format!("tensor '{}': {e}", rec.name)) + }) } } } @@ -818,15 +819,9 @@ impl ContainerWriter { let et = crate::extension::compact_table::compact_table_bytes(&self.extension_table)?; let mut prelude = Vec::new(); write_prelude(&self.prelude_entries, &mut prelude)?; - let registry = if self.chain_registry_bytes.is_empty() { - crate::ptwc::write_registry(&self.chain_registry, |op| { - Err(PtwmCoreError::InvalidContainer(format!( - "op {op:?} is not in the extension table" - ))) - })? - } else { - self.chain_registry_bytes.clone() - }; + // Encode the registry against the final table: the table may have + // been replaced after the registry was set. + let registry = compact_registry_bytes(&self.chain_registry, &self.extension_table)?; let header = crate::ptwc::header_bytes(&et, &prelude, ®istry); for part in [&header, &et, &prelude, ®istry] { self.writer @@ -1190,8 +1185,7 @@ fn check_extensions( // short-circuit anyway, but skip the lookup entirely. continue; } - let Some((pubkey, contrib_idx)) = - lookup_author_pubkey(&entry.canonical_id, &installed) + let Some((pubkey, contrib_idx)) = lookup_author_pubkey(&entry.canonical_id, &installed) else { // No matching installed manifest — MissingFlavor already // accounts for this entry. If skip_missing is true we keep @@ -1568,6 +1562,14 @@ impl<'a> ContainerReader<'a> { Ok(rec) } + /// Every tensor record with its name, in body (file) order. + pub fn tensor_records(&self) -> Result, PtwmCoreError> { + self.entries_in_body_order() + .into_iter() + .map(|e| self.entry_record(e)) + .collect() + } + /// Index entries in body (file) order. fn entries_in_body_order(&self) -> Vec<&IndexEntry> { let mut entries: Vec<&IndexEntry> = self.index.iter().collect(); @@ -1882,8 +1884,7 @@ impl<'a> ContainerReader<'a> { "chunk table decodes to more than {plane_bound} bytes (or, in \ a compact container, not to the plane's {} bytes) for a \ {}-byte tensor", - descriptor.length_bytes, - rec.orig_size + descriptor.length_bytes, rec.orig_size )) })? as usize; let decoded_chunks: Vec> = chunks @@ -3023,6 +3024,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -3076,6 +3078,7 @@ mod tests { forced_codec: Some(CodecId::Rans), allow_codec_ids: None, chunk_size: Some(256), + ..Default::default() }, ) .expect("compress succeeds"); @@ -3126,6 +3129,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -3269,6 +3273,7 @@ mod tests { forced_codec: Some(crate::codec::CodecId::Identity), allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -3670,6 +3675,7 @@ label = "test-author" forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -3734,6 +3740,7 @@ label = "test-author" forced_codec: Some(CodecId::Identity), allow_codec_ids: None, chunk_size, + ..Default::default() }, ) .unwrap(); diff --git a/crates/ptwm-core/src/extension/compact_table.rs b/crates/ptwm-core/src/extension/compact_table.rs index 7fc210b..ec43ddb 100644 --- a/crates/ptwm-core/src/extension/compact_table.rs +++ b/crates/ptwm-core/src/extension/compact_table.rs @@ -120,25 +120,17 @@ pub fn classify(entry: &ExtensionTableEntry) -> Result, PtwmCor /// Append the compact encoding of `table` to `out`. On error `out` is left /// unchanged. -pub fn write_compact_table( - table: &ExtensionTable, - out: &mut Vec, -) -> Result<(), PtwmCoreError> { +pub fn write_compact_table(table: &ExtensionTable, out: &mut Vec) -> Result<(), PtwmCoreError> { let n = table.entries.len() as u64; if n > MAX_COMPACT_ENTRIES { - return Err(err(format_args!( - "{n} entries (max {MAX_COMPACT_ENTRIES})" - ))); + return Err(err(format_args!("{n} entries (max {MAX_COMPACT_ENTRIES})"))); } let mut buf = Vec::with_capacity(1 + 3 * table.entries.len()); write_varint(&mut buf, n); let mut seen = HashSet::with_capacity(table.entries.len()); for entry in &table.entries { if !seen.insert(entry.canonical_id) { - return Err(err(format_args!( - "duplicate entry {:?}", - entry.human_label - ))); + return Err(err(format_args!("duplicate entry {:?}", entry.human_label))); } match classify(entry)? { CompactEntry::Op(op) => { @@ -174,7 +166,10 @@ pub fn compact_table_bytes(table: &ExtensionTable) -> Result, PtwmCoreEr } /// The built-in entry for a wire id or name, or an error naming `what`. -fn builtin_by_name(name: &str, what: impl std::fmt::Display) -> Result { +fn builtin_by_name( + name: &str, + what: impl std::fmt::Display, +) -> Result { builtin_registry() .entry_by_name(name) .cloned() @@ -187,7 +182,8 @@ fn read_entry(r: &mut Cursor<'_>) -> Result match tag { TAG_OP => { let id = r.varint_max(u64::from(u16::MAX), "op id")? as u16; - let op = OpId::from_u16(id).map_err(|_| r.err(format_args!("unknown op id {id:#06x}")))?; + let op = + OpId::from_u16(id).map_err(|_| r.err(format_args!("unknown op id {id:#06x}")))?; builtin_by_name(op.builtin_name(), format_args!("op {op:?}")) } TAG_CODEC => { @@ -205,7 +201,8 @@ fn read_entry(r: &mut Cursor<'_>) -> Result return Err(r.err("empty built-in name")); } let bytes = r.take(len as usize)?; - let name = std::str::from_utf8(bytes).map_err(|_| r.err("built-in name is not UTF-8"))?; + let name = + std::str::from_utf8(bytes).map_err(|_| r.err("built-in name is not UTF-8"))?; let reg = builtin_registry(); let id = reg .id_for_entry_name(name) @@ -368,9 +365,15 @@ mod tests { rejects(&[1, 0x04, 0], "unknown entry tag"); rejects(&[1, TAG_OP, 0x7F], "unknown op id"); rejects(&[1, TAG_CODEC, 0x7F], "unknown codec id"); - rejects(&[1, TAG_NAMED, 3, b'f', b'o', b'o'], "unknown built-in name"); + rejects( + &[1, TAG_NAMED, 3, b'f', b'o', b'o'], + "unknown built-in name", + ); // huffman_nibble has a codec id but no built-in entry. - rejects(&[1, TAG_CODEC, CodecId::HuffmanNibble.as_u16() as u8], "no built-in entry"); + rejects( + &[1, TAG_CODEC, CodecId::HuffmanNibble.as_u16() as u8], + "no built-in entry", + ); // A named entry for a built-in that has a wire id. rejects(&[1, TAG_NAMED, 4, b'z', b's', b't', b'd'], "non-canonical"); // Duplicates. diff --git a/crates/ptwm-core/src/ptwc.rs b/crates/ptwm-core/src/ptwc.rs index 17f21f7..61b043e 100644 --- a/crates/ptwm-core/src/ptwc.rs +++ b/crates/ptwm-core/src/ptwc.rs @@ -194,7 +194,11 @@ pub fn read_registry(buf: &[u8], table: &ExtensionTable) -> Result Ok(op), _ => Err(err(format_args!( @@ -292,3 +296,146 @@ pub fn footer_bytes(index_len: u64) -> [u8; FOOTER_LEN] { f[8..].copy_from_slice(&EOF_SENTINEL); f } + +#[cfg(test)] +mod tests { + use super::*; + + /// A file of the given sections, body and index (index bytes as given). + fn file(et: &[u8], prelude: &[u8], registry: &[u8], body: &[u8], index: &[u8]) -> Vec { + let mut f = header_bytes(et, prelude, registry); + for part in [et, prelude, registry, body, index] { + f.extend_from_slice(part); + } + f.extend_from_slice(&footer_bytes(index.len() as u64)); + f + } + + fn rejects(buf: &[u8], needle: &str) { + let msg = locate_sections(buf).unwrap_err().to_string(); + assert!(msg.contains(needle), "{needle:?} not in {msg:?}"); + } + + #[test] + fn sections_are_located() { + let index = write_index(&[("a".into(), 3)]).unwrap(); + let f = file(&[0], &[], &[0], &[1, 2, 3], &index); + let s = locate_sections(&f).unwrap(); + // magic, hflags, three one-byte lengths, hash. + assert_eq!(s.header, 0..9 + 2 + 3 + 8); + assert_eq!(s.extension_table.len(), 1); + assert_eq!(s.prelude.len(), 0); + assert_eq!(s.registry.len(), 1); + assert_eq!(&f[s.body.clone()], &[1, 2, 3]); + assert_eq!(&f[s.index.clone()], &index[..]); + let entries = read_index(&f[s.index], s.body.clone()).unwrap(); + assert_eq!(entries[0].tensor_offset, s.body.start as u64); + assert_eq!(entries[0].name.as_deref(), Some("a")); + } + + #[test] + fn header_rejections() { + let f = file(&[0], &[], &[0], &[], &[0]); + // Every hflags bit is reserved. + for bit in 0..16 { + let mut b = f.clone(); + let flags = 1u16 << bit; + b[9..11].copy_from_slice(&flags.to_le_bytes()); + rejects(&b, "unknown hflags"); + } + // Any flipped bit of the hashed bytes fails the hash. + let mut b = f.clone(); + b[11 + 3 + 8] ^= 1; // the extension table byte + rejects(&b, "meta_xxh64 mismatch"); + let mut b = f.clone(); + b[11 + 3] ^= 1; // the hash itself + rejects(&b, "meta_xxh64 mismatch"); + // Overlong section length. + let mut b = MAGIC_PTWC.to_vec(); + b.extend_from_slice(&[0, 0, 0x81, 0x00, 0, 0]); + b.extend_from_slice(&[0; 8 + 16]); + rejects(&b, "non-minimal"); + // A section that runs past the footer. + let mut b = MAGIC_PTWC.to_vec(); + b.extend_from_slice(&[0, 0, 0x7F, 0, 0]); + b.extend_from_slice(&[0; 8]); + b.extend_from_slice(&footer_bytes(0)); + rejects(&b, "exceeds the file"); + // index_len beyond the file, and an index that overlaps the registry. + let mut b = f.clone(); + let n = b.len(); + b[n - 16..n - 8].copy_from_slice(&u64::MAX.to_le_bytes()); + rejects(&b, "index_len exceeds the file"); + let mut b = f.clone(); + b[n - 16..n - 8].copy_from_slice(&2u64.to_le_bytes()); + rejects(&b, "index overlaps the registry"); + // Wrong magic or too short. + rejects(&f[1..], "bad magic"); + rejects(&MAGIC_PTWC, "bad magic"); + } + + fn index_rejects(index: &[u8], body_len: usize, needle: &str) { + let msg = read_index(index, 100..100 + body_len) + .unwrap_err() + .to_string(); + assert!(msg.contains(needle), "{needle:?} not in {msg:?}"); + } + + #[test] + fn index_rejections() { + let ok = write_index(&[("a".into(), 3), ("b".into(), 4)]).unwrap(); + assert_eq!(read_index(&ok, 100..107).unwrap().len(), 2); + // A gap after the last record, and records past the body. + index_rejects(&ok, 8, "records cover 7 of 8 body bytes"); + index_rejects(&ok, 6, "runs past the body"); + // Duplicate and empty names, short records, trailing bytes. + let dup = write_index(&[("a".into(), 3), ("a".into(), 4)]).unwrap(); + index_rejects(&dup, 7, "duplicate tensor name"); + index_rejects(&[1, 0, 3], 3, "empty tensor name"); + index_rejects(&[1, 1, b'a', 2], 2, "record of 2 bytes"); + let mut trailing = ok.clone(); + trailing.push(0); + index_rejects(&trailing, 7, "trailing"); + // Bad UTF-8 and a huge count. + index_rejects(&[1, 1, 0xFF, 3], 3, "not UTF-8"); + index_rejects( + &[0xFF, 0xFF, 0xFF, 0xFF, 0x0F], + 0, + "exceeds the index bytes", + ); + // Names must fit their length limit on write. + assert!(write_index(&[(String::new(), 3)]).is_err()); + assert!(write_index(&[("x".repeat(MAX_NAME_LEN + 1), 3)]).is_err()); + assert!(write_index(&[("a".into(), 2)]).is_err()); + } + + #[test] + fn index_is_sorted_for_lookup_and_keeps_body_offsets() { + let names = ["w", "b", "q", "a"]; + let entries: Vec<(String, u64)> = names.iter().map(|n| (n.to_string(), 3)).collect(); + let index = write_index(&entries).unwrap(); + let read = read_index(&index, 10..22).unwrap(); + assert!(read.windows(2).all(|w| w[0].name_hash <= w[1].name_hash)); + for (i, n) in names.iter().enumerate() { + let e = read.iter().find(|e| e.name.as_deref() == Some(*n)).unwrap(); + assert_eq!(e.tensor_offset, 10 + 3 * i as u64); + assert_eq!(e.name_hash, name_hash(n)); + } + } + + #[test] + fn registry_rejections() { + let table = ExtensionTable::empty(); + // Empty registry. + assert!(read_registry(&[0], &table).unwrap().chains.is_empty()); + // Too many chains for the bytes, above the limit, trailing bytes. + let msg = read_registry(&[10, 0, 0], &table).unwrap_err().to_string(); + assert!(msg.contains("exceeds the registry bytes"), "{msg}"); + let mut huge = Vec::new(); + write_varint(&mut huge, MAX_REGISTRY_CHAINS as u64 + 1); + let msg = read_registry(&huge, &table).unwrap_err().to_string(); + assert!(msg.contains("chain count"), "{msg}"); + let msg = read_registry(&[0, 0], &table).unwrap_err().to_string(); + assert!(msg.contains("trailing"), "{msg}"); + } +} diff --git a/crates/ptwm-core/src/tensor_record_compact.rs b/crates/ptwm-core/src/tensor_record_compact.rs index 530ed87..693a997 100644 --- a/crates/ptwm-core/src/tensor_record_compact.rs +++ b/crates/ptwm-core/src/tensor_record_compact.rs @@ -72,9 +72,7 @@ use crate::codec::{CodecId, StateSource}; use crate::dtype::canonical_dtype_name; use crate::error::PtwmCoreError; use crate::extension::compact_table::{read_compact_table, write_compact_table}; -use crate::extension::{ - BuiltinKind, ExtensionTable, Kind, builtin_registry, dispatch_builtin, -}; +use crate::extension::{BuiltinKind, ExtensionTable, Kind, builtin_registry, dispatch_builtin}; use crate::layout::PlaneLayout; use crate::metadata::{decode_shape_metadata, encode_shape_metadata}; use crate::plane_record::{ChunkEntry, ExternalStateRef, PlaneRecord}; @@ -168,7 +166,8 @@ fn classify_meta(meta: Option<&[u8]>, dtype_code: u16) -> Meta<'_> { Ok((shape, name)) if shape.len() <= MAX_NDIM && encode_shape_metadata(&shape, &name) == m => { - let stored_name = (canonical_dtype_name(dtype_code) != Some(name.as_str())).then_some(name); + let stored_name = + (canonical_dtype_name(dtype_code) != Some(name.as_str())).then_some(name); Meta::Shape { shape, stored_name } } _ => Meta::Raw(m), @@ -227,7 +226,10 @@ fn write_unchecked( ))); } if rec.terminals.len() as u64 > MAX_PLANES { - return Err(err(format_args!("{} planes (max {MAX_PLANES})", rec.terminals.len()))); + return Err(err(format_args!( + "{} planes (max {MAX_PLANES})", + rec.terminals.len() + ))); } if rec.dependencies.len() > u16::MAX as usize { return Err(err(format_args!( @@ -510,7 +512,9 @@ pub fn parse_compact_record( let n = ctx.registry_bound.len() as u64; let id = r.varint()?; if id >= n { - return Err(r.err(format_args!("chain id {id} is not in the registry ({n} chains)"))); + return Err(r.err(format_args!( + "chain id {id} is not in the registry ({n} chains)" + ))); } Some(id as usize) }; @@ -532,8 +536,7 @@ pub fn parse_compact_record( let canonical = canonical_dtype_name(dtype_code); let name: String = if tflags & TF_DTYPE_NAME != 0 { let bytes = r.len_bytes()?; - let name = - std::str::from_utf8(bytes).map_err(|_| r.err("dtype name is not UTF-8"))?; + let name = std::str::from_utf8(bytes).map_err(|_| r.err("dtype name is not UTF-8"))?; if canonical == Some(name) { return Err(r.err("stored dtype name equals the canonical name (non-canonical)")); } @@ -578,8 +581,8 @@ pub fn parse_compact_record( let inline_chain = if inline { let et_bytes = r.len_bytes()?; - let table = read_compact_table(et_bytes) - .map_err(|e| r.err(format_args!("inline chain: {e}")))?; + let table = + read_compact_table(et_bytes).map_err(|e| r.err(format_args!("inline chain: {e}")))?; let ops: Vec = table .entries .iter() @@ -759,7 +762,9 @@ fn read_plane(r: &mut Cursor<'_>, ctx: &RecordContext<'_>) -> Result, ctx: &RecordContext<'_>) -> Result] = &[None, Some(CodecId::Zstd), Some(CodecId::Identity)]; + /// Chain 0 is explicit, chain 1 has a bound Source. + const BOUND: &[bool] = &[false, true]; + + fn ctx() -> RecordContext<'static> { + RecordContext { + codec_ids: CODEC_IDS, + registry_bound: BOUND, + } + } + + fn plane(payload: &[u8]) -> PlaneRecord { + PlaneRecord { + role: PlaneRole::Value, + codec_id: CodecId::Zstd, + codec_table_idx: 1, + state_source: StateSource::None, + state_version: 0, + state_info: 0, + payload_len: payload.len() as u64, + crc32: None, + chunk_table: None, + external_state: None, + inline_state_bytes: Vec::new(), + payload_bytes: payload.to_vec(), + layout: PlaneLayout::Flat, + } + } + + fn record() -> TensorRecord { + TensorRecord { + chain_ref: 0, + dtype_code: 1, + input_format: 2, + flags: 0, + orig_size: 16, + name: String::new(), + payload_hash: None, + inline_chain: None, + dependencies: Vec::new(), + tensor_metadata: None, + terminals: vec![plane(&[1, 2, 3, 4])], + } + } + + /// A record with every optional field. + fn rich_record() -> TensorRecord { + let mut inline = plane(&[9; 7]); + inline.state_source = StateSource::Inline; + inline.inline_state_bytes = vec![5; 3]; + inline.state_version = 2; + inline.layout = PlaneLayout::rows(32).unwrap(); + inline.crc32 = Some(0xDEAD_BEEF); + let mut shared = plane(&[1; 20]); + shared.state_source = StateSource::Shared; + shared.state_info = 300; + shared.role = PlaneRole::Scale; + shared.chunk_table = Some(vec![ + ChunkEntry { + offset_in_payload: 0, + decoded_size: 70_000, + }, + ChunkEntry { + offset_in_payload: 10, + decoded_size: 5, + }, + ]); + let mut external = plane(&[]); + external.codec_id = CodecId::Identity; + external.codec_table_idx = 2; + external.state_source = StateSource::External; + external.external_state = Some(ExternalStateRef { + ref_kind: 3, + expected_hash: Some([4; 32]), + ref_bytes: b"ref".to_vec(), + state_format_version: 7, + }); + TensorRecord { + flags: TENSOR_FLAG_PAYLOAD_HASH | TENSOR_FLAG_DEPENDENCIES | TENSOR_FLAG_METADATA, + payload_hash: Some(0x0123_4567_89AB_CDEF), + dependencies: vec![ + Dependency { + ref_kind: RefKind::ExternalSafetensors, + expected_hash: Some([1; 32]), + ref_bytes: Vec::new(), + }, + Dependency { + ref_kind: RefKind::LocalTensor, + expected_hash: None, + ref_bytes: b"other".to_vec(), + }, + ], + // Not the shape map: stored as raw metadata. + tensor_metadata: Some(vec![0xA0]), + terminals: vec![inline, shared, external], + ..record() + } + } + + fn encode(rec: &TensorRecord) -> Vec { + let mut out = Vec::new(); + write_compact_record(rec, &ctx(), &mut out).unwrap(); + out + } + + fn round_trip(rec: &TensorRecord) -> Vec { + let bytes = encode(rec); + let back = parse_compact_record(&bytes, &ctx()).unwrap(); + assert_eq!(&back, rec); + assert_eq!(encode(&back), bytes); + bytes + } + + fn rejects(bytes: &[u8], needle: &str) { + let msg = parse_compact_record(bytes, &ctx()).unwrap_err().to_string(); + assert!(msg.contains(needle), "{needle:?} not in {msg:?}"); + } + + fn fp32_chain(source_params: Vec) -> Chain { + Chain { + nodes: vec![ + ChainNode { + op: OpId::Source, + params: source_params, + }, + ChainNode { + op: OpId::ByteSplit, + params: vec![4], + }, + ], + edges: vec![ChainEdge { + src_node: 0, + src_output_idx: 0, + dst_node: 1, + dst_input_idx: 0, + role_override: None, + vendor_bytes: Vec::new(), + }], + terminals: (0..4) + .map(|i| TerminalRef { + node_idx: 1, + output_idx: i, + role: Role::Raw, + }) + .collect(), + } + } + + #[test] + fn minimal_record_is_a_few_bytes() { + let bytes = round_trip(&record()); + // tflags, chain, dtype, format, size, n_planes; pflags, role, codec, + // payload_len, payload. + assert_eq!(bytes, [0, 0, 1, 2, 16, 1, 0, 1, 1, 4, 1, 2, 3, 4]); + } + + #[test] + fn every_optional_field_round_trips() { + let bytes = round_trip(&rich_record()); + assert_eq!(bytes[0], TF_PAYLOAD_HASH | TF_DEPS | TF_RAW_META, "tflags"); + } + + #[test] + fn shape_metadata_is_stored_as_a_shape() { + let mut rec = record(); + rec.flags = TENSOR_FLAG_METADATA; + rec.tensor_metadata = Some(encode_shape_metadata(&[2, 2], "float32")); + let bytes = round_trip(&rec); + assert_eq!(bytes[0], TF_SHAPE); + // A dtype name other than the canonical one is stored. + rec.tensor_metadata = Some(encode_shape_metadata(&[2, 2], "my_float")); + let bytes = round_trip(&rec); + assert_eq!(bytes[0], TF_SHAPE | TF_DTYPE_NAME); + // 0-d tensors keep their empty shape. + rec.tensor_metadata = Some(encode_shape_metadata(&[], "float32")); + round_trip(&rec); + // More than MAX_NDIM dimensions fall back to raw metadata. + rec.tensor_metadata = Some(encode_shape_metadata(&[1; MAX_NDIM + 1], "float32")); + assert_eq!(round_trip(&rec)[0], TF_RAW_META); + } + + #[test] + fn inline_chain_binds_its_source_to_the_shape() { + let params = bound_source_params(&[2, 2], 1).unwrap(); + let mut rec = record(); + rec.chain_ref = CHAIN_REF_INLINE_BIT; + rec.inline_chain = Some(fp32_chain(params.clone())); + rec.flags = TENSOR_FLAG_METADATA; + rec.tensor_metadata = Some(encode_shape_metadata(&[2, 2], "float32")); + let bound = round_trip(&rec); + // Without shape metadata the Source params are stored. + rec.flags = 0; + rec.tensor_metadata = None; + let explicit = round_trip(&rec); + assert!(explicit.len() > bound.len() + params.len() - 4); + } + + #[test] + fn bound_registry_chain_needs_a_binding_shape() { + let mut rec = record(); + rec.chain_ref = 1; + let mut out = Vec::new(); + let msg = write_compact_record(&rec, &ctx(), &mut out) + .unwrap_err() + .to_string(); + assert!(msg.contains("bound Source"), "{msg}"); + assert!(out.is_empty()); + rec.flags = TENSOR_FLAG_METADATA; + rec.tensor_metadata = Some(encode_shape_metadata(&[2, 2], "float32")); + let mut bytes = round_trip(&rec); + // Same record bytes without the shape: rejected on read. + assert_eq!(bytes[0], TF_SHAPE); + bytes[0] = 0; + let shape_len = 1 + 1 + 2; + bytes.drain(5..5 + shape_len); + rejects(&bytes, "bound Source"); + } + + #[test] + fn writer_rejects_inconsistent_records() { + let check = |rec: TensorRecord, needle: &str| { + let mut out = Vec::new(); + let msg = write_compact_record(&rec, &ctx(), &mut out) + .unwrap_err() + .to_string(); + assert!(msg.contains(needle), "{needle:?} not in {msg:?}"); + assert!(out.is_empty()); + }; + let mut rec = record(); + rec.terminals[0].codec_id = CodecId::Rans; + check(rec, "differs from Zstd"); + let mut rec = record(); + rec.terminals[0].codec_table_idx = 0; + check(rec, "not a plane codec"); + let mut rec = record(); + rec.terminals[0].codec_table_idx = 9; + check(rec, "outside the Extension Table"); + let mut rec = record(); + rec.terminals[0].payload_len = 3; + check(rec, "payload_len"); + let mut rec = record(); + rec.chain_ref = 5; + check(rec, "not in the registry"); + let mut rec = record(); + rec.chain_ref = CHAIN_REF_INLINE_BIT; + check(rec, "disagree"); + let mut rec = record(); + rec.terminals[0].inline_state_bytes = vec![1]; + check(rec, "inline state bytes"); + } + + #[test] + fn reader_rejects_reserved_bits_and_non_canonical_fields() { + let base = encode(&record()); + // Reserved tflags bits. + for bit in [0x40u8, 0x80] { + let mut b = base.clone(); + b[0] |= bit; + rejects(&b, "unknown tflags"); + } + let mut b = base.clone(); + b[0] = TF_DTYPE_NAME; + rejects(&b, "TF_DTYPE_NAME without TF_SHAPE"); + let mut b = base.clone(); + b[0] = TF_SHAPE | TF_RAW_META; + rejects(&b, "TF_RAW_META with TF_SHAPE"); + // Reserved pflags bits (pflags at offset 6). + for bit in [0x40u8, 0x80] { + let mut b = base.clone(); + b[6] |= bit; + rejects(&b, "unknown pflags"); + } + // Overlong varint (dtype_code 1 as two bytes). + let mut b = base.clone(); + b.splice(2..3, [0x81, 0x00]); + rejects(&b, "non-minimal"); + // Unknown role, codec_ref to a transform or outside the table. + let mut b = base.clone(); + b[7] = 0x7E; + rejects(&b, "unknown plane role"); + let mut b = base.clone(); + b[8] = 0; + rejects(&b, "not a plane codec"); + let mut b = base.clone(); + b[8] = 3; + rejects(&b, "outside the Extension Table"); + // Chain id outside the registry. + let mut b = base.clone(); + b[1] = 2; + rejects(&b, "not in the registry"); + // Trailing and missing bytes. + let mut b = base.clone(); + b.push(0); + rejects(&b, "trailing"); + rejects(&base[..base.len() - 1], ""); + // STATE_EXTRA with zero values. + let mut b = base.clone(); + b[6] |= PF_STATE_EXTRA; + b.splice(9..9, [0, 0]); + rejects(&b, "PF_STATE_EXTRA with zero"); + // Rows layout with row_len 0 (row_len follows payload_len). + let mut b = base.clone(); + b[6] |= PF_LAYOUT_ROWS; + b.insert(10, 0); + rejects(&b, "row_len 0"); + } + + #[test] + fn reader_rejects_non_canonical_metadata() { + let mut rec = record(); + rec.flags = TENSOR_FLAG_METADATA; + rec.tensor_metadata = Some(encode_shape_metadata(&[2], "my_float")); + let b = encode(&rec); + // The stored name (after ndim and dim at offsets 5-6) rewritten to + // the canonical name. + let mut c = b[..7].to_vec(); + c.push(7); + c.extend_from_slice(b"float32"); + c.extend_from_slice(&b[7 + 1 + 8..]); + rejects(&c, "equals the canonical name"); + // Shape metadata stored as raw bytes. + let meta = encode_shape_metadata(&[2], "float32"); + let mut c = b[..5].to_vec(); + c[0] = TF_RAW_META; + c.push(meta.len() as u8); + c.extend_from_slice(&meta); + c.extend_from_slice(&b[7 + 1 + 8..]); + rejects(&c, "raw metadata that TF_SHAPE stores"); + // A shape for a dtype without a canonical name needs a stored name. + let mut rec = record(); + rec.dtype_code = 200; + let b = encode(&rec); + let mut c = b[..5].to_vec(); + c[0] = TF_SHAPE; + c.extend_from_slice(&[1, 2]); + c.extend_from_slice(&b[5..]); + rejects(&c, "no canonical name"); + } + + #[test] + fn reader_rejects_a_wrong_inline_table_order() { + let mut rec = record(); + rec.chain_ref = CHAIN_REF_INLINE_BIT; + rec.inline_chain = Some(fp32_chain(bound_source_params(&[4], 1).unwrap())); + let b = encode(&rec); + // Inline table after orig_size (offset 4): et_len, then the table + // [Source, ByteSplit]. Rewrite the chain with the node op refs + // swapped: the table then no longer lists the ops in first-use + // order. + let at = 4; + let et_len = usize::from(b[at]); + let chain_at = at + 1 + et_len; + let chain_len = usize::from(b[chain_at]); + let swapped = Chain { + nodes: vec![ + ChainNode { + op: OpId::ByteSplit, + params: vec![1, 4, 0, 0, 0, 3, 0], + }, + ChainNode { + op: OpId::Source, + params: vec![4], + }, + ], + ..fp32_chain(Vec::new()) + }; + let mut chain = Vec::new(); + write_chain_compact( + &swapped, + None, + |op| Ok(if op == OpId::Source { 0 } else { 1 }), + &mut chain, + ) + .unwrap(); + let mut c = b[..chain_at].to_vec(); + write_varint(&mut c, chain.len() as u64); + c.extend_from_slice(&chain); + c.extend_from_slice(&b[chain_at + 1 + chain_len..]); + rejects(&c, "first-use order"); + } + + #[test] + fn reader_rejects_bad_dependencies() { + let mut rec = record(); + rec.flags = TENSOR_FLAG_DEPENDENCIES; + rec.dependencies = vec![Dependency { + ref_kind: RefKind::LocalTensor, + expected_hash: None, + ref_bytes: vec![1], + }]; + let b = encode(&rec); + // deps at offset 5: count, kind, flags, len, byte. + let mut c = b.clone(); + c[5] = 0; + c.drain(6..10); + rejects(&c, "TF_DEPS with no dependencies"); + let mut c = b.clone(); + c[7] = 0x02; + rejects(&c, "unknown dependency flag"); + let mut c = b.clone(); + c[6] = 0x7F; + rejects(&c, "unknown ref_kind"); + } + + #[test] + fn huge_counts_fail_before_allocating() { + let huge = |prefix: &[u8], value: u64, suffix: &[u8]| { + let mut b = prefix.to_vec(); + write_varint(&mut b, value); + b.extend_from_slice(suffix); + b + }; + // n_planes above 255 and above the remaining bytes. + rejects(&huge(&[0, 0, 1, 2, 16], 255, &[]), "plane count exceeds"); + rejects(&huge(&[0, 0, 1, 2, 16], 1 << 40, &[]), "exceeds 255"); + // ndim above the limit. + rejects(&huge(&[TF_SHAPE, 0, 1, 2, 16], 1 << 30, &[]), "ndim"); + // Dependency count above the remaining bytes. + rejects( + &huge(&[TF_DEPS, 0, 1, 2, 16], 60_000, &[]), + "dependency count", + ); + // Payload length above the remaining bytes. + rejects(&huge(&[0, 0, 1, 2, 16, 1, 0, 1, 1], u64::MAX >> 1, &[]), ""); + // Chunk count above the payload length. + rejects( + &huge(&[0, 0, 1, 2, 16, 1, PF_CHUNKS, 1, 1, 4], 5, &[0; 16]), + "chunk count", + ); + // Raw metadata length above the remaining bytes. + rejects(&huge(&[TF_RAW_META, 0, 1, 2, 16], u32::MAX.into(), &[]), ""); + } +} diff --git a/crates/ptwm-core/src/transcode/mod.rs b/crates/ptwm-core/src/transcode/mod.rs index 9507d4d..7d0d7ba 100644 --- a/crates/ptwm-core/src/transcode/mod.rs +++ b/crates/ptwm-core/src/transcode/mod.rs @@ -307,6 +307,11 @@ fn extension_table_bytes<'a>( /// and loose binary members. pub fn explode(blob: &[u8]) -> Result { let reader = ContainerReader::open_partial(blob)?; + if reader.layout != crate::container::ContainerLayout::Legacy { + return Err(err( + "compact (\\x89PTWC) containers are not supported by the transcoder yet", + )); + } let embedded_wasm = reader.header.flags & FLAG_PTWX != 0; if embedded_wasm { return Err(err( diff --git a/crates/ptwm-core/src/transforms/op.rs b/crates/ptwm-core/src/transforms/op.rs index e0547d8..9efaf4d 100644 --- a/crates/ptwm-core/src/transforms/op.rs +++ b/crates/ptwm-core/src/transforms/op.rs @@ -6,7 +6,7 @@ use std::sync::Arc; use crate::error::PtwmCoreError; use crate::types::descriptor::PlaneDescriptor; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] #[repr(u16)] pub enum OpId { Source = 0x0000, diff --git a/crates/ptwm-core/src/types/role.rs b/crates/ptwm-core/src/types/role.rs index 21bf9ba..1b0533c 100644 --- a/crates/ptwm-core/src/types/role.rs +++ b/crates/ptwm-core/src/types/role.rs @@ -3,7 +3,7 @@ use crate::error::PtwmCoreError; -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum ScaleFormat { E8M0, E4M3, @@ -36,7 +36,7 @@ impl ScaleFormat { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum ValueFormat { Fp4E2m1, IntN { bits: u8 }, @@ -75,7 +75,7 @@ impl ValueFormat { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum NibbleKind { Exponent, SignMantissa, @@ -102,7 +102,7 @@ impl NibbleKind { } } -#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] pub enum ResidualFormat { Xor, FloatDelta { dtype: u16 }, @@ -142,7 +142,7 @@ impl ResidualFormat { } } -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum Role { Scale { format: ScaleFormat }, Value { format: ValueFormat }, diff --git a/crates/ptwm-core/tests/decode_bomb.rs b/crates/ptwm-core/tests/decode_bomb.rs index 8dbc087..c369d4b 100644 --- a/crates/ptwm-core/tests/decode_bomb.rs +++ b/crates/ptwm-core/tests/decode_bomb.rs @@ -124,6 +124,7 @@ fn opts(chunk_size: Option) -> CompressorOptions { forced_codec: Some(CodecId::Zstd), allow_codec_ids: None, chunk_size, + ..Default::default() } } diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs index 8988c59..a886abe 100644 --- a/crates/ptwm-core/tests/legacy_fixtures.rs +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -458,6 +458,7 @@ fn main_options(chunk_size: Option) -> CompressorOptions { forced_codec: None, allow_codec_ids: None, chunk_size, + ..Default::default() } } diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs new file mode 100644 index 0000000..b2e7217 --- /dev/null +++ b/crates/ptwm-core/tests/ptwc.rs @@ -0,0 +1,725 @@ +//! Compact container (`\x89PTWC`) round trips, size bounds and old-reader +//! rejection. +//! +//! The legacy fixtures in `tests/fixtures/` hold every record feature the +//! writers produce (registry and inline chains, a shared prelude state, +//! chunked planes, a dependency, the Rows layout, bf16/fp16/fp32/fp8/int8). +//! Rewriting them in the compact layout and decoding both files covers the +//! compact record, plane, registry and Extension Table codecs on real data. +//! `compress_model` tests cover the compact writer path, bound Sources and +//! the overhead bounds of synthesis section 9 item 14. + +use std::collections::HashMap; +use std::io::Cursor; +use std::path::PathBuf; + +use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; +use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; +use ptwm_core::container::{ContainerLayout, ContainerReader, ContainerWriter}; +use ptwm_core::header::{HEADER_LEN, Header, MAGIC_PTWC}; +use ptwm_core::transforms::op::OpId; +use ptwm_core::types::{Role, ScaleFormat, ValueFormat}; + +const FIXTURES: &[&str] = &[ + "legacy_v1_1_0.ptwm", + "legacy_v1_0_0.ptwm", + "legacy_v1_0_0_native.ptwm", + "legacy_trimmed.ptwm", +]; + +fn load(name: &str) -> Vec { + let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("tests/fixtures") + .join(name); + std::fs::read(&path).unwrap_or_else(|e| panic!("read {}: {e}", path.display())) +} + +/// Rewrite the container `blob` (either layout) in the compact layout: +/// same Extension Table, registry, prelude and records, in body order. +fn to_compact(blob: &[u8]) -> Vec { + let reader = ContainerReader::open(blob).expect("open"); + let mut out = Vec::new(); + let mut w = + ContainerWriter::with_layout(Cursor::new(&mut out), ContainerLayout::Compact).unwrap(); + w.set_prelude(&reader.prelude).unwrap(); + w.set_compact_tables( + reader.extension_table.clone(), + reader.chain_registry.clone(), + ) + .unwrap(); + for rec in reader.tensor_records().unwrap() { + w.append_tensor(&rec).unwrap(); + } + w.finalize().unwrap(); + out +} + +fn decode_all(blob: &[u8]) -> HashMap> { + ContainerReader::open(blob) + .unwrap() + .decode_model() + .unwrap() + .into_iter() + .collect() +} + +#[test] +fn legacy_fixtures_round_trip_through_the_compact_layout() { + for fixture in FIXTURES { + let legacy = load(fixture); + let compact = to_compact(&legacy); + assert_eq!(compact[..9], MAGIC_PTWC, "{fixture}"); + assert!( + compact.len() < legacy.len(), + "{fixture}: compact {} B, legacy {} B", + compact.len(), + legacy.len() + ); + + let lr = ContainerReader::open(&legacy).unwrap(); + let cr = ContainerReader::open(&compact).unwrap(); + assert_eq!(cr.layout, ContainerLayout::Compact); + assert_eq!(lr.layout, ContainerLayout::Legacy); + assert_eq!(cr.prelude, lr.prelude, "{fixture}: prelude"); + assert_eq!(cr.chain_registry, lr.chain_registry, "{fixture}: registry"); + assert_eq!( + cr.extension_table, lr.extension_table, + "{fixture}: extension table" + ); + + // Every record parses back to the legacy record. + let lrecs = lr.tensor_records().unwrap(); + let crecs = cr.tensor_records().unwrap(); + assert_eq!(crecs, lrecs, "{fixture}: records"); + // Names come from the index in body order. + let names: Vec = lrecs.iter().map(|r| r.name.clone()).collect(); + assert_eq!(cr.tensor_names().unwrap(), names, "{fixture}: names"); + + // Per-tensor and whole-model decode agree with the legacy file. + let want = decode_all(&legacy); + assert_eq!(decode_all(&compact), want, "{fixture}: decode_model"); + for name in &names { + assert_eq!(&cr.decode_tensor(name).unwrap(), &want[name], "{fixture}"); + assert_eq!( + cr.tensor_shape(name).unwrap(), + lr.tensor_shape(name).unwrap(), + "{fixture}: {name}" + ); + } + assert!(cr.get_tensor_record_bytes("no.such.tensor").is_none()); + + // One encoding per content: rewriting the compact file reproduces it. + assert!( + to_compact(&compact) == compact, + "{fixture}: rewrite differs" + ); + } +} + +#[test] +fn every_byte_of_a_compact_fixture_is_checked() { + // Flipping any header, table, prelude, registry or index byte must fail + // the open (meta_xxh64 or a structural check), never decode quietly. + let compact = to_compact(&load("legacy_v1_1_0.ptwm")); + let reader = ContainerReader::open(&compact).unwrap(); + let body_start = reader.tensors_offset as usize; + let index_start = reader.header.tensor_index_offset as usize; + let checked = (0..body_start).chain(index_start..compact.len()); + for pos in checked { + let mut bad = compact.clone(); + bad[pos] ^= 0x01; + let opened = ContainerReader::open(&bad); + if let Ok(r) = opened { + // An index name byte can turn into another valid name; the file + // is then a different, still consistent container. + assert!( + pos >= index_start, + "byte {pos} of the metadata sections is not checked" + ); + assert_ne!( + r.tensor_names().unwrap(), + reader.tensor_names().unwrap(), + "byte {pos}" + ); + } + } +} + +#[test] +fn old_reader_rejects_compact_files_on_the_magic() { + let compact = to_compact(&load("legacy_trimmed.ptwm")); + assert!(compact.len() >= HEADER_LEN); + let err = Header::from_bytes(&compact[..HEADER_LEN]) + .unwrap_err() + .to_string(); + assert!(err.contains("bad magic"), "{err}"); +} + +#[test] +fn p4_source_slack_is_rejected_in_the_compact_layout() { + // Legacy files from v1.1 BYTE input declare a Source 4x the real size; + // the legacy reader tolerates it, the compact reader requires exact + // sizes. + let legacy = load("legacy_p4_byte_fp32.ptwm"); + let compact = to_compact(&legacy); + let reader = ContainerReader::open(&compact).unwrap(); + let err = reader.decode_model().unwrap_err().to_string(); + assert!(err.contains("Source declares"), "{err}"); +} + +// ── compress_model in the compact layout ─────────────────────────────────── + +struct Rng(u64); + +impl Rng { + 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 weight(&mut self) -> f32 { + let mut s = 0.0f32; + for _ in 0..4 { + s += (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32; + } + (s - 2.0) * 0.05 + } +} + +fn source(shape: &[u64], internal_dtype: u16) -> ChainNode { + let mut params = vec![shape.len() as u8]; + for d in shape { + params.extend_from_slice(&(*d as u32).to_le_bytes()); + } + params.extend_from_slice(&internal_dtype.to_le_bytes()); + ChainNode { + op: OpId::Source, + params, + } +} + +fn edge(src: u8, dst: u8, role_override: Option) -> ChainEdge { + ChainEdge { + src_node: src, + src_output_idx: 0, + dst_node: dst, + dst_input_idx: 0, + role_override, + vendor_bytes: vec![], + } +} + +fn term(node_idx: u8, output_idx: u8, role: Role) -> TerminalRef { + TerminalRef { + node_idx, + output_idx, + role, + } +} + +fn fp32_split_chain(shape: &[u64]) -> Chain { + Chain { + nodes: vec![ + source(shape, 0x0003), + ChainNode { + op: OpId::BitReorderIeee32, + params: vec![], + }, + ChainNode { + op: OpId::ByteSplit, + params: vec![4], + }, + ], + edges: vec![edge(0, 1, None), edge(1, 2, None)], + terminals: vec![ + term(2, 0, Role::MantissaByte { index: 0, of: 4 }), + term(2, 1, Role::MantissaByte { index: 1, of: 4 }), + term(2, 2, Role::MantissaByte { index: 2, of: 4 }), + term(2, 3, Role::ExponentByte), + ], + } +} + +fn bf16_split_chain(shape: &[u64]) -> Chain { + Chain { + nodes: vec![ + source(shape, 0x000F), + ChainNode { + op: OpId::BitReorderIeee16, + params: vec![], + }, + ChainNode { + op: OpId::ByteSplit, + params: vec![2], + }, + ], + edges: vec![edge(0, 1, None), edge(1, 2, None)], + terminals: vec![ + term(2, 0, Role::MantissaByte { index: 0, of: 2 }), + term(2, 1, Role::ExponentByte), + ], + } +} + +fn byte_chain(shape: &[u64], scale: bool) -> Chain { + let role = if scale { + Role::Scale { + format: ScaleFormat::E8M0, + } + } else { + Role::Value { + format: ValueFormat::IntN { bits: 8 }, + } + }; + Chain { + nodes: vec![ + source(shape, 0x0006), + ChainNode { + op: OpId::BytePassthrough, + params: vec![], + }, + ], + edges: vec![edge(0, 1, scale.then(|| role.clone()))], + terminals: vec![term(1, 0, role)], + } +} + +struct Tensor { + name: String, + dtype_code: u16, + dtype_name: String, + shape: Vec, + bytes: Vec, + chain: Chain, + delta: bool, +} + +fn fp32(name: &str, shape: Vec, rng: &mut Rng) -> Tensor { + let n: u64 = shape.iter().product(); + Tensor { + name: name.into(), + dtype_code: 1, + dtype_name: "float32".into(), + chain: fp32_split_chain(&shape), + bytes: (0..n).flat_map(|_| rng.weight().to_le_bytes()).collect(), + shape, + delta: false, + } +} + +/// One tensor per feature the compressor emits: dtypes, 0-d, a +/// non-canonical dtype name, a dependency, Rows scale planes that share a +/// prelude state. +fn feature_tensors() -> Vec { + let mut rng = Rng(0x5054_5743_5445_5354); + let mut out = vec![ + fp32("model.norm.weight", vec![4, 32], &mut rng), + fp32("model.scalar", vec![], &mut rng), + ]; + let mut custom = fp32("model.custom.weight", vec![16], &mut rng); + custom.dtype_name = "my_float".into(); + out.push(custom); + let mut delta = fp32("model.delta.weight", vec![8, 8], &mut rng); + delta.delta = true; + out.push(delta); + let shape = vec![8u64, 32]; + out.push(Tensor { + name: "model.layers.0.q.weight".into(), + dtype_code: 6, + dtype_name: "bfloat16".into(), + chain: bf16_split_chain(&shape), + bytes: (0..256) + .flat_map(|_| ((rng.weight().to_bits() >> 16) as u16).to_le_bytes()) + .collect(), + shape, + delta: false, + }); + for name in ["model.layers.0.q.scale", "model.layers.0.k.scale"] { + let mut bytes = Vec::with_capacity(2048); + for _ in 0..64 { + let mut prev: u8 = 0x80; + for _ in 0..32 { + prev = prev + .wrapping_add((rng.next_u64() % 3) as u8) + .wrapping_sub(1); + bytes.push(prev); + } + } + out.push(Tensor { + name: name.into(), + dtype_code: 17, + dtype_name: "int8".into(), + chain: byte_chain(&[2048], true), + bytes, + shape: vec![64, 32], + delta: false, + }); + } + out.push(Tensor { + name: "model.embed.int8".into(), + dtype_code: 17, + dtype_name: "int8".into(), + chain: byte_chain(&[256], false), + bytes: (0..256) + .map(|_| ((rng.next_u64() % 9) as i8 - 4) as u8) + .collect(), + shape: vec![256], + delta: false, + }); + out +} + +fn inputs(data: &[Tensor]) -> Vec> { + data.iter() + .map(|t| InputTensor { + name: t.name.clone(), + candidate_chains: vec![t.chain.clone()], + dtype_code: t.dtype_code, + input_format: 2, + raw_bytes: &t.bytes, + shape: Some(t.shape.clone()), + dtype_name: Some(t.dtype_name.clone()), + delta_reference_blake3: t.delta.then_some([7u8; 32]), + }) + .collect() +} + +fn options( + hash: bool, + crc: bool, + chunk_size: Option, + layout: ContainerLayout, +) -> CompressorOptions { + CompressorOptions { + method_hint: 1, + emit_payload_hash: hash, + emit_plane_crc: crc, + chunk_size, + container_layout: layout, + ..Default::default() + } +} + +fn compress(data: &[Tensor], opts: CompressorOptions) -> Vec { + let mut out = Cursor::new(Vec::new()); + compress_model(&mut out, &inputs(data), opts).expect("compress_model"); + out.into_inner() +} + +#[test] +fn compress_model_matrix_round_trips_in_both_layouts() { + let data = feature_tensors(); + for (hash, crc) in [(true, true), (false, false), (true, false), (false, true)] { + for chunk in [None, Some(256)] { + let mut decoded = Vec::new(); + for layout in [ContainerLayout::Legacy, ContainerLayout::Compact] { + let blob = compress(&data, options(hash, crc, chunk, layout)); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(reader.layout, layout); + for t in &data { + assert_eq!( + reader.decode_tensor(&t.name).unwrap(), + t.bytes, + "{layout:?} hash={hash} crc={crc} chunk={chunk:?}: {}", + t.name + ); + let (shape, name) = reader.tensor_shape(&t.name).unwrap().unwrap(); + assert_eq!((shape, name), (t.shape.clone(), t.dtype_name.clone())); + let rec = reader.tensor_record(&t.name).unwrap(); + assert_eq!(rec.name, t.name); + assert_eq!(rec.payload_hash.is_some(), hash); + assert_eq!(rec.dependencies.len(), usize::from(t.delta)); + for p in &rec.terminals { + assert_eq!(p.crc32.is_some(), crc); + assert_eq!(p.chunk_table.is_some(), chunk.is_some()); + } + } + let mut all = reader.decode_model().unwrap(); + all.sort(); + decoded.push(all); + if layout == ContainerLayout::Compact { + // parse then write is the identity. + assert!(to_compact(&blob) == blob, "rewrite differs"); + if chunk.is_none() { + assert!(!reader.prelude.is_empty(), "no shared state"); + } + } + } + assert_eq!(decoded[0], decoded[1]); + } + } +} + +#[test] +fn legacy_layout_stays_the_default() { + let data = feature_tensors(); + let default = compress(&data, options(true, true, None, ContainerLayout::default())); + let legacy = compress(&data, options(true, true, None, ContainerLayout::Legacy)); + assert!(default == legacy); + assert_eq!(default[..5], *b"\x89PTWM"); +} + +#[test] +fn bound_source_shares_one_chain_across_shapes() { + let mut rng = Rng(1); + let data: Vec = (1..=64) + .map(|i| fp32(&format!("t{i}"), vec![i, 2], &mut rng)) + .collect(); + let legacy = compress(&data, options(true, true, None, ContainerLayout::Legacy)); + let compact = compress(&data, options(true, true, None, ContainerLayout::Compact)); + let lr = ContainerReader::open(&legacy).unwrap(); + let cr = ContainerReader::open(&compact).unwrap(); + assert_eq!(lr.chain_registry.chains.len(), 64); + assert_eq!(cr.chain_registry.chains.len(), 1); + for t in &data { + assert_eq!(cr.decode_tensor(&t.name).unwrap(), t.bytes); + } + + // A record without shape metadata cannot use the bound chain: the writer + // refuses it. + let mut rec = cr.tensor_record("t1").unwrap(); + rec.tensor_metadata = None; + rec.flags &= !ptwm_core::tensor_record::TENSOR_FLAG_METADATA; + rec.name = "t_no_shape".into(); + let mut out = Vec::new(); + let mut w = + ContainerWriter::with_layout(Cursor::new(&mut out), ContainerLayout::Compact).unwrap(); + w.set_compact_tables(cr.extension_table.clone(), cr.chain_registry.clone()) + .unwrap(); + let err = w.append_tensor(&rec).unwrap_err().to_string(); + assert!(err.contains("bound Source"), "{err}"); +} + +#[test] +fn duplicate_names_are_rejected_and_lookup_is_exact() { + let compact = to_compact(&load("legacy_trimmed.ptwm")); + let reader = ContainerReader::open(&compact).unwrap(); + let rec = reader.tensor_records().unwrap().remove(0); + let mut out = Vec::new(); + let mut w = + ContainerWriter::with_layout(Cursor::new(&mut out), ContainerLayout::Compact).unwrap(); + w.set_prelude(&reader.prelude).unwrap(); + w.set_compact_tables( + reader.extension_table.clone(), + reader.chain_registry.clone(), + ) + .unwrap(); + w.append_tensor(&rec).unwrap(); + let err = w.append_tensor(&rec).unwrap_err().to_string(); + assert!(err.contains("duplicate tensor name"), "{err}"); + + // Lookup compares names, so a name that shares a prefix or differs in + // one byte is not found. + for name in reader.tensor_names().unwrap() { + assert!(reader.get_tensor_record_bytes(&name).is_some()); + let mut other = name.clone(); + other.push('x'); + assert!(reader.get_tensor_record_bytes(&other).is_none()); + assert!( + reader + .get_tensor_record_bytes(&name[..name.len() - 1]) + .is_none() + ); + } +} + +// ── Size bounds (synthesis section 9 item 14) ───────────────────────────── + +/// Bytes of a compact file that are not plane payload, inline state or +/// shared prelude state, per tensor. +fn overhead_per_tensor(blob: &[u8]) -> f64 { + let reader = ContainerReader::open(blob).unwrap(); + let mut data = 0usize; + for rec in reader.tensor_records().unwrap() { + for p in &rec.terminals { + data += p.payload_bytes.len() + p.inline_state_bytes.len(); + } + } + data += reader + .prelude + .iter() + .map(|e| e.state_bytes.len()) + .sum::(); + (blob.len() - data) as f64 / reader.index.len() as f64 +} + +fn case_a(distinct_shapes: bool) -> Vec { + let mut rng = Rng(0xA); + (0..64) + .map(|i| { + let shape = if distinct_shapes { + vec![64 + i as u64] + } else { + vec![64] + }; + fp32(&format!("model.layers.{i:02}.attn.q.bias"), shape, &mut rng) + }) + .collect() +} + +#[test] +fn compact_overhead_bounds() { + let a = compress( + &case_a(false), + options(true, true, None, ContainerLayout::Compact), + ); + let b = compress( + &case_a(true), + options(true, true, None, ContainerLayout::Compact), + ); + let bare = compress( + &case_a(false), + options(false, false, None, ContainerLayout::Compact), + ); + let legacy = compress( + &case_a(false), + options(true, true, None, ContainerLayout::Legacy), + ); + let (oa, ob, obare, olegacy) = ( + overhead_per_tensor(&a), + overhead_per_tensor(&b), + overhead_per_tensor(&bare), + overhead_per_tensor(&legacy), + ); + eprintln!("overhead B/tensor: A {oa:.1}, B {ob:.1}, bare {obare:.1}, legacy {olegacy:.1}"); + assert!(oa <= 95.0, "case A: {oa:.1} B/tensor"); + assert!(ob <= 100.0, "case B: {ob:.1} B/tensor"); + assert!(obare <= 70.0, "no hash/CRC: {obare:.1} B/tensor"); + assert!(oa < olegacy / 2.0, "compact {oa:.1} vs legacy {olegacy:.1}"); + + // A built-in-only Extension Table costs at most 4 B per entry. + let reader = ContainerReader::open(&a).unwrap(); + let et_len = reader.header.extension_table_length as usize; + let n = reader.extension_table.entries.len(); + assert!(et_len <= 4 * n, "{et_len} B for {n} entries"); +} + +// ── Mutations and fuzz seeds ─────────────────────────────────────────────── + +#[test] +fn mutated_records_are_rejected_or_re_encode_identically() { + let compact = to_compact(&load("legacy_v1_1_0.ptwm")); + let reader = ContainerReader::open(&compact).unwrap(); + let mut w = + ContainerWriter::with_layout(Cursor::new(Vec::new()), ContainerLayout::Compact).unwrap(); + w.set_compact_tables( + reader.extension_table.clone(), + reader.chain_registry.clone(), + ) + .unwrap(); + let encoder = w.record_encoder(); + let records: Vec<&[u8]> = reader + .tensor_names() + .unwrap() + .iter() + .map(|n| reader.get_tensor_record_bytes(n).unwrap()) + .collect(); + let mut rng = Rng(0x4D55_5441_5445); + let mut accepted = 0; + for i in 0..30_000 { + let base = records[i % records.len()]; + let mut m = base.to_vec(); + // Mutate the record head, where the structure is; payload bytes are + // opaque to the parser. + let head = m.len().min(64); + match rng.next_u64() % 4 { + 0 => { + let p = (rng.next_u64() as usize) % head; + m[p] ^= 1 << (rng.next_u64() % 8); + } + 1 => { + let p = (rng.next_u64() as usize) % head; + m[p] = rng.next_u64() as u8; + } + 2 => { + let p = (rng.next_u64() as usize) % head; + m.insert(p, rng.next_u64() as u8); + } + _ => { + let p = (rng.next_u64() as usize) % head; + m.remove(p); + } + } + if let Ok(rec) = reader.parse_record(&m) { + accepted += 1; + assert_eq!(encoder.encode(&rec).unwrap(), m, "mutation {i}"); + } + } + eprintln!("{accepted} of 30000 mutated records parsed"); +} + +#[test] +fn mutated_files_never_panic() { + let compact = to_compact(&load("legacy_trimmed.ptwm")); + let mut rng = Rng(0x4649_4C45); + for _ in 0..3_000 { + let mut m = compact.clone(); + for _ in 0..1 + rng.next_u64() % 3 { + let p = (rng.next_u64() as usize) % m.len(); + m[p] = rng.next_u64() as u8; + } + if let Ok(r) = ContainerReader::open(&m) { + let _ = r.decode_model(); + } + } +} + +/// Write PTWC seeds for the fuzz corpora: +/// +/// ```text +/// PTWM_WRITE_FUZZ_SEEDS=1 cargo test -p ptwm-core --test ptwc \ +/// -- --ignored write_fuzz_seeds +/// ``` +#[test] +#[ignore] +fn write_fuzz_seeds() { + if std::env::var_os("PTWM_WRITE_FUZZ_SEEDS").is_none() { + return; + } + let corpus = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("../../fuzz/corpus"); + let write = |dir: &str, name: &str, bytes: &[u8]| { + let d = corpus.join(dir); + std::fs::create_dir_all(&d).unwrap(); + std::fs::write(d.join(name), bytes).unwrap(); + }; + let trimmed = to_compact(&load("legacy_trimmed.ptwm")); + write("container_open", "ptwc_trimmed", &trimmed); + let mut rng = Rng(3); + let small: Vec = (1..=4) + .map(|i| fp32(&format!("t{i}"), vec![i, 2], &mut rng)) + .collect(); + let small = compress(&small, options(true, true, None, ContainerLayout::Compact)); + write("container_open", "ptwc_bound_source", &small); + + // Compact Extension Tables of both files. + for (name, blob) in [("trimmed", &trimmed), ("bound_source", &small)] { + let r = ContainerReader::open(blob).unwrap(); + let start = r.header.extension_table_offset as usize; + let len = r.header.extension_table_length as usize; + write("extension_table_compact", name, &blob[start..start + len]); + } + + // Records parsed against the fuzz target's fixed context: Extension + // Table entry 0 a transform, 1 Zstd, 2 Identity; registry chain 0 + // explicit, 1 bound. + use ptwm_core::codec::CodecId; + use ptwm_core::tensor_record_compact::{RecordContext, write_compact_record}; + let ctx = RecordContext { + codec_ids: &[None, Some(CodecId::Zstd), Some(CodecId::Identity)], + registry_bound: &[false, true], + }; + let reader = ContainerReader::open(&small).unwrap(); + for (i, mut rec) in reader.tensor_records().unwrap().into_iter().enumerate() { + rec.chain_ref = 1; + for p in &mut rec.terminals { + p.codec_id = CodecId::Zstd; + p.codec_table_idx = 1; + } + let mut out = Vec::new(); + write_compact_record(&rec, &ctx, &mut out).unwrap(); + write("parse_wire_compact", &format!("record_{i}"), &out); + } +} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 5876a2c..e155f15 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -438,6 +438,7 @@ fn compress_model<'py>( forced_codec, allow_codec_ids, chunk_size: streaming_chunk, + ..Default::default() }; let cursor = std::io::Cursor::new(&mut out); @@ -524,6 +525,7 @@ fn compress_compact<'py>( forced_codec, allow_codec_ids, chunk_size: None, + ..Default::default() }; ptwm_core::compact::compress_compact(&input, &opts) }) @@ -649,7 +651,6 @@ fn decode_tensor_info( skip_missing: bool, ) -> PyResult, String, u16, u8)>> { use ptwm_core::metadata::decode_shape_metadata; - use ptwm_core::tensor_record::parse_tensor_record; let data: Vec = blob.as_bytes().to_vec(); let name_owned = name.to_string(); @@ -661,11 +662,7 @@ fn decode_tensor_info( } else { ContainerReader::open(&data)? }; - let record_bytes = - reader.get_tensor_record_bytes(&name_owned).ok_or_else(|| { - PtwmCoreError::InvalidContainer(format!("tensor '{name_owned}' not found")) - })?; - let (rec, _) = parse_tensor_record(record_bytes)?; + let rec = reader.tensor_record(&name_owned)?; match rec.tensor_metadata { Some(meta) => { let (shape, dtype_name) = decode_shape_metadata(&meta)?; @@ -690,7 +687,7 @@ fn decode_tensor_delta_hash<'py>( name: &str, skip_missing: bool, ) -> PyResult>> { - use ptwm_core::tensor_record::{RefKind, parse_tensor_record}; + use ptwm_core::tensor_record::RefKind; let data: Vec = blob.as_bytes().to_vec(); let name_owned = name.to_string(); @@ -701,10 +698,7 @@ fn decode_tensor_delta_hash<'py>( } else { ContainerReader::open(&data)? }; - let record_bytes = reader.get_tensor_record_bytes(&name_owned).ok_or_else(|| { - PtwmCoreError::InvalidContainer(format!("tensor '{name_owned}' not found")) - })?; - let (rec, _) = parse_tensor_record(record_bytes)?; + let rec = reader.tensor_record(&name_owned)?; for dep in &rec.dependencies { if matches!(dep.ref_kind, RefKind::ExternalSafetensors) && let Some(h) = dep.expected_hash @@ -780,8 +774,6 @@ fn list_plane_summaries( blob: &Bound<'_, pyo3::types::PyBytes>, skip_missing: bool, ) -> PyResult> { - use ptwm_core::tensor_record::parse_tensor_record; - let data: Vec = blob.as_bytes().to_vec(); let summaries = py .allow_threads( @@ -794,10 +786,7 @@ fn list_plane_summaries( let names = reader.tensor_names()?; let mut out = Vec::new(); for name in names { - let record_bytes = reader.get_tensor_record_bytes(&name).ok_or_else(|| { - PtwmCoreError::InvalidContainer(format!("tensor '{name}' not found")) - })?; - let (rec, _) = parse_tensor_record(record_bytes)?; + let rec = reader.tensor_record(&name)?; for (idx, plane) in rec.terminals.iter().enumerate() { out.push(( name.clone(), diff --git a/fuzz/Cargo.toml b/fuzz/Cargo.toml index fe20860..a94df8a 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -107,3 +107,15 @@ name = "chain_wire_compact" path = "fuzz_targets/chain_wire_compact.rs" test = false doc = false + +[[bin]] +name = "parse_wire_compact" +path = "fuzz_targets/parse_wire_compact.rs" +test = false +doc = false + +[[bin]] +name = "extension_table_compact" +path = "fuzz_targets/extension_table_compact.rs" +test = false +doc = false diff --git a/fuzz/corpus/container_open/ptwc_bound_source b/fuzz/corpus/container_open/ptwc_bound_source new file mode 100644 index 0000000000000000000000000000000000000000..172f6cb95c409b2c1ec2f9c33eb8df4b2a66e207 GIT binary patch literal 365 zcmeAS2nl!Q<&xrJVBlmC-Fi{|QNaIM9xRLujDm~`Obj5vD9p^n0;Cw27+IJUA`~N) zqLd{$85o&3m>8M1AD*JV_I?!?3nMEtQ%*>NmG5&9>to5aM?ZW(Ea8Rfv1d*(GBY#X zH*qbjEC*>6U}9qWAEi3`|BU&!KpI(|WWIE?J-e6}#GdoOZ}l;)6V4#EN6mUcMUIs< zKu+ z3gRWb@jCO({@y0*g+OE29?j?qDzB<2uc!bSq`<_&lv2J=q>HgM1Y{70U3HkibK?Ye z2KCFcK>U)R%PAbIw@#nlIrSWfpS=6^OV+lH?I{Oe^#hIL$S#<}SPnF-qO!7_g{j2Q elBvYVk*UPklc~fcP!Jq`3=Ewi;l8f^Zc+fh8f`ZK literal 0 HcmV?d00001 diff --git a/fuzz/corpus/container_open/ptwc_trimmed b/fuzz/corpus/container_open/ptwc_trimmed new file mode 100644 index 0000000000000000000000000000000000000000..9368bc99a97a3627ae5b05155cabde6f70fc292b GIT binary patch literal 2457 zcmZ8i2{e^k8{XgdozYCmA?h5{Ij0O|SWdr?Ib<$EC5Mxoj@e1-mJbOTlPQ&iNQ0p+ zif~d%^b3`eA9Zoji0c|c?kE0#{rA3my=(2gpY`r{ulL>0vlAS4yPAp$$q4}*wcvWba#i7^^E6o|*KZ$ZwzIQ{nyT zh}Ar+UPUS+Be9)8WhUCy)p=@tv8mTPlhx2!e;C!6YbY8$-}>Y0d51evH&Sccow}P< zEgH2=H8=KD#G*QQ|7(n!ro(i~Tom%*paxcj{#23O!5FSc`9pzo_G?+( znx2!oZGwvF7i9%k)dpOR`iyk@B4xA;=p|u~MO?vLuP9e8U1RS1Y8SWwLFV0YoFU4GJe#xzbOLJ~)H4u&$ z&YP%ySaEQ|-QYhS3kGChLt<^!+peJLJ=qX?t6RkVIZowfqO&_5PX@S^&M1A8(X7>% z30>)o-M6w{=Sp>Cq5nJOb8pswwQj;Z{WMNlDif-^q-M(w-j?R*GwoZp%ix}qbOe3q z*`3^19qqm&Hh&7j%5`XR9^gVk8k$i)zTQk}NPLc*AOd z4w%64gDIn>+Ar=ry_UaY*N&Vg-dPq0Z=BLFC(%kr+}etc3vAbJ9_kq!KDP;~-2mTs zb>>Od?b7?NL-|L8f~#&plTYJb)GS)xep}#uWwSY*sCj``#_SWB$zV=g_83pkKUX&) zxp5u!Wi{h!H{(-jSl;%-hO|+%eNk~%$#!~opXWh-)1#6;594rn3r8$fDt5-K@R_!y zsk?}t9HmCJQsG)95s)Ehx%D{eg1!cyRm{u4Et3*5m@c0`(0Qt6RJu-5bJR1{F*2-! z6BHQ<@eVe*>@U0CX!n47s_;DlY7=PGoj?ub9XX0gTQ^XJ6Ty7ZR7geM-B zy@#0(eO!qqnI{+aU2R{x?o-JzbMCt<;Q^DmZ?>(^RyOm^vAv^~nUcsX;2Nlo8*2ZW zzIIxBD(ZW8f$7hSQtvmry zOdWU1$fT+GBZ~E1-R_wt6Bt34h{;gMb!^u3xa6&*wO(! zjtnkB!CVveN|`cofRLs{kYJn)lM|u>8ushRzv0)~=F1prk7RmviM2&+AeblSy{@-g zrxXipQtiX|_XB=8ERk2rO%tJ^O66nQF4a%6YrU%J-=gIyI;>(`wy}o=N%xD5N63g< z(5~D-Uc-x+nEn1PIZ2&(Y0s@2ojflW+<4WtCm&)2-U`pixp{KVTb|DOteu?f(F_Ki z02z_Q$^@bCnV0~U|#lIjJ#Wx6zZ}a=Hl&WiW!ATA*kF;x* zfqOoXjDuCpXCJ5x#K)yFXg7abWCalh>wfqO*ab3iWRRQpS@8mAu3}h)kl5ZRe9|Y0 zdGPLFrw)bN>F~(V+Vjw-=Isw&Sq~GeNP^$k63WhR5(VW^xdRWmSzo>Kg91i;q<|<6 z`zK+Hl76pG9P2QM6on87xXQ*~QBXn z-sgPQ0xW6NdlHB;0Yx8(H=Ia*QjJli`=~M( zUE_<|Y(y0XfG!sZ>!A2L)H4_r6Gssj)czGJR|KHfK%oW i&GFIj3kYKsL4QxM^?#|e;Vg_s%a*X))z;kJLhdh+U)FyB literal 0 HcmV?d00001 diff --git a/fuzz/corpus/extension_table_compact/bound_source b/fuzz/corpus/extension_table_compact/bound_source new file mode 100644 index 0000000000000000000000000000000000000000..3b3783782265d544271c8b7dda8138ed71bf5de5 GIT binary patch literal 9 QcmZQ!WMC9zRA6EN008&^I{*Lx literal 0 HcmV?d00001 diff --git a/fuzz/corpus/extension_table_compact/trimmed b/fuzz/corpus/extension_table_compact/trimmed new file mode 100644 index 0000000000000000000000000000000000000000..d767621a6d11a24125db737a8ba842565748f123 GIT binary patch literal 23 ecmd;PWMC9vRA3Zj6k=3lG+<(2Vq{`rVg~>U7y!!v literal 0 HcmV?d00001 diff --git a/fuzz/corpus/parse_wire_compact/record_0 b/fuzz/corpus/parse_wire_compact/record_0 new file mode 100644 index 0000000..ca848f9 --- /dev/null +++ b/fuzz/corpus/parse_wire_compact/record_0 @@ -0,0 +1,2 @@ + Ô'z +lT`:MtL']4Eqyw \ No newline at end of file diff --git a/fuzz/corpus/parse_wire_compact/record_1 b/fuzz/corpus/parse_wire_compact/record_1 new file mode 100644 index 0000000..7492337 --- /dev/null +++ b/fuzz/corpus/parse_wire_compact/record_1 @@ -0,0 +1 @@ + Z%iA=ͣ N*CH|!|ߔzyzy \ No newline at end of file diff --git a/fuzz/corpus/parse_wire_compact/record_2 b/fuzz/corpus/parse_wire_compact/record_2 new file mode 100644 index 0000000..d920023 --- /dev/null +++ b/fuzz/corpus/parse_wire_compact/record_2 @@ -0,0 +1 @@ + :fXZUgM,obJ?޲;☊Rwzxwxx \ No newline at end of file diff --git a/fuzz/corpus/parse_wire_compact/record_3 b/fuzz/corpus/parse_wire_compact/record_3 new file mode 100644 index 0000000000000000000000000000000000000000..4a27e84045a092426a216d2f5b6c9d434427fe39 GIT binary patch literal 81 zcmd;NWMoocVqr=t-zU<=SQ^5@$jZoJR~;ts+&ICVLH+V95Wghoatg=lt<$G>PCW@~%K?~$8u9=D literal 0 HcmV?d00001 diff --git a/fuzz/fuzz_targets/container_open.rs b/fuzz/fuzz_targets/container_open.rs index ed60c87..65e3d54 100644 --- a/fuzz/fuzz_targets/container_open.rs +++ b/fuzz/fuzz_targets/container_open.rs @@ -1,8 +1,21 @@ #![no_main] use libfuzzer_sys::fuzz_target; -use ptwm_core::container::ContainerReader; +use ptwm_core::container::{ContainerLayout, ContainerReader}; fuzz_target!(|data: &[u8]| { - let _ = ContainerReader::open(data); + // Both layouts (\x89PTWM/\x89PTWX and \x89PTWC): open, then decode every + // tensor. Neither may panic or allocate without bound. + let Ok(reader) = ContainerReader::open(data) else { + return; + }; + let _ = reader.decode_model(); + if reader.layout == ContainerLayout::Compact { + // Names come from the index; every record must parse or fail cleanly. + if let Ok(names) = reader.tensor_names() { + for name in names { + let _ = reader.tensor_record(&name); + } + } + } }); diff --git a/fuzz/fuzz_targets/extension_table_compact.rs b/fuzz/fuzz_targets/extension_table_compact.rs new file mode 100644 index 0000000..75b1712 --- /dev/null +++ b/fuzz/fuzz_targets/extension_table_compact.rs @@ -0,0 +1,14 @@ +#![no_main] + +//! Compact Extension Table parser. A table that parses must re-encode to +//! the same bytes: each table has one encoding. + +use libfuzzer_sys::fuzz_target; +use ptwm_core::extension::compact_table::{compact_table_bytes, read_compact_table}; + +fuzz_target!(|data: &[u8]| { + if let Ok(table) = read_compact_table(data) { + let out = compact_table_bytes(&table).expect("a parsed table re-encodes"); + assert_eq!(out, data, "parsed table re-encodes to other bytes"); + } +}); diff --git a/fuzz/fuzz_targets/parse_wire_compact.rs b/fuzz/fuzz_targets/parse_wire_compact.rs new file mode 100644 index 0000000..93381af --- /dev/null +++ b/fuzz/fuzz_targets/parse_wire_compact.rs @@ -0,0 +1,25 @@ +#![no_main] + +//! Compact (`\x89PTWC`) tensor record parser. A record that parses must +//! re-encode to the same bytes: each record has one encoding. + +use libfuzzer_sys::fuzz_target; +use ptwm_core::codec::CodecId; +use ptwm_core::tensor_record_compact::{RecordContext, parse_compact_record, write_compact_record}; + +/// Extension Table entry 0 is a transform, 1 and 2 are plane codecs. +const CODEC_IDS: &[Option] = &[None, Some(CodecId::Zstd), Some(CodecId::Identity)]; +/// Registry chain 0 is explicit, chain 1 has a bound Source. +const BOUND: &[bool] = &[false, true]; + +fuzz_target!(|data: &[u8]| { + let ctx = RecordContext { + codec_ids: CODEC_IDS, + registry_bound: BOUND, + }; + if let Ok(rec) = parse_compact_record(data, &ctx) { + let mut out = Vec::new(); + write_compact_record(&rec, &ctx, &mut out).expect("a parsed record re-encodes"); + assert_eq!(out, data, "parsed record re-encodes to other bytes"); + } +}); From 13faa3b0931508871dfebeca602f6b604bbf73d9 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 01:04:05 +0000 Subject: [PATCH 16/34] fix(S6-ptwc): address review findings - tensor_record_compact: document that the in-memory `flags` byte is recomputed, as the legacy writer does, instead of claiming the writer rejects mismatches; pin it with a unit test. - ptwc: document the registry limit (chain ids 0..=0x7FFF, so at most 0x8000 chains, the legacy ChainRegistry limit) as a deliberate reading of synthesis section 2. - tests/ptwc.rs: single-tensor overhead bound (<= 160 B), a 4 GiB bound Source in PTWC fails fast, BYTE-input tensor, TF_RAW_META metadata and an External state plane round-trip through written PTWC containers. - fuzz workflow: add chain_wire, chain_wire_compact, extension_table_compact and parse_wire_compact to the full sweep. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- .github/workflows/fuzz.yml | 6 +- crates/ptwm-core/src/ptwc.rs | 9 +- crates/ptwm-core/src/tensor_record_compact.rs | 34 +++- crates/ptwm-core/tests/ptwc.rs | 182 ++++++++++++++++++ 4 files changed, 222 insertions(+), 9 deletions(-) diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 67aee02..3e8202f 100644 --- a/.github/workflows/fuzz.yml +++ b/.github/workflows/fuzz.yml @@ -36,12 +36,12 @@ jobs: fail-fast: false matrix: # PRs run a 3-target smoke set (both entropy decoders + the wire - # parser); the full 12-target sweep runs on the weekly cron and on - # manual dispatch. Per-job whole-minute round-up × 12 targets on + # parser); the full 16-target sweep runs on the weekly cron and on + # manual dispatch. Per-job whole-minute round-up × 16 targets on # every core PR dominated this repo's Actions spend — the smoke set # keeps fast regression signal without paying for the full matrix on # every push. Edit both lists together when adding/removing targets. - target: ${{ github.event_name == 'pull_request' && fromJSON('["huffman_decode","rans_decode","parse_wire"]') || fromJSON('["huffman_decode","rans_decode","parse_header","parse_wire","container_open","order1_scale_ac_decode","order1_scale_ac_state_parse","extension_table","ptwx_wasm_extract","policy_resolve","raw_binary_schema","compact_decode"]') }} + target: ${{ github.event_name == 'pull_request' && fromJSON('["huffman_decode","rans_decode","parse_wire"]') || fromJSON('["huffman_decode","rans_decode","parse_header","parse_wire","container_open","order1_scale_ac_decode","order1_scale_ac_state_parse","extension_table","ptwx_wasm_extract","policy_resolve","raw_binary_schema","compact_decode","chain_wire","chain_wire_compact","extension_table_compact","parse_wire_compact"]') }} steps: - name: Checkout repository diff --git a/crates/ptwm-core/src/ptwc.rs b/crates/ptwm-core/src/ptwc.rs index 61b043e..81ce002 100644 --- a/crates/ptwm-core/src/ptwc.rs +++ b/crates/ptwm-core/src/ptwc.rs @@ -16,7 +16,8 @@ //! field, then the Extension Table, prelude and //! registry bytes //! REGISTRY -//! count varint at most 32768 +//! count varint at most 32768 (0x8000), so chain ids are +//! 0..=0x7FFF (container::MAX_CHAIN_ID) //! chain × count varint chain_len, then a CW2 chain (chain::wire_compact) //! whose op_ref is an Extension Table index of a //! built-in transform; chain id = position @@ -37,6 +38,12 @@ //! canonical unsigned LEB128 ([`crate::varint`]), so each file has one //! encoding for what it holds. //! +//! The registry limit is the same as the legacy `ChainRegistry` limit: the +//! largest chain id is 0x7FFF, so a registry holds at most 0x8000 chains. +//! Synthesis section 2 writes "count ≤ 0x7FFF"; this layout reads that as +//! the bound on chain ids, so a model that fits the legacy registry also +//! fits the compact one. +//! //! There is no version field. A future layout change sets an `hflags` bit, //! which current readers reject. diff --git a/crates/ptwm-core/src/tensor_record_compact.rs b/crates/ptwm-core/src/tensor_record_compact.rs index 693a997..ce30821 100644 --- a/crates/ptwm-core/src/tensor_record_compact.rs +++ b/crates/ptwm-core/src/tensor_record_compact.rs @@ -60,11 +60,19 @@ //! ([`bound_source_params`]). An inline chain binds its Source to the //! record the same way as the `\x89PTWS` frame does. //! -//! **Derived fields.** The in-memory `flags`, `chain_ref` (`0x8000` for an -//! inline chain), `payload_len`, and each plane's `codec_id` (from the -//! Extension Table entry its `codec_ref` names) are not stored. The writer -//! rejects a record whose stored values disagree with them, so each record -//! has one encoding and parsing then writing gives the same bytes. +//! **Derived fields.** The in-memory `chain_ref` (`0x8000` for an inline +//! chain), `payload_len`, and each plane's `codec_id` (from the Extension +//! Table entry its `codec_ref` names) are not stored. The writer rejects a +//! record whose values disagree with what the reader derives, so each +//! record has one encoding and parsing then writing gives the same bytes. +//! +//! The in-memory `flags` byte is not stored either. Like the legacy writer +//! ([`crate::tensor_record::write_tensor_record`]), the compact writer +//! ignores it and the reader recomputes it from `payload_hash`, +//! `dependencies` and `tensor_metadata`. A record whose `flags` disagree +//! with those fields (for example a legacy record parsed with the +//! dependency bit set and no dependencies) therefore reads back with the +//! derived flags. use crate::chain::Chain; use crate::chain::wire_compact::{read_chain_compact_exact, write_chain_compact}; @@ -964,6 +972,22 @@ mod tests { assert_eq!(bytes, [0, 0, 1, 2, 16, 1, 0, 1, 1, 4, 1, 2, 3, 4]); } + #[test] + fn in_memory_flags_are_recomputed_not_stored() { + // The writer ignores `flags`, as the legacy writer does, and the + // reader derives them, so stale or unknown bits never reach the wire. + for rec in [record(), rich_record()] { + let want = encode(&rec); + for flags in [0x00, 0xFF, TENSOR_FLAG_DEPENDENCIES, rec.flags ^ 0x07] { + let mut odd = rec.clone(); + odd.flags = flags; + assert_eq!(encode(&odd), want, "flags {flags:#04x}"); + let back = parse_compact_record(&want, &ctx()).unwrap(); + assert_eq!(back.flags, rec.flags); + } + } + } + #[test] fn every_optional_field_round_trips() { let bytes = round_trip(&rich_record()); diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs index b2e7217..531debf 100644 --- a/crates/ptwm-core/tests/ptwc.rs +++ b/crates/ptwm-core/tests/ptwc.rs @@ -17,6 +17,8 @@ use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; use ptwm_core::container::{ContainerLayout, ContainerReader, ContainerWriter}; use ptwm_core::header::{HEADER_LEN, Header, MAGIC_PTWC}; +use ptwm_core::tensor_record::TensorRecord; +use ptwm_core::tensor_record_compact::TF_RAW_META; use ptwm_core::transforms::op::OpId; use ptwm_core::types::{Role, ScaleFormat, ValueFormat}; @@ -526,6 +528,177 @@ fn duplicate_names_are_rejected_and_lookup_is_exact() { } } +/// Rewrite the compact or legacy container `blob` in the compact layout, +/// applying `f` to each record first. +fn rewrite_compact(blob: &[u8], mut f: impl FnMut(&mut TensorRecord)) -> Vec { + let reader = ContainerReader::open(blob).expect("open"); + let mut out = Vec::new(); + let mut w = + ContainerWriter::with_layout(Cursor::new(&mut out), ContainerLayout::Compact).unwrap(); + w.set_prelude(&reader.prelude).unwrap(); + w.set_compact_tables( + reader.extension_table.clone(), + reader.chain_registry.clone(), + ) + .unwrap(); + for mut rec in reader.tensor_records().unwrap() { + f(&mut rec); + w.append_tensor(&rec).unwrap(); + } + w.finalize().unwrap(); + out +} + +/// A BYTE-input fp32 tensor: no shape metadata, an exact Source over the +/// raw byte count (the P4-fixed Python byte path). +fn byte_input_blob(hash: bool, crc: bool) -> (Vec, Vec) { + let mut rng = Rng(0xB7E); + let t = fp32("model.raw", vec![32], &mut rng); + let input = InputTensor { + name: t.name.clone(), + candidate_chains: vec![t.chain.clone()], + dtype_code: t.dtype_code, + input_format: 1, + raw_bytes: &t.bytes, + shape: None, + dtype_name: None, + delta_reference_blake3: None, + }; + let mut out = Cursor::new(Vec::new()); + compress_model( + &mut out, + &[input], + options(hash, crc, None, ContainerLayout::Compact), + ) + .expect("compress_model"); + (out.into_inner(), t.bytes) +} + +#[test] +fn byte_input_tensor_round_trips_in_the_compact_layout() { + for (hash, crc) in [(true, true), (false, false)] { + let (blob, raw) = byte_input_blob(hash, crc); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(reader.layout, ContainerLayout::Compact); + let rec = reader.tensor_record("model.raw").unwrap(); + assert_eq!(rec.input_format, 1); + assert!(rec.tensor_metadata.is_none()); + assert_eq!(reader.tensor_shape("model.raw").unwrap(), None); + assert_eq!(reader.decode_tensor("model.raw").unwrap(), raw); + assert!(to_compact(&blob) == blob, "rewrite differs"); + } +} + +#[test] +fn raw_cbor_metadata_round_trips_in_a_compact_container() { + // Metadata that is not the canonical shape map is stored as TF_RAW_META + // and must come back byte for byte. The tensor uses an explicit + // registry chain, so its Source does not depend on the metadata. + let (blob, raw) = byte_input_blob(true, true); + let metas: [&[u8]; 3] = [ + // {"a": 1} + &[0xA1, 0x61, b'a', 0x01], + // Shape map with a non-canonical (two-byte) length for the shape. + &[ + 0xA2, 0x65, b's', b'h', b'a', b'p', b'e', 0x98, 0x01, 0x18, 0x20, 0x65, b'd', b't', + b'y', b'p', b'e', 0x67, b'f', b'l', b'o', b'a', b't', b'3', b'2', + ], + // Arbitrary bytes, not CBOR at all. + &[0xFF, 0x00, 0x13], + ]; + for meta in metas { + let compact = rewrite_compact(&blob, |rec| { + rec.tensor_metadata = Some(meta.to_vec()); + rec.flags |= ptwm_core::tensor_record::TENSOR_FLAG_METADATA; + }); + let reader = ContainerReader::open(&compact).unwrap(); + let rec = reader.tensor_record("model.raw").unwrap(); + assert_eq!(rec.tensor_metadata.as_deref(), Some(meta)); + let bytes = reader.get_tensor_record_bytes("model.raw").unwrap(); + assert_ne!(bytes[0] & TF_RAW_META, 0, "stored as TF_RAW_META"); + assert_eq!(reader.decode_tensor("model.raw").unwrap(), raw); + assert!(to_compact(&compact) == compact, "rewrite differs"); + } +} + +#[test] +fn external_state_plane_round_trips_in_a_compact_container() { + // The reader does not resolve external state at open time; the record + // and the file must still round-trip exactly. + use ptwm_core::codec::StateSource; + use ptwm_core::plane_record::ExternalStateRef; + let data = feature_tensors(); + let blob = compress(&data, options(true, true, None, ContainerLayout::Compact)); + let target = "model.norm.weight"; + let ext = ExternalStateRef { + ref_kind: 1, + expected_hash: Some([0x5A; 32]), + ref_bytes: b"dict/zstd-v1.bin".to_vec(), + state_format_version: 3, + }; + let compact = rewrite_compact(&blob, |rec| { + if rec.name == target { + let p = &mut rec.terminals[0]; + p.state_source = StateSource::External; + p.state_version = 0; + p.state_info = 0; + p.inline_state_bytes.clear(); + p.external_state = Some(ext.clone()); + } + }); + let reader = ContainerReader::open(&compact).unwrap(); + let rec = reader.tensor_record(target).unwrap(); + assert_eq!(rec.terminals[0].state_source, StateSource::External); + assert_eq!(rec.terminals[0].external_state.as_ref(), Some(&ext)); + assert!(to_compact(&compact) == compact, "rewrite differs"); + // The other tensors are unaffected. + for t in data.iter().filter(|t| t.name != target) { + assert_eq!( + reader.decode_tensor(&t.name).unwrap(), + t.bytes, + "{}", + t.name + ); + } +} + +#[test] +fn compact_source_of_4_gib_fails_without_allocating() { + // Shape metadata that declares [1 << 30] fp32 (4 GiB) binds the + // registry chain's Source to 4 GiB for a 256-byte tensor. Decoding must + // fail on the size check, not try to size buffers from the chain. + use ptwm_core::metadata::encode_shape_metadata; + let mut rng = Rng(0xB0B); + let one = vec![fp32("t", vec![64], &mut rng)]; + let blob = compress(&one, options(true, true, None, ContainerLayout::Compact)); + assert_eq!( + ContainerReader::open(&blob) + .unwrap() + .chain_registry + .chains + .len(), + 1 + ); + let bomb = rewrite_compact(&blob, |rec| { + rec.tensor_metadata = Some(encode_shape_metadata(&[1 << 30], "float32")); + rec.orig_size = 4 << 30; + }); + let reader = ContainerReader::open(&bomb).unwrap(); + assert_eq!(reader.tensor_shape("t").unwrap().unwrap().0, vec![1 << 30]); + let check = |err: String, t0: std::time::Instant| { + eprintln!("4 GiB Source: {err}"); + assert!(t0.elapsed() < std::time::Duration::from_secs(1), "{err}"); + assert!( + err.contains("1073741824") || err.contains("Source declares"), + "{err}" + ); + }; + let t0 = std::time::Instant::now(); + check(reader.decode_tensor("t").unwrap_err().to_string(), t0); + let t0 = std::time::Instant::now(); + check(reader.decode_model().unwrap_err().to_string(), t0); +} + // ── Size bounds (synthesis section 9 item 14) ───────────────────────────── /// Bytes of a compact file that are not plane payload, inline state or @@ -590,6 +763,15 @@ fn compact_overhead_bounds() { assert!(obare <= 70.0, "no hash/CRC: {obare:.1} B/tensor"); assert!(oa < olegacy / 2.0, "compact {oa:.1} vs legacy {olegacy:.1}"); + // A single tensor pays the whole fixed overhead (header, tables, + // registry, index, footer) alone. + let mut rng = Rng(0x51); + let one = vec![fp32("model.bias", vec![64], &mut rng)]; + let single = compress(&one, options(true, true, None, ContainerLayout::Compact)); + let osingle = overhead_per_tensor(&single); + eprintln!("single-tensor overhead: {osingle:.0} B"); + assert!(osingle <= 160.0, "single tensor: {osingle:.0} B"); + // A built-in-only Extension Table costs at most 4 B per entry. let reader = ContainerReader::open(&a).unwrap(); let et_len = reader.header.extension_table_length as usize; From 7e566dc6d6ccca03d4efbf7ff9d1731c528bf3c5 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 01:16:32 +0000 Subject: [PATCH 17/34] feat(transcode): compact manifest for PTWC containers Stage 7 of the issue #5 work (synthesis sections 6 and 10 step 7). - explode of a \x89PTWC container gives a "ptwm-manifest-c" version 2 registry manifest: container_flags {embedded_wasm: false, layout: "compact"}, registry.extension_table as a list of {"op"} / {"codec"} / {"named"} / {"full_b64"} entries in table order, shared_state, and chains in registry order with a bound Source as params_b64 null. No registry_hash_blake3. Tensors come in body order; tensor and plane manifests have the legacy form (codec_table_idx = codec_ref, codec is informational and checked by the writer). - implode dispatches on (ptwm_format, manifest_version): ("ptwm-manifest", 1) takes the legacy path unchanged, ("ptwm-manifest-c", 2) rebuilds the table through the compact table reader (unknown ids and names, duplicates, built-ins stored in full are rejected) and writes table, registry and order verbatim with ContainerWriter::with_layout(Compact) and set_compact_tables. Any other pair is rejected. The compact manifest structs use deny_unknown_fields. - Legacy explode/implode output is unchanged (shared helpers only). - ptwm-py compress_model takes a keyword-only container_layout ("legacy" default, "compact"; other values raise ValueError), so Python can write PTWC now. compress_ptwm_blob, write_ptwm_shard, write_lmdb, LmdbWeightStore.put_tensor and write_webdataset pass a keyword-only container_layout through (default "legacy" until the default flips). - stores/_common.py: manifest_container_layout() and ExplodedModel.container_layout; docs state that per-tensor implode uses the archive-wide tables. Tests (Rust, tests/ptwc.rs): every stage-6 PTWC case (legacy fixtures rewritten as PTWC, compress_model matrix over both layouts x hash/CRC x chunked, bound Source, BYTE input, raw CBOR metadata, external state, P4 slack) now also checks implode(explode(x)) == x, the manifest form, body order and a one-tensor implode per tensor that decodes like the full file. New: manifest lists the table and binds Sources; named and third-party entries round-trip; malformed compact manifests are rejected (version/format mismatch, layout, unknown fields, entry kinds, unknown ids/names, duplicates, built-in in full_b64, chain id != position, codec disagreeing with the table); a frozen v1.1 implode gate and a ("ptwm-manifest", 1)-only gate both reject compact manifests. Tests (Python, tests/ptwm/stores/test_compact_layout.py): layout keyword on compress_ptwm_blob/write_ptwm_shard, unknown layout rejected, compact manifest form and identity per method hint, Rust JCS == Python JCS, per-tensor raw_from_members in both layouts, LMDB blob and exploded stores, put_tensor (incl. delta) in either layout, WebDataset single and multi-shard stores with a compact registry. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/transcode/manifest.rs | 71 ++- crates/ptwm-core/src/transcode/mod.rs | 355 +++++++++++---- crates/ptwm-core/tests/ptwc.rs | 499 +++++++++++++++++++++ crates/ptwm-py/src/lib.rs | 12 +- python/ptwm/_rust/__init__.py | 18 +- python/ptwm/sharding/_writer.py | 8 +- python/ptwm/stores/_common.py | 27 +- python/ptwm/stores/_lmdb.py | 19 +- python/ptwm/stores/_webdataset.py | 9 +- tests/ptwm/stores/test_compact_layout.py | 214 +++++++++ 10 files changed, 1144 insertions(+), 88 deletions(-) create mode 100644 tests/ptwm/stores/test_compact_layout.py diff --git a/crates/ptwm-core/src/transcode/manifest.rs b/crates/ptwm-core/src/transcode/manifest.rs index 5c02725..08b20bc 100644 --- a/crates/ptwm-core/src/transcode/manifest.rs +++ b/crates/ptwm-core/src/transcode/manifest.rs @@ -5,6 +5,10 @@ //! large/opaque binaries (compressed plane payloads, externalized codec state, //! embedded modules) are referenced by member key and live outside the JSON. //! +//! The registry manifest has one form per container layout: +//! [`RegistryManifest`] for the legacy layout and [`CompactRegistryManifest`] +//! for the compact one. Per-tensor manifests ([`TensorManifest`]) are shared. +//! //! The schema restricts scalars to integers, strings, booleans, and `null`; //! opaque byte fields are carried as base64 (`*_b64`). See `super::jcs` for //! the canonicalization that makes manifest hashes reproducible. @@ -13,6 +17,12 @@ use serde::{Deserialize, Serialize}; pub const PTWM_MANIFEST_FORMAT: &str = "ptwm-manifest"; pub const MANIFEST_VERSION: u32 = 1; +/// `ptwm_format` of a compact-layout (`\x89PTWC`) registry manifest. +pub const PTWM_MANIFEST_FORMAT_COMPACT: &str = "ptwm-manifest-c"; +/// `manifest_version` of a compact-layout registry manifest. +pub const MANIFEST_VERSION_COMPACT: u32 = 2; +/// `container_flags.layout` of a compact-layout registry manifest. +pub const LAYOUT_COMPACT: &str = "compact"; /// Archive-scoped registry manifest. Lifts the binary `extension_table` and /// `shared_prelude` to archive scope so codecs/ops/codebooks/chains are defined @@ -54,6 +64,60 @@ pub struct Registry { pub extension_table_b64: Option, } +/// Archive-scoped registry manifest of a compact-layout (`\x89PTWC`) +/// container: `ptwm_format` [`PTWM_MANIFEST_FORMAT_COMPACT`], +/// `manifest_version` [`MANIFEST_VERSION_COMPACT`]. +/// +/// It lists the Extension Table entry by entry instead of carrying its hash, +/// because the compact table is rebuilt from that list. Per-tensor manifests +/// are [`TensorManifest`]s, as for the legacy layout. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CompactRegistryManifest { + pub ptwm_format: String, + pub manifest_version: u32, + pub container_flags: CompactContainerFlags, + pub registry: CompactRegistry, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CompactContainerFlags { + /// Always `false`: embedded WASM needs the legacy layout. + pub embedded_wasm: bool, + /// Always [`LAYOUT_COMPACT`]. + pub layout: String, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct CompactRegistry { + /// The Extension Table entries in table order. Plane `codec_table_idx` + /// values and the chains' op references index this list. + pub extension_table: Vec, + pub shared_state: Vec, + /// Registry chains in registry order; `id` is the position. A Source + /// node with `params_b64` `null` is bound: it takes its params from the + /// shape and dtype of each record that uses the chain. + pub chains: Vec, +} + +/// One compact Extension Table entry (see +/// [`crate::extension::compact_table`]). Serialized as a one-key object: +/// `{"op": u16}`, `{"codec": u16}`, `{"named": str}` or `{"full_b64": str}`. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case", deny_unknown_fields)] +pub enum ExtensionEntryJson { + /// Built-in transform by `OpId` wire id. + Op(u16), + /// Built-in plane codec by `CodecId` wire id. + Codec(u16), + /// Built-in without a wire id, by name. + Named(String), + /// Third-party entry: the legacy entry bytes, base64. + FullB64(String), +} + #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] #[serde(deny_unknown_fields)] pub struct OpEntry { @@ -181,9 +245,12 @@ pub struct PlaneJson { pub role: String, /// Exact `PlaneRole` discriminant, used for reconstruction. pub role_code: u8, - /// `CodecId` wire value. + /// `CodecId` wire value. The compact layout does not store it: it is + /// informational there, and `implode` rejects a value that differs from + /// the codec of the table entry `codec_table_idx` names. pub codec: u16, - /// Extension-table index for the codec (`0xFFFF` = absent). + /// Extension-table index for the codec (`0xFFFF` = absent). The compact + /// layout's `codec_ref`. pub codec_table_idx: u16, /// Where the codec state lives in the manifest: /// `none` | `inline` | `member` | `shared` | `external`. diff --git a/crates/ptwm-core/src/transcode/mod.rs b/crates/ptwm-core/src/transcode/mod.rs index 7d0d7ba..551e368 100644 --- a/crates/ptwm-core/src/transcode/mod.rs +++ b/crates/ptwm-core/src/transcode/mod.rs @@ -11,6 +11,18 @@ //! that `implode(explode(blob)) == blob` byte-for-byte for every container the //! compression engine produces. //! +//! Both container layouts are supported. A legacy (`\x89PTWM`) container gives +//! a [`RegistryManifest`] (`"ptwm-manifest"`, version 1); `implode` rebuilds +//! its Extension Table from the chains and codecs and checks the recorded +//! hash. A compact (`\x89PTWC`) container gives a [`CompactRegistryManifest`] +//! (`"ptwm-manifest-c"`, version 2) that lists the table entries, the +//! registry chains (a bound Source has `params_b64` `null`) and the tensors in +//! body order; `implode` writes them back as listed. Readers that predate the +//! compact manifest reject it on its format and version. Tensor manifests +//! have the same form for both layouts. `implode` with a subset of the tensor +//! manifests writes a container of those tensors with the archive-wide +//! tables, so each tensor decodes as in the full container. +//! //! The encoder may build a container transiently in memory; the stored artifact //! (WebDataset shard / LMDB exploded value) never contains one. @@ -24,10 +36,13 @@ use std::io::Cursor; use crate::chain::graph::{Chain, ChainEdge, ChainNode, TerminalRef}; use crate::codec::{CodecId, StateSource}; use crate::container::{ - BuiltExtensionTable, ChainRegistry, CodecSet, ContainerReader, ContainerWriter, - build_extension_table, encode_registry_for_table, + BuiltExtensionTable, ChainRegistry, CodecSet, ContainerLayout, ContainerReader, + ContainerWriter, build_extension_table, encode_registry_for_table, }; use crate::error::PtwmCoreError; +use crate::extension::compact_table::{ + CompactEntry, TAG_CODEC, TAG_FULL, TAG_NAMED, TAG_OP, classify, read_compact_table, +}; use crate::extension::{CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionTable}; use crate::header::FLAG_PTWX; use crate::layout::PlaneLayout; @@ -41,6 +56,7 @@ use crate::tensor_record::{ use crate::transforms::op::OpId; use crate::types::PlaneRole; use crate::types::role::Role; +use crate::varint::write_varint; use manifest::*; @@ -303,14 +319,12 @@ fn extension_table_bytes<'a>( // ── Explode ─────────────────────────────────────────────────────────────────── -/// Parse a `.ptwm` container into a registry manifest, per-tensor manifests, -/// and loose binary members. +/// Parse a `.ptwm` container (either layout) into a registry manifest, +/// per-tensor manifests in body order, and loose binary members. pub fn explode(blob: &[u8]) -> Result { let reader = ContainerReader::open_partial(blob)?; - if reader.layout != crate::container::ContainerLayout::Legacy { - return Err(err( - "compact (\\x89PTWC) containers are not supported by the transcoder yet", - )); + if reader.layout == ContainerLayout::Compact { + return explode_compact(&reader); } let embedded_wasm = reader.header.flags & FLAG_PTWX != 0; if embedded_wasm { @@ -322,35 +336,15 @@ pub fn explode(blob: &[u8]) -> Result { let mut members: Vec<(String, Vec)> = Vec::new(); // Shared state → registry + members. - let mut shared_state: Vec = Vec::with_capacity(reader.prelude.len()); - for entry in &reader.prelude { - let payload = shared_state_key(entry.shared_state_id); - members.push((payload.clone(), entry.state_bytes.clone())); - shared_state.push(SharedStateEntry { - id: entry.shared_state_id, - codec: entry.codec_id.as_u16(), - state_format_version: entry.state_format_version, - applies_to_mask: entry.applies_to_mask, - hash_xxh64: format!("{:016x}", entry.state_xxhash64), - name: entry.name.clone(), - payload, - }); - } + let shared_state = shared_state_to_manifest(&reader.prelude, &mut members); // Chains → registry + op projection. - let mut chains: Vec = Vec::with_capacity(reader.chain_registry.chains.len()); + let chains = chains_to_manifest(&reader.chain_registry); let mut op_ids: std::collections::BTreeSet = std::collections::BTreeSet::new(); - for (id, chain) in &reader.chain_registry.chains { - let (nodes, edges, terminals) = chain_to_json(chain); + for (_, chain) in &reader.chain_registry.chains { for n in &chain.nodes { op_ids.insert(n.op.as_u16()); } - chains.push(ChainEntry { - id: *id, - nodes, - edges, - terminals, - }); } // File order = ascending tensor offset. @@ -507,6 +501,147 @@ pub fn explode(blob: &[u8]) -> Result { }) } +/// Project the shared-state prelude to manifest entries, moving each state +/// into a member. +fn shared_state_to_manifest( + prelude: &[PreludeEntry], + members: &mut Vec<(String, Vec)>, +) -> Vec { + prelude + .iter() + .map(|entry| { + let payload = shared_state_key(entry.shared_state_id); + members.push((payload.clone(), entry.state_bytes.clone())); + SharedStateEntry { + id: entry.shared_state_id, + codec: entry.codec_id.as_u16(), + state_format_version: entry.state_format_version, + applies_to_mask: entry.applies_to_mask, + hash_xxh64: format!("{:016x}", entry.state_xxhash64), + name: entry.name.clone(), + payload, + } + }) + .collect() +} + +/// Project the chain registry to manifest entries, in registry order. +fn chains_to_manifest(registry: &ChainRegistry) -> Vec { + registry + .chains + .iter() + .map(|(id, chain)| { + let (nodes, edges, terminals) = chain_to_json(chain); + ChainEntry { + id: *id, + nodes, + edges, + terminals, + } + }) + .collect() +} + +/// Project one Extension Table entry to its compact manifest form. +fn extension_entry_to_json( + entry: &crate::extension::ExtensionTableEntry, +) -> Result { + Ok(match classify(entry)? { + CompactEntry::Op(op) => ExtensionEntryJson::Op(op.as_u16()), + CompactEntry::Codec(c) => ExtensionEntryJson::Codec(c.as_u16()), + CompactEntry::Named(name) => ExtensionEntryJson::Named(name.to_owned()), + CompactEntry::Full(e) => ExtensionEntryJson::FullB64(base64::encode(&e.to_bytes()?)), + }) +} + +/// Rebuild the Extension Table of a compact manifest. The entries are +/// encoded as a compact table and parsed back, so the manifest gets every +/// check the compact reader applies (unknown ids and names, duplicates, +/// built-in ids in full entries, canonical entry bytes). +fn extension_table_from_json( + entries: &[ExtensionEntryJson], +) -> Result { + let mut bytes = Vec::new(); + write_varint(&mut bytes, entries.len() as u64); + for entry in entries { + match entry { + ExtensionEntryJson::Op(id) => { + bytes.push(TAG_OP); + write_varint(&mut bytes, u64::from(*id)); + } + ExtensionEntryJson::Codec(id) => { + bytes.push(TAG_CODEC); + write_varint(&mut bytes, u64::from(*id)); + } + ExtensionEntryJson::Named(name) => { + bytes.push(TAG_NAMED); + write_varint(&mut bytes, name.len() as u64); + bytes.extend_from_slice(name.as_bytes()); + } + ExtensionEntryJson::FullB64(b64) => { + let entry = base64::decode(b64) + .ok_or_else(|| err("extension_table full_b64 not valid base64"))?; + bytes.push(TAG_FULL); + write_varint(&mut bytes, entry.len() as u64); + bytes.extend_from_slice(&entry); + } + } + } + read_compact_table(&bytes).map_err(|e| err(format!("registry.extension_table: {e}"))) +} + +/// [`explode`] for the compact layout (`\x89PTWC`). +/// +/// The compact reader accepts only canonical encodings, so writing the +/// parsed tables, prelude, registry and records again (as [`implode`] does) +/// reproduces the file. Tensors come in body order. +fn explode_compact(reader: &ContainerReader<'_>) -> Result { + let mut members: Vec<(String, Vec)> = Vec::new(); + let shared_state = shared_state_to_manifest(&reader.prelude, &mut members); + let chains = chains_to_manifest(&reader.chain_registry); + let extension_table = reader + .extension_table + .entries + .iter() + .map(extension_entry_to_json) + .collect::, _>>()?; + + let records = reader.tensor_records()?; + let mut tensors: Vec = Vec::with_capacity(records.len()); + for (file_idx, record) in records.iter().enumerate() { + let sample_key = format!("{file_idx:06}"); + let tm = tensor_to_manifest(record, &sample_key, &mut members)?; + let value = serde_json::to_value(&tm) + .map_err(|e| err(format!("serialize tensor manifest: {e}")))?; + tensors.push(ExplodedTensor { + name: record.name.clone(), + key: sample_key, + manifest_json: jcs::to_string(&value)?, + }); + } + + let registry = CompactRegistryManifest { + ptwm_format: PTWM_MANIFEST_FORMAT_COMPACT.to_string(), + manifest_version: MANIFEST_VERSION_COMPACT, + container_flags: CompactContainerFlags { + embedded_wasm: false, + layout: LAYOUT_COMPACT.to_string(), + }, + registry: CompactRegistry { + extension_table, + shared_state, + chains, + }, + }; + let registry_value = + serde_json::to_value(®istry).map_err(|e| err(format!("serialize registry: {e}")))?; + Ok(Exploded { + registry_json: jcs::to_string(®istry_value)?, + tensors, + members, + }) +} + /// Fold any inline-chain ops into the op projection so `registry.ops` is a /// superset of every op referenced anywhere in the archive. fn check_no_inline_chain_ops(record: &TensorRecord, op_ids: &mut std::collections::BTreeSet) { @@ -650,7 +785,9 @@ fn tensor_to_manifest( // ── Implode ─────────────────────────────────────────────────────────────────── /// Reconstruct a `.ptwm` container from a registry manifest, per-tensor -/// manifests (in file order), and loose binary members. +/// manifests (in file order), and loose binary members. The registry +/// manifest's `ptwm_format` and `manifest_version` select the layout; any +/// other pair is rejected. pub fn implode( registry_json: &str, tensor_jsons: &[String], @@ -662,11 +799,20 @@ pub fn implode( .map_err(|e| err(format!("parse registry manifest: {e}")))?; let format = value.get("ptwm_format").and_then(|v| v.as_str()); let version = value.get("manifest_version").and_then(|v| v.as_u64()); - if format != Some(PTWM_MANIFEST_FORMAT) || version != Some(u64::from(MANIFEST_VERSION)) { - return Err(err(format!( - "unsupported manifest: ptwm_format {format:?}, manifest_version {version:?} \ - (this reader handles {PTWM_MANIFEST_FORMAT:?} version {MANIFEST_VERSION})" - ))); + match (format, version) { + (Some(PTWM_MANIFEST_FORMAT), Some(v)) if v == u64::from(MANIFEST_VERSION) => {} + (Some(PTWM_MANIFEST_FORMAT_COMPACT), Some(v)) + if v == u64::from(MANIFEST_VERSION_COMPACT) => + { + return implode_compact(value, tensor_jsons, members); + } + _ => { + return Err(err(format!( + "unsupported manifest: ptwm_format {format:?}, manifest_version {version:?} \ + (this reader handles {PTWM_MANIFEST_FORMAT:?} version {MANIFEST_VERSION} \ + and {PTWM_MANIFEST_FORMAT_COMPACT:?} version {MANIFEST_VERSION_COMPACT})" + ))); + } } let registry: RegistryManifest = serde_json::from_value(value).map_err(|e| err(format!("parse registry manifest: {e}")))?; @@ -683,44 +829,9 @@ pub fn implode( .as_deref() .unwrap_or(CURRENT_BUILTIN_VERSION); - // Rebuild the shared-state prelude (order preserved). - let mut prelude_entries: Vec = - Vec::with_capacity(registry.registry.shared_state.len()); - for s in ®istry.registry.shared_state { - // Materialize the (potentially large) state bytes here — this copy now - // runs with the GIL released (the caller borrows the Python buffers). - let state_bytes = members - .remove(&s.payload) - .ok_or_else(|| err(format!("missing shared-state member {:?}", s.payload)))? - .to_vec(); - prelude_entries.push(PreludeEntry { - shared_state_id: s.id, - codec_id: CodecId::from_u16(s.codec) - .ok_or_else(|| err(format!("unknown shared-state codec id {:#06x}", s.codec)))?, - state_format_version: s.state_format_version, - applies_to_mask: s.applies_to_mask, - state_xxhash64: u64::from_str_radix(&s.hash_xxh64, 16) - .map_err(|_| err("shared-state hash_xxh64 not valid hex"))?, - name: s.name.clone(), - state_bytes, - }); - } - - // Rebuild the chain registry (order preserved). - let mut chain_registry = ChainRegistry::new(); - for c in ®istry.registry.chains { - chain_registry - .chains - .push((c.id, chain_from_json(&c.nodes, &c.edges, &c.terminals)?)); - } - - // Rebuild each tensor record. - let mut records: Vec = Vec::with_capacity(tensor_jsons.len()); - for json in tensor_jsons { - let tm: TensorManifest = - serde_json::from_str(json).map_err(|e| err(format!("parse tensor manifest: {e}")))?; - records.push(record_from_manifest(&tm, &mut members)?); - } + let prelude_entries = prelude_from_manifest(®istry.registry.shared_state, &mut members)?; + let chain_registry = chains_from_manifest(®istry.registry.chains)?; + let records = records_from_manifests(tensor_jsons, &mut members)?; // Re-drive the deterministic container writer. let mut cursor = Cursor::new(Vec::::new()); @@ -777,6 +888,100 @@ pub fn implode( Ok(cursor.into_inner()) } +/// [`implode`] for a compact-layout manifest: writes the Extension Table, +/// registry and tensor order of the manifest verbatim in the compact layout. +fn implode_compact( + value: serde_json::Value, + tensor_jsons: &[String], + mut members: HashMap, +) -> Result, PtwmCoreError> { + let registry: CompactRegistryManifest = + serde_json::from_value(value).map_err(|e| err(format!("parse registry manifest: {e}")))?; + if registry.container_flags.embedded_wasm { + return Err(err("embedded WASM needs the legacy layout")); + } + if registry.container_flags.layout != LAYOUT_COMPACT { + return Err(err(format!( + "container_flags.layout {:?} in a {PTWM_MANIFEST_FORMAT_COMPACT:?} manifest \ + (expected {LAYOUT_COMPACT:?})", + registry.container_flags.layout + ))); + } + let table = extension_table_from_json(®istry.registry.extension_table)?; + let prelude_entries = prelude_from_manifest(®istry.registry.shared_state, &mut members)?; + let chain_registry = chains_from_manifest(®istry.registry.chains)?; + let records = records_from_manifests(tensor_jsons, &mut members)?; + + let mut cursor = Cursor::new(Vec::::new()); + { + let mut writer = ContainerWriter::with_layout(&mut cursor, ContainerLayout::Compact)?; + writer.set_prelude(&prelude_entries)?; + // The table and registry are written as listed. The table is + // archive-wide, so a one-tensor implode keeps every plane's + // `codec_table_idx` valid. + writer.set_compact_tables(table, chain_registry)?; + for record in &records { + writer.append_tensor(record)?; + } + writer.finalize()?; + } + Ok(cursor.into_inner()) +} + +/// Rebuild the shared-state prelude (order preserved), taking each state +/// out of `members`. +fn prelude_from_manifest( + shared_state: &[SharedStateEntry], + members: &mut HashMap, +) -> Result, PtwmCoreError> { + let mut prelude_entries: Vec = Vec::with_capacity(shared_state.len()); + for s in shared_state { + // Materialize the (potentially large) state bytes here — this copy now + // runs with the GIL released (the caller borrows the Python buffers). + let state_bytes = members + .remove(&s.payload) + .ok_or_else(|| err(format!("missing shared-state member {:?}", s.payload)))? + .to_vec(); + prelude_entries.push(PreludeEntry { + shared_state_id: s.id, + codec_id: CodecId::from_u16(s.codec) + .ok_or_else(|| err(format!("unknown shared-state codec id {:#06x}", s.codec)))?, + state_format_version: s.state_format_version, + applies_to_mask: s.applies_to_mask, + state_xxhash64: u64::from_str_radix(&s.hash_xxh64, 16) + .map_err(|_| err("shared-state hash_xxh64 not valid hex"))?, + name: s.name.clone(), + state_bytes, + }); + } + Ok(prelude_entries) +} + +/// Rebuild the chain registry (order preserved). +fn chains_from_manifest(chains: &[ChainEntry]) -> Result { + let mut chain_registry = ChainRegistry::new(); + for c in chains { + chain_registry + .chains + .push((c.id, chain_from_json(&c.nodes, &c.edges, &c.terminals)?)); + } + Ok(chain_registry) +} + +/// Rebuild each tensor record, taking its members out of `members`. +fn records_from_manifests( + tensor_jsons: &[String], + members: &mut HashMap, +) -> Result, PtwmCoreError> { + let mut records: Vec = Vec::with_capacity(tensor_jsons.len()); + for json in tensor_jsons { + let tm: TensorManifest = + serde_json::from_str(json).map_err(|e| err(format!("parse tensor manifest: {e}")))?; + records.push(record_from_manifest(&tm, members)?); + } + Ok(records) +} + fn record_from_manifest( tm: &TensorManifest, members: &mut HashMap, diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs index 531debf..a9ac21f 100644 --- a/crates/ptwm-core/tests/ptwc.rs +++ b/crates/ptwm-core/tests/ptwc.rs @@ -14,11 +14,18 @@ use std::io::Cursor; use std::path::PathBuf; use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; +use ptwm_core::codec::CodecId; use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; use ptwm_core::container::{ContainerLayout, ContainerReader, ContainerWriter}; +use ptwm_core::extension::compact_table::{CompactEntry, classify}; +use ptwm_core::extension::{ + Attestation, CanonicalId, CapabilityMap, ExtensionTableEntry, Kind, Lifecycle, builtin_entries, + builtin_registry, +}; use ptwm_core::header::{HEADER_LEN, Header, MAGIC_PTWC}; use ptwm_core::tensor_record::TensorRecord; use ptwm_core::tensor_record_compact::TF_RAW_META; +use ptwm_core::transcode::{Exploded, explode, implode}; use ptwm_core::transforms::op::OpId; use ptwm_core::types::{Role, ScaleFormat, ValueFormat}; @@ -115,6 +122,7 @@ fn legacy_fixtures_round_trip_through_the_compact_layout() { to_compact(&compact) == compact, "{fixture}: rewrite differs" ); + assert_transcodes(&compact, fixture); } } @@ -167,6 +175,9 @@ fn p4_source_slack_is_rejected_in_the_compact_layout() { let reader = ContainerReader::open(&compact).unwrap(); let err = reader.decode_model().unwrap_err().to_string(); assert!(err.contains("Source declares"), "{err}"); + // The transcoder carries the file unchanged; per-tensor implodes fail to + // decode the same way. + assert_transcodes(&compact, "legacy_p4_byte_fp32.ptwm"); } // ── compress_model in the compact layout ─────────────────────────────────── @@ -442,6 +453,10 @@ fn compress_model_matrix_round_trips_in_both_layouts() { let mut all = reader.decode_model().unwrap(); all.sort(); decoded.push(all); + assert_transcodes( + &blob, + &format!("{layout:?} hash={hash} crc={crc} chunk={chunk:?}"), + ); if layout == ContainerLayout::Compact { // parse then write is the identity. assert!(to_compact(&blob) == blob, "rewrite differs"); @@ -479,6 +494,7 @@ fn bound_source_shares_one_chain_across_shapes() { for t in &data { assert_eq!(cr.decode_tensor(&t.name).unwrap(), t.bytes); } + assert_transcodes(&compact, "bound Source"); // A record without shape metadata cannot use the bound chain: the writer // refuses it. @@ -586,6 +602,7 @@ fn byte_input_tensor_round_trips_in_the_compact_layout() { assert_eq!(reader.tensor_shape("model.raw").unwrap(), None); assert_eq!(reader.decode_tensor("model.raw").unwrap(), raw); assert!(to_compact(&blob) == blob, "rewrite differs"); + assert_transcodes(&blob, "BYTE input"); } } @@ -618,6 +635,7 @@ fn raw_cbor_metadata_round_trips_in_a_compact_container() { assert_ne!(bytes[0] & TF_RAW_META, 0, "stored as TF_RAW_META"); assert_eq!(reader.decode_tensor("model.raw").unwrap(), raw); assert!(to_compact(&compact) == compact, "rewrite differs"); + assert_transcodes(&compact, "raw metadata"); } } @@ -651,6 +669,7 @@ fn external_state_plane_round_trips_in_a_compact_container() { assert_eq!(rec.terminals[0].state_source, StateSource::External); assert_eq!(rec.terminals[0].external_state.as_ref(), Some(&ext)); assert!(to_compact(&compact) == compact, "rewrite differs"); + assert_transcodes(&compact, "external state"); // The other tensors are unaffected. for t in data.iter().filter(|t| t.name != target) { assert_eq!( @@ -905,3 +924,483 @@ fn write_fuzz_seeds() { write("parse_wire_compact", &format!("record_{i}"), &out); } } + +// ── Transcode (synthesis section 6) ───────────────────────────────────────── + +fn members_of(exploded: &Exploded) -> HashMap { + exploded + .members + .iter() + .map(|(k, v)| (k.clone(), v.as_slice())) + .collect() +} + +fn manifests_of(exploded: &Exploded) -> Vec { + exploded + .tensors + .iter() + .map(|t| t.manifest_json.clone()) + .collect() +} + +/// The members one tensor needs: its own (`.`) and the archive-wide +/// shared state, as the exploded stores keep them. +fn tensor_members<'a>(exploded: &'a Exploded, key: &str) -> HashMap { + let prefix = format!("{key}."); + exploded + .members + .iter() + .filter(|(k, _)| k.starts_with(&prefix) || k.starts_with("shared/")) + .map(|(k, v)| (k.clone(), v.as_slice())) + .collect() +} + +/// `implode(explode(blob)) == blob`, the manifest has the layout's form, +/// tensors come in body order, and each tensor imploded on its own (as the +/// exploded stores do) decodes like the full file. +fn assert_transcodes(blob: &[u8], what: &str) { + let reader = ContainerReader::open_partial(blob).unwrap(); + let exploded = explode(blob).unwrap_or_else(|e| panic!("{what}: explode: {e}")); + let rebuilt = implode( + &exploded.registry_json, + &manifests_of(&exploded), + members_of(&exploded), + ) + .unwrap_or_else(|e| panic!("{what}: implode: {e}")); + assert!(rebuilt == blob, "{what}: implode(explode(x)) != x"); + + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + match reader.layout { + ContainerLayout::Compact => { + assert_eq!(reg["ptwm_format"], "ptwm-manifest-c", "{what}"); + assert_eq!(reg["manifest_version"], 2, "{what}"); + assert_eq!(reg["container_flags"]["layout"], "compact", "{what}"); + assert!(reg.get("registry_hash_blake3").is_none(), "{what}"); + let table = reg["registry"]["extension_table"].as_array().unwrap(); + assert_eq!(table.len(), reader.extension_table.entries.len(), "{what}"); + let chains = reg["registry"]["chains"].as_array().unwrap(); + assert_eq!(chains.len(), reader.chain_registry.chains.len(), "{what}"); + for (pos, c) in chains.iter().enumerate() { + assert_eq!(c["id"], pos, "{what}: chains in registry order"); + } + } + ContainerLayout::Legacy => { + assert_eq!(reg["ptwm_format"], "ptwm-manifest", "{what}"); + assert_eq!(reg["manifest_version"], 1, "{what}"); + } + } + + // Tensors in body order. + let names: Vec = reader + .tensor_records() + .unwrap() + .into_iter() + .map(|r| r.name) + .collect(); + let exploded_names: Vec<&str> = exploded.tensors.iter().map(|t| t.name.as_str()).collect(); + assert_eq!(exploded_names, names, "{what}: tensor order"); + + // Per-tensor implode with the archive-wide registry. + for t in &exploded.tensors { + let single = implode( + &exploded.registry_json, + std::slice::from_ref(&t.manifest_json), + tensor_members(&exploded, &t.key), + ) + .unwrap_or_else(|e| panic!("{what}: implode {}: {e}", t.name)); + let one = ContainerReader::open_partial(&single).unwrap(); + assert_eq!(one.layout, reader.layout, "{what}"); + assert_eq!(one.tensor_names().unwrap(), vec![t.name.clone()], "{what}"); + match reader.decode_tensor(&t.name) { + Ok(want) => assert_eq!( + one.decode_tensor(&t.name).unwrap(), + want, + "{what}: subset implode of {}", + t.name + ), + Err(_) => assert!(one.decode_tensor(&t.name).is_err(), "{what}: {}", t.name), + } + } +} + +fn compact_feature_blob() -> Vec { + compress( + &feature_tensors(), + options(true, true, None, ContainerLayout::Compact), + ) +} + +#[test] +fn compact_manifest_lists_the_table_and_binds_sources() { + let blob = compact_feature_blob(); + let reader = ContainerReader::open(&blob).unwrap(); + let exploded = explode(&blob).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + assert_eq!( + reg["container_flags"], + serde_json::json!({"embedded_wasm": false, "layout": "compact"}) + ); + let keys: Vec<&String> = reg["registry"].as_object().unwrap().keys().collect(); + assert_eq!(keys, ["chains", "extension_table", "shared_state"]); + + // Every entry is a one-key object naming a built-in by wire id. + let table = reg["registry"]["extension_table"].as_array().unwrap(); + for (i, entry) in table.iter().enumerate() { + let obj = entry.as_object().unwrap(); + assert_eq!(obj.len(), 1, "{entry}"); + let (kind, value) = obj.iter().next().unwrap(); + let e = &reader.extension_table.entries[i]; + match kind.as_str() { + "op" => { + let op = OpId::from_u16(value.as_u64().unwrap() as u16).unwrap(); + assert_eq!( + e.human_label, + format!( + "io.ptwm.builtin.{}@{}", + op.builtin_name(), + env!("CARGO_PKG_VERSION") + ) + ); + } + "codec" => assert!(value.is_u64()), + other => panic!("unexpected entry kind {other} in {entry}"), + } + } + + // Registry chains with a bound Source carry params_b64 null. + let chains = reg["registry"]["chains"].as_array().unwrap(); + let bound = chains + .iter() + .filter(|c| { + c["nodes"] + .as_array() + .unwrap() + .iter() + .any(|n| n["op"] == OpId::Source.as_u16() && n["params_b64"].is_null()) + }) + .count(); + assert!(bound > 0, "no bound Source in {chains:?}"); + + // Plane codec_table_idx names a codec entry of the listed table. + for t in &exploded.tensors { + let tm: serde_json::Value = serde_json::from_str(&t.manifest_json).unwrap(); + for p in tm["planes"].as_array().unwrap() { + let idx = p["codec_table_idx"].as_u64().unwrap() as usize; + assert_eq!( + table[idx], + serde_json::json!({"codec": p["codec"]}), + "{}", + t.name + ); + } + } +} + +/// Rewrite `blob` in the compact layout with `extra` appended to its +/// Extension Table. Existing table indices stay valid. +fn with_extra_entries(blob: &[u8], extra: Vec) -> Vec { + let reader = ContainerReader::open(blob).unwrap(); + let mut table = reader.extension_table.clone(); + table.entries.extend(extra); + let mut out = Vec::new(); + let mut w = + ContainerWriter::with_layout(Cursor::new(&mut out), ContainerLayout::Compact).unwrap(); + w.set_prelude(&reader.prelude).unwrap(); + w.set_compact_tables(table, reader.chain_registry.clone()) + .unwrap(); + for rec in reader.tensor_records().unwrap() { + w.append_tensor(&rec).unwrap(); + } + w.finalize().unwrap(); + out +} + +fn third_party_entry() -> ExtensionTableEntry { + ExtensionTableEntry { + canonical_id: CanonicalId::derive(&[7; 32], "acme_codec", "0.1.0"), + human_label: "com.acme.codec@0.1.0".into(), + kind: Kind::PlaneCodec, + abi_version: 1, + lifecycle: Lifecycle::Thread, + flavor_hints: ptwm_core::extension::table::FLAVOR_NATIVE, + capabilities: CapabilityMap::new(), + attestation: Attestation::PgpSignature(vec![1, 2, 3]), + install_hint: Some("pip install acme".into()), + embedded_wasm_offset: None, + embedded_wasm_length: None, + } +} + +#[test] +fn compact_manifest_carries_named_and_third_party_entries() { + let blob = compact_feature_blob(); + let reader = ContainerReader::open(&blob).unwrap(); + // A built-in without a wire id (stored by name). + let named: Vec = builtin_entries() + .into_iter() + .filter(|e| { + matches!(classify(e), Ok(CompactEntry::Named(_))) + && !reader.extension_table.entries.contains(e) + }) + .take(1) + .collect(); + assert_eq!(named.len(), 1, "no built-in without a wire id"); + let mut extra = named.clone(); + extra.push(third_party_entry()); + let file = with_extra_entries(&blob, extra); + + let exploded = explode(&file).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + let table = reg["registry"]["extension_table"].as_array().unwrap(); + let last = table.last().unwrap(); + assert!(last["full_b64"].is_string(), "{last}"); + assert!(table[table.len() - 2]["named"].is_string(), "{table:?}"); + let rebuilt = implode( + &exploded.registry_json, + &manifests_of(&exploded), + members_of(&exploded), + ) + .unwrap(); + assert!(rebuilt == file, "implode(explode(x)) != x"); +} + +/// Implode `registry` (a JSON value) with the tensors of `exploded`. +fn implode_with( + exploded: &Exploded, + registry: &serde_json::Value, +) -> Result, ptwm_core::error::PtwmCoreError> { + implode( + ®istry.to_string(), + &manifests_of(exploded), + members_of(exploded), + ) +} + +#[test] +fn implode_rejects_malformed_compact_manifests() { + let blob = compact_feature_blob(); + let exploded = explode(&blob).unwrap(); + let reg: serde_json::Value = serde_json::from_str(&exploded.registry_json).unwrap(); + assert!(implode_with(&exploded, ®).unwrap() == blob); + + let cases: Vec<(&str, Box, &str)> = vec![ + ( + "compact format, legacy version", + Box::new(|r| r["manifest_version"] = 1.into()), + "unsupported manifest", + ), + ( + "compact format, version 3", + Box::new(|r| r["manifest_version"] = 3.into()), + "unsupported manifest", + ), + ( + "legacy format, compact version", + Box::new(|r| r["ptwm_format"] = "ptwm-manifest".into()), + "unsupported manifest", + ), + ( + "legacy layout name", + Box::new(|r| r["container_flags"]["layout"] = "legacy".into()), + "container_flags.layout", + ), + ( + "missing layout", + Box::new(|r| { + r["container_flags"] + .as_object_mut() + .unwrap() + .remove("layout"); + }), + "missing field", + ), + ( + "embedded WASM", + Box::new(|r| r["container_flags"]["embedded_wasm"] = true.into()), + "embedded WASM", + ), + ( + "legacy registry hash", + Box::new(|r| r["registry_hash_blake3"] = "00".repeat(32).into()), + "unknown field", + ), + ( + "legacy op list", + Box::new(|r| r["registry"]["ops"] = serde_json::json!([])), + "unknown field", + ), + ( + "unknown entry kind", + Box::new(|r| r["registry"]["extension_table"][0] = serde_json::json!({"wasm": 1})), + "unknown variant", + ), + ( + "two keys in an entry", + Box::new(|r| { + r["registry"]["extension_table"][0] = serde_json::json!({"op": 1, "codec": 1}) + }), + "parse registry manifest", + ), + ( + "unknown op id", + Box::new(|r| r["registry"]["extension_table"][0] = serde_json::json!({"op": 0x7FFF})), + "unknown op id", + ), + ( + "unknown built-in name", + Box::new(|r| { + r["registry"]["extension_table"][0] = serde_json::json!({"named": "no_such"}) + }), + "unknown built-in name", + ), + ( + "duplicate entry", + Box::new(|r| { + let first = r["registry"]["extension_table"][0].clone(); + r["registry"]["extension_table"] + .as_array_mut() + .unwrap() + .push(first); + }), + "duplicate entry", + ), + ( + "chain id not its position", + Box::new(|r| r["registry"]["chains"][0]["id"] = 7.into()), + "compact chain ids are positions", + ), + ]; + for (what, mutate, needle) in cases { + let mut bad = reg.clone(); + mutate(&mut bad); + let e = implode_with(&exploded, &bad).unwrap_err().to_string(); + assert!(e.contains(needle), "{what}: {e}"); + } + + // A built-in stored as a full entry is not canonical. + let mut bad = reg.clone(); + let zstd = builtin_registry() + .entry_by_name("zstd") + .unwrap() + .to_bytes() + .unwrap(); + bad["registry"]["extension_table"] + .as_array_mut() + .unwrap() + .push(serde_json::json!({ "full_b64": b64(&zstd) })); + let e = implode_with(&exploded, &bad).unwrap_err().to_string(); + assert!(e.contains("built-in id"), "{e}"); + + // codec is informational, but must agree with the table entry that + // codec_table_idx names. + let mut tm: serde_json::Value = + serde_json::from_str(&exploded.tensors[0].manifest_json).unwrap(); + tm["planes"][0]["codec"] = 0x7FFE.into(); + let e = implode( + &exploded.registry_json, + &[tm.to_string()], + tensor_members(&exploded, &exploded.tensors[0].key), + ) + .unwrap_err() + .to_string(); + assert!(e.contains("codec"), "{e}"); + let mut tm: serde_json::Value = + serde_json::from_str(&exploded.tensors[0].manifest_json).unwrap(); + let other = if tm["planes"][0]["codec"] == CodecId::Zstd.as_u16() { + CodecId::Identity + } else { + CodecId::Zstd + }; + tm["planes"][0]["codec"] = other.as_u16().into(); + let e = implode( + &exploded.registry_json, + &[tm.to_string()], + tensor_members(&exploded, &exploded.tensors[0].key), + ) + .unwrap_err() + .to_string(); + assert!(e.contains("codec"), "{e}"); +} + +/// Standard base64 with padding, for building manifests by hand. +fn b64(bytes: &[u8]) -> String { + const A: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/"; + let mut out = String::new(); + for chunk in bytes.chunks(3) { + let b = [ + chunk[0], + *chunk.get(1).unwrap_or(&0), + *chunk.get(2).unwrap_or(&0), + ]; + let n = (u32::from(b[0]) << 16) | (u32::from(b[1]) << 8) | u32::from(b[2]); + for i in 0..4 { + if i <= chunk.len() { + out.push(A[(n >> (18 - 6 * i)) as usize & 63] as char); + } else { + out.push('='); + } + } + } + out +} + +/// The registry-manifest gate of the v1.1 `implode`, frozen: the v1.1 +/// manifest struct (no `deny_unknown_fields`) and its `ptwm_format` check. +fn v1_1_implode_gate(registry_json: &str) -> Result<(), String> { + #[derive(serde::Deserialize)] + #[allow(dead_code)] + struct Flags { + embedded_wasm: bool, + } + #[derive(serde::Deserialize)] + #[allow(dead_code)] + struct Registry { + ops: Vec, + codecs: Vec, + shared_state: Vec, + chains: Vec, + } + #[derive(serde::Deserialize)] + #[allow(dead_code)] + struct Manifest { + ptwm_format: String, + manifest_version: u32, + container_flags: Flags, + registry_hash_blake3: String, + registry: Registry, + } + let m: Manifest = serde_json::from_str(registry_json).map_err(|e| e.to_string())?; + if m.ptwm_format != "ptwm-manifest" { + return Err(format!("unexpected ptwm_format {:?}", m.ptwm_format)); + } + Ok(()) +} + +/// The gate of an implode that accepts only `("ptwm-manifest", 1)`. +fn legacy_only_implode_gate(registry_json: &str) -> Result<(), String> { + let v: serde_json::Value = serde_json::from_str(registry_json).map_err(|e| e.to_string())?; + if v["ptwm_format"] != "ptwm-manifest" || v["manifest_version"] != 1 { + return Err(format!( + "unsupported manifest {} version {}", + v["ptwm_format"], v["manifest_version"] + )); + } + Ok(()) +} + +#[test] +fn old_implode_rejects_compact_manifests() { + let data = feature_tensors(); + let legacy = compress(&data, options(true, true, None, ContainerLayout::Legacy)); + let legacy_reg = explode(&legacy).unwrap().registry_json; + let compact_reg = explode(&compact_feature_blob()).unwrap().registry_json; + + for gate in [v1_1_implode_gate, legacy_only_implode_gate] { + assert_eq!(gate(&legacy_reg), Ok(())); + let e = gate(&compact_reg).unwrap_err(); + assert!(!e.is_empty()); + } + // The v1.1 gate fails on the manifest structure before the format check. + let e = v1_1_implode_gate(&compact_reg).unwrap_err(); + assert!(e.contains("missing field"), "{e}"); +} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index e155f15..8378b36 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -317,7 +317,7 @@ fn py_dtype_element_size(dtype_code: u8) -> PyResult { /// Returns the entire container as a ``bytes`` object. #[allow(clippy::type_complexity, clippy::too_many_arguments)] #[pyfunction] -#[pyo3(signature = (records, chains_per_tensor, emit_payload_hash, emit_plane_crc, method_hint=3, streaming_chunk=None, codec_menu=None))] +#[pyo3(signature = (records, chains_per_tensor, emit_payload_hash, emit_plane_crc, method_hint=3, streaming_chunk=None, codec_menu=None, *, container_layout="legacy"))] fn compress_model<'py>( py: Python<'py>, records: Vec<( @@ -338,7 +338,15 @@ fn compress_model<'py>( // trial-encode menu is filtered to codecs in this set on every plane; // empty intersection on a plane raises `InvalidContainer`. codec_menu: Option>, + // Container layout: "legacy" (\x89PTWM) or "compact" (\x89PTWC). + container_layout: &str, ) -> PyResult> { + let container_layout = ptwm_core::container::ContainerLayout::parse(container_layout) + .ok_or_else(|| { + pyo3::exceptions::PyValueError::new_err(format!( + "unknown container_layout {container_layout:?} (expected \"legacy\" or \"compact\")" + )) + })?; if chains_per_tensor.len() != records.len() { return Err(pyo3::exceptions::PyValueError::new_err(format!( "chains_per_tensor length {} != records length {}", @@ -438,7 +446,7 @@ fn compress_model<'py>( forced_codec, allow_codec_ids, chunk_size: streaming_chunk, - ..Default::default() + container_layout, }; let cursor = std::io::Cursor::new(&mut out); diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index 8045e21..58fefec 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -190,8 +190,10 @@ def compress_model( method_hint: int = 3, streaming_chunk: int | None = None, codec_menu: list[int] | None = None, + *, + container_layout: str = "legacy", ) -> bytes: - """Compress a list of tensors into a ``.ptwm`` container. + r"""Compress a list of tensors into a ``.ptwm`` container. Each element of ``records`` is a 7-tuple: ``(name, dtype_code, input_format, raw_bytes, shape, @@ -204,6 +206,9 @@ def compress_model( ``codec_menu``: optional list of wire-stable ``CodecId`` integers restricting the trial-encode menu. ``None`` uses the full menu; empty intersection on a plane raises ``InvalidContainer``. + ``container_layout``: ``"legacy"`` writes ``\x89PTWM`` (readable by every + release), ``"compact"`` writes ``\x89PTWC`` (readable by v1.2 and later). + Other values raise ``ValueError``. """ return bytes( _ext.compress_model( @@ -214,6 +219,7 @@ def compress_model( method_hint, streaming_chunk, codec_menu, + container_layout=container_layout, ) ) @@ -352,13 +358,19 @@ def decode_tensor_info( def explode_ptwm( blob: bytes, ) -> tuple[str, list[tuple[str, str, str]], dict[str, bytes]]: - """Explode a ``.ptwm`` container into its canonical PTWM-JSON projection. + r"""Explode a ``.ptwm`` container into its canonical PTWM-JSON projection. Returns ``(registry_json, tensors, members)`` where ``registry_json`` is the canonical (JCS) archive-scoped registry manifest, ``tensors`` is a list of ``(name, key, tensor_manifest_json)`` in file order, and ``members`` maps member name → bytes. ``implode_ptwm`` reproduces ``blob`` byte-for-byte. + A legacy (``\x89PTWM``) container gives a ``"ptwm-manifest"`` version 1 + registry manifest; a compact (``\x89PTWC``) container gives a + ``"ptwm-manifest-c"`` version 2 manifest that lists the Extension Table + entries and records ``container_flags.layout == "compact"``. Tensors are + in body order for both. + The extension already returns ``str``/``bytes`` of the correct type, so the result is returned without re-wrapping (avoids copying large payloads). """ @@ -374,5 +386,7 @@ def implode_ptwm( """Reconstruct a ``.ptwm`` container from a manifest projection. ``tensor_jsons`` must be in file order. Inverse of :func:`explode_ptwm`. + The registry manifest selects the layout of the result. A subset of the + tensors gives a container of those tensors with the archive-wide tables. """ return _ext.implode_ptwm(registry_json, tensor_jsons, members) diff --git a/python/ptwm/sharding/_writer.py b/python/ptwm/sharding/_writer.py index 405a415..6a7c87f 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -40,14 +40,17 @@ def compress_ptwm_blob( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, + container_layout: str = "legacy", ) -> bytes: - """Compress ``tensors`` into a single in-memory ``.ptwm`` container blob. + r"""Compress ``tensors`` into a single in-memory ``.ptwm`` container blob. ``tensors`` is a sequence of ``(name, raw_bytes, dtype, shape)``. ``chains_per_tensor`` — one list of chain wire-byte blobs per tensor; each blob is produced by ``Chain.to_bytes()``. When ``None``, chains are auto-derived from each tensor's dtype via the STANDARD production table. ``method_hint``: 1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY. + ``container_layout``: ``"legacy"`` (``\x89PTWM``) or ``"compact"`` + (``\x89PTWC``). """ records = [] auto_chains: list[list[bytes]] = [] @@ -77,6 +80,7 @@ def compress_ptwm_blob( emit_plane_crc, method_hint, streaming_chunk, + container_layout=container_layout, ) @@ -89,6 +93,7 @@ def write_ptwm_shard( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, + container_layout: str = "legacy", ) -> None: """Write a single ``.ptwm`` container holding ``tensors``. @@ -101,5 +106,6 @@ def write_ptwm_shard( emit_payload_hash=emit_payload_hash, emit_plane_crc=emit_plane_crc, streaming_chunk=streaming_chunk, + container_layout=container_layout, ) Path(out_path).write_bytes(blob) diff --git a/python/ptwm/stores/_common.py b/python/ptwm/stores/_common.py index 89ddcbb..af52106 100644 --- a/python/ptwm/stores/_common.py +++ b/python/ptwm/stores/_common.py @@ -25,6 +25,7 @@ "decode_member_tensor", "raw_from_members", "jcs_canonicalize", + "manifest_container_layout", "shared_member_keys", "tensor_member_keys", ] @@ -66,6 +67,11 @@ class ExplodedModel: __slots__ = ("registry_json", "tensors", "members", "_by_name") + @property + def container_layout(self) -> str: + """The source container's layout, ``"legacy"`` or ``"compact"``.""" + return manifest_container_layout(self.registry_json) + def __init__( self, registry_json: str, @@ -99,6 +105,22 @@ def explode_model(blob: bytes) -> ExplodedModel: return ExplodedModel(registry_json, list(tensors), dict(members)) +def manifest_container_layout(registry_json: str) -> str: + """Return the container layout a registry manifest rebuilds. + + ``"compact"`` for a ``"ptwm-manifest-c"`` manifest, ``"legacy"`` for a + ``"ptwm-manifest"`` one; raises ``ValueError`` for anything else. + """ + reg = json.loads(registry_json) + fmt = reg.get("ptwm_format") + if fmt == "ptwm-manifest-c": + return "compact" + if fmt == "ptwm-manifest": + return "legacy" + msg = f"unknown registry manifest format {fmt!r}" + raise ValueError(msg) + + def tensor_member_keys(members: dict[str, bytes], sample_key: str) -> dict[str, bytes]: """Return the subset of ``members`` owned by the tensor ``sample_key``.""" prefix = f"{sample_key}." @@ -121,7 +143,10 @@ def raw_from_members( """Reconstruct a one-tensor container and return its raw decompressed bytes. ``members`` must contain the tensor's plane/state members plus any - archive-global shared state referenced by ``registry_json``. + archive-global shared state referenced by ``registry_json``. The container + has the layout of ``registry_json`` (legacy or compact) and the archive-wide + Extension Table and chain registry, so the tensor's table indices and chain + reference stay valid. """ blob = _rust.implode_ptwm(registry_json, [tensor_json], members) name = json.loads(tensor_json)["name"] diff --git a/python/ptwm/stores/_lmdb.py b/python/ptwm/stores/_lmdb.py index 4fc56f3..1ede0ae 100644 --- a/python/ptwm/stores/_lmdb.py +++ b/python/ptwm/stores/_lmdb.py @@ -76,10 +76,14 @@ def write_lmdb( map_size: int = _DEFAULT_MAP_SIZE, method_hint: int = 3, chains_per_tensor: list[list[bytes]] | None = None, + container_layout: str = "legacy", ) -> Path: """Write ``tensors`` into an LMDB environment at ``path``. - ``layout`` is ``"blob"`` (default) or ``"exploded"``. Returns ``path``. + ``layout`` is ``"blob"`` (default) or ``"exploded"``. ``container_layout`` + (``"legacy"`` or ``"compact"``) selects the container layout of the + per-tensor blobs (blob layout) or of the registry manifest the tensors are + rebuilt from (exploded layout). Returns ``path``. """ if layout not in ("blob", "exploded"): msg = f"unknown LMDB layout {layout!r} (expected 'blob' or 'exploded')" @@ -88,7 +92,10 @@ def write_lmdb( lmdb = _require_lmdb() path = Path(path) blob = compress_ptwm_blob( - tensors, chains_per_tensor=chains_per_tensor, method_hint=method_hint + tensors, + chains_per_tensor=chains_per_tensor, + method_hint=method_hint, + container_layout=container_layout, ) model = explode_model(blob) @@ -359,13 +366,15 @@ def put_tensor( *, delta_reference: str | None = None, method_hint: int = 3, + container_layout: str = "legacy", ) -> None: """Atomically insert or replace tensor ``name`` (blob layout only). When ``delta_reference`` names another tensor already in the store, the new tensor is stored as an XOR residual against it (a ``local_tensor`` reference with BLAKE3 verification), turning small per-step changes into - small diffs. + small diffs. ``container_layout`` (``"legacy"`` or ``"compact"``) + selects the layout of the stored one-tensor container. """ if self._readonly: msg = "store opened read-only; reopen with readonly=False to mutate" @@ -392,7 +401,9 @@ def put_tensor( stored, store_shape, store_dtype = raw, shape, dtype blob = compress_ptwm_blob( - [(name, stored, store_dtype, store_shape)], method_hint=method_hint + [(name, stored, store_dtype, store_shape)], + method_hint=method_hint, + container_layout=container_layout, ) entry["comp_size"] = len(blob) with self._env.begin(write=True) as txn: diff --git a/python/ptwm/stores/_webdataset.py b/python/ptwm/stores/_webdataset.py index 4bf3a7c..678dbd4 100644 --- a/python/ptwm/stores/_webdataset.py +++ b/python/ptwm/stores/_webdataset.py @@ -70,18 +70,25 @@ def write_webdataset( max_shard_size: int = _DEFAULT_MAX_SHARD_SIZE, method_hint: int = 3, chains_per_tensor: list[list[bytes]] | None = None, + container_layout: str = "legacy", ) -> Path: """Write ``tensors`` as a WebDataset shard set rooted at ``out_dir``. Returns the directory path. ``tensors`` is a sequence of ``(name, raw_bytes, dtype, shape)``; shards are packed by compressed size via :func:`ptwm.sharding.plan_shards` so a tensor never straddles a shard. + ``container_layout`` (``"legacy"`` or ``"compact"``) selects the + container the tensors are compressed into; the registry manifest records + it, and readers rebuild each tensor in that layout. """ out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) blob = compress_ptwm_blob( - tensors, chains_per_tensor=chains_per_tensor, method_hint=method_hint + tensors, + chains_per_tensor=chains_per_tensor, + method_hint=method_hint, + container_layout=container_layout, ) model = explode_model(blob) _write_exploded(out_dir, model, max(1, max_shard_size)) diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py new file mode 100644 index 0000000..f662a57 --- /dev/null +++ b/tests/ptwm/stores/test_compact_layout.py @@ -0,0 +1,214 @@ +"""Stores and the transcoder over compact (``\\x89PTWC``) containers. + +The LMDB and WebDataset stores rebuild tensors through ``implode_ptwm``. A +compact container explodes to a ``"ptwm-manifest-c"`` registry manifest, and +every store path must rebuild tensors from it as it does from a legacy one. +""" + +from __future__ import annotations + +import json +import tarfile + +import pytest +import torch +from ptwm import _rust +from ptwm._exceptions import Error +from ptwm.sharding import compress_ptwm_blob, write_ptwm_shard +from ptwm.stores import ( + LmdbWeightStore, + WebDatasetWeightStore, + explode_model, + jcs_canonicalize, + raw_from_members, + write_lmdb, + write_webdataset, +) +from ptwm.stores._common import ( + manifest_container_layout, + shared_member_keys, + tensor_member_keys, +) +from ptwm.stores._lmdb import KEY_REG +from ptwm.stores._webdataset import INDEX_NAME, SHARDMETA_MEMBER + +MAGIC = {"legacy": b"\x89PTWM", "compact": b"\x89PTWC"} +LAYOUTS = ("legacy", "compact") + + +def _reference_raw(store_tensors, layout: str) -> dict[str, bytes]: + blob = compress_ptwm_blob(store_tensors, method_hint=3, container_layout=layout) + return {name: bytes(_rust.decode_tensor(blob, name)) for name, *_ in store_tensors} + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_compress_ptwm_blob_honours_the_layout(tmp_path, store_tensors, layout): + blob = compress_ptwm_blob(store_tensors, method_hint=3, container_layout=layout) + assert blob[:9].startswith(MAGIC[layout]) + out = tmp_path / "m.ptwm" + write_ptwm_shard(out, store_tensors, container_layout=layout) + assert out.read_bytes()[:5] == MAGIC[layout] + # Both layouts decode to the same bytes. + assert _reference_raw(store_tensors, layout) == _reference_raw( + store_tensors, "legacy" + ) + + +def test_unknown_container_layout_is_rejected(store_tensors): + with pytest.raises(Error, match="container_layout"): + compress_ptwm_blob(store_tensors, container_layout="bogus") + + +def test_compact_manifest_form_and_identity(store_tensors): + blob = compress_ptwm_blob(store_tensors, method_hint=3, container_layout="compact") + registry_json, tensors, members = _rust.explode_ptwm(blob) + reg = json.loads(registry_json) + assert reg["ptwm_format"] == "ptwm-manifest-c" + assert reg["manifest_version"] == 2 + assert reg["container_flags"] == {"embedded_wasm": False, "layout": "compact"} + assert "registry_hash_blake3" not in reg + assert set(reg["registry"]) == {"chains", "extension_table", "shared_state"} + for entry in reg["registry"]["extension_table"]: + assert len(entry) == 1 + assert next(iter(entry)) in {"op", "codec", "named", "full_b64"} + # The Rust JCS output and the Python canonicalizer agree. + assert jcs_canonicalize(reg) == registry_json + # Tensors come in body (write) order. + assert [name for (name, _, _) in tensors] == [n for n, *_ in store_tensors] + + rebuilt = _rust.implode_ptwm(registry_json, [tj for (_, _, tj) in tensors], members) + assert rebuilt == blob + + model = explode_model(blob) + assert model.container_layout == "compact" + assert manifest_container_layout(registry_json) == "compact" + legacy = compress_ptwm_blob(store_tensors, method_hint=3) + assert explode_model(legacy).container_layout == "legacy" + + +@pytest.mark.parametrize("method_hint", [1, 2, 3, 4, 5]) +def test_compact_identity_per_method(store_tensors, method_hint): + blob = compress_ptwm_blob( + store_tensors, method_hint=method_hint, container_layout="compact" + ) + registry_json, tensors, members = _rust.explode_ptwm(blob) + rebuilt = _rust.implode_ptwm(registry_json, [tj for (_, _, tj) in tensors], members) + assert rebuilt == blob + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_raw_from_members_rebuilds_one_tensor(store_tensors, layout): + blob = compress_ptwm_blob(store_tensors, method_hint=3, container_layout=layout) + want = _reference_raw(store_tensors, layout) + model = explode_model(blob) + shared = shared_member_keys(model.members) + for name in model.names(): + members = tensor_member_keys(model.members, model.key(name)) + members.update(shared) + single = _rust.implode_ptwm( + model.registry_json, [model.tensor_json(name)], members + ) + assert single[:5] == MAGIC[layout] + raw = raw_from_members(model.registry_json, model.tensor_json(name), members) + assert raw == want[name] + + +def test_implode_rejects_a_compact_manifest_with_another_layout(store_tensors): + blob = compress_ptwm_blob(store_tensors, container_layout="compact") + registry_json, tensors, members = _rust.explode_ptwm(blob) + reg = json.loads(registry_json) + reg["container_flags"]["layout"] = "legacy" + with pytest.raises(Error, match="container_flags.layout"): + _rust.implode_ptwm(json.dumps(reg), [tj for (_, _, tj) in tensors], members) + + +# ── LMDB ────────────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("store_layout", ["blob", "exploded"]) +def test_lmdb_compact_round_trip(tmp_path, store_tensors, torch_tensors, store_layout): + path = tmp_path / "m.lmdb" + write_lmdb(path, store_tensors, layout=store_layout, container_layout="compact") + want = _reference_raw(store_tensors, "compact") + store = LmdbWeightStore.open(path) + try: + assert set(store.names()) == {n for n, *_ in store_tensors} + if store_layout == "blob": + for name, *_ in store_tensors: + assert store._get(b"t/" + name.encode())[:5] == MAGIC["compact"] + else: + reg = json.loads(store._get(KEY_REG)) + assert reg["ptwm_format"] == "ptwm-manifest-c" + for name, *_ in store_tensors: + assert store.get_bytes(name) == want[name] + for name, expected in torch_tensors.items(): + got = store.get_tensor(name) + assert got.dtype == expected.dtype + assert torch.equal(got.view(torch.uint8), expected.view(torch.uint8)) + finally: + store.close() + + +def test_lmdb_put_tensor_in_either_layout(tmp_path, raw_of): + tensors = [("w", raw_of(torch.zeros(8).float()), torch.float32, (8,))] + path = tmp_path / "m.lmdb" + write_lmdb(path, tensors, layout="blob", container_layout="compact") + + new = torch.arange(8, dtype=torch.float32) + step = new + 1 + with LmdbWeightStore.open(path, readonly=False) as store: + store.put_tensor("w", raw_of(new), torch.float32, (8,)) + assert store._get(b"t/w")[:5] == MAGIC["legacy"] + store.put_tensor( + "v", raw_of(new), torch.float32, (8,), container_layout="compact" + ) + assert store._get(b"t/v")[:5] == MAGIC["compact"] + store.put_tensor( + "d", + raw_of(step), + torch.float32, + (8,), + delta_reference="v", + container_layout="compact", + ) + assert store._get(b"t/d")[:5] == MAGIC["compact"] + with LmdbWeightStore.open(path) as store: + assert torch.equal(store.get_tensor("w"), new) + assert torch.equal(store.get_tensor("v"), new) + assert torch.equal(store.get_tensor("d"), step) + + +# ── WebDataset ──────────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("max_shard_size", [2 * 1024**3, 1]) +def test_webdataset_compact_round_trip(tmp_path, store_tensors, max_shard_size): + wd = tmp_path / "m.wds" + write_webdataset( + wd, store_tensors, max_shard_size=max_shard_size, container_layout="compact" + ) + want = _reference_raw(store_tensors, "compact") + + index = json.loads((wd / INDEX_NAME).read_text()) + assert index["manifest_version"] == 1 # describes the store + assert index["registry"]["ptwm_format"] == "ptwm-manifest-c" + shards = sorted(wd.glob("weights-*-of-*.tar")) + assert (len(shards) > 1) == (max_shard_size == 1) + for shard in shards: + with tarfile.open(shard, "r") as tar: + f = tar.extractfile(SHARDMETA_MEMBER) + assert f is not None + meta = json.loads(f.read()) + assert meta["container_flags"]["layout"] == "compact" + + store = WebDatasetWeightStore.open(wd) + for name, *_ in store_tensors: + assert store.get_bytes(name) == want[name] + streamed = {name for name, _ in store.stream_tensors()} + assert streamed == {n for n, *_ in store_tensors} + + # One shard opened on its own streams with its own registry copy. + single = WebDatasetWeightStore.open(shards[0]) + for name, tensor in single.stream_tensors(): + assert tensor.numel() > 0 + assert name in want From f7ba328410abe0483714b2088e1546264b25150f Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 01:25:23 +0000 Subject: [PATCH 18/34] fix(S7-transcode-c): address review findings - compress_model docstring: an unknown container_layout raises ptwm.Error (translate_errors wraps the binding's ValueError), not ValueError. - LMDB: write_lmdb(container_layout="compact") records the layout in the \x00fmt record; LmdbWeightStore.put_tensor now defaults container_layout to None and uses the store's recorded layout ("legacy" when none is recorded). The fmt record of a legacy store is unchanged. - Tests: put_tensor inherits the store layout (legacy and compact, incl. delta and reopen); compress_model error type for an unknown layout. test_lmdb_put_tensor_in_either_layout now passes container_layout="legacy" explicitly, because the default it pinned changed deliberately. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- python/ptwm/_rust/__init__.py | 3 +- python/ptwm/stores/_lmdb.py | 36 +++++++++++-------- tests/ptwm/stores/test_compact_layout.py | 44 ++++++++++++++++++++++-- 3 files changed, 66 insertions(+), 17 deletions(-) diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index 58fefec..89b4050 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -208,7 +208,8 @@ def compress_model( empty intersection on a plane raises ``InvalidContainer``. ``container_layout``: ``"legacy"`` writes ``\x89PTWM`` (readable by every release), ``"compact"`` writes ``\x89PTWC`` (readable by v1.2 and later). - Other values raise ``ValueError``. + Other values raise :class:`ptwm.Error` (the binding's ``ValueError`` + passes through ``translate_errors``). """ return bytes( _ext.compress_model( diff --git a/python/ptwm/stores/_lmdb.py b/python/ptwm/stores/_lmdb.py index 1ede0ae..8e48b49 100644 --- a/python/ptwm/stores/_lmdb.py +++ b/python/ptwm/stores/_lmdb.py @@ -78,12 +78,14 @@ def write_lmdb( chains_per_tensor: list[list[bytes]] | None = None, container_layout: str = "legacy", ) -> Path: - """Write ``tensors`` into an LMDB environment at ``path``. + r"""Write ``tensors`` into an LMDB environment at ``path``. ``layout`` is ``"blob"`` (default) or ``"exploded"``. ``container_layout`` (``"legacy"`` or ``"compact"``) selects the container layout of the per-tensor blobs (blob layout) or of the registry manifest the tensors are - rebuilt from (exploded layout). Returns ``path``. + rebuilt from (exploded layout). A ``"compact"`` store records the layout + in its ``\x00fmt`` record so that later :meth:`LmdbWeightStore.put_tensor` + calls use it; a ``"legacy"`` store's record is unchanged. Returns ``path``. """ if layout not in ("blob", "exploded"): msg = f"unknown LMDB layout {layout!r} (expected 'blob' or 'exploded')" @@ -102,16 +104,16 @@ def write_lmdb( env = lmdb.open(str(path), map_size=map_size, subdir=True) try: with env.begin(write=True) as txn: - txn.put( - KEY_FMT, - jcs_canonicalize( - { - "ptwm_format": LMDB_FORMAT, - "manifest_version": 1, - "layout": layout, - } - ).encode("utf-8"), - ) + fmt_obj: dict[str, Any] = { + "ptwm_format": LMDB_FORMAT, + "manifest_version": 1, + "layout": layout, + } + # Recorded only for the compact layout: a legacy store's record + # stays byte-identical to the one older releases write. + if container_layout != "legacy": + fmt_obj["container_layout"] = container_layout + txn.put(KEY_FMT, jcs_canonicalize(fmt_obj).encode("utf-8")) if layout == "blob": _write_blob_layout(txn, model) else: @@ -220,6 +222,8 @@ def __init__( msg = f"{self._path}: not a {LMDB_FORMAT} environment" raise ValueError(msg) self._layout = fmt_obj["layout"] + # Stores written before the field existed hold legacy blobs. + self._container_layout = fmt_obj.get("container_layout", "legacy") idx = self._get(KEY_IDX) self._index = json.loads(idx) if idx is not None else {"weight_map": {}} self._weight_map = dict(self._index.get("weight_map", {})) @@ -366,7 +370,7 @@ def put_tensor( *, delta_reference: str | None = None, method_hint: int = 3, - container_layout: str = "legacy", + container_layout: str | None = None, ) -> None: """Atomically insert or replace tensor ``name`` (blob layout only). @@ -374,7 +378,9 @@ def put_tensor( new tensor is stored as an XOR residual against it (a ``local_tensor`` reference with BLAKE3 verification), turning small per-step changes into small diffs. ``container_layout`` (``"legacy"`` or ``"compact"``) - selects the layout of the stored one-tensor container. + selects the layout of the stored one-tensor container; ``None`` + (default) uses the layout the store was written with (see + :func:`write_lmdb`), or ``"legacy"`` for a store that records none. """ if self._readonly: msg = "store opened read-only; reopen with readonly=False to mutate" @@ -400,6 +406,8 @@ def put_tensor( else: stored, store_shape, store_dtype = raw, shape, dtype + if container_layout is None: + container_layout = self._container_layout blob = compress_ptwm_blob( [(name, stored, store_dtype, store_shape)], method_hint=method_hint, diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py index f662a57..1192415 100644 --- a/tests/ptwm/stores/test_compact_layout.py +++ b/tests/ptwm/stores/test_compact_layout.py @@ -29,7 +29,7 @@ shared_member_keys, tensor_member_keys, ) -from ptwm.stores._lmdb import KEY_REG +from ptwm.stores._lmdb import KEY_FMT, KEY_REG from ptwm.stores._webdataset import INDEX_NAME, SHARDMETA_MEMBER MAGIC = {"legacy": b"\x89PTWM", "compact": b"\x89PTWC"} @@ -157,7 +157,9 @@ def test_lmdb_put_tensor_in_either_layout(tmp_path, raw_of): new = torch.arange(8, dtype=torch.float32) step = new + 1 with LmdbWeightStore.open(path, readonly=False) as store: - store.put_tensor("w", raw_of(new), torch.float32, (8,)) + store.put_tensor( + "w", raw_of(new), torch.float32, (8,), container_layout="legacy" + ) assert store._get(b"t/w")[:5] == MAGIC["legacy"] store.put_tensor( "v", raw_of(new), torch.float32, (8,), container_layout="compact" @@ -178,6 +180,44 @@ def test_lmdb_put_tensor_in_either_layout(tmp_path, raw_of): assert torch.equal(store.get_tensor("d"), step) +@pytest.mark.parametrize("layout", ["legacy", "compact"]) +def test_lmdb_put_tensor_inherits_store_layout(tmp_path, raw_of, layout): + # put_tensor without container_layout uses the layout the store was + # written with, so a compact store does not silently gain legacy blobs. + tensors = [("w", raw_of(torch.zeros(8).float()), torch.float32, (8,))] + path = tmp_path / "m.lmdb" + write_lmdb(path, tensors, layout="blob", container_layout=layout) + + new = torch.arange(8, dtype=torch.float32) + with LmdbWeightStore.open(path, readonly=False) as store: + fmt = json.loads(store._get(KEY_FMT)) + if layout == "legacy": + # The legacy fmt record is unchanged from earlier releases. + assert "container_layout" not in fmt + else: + assert fmt["container_layout"] == "compact" + store.put_tensor("w", raw_of(new), torch.float32, (8,)) + store.put_tensor("d", raw_of(new + 1), torch.float32, (8,), delta_reference="w") + assert store._get(b"t/w")[:5] == MAGIC[layout] + assert store._get(b"t/d")[:5] == MAGIC[layout] + # The layout survives a reopen. + with LmdbWeightStore.open(path, readonly=False) as store: + store.put_tensor("v", raw_of(new), torch.float32, (8,)) + assert store._get(b"t/v")[:5] == MAGIC[layout] + with LmdbWeightStore.open(path) as store: + assert torch.equal(store.get_tensor("w"), new) + assert torch.equal(store.get_tensor("v"), new) + assert torch.equal(store.get_tensor("d"), new + 1) + + +def test_compress_model_unknown_layout_raises_ptwm_error(): + # The public wrapper translates the binding's ValueError to ptwm.Error. + with pytest.raises(Error, match="container_layout"): + _rust.compress_model([], [], True, True, container_layout="bogus") + with pytest.raises(ValueError, match="container_layout"): + _rust._ext.compress_model([], [], True, True, container_layout="bogus") + + # ── WebDataset ──────────────────────────────────────────────────────────────── From d20f92cbf3b79d1d985dab430cc28d6e9f0a4b5b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 01:32:37 +0000 Subject: [PATCH 19/34] wip(S8-python): snapshot of an interrupted implementation Partial Python and binding work (ContainerHandle, _magic, layout plumbing), saved so a container restart cannot lose it. May not build yet; the resumed Stage 8 implementer completes it. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/container.rs | 123 ++++++------ crates/ptwm-core/tests/ptwc.rs | 98 +++++++++ crates/ptwm-py/src/handle.rs | 273 ++++++++++++++++++++++++++ crates/ptwm-py/src/inspect.rs | 65 +----- crates/ptwm-py/src/lib.rs | 200 ++++--------------- python/ptwm/__init__.py | 9 +- python/ptwm/_config.py | 59 +++++- python/ptwm/_magic.py | 72 +++++++ python/ptwm/_rust/__init__.py | 101 +++++++++- python/ptwm/core/_compressor.py | 1 + python/ptwm/core/_decompressor.py | 42 ++-- python/ptwm/random_access/__init__.py | 37 ++-- python/ptwm/stores/__init__.py | 6 +- 13 files changed, 768 insertions(+), 318 deletions(-) create mode 100644 crates/ptwm-py/src/handle.rs create mode 100644 python/ptwm/_magic.py diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 6b43609..8909365 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -1222,6 +1222,58 @@ fn check_extensions( Ok((installed, needs)) } +/// Parse the Extension Table of a legacy container whose header is +/// `header`. The region is bounds-checked and its BLAKE3 hash verified. +/// Built-in ids are not normalized. +fn legacy_extension_table(buf: &[u8], header: &Header) -> Result { + if header.extension_table_length == 0 { + return Ok(ExtensionTable::empty()); + } + let table_bytes = usize::try_from(header.extension_table_offset) + .ok() + .zip(usize::try_from(header.extension_table_length).ok()) + .and_then(|(start, len)| buf.get(start..start.checked_add(len)?)) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "extension table region at offset {} with length {} exceeds the buffer ({} bytes)", + header.extension_table_offset, + header.extension_table_length, + buf.len() + )) + })?; + if ExtensionTable::hash(table_bytes) != header.extension_table_hash { + return Err(PtwmCoreError::ExtensionTableHashMismatch); + } + ExtensionTable::from_bytes(table_bytes) + .map_err(|e| PtwmCoreError::InvalidContainer(format!("extension table: {e}"))) +} + +/// Read the Extension Table of a container in either layout without opening +/// the rest of it: no extension check, prelude, registry or index parse. +/// +/// A legacy (`\x89PTWM` / `\x89PTWX`) file has its table region +/// bounds-checked and its hash verified, and built-in ids written by other +/// releases are mapped to the current ones (see +/// [`crate::extension::normalize_builtin_ids`]). A compact (`\x89PTWC`) file +/// has its EOF sentinel, sections and `meta_xxh64` checked, and its compact +/// entries expanded to full entries. Any other magic is an error. +pub fn read_extension_table(buf: &[u8]) -> Result { + if sniff(buf) == Some(ContainerMagic::Ptwc) { + verify_eof_sentinel(buf)?; + let sec = crate::ptwc::locate_sections(buf)?; + return crate::extension::compact_table::read_compact_table(&buf[sec.extension_table]); + } + if buf.len() < HEADER_LEN { + return Err(PtwmCoreError::InvalidContainer( + "file shorter than header".into(), + )); + } + let header = Header::from_bytes(&buf[..HEADER_LEN])?; + let mut table = legacy_extension_table(buf, &header)?; + crate::extension::normalize_builtin_ids(&mut table); + Ok(table) +} + /// Zero-copy reader for a `.ptwm` container backed by an in-memory byte slice. #[derive(Debug)] pub struct ContainerReader<'a> { @@ -1318,35 +1370,7 @@ impl<'a> ContainerReader<'a> { )); } - // Parse the Extension Table (sits at extension_table_offset for - // extension_table_length bytes, immediately after the header). - let mut extension_table = if header.extension_table_length == 0 { - ExtensionTable::empty() - } else { - let et_start = usize::try_from(header.extension_table_offset).map_err(|_| { - PtwmCoreError::InvalidContainer("extension_table_offset overflows usize".into()) - })?; - let et_len = usize::try_from(header.extension_table_length).map_err(|_| { - PtwmCoreError::InvalidContainer("extension_table_length overflows usize".into()) - })?; - let et_end = et_start.checked_add(et_len).ok_or_else(|| { - PtwmCoreError::InvalidContainer( - "extension_table_offset + extension_table_length overflows".into(), - ) - })?; - if et_end > buf.len() { - return Err(PtwmCoreError::InvalidContainer( - "extension table region exceeds buffer".into(), - )); - } - let table_bytes = &buf[et_start..et_end]; - let actual_hash = ExtensionTable::hash(table_bytes); - if actual_hash != header.extension_table_hash { - return Err(PtwmCoreError::ExtensionTableHashMismatch); - } - ExtensionTable::from_bytes(table_bytes) - .map_err(|e| PtwmCoreError::InvalidContainer(format!("extension table: {e}")))? - }; + let mut extension_table = legacy_extension_table(buf, &header)?; // Map built-in ids written by other releases to the current ones. // This runs after the hash check, which covers the on-disk bytes. let builtin_version = crate::extension::normalize_builtin_ids(&mut extension_table); @@ -1620,37 +1644,20 @@ impl<'a> ContainerReader<'a> { } } - /// Return every tensor name. For the legacy layout the names come in - /// index order (sorted by name hash, which is how the on-disk index is - /// stored) and each record header is parsed; use sparingly on very large - /// containers. For the compact layout they come from the index alone, in - /// body (write) order. + /// Return every tensor name in body (write) order: the order of the + /// records in the file, which is the order the writer appended them. The + /// compact layout reads the names from the index alone; the legacy layout + /// parses the name at the start of each record. pub fn tensor_names(&self) -> Result, PtwmCoreError> { - use crate::tensor_record::parse_tensor_record; + use crate::tensor_record::peek_tensor_record_name; - if self.layout == ContainerLayout::Compact { - return Ok(self - .entries_in_body_order() - .into_iter() - .map(|e| e.name.as_deref().unwrap_or_default().to_owned()) - .collect()); - } - let mut names = Vec::with_capacity(self.index.len()); - for e in &self.index { - let start = e.tensor_offset as usize; - let end = start.checked_add(e.tensor_len as usize).ok_or_else(|| { - PtwmCoreError::InvalidContainer("tensor_offset + tensor_len overflows usize".into()) - })?; - if end > self.buf.len() { - return Err(PtwmCoreError::InvalidContainer(format!( - "tensor index entry ({start}..{end}) exceeds buffer ({} bytes)", - self.buf.len() - ))); - } - let (rec, _) = parse_tensor_record(&self.buf[start..end])?; - names.push(rec.name); - } - Ok(names) + self.entries_in_body_order() + .into_iter() + .map(|e| match &e.name { + Some(name) => Ok(name.to_string()), + None => peek_tensor_record_name(self.entry_bytes(e)?).map(str::to_owned), + }) + .collect() } /// Decode every tensor in the container, in index order, across the rayon diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs index a9ac21f..eb855ec 100644 --- a/crates/ptwm-core/tests/ptwc.rs +++ b/crates/ptwm-core/tests/ptwc.rs @@ -1404,3 +1404,101 @@ fn old_implode_rejects_compact_manifests() { let e = v1_1_implode_gate(&compact_reg).unwrap_err(); assert!(e.contains("missing field"), "{e}"); } + +// ── read_extension_table and tensor_names (synthesis section 3, E3 / E14) ── + +#[test] +fn read_extension_table_dispatches_on_the_magic() { + use ptwm_core::container::read_extension_table; + + for name in FIXTURES { + let legacy = load(name); + let compact = to_compact(&legacy); + for blob in [&legacy, &compact] { + let reader = ContainerReader::open(blob).unwrap(); + let table = read_extension_table(blob).unwrap(); + assert_eq!(table.entries, reader.extension_table.entries, "{name}"); + } + } + + let data = feature_tensors(); + let compact = compress(&data, options(true, true, None, ContainerLayout::Compact)); + let table = read_extension_table(&compact).unwrap(); + assert!(!table.entries.is_empty()); + assert!( + table + .entries + .iter() + .all(|e| ptwm_core::extension::is_builtin(&e.canonical_id)) + ); + + assert!(read_extension_table(b"not a container").is_err()); + assert!(read_extension_table(b"").is_err()); +} + +#[test] +fn read_extension_table_checks_the_legacy_region_and_hash() { + use ptwm_core::container::read_extension_table; + + let legacy = load("legacy_v1_1_0.ptwm"); + let header = Header::from_bytes(&legacy[..HEADER_LEN]).unwrap(); + assert!(header.extension_table_length > 0); + + let patched = |h: &Header| { + let mut out = legacy.clone(); + out[..HEADER_LEN].copy_from_slice(&h.to_bytes().unwrap()); + out + }; + + // offset + length overflows u64 / usize. + let mut h = header.clone(); + h.extension_table_offset = u64::MAX - 4; + let err = read_extension_table(&patched(&h)).unwrap_err().to_string(); + assert!(err.contains("exceeds the buffer"), "{err}"); + + let mut h = header.clone(); + h.extension_table_length = u64::MAX; + let err = read_extension_table(&patched(&h)).unwrap_err().to_string(); + assert!(err.contains("exceeds the buffer"), "{err}"); + + // A region that runs past the end of the file. + let mut h = header.clone(); + h.extension_table_length = legacy.len() as u64; + assert!(read_extension_table(&patched(&h)).is_err()); + + // A table byte that no longer matches the header hash. + let mut corrupt = legacy.clone(); + corrupt[header.extension_table_offset as usize + 2] ^= 0x01; + let err = read_extension_table(&corrupt).unwrap_err(); + assert!( + matches!(err, ptwm_core::PtwmCoreError::ExtensionTableHashMismatch), + "{err}" + ); +} + +#[test] +fn read_extension_table_checks_the_compact_meta_hash() { + use ptwm_core::container::read_extension_table; + + let compact = to_compact(&load("legacy_v1_1_0.ptwm")); + // meta_xxh64 covers the Extension Table: flip its first byte. + let sec = ptwm_core::ptwc::locate_sections(&compact).unwrap(); + let mut corrupt = compact.clone(); + corrupt[sec.extension_table.start] ^= 0x01; + assert!(read_extension_table(&corrupt).is_err()); +} + +#[test] +fn tensor_names_come_in_write_order_in_both_layouts() { + let mut rng = Rng(7); + let names = ["zeta", "alpha", "mid", "beta", "omega", "a.b.c"]; + let data: Vec = names + .iter() + .map(|n| fp32(n, vec![8, 4], &mut rng)) + .collect(); + for layout in [ContainerLayout::Legacy, ContainerLayout::Compact] { + let blob = compress(&data, options(true, true, None, layout)); + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(reader.tensor_names().unwrap(), names, "{layout:?}"); + } +} diff --git a/crates/ptwm-py/src/handle.rs b/crates/ptwm-py/src/handle.rs new file mode 100644 index 0000000..4c5629b --- /dev/null +++ b/crates/ptwm-py/src/handle.rs @@ -0,0 +1,273 @@ +//! `ContainerHandle`: one opened container reader over borrowed bytes. +//! +//! Opening a container parses its header, Extension Table, prelude, chain +//! registry and index and checks the referenced extensions. A handle does +//! this once and answers any number of queries from the parsed state. The +//! free functions in `lib.rs` (`decode_tensor`, `list_tensor_names`, ...) +//! open a temporary handle per call. +//! +//! A read-only, C-contiguous buffer (``bytes``, a read-only ``memoryview`` or +//! ``mmap``) is borrowed without a copy: the handle holds the buffer export, +//! which keeps the object alive and its memory pinned. A writable or +//! non-contiguous buffer (``bytearray``, a writable ``mmap`` or numpy array) +//! is copied once, because the caller could change it while the reader +//! holds slices into it. + +use pyo3::buffer::PyBuffer; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use ptwm_core::PtwmCoreError; +use ptwm_core::container::{ContainerLayout, ContainerReader}; +use ptwm_core::header::{ContainerMagic, sniff}; +use ptwm_core::metadata::decode_shape_metadata; +use ptwm_core::tensor_record::RefKind; + +/// `(shape, dtype_name, dtype_code, input_format)` of a tensor. +pub type TensorInfo = (Vec, String, u16, u8); + +/// `(tensor_name, plane_index, role, codec_id, state_source, payload_len)`. +pub type PlaneSummary = (String, u32, u8, u16, u8, u64); + +/// The memory a handle's reader points into. +enum Backing { + /// A read-only, C-contiguous buffer export. Holding it keeps the + /// exporting object alive and its memory unchanged. + Borrowed(#[allow(dead_code)] PyBuffer), + /// A private copy of a buffer that could change under the reader. + Owned(#[allow(dead_code)] Box<[u8]>), +} + +/// An opened `.ptwm` container (either layout). +/// +/// ``ContainerHandle(buffer, skip_missing=False)`` opens the container once; +/// the methods then query it without re-parsing. ``skip_missing`` opens the +/// container even when non-builtin extensions it references are not +/// installed; tensors that need them fail at decode time. +#[pyclass(frozen, module = "ptwm._core")] +pub struct ContainerHandle { + // Field order matters: `reader` borrows from `_backing` and must drop + // first. Struct fields drop in declaration order. + reader: ContainerReader<'static>, + _backing: Backing, +} + +fn runtime_err(e: PtwmCoreError) -> PyErr { + pyo3::exceptions::PyRuntimeError::new_err(e.to_string()) +} + +impl ContainerHandle { + /// Open `buffer`, borrowing it when it is read-only and C-contiguous. + pub fn open(py: Python<'_>, buffer: PyBuffer, skip_missing: bool) -> PyResult { + let (bytes, backing): (&[u8], Backing) = + if buffer.readonly() && buffer.is_c_contiguous() { + let len = buffer.len_bytes(); + let bytes: &[u8] = if len == 0 { + &[] + } else { + // SAFETY: the buffer is C-contiguous and `len` bytes long. + // The export stays alive in `Backing::Borrowed` for as + // long as the reader, and a read-only export cannot be + // written through. + unsafe { std::slice::from_raw_parts(buffer.buf_ptr() as *const u8, len) } + }; + (bytes, Backing::Borrowed(buffer)) + } else { + let owned: Box<[u8]> = buffer.to_vec(py)?.into_boxed_slice(); + // SAFETY: the boxed slice's heap memory does not move when the + // box moves into `Backing::Owned`, and it lives as long as the + // reader. + let bytes = unsafe { std::slice::from_raw_parts(owned.as_ptr(), owned.len()) }; + (bytes, Backing::Owned(owned)) + }; + // SAFETY: `bytes` points into `backing`, which the handle owns and + // drops after the reader (see the field order of `ContainerHandle`). + let bytes: &'static [u8] = unsafe { std::mem::transmute::<&[u8], &'static [u8]>(bytes) }; + let reader = py + .allow_threads(|| { + if skip_missing { + ContainerReader::open_partial(bytes) + } else { + ContainerReader::open(bytes) + } + }) + .map_err(runtime_err)?; + Ok(Self { + reader, + _backing: backing, + }) + } + + pub fn names_vec(&self, py: Python<'_>) -> PyResult> { + py.allow_threads(|| self.reader.tensor_names()) + .map_err(runtime_err) + } + + pub fn decode_bytes<'py>(&self, py: Python<'py>, name: &str) -> PyResult> { + let out = py + .allow_threads(|| self.reader.decode_tensor(name)) + .map_err(runtime_err)?; + Ok(PyBytes::new(py, &out)) + } + + pub fn decode_all<'py>(&self, py: Python<'py>) -> PyResult)>> { + let decoded = py + .allow_threads(|| self.reader.decode_model()) + .map_err(runtime_err)?; + Ok(decoded + .into_iter() + .map(|(name, bytes)| (name, PyBytes::new(py, &bytes))) + .collect()) + } + + pub fn tensor_info(&self, py: Python<'_>, name: &str) -> PyResult> { + py.allow_threads(|| -> Result, PtwmCoreError> { + let rec = self.reader.tensor_record(name)?; + match rec.tensor_metadata { + Some(meta) => { + let (shape, dtype_name) = decode_shape_metadata(&meta)?; + Ok(Some((shape, dtype_name, rec.dtype_code, rec.input_format))) + } + None => Ok(None), + } + }) + .map_err(runtime_err) + } + + pub fn tensor_delta_hash<'py>( + &self, + py: Python<'py>, + name: &str, + ) -> PyResult>> { + let hash = py + .allow_threads(|| -> Result, PtwmCoreError> { + let rec = self.reader.tensor_record(name)?; + Ok(rec + .dependencies + .iter() + .filter(|d| matches!(d.ref_kind, RefKind::ExternalSafetensors)) + .find_map(|d| d.expected_hash)) + }) + .map_err(runtime_err)?; + Ok(hash.map(|h| PyBytes::new(py, &h))) + } + + pub fn tensor_shape(&self, py: Python<'_>, name: &str) -> PyResult, String)>> { + py.allow_threads(|| self.reader.tensor_shape(name)) + .map_err(runtime_err) + } + + pub fn plane_summary_rows(&self, py: Python<'_>) -> PyResult> { + py.allow_threads(|| -> Result, PtwmCoreError> { + let mut out = Vec::new(); + for rec in self.reader.tensor_records()? { + for (idx, plane) in rec.terminals.iter().enumerate() { + out.push(( + rec.name.clone(), + idx as u32, + plane.role as u8, + plane.codec_id.as_u16(), + plane.state_source as u8, + plane.payload_len, + )); + } + } + Ok(out) + }) + .map_err(runtime_err) + } +} + +#[pymethods] +impl ContainerHandle { + #[new] + #[pyo3(signature = (buffer, skip_missing = false))] + fn py_new(py: Python<'_>, buffer: PyBuffer, skip_missing: bool) -> PyResult { + Self::open(py, buffer, skip_missing) + } + + /// ``"legacy"`` (``\x89PTWM`` / ``\x89PTWX``) or ``"compact"`` + /// (``\x89PTWC``). + #[getter] + fn layout(&self) -> &'static str { + self.reader.layout.as_str() + } + + /// Whether the buffer is borrowed (``True``) or was copied (``False``). + #[getter] + fn zero_copy(&self) -> bool { + matches!(self._backing, Backing::Borrowed(_)) + } + + /// Every tensor name in write (body) order. + fn names(&self, py: Python<'_>) -> PyResult> { + self.names_vec(py) + } + + /// Decode one tensor and return its raw bytes. + fn decode<'py>(&self, py: Python<'py>, name: &str) -> PyResult> { + self.decode_bytes(py, name) + } + + /// Decode every tensor in parallel; ``(name, bytes)`` in index order. + fn decode_model<'py>(&self, py: Python<'py>) -> PyResult)>> { + self.decode_all(py) + } + + /// ``(shape, dtype_name, dtype_code, input_format)``, or ``None`` when + /// the tensor has no shape metadata. + fn info(&self, py: Python<'_>, name: &str) -> PyResult> { + self.tensor_info(py, name) + } + + /// The BLAKE3 hash of the tensor's delta reference, or ``None``. + fn delta_hash<'py>(&self, py: Python<'py>, name: &str) -> PyResult>> { + self.tensor_delta_hash(py, name) + } + + /// ``(shape, dtype_name)``, or ``None`` when the tensor has no shape + /// metadata. + fn shape(&self, py: Python<'_>, name: &str) -> PyResult, String)>> { + self.tensor_shape(py, name) + } + + /// One ``(tensor_name, plane_index, role, codec_id, state_source, + /// payload_len)`` row per plane, tensors in write order. + fn plane_summaries(&self, py: Python<'_>) -> PyResult> { + self.plane_summary_rows(py) + } + + fn __len__(&self) -> usize { + self.reader.index.len() + } + + fn __contains__(&self, name: &str) -> bool { + self.reader.get_tensor_record_bytes(name).is_some() + } +} + +/// The multi-tensor layout of `blob` from its magic: ``"legacy"`` for +/// ``\x89PTWM`` / ``\x89PTWX``, ``"compact"`` for ``\x89PTWC``, ``None`` for +/// anything else (including a ``\x89PTWS`` single-tensor frame). Only the +/// first nine bytes are read. +#[pyfunction] +pub fn container_layout(py: Python<'_>, blob: PyBuffer) -> PyResult> { + let owned; + let bytes: &[u8] = if blob.is_c_contiguous() { + crate::contiguous_slice(&blob, "blob")? + } else { + owned = blob.to_vec(py)?; + &owned + }; + Ok(match sniff(bytes) { + Some(ContainerMagic::Ptwm | ContainerMagic::Ptwx) => Some(ContainerLayout::Legacy.as_str()), + Some(ContainerMagic::Ptwc) => Some(ContainerLayout::Compact.as_str()), + None => None, + }) +} + +pub fn register(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_function(pyo3::wrap_pyfunction!(container_layout, m)?)?; + Ok(()) +} diff --git a/crates/ptwm-py/src/inspect.rs b/crates/ptwm-py/src/inspect.rs index 3051de7..0f55366 100644 --- a/crates/ptwm-py/src/inspect.rs +++ b/crates/ptwm-py/src/inspect.rs @@ -5,33 +5,27 @@ use std::io::Read; use pyo3::prelude::*; -use ptwm_core::extension::{ExtensionTable, normalize_builtin_ids}; -use ptwm_core::header::{HEADER_LEN, Header}; +use ptwm_core::container::read_extension_table; fn to_pyerr(e: E) -> PyErr { pyo3::exceptions::PyValueError::new_err(e.to_string()) } -/// Return `buf[off..off + len]`, or `None` when the range overflows or runs -/// past the end of `buf`. -fn extension_table_region(buf: &[u8], off: u64, len: u64) -> Option<&[u8]> { - let start = usize::try_from(off).ok()?; - let end = start.checked_add(usize::try_from(len).ok()?)?; - buf.get(start..end) -} - /// List the extension-table entries recorded in a `.ptwm` file. /// /// Returns a list of ``(canonical_id_hex, human_label, kind_u16, flavor_hints)`` /// tuples -- one per entry in the file's extension table. Returns an empty -/// list when the file carries no extension table (``extension_table_length == 0``). +/// list when the file carries no extension table. /// -/// Built-in entries written by another release are reported with the ids and -/// labels of the current release (see ``normalize_builtin_ids``), so the ids -/// of files written by different releases compare equal. +/// Both layouts are read: a legacy (``\x89PTWM``) file has its table hash +/// verified, and a compact (``\x89PTWC``) file has its compact entries +/// expanded to full entries. Built-in entries written by another release are +/// reported with the ids and labels of the current release (see +/// ``normalize_builtin_ids``), so the ids of files written by different +/// releases compare equal. /// /// Raises ``ValueError`` for any parse error (bad magic, truncated header, -/// out-of-bounds extension table region, ...). +/// out-of-bounds extension table region, hash mismatch, ...). #[pyfunction] pub fn list_extension_table_entries(path: String) -> PyResult> { let mut f = @@ -39,31 +33,7 @@ pub fn list_extension_table_entries(path: String) -> PyResult) -> PyResult<()> { m.add_function(pyo3::wrap_pyfunction!(list_extension_table_entries, m)?)?; Ok(()) } - -#[cfg(test)] -mod tests { - use super::extension_table_region; - - #[test] - fn region_rejects_overflow_and_out_of_bounds() { - let buf = [0u8; 16]; - assert_eq!(extension_table_region(&buf, 4, 8), Some(&buf[4..12])); - assert_eq!(extension_table_region(&buf, 16, 0), Some(&buf[16..16])); - assert!(extension_table_region(&buf, 8, 9).is_none()); - assert!(extension_table_region(&buf, u64::MAX - 4, 8).is_none()); - assert!(extension_table_region(&buf, 4, u64::MAX).is_none()); - } -} diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 8378b36..26ceb57 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -13,12 +13,14 @@ use ptwm_core::chain::wire::read_chain_blob; use ptwm_core::compressor::{ CompressorOptions, InputTensor, compress_model as core_compress_model, }; -use ptwm_core::container::ContainerReader; use ptwm_core::{PtwmCoreError, codec_tagged, delta, entropy}; +use handle::ContainerHandle; + mod delta_scheme; mod device_buffer; mod ext; +mod handle; mod hardware; mod host_flavor; mod inspect; @@ -66,6 +68,7 @@ fn _core(m: &Bound<'_, PyModule>) -> PyResult<()> { host_flavor::register(m)?; ext::register(m)?; inspect::register(m)?; + handle::register(m)?; delta_scheme::register(m)?; hardware::register(m)?; plane_codec_cuda::register(m)?; @@ -577,41 +580,33 @@ fn decode_compact<'py>( )) } -/// Decode a single tensor from a container blob. +// The functions below open a temporary `ContainerHandle` per call (see +// `handle.rs`): `blob` is any buffer, borrowed when read-only and +// C-contiguous. Open one `ContainerHandle` instead for several queries. + +/// Decode a single tensor from a container blob (either layout). /// -/// `blob` is the full container ``bytes`` returned by :func:`compress_model`. -/// `name` is the tensor name string. Returns the raw decompressed ``bytes`` -/// via the chain-based decode path. +/// `blob` is the full container (any buffer; ``bytes`` is borrowed without a +/// copy). `name` is the tensor name string. Returns the raw decompressed +/// ``bytes`` via the chain-based decode path. /// /// `skip_missing` — when ``True`` open the container even if non-builtin /// extensions referenced in its Extension Table are not installed (mirrors /// ``ContainerReader::open_partial``). Tensors that actually need those /// missing extensions will still fail at decode time. /// -/// Raises ``RuntimeError`` when the container is malformed or the tensor -/// name is missing. Raises ``ValueError`` when non-builtin extensions are -/// missing and ``skip_missing`` is ``False`` (the default). +/// Raises ``RuntimeError`` when the container is malformed, the tensor name +/// is missing, or non-builtin extensions are missing and ``skip_missing`` is +/// ``False`` (the default). #[pyfunction] #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor<'py>( py: Python<'py>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: PyBuffer, name: &str, skip_missing: bool, ) -> PyResult> { - let data: Vec = blob.as_bytes().to_vec(); - let name_owned = name.to_string(); - let out = py - .allow_threads(|| -> Result, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } else { - ContainerReader::open(&data)? - }; - reader.decode_tensor(&name_owned) - }) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - Ok(pyo3::types::PyBytes::new(py, &out)) + ContainerHandle::open(py, blob, skip_missing)?.decode_bytes(py, name) } /// Decode every tensor in `blob` at once, in parallel across the rayon pool @@ -623,30 +618,13 @@ fn decode_tensor<'py>( #[pyo3(signature = (blob, skip_missing = false))] fn decode_model<'py>( py: Python<'py>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: PyBuffer, skip_missing: bool, ) -> PyResult)>> { - // Borrow the blob's bytes directly — no full-model copy. `blob` outlives - // this call, and the bytes are plain immutable memory, so reading them with - // the GIL released is sound. - let data = blob.as_bytes(); - let decoded = py - .allow_threads(|| -> Result)>, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(data)? - } else { - ContainerReader::open(data)? - }; - reader.decode_model() - }) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - decoded - .into_iter() - .map(|(name, bytes)| Ok((name, pyo3::types::PyBytes::new(py, &bytes)))) - .collect() + ContainerHandle::open(py, blob, skip_missing)?.decode_all(py) } -/// Return `(shape, dtype_name, dtype_code, input_format)` for a v3 tensor. +/// Return `(shape, dtype_name, dtype_code, input_format)` for a tensor. /// /// Returns `None` when the tensor has no CBOR shape metadata embedded. /// `skip_missing` has the same semantics as for :func:`decode_tensor`. @@ -654,96 +632,36 @@ fn decode_model<'py>( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_info( py: Python<'_>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: PyBuffer, name: &str, skip_missing: bool, -) -> PyResult, String, u16, u8)>> { - use ptwm_core::metadata::decode_shape_metadata; - - let data: Vec = blob.as_bytes().to_vec(); - let name_owned = name.to_string(); - let info = py - .allow_threads( - || -> Result, String, u16, u8)>, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } else { - ContainerReader::open(&data)? - }; - let rec = reader.tensor_record(&name_owned)?; - match rec.tensor_metadata { - Some(meta) => { - let (shape, dtype_name) = decode_shape_metadata(&meta)?; - Ok(Some((shape, dtype_name, rec.dtype_code, rec.input_format))) - } - None => Ok(None), - } - }, - ) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - Ok(info) +) -> PyResult> { + ContainerHandle::open(py, blob, skip_missing)?.tensor_info(py, name) } -/// Return the BLAKE3 hash of the delta reference recorded for a v3 tensor, +/// Return the BLAKE3 hash of the delta reference recorded for a tensor, /// or `None` if the tensor carries no delta dependency. /// `skip_missing` has the same semantics as for :func:`decode_tensor`. #[pyfunction] #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_delta_hash<'py>( py: Python<'py>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: PyBuffer, name: &str, skip_missing: bool, ) -> PyResult>> { - use ptwm_core::tensor_record::RefKind; - - let data: Vec = blob.as_bytes().to_vec(); - let name_owned = name.to_string(); - let hash_opt = py - .allow_threads(|| -> Result, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } else { - ContainerReader::open(&data)? - }; - let rec = reader.tensor_record(&name_owned)?; - for dep in &rec.dependencies { - if matches!(dep.ref_kind, RefKind::ExternalSafetensors) - && let Some(h) = dep.expected_hash - { - return Ok(Some(h)); - } - } - Ok(None) - }) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - Ok(hash_opt.map(|h| pyo3::types::PyBytes::new(py, &h))) + ContainerHandle::open(py, blob, skip_missing)?.tensor_delta_hash(py, name) } -/// List every tensor name stored in a v3 container blob, in index order. -/// `skip_missing` has the same semantics as for :func:`decode_tensor`. +/// List every tensor name stored in a container blob, in write (body) +/// order. `skip_missing` has the same semantics as for :func:`decode_tensor`. #[pyfunction] #[pyo3(signature = (blob, skip_missing = false))] -fn list_tensor_names( - py: Python<'_>, - blob: &Bound<'_, pyo3::types::PyBytes>, - skip_missing: bool, -) -> PyResult> { - let data: Vec = blob.as_bytes().to_vec(); - let names = py - .allow_threads(|| -> Result, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } else { - ContainerReader::open(&data)? - }; - reader.tensor_names() - }) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - Ok(names) +fn list_tensor_names(py: Python<'_>, blob: PyBuffer, skip_missing: bool) -> PyResult> { + ContainerHandle::open(py, blob, skip_missing)?.names_vec(py) } -/// Decode the shape-metadata map for a v3 tensor, if present. +/// Decode the shape-metadata map for a tensor, if present. /// /// Returns ``(shape, dtype_name)`` or ``None`` when no CBOR shape metadata. /// `skip_missing` has the same semantics as for :func:`decode_tensor`. @@ -751,66 +669,26 @@ fn list_tensor_names( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_shape( py: Python<'_>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: PyBuffer, name: &str, skip_missing: bool, ) -> PyResult, String)>> { - let data: Vec = blob.as_bytes().to_vec(); - let name_owned = name.to_string(); - let got = py - .allow_threads(|| -> Result, String)>, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } else { - ContainerReader::open(&data)? - }; - reader.tensor_shape(&name_owned) - }) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - Ok(got) + ContainerHandle::open(py, blob, skip_missing)?.tensor_shape(py, name) } -/// Per-plane summary for every tensor in a v3 container blob. +/// Per-plane summary for every tensor in a container blob. /// /// Returns a list of (tensor_name, plane_index, role, codec_id, -/// state_source, payload_len) tuples — one row per plane. -/// `skip_missing` has the same semantics as for :func:`decode_tensor`. +/// state_source, payload_len) tuples — one row per plane, tensors in write +/// order. `skip_missing` has the same semantics as for :func:`decode_tensor`. #[pyfunction] #[pyo3(signature = (blob, skip_missing = false))] fn list_plane_summaries( py: Python<'_>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: PyBuffer, skip_missing: bool, -) -> PyResult> { - let data: Vec = blob.as_bytes().to_vec(); - let summaries = py - .allow_threads( - || -> Result, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } else { - ContainerReader::open(&data)? - }; - let names = reader.tensor_names()?; - let mut out = Vec::new(); - for name in names { - let rec = reader.tensor_record(&name)?; - for (idx, plane) in rec.terminals.iter().enumerate() { - out.push(( - name.clone(), - idx as u32, - plane.role as u8, - plane.codec_id.as_u16(), - plane.state_source as u8, - plane.payload_len, - )); - } - } - Ok(out) - }, - ) - .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; - Ok(summaries) +) -> PyResult> { + ContainerHandle::open(py, blob, skip_missing)?.plane_summary_rows(py) } /// Explode a `.ptwm` container into its canonical PTWM-JSON projection. diff --git a/python/ptwm/__init__.py b/python/ptwm/__init__.py index 649b35a..5f774fc 100644 --- a/python/ptwm/__init__.py +++ b/python/ptwm/__init__.py @@ -7,15 +7,17 @@ Public entry points: :class:`Compressor`, :class:`Decompressor`, :class:`CompressionConfig`, :class:`DecompressionConfig`, -:class:`Method`, :class:`Format`, :class:`Framing`. Third-party integrations -(HuggingFace `transformers`, `safetensors`) live in :mod:`ptwm.integrations` -and are imported from there explicitly, not re-exported here. +:class:`Method`, :class:`Format`, :class:`Framing`, +:class:`ContainerLayout`. Third-party integrations (HuggingFace +`transformers`, `safetensors`) live in :mod:`ptwm.integrations` and are +imported from there explicitly, not re-exported here. """ from importlib.metadata import PackageNotFoundError, version from ._config import ( CompressionConfig, + ContainerLayout, DecompressionConfig, Format, Framing, @@ -40,6 +42,7 @@ "CompressionConfig", "CompressionMethodNotSupportedError", "Compressor", + "ContainerLayout", "DecompressionConfig", "Decompressor", "Format", diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index c3d0e1c..36e47f4 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -49,9 +49,13 @@ class Framing(Enum): with a fixed cost of at most about 70 bytes, and stores the raw bytes when no codec makes them smaller. It falls back to ``CONTAINER`` when the tensor cannot use the frame (streaming, or a plane that selects shared - codec state). ``CONTAINER`` writes a full ``.ptwm`` container (magic - ``\x89PTWM``), which readers before v1.2 require, and which the + codec state). ``CONTAINER`` writes a full ``.ptwm`` container in the + layout :attr:`CompressionConfig.container_layout` selects, which the random-access, store and transcode tools open. + + Readers before v1.2 read only ``CONTAINER`` output in the + :attr:`ContainerLayout.LEGACY` layout (magic ``\x89PTWM``); they reject + the ``\x89PTWS`` frame and the ``\x89PTWC`` compact container. """ COMPACT = 1 @@ -66,6 +70,45 @@ def _missing_(cls, value): return None +class ContainerLayout(Enum): + r"""On-disk layout of a multi-tensor ``.ptwm`` container. + + ``COMPACT`` (magic ``\x89PTWC``) stores built-in extensions by id, + integers as varints and tensor names only in the index; it needs v1.2 or + later to read. ``LEGACY`` (magic ``\x89PTWM``) is the fixed-width layout + that every release reads. + + Strings coerce case-insensitively: ``ContainerLayout("compact")`` and + ``ContainerLayout("LEGACY")`` work, and the bindings take the lower-case + :attr:`value`. + """ + + COMPACT = "compact" + LEGACY = "legacy" + + @classmethod + def _missing_(cls, value): + if isinstance(value, str): + value = value.upper() + if value in cls.__members__: + return cls.__members__[value] + return None + + +def container_layout_name(layout: ContainerLayout | str) -> str: + """Return the binding name (``"compact"`` or ``"legacy"``) of ``layout``. + + A string that names no layout is returned unchanged, so the binding + rejects it with an error that names ``container_layout``. + """ + if isinstance(layout, ContainerLayout): + return layout.value + try: + return ContainerLayout(layout).value + except ValueError: + return layout + + class Lossy(Enum): NONE = 0 INTEGER = 1 @@ -113,6 +156,10 @@ class CompressionConfig: device: int | None = None framing: Framing = Framing.COMPACT """Output layout; see :class:`Framing`.""" + container_layout: ContainerLayout = ContainerLayout.LEGACY + """Layout of a full container written with ``Framing.CONTAINER`` (or as + the ``Framing.COMPACT`` fallback); see :class:`ContainerLayout`. A string + such as ``"compact"`` is coerced on construction.""" emit_payload_hash: bool = True """Store an xxh64 of the raw bytes (8 B per tensor), checked on decode.""" emit_plane_crc: bool = True @@ -122,6 +169,14 @@ class CompressionConfig: the payload hash already covers every plane of its one tensor. """ + def __post_init__(self) -> None: + # Accept "compact" / "legacy" (any case); unknown names raise + # ValueError here rather than when the container is written. + if not isinstance(self.container_layout, ContainerLayout): + object.__setattr__( + self, "container_layout", ContainerLayout(self.container_layout) + ) + @classmethod def from_resolved_policy( cls, diff --git a/python/ptwm/_magic.py b/python/ptwm/_magic.py new file mode 100644 index 0000000..3d48dd7 --- /dev/null +++ b/python/ptwm/_magic.py @@ -0,0 +1,72 @@ +r"""Magic numbers of the PTWM file formats. + +Four formats share the ``\x89PTW`` prefix: + +- ``\x89PTWM\r\n\x1a\n``: legacy multi-tensor container (every release reads it). +- ``\x89PTWX\r\n\x1a\n``: legacy container with embedded WASM extensions. +- ``\x89PTWC\r\n\x1a\n``: compact multi-tensor container (v1.2 and later). +- ``\x89PTWS``: compact single-tensor frame written by + :meth:`ptwm.Compressor.compress` (v1.2 and later). + +Compare the full magic, not a 5-byte prefix: the fifth byte is what tells the +formats apart, and readers before v1.2 that check only ``\x89PTWM`` reject the +newer formats by design. +""" + +from __future__ import annotations + +from typing import Literal + +__all__ = [ + "MAGIC_LEN", + "MAGIC_PTWC", + "MAGIC_PTWM", + "MAGIC_PTWS", + "MAGIC_PTWX", + "ContainerKind", + "container_kind", + "is_container", +] + +MAGIC_PTWM = b"\x89PTWM\r\n\x1a\n" +MAGIC_PTWX = b"\x89PTWX\r\n\x1a\n" +MAGIC_PTWC = b"\x89PTWC\r\n\x1a\n" +MAGIC_PTWS = b"\x89PTWS" + +MAGIC_LEN = len(MAGIC_PTWM) +"""Bytes to read from the start of a file to tell every format apart.""" + +ContainerKind = Literal["ptwm", "ptwx", "ptwc", "ptws"] + +_CONTAINERS: dict[bytes, ContainerKind] = { + MAGIC_PTWM: "ptwm", + MAGIC_PTWX: "ptwx", + MAGIC_PTWC: "ptwc", +} + + +def container_kind(buf: bytes | bytearray | memoryview) -> ContainerKind | None: + """Return the format of ``buf`` from its magic, or ``None``. + + ``buf`` needs only its first :data:`MAGIC_LEN` bytes. Returns ``"ptwm"``, + ``"ptwx"`` or ``"ptwc"`` for a multi-tensor container, ``"ptws"`` for a + single-tensor frame, and ``None`` for anything else (including a buffer + too short to hold a magic). + """ + head = bytes(memoryview(buf)[:MAGIC_LEN]) + kind = _CONTAINERS.get(head) + if kind is not None: + return kind + if head[: len(MAGIC_PTWS)] == MAGIC_PTWS: + return "ptws" + return None + + +def is_container(buf: bytes | bytearray | memoryview) -> bool: + """Whether ``buf`` starts with a multi-tensor container magic. + + True for ``\\x89PTWM``, ``\\x89PTWX`` and ``\\x89PTWC``: the formats that + :class:`ptwm.random_access.TensorIndex`, the stores and the transcoder + open. False for a ``\\x89PTWS`` single-tensor frame. + """ + return container_kind(buf) in ("ptwm", "ptwx", "ptwc") diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index 89b4050..4acd190 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -13,11 +13,13 @@ from ._errors import translate_errors __all__ = [ + "ContainerHandle", "blake3_hash", "builtin_canonical_id", "codec_id_registry", "compress_compact", "compress_model", + "container_layout", "decode_compact", "decode_model", "decode_tensor", @@ -328,7 +330,10 @@ def decode_tensor_delta_hash( @translate_errors def list_tensor_names(blob: bytes, skip_missing: bool = False) -> list[str]: - """Return every tensor name stored in a ``.ptwm`` container.""" + """Return every tensor name stored in a ``.ptwm`` container. + + Names come in write (body) order for both layouts. + """ return list(_ext.list_tensor_names(blob, skip_missing)) @@ -352,6 +357,100 @@ def decode_tensor_info( return _ext.decode_tensor_info(blob, name, skip_missing) +def container_layout(blob: bytes | bytearray | memoryview) -> str | None: + r"""Return the multi-tensor layout of ``blob`` from its magic. + + ``"legacy"`` for ``\x89PTWM`` / ``\x89PTWX``, ``"compact"`` for + ``\x89PTWC``, ``None`` for anything else (including a ``\x89PTWS`` + single-tensor frame). Only the magic is read; the container is not + validated. + """ + return _ext.container_layout(blob) + + +class ContainerHandle: + """One opened ``.ptwm`` container (either layout). + + Opening parses the header, Extension Table, prelude, chain registry and + index once; the methods then query that state. Use a handle instead of + the free functions (:func:`decode_tensor`, :func:`list_tensor_names`, + ...) when a container is queried more than once: each free function + opens the container again. + + ``blob`` is any buffer. ``bytes``, a read-only ``memoryview`` and a + read-only ``mmap`` are borrowed without a copy (the handle keeps them + alive); a writable buffer such as ``bytearray`` is copied once. + ``skip_missing`` opens the container even when non-builtin extensions it + references are not installed; tensors that need them fail at decode. + + Errors from the binding surface as :class:`ptwm.Error`. + """ + + __slots__ = ("_h",) + + @translate_errors + def __init__( + self, blob: bytes | bytearray | memoryview, skip_missing: bool = False + ) -> None: + self._h = _ext.ContainerHandle(blob, skip_missing) + + @property + def layout(self) -> str: + """``"legacy"`` or ``"compact"``.""" + return self._h.layout + + @property + def zero_copy(self) -> bool: + """Whether the handle borrows ``blob`` (``False``: it holds a copy).""" + return self._h.zero_copy + + @translate_errors + def names(self) -> list[str]: + """Every tensor name, in write (body) order.""" + return list(self._h.names()) + + @translate_errors + def decode(self, name: str) -> bytes: + """Decode tensor ``name`` and return its raw bytes.""" + return self._h.decode(name) + + @translate_errors + def decode_model(self) -> list[tuple[str, bytes]]: + """Decode every tensor in parallel; ``(name, raw_bytes)`` pairs.""" + return self._h.decode_model() + + @translate_errors + def info(self, name: str) -> tuple[list[int], str, int, int] | None: + """``(shape, dtype_name, dtype_code, input_format)`` or ``None``.""" + return self._h.info(name) + + @translate_errors + def delta_hash(self, name: str) -> bytes | None: + """BLAKE3 hash of the delta reference of ``name``, or ``None``.""" + got = self._h.delta_hash(name) + return bytes(got) if got is not None else None + + @translate_errors + def shape(self, name: str) -> tuple[list[int], str] | None: + """``(shape, dtype_name)`` of ``name``, or ``None``.""" + return self._h.shape(name) + + @translate_errors + def plane_summaries(self) -> list[tuple[str, int, int, int, int, int]]: + """One ``(tensor_name, plane_index, role, codec_id, state_source, + payload_len)`` row per plane, tensors in write order.""" + return list(self._h.plane_summaries()) + + def __len__(self) -> int: # noqa: D105 + return len(self._h) + + def __contains__(self, name: object) -> bool: # noqa: D105 + return isinstance(name, str) and name in self._h + + def __repr__(self) -> str: # noqa: D105 + return f"ContainerHandle(layout={self.layout!r}, n={len(self)})" + + # --- Container ⇄ canonical PTWM-JSON manifest transcoder ------------------- diff --git a/python/ptwm/core/_compressor.py b/python/ptwm/core/_compressor.py index bdcca53..a0d3fe1 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -222,4 +222,5 @@ def _compress_bin( method_hint, streaming_chunk, codec_menu_ids, + container_layout=self.config.container_layout.value, ) diff --git a/python/ptwm/core/_decompressor.py b/python/ptwm/core/_decompressor.py index cade22f..b66bbf4 100644 --- a/python/ptwm/core/_decompressor.py +++ b/python/ptwm/core/_decompressor.py @@ -7,14 +7,10 @@ from .. import _rust from .._config import DecompressionConfig, Format from .._exceptions import HeaderParseError, InvalidDTypeError +from .._magic import MAGIC_PTWS, container_kind from ..delta import decode as delta_decode from ..utils import DType -# Full multi-tensor container and compact single-tensor frame. -_MAGIC_CONTAINER = b"\x89PTWM" -_MAGIC_COMPACT = b"\x89PTWS" - - def _torch_dtype_for_code(dtype_code: int) -> torch.dtype | None: """Map a wire ``dtype_code`` to a concrete ``torch.dtype`` if supported.""" for m in DType: @@ -42,8 +38,8 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten r"""Decompress a ``.ptwm`` blob. Accepts a compact single-tensor frame (magic ``\x89PTWS``) or a - full container (magic ``\x89PTWM``). The single-tensor container - uses name ``"_"``. Raw decompressed bytes + full container in either layout (magic ``\x89PTWM``, ``\x89PTWX`` + or ``\x89PTWC``). The single-tensor container uses name ``"_"``. Raw decompressed bytes come from ``decode_tensor``; shape/dtype/input_format (when present) come from the CBOR shape-metadata map. Format recovery: @@ -79,14 +75,14 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten raise ValueError(msg) mv_data = memoryview(data) - magic = bytes(mv_data[: len(_MAGIC_CONTAINER)]) - if len(magic) < len(_MAGIC_CONTAINER): + if mv_data.nbytes < len(MAGIC_PTWS): msg = "too short to be a .ptwm container" raise HeaderParseError(msg) - blob = bytes(mv_data) + blob = data if isinstance(data, bytes) else bytes(mv_data) + kind = container_kind(blob) info: tuple[list[int], str, int, int] | None - if magic == _MAGIC_COMPACT: + if kind == "ptws": raw, shape, dtype_name, dtype_code, input_format, expected_hash, _ = ( _rust.decode_compact(blob) ) @@ -95,11 +91,12 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten if shape is not None else None ) - elif magic == _MAGIC_CONTAINER: - skip = self.config.skip_missing - raw = _rust.decode_tensor(blob, "_", skip) - expected_hash = _rust.decode_tensor_delta_hash(blob, "_", skip) - info = _rust.decode_tensor_info(blob, "_", skip) + elif kind is not None: + # One open serves the decode and both metadata lookups. + handle = _rust.ContainerHandle(blob, self.config.skip_missing) + raw = handle.decode("_") + expected_hash = handle.delta_hash("_") + info = handle.info("_") else: msg = "not a .ptwm container" raise HeaderParseError(msg) @@ -174,15 +171,14 @@ def _bytes_to_numpy( return arr.reshape(tuple(shape)) if shape else arr.reshape(()) def decompress_container_to_state_dict(self, blob: bytes) -> dict[str, bytes]: - """Decompress a container into a ``{name: raw_bytes}`` map. + """Decompress a container (either layout) into ``{name: raw_bytes}``. - Iterates the embedded tensor index and decodes each tensor in input - order. Result keys preserve the tensor names from the source - safetensors archive. + Opens the container once and decodes each tensor in write order, so + the keys come in the order the tensors were compressed. Result keys + preserve the tensor names from the source safetensors archive. """ - skip = self.config.skip_missing - names = _rust.list_tensor_names(blob, skip) - return {name: bytes(_rust.decode_tensor(blob, name, skip)) for name in names} + handle = _rust.ContainerHandle(blob, self.config.skip_missing) + return {name: handle.decode(name) for name in handle.names()} def decompress_file(self, path: str) -> bytes | np.ndarray | torch.Tensor: """Decompress data from a file.""" diff --git a/python/ptwm/random_access/__init__.py b/python/ptwm/random_access/__init__.py index e43698f..a4c8927 100644 --- a/python/ptwm/random_access/__init__.py +++ b/python/ptwm/random_access/__init__.py @@ -1,9 +1,10 @@ -"""Random-access loader for ``.ptwm`` containers. +r"""Random-access loader for ``.ptwm`` containers. -:class:`TensorIndex` opens a compressed bundle and decompresses individual -tensors on demand. The index is tiny (a manifest of name → offset/length -records), so opening a bundle and fetching one tensor touches only that -tensor's bytes. +:class:`TensorIndex` opens a compressed bundle in either layout (magic +``\x89PTWM``, ``\x89PTWX`` or ``\x89PTWC``) and decompresses individual +tensors on demand. The container is parsed once into a +:class:`ptwm._rust.ContainerHandle`; fetching one tensor then decodes only +that tensor's bytes. idx = TensorIndex.open("model.safetensors.ptwm") weight = idx.get_tensor("layer_5.weight") @@ -21,10 +22,9 @@ import torch from .. import _rust +from .._magic import is_container from ..utils import raw_to_tensor -_MAGIC = b"\x89PTWM" - @dataclass(frozen=True, slots=True) class TensorEntry: @@ -39,10 +39,11 @@ class TensorIndex: def __init__(self, path: str | Path) -> None: self._path = Path(path) self._blob: bytes = self._path.read_bytes() - if self._blob[: len(_MAGIC)] != _MAGIC: + if not is_container(self._blob): msg = f"{self._path} is not a .ptwm container" raise ValueError(msg) - self._names: list[str] = list(_rust.list_tensor_names(self._blob)) + self._handle = _rust.ContainerHandle(self._blob) + self._names: list[str] = self._handle.names() self._entries_by_name: dict[str, TensorEntry] = { n: TensorEntry(name=n) for n in self._names } @@ -59,12 +60,22 @@ def path(self) -> Path: """Filesystem path this index was opened from.""" return self._path + @property + def layout(self) -> str: + """Container layout: ``"legacy"`` or ``"compact"``.""" + return self._handle.layout + def names(self) -> list[str]: - """Return tensor names in manifest (write) order.""" + """Return tensor names in write order. + + This is the order the writer appended the tensors (the order of the + records in the file), for both layouts. It is not the order of the + legacy on-disk index, which is sorted by name hash. + """ return list(self._names) def entries(self) -> list[TensorEntry]: - """Return every :class:`TensorEntry` in manifest order.""" + """Return every :class:`TensorEntry` in write order.""" return [self._entries_by_name[n] for n in self._names] def entry(self, name: str) -> TensorEntry: @@ -80,10 +91,10 @@ def get_bytes(self, name: str) -> bytes: """Return the raw decompressed bytes for ``name``.""" if name not in self._entries_by_name: raise KeyError(name) - return _rust.decode_tensor(self._blob, name) + return self._handle.decode(name) def _get_shape_dtype(self, name: str) -> tuple[list[int], str] | None: - return _rust.decode_tensor_shape(self._blob, name) + return self._handle.shape(name) def get_tensor(self, name: str) -> torch.Tensor: """Decompress ``name`` and return a :class:`torch.Tensor`. diff --git a/python/ptwm/stores/__init__.py b/python/ptwm/stores/__init__.py index 8bd3ebe..871d9b8 100644 --- a/python/ptwm/stores/__init__.py +++ b/python/ptwm/stores/__init__.py @@ -14,6 +14,7 @@ import tarfile from pathlib import Path +from .._magic import MAGIC_LEN, is_container from ._common import ( ExplodedModel, WeightStore, @@ -52,7 +53,8 @@ def detect_format(path: str | Path) -> str | None: Detection follows the resolved spec: WebDataset via ``index.json.ptwm_format == "ptwm-wds"`` (or a ``__shardmeta__.json`` member in a single shard); LMDB via the ``\\x00fmt`` key (an env directory contains - ``data.mdb``); native ``.ptwm`` via the container magic. + ``data.mdb``); native ``.ptwm`` via the container magic of either layout + (``\x89PTWM``, ``\x89PTWX`` or ``\x89PTWC``; see :mod:`ptwm._magic`). """ p = Path(path) if p.is_dir(): @@ -78,7 +80,7 @@ def detect_format(path: str | Path) -> str | None: pass try: with p.open("rb") as fh: - if fh.read(5) == b"\x89PTWM": + if is_container(fh.read(MAGIC_LEN)): return "ptwm" except OSError: pass From 18c242118bfd42560cc4633f649c75774e5f5ce8 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 01:51:16 +0000 Subject: [PATCH 20/34] feat(py): ContainerHandle, container_layout plumbing and magic helpers Stage 8 of the PTWC work (synthesis sections 3 E3/E10/E14 and 7). Bindings (ptwm-py): - ContainerHandle pyclass: one opened ContainerReader over a borrowed read-only buffer (bytes, read-only memoryview or mmap); writable buffers are copied once. Methods: names, decode, decode_model, info, delta_hash, shape, plane_summaries, plus layout and zero_copy. - decode_tensor, decode_model, decode_tensor_info, decode_tensor_delta_hash, list_tensor_names, decode_tensor_shape and list_plane_summaries are thin wrappers over a handle and take any buffer instead of copying it; all honour skip_missing. - _core.container_layout(blob) maps the magic to "legacy"/"compact". - list_extension_table_entries reads both layouts through the new core read_extension_table (bounds-checked region, verified hash, compact entries expanded). Core: - read_extension_table(buf) dispatches on the magic. - tensor_names and decode_model return body (write) order in both layouts (E14). Python: - ptwm._magic: MAGIC_* constants, container_kind, is_container; used by the Decompressor, TensorIndex and stores.detect_format. - ContainerLayout enum (string coercion) and CompressionConfig.container_layout (default stays LEGACY until the default switch); container_layout threaded through the Compressor, compress_ptwm_blob / write_ptwm_shard, the LMDB and WebDataset writers and compress_safetensors_file; the enum is accepted everywhere. - ContainerHandle used by TensorIndex, Decompressor.decompress/decompress_container_to_state_dict, utils.decode_ptwm_tensor, the HF loader and SafeOpen. - Framing docstring: only LEGACY container output is v1.1-readable. Tests: tests/ptwm/core/test_container_handle.py (new); handle skip_missing and inspect tests in test_legacy_hardening.py; layout forwarding in test_compress_safetensors.py; enum test in test_compact_layout.py; store tests check magics through _magic. The two SafeOpen laziness tests now spy on ContainerHandle.decode, because SafeOpen no longer calls _rust.decode_tensor. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/container.rs | 19 +- crates/ptwm-core/tests/ptwc.rs | 9 +- crates/ptwm-py/src/handle.rs | 45 ++- crates/ptwm-py/src/lib.rs | 6 +- python/ptwm/_config.py | 7 +- python/ptwm/_magic.py | 6 +- python/ptwm/_rust/__init__.py | 44 +- python/ptwm/core/_decompressor.py | 8 +- .../integrations/_compress_safetensors.py | 14 +- python/ptwm/integrations/_hf.py | 51 ++- python/ptwm/integrations/_safetensors.py | 28 +- python/ptwm/sharding/_writer.py | 10 +- python/ptwm/stores/_lmdb.py | 11 +- python/ptwm/stores/_webdataset.py | 3 +- python/ptwm/utils/_decode.py | 17 +- tests/ptwm/core/test_container_handle.py | 378 ++++++++++++++++++ .../integrations/test_compress_safetensors.py | 31 ++ .../integrations/test_mode_a_safe_open.py | 16 +- tests/ptwm/integrations/test_mode_b_shell.py | 16 +- tests/ptwm/stores/test_compact_layout.py | 48 ++- tests/ptwm/stores/test_lmdb.py | 3 +- tests/ptwm/stores/test_transcoder.py | 5 +- tests/ptwm/stores/test_webdataset.py | 3 +- .../ptwm/wire_format/test_legacy_hardening.py | 57 +++ 24 files changed, 704 insertions(+), 131 deletions(-) create mode 100644 tests/ptwm/core/test_container_handle.py diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 8909365..4c3bfd2 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -1660,19 +1660,20 @@ impl<'a> ContainerReader<'a> { .collect() } - /// Decode every tensor in the container, in index order, across the rayon - /// pool. `decode_tensor` is `&self` and reads only shared state, so the - /// per-tensor decodes are independent — this is the primary multi-core - /// scaling lever for full-model loads (the per-tensor `decode_tensor` stays - /// available for random access). Returns `(name, raw_bytes)` per tensor. + /// Decode every tensor in the container across the rayon pool. Returns + /// `(name, raw_bytes)` per tensor in body (write) order, the order of + /// [`tensor_names`](Self::tensor_names). `decode_tensor` is `&self` and + /// reads only shared state, so the per-tensor decodes are independent; + /// this is the main multi-core lever for full-model loads (the + /// per-tensor `decode_tensor` stays available for random access). pub fn decode_model(&self) -> Result)>, PtwmCoreError> { if self.index.is_empty() { return Ok(Vec::new()); } - // Parse each record header and decode it in the same parallel pass — - // no sequential tensor_names() pre-pass. - self.index - .par_iter() + // Parse each record header and decode it in the same parallel pass, + // with no sequential tensor_names() pre-pass. + self.entries_in_body_order() + .into_par_iter() .map(|e| { let rec = self.entry_record(e)?; let bytes = self.decode_tensor_record(&rec)?; diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs index eb855ec..99c02b6 100644 --- a/crates/ptwm-core/tests/ptwc.rs +++ b/crates/ptwm-core/tests/ptwc.rs @@ -1489,7 +1489,7 @@ fn read_extension_table_checks_the_compact_meta_hash() { } #[test] -fn tensor_names_come_in_write_order_in_both_layouts() { +fn tensor_names_and_decode_model_come_in_write_order_in_both_layouts() { let mut rng = Rng(7); let names = ["zeta", "alpha", "mid", "beta", "omega", "a.b.c"]; let data: Vec = names @@ -1500,5 +1500,12 @@ fn tensor_names_come_in_write_order_in_both_layouts() { let blob = compress(&data, options(true, true, None, layout)); let reader = ContainerReader::open(&blob).unwrap(); assert_eq!(reader.tensor_names().unwrap(), names, "{layout:?}"); + let decoded: Vec = reader + .decode_model() + .unwrap() + .into_iter() + .map(|(name, _)| name) + .collect(); + assert_eq!(decoded, names, "{layout:?}"); } } diff --git a/crates/ptwm-py/src/handle.rs b/crates/ptwm-py/src/handle.rs index 4c5629b..05cb14c 100644 --- a/crates/ptwm-py/src/handle.rs +++ b/crates/ptwm-py/src/handle.rs @@ -59,27 +59,26 @@ fn runtime_err(e: PtwmCoreError) -> PyErr { impl ContainerHandle { /// Open `buffer`, borrowing it when it is read-only and C-contiguous. pub fn open(py: Python<'_>, buffer: PyBuffer, skip_missing: bool) -> PyResult { - let (bytes, backing): (&[u8], Backing) = - if buffer.readonly() && buffer.is_c_contiguous() { - let len = buffer.len_bytes(); - let bytes: &[u8] = if len == 0 { - &[] - } else { - // SAFETY: the buffer is C-contiguous and `len` bytes long. - // The export stays alive in `Backing::Borrowed` for as - // long as the reader, and a read-only export cannot be - // written through. - unsafe { std::slice::from_raw_parts(buffer.buf_ptr() as *const u8, len) } - }; - (bytes, Backing::Borrowed(buffer)) + let (bytes, backing): (&[u8], Backing) = if buffer.readonly() && buffer.is_c_contiguous() { + let len = buffer.len_bytes(); + let bytes: &[u8] = if len == 0 { + &[] } else { - let owned: Box<[u8]> = buffer.to_vec(py)?.into_boxed_slice(); - // SAFETY: the boxed slice's heap memory does not move when the - // box moves into `Backing::Owned`, and it lives as long as the - // reader. - let bytes = unsafe { std::slice::from_raw_parts(owned.as_ptr(), owned.len()) }; - (bytes, Backing::Owned(owned)) + // SAFETY: the buffer is C-contiguous and `len` bytes long. + // The export stays alive in `Backing::Borrowed` for as + // long as the reader, and a read-only export cannot be + // written through. + unsafe { std::slice::from_raw_parts(buffer.buf_ptr() as *const u8, len) } }; + (bytes, Backing::Borrowed(buffer)) + } else { + let owned: Box<[u8]> = buffer.to_vec(py)?.into_boxed_slice(); + // SAFETY: the boxed slice's heap memory does not move when the + // box moves into `Backing::Owned`, and it lives as long as the + // reader. + let bytes = unsafe { std::slice::from_raw_parts(owned.as_ptr(), owned.len()) }; + (bytes, Backing::Owned(owned)) + }; // SAFETY: `bytes` points into `backing`, which the handle owns and // drops after the reader (see the field order of `ContainerHandle`). let bytes: &'static [u8] = unsafe { std::mem::transmute::<&[u8], &'static [u8]>(bytes) }; @@ -209,7 +208,7 @@ impl ContainerHandle { self.decode_bytes(py, name) } - /// Decode every tensor in parallel; ``(name, bytes)`` in index order. + /// Decode every tensor in parallel; ``(name, bytes)`` in write order. fn decode_model<'py>(&self, py: Python<'py>) -> PyResult)>> { self.decode_all(py) } @@ -221,7 +220,11 @@ impl ContainerHandle { } /// The BLAKE3 hash of the tensor's delta reference, or ``None``. - fn delta_hash<'py>(&self, py: Python<'py>, name: &str) -> PyResult>> { + fn delta_hash<'py>( + &self, + py: Python<'py>, + name: &str, + ) -> PyResult>> { self.tensor_delta_hash(py, name) } diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 26ceb57..4f7b183 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -657,7 +657,11 @@ fn decode_tensor_delta_hash<'py>( /// order. `skip_missing` has the same semantics as for :func:`decode_tensor`. #[pyfunction] #[pyo3(signature = (blob, skip_missing = false))] -fn list_tensor_names(py: Python<'_>, blob: PyBuffer, skip_missing: bool) -> PyResult> { +fn list_tensor_names( + py: Python<'_>, + blob: PyBuffer, + skip_missing: bool, +) -> PyResult> { ContainerHandle::open(py, blob, skip_missing)?.names_vec(py) } diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 36e47f4..51dc7d7 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -98,15 +98,16 @@ def _missing_(cls, value): def container_layout_name(layout: ContainerLayout | str) -> str: """Return the binding name (``"compact"`` or ``"legacy"``) of ``layout``. - A string that names no layout is returned unchanged, so the binding - rejects it with an error that names ``container_layout``. + ``layout`` is a :class:`ContainerLayout` or a string that names one (any + case). Anything else raises :class:`ValueError`. """ if isinstance(layout, ContainerLayout): return layout.value try: return ContainerLayout(layout).value except ValueError: - return layout + msg = f"unknown container_layout {layout!r} (expected 'compact' or 'legacy')" + raise ValueError(msg) from None class Lossy(Enum): diff --git a/python/ptwm/_magic.py b/python/ptwm/_magic.py index 3d48dd7..7f7c45a 100644 --- a/python/ptwm/_magic.py +++ b/python/ptwm/_magic.py @@ -63,10 +63,10 @@ def container_kind(buf: bytes | bytearray | memoryview) -> ContainerKind | None: def is_container(buf: bytes | bytearray | memoryview) -> bool: - """Whether ``buf`` starts with a multi-tensor container magic. + r"""Whether ``buf`` starts with a multi-tensor container magic. - True for ``\\x89PTWM``, ``\\x89PTWX`` and ``\\x89PTWC``: the formats that + True for ``\x89PTWM``, ``\x89PTWX`` and ``\x89PTWC``: the formats that :class:`ptwm.random_access.TensorIndex`, the stores and the transcoder - open. False for a ``\\x89PTWS`` single-tensor frame. + open. False for a ``\x89PTWS`` single-tensor frame. """ return container_kind(buf) in ("ptwm", "ptwx", "ptwc") diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index 4acd190..b2d6e8f 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -10,8 +10,13 @@ from ptwm import _core as _ext +from .._config import ContainerLayout, container_layout_name from ._errors import translate_errors +Buffer = bytes | bytearray | memoryview +"""A byte buffer the container functions accept. ``bytes`` and read-only +buffers are borrowed without a copy; writable ones are copied once.""" + __all__ = [ "ContainerHandle", "blake3_hash", @@ -193,7 +198,7 @@ def compress_model( streaming_chunk: int | None = None, codec_menu: list[int] | None = None, *, - container_layout: str = "legacy", + container_layout: ContainerLayout | str = "legacy", ) -> bytes: r"""Compress a list of tensors into a ``.ptwm`` container. @@ -210,9 +215,11 @@ def compress_model( empty intersection on a plane raises ``InvalidContainer``. ``container_layout``: ``"legacy"`` writes ``\x89PTWM`` (readable by every release), ``"compact"`` writes ``\x89PTWC`` (readable by v1.2 and later). - Other values raise :class:`ptwm.Error` (the binding's ``ValueError`` - passes through ``translate_errors``). + A :class:`ptwm.ContainerLayout` is accepted too. Other values raise + :class:`ptwm.Error` before anything is compressed (the binding itself + raises ``ValueError``). """ + layout = container_layout_name(container_layout) return bytes( _ext.compress_model( records, @@ -222,7 +229,7 @@ def compress_model( method_hint, streaming_chunk, codec_menu, - container_layout=container_layout, + container_layout=layout, ) ) @@ -289,7 +296,7 @@ def decode_compact( @translate_errors -def decode_tensor(blob: bytes, name: str, skip_missing: bool = False) -> bytes: +def decode_tensor(blob: Buffer, name: str, skip_missing: bool = False) -> bytes: """Decode a single tensor from a ``.ptwm`` container blob. Returns raw decompressed bytes. @@ -298,10 +305,10 @@ def decode_tensor(blob: bytes, name: str, skip_missing: bool = False) -> bytes: @translate_errors -def decode_model(blob: bytes, skip_missing: bool = False) -> list[tuple[str, bytes]]: +def decode_model(blob: Buffer, skip_missing: bool = False) -> list[tuple[str, bytes]]: """Decode every tensor from a ``.ptwm`` container blob in parallel. - Returns ``(name, raw_bytes)`` per tensor in container index order. Prefer + Returns ``(name, raw_bytes)`` per tensor in write order. Prefer this over per-tensor :func:`decode_tensor` for a full-model load — the decodes run across the rayon pool with the GIL released. """ @@ -310,7 +317,7 @@ def decode_model(blob: bytes, skip_missing: bool = False) -> list[tuple[str, byt @translate_errors def decode_tensor_shape( - blob: bytes, name: str, skip_missing: bool = False + blob: Buffer, name: str, skip_missing: bool = False ) -> tuple[list[int], str] | None: """Return ``(shape, dtype_name)`` for a tensor, or ``None``.""" return _ext.decode_tensor_shape(blob, name, skip_missing) @@ -318,7 +325,7 @@ def decode_tensor_shape( @translate_errors def decode_tensor_delta_hash( - blob: bytes, name: str, skip_missing: bool = False + blob: Buffer, name: str, skip_missing: bool = False ) -> bytes | None: """Return the 32-byte BLAKE3 hash of the delta reference for ``name``. @@ -329,7 +336,7 @@ def decode_tensor_delta_hash( @translate_errors -def list_tensor_names(blob: bytes, skip_missing: bool = False) -> list[str]: +def list_tensor_names(blob: Buffer, skip_missing: bool = False) -> list[str]: """Return every tensor name stored in a ``.ptwm`` container. Names come in write (body) order for both layouts. @@ -339,7 +346,7 @@ def list_tensor_names(blob: bytes, skip_missing: bool = False) -> list[str]: @translate_errors def list_plane_summaries( - blob: bytes, skip_missing: bool = False + blob: Buffer, skip_missing: bool = False ) -> list[tuple[str, int, int, int, int, int]]: """Return per-plane summaries for every tensor in a container. @@ -351,13 +358,13 @@ def list_plane_summaries( @translate_errors def decode_tensor_info( - blob: bytes, name: str, skip_missing: bool = False + blob: Buffer, name: str, skip_missing: bool = False ) -> tuple[list[int], str, int, int] | None: """Return ``(shape, dtype_name, dtype_code, input_format)`` or ``None``.""" return _ext.decode_tensor_info(blob, name, skip_missing) -def container_layout(blob: bytes | bytearray | memoryview) -> str | None: +def container_layout(blob: Buffer) -> str | None: r"""Return the multi-tensor layout of ``blob`` from its magic. ``"legacy"`` for ``\x89PTWM`` / ``\x89PTWX``, ``"compact"`` for @@ -389,9 +396,7 @@ class ContainerHandle: __slots__ = ("_h",) @translate_errors - def __init__( - self, blob: bytes | bytearray | memoryview, skip_missing: bool = False - ) -> None: + def __init__(self, blob: Buffer, skip_missing: bool = False) -> None: self._h = _ext.ContainerHandle(blob, skip_missing) @property @@ -437,8 +442,11 @@ def shape(self, name: str) -> tuple[list[int], str] | None: @translate_errors def plane_summaries(self) -> list[tuple[str, int, int, int, int, int]]: - """One ``(tensor_name, plane_index, role, codec_id, state_source, - payload_len)`` row per plane, tensors in write order.""" + """One row per plane, tensors in write order. + + Each row is ``(tensor_name, plane_index, role, codec_id, + state_source, payload_len)``. + """ return list(self._h.plane_summaries()) def __len__(self) -> int: # noqa: D105 diff --git a/python/ptwm/core/_decompressor.py b/python/ptwm/core/_decompressor.py index b66bbf4..0c8bd9e 100644 --- a/python/ptwm/core/_decompressor.py +++ b/python/ptwm/core/_decompressor.py @@ -11,6 +11,7 @@ from ..delta import decode as delta_decode from ..utils import DType + def _torch_dtype_for_code(dtype_code: int) -> torch.dtype | None: """Map a wire ``dtype_code`` to a concrete ``torch.dtype`` if supported.""" for m in DType: @@ -39,9 +40,10 @@ def decompress(self, data: bytes | memoryview) -> bytes | np.ndarray | torch.Ten Accepts a compact single-tensor frame (magic ``\x89PTWS``) or a full container in either layout (magic ``\x89PTWM``, ``\x89PTWX`` - or ``\x89PTWC``). The single-tensor container uses name ``"_"``. Raw decompressed bytes - come from ``decode_tensor``; shape/dtype/input_format (when present) - come from the CBOR shape-metadata map. Format recovery: + or ``\x89PTWC``). The single-tensor container uses name ``"_"``. + Raw decompressed bytes come from one + :class:`ptwm._rust.ContainerHandle`; shape/dtype/input_format (when + present) come from the CBOR shape-metadata map. Format recovery: - No metadata or ``input_format == BYTE`` → raw ``bytes``. - ``input_format == TORCH`` → :class:`torch.Tensor` reconstructed diff --git a/python/ptwm/integrations/_compress_safetensors.py b/python/ptwm/integrations/_compress_safetensors.py index d5c71bd..81da681 100644 --- a/python/ptwm/integrations/_compress_safetensors.py +++ b/python/ptwm/integrations/_compress_safetensors.py @@ -10,6 +10,7 @@ import torch from safetensors import safe_open +from .._config import ContainerLayout, container_layout_name from .._rust import shannon from ..classify import ( AuditLog, @@ -190,8 +191,9 @@ def compress_safetensors_file( use_user_cache: bool = True, fast: bool = False, fast_chunk: int = 512 * 1024, + container_layout: ContainerLayout | str = "legacy", ) -> AuditLog: - """Compress a safetensors file into a sharded ``.ptwm`` directory. + r"""Compress a safetensors file into a sharded ``.ptwm`` directory. Parameters ---------- @@ -211,12 +213,19 @@ def compress_safetensors_file( method_hint: Method hint forwarded to ``write_ptwm_shard`` (1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY). + container_layout: + Layout of every shard container: ``"legacy"`` (``\x89PTWM``, readable + by every release) or ``"compact"`` (``\x89PTWC``, v1.2 and later). + Forwarded to ``write_ptwm_shard``; a string or + :class:`ptwm.ContainerLayout`. Other values raise ``ValueError`` + before any tensor is read. Returns ------- AuditLog Per-tensor classification audit log. """ + container_layout = container_layout_name(container_layout) src = Path(src_path) out = Path(out_dir) out.mkdir(parents=True, exist_ok=True) @@ -339,6 +348,7 @@ def _build_chains_and_record( chains_per_tensor=shard_chains, method_hint=eff_method, streaming_chunk=eff_stream, + container_layout=container_layout, ) _persist_discovered(_new_templates, use_user_cache) return audit @@ -373,6 +383,7 @@ def _build_chains_and_record( chains_per_tensor=shard_chains, method_hint=eff_method, streaming_chunk=eff_stream, + container_layout=container_layout, ) else: if _write_mode_b_shell is None: @@ -388,6 +399,7 @@ def _build_chains_and_record( chains_per_tensor=shard_chains, method_hint=eff_method, streaming_chunk=eff_stream, + container_layout=container_layout, ) ptwm_bytes = tmp_ptwm.read_bytes() tmp_ptwm.unlink() diff --git a/python/ptwm/integrations/_hf.py b/python/ptwm/integrations/_hf.py index f05adb7..958208e 100644 --- a/python/ptwm/integrations/_hf.py +++ b/python/ptwm/integrations/_hf.py @@ -32,6 +32,7 @@ from packaging.version import Version from .. import _rust +from .._magic import is_container from ..core import Decompressor try: @@ -91,34 +92,46 @@ def _decompress_alongside(path: Path) -> Path: return out_path -def _is_multi_tensor(blob: bytes) -> bool: - """Return True if ``blob`` is a multi-tensor v3 ``.ptwm`` container. +def _open_multi_tensor(blob: bytes) -> _rust.ContainerHandle | None: + """Open ``blob`` if it is a multi-tensor ``.ptwm`` container. - A multi-tensor v3 container has a non-empty tensor name list, and none of - the names is ``"_"`` (the BYTE-format sentinel for single-tensor blobs). + A multi-tensor container (either layout) has a non-empty tensor name + list, and none of the names is ``"_"`` (the BYTE-format sentinel for + single-tensor blobs). Returns the open handle, or ``None`` for anything + else (a single-tensor blob, a compact frame or unreadable bytes). """ + if not is_container(blob): + return None try: - names = _rust.list_tensor_names(blob) + handle = _rust.ContainerHandle(blob) + names = handle.names() except Exception: # noqa: BLE001 - return False - return bool(names) and "_" not in names + return None + return handle if names and "_" not in names else None + + +def _is_multi_tensor(blob: bytes) -> bool: + """Return True if ``blob`` is a multi-tensor ``.ptwm`` container.""" + return _open_multi_tensor(blob) is not None -def _load_multi_tensor_ptwm(path: Path) -> dict[str, Any]: - """Decompress a multi-tensor v3 ``.ptwm`` container to a state dict. +def _load_multi_tensor_ptwm( + path: Path, handle: _rust.ContainerHandle | None = None +) -> dict[str, Any]: + """Decompress a multi-tensor ``.ptwm`` container to a state dict. - Reads ``path``, lists every tensor name via the embedded index, decodes - each tensor's raw bytes and shape metadata, and reconstructs the original - :class:`torch.Tensor` objects. + Opens ``path`` once (or uses ``handle``, already opened over its bytes), + decodes every tensor's raw bytes and shape metadata in write order, and + reconstructs the original :class:`torch.Tensor` objects. BF16 round-trip: the Rust layer stores BF16 data as raw ``uint16`` bytes; recovery uses the shared :func:`ptwm.utils.decode_ptwm_tensor` helper. """ from ..utils import decode_ptwm_tensor - blob = path.read_bytes() - names = _rust.list_tensor_names(blob) - return {name: decode_ptwm_tensor(blob, name) for name in names} + if handle is None: + handle = _rust.ContainerHandle(path.read_bytes()) + return {name: decode_ptwm_tensor(handle, name) for name in handle.names()} def materialize_hf_cache( @@ -279,12 +292,12 @@ def custom_load_state_dict( ) -> Any: path = Path(checkpoint_file) if any(path.name.endswith(s) for s in _COMPRESSED_SUFFIXES): - blob = path.read_bytes() - if _is_multi_tensor(blob): - # Multi-tensor v3 .ptwm: decode via tensor index and return + handle = _open_multi_tensor(path.read_bytes()) + if handle is not None: + # Multi-tensor .ptwm: decode via the tensor index and return # the state dict directly; the original load_state_dict does # not understand this format. - return _load_multi_tensor_ptwm(path) + return _load_multi_tensor_ptwm(path, handle) # Single-tensor BYTE-format blob (legacy path): decompress to a # sibling file and hand off to the original loader. path = _decompress_alongside(path) diff --git a/python/ptwm/integrations/_safetensors.py b/python/ptwm/integrations/_safetensors.py index ca98dfd..0efdabb 100644 --- a/python/ptwm/integrations/_safetensors.py +++ b/python/ptwm/integrations/_safetensors.py @@ -29,8 +29,8 @@ def decompress_safetensors_tensor(tensor: torch.Tensor) -> torch.Tensor: return decompressor.decompress(tensor.cpu().contiguous().numpy()) -def _decode_ptwm_tensor(blob: bytes, name: str) -> torch.Tensor: - """Decode a single named tensor from a v3 .ptwm blob. +def _decode_ptwm_tensor(blob: bytes | _rust.ContainerHandle, name: str) -> torch.Tensor: + """Decode a single named tensor from a .ptwm blob or open handle. Thin wrapper over the shared :func:`ptwm.utils.decode_ptwm_tensor` helper, kept for backward compatibility with existing call sites. @@ -54,15 +54,17 @@ def __init__(self, filename: str, framework: str, device: str = "cpu") -> None: path = Path(filename) self._mode_a = path.suffix == ".ptwm" self._mode_b = False - # Compressed blob for Mode A / Mode B — held for lazy per-tensor decode. - self._ptwm_blob: bytes | None = None + # Opened container for Mode A / Mode B, held for lazy per-tensor + # decode. The handle keeps the compressed bytes alive and parses the + # container once. + self._ptwm: _rust.ContainerHandle | None = None self._ptwm_keys: tuple[str, ...] = () self._underlying = None self.compressed_tensors_metadata: dict[str, Any] = {} if self._mode_a: - self._ptwm_blob = path.read_bytes() - self._ptwm_keys = tuple(_rust.list_tensor_names(self._ptwm_blob)) + self._ptwm = _rust.ContainerHandle(path.read_bytes()) + self._ptwm_keys = tuple(self._ptwm.names()) return # Mode B detection: open as safetensors first, sniff metadata. @@ -76,13 +78,15 @@ def __init__(self, filename: str, framework: str, device: str = "cpu") -> None: # __ptwm_payload__ tensor eagerly; that is safetensors' own # behaviour and cannot be deferred further at this layer. payload = self._underlying.get_tensor("__ptwm_payload__") - self._ptwm_blob = bytes(memoryview(payload.cpu().numpy()).cast("B")) - self._ptwm_keys = tuple(_rust.list_tensor_names(self._ptwm_blob)) + self._ptwm = _rust.ContainerHandle( + bytes(memoryview(payload.cpu().numpy()).cast("B")) + ) + self._ptwm_keys = tuple(self._ptwm.names()) else: self.compressed_tensors_metadata = get_compressed_tensors_metadata(meta) def keys(self) -> list[str]: - if self._ptwm_blob is not None: + if self._ptwm is not None: return list(self._ptwm_keys) return list(self._underlying.keys()) @@ -92,14 +96,14 @@ def metadata(self) -> dict[str, str] | None: return self._underlying.metadata() def get_tensor(self, name: str) -> torch.Tensor: - if self._ptwm_blob is not None: - return _decode_ptwm_tensor(self._ptwm_blob, name) + if self._ptwm is not None: + return _decode_ptwm_tensor(self._ptwm, name) if name not in self.compressed_tensors_metadata: return self._underlying.get_tensor(name) return decompress_safetensors_tensor(self._underlying.get_tensor(name)) def get_slice(self, name: str) -> Any: - if self._ptwm_blob is not None: + if self._ptwm is not None: msg = "get_slice is not supported on .ptwm/ptwm-shell containers" raise NotImplementedError(msg) if name not in self.compressed_tensors_metadata: diff --git a/python/ptwm/sharding/_writer.py b/python/ptwm/sharding/_writer.py index 6a7c87f..936cb91 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -7,6 +7,7 @@ from typing import Any from .. import _rust +from .._config import ContainerLayout from ..preprocessing._chains import ( ClassifierRole, build_chain_blobs, @@ -40,7 +41,7 @@ def compress_ptwm_blob( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, - container_layout: str = "legacy", + container_layout: ContainerLayout | str = "legacy", ) -> bytes: r"""Compress ``tensors`` into a single in-memory ``.ptwm`` container blob. @@ -49,8 +50,9 @@ def compress_ptwm_blob( each blob is produced by ``Chain.to_bytes()``. When ``None``, chains are auto-derived from each tensor's dtype via the STANDARD production table. ``method_hint``: 1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY. - ``container_layout``: ``"legacy"`` (``\x89PTWM``) or ``"compact"`` - (``\x89PTWC``). + ``container_layout``: ``"legacy"`` (``\x89PTWM``, readable by every + release) or ``"compact"`` (``\x89PTWC``, v1.2 and later), as a string or + :class:`ptwm.ContainerLayout`. Other values raise :class:`ptwm.Error`. """ records = [] auto_chains: list[list[bytes]] = [] @@ -93,7 +95,7 @@ def write_ptwm_shard( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, - container_layout: str = "legacy", + container_layout: ContainerLayout | str = "legacy", ) -> None: """Write a single ``.ptwm`` container holding ``tensors``. diff --git a/python/ptwm/stores/_lmdb.py b/python/ptwm/stores/_lmdb.py index 8e48b49..21509d9 100644 --- a/python/ptwm/stores/_lmdb.py +++ b/python/ptwm/stores/_lmdb.py @@ -28,6 +28,7 @@ import torch from .. import _rust +from .._config import ContainerLayout, container_layout_name from ..sharding import compress_ptwm_blob from ..utils import DType, decode_ptwm_tensor, raw_to_tensor from ._common import ( @@ -76,7 +77,7 @@ def write_lmdb( map_size: int = _DEFAULT_MAP_SIZE, method_hint: int = 3, chains_per_tensor: list[list[bytes]] | None = None, - container_layout: str = "legacy", + container_layout: ContainerLayout | str = "legacy", ) -> Path: r"""Write ``tensors`` into an LMDB environment at ``path``. @@ -91,6 +92,7 @@ def write_lmdb( msg = f"unknown LMDB layout {layout!r} (expected 'blob' or 'exploded')" raise ValueError(msg) + container_layout = container_layout_name(container_layout) lmdb = _require_lmdb() path = Path(path) blob = compress_ptwm_blob( @@ -370,7 +372,7 @@ def put_tensor( *, delta_reference: str | None = None, method_hint: int = 3, - container_layout: str | None = None, + container_layout: ContainerLayout | str | None = None, ) -> None: """Atomically insert or replace tensor ``name`` (blob layout only). @@ -406,8 +408,9 @@ def put_tensor( else: stored, store_shape, store_dtype = raw, shape, dtype - if container_layout is None: - container_layout = self._container_layout + container_layout = container_layout_name( + self._container_layout if container_layout is None else container_layout + ) blob = compress_ptwm_blob( [(name, stored, store_dtype, store_shape)], method_hint=method_hint, diff --git a/python/ptwm/stores/_webdataset.py b/python/ptwm/stores/_webdataset.py index 678dbd4..85edc23 100644 --- a/python/ptwm/stores/_webdataset.py +++ b/python/ptwm/stores/_webdataset.py @@ -21,6 +21,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Self +from .._config import ContainerLayout from ..sharding import compress_ptwm_blob, plan_shards from ._common import ( ExplodedModel, @@ -70,7 +71,7 @@ def write_webdataset( max_shard_size: int = _DEFAULT_MAX_SHARD_SIZE, method_hint: int = 3, chains_per_tensor: list[list[bytes]] | None = None, - container_layout: str = "legacy", + container_layout: ContainerLayout | str = "legacy", ) -> Path: """Write ``tensors`` as a WebDataset shard set rooted at ``out_dir``. diff --git a/python/ptwm/utils/_decode.py b/python/ptwm/utils/_decode.py index 4472b2f..9de4714 100644 --- a/python/ptwm/utils/_decode.py +++ b/python/ptwm/utils/_decode.py @@ -74,14 +74,23 @@ def raw_to_tensor( raise -def decode_ptwm_tensor(blob: bytes, name: str) -> torch.Tensor: - """Decode a single named tensor from a ``.ptwm`` container blob. +def decode_ptwm_tensor( + blob: bytes | bytearray | memoryview | _rust.ContainerHandle, name: str +) -> torch.Tensor: + """Decode a single named tensor from a ``.ptwm`` container. + + ``blob`` is the container bytes (either layout) or an open + :class:`ptwm._rust.ContainerHandle`. Pass a handle when decoding several + tensors from one container: bytes are parsed again on every call. Falls back to a flat ``uint8`` tensor (with a warning) when the container has no shape metadata for ``name``. """ - raw = bytes(_rust.decode_tensor(blob, name)) - info = _rust.decode_tensor_shape(blob, name) + handle = ( + blob if isinstance(blob, _rust.ContainerHandle) else _rust.ContainerHandle(blob) + ) + raw = handle.decode(name) + info = handle.shape(name) if info is None: warnings.warn( f"weights: no shape metadata for tensor {name!r}; " diff --git a/tests/ptwm/core/test_container_handle.py b/tests/ptwm/core/test_container_handle.py new file mode 100644 index 0000000..a3ffd0d --- /dev/null +++ b/tests/ptwm/core/test_container_handle.py @@ -0,0 +1,378 @@ +r"""ContainerHandle, the container_layout setting and the ``_magic`` helpers. + +Covers synthesis section 9 item 13: layout coercion, the handle agreeing with +the free functions in both layouts, write-order names (E14), and the +Decompressor / TensorIndex accepting ``\x89PTWM`` and ``\x89PTWC``. +""" + +from __future__ import annotations + +import mmap + +import pytest +import torch +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + Decompressor, + Format, + Framing, + Method, + _rust, +) +from ptwm._config import container_layout_name +from ptwm._exceptions import Error +from ptwm._magic import ( + MAGIC_LEN, + MAGIC_PTWC, + MAGIC_PTWM, + MAGIC_PTWS, + MAGIC_PTWX, + container_kind, + is_container, +) +from ptwm.integrations._hf import _load_multi_tensor_ptwm, _open_multi_tensor +from ptwm.integrations._safetensors import SafeOpen +from ptwm.random_access import TensorIndex +from ptwm.sharding import compress_ptwm_blob, write_ptwm_shard +from ptwm.sharding._writer import _default_chain_bytes +from ptwm.utils import DType, decode_ptwm_tensor + +LAYOUTS = ("legacy", "compact") +MAGIC = {"legacy": MAGIC_PTWM, "compact": MAGIC_PTWC} + +# Deliberately not sorted, and not in name-hash order. +NAMES = ("zeta.weight", "alpha.bias", "mid.weight", "beta", "omega.scale") + + +def _tensors() -> dict[str, torch.Tensor]: + torch.manual_seed(3) + return { + "zeta.weight": torch.randn(24, 16).bfloat16(), + "alpha.bias": torch.randn(16).float(), + "mid.weight": torch.randn(8, 12).to(torch.float16), + "beta": torch.randn(5, 3).float(), + "omega.scale": torch.randint(0, 256, (40,), dtype=torch.uint8), + } + + +def _raw(t: torch.Tensor) -> bytes: + view = {1: torch.uint8, 2: torch.uint16, 4: torch.uint32}[t.element_size()] + return t.flatten().contiguous().view(view).numpy().tobytes() + + +def _blob(layout: str) -> bytes: + tensors = _tensors() + items = [ + (n, _raw(tensors[n]), tensors[n].dtype, tuple(tensors[n].shape)) for n in NAMES + ] + return compress_ptwm_blob(items, method_hint=3, container_layout=layout) + + +def _delta_blob(layout: str) -> tuple[bytes, bytes]: + """A one-tensor container whose record carries a delta reference hash.""" + raw = bytes(range(256)) * 4 + ref_hash = bytes(range(32)) + dt = DType.from_dtype(torch.uint8) + chains = _default_chain_bytes(int(dt.code), (1024,), len(raw)) + blob = _rust.compress_model( + [("d", int(dt.code), 0, raw, [1024], dt.dtype_str, ref_hash)], + [chains], + True, + True, + container_layout=layout, + ) + return blob, ref_hash + + +# ── ContainerLayout coercion ───────────────────────────────────────────────── + + +def test_container_layout_enum_coerces_strings() -> None: + assert ContainerLayout("compact") is ContainerLayout.COMPACT + assert ContainerLayout("LEGACY") is ContainerLayout.LEGACY + assert ContainerLayout("Compact") is ContainerLayout.COMPACT + with pytest.raises(ValueError, match="ptwc"): + ContainerLayout("ptwc") + with pytest.raises(ValueError, match="1"): + ContainerLayout(1) + + +def test_container_layout_name() -> None: + assert container_layout_name(ContainerLayout.COMPACT) == "compact" + assert container_layout_name("LEGACY") == "legacy" + with pytest.raises(ValueError, match="container_layout"): + container_layout_name("bogus") + + +def test_compression_config_coerces_container_layout() -> None: + # The default stays LEGACY until the switch to the compact layout. + assert CompressionConfig().container_layout is ContainerLayout.LEGACY + cfg = CompressionConfig(container_layout="compact") + assert cfg.container_layout is ContainerLayout.COMPACT + cfg = CompressionConfig(container_layout="Legacy") + assert cfg.container_layout is ContainerLayout.LEGACY + with pytest.raises(ValueError, match="bogus"): + CompressionConfig(container_layout="bogus") + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_compressor_honours_container_layout(layout: str) -> None: + tensor = torch.randn(32, 16) + cfg = CompressionConfig( + input_format=Format.TORCH, + method=Method.HUFFMAN, + framing=Framing.CONTAINER, + container_layout=layout, + ) + blob = Compressor(cfg).compress(tensor) + assert blob[:MAGIC_LEN] == MAGIC[layout] + # Decompressor reads both layouts. + out = Decompressor().decompress(blob) + assert torch.equal(out, tensor) + + +def test_compact_frame_ignores_container_layout() -> None: + tensor = torch.randn(32, 16) + cfg = CompressionConfig(input_format=Format.TORCH, container_layout="compact") + blob = Compressor(cfg).compress(tensor) + assert blob[: len(MAGIC_PTWS)] == MAGIC_PTWS + assert torch.equal(Decompressor().decompress(blob), tensor) + + +@pytest.mark.parametrize( + "layout", [ContainerLayout.LEGACY, ContainerLayout.COMPACT, "COMPACT"] +) +def test_writers_accept_the_enum_and_any_case(tmp_path, layout) -> None: + expected = MAGIC[container_layout_name(layout)] + assert _rust.container_layout(_blob(container_layout_name(layout))) is not None + items = [("w", bytes(64), torch.float32, (16,))] + assert compress_ptwm_blob(items, container_layout=layout)[:MAGIC_LEN] == expected + out = tmp_path / "s.ptwm" + write_ptwm_shard(out, items, container_layout=layout) + assert out.read_bytes()[:MAGIC_LEN] == expected + dt = DType.from_dtype(torch.float32) + chains = _default_chain_bytes(int(dt.code), (16,), 64) + blob = _rust.compress_model( + [("w", int(dt.code), 0, bytes(64), [16], dt.dtype_str, None)], + [chains], + True, + True, + container_layout=layout, + ) + assert blob[:MAGIC_LEN] == expected + + +def test_binding_rejects_unknown_layout_with_value_error() -> None: + with pytest.raises(ValueError, match="container_layout"): + _rust._ext.compress_model([], [], True, True, container_layout="ptwc") + with pytest.raises(Error, match="container_layout"): + _rust.compress_model([], [], True, True, container_layout="ptwc") + + +# ── _magic ─────────────────────────────────────────────────────────────────── + + +def test_container_kind() -> None: + assert container_kind(MAGIC_PTWM + b"rest") == "ptwm" + assert container_kind(MAGIC_PTWX) == "ptwx" + assert container_kind(bytearray(MAGIC_PTWC)) == "ptwc" + assert container_kind(memoryview(MAGIC_PTWS + b"\x00\x01")) == "ptws" + # A 5-byte prefix is not enough for the 9-byte magics. + assert container_kind(MAGIC_PTWM[:5]) is None + assert container_kind(b"\x89PTWQ\r\n\x1a\n") is None + assert container_kind(b"") is None + assert container_kind(b"\x89PTW") is None + + +def test_is_container() -> None: + assert is_container(MAGIC_PTWM) + assert is_container(MAGIC_PTWX) + assert is_container(MAGIC_PTWC) + assert not is_container(MAGIC_PTWS + bytes(8)) + assert not is_container(b"not a container") + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_rust_container_layout(layout: str) -> None: + blob = _blob(layout) + assert _rust.container_layout(blob) == layout + assert _rust.container_layout(bytearray(blob)) == layout + assert container_kind(blob) == ("ptwm" if layout == "legacy" else "ptwc") + + +def test_rust_container_layout_of_other_bytes() -> None: + assert _rust.container_layout(MAGIC_PTWX) == "legacy" + assert _rust.container_layout(MAGIC_PTWS + bytes(16)) is None + assert _rust.container_layout(b"") is None + assert _rust.container_layout(b"\x89PTWM") is None + + +# ── ContainerHandle ────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_handle_matches_the_free_functions(layout: str) -> None: + blob = _blob(layout) + h = _rust.ContainerHandle(blob) + assert h.layout == layout + assert len(h) == len(NAMES) + assert "alpha.bias" in h + assert "missing" not in h + assert 3 not in h + assert h.names() == _rust.list_tensor_names(blob) + assert h.decode_model() == _rust.decode_model(blob) + assert h.plane_summaries() == _rust.list_plane_summaries(blob) + for name in NAMES: + assert h.decode(name) == _rust.decode_tensor(blob, name) + assert h.info(name) == _rust.decode_tensor_info(blob, name) + assert h.shape(name) == _rust.decode_tensor_shape(blob, name) + assert h.delta_hash(name) == _rust.decode_tensor_delta_hash(blob, name) + assert h.delta_hash(name) is None + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_handle_names_and_decode_model_in_write_order(layout: str) -> None: + blob = _blob(layout) + h = _rust.ContainerHandle(blob) + assert h.names() == list(NAMES) + assert [n for n, _ in h.decode_model()] == list(NAMES) + rows = h.plane_summaries() + assert list(dict.fromkeys(r[0] for r in rows)) == list(NAMES) + assert _rust.list_tensor_names(blob) == list(NAMES) + + +def test_handle_decodes_the_same_bytes_in_both_layouts() -> None: + legacy = _rust.ContainerHandle(_blob("legacy")) + compact = _rust.ContainerHandle(_blob("compact")) + tensors = _tensors() + for name in NAMES: + assert legacy.decode(name) == compact.decode(name) == _raw(tensors[name]) + assert legacy.shape(name) == compact.shape(name) + assert legacy.info(name) == compact.info(name) + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_handle_delta_hash(layout: str) -> None: + blob, ref_hash = _delta_blob(layout) + h = _rust.ContainerHandle(blob) + assert h.delta_hash("d") == ref_hash + assert _rust.decode_tensor_delta_hash(blob, "d") == ref_hash + + +def test_handle_borrows_read_only_buffers_and_copies_writable_ones(tmp_path) -> None: + blob = _blob("compact") + assert _rust.ContainerHandle(blob).zero_copy + assert _rust.ContainerHandle(memoryview(blob)).zero_copy + + ba = bytearray(blob) + h = _rust.ContainerHandle(ba) + assert not h.zero_copy + expected = h.decode("beta") + # The copy is independent of later writes to the bytearray. + ba[:] = bytes(len(ba)) + assert h.decode("beta") == expected + + path = tmp_path / "m.ptwm" + path.write_bytes(blob) + with ( + path.open("rb") as fh, + mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm, + ): + h = _rust.ContainerHandle(mm) + assert h.zero_copy + assert h.decode("beta") == expected + del h + + +def test_handle_keeps_the_buffer_alive() -> None: + h = _rust.ContainerHandle(bytes(_blob("legacy"))) + # The only reference to the bytes object is the handle's buffer export. + assert h.decode("alpha.bias") == _raw(_tensors()["alpha.bias"]) + + +def test_handle_errors_surface_as_ptwm_error() -> None: + with pytest.raises(Error): + _rust.ContainerHandle(b"not a container") + h = _rust.ContainerHandle(_blob("compact")) + with pytest.raises(Error, match="not found"): + h.decode("missing") + with pytest.raises(Error, match="not found"): + h.shape("missing") + + +# ── Python callers ─────────────────────────────────────────────────────────── + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_decompress_container_to_state_dict(layout: str) -> None: + blob = _blob(layout) + got = Decompressor().decompress_container_to_state_dict(blob) + assert list(got) == list(NAMES) + tensors = _tensors() + for name in NAMES: + assert got[name] == _raw(tensors[name]) + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_tensor_index_opens_both_layouts(tmp_path, layout: str) -> None: + path = tmp_path / "m.ptwm" + path.write_bytes(_blob(layout)) + idx = TensorIndex.open(path) + assert idx.layout == layout + assert idx.names() == list(NAMES) + tensors = _tensors() + for name in NAMES: + assert torch.equal(idx.get_tensor(name), tensors[name]) + + +def test_tensor_index_rejects_a_compact_frame(tmp_path) -> None: + path = tmp_path / "t.ptwm" + path.write_bytes( + Compressor(CompressionConfig(input_format=Format.TORCH)).compress( + torch.randn(8) + ) + ) + with pytest.raises(ValueError, match="not a .ptwm container"): + TensorIndex.open(path) + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_decode_ptwm_tensor_takes_bytes_or_a_handle(layout: str) -> None: + blob = _blob(layout) + h = _rust.ContainerHandle(blob) + tensors = _tensors() + for name in NAMES: + assert torch.equal(decode_ptwm_tensor(blob, name), tensors[name]) + assert torch.equal(decode_ptwm_tensor(h, name), tensors[name]) + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_hf_loader_reads_both_layouts(tmp_path, layout: str) -> None: + path = tmp_path / "model.safetensors.ptwm" + blob = _blob(layout) + path.write_bytes(blob) + assert _open_multi_tensor(blob) is not None + state = _load_multi_tensor_ptwm(path) + assert list(state) == list(NAMES) + tensors = _tensors() + for name in NAMES: + assert torch.equal(state[name], tensors[name]) + # A single-tensor blob is not a multi-tensor container. + single = Compressor( + CompressionConfig(framing=Framing.CONTAINER, container_layout=layout) + ).compress(b"\x00" * 64) + assert _open_multi_tensor(single) is None + assert _open_multi_tensor(b"\x89PTWS" + bytes(32)) is None + + +@pytest.mark.parametrize("layout", LAYOUTS) +def test_safe_open_mode_a_reads_both_layouts(tmp_path, layout: str) -> None: + path = tmp_path / "model.ptwm" + path.write_bytes(_blob(layout)) + tensors = _tensors() + with SafeOpen(str(path), framework="pt") as f: + assert f.keys() == list(NAMES) + for name in NAMES: + assert torch.equal(f.get_tensor(name), tensors[name]) diff --git a/tests/ptwm/integrations/test_compress_safetensors.py b/tests/ptwm/integrations/test_compress_safetensors.py index 6d07b50..57d7be0 100644 --- a/tests/ptwm/integrations/test_compress_safetensors.py +++ b/tests/ptwm/integrations/test_compress_safetensors.py @@ -456,3 +456,34 @@ def test_audit_log_cbor_roundtrip_with_discovered_chains(tmp_path): assert recovered.entries() == audit.entries() assert recovered.discovered_chains() == audit.discovered_chains() + + +@pytest.mark.parametrize("layout", ["legacy", "compact"]) +@pytest.mark.parametrize("max_shard_size", [None, 4096]) +def test_container_layout_is_forwarded_to_every_shard(tmp_path, layout, max_shard_size): + from ptwm._magic import MAGIC_LEN, MAGIC_PTWC, MAGIC_PTWM + + src = _save_nvfp4_safetensors(tmp_path) + out_dir = tmp_path / "out" + kwargs = {} if max_shard_size is None else {"max_shard_size": max_shard_size} + compress_safetensors_file(src, out_dir, mode="a", container_layout=layout, **kwargs) + shards = sorted(out_dir.glob("*.ptwm")) + assert shards + if max_shard_size is not None: + assert len(shards) > 1 + expected = MAGIC_PTWC if layout == "compact" else MAGIC_PTWM + for shard in shards: + assert shard.read_bytes()[:MAGIC_LEN] == expected + state = read_sharded_ptwm(out_dir) + ref = _build_synthetic_nvfp4() + assert set(state) == set(ref) + for name, arr in ref.items(): + assert bytes(state[name]) == arr.tobytes() + + +def test_unknown_container_layout_is_rejected_before_writing(tmp_path): + src = _save_nvfp4_safetensors(tmp_path) + out_dir = tmp_path / "out" + with pytest.raises(ValueError, match="container_layout"): + compress_safetensors_file(src, out_dir, container_layout="ptwc") + assert not out_dir.exists() diff --git a/tests/ptwm/integrations/test_mode_a_safe_open.py b/tests/ptwm/integrations/test_mode_a_safe_open.py index 63840e3..be7b2b4 100644 --- a/tests/ptwm/integrations/test_mode_a_safe_open.py +++ b/tests/ptwm/integrations/test_mode_a_safe_open.py @@ -92,19 +92,23 @@ def test_safe_open_ptwm_float32_tensor(tmp_path): def test_lazy_decoding_only_decodes_requested_tensor(tmp_path): - """decode_tensor must NOT be called during __init__, only from get_tensor.""" + """No tensor is decoded during __init__, only from get_tensor. + + SafeOpen decodes through one ``ContainerHandle``, so the spy wraps + ``ContainerHandle.decode``. + """ ptwm_path = _make_ptwm(tmp_path) import ptwm._rust as _rust_mod - original = _rust_mod.decode_tensor + original = _rust_mod.ContainerHandle.decode call_log: list[str] = [] - def recording_decode(blob: bytes, name: str) -> bytes: + def recording_decode(self: _rust_mod.ContainerHandle, name: str) -> bytes: call_log.append(name) - return original(blob, name) + return original(self, name) - _rust_mod.decode_tensor = recording_decode + _rust_mod.ContainerHandle.decode = recording_decode try: f = SafeOpen(str(ptwm_path), framework="pt") # __init__ must not have triggered any decode. @@ -115,7 +119,7 @@ def recording_decode(blob: bytes, name: str) -> bytes: # Only the requested tensor should have been decoded. assert call_log == [target_key], f"unexpected decode calls: {call_log}" finally: - _rust_mod.decode_tensor = original + _rust_mod.ContainerHandle.decode = original def test_safe_open_bf16_round_trip(tmp_path): diff --git a/tests/ptwm/integrations/test_mode_b_shell.py b/tests/ptwm/integrations/test_mode_b_shell.py index 5dfc16e..92ba498 100644 --- a/tests/ptwm/integrations/test_mode_b_shell.py +++ b/tests/ptwm/integrations/test_mode_b_shell.py @@ -112,21 +112,25 @@ def test_mode_b_get_tensor_dtype_and_shape(tmp_path): def test_mode_b_lazy_decoding(tmp_path): - """decode_tensor must NOT be called during SafeOpen.__init__ for Mode B.""" + """No tensor is decoded during SafeOpen.__init__ for Mode B. + + SafeOpen decodes through one ``ContainerHandle``, so the spy wraps + ``ContainerHandle.decode``. + """ blob = _make_ptwm_bytes(tmp_path) out = tmp_path / "shell.safetensors" write_mode_b_shell(out, blob) import ptwm._rust as _rust_mod - original = _rust_mod.decode_tensor + original = _rust_mod.ContainerHandle.decode call_log: list[str] = [] - def recording_decode(b: bytes, name: str) -> bytes: + def recording_decode(self: _rust_mod.ContainerHandle, name: str) -> bytes: call_log.append(name) - return original(b, name) + return original(self, name) - _rust_mod.decode_tensor = recording_decode + _rust_mod.ContainerHandle.decode = recording_decode try: f = SafeOpen(str(out), framework="pt") assert call_log == [], f"decode called during __init__: {call_log}" @@ -135,7 +139,7 @@ def recording_decode(b: bytes, name: str) -> bytes: _t = f.get_tensor(target_key) assert call_log == [target_key], f"unexpected decode calls: {call_log}" finally: - _rust_mod.decode_tensor = original + _rust_mod.ContainerHandle.decode = original def test_mode_b_bf16_round_trip(tmp_path): diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py index 1192415..0b0bb03 100644 --- a/tests/ptwm/stores/test_compact_layout.py +++ b/tests/ptwm/stores/test_compact_layout.py @@ -12,8 +12,9 @@ import pytest import torch -from ptwm import _rust +from ptwm import ContainerLayout, _rust from ptwm._exceptions import Error +from ptwm._magic import MAGIC_LEN, MAGIC_PTWC, MAGIC_PTWM, container_kind from ptwm.sharding import compress_ptwm_blob, write_ptwm_shard from ptwm.stores import ( LmdbWeightStore, @@ -32,7 +33,7 @@ from ptwm.stores._lmdb import KEY_FMT, KEY_REG from ptwm.stores._webdataset import INDEX_NAME, SHARDMETA_MEMBER -MAGIC = {"legacy": b"\x89PTWM", "compact": b"\x89PTWC"} +MAGIC = {"legacy": MAGIC_PTWM, "compact": MAGIC_PTWC} LAYOUTS = ("legacy", "compact") @@ -47,7 +48,7 @@ def test_compress_ptwm_blob_honours_the_layout(tmp_path, store_tensors, layout): assert blob[:9].startswith(MAGIC[layout]) out = tmp_path / "m.ptwm" write_ptwm_shard(out, store_tensors, container_layout=layout) - assert out.read_bytes()[:5] == MAGIC[layout] + assert out.read_bytes()[:MAGIC_LEN] == MAGIC[layout] # Both layouts decode to the same bytes. assert _reference_raw(store_tensors, layout) == _reference_raw( store_tensors, "legacy" @@ -108,7 +109,7 @@ def test_raw_from_members_rebuilds_one_tensor(store_tensors, layout): single = _rust.implode_ptwm( model.registry_json, [model.tensor_json(name)], members ) - assert single[:5] == MAGIC[layout] + assert single[:MAGIC_LEN] == MAGIC[layout] raw = raw_from_members(model.registry_json, model.tensor_json(name), members) assert raw == want[name] @@ -135,7 +136,7 @@ def test_lmdb_compact_round_trip(tmp_path, store_tensors, torch_tensors, store_l assert set(store.names()) == {n for n, *_ in store_tensors} if store_layout == "blob": for name, *_ in store_tensors: - assert store._get(b"t/" + name.encode())[:5] == MAGIC["compact"] + assert store._get(b"t/" + name.encode())[:MAGIC_LEN] == MAGIC["compact"] else: reg = json.loads(store._get(KEY_REG)) assert reg["ptwm_format"] == "ptwm-manifest-c" @@ -160,11 +161,11 @@ def test_lmdb_put_tensor_in_either_layout(tmp_path, raw_of): store.put_tensor( "w", raw_of(new), torch.float32, (8,), container_layout="legacy" ) - assert store._get(b"t/w")[:5] == MAGIC["legacy"] + assert store._get(b"t/w")[:MAGIC_LEN] == MAGIC["legacy"] store.put_tensor( "v", raw_of(new), torch.float32, (8,), container_layout="compact" ) - assert store._get(b"t/v")[:5] == MAGIC["compact"] + assert store._get(b"t/v")[:MAGIC_LEN] == MAGIC["compact"] store.put_tensor( "d", raw_of(step), @@ -173,7 +174,7 @@ def test_lmdb_put_tensor_in_either_layout(tmp_path, raw_of): delta_reference="v", container_layout="compact", ) - assert store._get(b"t/d")[:5] == MAGIC["compact"] + assert store._get(b"t/d")[:MAGIC_LEN] == MAGIC["compact"] with LmdbWeightStore.open(path) as store: assert torch.equal(store.get_tensor("w"), new) assert torch.equal(store.get_tensor("v"), new) @@ -198,12 +199,12 @@ def test_lmdb_put_tensor_inherits_store_layout(tmp_path, raw_of, layout): assert fmt["container_layout"] == "compact" store.put_tensor("w", raw_of(new), torch.float32, (8,)) store.put_tensor("d", raw_of(new + 1), torch.float32, (8,), delta_reference="w") - assert store._get(b"t/w")[:5] == MAGIC[layout] - assert store._get(b"t/d")[:5] == MAGIC[layout] + assert store._get(b"t/w")[:MAGIC_LEN] == MAGIC[layout] + assert store._get(b"t/d")[:MAGIC_LEN] == MAGIC[layout] # The layout survives a reopen. with LmdbWeightStore.open(path, readonly=False) as store: store.put_tensor("v", raw_of(new), torch.float32, (8,)) - assert store._get(b"t/v")[:5] == MAGIC[layout] + assert store._get(b"t/v")[:MAGIC_LEN] == MAGIC[layout] with LmdbWeightStore.open(path) as store: assert torch.equal(store.get_tensor("w"), new) assert torch.equal(store.get_tensor("v"), new) @@ -252,3 +253,28 @@ def test_webdataset_compact_round_trip(tmp_path, store_tensors, max_shard_size): for name, tensor in single.stream_tensors(): assert tensor.numel() > 0 assert name in want + + +def test_stores_accept_the_layout_enum(tmp_path, store_tensors, raw_of): + path = tmp_path / "m.lmdb" + write_lmdb( + path, store_tensors, layout="blob", container_layout=ContainerLayout.COMPACT + ) + with LmdbWeightStore.open(path, readonly=False) as store: + fmt = json.loads(store._get(KEY_FMT)) + assert fmt["container_layout"] == "compact" + for name, *_ in store_tensors: + assert container_kind(store._get(b"t/" + name.encode())) == "ptwc" + store.put_tensor( + "x", + raw_of(torch.ones(4)), + torch.float32, + (4,), + container_layout=ContainerLayout.LEGACY, + ) + assert container_kind(store._get(b"t/x")) == "ptwm" + + wd = tmp_path / "m.wds" + write_webdataset(wd, store_tensors, container_layout=ContainerLayout.COMPACT) + index = json.loads((wd / INDEX_NAME).read_text()) + assert index["registry"]["ptwm_format"] == "ptwm-manifest-c" diff --git a/tests/ptwm/stores/test_lmdb.py b/tests/ptwm/stores/test_lmdb.py index 201eec3..26060a9 100644 --- a/tests/ptwm/stores/test_lmdb.py +++ b/tests/ptwm/stores/test_lmdb.py @@ -7,6 +7,7 @@ import pytest import torch from ptwm import _rust +from ptwm._magic import is_container from ptwm.stores import LmdbWeightStore, detect_format, write_lmdb from ptwm.stores._lmdb import KEY_FMT, KEY_IDX, KEY_REG, LMDB_FORMAT @@ -25,7 +26,7 @@ def test_blob_layout_keys(tmp_path, store_tensors): # Every tensor is a self-contained single-tensor container. for name, *_ in store_tensors: blob = store._get(b"t/" + name.encode()) - assert blob[:5] == b"\x89PTWM" + assert is_container(blob) def test_exploded_layout_keys(tmp_path, store_tensors): diff --git a/tests/ptwm/stores/test_transcoder.py b/tests/ptwm/stores/test_transcoder.py index 0f3e6f5..8774cc3 100644 --- a/tests/ptwm/stores/test_transcoder.py +++ b/tests/ptwm/stores/test_transcoder.py @@ -12,6 +12,7 @@ import numpy as np import torch from ptwm import _rust +from ptwm._magic import is_container from ptwm.integrations import compress_safetensors_file from ptwm.sharding import compress_ptwm_blob from ptwm.stores import jcs_canonicalize @@ -19,12 +20,12 @@ def _ptwm_blobs_under(root: Path) -> list[bytes]: - """Collect every native ``.ptwm`` container blob under ``root``.""" + """Collect every native ``.ptwm`` container blob (either layout) under ``root``.""" blobs: list[bytes] = [] for p in sorted(root.rglob("*")): if p.is_file(): data = p.read_bytes() - if data[:5] == b"\x89PTWM": + if is_container(data): blobs.append(data) return blobs diff --git a/tests/ptwm/stores/test_webdataset.py b/tests/ptwm/stores/test_webdataset.py index 25a8244..c16db94 100644 --- a/tests/ptwm/stores/test_webdataset.py +++ b/tests/ptwm/stores/test_webdataset.py @@ -6,6 +6,7 @@ import tarfile from ptwm import _rust +from ptwm._magic import MAGIC_LEN, is_container from ptwm.sharding import compress_ptwm_blob from ptwm.stores import WebDatasetWeightStore, detect_format, write_webdataset from ptwm.stores._webdataset import INDEX_NAME, SHARDMETA_MEMBER, WDS_FORMAT @@ -31,7 +32,7 @@ def test_shard_and_index_layout(tmp_path, store_tensors): if info.name.endswith(".bin"): f = tar.extractfile(info) assert f is not None - assert f.read(5) != b"\x89PTWM" + assert not is_container(f.read(MAGIC_LEN)) index = json.loads((wd / INDEX_NAME).read_text()) assert index["ptwm_format"] == WDS_FORMAT diff --git a/tests/ptwm/wire_format/test_legacy_hardening.py b/tests/ptwm/wire_format/test_legacy_hardening.py index b80f410..659c0d0 100644 --- a/tests/ptwm/wire_format/test_legacy_hardening.py +++ b/tests/ptwm/wire_format/test_legacy_hardening.py @@ -100,6 +100,63 @@ def test_shape_and_plane_summaries_honour_skip_missing() -> None: ) +def test_container_handle_honours_skip_missing() -> None: + blob = _legacy_blob() + patched = _with_missing_extension(blob) + with pytest.raises(Exception, match="(?i)missing|flavor"): + _rust.ContainerHandle(patched) + h = _rust.ContainerHandle(patched, skip_missing=True) + ref = _rust.ContainerHandle(blob) + (name,) = ref.names() + assert h.names() == [name] + assert h.shape(name) == ref.shape(name) + assert h.info(name) == ref.info(name) + assert h.plane_summaries() == ref.plane_summaries() + assert h.decode(name) == ref.decode(name) + + +def test_inspect_verifies_the_extension_table_hash(tmp_path) -> None: + from ptwm._core import list_extension_table_entries + + blob = bytearray(_legacy_blob()) + # Flip one byte inside the table: the header hash no longer matches. + blob[HEADER_LEN + 40] ^= 0x01 + path = tmp_path / "bad_hash.ptwm" + path.write_bytes(bytes(blob)) + with pytest.raises(ValueError, match="(?i)hash"): + list_extension_table_entries(str(path)) + + +def test_inspect_reads_both_layouts(tmp_path) -> None: + from ptwm._core import list_extension_table_entries + + tensor = torch.arange(64 * 32, dtype=torch.float32).reshape(64, 32) + entries = {} + for layout in ("legacy", "compact"): + blob = Compressor( + CompressionConfig( + input_format=Format.TORCH, + method=Method.HUFFMAN, + framing=Framing.CONTAINER, + container_layout=layout, + ) + ).compress(tensor) + path = tmp_path / f"{layout}.ptwm" + path.write_bytes(blob) + entries[layout] = list_extension_table_entries(str(path)) + assert entries["legacy"] + # The compact table expands to the same built-in entries. + assert {e[0] for e in entries["compact"]} <= {e[0] for e in entries["legacy"]} + assert entries["compact"] + + path = tmp_path / "frame.ptwm" + path.write_bytes( + Compressor(CompressionConfig(input_format=Format.TORCH)).compress(tensor) + ) + with pytest.raises(ValueError, match="(?i)magic"): + list_extension_table_entries(str(path)) + + def test_inspect_rejects_extension_table_offset_overflow(tmp_path) -> None: from ptwm._core import list_extension_table_entries From 8edf8273cd21e444e4767cc63bfe91b11ea675c7 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 02:02:19 +0000 Subject: [PATCH 21/34] fix(S8-python): address review findings - ContainerHandle borrows a buffer only when its owner is immutable: an exact bytes object or an mmap opened with ACCESS_READ, directly or through memoryview links. Read-only views of mutable memory (memoryview(bytearray).toreadonly(), numpy arrays with writeable=False) are now copied, because the owner could still write them while the reader runs with the GIL released. The free container functions use the same rule. SAFETY comments and docstrings corrected. - compress_ptwm_blob, write_ptwm_shard and write_webdataset validate container_layout first and raise ValueError, before any file or directory is created. - test_lmdb pins the legacy magic of the default store layout again. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-py/src/handle.rs | 83 +++++++++++++++++++----- crates/ptwm-py/src/lib.rs | 20 +++--- python/ptwm/_rust/__init__.py | 14 ++-- python/ptwm/sharding/_writer.py | 12 ++-- python/ptwm/stores/_webdataset.py | 7 +- tests/ptwm/core/test_container_handle.py | 57 ++++++++++++++++ tests/ptwm/stores/test_compact_layout.py | 13 +++- tests/ptwm/stores/test_lmdb.py | 5 +- 8 files changed, 170 insertions(+), 41 deletions(-) diff --git a/crates/ptwm-py/src/handle.rs b/crates/ptwm-py/src/handle.rs index 05cb14c..0663c53 100644 --- a/crates/ptwm-py/src/handle.rs +++ b/crates/ptwm-py/src/handle.rs @@ -6,16 +6,24 @@ //! free functions in `lib.rs` (`decode_tensor`, `list_tensor_names`, ...) //! open a temporary handle per call. //! -//! A read-only, C-contiguous buffer (``bytes``, a read-only ``memoryview`` or -//! ``mmap``) is borrowed without a copy: the handle holds the buffer export, -//! which keeps the object alive and its memory pinned. A writable or -//! non-contiguous buffer (``bytearray``, a writable ``mmap`` or numpy array) -//! is copied once, because the caller could change it while the reader -//! holds slices into it. +//! A buffer is borrowed without a copy only when its memory cannot change +//! while the handle exists: a C-contiguous export of ``bytes``, or of an +//! ``mmap`` opened with ``ACCESS_READ``, directly or through any chain of +//! ``memoryview`` objects. The handle holds the buffer export, which keeps +//! the object alive and its memory pinned. Every other buffer is copied +//! once, including read-only views of mutable objects (for example +//! ``memoryview(bytearray).toreadonly()`` or a numpy array with +//! ``writeable=False`` over a writable base): a read-only export only says +//! that the export cannot be written through, not that the owner cannot +//! write the memory, and the reader runs with the GIL released. +//! +//! A read-only ``mmap`` maps a file with ``MAP_SHARED``. Another process (or +//! a second, writable mapping of the same file) can still change the file; +//! the caller must not modify a file while a handle borrows it. use pyo3::buffer::PyBuffer; use pyo3::prelude::*; -use pyo3::types::PyBytes; +use pyo3::types::{PyBytes, PyMemoryView}; use ptwm_core::PtwmCoreError; use ptwm_core::container::{ContainerLayout, ContainerReader}; @@ -31,8 +39,9 @@ pub type PlaneSummary = (String, u32, u8, u16, u8, u64); /// The memory a handle's reader points into. enum Backing { - /// A read-only, C-contiguous buffer export. Holding it keeps the - /// exporting object alive and its memory unchanged. + /// A C-contiguous export of immutable memory (see + /// [`immutable_backing`]). Holding it keeps the exporting object alive + /// and its memory pinned. Borrowed(#[allow(dead_code)] PyBuffer), /// A private copy of a buffer that could change under the reader. Owned(#[allow(dead_code)] Box<[u8]>), @@ -56,18 +65,60 @@ fn runtime_err(e: PtwmCoreError) -> PyErr { pyo3::exceptions::PyRuntimeError::new_err(e.to_string()) } +/// Whether the memory behind `obj`'s buffer cannot change for as long as +/// an export of it is held. +/// +/// True for an exact ``bytes`` object and for an exact ``mmap.mmap`` whose +/// own export is read-only (``ACCESS_READ``), reached directly or through +/// ``memoryview.obj`` links. Everything else (``bytearray``, writable +/// ``mmap``, numpy arrays, subclasses that could override the buffer +/// protocol, memoryviews without an owner object) is treated as mutable. +fn immutable_backing(obj: &Bound<'_, PyAny>) -> PyResult { + let py = obj.py(); + let mut cur = obj.clone(); + // `memoryview.obj` of a memoryview over a memoryview names the base + // object, so the loop ends after one step in practice; the bound guards + // against an unexpected cycle. + for _ in 0..16 { + if cur.is_exact_instance_of::() { + let base = cur.getattr("obj")?; + if base.is_none() { + return Ok(false); + } + cur = base; + continue; + } + if cur.is_exact_instance_of::() { + return Ok(true); + } + let mmap_type = py.import("mmap")?.getattr("mmap")?; + if cur.get_type().is(&mmap_type) { + return Ok(PyBuffer::::get(&cur)?.readonly()); + } + return Ok(false); + } + Ok(false) +} + impl ContainerHandle { - /// Open `buffer`, borrowing it when it is read-only and C-contiguous. - pub fn open(py: Python<'_>, buffer: PyBuffer, skip_missing: bool) -> PyResult { - let (bytes, backing): (&[u8], Backing) = if buffer.readonly() && buffer.is_c_contiguous() { + /// Open the buffer `obj` exports. The buffer is borrowed when its memory + /// cannot change (see [`immutable_backing`]) and is C-contiguous; + /// otherwise it is copied once. + pub fn open(py: Python<'_>, obj: &Bound<'_, PyAny>, skip_missing: bool) -> PyResult { + let buffer = PyBuffer::::get(obj)?; + let borrow = buffer.is_c_contiguous() && immutable_backing(obj)?; + let (bytes, backing): (&[u8], Backing) = if borrow { let len = buffer.len_bytes(); let bytes: &[u8] = if len == 0 { &[] } else { // SAFETY: the buffer is C-contiguous and `len` bytes long. - // The export stays alive in `Backing::Borrowed` for as - // long as the reader, and a read-only export cannot be - // written through. + // The export stays alive in `Backing::Borrowed` for as long + // as the reader, which keeps the owner alive and (for + // `mmap`) prevents it from being closed or resized. The + // owner is `bytes` or a read-only `mmap`, so no Python code + // can write the memory while the reader holds slices into + // it, even with the GIL released. unsafe { std::slice::from_raw_parts(buffer.buf_ptr() as *const u8, len) } }; (bytes, Backing::Borrowed(buffer)) @@ -181,7 +232,7 @@ impl ContainerHandle { impl ContainerHandle { #[new] #[pyo3(signature = (buffer, skip_missing = false))] - fn py_new(py: Python<'_>, buffer: PyBuffer, skip_missing: bool) -> PyResult { + fn py_new(py: Python<'_>, buffer: &Bound<'_, PyAny>, skip_missing: bool) -> PyResult { Self::open(py, buffer, skip_missing) } diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 4f7b183..2fed6ce 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -581,13 +581,13 @@ fn decode_compact<'py>( } // The functions below open a temporary `ContainerHandle` per call (see -// `handle.rs`): `blob` is any buffer, borrowed when read-only and -// C-contiguous. Open one `ContainerHandle` instead for several queries. +// `handle.rs`): `blob` is any buffer, borrowed when its memory cannot +// change (see `ContainerHandle::open`) and copied otherwise. Open one `ContainerHandle` instead for several queries. /// Decode a single tensor from a container blob (either layout). /// /// `blob` is the full container (any buffer; ``bytes`` is borrowed without a -/// copy). `name` is the tensor name string. Returns the raw decompressed +/// copy, other buffers may be copied). `name` is the tensor name string. Returns the raw decompressed /// ``bytes`` via the chain-based decode path. /// /// `skip_missing` — when ``True`` open the container even if non-builtin @@ -602,7 +602,7 @@ fn decode_compact<'py>( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor<'py>( py: Python<'py>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, name: &str, skip_missing: bool, ) -> PyResult> { @@ -618,7 +618,7 @@ fn decode_tensor<'py>( #[pyo3(signature = (blob, skip_missing = false))] fn decode_model<'py>( py: Python<'py>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, skip_missing: bool, ) -> PyResult)>> { ContainerHandle::open(py, blob, skip_missing)?.decode_all(py) @@ -632,7 +632,7 @@ fn decode_model<'py>( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_info( py: Python<'_>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, name: &str, skip_missing: bool, ) -> PyResult> { @@ -646,7 +646,7 @@ fn decode_tensor_info( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_delta_hash<'py>( py: Python<'py>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, name: &str, skip_missing: bool, ) -> PyResult>> { @@ -659,7 +659,7 @@ fn decode_tensor_delta_hash<'py>( #[pyo3(signature = (blob, skip_missing = false))] fn list_tensor_names( py: Python<'_>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, skip_missing: bool, ) -> PyResult> { ContainerHandle::open(py, blob, skip_missing)?.names_vec(py) @@ -673,7 +673,7 @@ fn list_tensor_names( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_shape( py: Python<'_>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, name: &str, skip_missing: bool, ) -> PyResult, String)>> { @@ -689,7 +689,7 @@ fn decode_tensor_shape( #[pyo3(signature = (blob, skip_missing = false))] fn list_plane_summaries( py: Python<'_>, - blob: PyBuffer, + blob: &Bound<'_, PyAny>, skip_missing: bool, ) -> PyResult> { ContainerHandle::open(py, blob, skip_missing)?.plane_summary_rows(py) diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index b2d6e8f..7bcf7ed 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -14,8 +14,8 @@ from ._errors import translate_errors Buffer = bytes | bytearray | memoryview -"""A byte buffer the container functions accept. ``bytes`` and read-only -buffers are borrowed without a copy; writable ones are copied once.""" +"""A byte buffer the container functions accept. ``bytes`` and a read-only +``mmap`` are borrowed without a copy; every other buffer is copied once.""" __all__ = [ "ContainerHandle", @@ -384,9 +384,13 @@ class ContainerHandle: ...) when a container is queried more than once: each free function opens the container again. - ``blob`` is any buffer. ``bytes``, a read-only ``memoryview`` and a - read-only ``mmap`` are borrowed without a copy (the handle keeps them - alive); a writable buffer such as ``bytearray`` is copied once. + ``blob`` is any buffer. ``bytes`` and an ``mmap`` opened with + ``ACCESS_READ`` (directly or through a ``memoryview``) are borrowed + without a copy; the handle keeps them alive. Every other buffer is + copied once, including a read-only view of mutable memory such as + ``memoryview(bytearray).toreadonly()``, because its owner could still + change it. The file behind a borrowed ``mmap`` must not change while + the handle exists. ``skip_missing`` opens the container even when non-builtin extensions it references are not installed; tensors that need them fail at decode. diff --git a/python/ptwm/sharding/_writer.py b/python/ptwm/sharding/_writer.py index 936cb91..fc4d40b 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -7,7 +7,7 @@ from typing import Any from .. import _rust -from .._config import ContainerLayout +from .._config import ContainerLayout, container_layout_name from ..preprocessing._chains import ( ClassifierRole, build_chain_blobs, @@ -52,8 +52,10 @@ def compress_ptwm_blob( ``method_hint``: 1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY. ``container_layout``: ``"legacy"`` (``\x89PTWM``, readable by every release) or ``"compact"`` (``\x89PTWC``, v1.2 and later), as a string or - :class:`ptwm.ContainerLayout`. Other values raise :class:`ptwm.Error`. + :class:`ptwm.ContainerLayout`. Other values raise :class:`ValueError` + before anything is compressed. """ + layout = container_layout_name(container_layout) records = [] auto_chains: list[list[bytes]] = [] for name, raw, dtype_in, shape in tensors: @@ -82,7 +84,7 @@ def compress_ptwm_blob( emit_plane_crc, method_hint, streaming_chunk, - container_layout=container_layout, + container_layout=layout, ) @@ -99,7 +101,9 @@ def write_ptwm_shard( ) -> None: """Write a single ``.ptwm`` container holding ``tensors``. - See :func:`compress_ptwm_blob` for the parameter semantics. + See :func:`compress_ptwm_blob` for the parameter semantics. An unknown + ``container_layout`` raises :class:`ValueError` before ``out_path`` is + written. """ blob = compress_ptwm_blob( tensors, diff --git a/python/ptwm/stores/_webdataset.py b/python/ptwm/stores/_webdataset.py index 85edc23..f6137dc 100644 --- a/python/ptwm/stores/_webdataset.py +++ b/python/ptwm/stores/_webdataset.py @@ -21,7 +21,7 @@ from pathlib import Path from typing import TYPE_CHECKING, Any, Self -from .._config import ContainerLayout +from .._config import ContainerLayout, container_layout_name from ..sharding import compress_ptwm_blob, plan_shards from ._common import ( ExplodedModel, @@ -80,8 +80,11 @@ def write_webdataset( via :func:`ptwm.sharding.plan_shards` so a tensor never straddles a shard. ``container_layout`` (``"legacy"`` or ``"compact"``) selects the container the tensors are compressed into; the registry manifest records - it, and readers rebuild each tensor in that layout. + it, and readers rebuild each tensor in that layout. An unknown + ``container_layout`` raises :class:`ValueError` before ``out_dir`` is + created. """ + container_layout = container_layout_name(container_layout) out_dir = Path(out_dir) out_dir.mkdir(parents=True, exist_ok=True) diff --git a/tests/ptwm/core/test_container_handle.py b/tests/ptwm/core/test_container_handle.py index a3ffd0d..d3b1dda 100644 --- a/tests/ptwm/core/test_container_handle.py +++ b/tests/ptwm/core/test_container_handle.py @@ -9,6 +9,7 @@ import mmap +import numpy as np import pytest import torch from ptwm import ( @@ -286,6 +287,62 @@ def test_handle_borrows_read_only_buffers_and_copies_writable_ones(tmp_path) -> del h +def test_handle_copies_read_only_views_of_mutable_memory(tmp_path) -> None: + # A read-only export does not make the owner immutable: the handle must + # copy, so later writes to the owner cannot reach the reader. + blob = _blob("compact") + expected = _rust.ContainerHandle(blob).decode("beta") + + ba = bytearray(blob) + h = _rust.ContainerHandle(memoryview(ba).toreadonly()) + assert not h.zero_copy + ba[:] = bytes(len(ba)) + assert h.decode("beta") == expected + assert "beta" in h.names() + + arr = np.frombuffer(bytearray(blob), dtype=np.uint8) + view = arr.view() + view.flags.writeable = False + h = _rust.ContainerHandle(view) + assert not h.zero_copy + arr[:] = 0 + assert h.decode("beta") == expected + + # Views of immutable memory are still borrowed, through any chain of + # memoryviews. + assert _rust.ContainerHandle(memoryview(memoryview(blob))).zero_copy + assert _rust.ContainerHandle(memoryview(blob).toreadonly()).zero_copy + + # A writable mmap is copied, also behind a read-only memoryview. + path = tmp_path / "w.ptwm" + path.write_bytes(blob) + with path.open("r+b") as fh, mmap.mmap(fh.fileno(), 0) as mm: + h = _rust.ContainerHandle(memoryview(mm).toreadonly()) + assert not h.zero_copy + mm[:] = bytes(len(mm)) + assert h.decode("beta") == expected + del h + + # A read-only mmap behind a memoryview is borrowed. + path.write_bytes(blob) + with ( + path.open("rb") as fh, + mmap.mmap(fh.fileno(), 0, access=mmap.ACCESS_READ) as mm, + memoryview(mm) as mv, + ): + h = _rust.ContainerHandle(mv) + assert h.zero_copy + assert h.decode("beta") == expected + del h + + +def test_free_functions_copy_read_only_views_of_mutable_memory() -> None: + blob = _blob("legacy") + view = memoryview(bytearray(blob)).toreadonly() + assert _rust.list_tensor_names(view) == list(NAMES) + assert _rust.decode_tensor(view, "beta") == _rust.decode_tensor(blob, "beta") + + def test_handle_keeps_the_buffer_alive() -> None: h = _rust.ContainerHandle(bytes(_blob("legacy"))) # The only reference to the bytes object is the handle's buffer export. diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py index 0b0bb03..bb3d098 100644 --- a/tests/ptwm/stores/test_compact_layout.py +++ b/tests/ptwm/stores/test_compact_layout.py @@ -55,9 +55,18 @@ def test_compress_ptwm_blob_honours_the_layout(tmp_path, store_tensors, layout): ) -def test_unknown_container_layout_is_rejected(store_tensors): - with pytest.raises(Error, match="container_layout"): +def test_unknown_container_layout_is_rejected(store_tensors, tmp_path): + # Every public writer raises ValueError, before touching the filesystem. + with pytest.raises(ValueError, match="container_layout"): compress_ptwm_blob(store_tensors, container_layout="bogus") + out = tmp_path / "x.ptwm" + with pytest.raises(ValueError, match="container_layout"): + write_ptwm_shard(out, store_tensors, container_layout="bogus") + assert not out.exists() + out_dir = tmp_path / "wds" + with pytest.raises(ValueError, match="container_layout"): + write_webdataset(out_dir, store_tensors, container_layout="bogus") + assert not out_dir.exists() def test_compact_manifest_form_and_identity(store_tensors): diff --git a/tests/ptwm/stores/test_lmdb.py b/tests/ptwm/stores/test_lmdb.py index 26060a9..5d8e98c 100644 --- a/tests/ptwm/stores/test_lmdb.py +++ b/tests/ptwm/stores/test_lmdb.py @@ -7,7 +7,7 @@ import pytest import torch from ptwm import _rust -from ptwm._magic import is_container +from ptwm._magic import container_kind from ptwm.stores import LmdbWeightStore, detect_format, write_lmdb from ptwm.stores._lmdb import KEY_FMT, KEY_IDX, KEY_REG, LMDB_FORMAT @@ -26,7 +26,8 @@ def test_blob_layout_keys(tmp_path, store_tensors): # Every tensor is a self-contained single-tensor container. for name, *_ in store_tensors: blob = store._get(b"t/" + name.encode()) - assert is_container(blob) + # The default container layout is still legacy. + assert container_kind(blob) == "ptwm" def test_exploded_layout_keys(tmp_path, store_tensors): From ebe551ec1c5c17bf4e944ebf0fb4bf679f90aad0 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 02:15:42 +0000 Subject: [PATCH 22/34] feat: write the compact (PTWC) container layout by default Stage 9 of the PTWC work (synthesis section 10 step 9). Defaults: - Rust: ContainerLayout::default() is Compact, so CompressorOptions::default() writes \x89PTWC. ContainerWriter::new still writes Legacy. - ptwm-py compress_model: container_layout defaults to "compact". - Python: CompressionConfig.container_layout defaults to COMPACT; _rust.compress_model, compress_ptwm_blob, write_ptwm_shard, write_lmdb, write_webdataset and integrations.compress_safetensors_file default to "compact". LmdbWeightStore.put_tensor still uses the store's recorded layout, and "legacy" for a store that records none. - Single-tensor Compressor.compress keeps the \x89PTWS frame by default and falls back to a PTWC container; with container_layout=LEGACY the fallback is \x89PTWM. Framing.CONTAINER writes the configured layout, so CONTAINER + COMPACT (the default) gives PTWC. Documented in the Framing, ContainerLayout and CompressionConfig docstrings. CLI: `ptwm compress` and the standalone compress entry point take --container-layout {compact,legacy} (default compact). It reaches raw-file, delta, per-tensor safetensors, directory, mode a/b, LMDB and WebDataset writers. The random-access, store and inspect tools open both layouts. cli/README.md documents the flag. Tests pinned to LEGACY because they read legacy bytes (deliberate default change): compressor/container unit tests that parse legacy records or patch legacy offsets, legacy_fixtures main_options (golden legacy hash), test_extension_table.py, test_v1_header.py, test_legacy_hardening.py _legacy_blob, test_transcoder registry_hash test (legacy manifest only), the legacy container floor test in test_container_overhead.py (benchmark measure/compress_checked take a layout). Tests updated for the new default: ptwc.rs legacy_layout_stays_the_default became compact_layout_is_the_default; test_compression_config default is COMPACT; test_compact_frame container/streaming fallback, test_zstd outputs and test_lmdb blob keys now cover default, compact and legacy; decode_bomb container tests run both layouts. New: test_v1_header compact default test, test_new_writes_default_to_the_compact_layout, tests/ptwm/cli/test_container_layout_flag.py. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- benchmarks/container_overhead.py | 28 ++++- crates/ptwm-core/src/compressor.rs | 13 +- crates/ptwm-core/src/container.rs | 10 +- crates/ptwm-core/tests/decode_bomb.rs | 70 +++++------ crates/ptwm-core/tests/legacy_fixtures.rs | 7 +- crates/ptwm-core/tests/ptwc.rs | 12 +- crates/ptwm-py/src/lib.rs | 5 +- python/ptwm/_config.py | 41 ++++--- python/ptwm/_rust/__init__.py | 7 +- python/ptwm/cli/README.md | 6 + python/ptwm/cli/compress.py | 70 ++++++++++- .../integrations/_compress_safetensors.py | 7 +- python/ptwm/sharding/_writer.py | 8 +- python/ptwm/stores/_lmdb.py | 4 +- python/ptwm/stores/_webdataset.py | 4 +- .../benchmarks/test_container_overhead.py | 11 +- tests/ptwm/cli/test_container_layout_flag.py | 115 ++++++++++++++++++ tests/ptwm/core/test_compact_frame.py | 40 ++++-- tests/ptwm/core/test_container_handle.py | 4 +- tests/ptwm/core/test_zstd.py | 18 ++- tests/ptwm/stores/test_compact_layout.py | 37 +++++- tests/ptwm/stores/test_lmdb.py | 22 +++- tests/ptwm/stores/test_transcoder.py | 3 +- .../ptwm/wire_format/test_extension_table.py | 24 +++- .../ptwm/wire_format/test_legacy_hardening.py | 11 +- tests/ptwm/wire_format/test_v1_header.py | 33 ++++- 26 files changed, 497 insertions(+), 113 deletions(-) create mode 100644 tests/ptwm/cli/test_container_layout_flag.py diff --git a/benchmarks/container_overhead.py b/benchmarks/container_overhead.py index 85b1158..a91c588 100644 --- a/benchmarks/container_overhead.py +++ b/benchmarks/container_overhead.py @@ -26,7 +26,7 @@ import torch import zstandard from ptwm import _rust -from ptwm._config import CompressionConfig, Format, Framing +from ptwm._config import CompressionConfig, ContainerLayout, Format, Framing from ptwm.core import Compressor, Decompressor SIZES = (64, 256, 1024, 4096, 16384, 65536, 1 << 20) @@ -87,10 +87,20 @@ def non_payload_bytes(blob: bytes) -> int: return len(blob) - payload -def compress_checked(t: torch.Tensor, framing: Framing) -> bytes: - """Compress ``t`` and assert a bit-exact round trip with shape and dtype.""" +def compress_checked( + t: torch.Tensor, + framing: Framing, + layout: ContainerLayout = ContainerLayout.COMPACT, +) -> bytes: + """Compress ``t`` and assert a bit-exact round trip with shape and dtype. + + ``layout`` is the layout of a full container (``Framing.CONTAINER`` or the + compact-frame fallback). + """ blob = Compressor( - CompressionConfig(input_format=Format.TORCH, framing=framing) + CompressionConfig( + input_format=Format.TORCH, framing=framing, container_layout=layout + ) ).compress(t) back = Decompressor().decompress(blob) assert isinstance(back, torch.Tensor) @@ -100,9 +110,15 @@ def compress_checked(t: torch.Tensor, framing: Framing) -> bytes: return blob -def measure(t: torch.Tensor, framing: Framing, dtype: str, kind: str) -> Row: +def measure( + t: torch.Tensor, + framing: Framing, + dtype: str, + kind: str, + layout: ContainerLayout = ContainerLayout.COMPACT, +) -> Row: raw = raw_bytes(t) - blob = compress_checked(t, framing) + blob = compress_checked(t, framing, layout) z = zstandard.ZstdCompressor(level=19).compress(raw) return Row( framing=framing.name.lower(), diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 9050817..b2b540b 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -90,7 +90,9 @@ pub struct InputTensor<'a> { } /// Compressor options. `Default` gives no hint, no hash or CRC, dispatcher -/// codec selection, whole planes and the legacy container layout. +/// codec selection, whole planes and the compact container layout +/// (`\x89PTWC`, readable by v1.2 and later). Set `container_layout` to +/// [`ContainerLayout::Legacy`] for output that every release reads. #[derive(Default)] pub struct CompressorOptions { /// Header hint to readers about the dominant codec. @@ -109,7 +111,8 @@ pub struct CompressorOptions { pub allow_codec_ids: Option>, /// When `Some`, encode each plane in fixed-size chunks. pub chunk_size: Option, - /// On-disk layout of the written container. + /// On-disk layout of the written container. Defaults to + /// [`ContainerLayout::Compact`]. pub container_layout: ContainerLayout, } @@ -1397,6 +1400,8 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + // Parses legacy records, so the layout is pinned. + container_layout: ContainerLayout::Legacy, ..Default::default() }, ) @@ -1801,6 +1806,8 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + // Parses legacy records, so the layout is pinned. + container_layout: ContainerLayout::Legacy, ..Default::default() }, ) @@ -1884,6 +1891,8 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + // Parses legacy records, so the layout is pinned. + container_layout: ContainerLayout::Legacy, ..Default::default() }, ) diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 4c3bfd2..531d4b2 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -302,14 +302,18 @@ fn compact_registry_bytes( } /// On-disk layout of a container. +/// +/// `Default` gives `Compact`, the layout of new multi-tensor writes +/// ([`crate::compressor::CompressorOptions`]). [`ContainerWriter::new`] +/// still writes `Legacy`; use [`ContainerWriter::with_layout`] to choose. #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] pub enum ContainerLayout { /// `\x89PTWM`: fixed-width header, full Extension Table, hash-sorted /// index. Readable by every release. - #[default] Legacy, /// `\x89PTWC`: varint fields, built-ins stored by id, names in the index /// ([`crate::ptwc`]). Readable by v1.2 and later. + #[default] Compact, } @@ -3281,6 +3285,8 @@ mod tests { forced_codec: Some(crate::codec::CodecId::Identity), allow_codec_ids: None, chunk_size: None, + // Reads the legacy index and records, so the layout is pinned. + container_layout: ContainerLayout::Legacy, ..Default::default() }, ) @@ -3748,6 +3754,8 @@ label = "test-author" forced_codec: Some(CodecId::Identity), allow_codec_ids: None, chunk_size, + // Parses legacy records, so the layout is pinned. + container_layout: ContainerLayout::Legacy, ..Default::default() }, ) diff --git a/crates/ptwm-core/tests/decode_bomb.rs b/crates/ptwm-core/tests/decode_bomb.rs index c369d4b..8a24bca 100644 --- a/crates/ptwm-core/tests/decode_bomb.rs +++ b/crates/ptwm-core/tests/decode_bomb.rs @@ -20,7 +20,7 @@ use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; use ptwm_core::codec::{CodecId, codec_for}; use ptwm_core::compact::{compress_compact, decode_compact}; use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; -use ptwm_core::container::ContainerReader; +use ptwm_core::container::{ContainerLayout, ContainerReader}; use ptwm_core::layout::PlaneLayout; use ptwm_core::transforms::op::OpId; use ptwm_core::types::{Role, ValueFormat}; @@ -194,11 +194,15 @@ fn compact_frame_bomb_is_bounded() { assert_bounded(err, t0); } -#[test] -fn legacy_plane_bomb_is_bounded() { - let raw = raw(); +/// A full container in `layout` holding `raw` with one Zstd plane (or one +/// Zstd chunk per `chunk_size` bytes). +fn container(raw: &[u8], chunk_size: Option, layout: ContainerLayout) -> Vec { let mut buf = Vec::new(); - compress_model(Cursor::new(&mut buf), &[input(&raw)], opts(None)).unwrap(); + let opts = CompressorOptions { + container_layout: layout, + ..opts(chunk_size) + }; + compress_model(Cursor::new(&mut buf), &[input(raw)], opts).unwrap(); assert_eq!( ContainerReader::open(&buf) .unwrap() @@ -206,40 +210,38 @@ fn legacy_plane_bomb_is_bounded() { .unwrap(), raw ); - plant_bomb(&mut buf, &zstd_payload(&raw)); - let t0 = Instant::now(); - let err = ContainerReader::open(&buf) - .unwrap() - .decode_tensor("t") - .unwrap_err() - .to_string(); - assert_bounded(err, t0); + buf } #[test] -fn legacy_chunk_bomb_is_bounded() { +fn container_plane_bomb_is_bounded() { + let raw = raw(); + for layout in [ContainerLayout::Legacy, ContainerLayout::Compact] { + let mut buf = container(&raw, None, layout); + plant_bomb(&mut buf, &zstd_payload(&raw)); + let t0 = Instant::now(); + let err = ContainerReader::open(&buf) + .unwrap() + .decode_tensor("t") + .unwrap_err() + .to_string(); + assert_bounded(err, t0); + } +} + +#[test] +fn container_chunk_bomb_is_bounded() { const CHUNK: usize = 64 * 1024; let raw = raw(); - let mut buf = Vec::new(); - compress_model( - Cursor::new(&mut buf), - &[input(&raw)], - opts(Some(CHUNK as u32)), - ) - .unwrap(); - assert_eq!( - ContainerReader::open(&buf) + for layout in [ContainerLayout::Legacy, ContainerLayout::Compact] { + let mut buf = container(&raw, Some(CHUNK as u32), layout); + plant_bomb(&mut buf, &zstd_payload(&raw[..CHUNK])); + let t0 = Instant::now(); + let err = ContainerReader::open(&buf) .unwrap() .decode_tensor("t") - .unwrap(), - raw - ); - plant_bomb(&mut buf, &zstd_payload(&raw[..CHUNK])); - let t0 = Instant::now(); - let err = ContainerReader::open(&buf) - .unwrap() - .decode_tensor("t") - .unwrap_err() - .to_string(); - assert_bounded(err, t0); + .unwrap_err() + .to_string(); + assert_bounded(err, t0); + } } diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs index a886abe..35fd954 100644 --- a/crates/ptwm-core/tests/legacy_fixtures.rs +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -45,7 +45,7 @@ use std::path::PathBuf; use ptwm_core::chain::{Chain, ChainEdge, ChainNode, TerminalRef}; use ptwm_core::compressor::{CompressorOptions, InputTensor, compress_model}; -use ptwm_core::container::{ChainRegistry, ContainerReader, ContainerWriter}; +use ptwm_core::container::{ChainRegistry, ContainerLayout, ContainerReader, ContainerWriter}; use ptwm_core::extension::{BUILTIN_PUBKEY, CanonicalId}; use ptwm_core::tensor_record::{CHAIN_REF_INLINE_BIT, TensorRecord, parse_tensor_record}; use ptwm_core::transcode::{explode, implode}; @@ -449,7 +449,8 @@ fn inputs(data: &[TensorData]) -> Vec> { .collect() } -/// Options of the Python default (`emit_payload_hash`, `emit_plane_crc`). +/// Options of the Python default (`emit_payload_hash`, `emit_plane_crc`), +/// pinned to the legacy layout that these fixtures freeze. fn main_options(chunk_size: Option) -> CompressorOptions { CompressorOptions { method_hint: 1, @@ -458,7 +459,7 @@ fn main_options(chunk_size: Option) -> CompressorOptions { forced_codec: None, allow_codec_ids: None, chunk_size, - ..Default::default() + container_layout: ContainerLayout::Legacy, } } diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs index 99c02b6..1d543de 100644 --- a/crates/ptwm-core/tests/ptwc.rs +++ b/crates/ptwm-core/tests/ptwc.rs @@ -470,13 +470,19 @@ fn compress_model_matrix_round_trips_in_both_layouts() { } } +/// New multi-tensor writes default to the compact layout; the legacy +/// layout is still available on request. #[test] -fn legacy_layout_stays_the_default() { +fn compact_layout_is_the_default() { let data = feature_tensors(); let default = compress(&data, options(true, true, None, ContainerLayout::default())); + let compact = compress(&data, options(true, true, None, ContainerLayout::Compact)); + assert!(default == compact); + assert_eq!(default[..5], *b"\x89PTWC"); let legacy = compress(&data, options(true, true, None, ContainerLayout::Legacy)); - assert!(default == legacy); - assert_eq!(default[..5], *b"\x89PTWM"); + assert_eq!(legacy[..5], *b"\x89PTWM"); + let from_default = compress(&data, CompressorOptions::default()); + assert_eq!(from_default[..5], *b"\x89PTWC"); } #[test] diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 2fed6ce..86a7ede 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -317,10 +317,13 @@ fn py_dtype_element_size(dtype_code: u8) -> PyResult { /// `method_hint` selects the dominant codec for the container header /// (1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY). /// +/// `container_layout` is ``"compact"`` (default, ``\x89PTWC``, readable by +/// v1.2 and later) or ``"legacy"`` (``\x89PTWM``, readable by every release). +/// /// Returns the entire container as a ``bytes`` object. #[allow(clippy::type_complexity, clippy::too_many_arguments)] #[pyfunction] -#[pyo3(signature = (records, chains_per_tensor, emit_payload_hash, emit_plane_crc, method_hint=3, streaming_chunk=None, codec_menu=None, *, container_layout="legacy"))] +#[pyo3(signature = (records, chains_per_tensor, emit_payload_hash, emit_plane_crc, method_hint=3, streaming_chunk=None, codec_menu=None, *, container_layout="compact"))] fn compress_model<'py>( py: Python<'py>, records: Vec<( diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 51dc7d7..2d40c06 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -45,17 +45,25 @@ def _missing_(cls, value): class Framing(Enum): r"""Byte layout of a single-tensor :meth:`Compressor.compress` output. - ``COMPACT`` writes a compact single-tensor frame (magic ``\x89PTWS``) - with a fixed cost of at most about 70 bytes, and stores the raw bytes when - no codec makes them smaller. It falls back to ``CONTAINER`` when the - tensor cannot use the frame (streaming, or a plane that selects shared - codec state). ``CONTAINER`` writes a full ``.ptwm`` container in the - layout :attr:`CompressionConfig.container_layout` selects, which the - random-access, store and transcode tools open. + ``COMPACT`` (default) writes a compact single-tensor frame (magic + ``\x89PTWS``) with a fixed cost of at most about 70 bytes, and stores the + raw bytes when no codec makes them smaller. It falls back to a full + container when the tensor cannot use the frame (streaming, or a plane that + selects shared codec state). ``CONTAINER`` always writes a full ``.ptwm`` + container, which the random-access, store and transcode tools open. + + A full container (from ``CONTAINER`` or the ``COMPACT`` fallback) has the + layout :attr:`CompressionConfig.container_layout` selects: ``\x89PTWC`` + for :attr:`ContainerLayout.COMPACT` (the default) and ``\x89PTWM`` for + :attr:`ContainerLayout.LEGACY`. So ``Framing.CONTAINER`` with the default + layout gives a ``\x89PTWC`` container. Readers before v1.2 read only ``CONTAINER`` output in the :attr:`ContainerLayout.LEGACY` layout (magic ``\x89PTWM``); they reject - the ``\x89PTWS`` frame and the ``\x89PTWC`` compact container. + the ``\x89PTWS`` frame and the ``\x89PTWC`` compact container. Set + ``framing=Framing.CONTAINER`` and + ``container_layout=ContainerLayout.LEGACY`` for output that older releases + read. """ COMPACT = 1 @@ -73,10 +81,11 @@ def _missing_(cls, value): class ContainerLayout(Enum): r"""On-disk layout of a multi-tensor ``.ptwm`` container. - ``COMPACT`` (magic ``\x89PTWC``) stores built-in extensions by id, - integers as varints and tensor names only in the index; it needs v1.2 or - later to read. ``LEGACY`` (magic ``\x89PTWM``) is the fixed-width layout - that every release reads. + ``COMPACT`` (magic ``\x89PTWC``, the default for new writes) stores + built-in extensions by id, integers as varints and tensor names only in + the index; it needs v1.2 or later to read. ``LEGACY`` (magic + ``\x89PTWM``) is the fixed-width layout that every release reads; its + output is byte-identical to the v1.1 writer. Strings coerce case-insensitively: ``ContainerLayout("compact")`` and ``ContainerLayout("LEGACY")`` work, and the bindings take the lower-case @@ -157,10 +166,12 @@ class CompressionConfig: device: int | None = None framing: Framing = Framing.COMPACT """Output layout; see :class:`Framing`.""" - container_layout: ContainerLayout = ContainerLayout.LEGACY + container_layout: ContainerLayout = ContainerLayout.COMPACT """Layout of a full container written with ``Framing.CONTAINER`` (or as - the ``Framing.COMPACT`` fallback); see :class:`ContainerLayout`. A string - such as ``"compact"`` is coerced on construction.""" + the ``Framing.COMPACT`` fallback): ``COMPACT`` (default, ``\x89PTWC``) or + ``LEGACY`` (``\x89PTWM``, readable before v1.2); see + :class:`ContainerLayout`. A string such as ``"legacy"`` is coerced on + construction.""" emit_payload_hash: bool = True """Store an xxh64 of the raw bytes (8 B per tensor), checked on decode.""" emit_plane_crc: bool = True diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index 7bcf7ed..b0f7740 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -198,7 +198,7 @@ def compress_model( streaming_chunk: int | None = None, codec_menu: list[int] | None = None, *, - container_layout: ContainerLayout | str = "legacy", + container_layout: ContainerLayout | str = "compact", ) -> bytes: r"""Compress a list of tensors into a ``.ptwm`` container. @@ -213,8 +213,9 @@ def compress_model( ``codec_menu``: optional list of wire-stable ``CodecId`` integers restricting the trial-encode menu. ``None`` uses the full menu; empty intersection on a plane raises ``InvalidContainer``. - ``container_layout``: ``"legacy"`` writes ``\x89PTWM`` (readable by every - release), ``"compact"`` writes ``\x89PTWC`` (readable by v1.2 and later). + ``container_layout``: ``"compact"`` (default) writes ``\x89PTWC`` + (readable by v1.2 and later), ``"legacy"`` writes ``\x89PTWM`` (readable + by every release). A :class:`ptwm.ContainerLayout` is accepted too. Other values raise :class:`ptwm.Error` before anything is compressed (the binding itself raises ``ValueError``). diff --git a/python/ptwm/cli/README.md b/python/ptwm/cli/README.md index 94eb07d..6bd00fe 100644 --- a/python/ptwm/cli/README.md +++ b/python/ptwm/cli/README.md @@ -79,3 +79,9 @@ Both `compress` and `decompress` support these flags: - `--verification`: verify the compressed output decompresses to the exact input. - `--hf_cache`: preserve Hugging Face cache symlink layout. + +`compress` also takes `--container-layout {compact,legacy}`. New `.ptwm` +containers use the compact layout (magic `\x89PTWC`) by default; ptwm +v1.2 and later read it, including the random-access, store and inspect +tools. `--container-layout legacy` writes `\x89PTWM` containers that +earlier releases also read. diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index f8881b3..65c4dd2 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -22,6 +22,25 @@ } +_CONTAINER_LAYOUT_HELP = ( + "Layout of written .ptwm containers: 'compact' (default; magic " + "\\x89PTWC, readable by ptwm v1.2 and later) or 'legacy' (magic " + "\\x89PTWM, readable by every release)." +) + + +def _add_container_layout_argument(parser) -> None: # noqa: ANN001 + """Add ``--container-layout {compact,legacy}`` to ``parser``.""" + parser.add_argument( + "--container-layout", + type=str.lower, + choices=["compact", "legacy"], + default="compact", + dest="container_layout", + help=_CONTAINER_LAYOUT_HELP, + ) + + def _parse_codec_menu(arg: str | None) -> list[CodecId] | None: """Parse a comma-separated list of codec names into a list of CodecId. @@ -131,8 +150,13 @@ def compress_file( codec_menu: list[CodecId] | None = None, codec: str | None = None, device: int | None = None, + container_layout: str = "compact", ) -> None: - """Compress a single file.""" + """Compress a single file. + + The output is a full ``.ptwm`` container in ``container_layout`` + (``"compact"`` or ``"legacy"``). + """ from ptwm import ( # noqa: PLC0415 CompressionConfig, Compressor, @@ -164,8 +188,9 @@ def compress_file( CompressionConfig( bytearray_dtype=dtype, # `.ptwm` files stay full containers so the random-access, store - # and inspect tools can open them. + # and inspect tools can open them (they read both layouts). framing=Framing.CONTAINER, + container_layout=container_layout, is_streaming=is_streaming, streaming_chunk=chunk_size, method=Method(method), @@ -250,8 +275,12 @@ def compress_file_delta( codec_menu: list[CodecId] | None = None, codec: str | None = None, device: int | None = None, + container_layout: str = "compact", ) -> None: - """Compress a file using delta compression.""" + """Compress a file using delta compression. + + The output is a full ``.ptwm`` container in ``container_layout``. + """ from ptwm import ( # noqa: PLC0415 CompressionConfig, Compressor, @@ -291,6 +320,7 @@ def compress_file_delta( CompressionConfig( bytearray_dtype="float32" if dtype else "bfloat16", framing=Framing.CONTAINER, + container_layout=container_layout, is_streaming=is_streaming, streaming_chunk=chunk_size, delta_compressed_type="file", @@ -350,8 +380,13 @@ def compress_safetensors_file( codec_menu: list[CodecId] | None = None, codec: str | None = None, device: int | None = None, + container_layout: str = "compact", ) -> None: - """Compress a safetensors file.""" + """Compress a safetensors file. + + ``container_layout`` applies to tensors that do not fit the compact + single-tensor frame and are written as full containers. + """ import torch # noqa: PLC0415 from safetensors import safe_open # noqa: PLC0415 from safetensors.torch import save_file # noqa: PLC0415 @@ -415,6 +450,7 @@ def compress_safetensors_file( codec_menu=codec_menu, codec=codec, device=device, + container_layout=container_layout, ) ) compressor = compressor_cache[dtype_str] @@ -473,6 +509,7 @@ def compress_path( codec_menu: list[CodecId] | None = None, codec: str | None = None, device: int | None = None, + container_layout: str = "compact", ) -> None: """Compress all files with the given suffix in the specified path.""" overwrite_first = True @@ -569,6 +606,7 @@ def compress_path( codec_menu, codec, device, + container_layout, ) else: compression_func = compress_safetensors_file @@ -582,6 +620,7 @@ def compress_path( codec_menu, codec, device, + container_layout, ) failures: list[tuple[str, BaseException]] = [] @@ -701,6 +740,7 @@ def main(argv: list[str] | None = None) -> int: dest="classify_heuristic", help="Enable the built-in name-pattern heuristic classifier.", ) + _add_container_layout_argument(parser) args = parser.parse_args(argv) @@ -742,11 +782,12 @@ def main(argv: list[str] | None = None) -> int: mode=args.mode, classifier=chain, max_shard_size=_parse_size(args.max_shard_size), + container_layout=args.container_layout, ) return 0 # Fallback: raw-bytes compression via the existing helper. - compress_file(input_file=str(input_path)) + compress_file(input_file=str(input_path), container_layout=args.container_layout) return 0 @@ -961,6 +1002,7 @@ def add_compress_parser(subparsers): dest="lmdb_map_size", help="LMDB environment map size (e.g. 8GB). --format lmdb only.", ) + _add_container_layout_argument(parser) return parser @@ -992,6 +1034,7 @@ def _compress_safetensors_store(args, store_format: str) -> int: # noqa: ANN001 layout=args.lmdb_layout, map_size=map_size, method_hint=method_hint, + container_layout=_container_layout_arg(args), ) else: from ptwm.stores import write_webdataset # noqa: PLC0415 @@ -1008,11 +1051,20 @@ def _compress_safetensors_store(args, store_format: str) -> int: # noqa: ANN001 ) print(f"Compressing {input_path} → {out_path} (webdataset)") # noqa: T201 write_webdataset( - out_path, tensors, max_shard_size=max_shard, method_hint=method_hint + out_path, + tensors, + max_shard_size=max_shard, + method_hint=method_hint, + container_layout=_container_layout_arg(args), ) return 0 +def _container_layout_arg(args) -> str: # noqa: ANN001 + """Return ``args.container_layout``, or ``"compact"`` when it is unset.""" + return getattr(args, "container_layout", None) or "compact" + + def _method_hint_from_name(name: str) -> int: """Map a ``--method`` name to the wire ``method_hint`` integer.""" return { @@ -1084,6 +1136,7 @@ def _compress_safetensors_integration(args, mode: str) -> int: # noqa: ANN001 classifier=ClassifierChain([HeuristicClassifier()]), explore_options=explore_options, use_user_cache=not getattr(args, "no_cache", False), + container_layout=_container_layout_arg(args), ) output_bytes = sum(p.stat().st_size for p in out_dir.rglob("*") if p.is_file()) @@ -1162,6 +1215,7 @@ def handle_compress(args): codec_menu=codec_menu, codec=args.codec, device=args.device, + container_layout=_container_layout_arg(args), ) elif path.is_file(): if args.delta: @@ -1181,6 +1235,7 @@ def handle_compress(args): codec_menu=codec_menu, codec=args.codec, device=args.device, + container_layout=_container_layout_arg(args), ) elif path.suffix == ".safetensors" and not args.file_compression: compress_safetensors_file( @@ -1193,6 +1248,7 @@ def handle_compress(args): codec_menu=codec_menu, codec=args.codec, device=args.device, + container_layout=_container_layout_arg(args), ) else: compress_file( @@ -1210,6 +1266,7 @@ def handle_compress(args): codec_menu=codec_menu, codec=args.codec, device=args.device, + container_layout=_container_layout_arg(args), ) elif args.hf_cache and args.model: compress_path( @@ -1233,6 +1290,7 @@ def handle_compress(args): codec_menu=codec_menu, codec=args.codec, device=args.device, + container_layout=_container_layout_arg(args), ) else: print( # noqa: T201 diff --git a/python/ptwm/integrations/_compress_safetensors.py b/python/ptwm/integrations/_compress_safetensors.py index 81da681..b087f82 100644 --- a/python/ptwm/integrations/_compress_safetensors.py +++ b/python/ptwm/integrations/_compress_safetensors.py @@ -191,7 +191,7 @@ def compress_safetensors_file( use_user_cache: bool = True, fast: bool = False, fast_chunk: int = 512 * 1024, - container_layout: ContainerLayout | str = "legacy", + container_layout: ContainerLayout | str = "compact", ) -> AuditLog: r"""Compress a safetensors file into a sharded ``.ptwm`` directory. @@ -214,8 +214,9 @@ def compress_safetensors_file( Method hint forwarded to ``write_ptwm_shard`` (1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY). container_layout: - Layout of every shard container: ``"legacy"`` (``\x89PTWM``, readable - by every release) or ``"compact"`` (``\x89PTWC``, v1.2 and later). + Layout of every shard container: ``"compact"`` (default; + ``\x89PTWC``, v1.2 and later) or ``"legacy"`` (``\x89PTWM``, readable + by every release). Forwarded to ``write_ptwm_shard``; a string or :class:`ptwm.ContainerLayout`. Other values raise ``ValueError`` before any tensor is read. diff --git a/python/ptwm/sharding/_writer.py b/python/ptwm/sharding/_writer.py index fc4d40b..f152f6d 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -41,7 +41,7 @@ def compress_ptwm_blob( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, - container_layout: ContainerLayout | str = "legacy", + container_layout: ContainerLayout | str = "compact", ) -> bytes: r"""Compress ``tensors`` into a single in-memory ``.ptwm`` container blob. @@ -50,8 +50,8 @@ def compress_ptwm_blob( each blob is produced by ``Chain.to_bytes()``. When ``None``, chains are auto-derived from each tensor's dtype via the STANDARD production table. ``method_hint``: 1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY. - ``container_layout``: ``"legacy"`` (``\x89PTWM``, readable by every - release) or ``"compact"`` (``\x89PTWC``, v1.2 and later), as a string or + ``container_layout``: ``"compact"`` (default; ``\x89PTWC``, v1.2 and + later) or ``"legacy"`` (``\x89PTWM``, readable by every release), as a string or :class:`ptwm.ContainerLayout`. Other values raise :class:`ValueError` before anything is compressed. """ @@ -97,7 +97,7 @@ def write_ptwm_shard( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, - container_layout: ContainerLayout | str = "legacy", + container_layout: ContainerLayout | str = "compact", ) -> None: """Write a single ``.ptwm`` container holding ``tensors``. diff --git a/python/ptwm/stores/_lmdb.py b/python/ptwm/stores/_lmdb.py index 21509d9..b9ead75 100644 --- a/python/ptwm/stores/_lmdb.py +++ b/python/ptwm/stores/_lmdb.py @@ -77,12 +77,12 @@ def write_lmdb( map_size: int = _DEFAULT_MAP_SIZE, method_hint: int = 3, chains_per_tensor: list[list[bytes]] | None = None, - container_layout: ContainerLayout | str = "legacy", + container_layout: ContainerLayout | str = "compact", ) -> Path: r"""Write ``tensors`` into an LMDB environment at ``path``. ``layout`` is ``"blob"`` (default) or ``"exploded"``. ``container_layout`` - (``"legacy"`` or ``"compact"``) selects the container layout of the + (``"compact"``, the default, or ``"legacy"``) selects the container layout of the per-tensor blobs (blob layout) or of the registry manifest the tensors are rebuilt from (exploded layout). A ``"compact"`` store records the layout in its ``\x00fmt`` record so that later :meth:`LmdbWeightStore.put_tensor` diff --git a/python/ptwm/stores/_webdataset.py b/python/ptwm/stores/_webdataset.py index f6137dc..fe6dcf2 100644 --- a/python/ptwm/stores/_webdataset.py +++ b/python/ptwm/stores/_webdataset.py @@ -71,14 +71,14 @@ def write_webdataset( max_shard_size: int = _DEFAULT_MAX_SHARD_SIZE, method_hint: int = 3, chains_per_tensor: list[list[bytes]] | None = None, - container_layout: ContainerLayout | str = "legacy", + container_layout: ContainerLayout | str = "compact", ) -> Path: """Write ``tensors`` as a WebDataset shard set rooted at ``out_dir``. Returns the directory path. ``tensors`` is a sequence of ``(name, raw_bytes, dtype, shape)``; shards are packed by compressed size via :func:`ptwm.sharding.plan_shards` so a tensor never straddles a shard. - ``container_layout`` (``"legacy"`` or ``"compact"``) selects the + ``container_layout`` (``"compact"``, the default, or ``"legacy"``) selects the container the tensors are compressed into; the registry manifest records it, and readers rebuild each tensor in that layout. An unknown ``container_layout`` raises :class:`ValueError` before ``out_dir`` is diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py index 0f033a5..e39da58 100644 --- a/tests/ptwm/benchmarks/test_container_overhead.py +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -11,7 +11,7 @@ import pytest import torch from ptwm import CompressionConfig, Compressor, Decompressor, Format -from ptwm._config import Framing +from ptwm._config import ContainerLayout, Framing from benchmarks.container_overhead import ( DTYPES, @@ -103,8 +103,13 @@ def test_tiled_rows_ratio_is_at_least_60() -> None: @pytest.mark.parametrize(("dname", "dtype"), DTYPES) def test_container_framing_floor_stays_trimmed(dname: str, dtype: torch.dtype) -> None: # v1.1 wrote about 2.5 KB of non-payload bytes per single-tensor - # container, 2.1 KB of it unreferenced codec entries. + # container, 2.1 KB of it unreferenced codec entries. The bound is for + # the legacy layout, so the layout is pinned. row = measure( - make_tensor(dtype, "normal", 4096), Framing.CONTAINER, dname, "normal" + make_tensor(dtype, "normal", 4096), + Framing.CONTAINER, + dname, + "normal", + ContainerLayout.LEGACY, ) assert row.non_payload < 1100, row diff --git a/tests/ptwm/cli/test_container_layout_flag.py b/tests/ptwm/cli/test_container_layout_flag.py new file mode 100644 index 0000000..17f85e2 --- /dev/null +++ b/tests/ptwm/cli/test_container_layout_flag.py @@ -0,0 +1,115 @@ +"""CLI: ``--container-layout`` selects the layout of written containers. + +New ``.ptwm`` files default to the compact layout (``\\x89PTWC``), and the +random-access tools open them. ``--container-layout legacy`` keeps writing +``\\x89PTWM`` files for readers before v1.2. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import pytest +import torch +from ptwm import Decompressor +from ptwm._magic import container_kind +from ptwm.cli.compress import add_compress_parser, compress_file, handle_compress +from ptwm.cli.compress import main as compress_main +from ptwm.random_access import TensorIndex +from safetensors.torch import save_file + +KIND = {"compact": "ptwc", "legacy": "ptwm"} + + +def _parse(argv: list[str]) -> argparse.Namespace: + parser = argparse.ArgumentParser() + add_compress_parser(parser.add_subparsers(dest="command")) + return parser.parse_args(["compress", *argv]) + + +def _raw_file(tmp_path: Path) -> tuple[Path, bytes]: + data = np.random.default_rng(0).standard_normal(2048).astype(np.float32) + path = tmp_path / "weights.bin" + path.write_bytes(data.tobytes()) + return path, data.tobytes() + + +def _safetensors_file(tmp_path: Path) -> Path: + g = torch.Generator().manual_seed(0) + path = tmp_path / "model.safetensors" + save_file( + { + "a.weight": torch.randn(64, 32, generator=g).to(torch.bfloat16), + "b.weight": torch.randn(128, generator=g), + }, + str(path), + ) + return path + + +def test_parser_defaults_to_compact() -> None: + assert _parse(["x"]).container_layout == "compact" + assert _parse(["x", "--container-layout", "LEGACY"]).container_layout == "legacy" + with pytest.raises(SystemExit): + _parse(["x", "--container-layout", "bogus"]) + + +@pytest.mark.parametrize("layout", [None, "compact", "legacy"]) +def test_file_compression_writes_an_openable_container( + tmp_path: Path, layout: str | None +) -> None: + src, raw = _raw_file(tmp_path) + argv = [str(src), "--file_compression", "--force", "--dtype", "float32"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + + out = src.with_suffix(".bin.ptwm") + blob = out.read_bytes() + assert container_kind(blob) == KIND[layout or "compact"] + assert Decompressor().decompress(blob) == raw + index = TensorIndex.open(out) + assert index.layout == (layout or "compact") + assert len(index.names()) == 1 + + +@pytest.mark.parametrize("layout", ["compact", "legacy"]) +def test_compress_file_function_takes_the_layout(tmp_path: Path, layout: str) -> None: + src, raw = _raw_file(tmp_path) + compress_file(str(src), dtype="float32", force=True, container_layout=layout) + blob = src.with_suffix(".bin.ptwm").read_bytes() + assert container_kind(blob) == KIND[layout] + assert Decompressor().decompress(blob) == raw + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_standalone_main_mode_a_shards(tmp_path: Path, layout: str | None) -> None: + src = _safetensors_file(tmp_path) + out = tmp_path / "out" + argv = ["--mode", "a", "--out", str(out), str(src)] + if layout is not None: + argv += ["--container-layout", layout] + assert compress_main(argv) == 0 + shards = sorted(out.glob("*.ptwm")) + assert shards + for shard in shards: + assert container_kind(shard.read_bytes()) == KIND[layout or "compact"] + assert set(TensorIndex.open(shard).names()) <= {"a.weight", "b.weight"} + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_store_formats_take_the_layout(tmp_path: Path, layout: str | None) -> None: + pytest.importorskip("lmdb") + from ptwm.stores import LmdbWeightStore # noqa: PLC0415 + + src = _safetensors_file(tmp_path) + out = tmp_path / "m.lmdb" + argv = [str(src), "--format", "lmdb", "--out", str(out)] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + store = LmdbWeightStore.open(out) + blob = store._get(b"t/a.weight") + assert container_kind(blob) == KIND[layout or "compact"] diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index a922705..ac3ccfb 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -8,6 +8,7 @@ from ptwm import ( CompressionConfig, Compressor, + ContainerLayout, DecompressionConfig, Decompressor, Format, @@ -17,7 +18,10 @@ from ptwm._exceptions import Error, HeaderParseError COMPACT = b"\x89PTWS" -CONTAINER = b"\x89PTWM" +# A full container has the default (compact) layout unless the config pins +# ContainerLayout.LEGACY. +CONTAINER = b"\x89PTWC" +LEGACY_CONTAINER = b"\x89PTWM" def _torch_cfg(**kw) -> CompressionConfig: @@ -39,17 +43,39 @@ def test_framing_accepts_names() -> None: assert Framing("COMPACT") is Framing.COMPACT -def test_container_framing_still_decodes() -> None: +@pytest.mark.parametrize( + ("layout", "magic"), + [ + (None, CONTAINER), + (ContainerLayout.COMPACT, CONTAINER), + (ContainerLayout.LEGACY, LEGACY_CONTAINER), + ], +) +def test_container_framing_still_decodes(layout, magic: bytes) -> None: t = torch.randn(8, 32, dtype=torch.float32) - blob = Compressor(_torch_cfg(framing=Framing.CONTAINER)).compress(t) - assert blob[:5] == CONTAINER + kw = {} if layout is None else {"container_layout": layout} + blob = Compressor(_torch_cfg(framing=Framing.CONTAINER, **kw)).compress(t) + assert blob[:5] == magic assert torch.equal(Decompressor().decompress(blob), t) -def test_streaming_falls_back_to_container() -> None: +@pytest.mark.parametrize( + ("layout", "magic"), + [ + (None, CONTAINER), + (ContainerLayout.COMPACT, CONTAINER), + (ContainerLayout.LEGACY, LEGACY_CONTAINER), + ], +) +def test_streaming_falls_back_to_container(layout, magic: bytes) -> None: + """The compact frame has no chunk table: streaming falls back to a full + container, compact by default and legacy only when pinned.""" t = torch.randn(4096).to(torch.bfloat16) - blob = Compressor(_torch_cfg(is_streaming=True, streaming_chunk=1024)).compress(t) - assert blob[:5] == CONTAINER + kw = {} if layout is None else {"container_layout": layout} + blob = Compressor( + _torch_cfg(is_streaming=True, streaming_chunk=1024, **kw) + ).compress(t) + assert blob[:5] == magic back = Decompressor().decompress(blob) assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) diff --git a/tests/ptwm/core/test_container_handle.py b/tests/ptwm/core/test_container_handle.py index d3b1dda..bf6dbf6 100644 --- a/tests/ptwm/core/test_container_handle.py +++ b/tests/ptwm/core/test_container_handle.py @@ -108,8 +108,8 @@ def test_container_layout_name() -> None: def test_compression_config_coerces_container_layout() -> None: - # The default stays LEGACY until the switch to the compact layout. - assert CompressionConfig().container_layout is ContainerLayout.LEGACY + # New writes default to the compact layout. + assert CompressionConfig().container_layout is ContainerLayout.COMPACT cfg = CompressionConfig(container_layout="compact") assert cfg.container_layout is ContainerLayout.COMPACT cfg = CompressionConfig(container_layout="Legacy") diff --git a/tests/ptwm/core/test_zstd.py b/tests/ptwm/core/test_zstd.py index 3579beb..a3ad9d2 100644 --- a/tests/ptwm/core/test_zstd.py +++ b/tests/ptwm/core/test_zstd.py @@ -15,6 +15,7 @@ from ptwm import ( CompressionConfig, Compressor, + ContainerLayout, Decompressor, Format, Framing, @@ -70,14 +71,23 @@ def test_zstd_numpy_roundtrip() -> None: assert np.array_equal(arr, recovered) -def test_zstd_outputs_ptwm() -> None: - """ZSTD produces a .ptwm container with a forced Zstd plane codec.""" +@pytest.mark.parametrize( + ("layout", "magic"), + [(ContainerLayout.LEGACY, b"\x89PTWM"), (ContainerLayout.COMPACT, b"\x89PTWC")], +) +def test_zstd_outputs_ptwm(layout: ContainerLayout, magic: bytes) -> None: + """ZSTD produces a .ptwm container with a forced Zstd plane codec, in + either layout.""" arr = np.zeros(256, dtype=np.float32).tobytes() config = CompressionConfig( - method=Method.ZSTD, bytearray_dtype="float32", framing=Framing.CONTAINER + method=Method.ZSTD, + bytearray_dtype="float32", + framing=Framing.CONTAINER, + container_layout=layout, ) compressed = Compressor(config).compress(arr) - assert compressed[:5] == b"\x89PTWM" + assert compressed[:5] == magic + assert Decompressor().decompress(compressed) == arr def test_zstd_empty_input() -> None: diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py index bb3d098..8a0035e 100644 --- a/tests/ptwm/stores/test_compact_layout.py +++ b/tests/ptwm/stores/test_compact_layout.py @@ -16,6 +16,7 @@ from ptwm._exceptions import Error from ptwm._magic import MAGIC_LEN, MAGIC_PTWC, MAGIC_PTWM, container_kind from ptwm.sharding import compress_ptwm_blob, write_ptwm_shard +from ptwm.sharding._writer import _default_chain_bytes from ptwm.stores import ( LmdbWeightStore, WebDatasetWeightStore, @@ -32,6 +33,7 @@ ) from ptwm.stores._lmdb import KEY_FMT, KEY_REG from ptwm.stores._webdataset import INDEX_NAME, SHARDMETA_MEMBER +from ptwm.utils import DType MAGIC = {"legacy": MAGIC_PTWM, "compact": MAGIC_PTWC} LAYOUTS = ("legacy", "compact") @@ -92,8 +94,11 @@ def test_compact_manifest_form_and_identity(store_tensors): model = explode_model(blob) assert model.container_layout == "compact" assert manifest_container_layout(registry_json) == "compact" - legacy = compress_ptwm_blob(store_tensors, method_hint=3) + legacy = compress_ptwm_blob(store_tensors, method_hint=3, container_layout="legacy") assert explode_model(legacy).container_layout == "legacy" + # The default layout is compact. + default = compress_ptwm_blob(store_tensors, method_hint=3) + assert explode_model(default).container_layout == "compact" @pytest.mark.parametrize("method_hint", [1, 2, 3, 4, 5]) @@ -287,3 +292,33 @@ def test_stores_accept_the_layout_enum(tmp_path, store_tensors, raw_of): write_webdataset(wd, store_tensors, container_layout=ContainerLayout.COMPACT) index = json.loads((wd / INDEX_NAME).read_text()) assert index["registry"]["ptwm_format"] == "ptwm-manifest-c" + + +# ── Defaults ────────────────────────────────────────────────────────────────── + + +def test_new_writes_default_to_the_compact_layout(tmp_path, store_tensors): + """Every multi-tensor writer defaults to ``\\x89PTWC``; ``"legacy"`` is + still available on request (see the layout-parametrized tests above).""" + assert container_kind(compress_ptwm_blob(store_tensors)) == "ptwc" + out = tmp_path / "m.ptwm" + write_ptwm_shard(out, store_tensors) + assert container_kind(out.read_bytes()) == "ptwc" + + # The binding wrapper defaults to the compact layout too. + name, raw, dtype, shape = store_tensors[0] + dt = DType.from_dtype(dtype) + records = [(name, int(dt.code), 0, raw, list(shape), dt.dtype_str, None)] + chains = [_default_chain_bytes(int(dt.code), shape, len(raw))] + blob = _rust.compress_model(records, chains, True, True) + assert container_kind(blob) == "ptwc" + + path = tmp_path / "m.lmdb" + write_lmdb(path, store_tensors, layout="blob") + with LmdbWeightStore.open(path) as store: + assert json.loads(store._get(KEY_FMT))["container_layout"] == "compact" + + wd = tmp_path / "m.wds" + write_webdataset(wd, store_tensors) + index = json.loads((wd / INDEX_NAME).read_text()) + assert index["registry"]["ptwm_format"] == "ptwm-manifest-c" diff --git a/tests/ptwm/stores/test_lmdb.py b/tests/ptwm/stores/test_lmdb.py index 5d8e98c..d239c0d 100644 --- a/tests/ptwm/stores/test_lmdb.py +++ b/tests/ptwm/stores/test_lmdb.py @@ -12,9 +12,14 @@ from ptwm.stores._lmdb import KEY_FMT, KEY_IDX, KEY_REG, LMDB_FORMAT -def test_blob_layout_keys(tmp_path, store_tensors): +@pytest.mark.parametrize( + ("container_layout", "kind"), + [(None, "ptwc"), ("compact", "ptwc"), ("legacy", "ptwm")], +) +def test_blob_layout_keys(tmp_path, store_tensors, container_layout, kind): path = tmp_path / "m.lmdb" - write_lmdb(path, store_tensors, layout="blob") + kw = {} if container_layout is None else {"container_layout": container_layout} + write_lmdb(path, store_tensors, layout="blob", **kw) assert detect_format(path) == "lmdb" store = LmdbWeightStore.open(path) assert store.layout == "blob" @@ -22,12 +27,21 @@ def test_blob_layout_keys(tmp_path, store_tensors): fmt = json.loads(store._get(KEY_FMT)) assert fmt["ptwm_format"] == LMDB_FORMAT assert fmt["layout"] == "blob" + if kind == "ptwm": + # A legacy store's record is the one older releases write. + assert fmt == { + "layout": "blob", + "manifest_version": 1, + "ptwm_format": LMDB_FORMAT, + } + else: + # The default container layout is compact, and the store records it. + assert fmt["container_layout"] == "compact" assert store._get(KEY_REG) is None # blob layout has no shared registry key # Every tensor is a self-contained single-tensor container. for name, *_ in store_tensors: blob = store._get(b"t/" + name.encode()) - # The default container layout is still legacy. - assert container_kind(blob) == "ptwm" + assert container_kind(blob) == kind def test_exploded_layout_keys(tmp_path, store_tensors): diff --git a/tests/ptwm/stores/test_transcoder.py b/tests/ptwm/stores/test_transcoder.py index 8774cc3..2b99b3e 100644 --- a/tests/ptwm/stores/test_transcoder.py +++ b/tests/ptwm/stores/test_transcoder.py @@ -93,7 +93,8 @@ def test_manifests_are_canonical_jcs(store_tensors): def test_registry_hash_is_hex(store_tensors): - blob = compress_ptwm_blob(store_tensors, method_hint=3) + # Only the legacy manifest carries registry_hash_blake3. + blob = compress_ptwm_blob(store_tensors, method_hint=3, container_layout="legacy") registry_json, _tensors, _members = _rust.explode_ptwm(blob) reg = json.loads(registry_json) assert reg["ptwm_format"] == "ptwm-manifest" diff --git a/tests/ptwm/wire_format/test_extension_table.py b/tests/ptwm/wire_format/test_extension_table.py index ece0c84..306dacd 100644 --- a/tests/ptwm/wire_format/test_extension_table.py +++ b/tests/ptwm/wire_format/test_extension_table.py @@ -1,7 +1,18 @@ -"""Verify the v1 file carries an Extension Table immediately after the header.""" +"""Verify the v1 file carries an Extension Table immediately after the header. + +These tests read the fixed-width legacy header, so they pin +``ContainerLayout.LEGACY``; the default layout is compact (``\x89PTWC``). +""" import torch -from ptwm import CompressionConfig, Compressor, Format, Framing, Method +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + Format, + Framing, + Method, +) def test_extension_table_is_present_for_huffman_chain(): @@ -11,8 +22,10 @@ def test_extension_table_is_present_for_huffman_chain(): input_format=Format.TORCH, method=Method.HUFFMAN, framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY, ) ).compress(tensor) + assert blob[:5] == b"\x89PTWM" # Header layout (little-endian): # magic: [u8; 9] bytes 0.. 9 @@ -43,8 +56,13 @@ def test_extension_table_holds_only_referenced_codecs(): v1.1 layout interned all 16 built-in plane codecs, about 2 KB.""" tensor = torch.full((32,), 0.5, dtype=torch.bfloat16) blob = Compressor( - CompressionConfig(input_format=Format.TORCH, framing=Framing.CONTAINER) + CompressionConfig( + input_format=Format.TORCH, + framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY, + ) ).compress(tensor) + assert blob[:5] == b"\x89PTWM" table_off = int.from_bytes(blob[11:19], "little") count = int.from_bytes(blob[table_off : table_off + 4], "little") assert count < 16, f"extension table has {count} entries" diff --git a/tests/ptwm/wire_format/test_legacy_hardening.py b/tests/ptwm/wire_format/test_legacy_hardening.py index 659c0d0..29e4a38 100644 --- a/tests/ptwm/wire_format/test_legacy_hardening.py +++ b/tests/ptwm/wire_format/test_legacy_hardening.py @@ -6,7 +6,15 @@ import pytest import torch -from ptwm import CompressionConfig, Compressor, Format, Framing, Method, _rust +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + Format, + Framing, + Method, + _rust, +) HEADER_LEN = 91 # Header field offsets (see crates/ptwm-core/src/header.rs). @@ -21,6 +29,7 @@ def _legacy_blob() -> bytes: input_format=Format.TORCH, method=Method.HUFFMAN, framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY, ) ).compress(tensor) diff --git a/tests/ptwm/wire_format/test_v1_header.py b/tests/ptwm/wire_format/test_v1_header.py index 3a24bd0..908e2ef 100644 --- a/tests/ptwm/wire_format/test_v1_header.py +++ b/tests/ptwm/wire_format/test_v1_header.py @@ -1,13 +1,26 @@ """Round-trip the v1 wire format from Python.""" +import pytest import torch -from ptwm import CompressionConfig, Compressor, Decompressor, Format, Framing +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + Decompressor, + Format, + Framing, +) def test_compress_decompress_round_trip_v1_format(): + """The legacy layout, pinned explicitly, keeps the v1 magic.""" tensor = torch.randn(64, 64, dtype=torch.bfloat16) compressor = Compressor( - CompressionConfig(input_format=Format.TORCH, framing=Framing.CONTAINER) + CompressionConfig( + input_format=Format.TORCH, + framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY, + ) ) decompressor = Decompressor() @@ -17,3 +30,19 @@ def test_compress_decompress_round_trip_v1_format(): expected_magic = b"\x89PTWM\r\n\x1a\n" assert blob[: len(expected_magic)] == expected_magic, "missing PTWM magic" assert torch.equal(tensor, restored) + + +@pytest.mark.parametrize("layout", [None, ContainerLayout.COMPACT, "compact"]) +def test_container_framing_defaults_to_compact_layout(layout): + """``Framing.CONTAINER`` with the default (or an explicit compact) layout + writes a ``\x89PTWC`` container.""" + tensor = torch.randn(64, 64, dtype=torch.bfloat16) + kwargs = {} if layout is None else {"container_layout": layout} + blob = Compressor( + CompressionConfig( + input_format=Format.TORCH, framing=Framing.CONTAINER, **kwargs + ) + ).compress(tensor) + + assert blob[:9] == b"\x89PTWC\r\n\x1a\n" + assert torch.equal(tensor, Decompressor().decompress(blob)) From a46f1eca7cd62e7e57defe285733eed749b86e85 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 02:47:02 +0000 Subject: [PATCH 23/34] fix(S9-defaults): address review findings - Docs: legacy (\x89PTWM) output is readable by v1.1 and later, not by every release; the legacy writer is byte-identical to the current trimmed-table writer, not to the v1.1.0 writer. - CLI --delta passes the reference file's bytes instead of its path (it always failed with TypeError) and honours --dtype. - CLI --mode b --out F writes F (several shards go to a directory named after F without '.safetensors') instead of leaving a staging directory. - Directory compression runs in-process for one worker and uses spawned children for several; a forked child could deadlock on locks held by the parent's Rust or torch threads. - Tests: cross-release chain-blob test covers both layouts; new CLI tests for --delta, WebDataset, mode a/b, directory and multi-process runs in both layouts. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/compressor.rs | 2 +- crates/ptwm-core/src/container.rs | 2 +- crates/ptwm-py/src/lib.rs | 2 +- python/ptwm/_config.py | 6 +- python/ptwm/_magic.py | 2 +- python/ptwm/_rust/__init__.py | 2 +- python/ptwm/cli/README.md | 2 +- python/ptwm/cli/compress.py | 115 ++++++++++---- .../integrations/_compress_safetensors.py | 2 +- python/ptwm/sharding/_writer.py | 2 +- tests/ptwm/cli/test_container_layout_flag.py | 147 ++++++++++++++++++ .../test_cross_version_builtins.py | 13 +- 12 files changed, 254 insertions(+), 43 deletions(-) diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index b2b540b..1a2f0e0 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -92,7 +92,7 @@ pub struct InputTensor<'a> { /// Compressor options. `Default` gives no hint, no hash or CRC, dispatcher /// codec selection, whole planes and the compact container layout /// (`\x89PTWC`, readable by v1.2 and later). Set `container_layout` to -/// [`ContainerLayout::Legacy`] for output that every release reads. +/// [`ContainerLayout::Legacy`] for output that v1.1 and later read. #[derive(Default)] pub struct CompressorOptions { /// Header hint to readers about the dominant codec. diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 531d4b2..37d58d3 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -309,7 +309,7 @@ fn compact_registry_bytes( #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash)] pub enum ContainerLayout { /// `\x89PTWM`: fixed-width header, full Extension Table, hash-sorted - /// index. Readable by every release. + /// index. Readable by v1.1 and later. Legacy, /// `\x89PTWC`: varint fields, built-ins stored by id, names in the index /// ([`crate::ptwc`]). Readable by v1.2 and later. diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 86a7ede..c7fd17c 100644 --- a/crates/ptwm-py/src/lib.rs +++ b/crates/ptwm-py/src/lib.rs @@ -318,7 +318,7 @@ fn py_dtype_element_size(dtype_code: u8) -> PyResult { /// (1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY). /// /// `container_layout` is ``"compact"`` (default, ``\x89PTWC``, readable by -/// v1.2 and later) or ``"legacy"`` (``\x89PTWM``, readable by every release). +/// v1.2 and later) or ``"legacy"`` (``\x89PTWM``, readable by v1.1 and later). /// /// Returns the entire container as a ``bytes`` object. #[allow(clippy::type_complexity, clippy::too_many_arguments)] diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 2d40c06..44cc6ce 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -84,8 +84,10 @@ class ContainerLayout(Enum): ``COMPACT`` (magic ``\x89PTWC``, the default for new writes) stores built-in extensions by id, integers as varints and tensor names only in the index; it needs v1.2 or later to read. ``LEGACY`` (magic - ``\x89PTWM``) is the fixed-width layout that every release reads; its - output is byte-identical to the v1.1 writer. + ``\x89PTWM``) is the fixed-width layout that v1.1 and later read. Its + output is byte-identical to the current pre-PTWC writer, whose Extension + Table holds only the referenced codecs; the v1.1.0 writer interned all 16 + built-in codecs, so its bytes differ. Strings coerce case-insensitively: ``ContainerLayout("compact")`` and ``ContainerLayout("LEGACY")`` work, and the bindings take the lower-case diff --git a/python/ptwm/_magic.py b/python/ptwm/_magic.py index 7f7c45a..58477b1 100644 --- a/python/ptwm/_magic.py +++ b/python/ptwm/_magic.py @@ -2,7 +2,7 @@ Four formats share the ``\x89PTW`` prefix: -- ``\x89PTWM\r\n\x1a\n``: legacy multi-tensor container (every release reads it). +- ``\x89PTWM\r\n\x1a\n``: legacy multi-tensor container (v1.1 and later read it). - ``\x89PTWX\r\n\x1a\n``: legacy container with embedded WASM extensions. - ``\x89PTWC\r\n\x1a\n``: compact multi-tensor container (v1.2 and later). - ``\x89PTWS``: compact single-tensor frame written by diff --git a/python/ptwm/_rust/__init__.py b/python/ptwm/_rust/__init__.py index b0f7740..9eaec13 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -215,7 +215,7 @@ def compress_model( empty intersection on a plane raises ``InvalidContainer``. ``container_layout``: ``"compact"`` (default) writes ``\x89PTWC`` (readable by v1.2 and later), ``"legacy"`` writes ``\x89PTWM`` (readable - by every release). + by v1.1 and later). A :class:`ptwm.ContainerLayout` is accepted too. Other values raise :class:`ptwm.Error` before anything is compressed (the binding itself raises ``ValueError``). diff --git a/python/ptwm/cli/README.md b/python/ptwm/cli/README.md index 6bd00fe..1222d5a 100644 --- a/python/ptwm/cli/README.md +++ b/python/ptwm/cli/README.md @@ -84,4 +84,4 @@ Both `compress` and `decompress` support these flags: containers use the compact layout (magic `\x89PTWC`) by default; ptwm v1.2 and later read it, including the random-access, store and inspect tools. `--container-layout legacy` writes `\x89PTWM` containers that -earlier releases also read. +v1.1 and later read. diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index 65c4dd2..c661065 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -1,3 +1,5 @@ +import multiprocessing +import shutil import sys from collections import deque from concurrent.futures import ProcessPoolExecutor, as_completed @@ -25,7 +27,7 @@ _CONTAINER_LAYOUT_HELP = ( "Layout of written .ptwm containers: 'compact' (default; magic " "\\x89PTWC, readable by ptwm v1.2 and later) or 'legacy' (magic " - "\\x89PTWM, readable by every release)." + "\\x89PTWM, readable by v1.1 and later)." ) @@ -318,7 +320,7 @@ def compress_file_delta( compressor = Compressor( CompressionConfig( - bytearray_dtype="float32" if dtype else "bfloat16", + bytearray_dtype=dtype, framing=Framing.CONTAINER, container_layout=container_layout, is_streaming=is_streaming, @@ -334,21 +336,20 @@ def compress_file_delta( with full_path.open("rb") as f: file_data = f.read() - compressed_data = compressor.compress(file_data, delta_second_data=delta_file) + reference = delta_path.read_bytes() + compressed_data = compressor.compress(file_data, delta_second_data=reference) if verification: # `device` is not forwarded here either; see the note in # `compress_file`. decompressor = Decompressor( DecompressionConfig( - delta_second_data=delta_path.read_bytes(), + delta_second_data=reference, threads=threads, ) ) - with full_path.open("rb") as f: - file_data2 = f.read() decompressed_data = decompressor.decompress(compressed_data) - if decompressed_data != file_data2: + if decompressed_data != file_data: msg = "Decompressed file should be equal to original file." raise RuntimeError(msg) @@ -624,28 +625,43 @@ def compress_path( ) failures: list[tuple[str, BaseException]] = [] - with ProcessPoolExecutor(max_workers=max_processes) as executor: - future_to_file = { - executor.submit(compression_func, file, *comp_args): file - for file in file_list[:max_processes] - } - remaining_files = deque(file_list[max_processes:]) - while future_to_file: - for future in as_completed(future_to_file): - file = future_to_file.pop(future) - try: - future.result() - except Exception as exc: # noqa: BLE001 - failures.append((file, exc)) - print( - f"{RED}error compressing {file}: {exc!r}{RESET}", - file=sys.stderr, - ) - if remaining_files: - next_file = remaining_files.popleft() - future_to_file[ - executor.submit(compression_func, next_file, *comp_args) - ] = next_file + + def record(file: str, exc: BaseException) -> None: + failures.append((file, exc)) + print(f"{RED}error compressing {file}: {exc!r}{RESET}", file=sys.stderr) + + if max_processes <= 1: + # One worker: compress in this process. A forked pool child can + # deadlock on locks that the parent's Rust or torch threads held at + # fork time. + for file in file_list: + try: + compression_func(file, *comp_args) + except Exception as exc: # noqa: BLE001 + record(file, exc) + else: + # Several workers run in "spawn" children for the same reason. + with ProcessPoolExecutor( + max_workers=max_processes, + mp_context=multiprocessing.get_context("spawn"), + ) as executor: + future_to_file = { + executor.submit(compression_func, file, *comp_args): file + for file in file_list[:max_processes] + } + remaining_files = deque(file_list[max_processes:]) + while future_to_file: + for future in as_completed(list(future_to_file)): + file = future_to_file.pop(future) + try: + future.result() + except Exception as exc: # noqa: BLE001 + record(file, exc) + if remaining_files: + next_file = remaining_files.popleft() + future_to_file[ + executor.submit(compression_func, next_file, *comp_args) + ] = next_file if failures: print( f"{RED}{len(failures)} of {len(file_list)} file(s) failed to compress{RESET}", @@ -950,7 +966,10 @@ def add_compress_parser(subparsers): help=( "Output path for .safetensors inputs. With --mode a, the output " "is a directory (default .ptwm/). With --mode b, the " - "output is a single .safetensors-shaped file." + "output is a single .safetensors-shaped file (default " + ".ptwm.safetensors); output that needs several " + "shards goes to a directory named after it without the " + "'.safetensors' suffix." ), ) parser.add_argument( @@ -1109,8 +1128,9 @@ def _compress_safetensors_integration(args, mode: str) -> int: # noqa: ANN001 out_path = ( Path(args.out) if args.out else input_path.with_suffix(".ptwm.safetensors") ) + # The shards are written to a staging directory first and moved + # into place afterwards (see `_finish_mode_b`). out_dir = out_path.parent / (out_path.stem + "__staging") - out_dir.mkdir(parents=True, exist_ok=True) is_existing_non_empty = out_path.exists() if not args.force and is_existing_non_empty: @@ -1120,6 +1140,11 @@ def _compress_safetensors_integration(args, mode: str) -> int: # noqa: ANN001 if user_input not in ("yes", "y"): return 0 + if mode == "b": + if out_dir.exists(): + shutil.rmtree(out_dir) + out_dir.mkdir(parents=True) + input_bytes = input_path.stat().st_size if explore_options is not None: print( # noqa: T201 @@ -1140,11 +1165,39 @@ def _compress_safetensors_integration(args, mode: str) -> int: # noqa: ANN001 ) output_bytes = sum(p.stat().st_size for p in out_dir.rglob("*") if p.is_file()) + if mode == "b": + final = _finish_mode_b(out_dir, out_path) + if final != out_path: + print( # noqa: T201 + f"The output needs more than one shard; wrote {final}/" + ) if explore_options is not None: _print_explore_summary(audit, input_bytes, output_bytes) return 0 +def _finish_mode_b(staging: Path, out_path: Path) -> Path: + """Move mode-b shards from ``staging`` to their final place. + + A single shard becomes the file ``out_path``; its index is dropped, since + a single ``.safetensors`` file needs none. Several shards and their index + move to the directory ``out_path`` without its last suffix (for example + ``model.ptwm.safetensors`` -> ``model.ptwm/``). Returns the final path. + """ + shards = sorted(staging.glob("*.safetensors")) + if len(shards) == 1: + if out_path.is_dir(): + shutil.rmtree(out_path) + shards[0].replace(out_path) + shutil.rmtree(staging) + return out_path + final_dir = out_path.with_suffix("") + if final_dir.exists(): + shutil.rmtree(final_dir) + staging.rename(final_dir) + return final_dir + + def _resolve_compress_mode(args) -> str | None: # noqa: ANN001 """Pick the effective compress mode for a .safetensors input. diff --git a/python/ptwm/integrations/_compress_safetensors.py b/python/ptwm/integrations/_compress_safetensors.py index b087f82..ed6fee6 100644 --- a/python/ptwm/integrations/_compress_safetensors.py +++ b/python/ptwm/integrations/_compress_safetensors.py @@ -216,7 +216,7 @@ def compress_safetensors_file( container_layout: Layout of every shard container: ``"compact"`` (default; ``\x89PTWC``, v1.2 and later) or ``"legacy"`` (``\x89PTWM``, readable - by every release). + by v1.1 and later). Forwarded to ``write_ptwm_shard``; a string or :class:`ptwm.ContainerLayout`. Other values raise ``ValueError`` before any tensor is read. diff --git a/python/ptwm/sharding/_writer.py b/python/ptwm/sharding/_writer.py index f152f6d..b827274 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -51,7 +51,7 @@ def compress_ptwm_blob( auto-derived from each tensor's dtype via the STANDARD production table. ``method_hint``: 1=HUFFMAN, 2=ZSTD, 3=MICROSCALE, 4=RANS, 5=IDENTITY. ``container_layout``: ``"compact"`` (default; ``\x89PTWC``, v1.2 and - later) or ``"legacy"`` (``\x89PTWM``, readable by every release), as a string or + later) or ``"legacy"`` (``\x89PTWM``, readable by v1.1 and later), as a string or :class:`ptwm.ContainerLayout`. Other values raise :class:`ValueError` before anything is compressed. """ diff --git a/tests/ptwm/cli/test_container_layout_flag.py b/tests/ptwm/cli/test_container_layout_flag.py index 17f85e2..2836d0c 100644 --- a/tests/ptwm/cli/test_container_layout_flag.py +++ b/tests/ptwm/cli/test_container_layout_flag.py @@ -113,3 +113,150 @@ def test_store_formats_take_the_layout(tmp_path: Path, layout: str | None) -> No store = LmdbWeightStore.open(out) blob = store._get(b"t/a.weight") assert container_kind(blob) == KIND[layout or "compact"] + + +def _assert_tensors_match(src: Path, out: Path) -> None: + from ptwm.integrations import SafeOpen # noqa: PLC0415 + from safetensors.torch import load_file # noqa: PLC0415 + + want = load_file(str(src)) + got = SafeOpen(str(out), "pt") + for name, tensor in want.items(): + assert torch.equal(got.get_tensor(name), tensor) + + +@pytest.mark.parametrize("layout", [None, "compact", "legacy"]) +def test_delta_compression_round_trips(tmp_path: Path, layout: str | None) -> None: + """``--delta`` reads the reference file; it used to pass the reference + path string to the compressor and fail with ``TypeError``.""" + from ptwm import DecompressionConfig # noqa: PLC0415 + + rng = np.random.default_rng(1) + ref = rng.standard_normal(2048).astype(np.float32) + cur = ref + rng.standard_normal(2048).astype(np.float32) * 1e-3 + src = tmp_path / "weights.bin" + src.write_bytes(cur.tobytes()) + ref_path = tmp_path / "base.bin" + ref_path.write_bytes(ref.tobytes()) + + argv = [str(src), "--delta", str(ref_path), "--force", "--dtype", "float32"] + argv += ["--verification"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + + out = tmp_path / "weights_delta_base.bin.ptwm" + blob = out.read_bytes() + assert container_kind(blob) == KIND[layout or "compact"] + dec = Decompressor(DecompressionConfig(delta_second_data=ref_path.read_bytes())) + assert dec.decompress(blob) == cur.tobytes() + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_webdataset_store_takes_the_layout(tmp_path: Path, layout: str | None) -> None: + from ptwm.stores import WebDatasetWeightStore # noqa: PLC0415 + + src = _safetensors_file(tmp_path) + out = tmp_path / "m.wds" + argv = [str(src), "--format", "webdataset", "--out", str(out)] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + from ptwm.stores._common import manifest_container_layout # noqa: PLC0415 + + with WebDatasetWeightStore.open(out) as store: + assert manifest_container_layout(store._registry_json) == (layout or "compact") + want = torch.randn(64, 32, generator=torch.Generator().manual_seed(0)) + assert torch.equal(store.get_tensor("a.weight"), want.to(torch.bfloat16)) + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_mode_b_writes_the_out_file(tmp_path: Path, layout: str | None) -> None: + """``--mode b --out F`` writes F itself and leaves no staging directory.""" + from safetensors import safe_open # noqa: PLC0415 + + src = _safetensors_file(tmp_path) + out = tmp_path / "b.ptwm.safetensors" + argv = [str(src), "--mode", "b", "--out", str(out), "--force"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + + assert out.is_file() + assert not (tmp_path / "b.ptwm__staging").exists() + with safe_open(str(out), "pt") as f: + payload = f.get_tensor("__ptwm_payload__").numpy().tobytes() + assert container_kind(payload) == KIND[layout or "compact"] + _assert_tensors_match(src, out) + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_mode_a_writes_the_out_directory(tmp_path: Path, layout: str | None) -> None: + src = _safetensors_file(tmp_path) + out = tmp_path / "a_out" + argv = [str(src), "--mode", "a", "--out", str(out), "--force"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + shard = out / "model.ptwm" + assert container_kind(shard.read_bytes()) == KIND[layout or "compact"] + _assert_tensors_match(src, shard) + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_directory_file_compression_takes_the_layout( + tmp_path: Path, layout: str | None +) -> None: + src = _safetensors_file(tmp_path) + argv = [str(tmp_path), "--file_compression", "--force"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + blob = (tmp_path / (src.name + ".ptwm")).read_bytes() + assert container_kind(blob) == KIND[layout or "compact"] + assert Decompressor().decompress(blob) == src.read_bytes() + + +@pytest.mark.parametrize("layout", [None, "legacy"]) +def test_directory_per_tensor_shell_takes_the_layout( + tmp_path: Path, layout: str | None +) -> None: + src = _safetensors_file(tmp_path) + argv = [str(tmp_path), "--force"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + out = tmp_path / "model.ptwm.safetensors" + + from ptwm.integrations._safetensors import ( # noqa: PLC0415 + decompress_safetensors_tensor, + ) + from safetensors import safe_open # noqa: PLC0415 + from safetensors.torch import load_file # noqa: PLC0415 + + want = load_file(str(src)) + with safe_open(str(out), "pt") as f: + assert set(f.keys()) == set(want) + for name, tensor in want.items(): + got = decompress_safetensors_tensor(f.get_tensor(name)) + raw = tensor.contiguous().view(torch.uint8).numpy().tobytes() + if isinstance(got, torch.Tensor): + got = got.contiguous().view(torch.uint8).numpy().tobytes() + assert bytes(got) == raw + + +def test_directory_compression_with_several_processes(tmp_path: Path) -> None: + """``--max_processes 2`` runs the workers in spawned children, so a + parent that already ran the Rust and torch thread pools does not hang.""" + torch.randn(256, 256) @ torch.randn(256, 256) + srcs = [] + for i in range(3): + sub = _safetensors_file(tmp_path) + srcs.append(sub.rename(tmp_path / f"m{i}.safetensors")) + argv = [str(tmp_path), "--file_compression", "--force", "--max_processes", "2"] + argv += ["--container-layout", "legacy"] + handle_compress(_parse(argv)) + for src in srcs: + blob = (tmp_path / (src.name + ".ptwm")).read_bytes() + assert container_kind(blob) == "ptwm" + assert Decompressor().decompress(blob) == src.read_bytes() diff --git a/tests/ptwm/wire_format/test_cross_version_builtins.py b/tests/ptwm/wire_format/test_cross_version_builtins.py index 3aae3e9..fdb907a 100644 --- a/tests/ptwm/wire_format/test_cross_version_builtins.py +++ b/tests/ptwm/wire_format/test_cross_version_builtins.py @@ -66,9 +66,16 @@ def test_chain_blob_info_maps_old_builtin_ids() -> None: assert chain_blob_info(old) == chain_blob_info(blob) -def test_compress_model_accepts_chain_blobs_of_other_releases() -> None: +@pytest.mark.parametrize( + ("layout", "magic"), [("legacy", b"\x89PTWM"), ("compact", b"\x89PTWC")] +) +def test_compress_model_accepts_chain_blobs_of_other_releases( + layout: str, magic: bytes +) -> None: """A chain blob cached by another release compresses to the same - container as the current blob.""" + container as the current blob. The legacy layout writes full canonical + ids and labels into the Extension Table, so it is the case that needs the + ids normalized; the compact layout stores built-ins by wire id.""" import torch shape = [8, 16] @@ -82,9 +89,11 @@ def compress(chain: bytes) -> bytes: chains_per_tensor=[[chain]], emit_payload_hash=True, emit_plane_crc=True, + container_layout=layout, ) current = compress(blob) + assert current.startswith(magic) assert compress(old) == current assert _rust.decode_tensor(current, "w") == raw From 1f83fda69e9f6056fa8980e9668a0bcc9d0be2c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 03:02:02 +0000 Subject: [PATCH 24/34] docs: rewrite the wire-format docs and fix the overhead benchmark Stage 10 of the issue #5 work (synthesis section 3 E13/E15, section 10 step 10). Docs: - docs/architecture/wire-format.md rewritten against the code: correct EOF sentinel bytes (89 54 57 4D 45 4F 46 1A), chain registry inside the shared-prelude section, the magic rule runs old-reader-rejects-new, real legacy plane record fields (no per-plane shape or hash; codec is dispatched by codec_table_idx), no payload alignment. New sections: conventions, versioning (built-in semantics keyed by wire id; a semantic change needs a new OpId/CodecId), legacy Extension Table, built-in ids, prelude, registry and chain wire, tensor record, index, what the legacy reader checks, PTWC (header, table, registry, record, plane, index, footer, measured sizes), CW2, PTWS, compatibility and reader limits. - Site: concepts/5.container.md (both layouts, compatibility, correct hash and error claims), guides/02.python-api.md (container_layout knob, pre-v1.2 readers), guides/01.cli.md (--container-layout), new 70.release-notes.md: default output needs >= v1.2 readers, how to write v1.1-readable files and the built-in id caveat. GLOSSARY and the architecture README updated. Benchmark (E15): - container_overhead counts plane payloads, inline state and shared prelude state as payload for every format (via explode_ptwm), reports ptws, ptwc and ptwm rows, labels rows by the magic written, and adds a multi-tensor per-tensor-overhead sweep over both layouts. - Tests: codec data counted alike in all three formats, format labels, single-tensor PTWC bound (<= 160 B), multi-tensor per-tensor bounds (<= 95/100/70 B and under half the legacy cost). - benchmarks/README.md updated. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- GLOSSARY.md | 2 + benchmarks/README.md | 2 +- benchmarks/container_overhead.py | 207 +++- crates/ptwm-core/src/chain/wire_compact.rs | 2 +- docs/architecture/README.md | 7 +- docs/architecture/wire-format.md | 999 ++++++++++++------ docs/site/content/20.concepts/5.container.md | 128 ++- docs/site/content/30.guides/01.cli.md | 12 + docs/site/content/30.guides/02.python-api.md | 25 +- docs/site/content/70.release-notes.md | 85 ++ .../benchmarks/test_container_overhead.py | 89 +- 11 files changed, 1178 insertions(+), 380 deletions(-) create mode 100644 docs/site/content/70.release-notes.md diff --git a/GLOSSARY.md b/GLOSSARY.md index 890dadc..22f8a8e 100644 --- a/GLOSSARY.md +++ b/GLOSSARY.md @@ -38,6 +38,8 @@ | Term | Definition | Aliases to avoid | | --- | --- | --- | | **`.ptwm` container** | The multi-tensor on-disk format: shared-state prelude, tensor index, and per-plane codec dispatch | ptwm file, container format | +| **Container layout** | One of the two byte layouts of a `.ptwm` container: compact (`\x89PTWC`, the default) or legacy (`\x89PTWM`, readable before v1.2). Selected with `container_layout` | format version | +| **Compact frame** | The single-tensor format (`\x89PTWS`) that `Compressor.compress` writes by default | PTWS container | | **Header** | The wire-format structure describing a `.ptwm` container's layout | wire header | | **Format** | The input/output data kind a tensor is read from or written to: `BYTE`, `TORCH`, `NUMPY`, or `FILE` | data format (only when meaning this specific enum) | | **Mode `a` / Mode `b`** | For safetensors input: `a` writes a native `.ptwm` directory; `b` writes a `.safetensors` shell wrapping a `.ptwm` blob | native mode / shell mode | diff --git a/benchmarks/README.md b/benchmarks/README.md index d69c99d..07d1e42 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -49,7 +49,7 @@ BF16 tensor ≥ 50 MB; the full-model scripts use BF16 tensors up to a 3 GB cap. | Module | Needs a shard? | What it measures | | --- | --- | --- | -| `container_overhead` | no (synthetic) | Fixed per-output bytes of single-tensor `Compressor.compress` output, compact frame vs full container, bf16/fp32, constant/normal, 64 B to 1 MB, against zstd-19. `tests/ptwm/benchmarks/test_container_overhead.py` enforces the size limits in CI. | +| `container_overhead` | no (synthetic) | Non-payload bytes of PTWM output. Single tensor: `Compressor.compress` as compact frame (`ptws`), compact container (`ptwc`) and legacy container (`ptwm`), bf16/fp32, constant/normal, 64 B to 1 MB, against zstd-19. Multi-tensor: bytes per tensor of `compress_ptwm_blob` output in both container layouts for 1 to 512 small tensors, with and without hash and CRC. Coded data (plane payloads, inline and shared codec state) counts as payload in every format. `tests/ptwm/benchmarks/test_container_overhead.py` enforces the size limits in CI. | | `validate_dtype` | no (synthetic) | Losslessness + ratio across BF16/FP16/FP32/FP8 through the production streaming path. | | `validate_scale {small,large,edge}` | `large` only | `small`: many tiny tensors; `large`: a 72B-class shard (scaling + RSS); `edge`: single-tensor files, non-float dtypes, degenerate shapes. Every mode verifies a bit-exact roundtrip. | | `validate_realmodel` | yes (real shards) | Real-model validation matrix: ratio / encode / decode / losslessness across Mistral-7B (BF16), gpt-oss-20b (MXFP4), Nemotron (FP8), Qwen2.5-72B (BF16). | diff --git a/benchmarks/container_overhead.py b/benchmarks/container_overhead.py index a91c588..b458a4b 100644 --- a/benchmarks/container_overhead.py +++ b/benchmarks/container_overhead.py @@ -1,38 +1,91 @@ # ruff: noqa: T201 -"""Fixed per-output overhead of single-tensor ``Compressor.compress`` output. +r"""Fixed overhead of PTWM output: single-tensor and multi-tensor. -Sweeps bf16 and fp32 tensors from 64 B to 1 MB, constant and random-normal, -and reports for each framing (compact frame and full container): output -size, non-payload bytes, ratio (raw / output), and zstd-19 on the same raw -bytes. Every output is decoded and checked bit-exact, including shape and -dtype. Needs no shards. +Single tensor: sweeps bf16 and fp32 tensors from 64 B to 1 MB, constant and +random-normal, through ``Compressor.compress`` in three output formats: + +- ``ptws``: the default compact frame (``Framing.COMPACT``); +- ``ptwc``: a full container in the compact layout (``Framing.CONTAINER``, + ``ContainerLayout.COMPACT``, the default layout); +- ``ptwm``: a full container in the legacy layout (``Framing.CONTAINER``, + ``ContainerLayout.LEGACY``), the format readers before v1.2 need. + +For each it reports output size, non-payload bytes, ratio (raw / output) and +zstd-19 on the same raw bytes. Every output is decoded and checked bit-exact, +including shape and dtype. + +Multi-tensor: writes ``n`` small tensors into one container with +``compress_ptwm_blob`` in both layouts and reports the non-payload bytes per +tensor (index entry, tensor record, plane records, and the file's fixed +sections divided by ``n``). Every tensor is decoded and checked. + +Needs no shards:: python -m benchmarks.container_overhead -"Non-payload" is the output size minus the codec payload. For the compact -frame the payload includes inline codec state; for the container it is the -sum of the plane payload lengths. A stored compact frame counts the raw bytes -as payload. +"Non-payload" is the output size minus the data a codec produced: plane +payloads, inline codec state and shared (prelude) codec state. The same rule +applies to every format. A stored compact frame counts the raw bytes as +payload. -``tests/ptwm/benchmarks/test_container_overhead.py`` runs the same sweep in +``tests/ptwm/benchmarks/test_container_overhead.py`` runs the same sweeps in CI and asserts the size limits, so a regression in the fixed overhead fails the build. """ from __future__ import annotations +import base64 +import json +from collections.abc import Callable from dataclasses import dataclass +import numpy as np import torch import zstandard from ptwm import _rust from ptwm._config import CompressionConfig, ContainerLayout, Format, Framing +from ptwm._magic import container_kind from ptwm.core import Compressor, Decompressor +from ptwm.sharding._writer import compress_ptwm_blob SIZES = (64, 256, 1024, 4096, 16384, 65536, 1 << 20) DTYPES = (("bf16", torch.bfloat16), ("fp32", torch.float32)) KINDS = ("constant", "normal") +# (framing, container layout) for each single-tensor output format, in report +# order. The layout applies to a full container: ``Framing.CONTAINER``, or the +# container the compact framing falls back to. +FORMATS: tuple[tuple[Framing, ContainerLayout], ...] = ( + (Framing.COMPACT, ContainerLayout.COMPACT), + (Framing.CONTAINER, ContainerLayout.COMPACT), + (Framing.CONTAINER, ContainerLayout.LEGACY), +) + +ShapeFn = Callable[[int], tuple[int, ...]] +"""Shape of tensor ``i`` in a multi-tensor case.""" + +# Tensor counts of the multi-tensor sweep. +MULTI_COUNTS = (1, 8, 64, 512) + + +def same_shape(_i: int) -> tuple[int, ...]: + """Shape of tensor ``i`` when every tensor has the same shape.""" + return (64,) + + +def distinct_shapes(i: int) -> tuple[int, ...]: + """Shape of tensor ``i`` when every tensor has its own shape.""" + return (64 + i,) + + +# Multi-tensor cases: (label, dtype, shape of tensor i). +MULTI_CASES = ( + ("fp32 same", "float32", same_shape), + ("fp32 distinct", "float32", distinct_shapes), + ("bf16 same", "bfloat16", same_shape), +) + @dataclass(frozen=True) class Row: @@ -57,6 +110,20 @@ def payload(self) -> int: return self.out - self.non_payload +@dataclass(frozen=True) +class MultiRow: + layout: str + case: str + hashes: bool + n: int + out: int + non_payload: int + + @property + def per_tensor(self) -> float: + return self.non_payload / self.n + + def make_tensor( dtype: torch.dtype, kind: str, nbytes: int, seed: int = 1 ) -> torch.Tensor: @@ -79,12 +146,33 @@ def raw_bytes(t: torch.Tensor) -> bytes: return t.contiguous().view(torch.uint8).numpy().tobytes() +def codec_data_bytes(blob: bytes) -> int: + """Bytes of ``blob`` that a codec produced (see the module docstring). + + For a container: every plane payload, every inline codec state and every + shared prelude state, read from the transcoder's manifest projection. + """ + if container_kind(blob) == "ptws": + # Payload plus inline state; the raw size for a stored frame. + return _rust.decode_compact(blob)[6] + registry_json, tensors, members = _rust.explode_ptwm(blob) + total = sum( + len(members[s["payload"]]) + for s in json.loads(registry_json).get("shared_state", []) + ) + for _name, _key, tensor_json in tensors: + for plane in json.loads(tensor_json)["planes"]: + total += plane["comp_len"] + if plane.get("state_b64"): + total += len(base64.b64decode(plane["state_b64"])) + elif plane.get("state_payload"): + total += len(members[plane["state_payload"]]) + return total + + def non_payload_bytes(blob: bytes) -> int: - """Output size minus codec payload (see the module docstring).""" - if blob[:5] == b"\x89PTWS": - return len(blob) - _rust.decode_compact(blob)[6] - payload = sum(s[5] for s in _rust.list_plane_summaries(blob)) - return len(blob) - payload + """Output size minus codec data (see the module docstring).""" + return len(blob) - codec_data_bytes(blob) def compress_checked( @@ -117,11 +205,12 @@ def measure( kind: str, layout: ContainerLayout = ContainerLayout.COMPACT, ) -> Row: + """Compress ``t`` once; ``Row.framing`` names the format written.""" raw = raw_bytes(t) blob = compress_checked(t, framing, layout) z = zstandard.ZstdCompressor(level=19).compress(raw) return Row( - framing=framing.name.lower(), + framing=container_kind(blob) or "?", dtype=dtype, kind=kind, raw=len(raw), @@ -132,42 +221,108 @@ def measure( def sweep( - framings: tuple[Framing, ...] = (Framing.COMPACT, Framing.CONTAINER), + formats: tuple[tuple[Framing, ContainerLayout], ...] = FORMATS, ) -> list[Row]: rows = [] - for framing in framings: + for framing, layout in formats: for dname, dt in DTYPES: for nbytes in SIZES: for kind in KINDS: t = make_tensor(dt, kind, nbytes) - rows.append(measure(t, framing, dname, kind)) + rows.append(measure(t, framing, dname, kind, layout)) + return rows + + +def multi_tensors( + n: int, dtype: str, shape_of: ShapeFn, seed: int = 7 +) -> list[tuple[str, bytes, str, tuple[int, ...]]]: + """``n`` random-normal tensors named like transformer bias vectors.""" + rng = np.random.default_rng(seed) + out = [] + for i in range(n): + shape = shape_of(i) + values = torch.from_numpy(rng.standard_normal(shape, dtype=np.float32)) + t = values.to(torch.bfloat16 if dtype == "bfloat16" else torch.float32) + out.append((f"model.layers.{i:03d}.attn.q.bias", raw_bytes(t), dtype, shape)) + return out + + +def measure_multi( + case: str, + dtype: str, + shape_of: ShapeFn, + n: int, + layout: ContainerLayout, + hashes: bool = True, +) -> MultiRow: + """Write ``n`` tensors into one container and check every tensor.""" + tensors = multi_tensors(n, dtype, shape_of) + blob = compress_ptwm_blob( + tensors, + emit_payload_hash=hashes, + emit_plane_crc=hashes, + container_layout=layout, + ) + handle = _rust.ContainerHandle(blob) + assert handle.names() == [name for name, *_ in tensors] + for name, raw, _dtype, _shape in tensors: + assert handle.decode(name) == raw, name + return MultiRow( + layout=container_kind(blob) or "?", + case=case, + hashes=hashes, + n=n, + out=len(blob), + non_payload=non_payload_bytes(blob), + ) + + +def multi_sweep(counts: tuple[int, ...] = MULTI_COUNTS) -> list[MultiRow]: + rows = [] + for case, dtype, shape_of in MULTI_CASES: + for hashes in (True, False): + for layout in (ContainerLayout.COMPACT, ContainerLayout.LEGACY): + for n in counts: + rows.append(measure_multi(case, dtype, shape_of, n, layout, hashes)) return rows def main() -> None: header = ( - f"{'framing':<10} {'dtype':<5} {'data':<9} {'raw B':>9} {'out B':>9} " + f"{'format':<7} {'dtype':<5} {'data':<9} {'raw B':>9} {'out B':>9} " f"{'non-pay B':>9} {'ratio':>8} {'zstd19 B':>9} {'zstd ratio':>10}" ) print(header) print("-" * len(header)) for r in sweep(): print( - f"{r.framing:<10} {r.dtype:<5} {r.kind:<9} {r.raw:>9,} {r.out:>9,} " + f"{r.framing:<7} {r.dtype:<5} {r.kind:<9} {r.raw:>9,} {r.out:>9,} " f"{r.non_payload:>9,} {r.ratio:>8.3f} {r.zstd19:>9,} {r.zstd_ratio:>10.2f}" ) print( - f"\n{'framing':<10} {'rows':>5} {'raw B':>9} {'out B':>7} {'ratio':>7} {'zstd19 B':>9}" + f"\n{'format':<7} {'rows':>5} {'raw B':>9} {'out B':>7} {'ratio':>7} {'zstd19 B':>9}" ) - for framing in (Framing.COMPACT, Framing.CONTAINER): + for framing, layout in FORMATS: for rows in (1, 8, 64): - r = measure(tiled_rows(rows), framing, "bf16", "tiled") + r = measure(tiled_rows(rows), framing, "bf16", "tiled", layout) print( - f"{r.framing:<10} {rows:>5} {r.raw:>9,} {r.out:>7,} " + f"{r.framing:<7} {rows:>5} {r.raw:>9,} {r.out:>7,} " f"{r.ratio:>7.1f} {r.zstd19:>9,}" ) + header = ( + f"\n{'layout':<6} {'case':<14} {'hash+crc':<8} {'tensors':>7} " + f"{'out B':>9} {'non-pay B':>9} {'B/tensor':>8}" + ) + print(header) + print("-" * (len(header) - 1)) + for m in multi_sweep(): + print( + f"{m.layout:<6} {m.case:<14} {'on' if m.hashes else 'off':<8} {m.n:>7} " + f"{m.out:>9,} {m.non_payload:>9,} {m.per_tensor:>8.1f}" + ) + if __name__ == "__main__": main() diff --git a/crates/ptwm-core/src/chain/wire_compact.rs b/crates/ptwm-core/src/chain/wire_compact.rs index 27d5dcd..f98cdc5 100644 --- a/crates/ptwm-core/src/chain/wire_compact.rs +++ b/crates/ptwm-core/src/chain/wire_compact.rs @@ -3,7 +3,7 @@ //! [`crate::chain::wire`] stores a chain with fixed-width fields and section //! lengths, which costs 50-70 bytes for a typical production chain. CW2 //! stores the same [`Chain`] in about 15-20 bytes. The `\x89PTWS` compact -//! frame uses it, and the compact container will too. The legacy wire stays +//! frame and the `\x89PTWC` compact container use it. The legacy wire stays //! unchanged for `\x89PTWM` containers and Python chain blobs. //! //! ```text diff --git a/docs/architecture/README.md b/docs/architecture/README.md index 6242b48..5051d6e 100644 --- a/docs/architecture/README.md +++ b/docs/architecture/README.md @@ -13,9 +13,10 @@ who want to extend, refactor, or audit it. A first-time reader can take these in order. Each builds on the previous one. -1. `wire-format.md` covers the on-disk container, since the wire - format is the most stable surface in the codebase and every other - subsystem either writes to it or reads from it. +1. `wire-format.md` covers the on-disk formats (the legacy and compact + containers and the single-tensor frame), since the wire format is + the most stable surface in the codebase and every other subsystem + either writes to it or reads from it. 2. `ppg.md` describes the preprocessing graph, which decides what bytes the entropy coder ever sees. The compression-ratio story lives almost entirely in this stage. diff --git a/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index af4b709..55db4cb 100644 --- a/docs/architecture/wire-format.md +++ b/docs/architecture/wire-format.md @@ -1,10 +1,25 @@ -# The `.ptwm` wire format +# The `.ptwm` wire formats -The on-disk container is the most stable surface in PTWM. Every -change to it must be backwards compatible with every file that an -older PTWM version may have produced. This document explains the -overall shape and the reasoning behind the choices that are not -obvious from the bytes themselves. +The on-disk formats are the most stable surface in PTWM. A reader must +keep decoding every file that an older PTWM release wrote, and a format +change must make older readers fail cleanly instead of misreading the +new bytes. This document describes the bytes of each format and the +reasoning behind the choices that the bytes do not explain. + +PTWM writes three formats. All of them start with the bytes `\x89PTW`, +and the fifth byte tells them apart: + +| Magic | Format | Writer | Readers | +|---|---|---|---| +| `\x89PTWM\r\n\x1a\n` | legacy multi-tensor container | `ContainerLayout::Legacy` / `container_layout="legacy"` | v1.2 and later; older releases with a caveat (see [Compatibility](#compatibility)) | +| `\x89PTWX\r\n\x1a\n` | legacy container with embedded WASM extensions | extension bundling | as `\x89PTWM` | +| `\x89PTWC\r\n\x1a\n` | compact multi-tensor container | `ContainerLayout::Compact`, the default | v1.2 and later | +| `\x89PTWS` | compact single-tensor frame | `Compressor.compress`, the default | v1.2 and later | + +The three container magics are nine bytes long. The frame magic is five +bytes long. Compare the whole magic, not a prefix: readers before v1.2 +check only for `\x89PTWM`, and that check is what makes them reject the +newer formats. ## Why a custom container at all @@ -14,325 +29,709 @@ the integration adapter (`compress_safetensors_file` with `mode="b"`), which produces a `.safetensors` file whose data section is a `.ptwm` blob. The native `.ptwm` container exists alongside the integration because three properties are hard to achieve in -`.safetensors` directly: +`.safetensors` directly. First, per-plane codec dispatch. A trained-weight tensor compresses best when its exponent byte and its mantissa bytes flow through -different codecs. The plane records in PTWM let the decoder know -which codec was used for which slice of bytes, with the slice -boundaries recorded by the preprocessing chain rather than guessed -at decode time. `.safetensors` has no notion of per-tensor metadata -beyond shape and dtype. - -Second, shared state. The microscaling codecs (`Order1ScaleAC` and -`PerGroupCodebook`) need a small probability table or codebook that -multiple tensors share. The `.ptwm` prelude stores these once per -file rather than once per tensor. - -Third, the extension table. Third-party codecs need a stable way to +different codecs. The plane records in PTWM tell the decoder which +codec produced each plane, and the preprocessing chain fixes the plane +boundaries, so the decoder does not guess them. `.safetensors` has no +notion of per-tensor metadata beyond shape and dtype. + +Second, shared state. Some codecs (`Order1ScaleAC`, for example) use a +probability table that several tensors can share. The prelude stores +such state once per file instead of once per tensor. + +Third, the Extension Table. Third-party codecs need a stable way to say "the bytes in this plane were produced by the codec with this -canonical id". The extension table makes that explicit and lets the -verifier check the contribution chain before any third-party code -runs. +canonical id". The Extension Table makes that explicit and lets the +verifier check the contribution before any third-party code runs. + +## Conventions + +These rules hold for every format unless a section says otherwise. + +- Fixed-width integers are little-endian. +- A *varint* is unsigned LEB128 (`crate::varint`). It is at most 10 + bytes long and must be minimal: a reader rejects a varint with + redundant trailing `0x80` groups, so every value has one encoding. + Every varint field also has a range, given in its section, which the + reader checks. +- A reader rejects bytes it does not understand (unknown flag bits, + unknown tags, unknown enum values) in the compact formats. The legacy + container is older and less strict; [What the legacy reader + checks](#what-the-legacy-reader-checks) lists the gaps. +- Every container ends with the 8-byte EOF sentinel + `89 54 57 4D 45 4F 46 1A` (`\x89TWMEOF\x1a`, `index::EOF_SENTINEL`). + The reader checks it first, so a truncated file fails in constant + time. The first byte has the high bit set, like the magic, so a + transfer that strips the eighth bit also breaks the sentinel. + +## Versioning + +No format stores a numeric version. A version field tends to encourage +a v1-then-v2 mindset: every reader becomes a switch statement over the +version, and every breaking change adds another arm. PTWM uses two +other mechanisms instead. + +1. **A new magic for an incompatible layout.** `\x89PTWX`, + `\x89PTWC` and `\x89PTWS` are examples. A reader that predates a + magic rejects the file at the magic check, so an old reader never + applies its rules to new bytes. The rule runs in that direction: a + new reader keeps reading the old magics. +2. **Reserved flag bits for compatible additions.** In the compact + formats every flag bit that is not defined must be zero, and the + reader rejects the file otherwise. A future feature takes a free + bit, and readers that predate it reject files that use it with a + clear error. The PTWC header reserves all 16 bits of `hflags` for + this. + +Built-in operations and codecs are identified on the wire by the +`OpId` and `CodecId` wire ids in PTWC and PTWS, and by the built-in +name in the legacy Extension Table, whose ids are derived from the name +and the release version (see [Built-in entries](#built-in-entries)). +The legacy plane record and prelude also store `CodecId` wire ids. A +reader decodes a built-in with the algorithm of its own release. **The meaning of a built-in wire id must +never change.** A change to what a built-in op or codec does to the +bytes, even a bug fix, needs a new `OpId` or `CodecId` and a new +built-in name, with the old ones kept for decoding existing files. +The same applies to a change to the layout of a codec's state or +params: a new state version or a new id, never a reinterpretation. + +## Legacy container (`\x89PTWM`, `\x89PTWX`) + +`ContainerLayout::Legacy`. Every PTWM release reads this layout; +readers before v1.2 also need the file's built-in ids to be those of +their own version (see [Compatibility](#compatibility)). + +``` +[ header ] 91 bytes +[ Extension Table ] header: extension_table_offset/length +[ prelude ] header: shared_prelude_offset; the +[ chain registry ] shared-prelude section covers both +[ tensor record 1 ] +[ tensor record 2 ] +... +[ tensor index ] header: tensor_index_offset/length +[ EOF sentinel ] 8 bytes +``` + +The writer streams the records and then seeks back to fill in the +header, so it needs a seekable output. + +### Header + +``` +magic [u8; 9] \x89PTWM\r\n\x1a\n or \x89PTWX\r\n\x1a\n +flags u16 bit 0: PTWX (embedded WASM) +extension_table_offset u64 91 in files current writers produce +extension_table_length u64 +extension_table_hash [u8; 32] blake3 of the Extension Table bytes +tensor_index_offset u64 +tensor_index_length u64 +shared_prelude_offset u64 +shared_prelude_length u64 prelude and chain registry together +``` + +The header is fixed-size, so a reader can map the file, read 91 bytes +and go straight to the section it needs. There is no field for the +chain registry: it starts where the prelude ends, inside the +shared-prelude section, and the reader parses both from that section +alone. Tensor records start at the end of the shared-prelude section. + +The Extension Table hash sits in the header rather than next to the +table, so a reader checks the table against bytes it has already +read. Bit 0 of `flags` must be set exactly when the magic is +`\x89PTWX`. + +### Extension Table + +``` +count u32 at most 65,536 +entry × count: + canonical_id [u8; 32] + human_label u16 length + UTF-8 + kind u16 Transform, PlaneCodec, ... + abi_version u16 + lifecycle u8 + flavor_hints u8 bit 0 WASM, bit 1 native, bit 2 host + capabilities u32 length (at most 1 MiB) + canonical CBOR + attestation_kind u8 1 = PGP signature + attestation u16 length + bytes + install_hint u16 length + UTF-8 (length 0: none) + embedded_wasm_offset u64 0: none + embedded_wasm_length u64 0: none +``` + +The table lists every contribution the file needs at decode time: +first the ops of every registry chain, then the plane codecs that the +plane records reference, in built-in registry order. Chain nodes and +plane records refer to entries by position. A built-in entry costs +about 130 bytes (six entries take 779 bytes in a one-tensor fp32 +file), which is why the compact container stores built-ins by wire id +instead. + +Writers up to the v1.1.0 release interned all 16 built-in plane +codecs, about 2 KB per file, whether or not a plane used them. Readers resolve +`codec_table_idx` by position and check the table hash against the +bytes present, so they read both forms. The transcoder rebuilds the +table from the chain registry and the codec list and keeps whichever +form reproduces the recorded hash. + +The flavor hints are a bitfield rather than an enum because a +contribution can ship in several flavours and the reader picks one by +host capability. The capabilities are CBOR so the reader needs no TOML +parser. + +#### Built-in entries + +A built-in entry has the canonical id +`blake3(BUILTIN_PUBKEY ‖ 0x00 ‖ name ‖ 0x00 ‖ version)` and the label +`io.ptwm.builtin.@`, where `version` is the crate +version of the writer (`CURRENT_BUILTIN_VERSION`). The id therefore +changes with every release, although the entry means the same thing. + +From v1.2 the reader handles this with `normalize_builtin_ids`: after +the hash check, an entry whose label names a built-in and whose id +re-derives from that label is mapped to the current built-in. An entry +whose id does not re-derive keeps its id and is treated as third-party, +so a forged label gains no trust. Readers before v1.2 recognise only +the ids of their own version. + +### Prelude + +``` +count u32 +entry × count: + shared_state_id u16 + codec_id u16 CodecId wire id + state_format_version u8 + applies_to_mask u8 bit (r − 1) for PlaneRole r in 1..=8 + reserved u16 + state_xxhash64 u64 xxh64 (seed 0) of state_bytes + name u16 length + UTF-8, zero-padded to a multiple of 4 + state_bytes u32 length + bytes +``` + +A plane with state source Shared names an entry by `shared_state_id` +(its `state_info`). The reader checks the version, the role mask and +the hash before it hands the state to the codec. + +### Chain registry + +``` +num_chains u16 at most 0x8000 +chain × num_chains: + chain_id u16 at most 0x7FFF; writers use the position + chain legacy chain wire +``` + +The legacy chain wire (`chain::wire`) is: + +``` +num_nodes, num_edges, num_terminals u8 each +nodes_len, edges_len, terminals_len u16 each, section byte lengths +node: op u16 (Extension Table index), param_len u8, params +edge: src u8, src_out u8, dst u8, dst_in u8, has_role u8, + [role], vendor_len u16, vendor bytes +terminal: node u8, out u8, role +``` + +A tensor record addresses a registry chain with 15 bits, so a writer +fails on a registry with more than 0x8000 chains or a chain id above +0x7FFF. + +### Tensor record + +``` +chain_ref u16 bit 15: inline chain; else bits 0-14 are the chain_id +dtype_code u16 +input_format u8 +num_planes u8 +flags u8 0x01 dependencies, 0x02 metadata, 0x04 payload hash +reserved u8 +orig_size u64 raw tensor bytes +name u16 length + UTF-8, zero-padded to a multiple of 4 +[payload_hash] u64 xxh64 (seed 0) of the raw tensor (0x04) +[inline chain] Extension Table (u32 count + entries as above), then a + legacy chain whose ops index that table (bit 15) +[dependencies] u16 count; per dependency: ref_kind u8, ref_flags u8 + (bit 0: hash), u16 length + ref bytes, [blake3 32 B] + (0x01) +[metadata] u32 length + CBOR shape map (shape, dtype name) (0x02) +plane × num_planes +``` + +The fixed part is 18 bytes. With the name, the shape metadata and four +planes, an fp32 tensor record carries 166 bytes besides its payload, or +190 bytes with the payload hash and plane CRCs (name of 28 bytes). + +### Plane record + +``` +role u8 PlaneRole (diagnostic label) +plane_flags u8 0x01 chunk table, 0x02 CRC32, 0x04 codec_table_idx + present (always set) +codec_id u16 CodecId wire id +codec_table_idx u16 Extension Table index of the plane codec +state_source u8 0 none, 1 inline, 2 shared, 3 external +state_version u8 +state_info u16 shared: shared_state_id +inline_state_len u32 +payload_len u64 +[crc32] u32 CRC32 of the decoded plane (0x02) +[chunk table] u32 count; per chunk: offset_in_payload u32, + decoded_size u32 (0x01) +[external state] ref_kind u8, ref_flags u8 (bit 0: hash), u16 length + + ref bytes, state_format_version u8, [blake3 32 B] + (state_source 3) +layout_kind u8 0 flat, 1 rows +[row_len] u32 non-zero (rows) +inline state inline_state_len bytes +payload payload_len bytes +``` + +A flat plane without CRC costs 23 bytes plus its payload. The decoder +finds the codec through `codec_table_idx` alone. `codec_id` duplicates +it; every writer fills both, and the reader requires `codec_id` to be a +known `CodecId` but does not dispatch on it. The reader reads +`codec_table_idx` whether or not flag 0x04 is set. + +The plane record stores no shape and no decoded length. The decoder +derives each plane's length from the chain (see [Reader +limits](#reader-limits)). + +### Tensor index + +``` +count u64 +entry × count: + name_hash u64 xxh64 (seed 0) of the UTF-8 name + tensor_offset u64 absolute offset of the record + tensor_len u64 record length +``` + +Entries are sorted by `name_hash`. A lookup binary-searches the hash, +then compares the name in each record of the run of equal hashes, so +two names with the same hash resolve correctly. Writers reject +duplicate tensor names. Tools that list or decode every tensor return +them in record (write) order, not index order. + +### What the legacy reader checks + +The reader checks the EOF sentinel, the magic, the PTWX flag against +the magic, the bounds of every section, the Extension Table hash, the +prelude and registry within their section, every index entry's range +before it reads a record, the decoded plane lengths against the plane +bound (see [Reader limits](#reader-limits)), the plane CRCs and the +payload hash. + +It does not check: + +- that the index is sorted (an unsorted index makes lookups miss); +- the reserved bytes and unknown flag bits of records and the prelude; +- that records are contiguous or that every byte belongs to a section. + +None of these can make the reader read out of bounds; they are why the +legacy layout has more than one encoding for the same content, and why +the compact layout rejects them. + +Payloads are not aligned. Only the names in the tensor record and the +prelude are padded to a multiple of 4 bytes, and the rest of the file +has no alignment. Codecs read their input as bytes and do not assume +alignment. + +## Compact container (`\x89PTWC`) + +`ContainerLayout::Compact`, the default for multi-tensor output since +v1.2. It holds the same information as the legacy container, with +varints, built-ins stored by wire id, names moved into the index, and +fields the reader can derive left out. The reader parses it into the +same in-memory records as a legacy file, so decoding, trust checks and +the transcoder share one code path. + +``` +FILE = HEADER | EXTENSION TABLE | PRELUDE | REGISTRY | BODY | INDEX | FOOTER +``` + +All sections are contiguous. The writer streams them and never seeks: +the header comes first because the table, prelude and registry are +known before the first record, and the index length goes in the +footer. -## Top-level layout +Every section parser gets exactly its slice and must consume all of +it. Together with minimal varints and zero reserved bits, this gives +each file one encoding for its content, so the transcoder's +`implode(explode(x))` reproduces `x` byte for byte. -The container is a flat byte stream with the following sections in -order. All offsets and lengths are stored in little-endian. +### Header and footer ``` -[ 91-byte header ] -[ extension table ] variable length -[ shared-state prelude ] variable length -[ tensor 1 record ] -[ tensor 2 record ] +magic [u8; 9] \x89PTWC\r\n\x1a\n +hflags u16 must be 0 (all bits reserved) +et_len varint Extension Table bytes +prelude_len varint prelude bytes +registry_len varint registry bytes +meta_xxh64 u64 xxh64 (seed 0) of the header bytes before this + field, then the Extension Table, prelude and + registry bytes ... -[ tensor index ] sorted by name hash -[ 8-byte EOF sentinel: 0x89 P T W M EOF] -``` - -Three of the four "variable length" sections (extension table, -prelude, tensor index) carry their offset and length in the header. -The header is fixed-size, so a reader can mmap the container, parse -91 bytes, and jump to any section without scanning the rest. - -The EOF sentinel exists so a reader can detect a truncated file in -constant time. The sentinel is `\x89PTWMEOF\x1a` plus padding; the -first four bytes are PNG-style so that a file misnamed as an image -type still trips an early decode error. - -## Header (91 bytes) - -The choice of 91 bytes is the result of two constraints meeting. A -header that fits in one or two cache lines on a modern x86 machine -keeps header parsing branch-predictable; the natural targets are 64 -or 128 bytes. The integrity field needs a cryptographic-strength -hash so that a malformed extension table is detected before any of -its entries are read; 32 bytes is the natural width for that. The -remaining offset and length fields claim another 56 bytes between -them, which puts the total above the 64-byte cache line. The 91 -bytes that result fit comfortably in two cache lines and leave room -for the offset fields without a numeric version field competing for -space. - -The header carries: - -``` -magic [u8;9] PTWM or PTWX -flags u16 -extension_table_offset u64 -extension_table_length u64 -extension_table_hash [u8;32] blake3 -tensor_index_offset u64 -tensor_index_length u64 -shared_prelude_offset u64 -shared_prelude_length u64 -``` - -The extension table hash is committed in the header rather than next -to the table itself. A reader that wants to check integrity does it -without trusting any byte that lives outside the header. - -## No numeric version field - -PTWM does not store a version `u16` in the header. A version field -tends to encourage a v1-then-v2 mindset: every reader becomes a -switch statement over the version, and every breaking change adds -another arm. The alternative used here is per-contribution version -negotiation through the extension table, with a fresh magic for any -wire-format change that crosses the contribution boundary. The -PTWX magic (for the embedded-WASM variant) is one example of the -pattern; a future format-incompatible variant would pick a new -magic. Files with the old magic are rejected by new readers, so -there is no path by which a reader silently applies new rules to -old bytes. - -This is the kind of architecture choice that is hard to walk back. -Adding a version field later is possible but the files in the wild -at the time of the addition cannot benefit from it. - -## Extension table - -The extension table is the per-file manifest of every contribution -the container needs at decode time. Each entry records: - -``` -canonical_id [u8;32] blake3-derived -human_label utf8 length-prefixed -kind u16 Transform, PlaneCodec, ... -abi_version u16 -lifecycle u8 none, thread, process -flavor_hints u8 bitfield: wasm, native, host -capabilities cbor -attestation Attestation -install_hint Option -embedded_wasm_offset Option -embedded_wasm_length Option -``` - -The canonical id is the cryptographic identifier the verifier checks -against the keyring. The human label is informational. The kind and -abi_version together tell the host which ABI surface the -contribution expects. - -Two design choices here are worth flagging. The flavor hints are a -bitfield rather than an enum because a contribution can ship in -multiple flavours and the reader chooses based on host capability. -The capabilities are CBOR rather than TOML because the format is -parsed in Rust without an extra dependency on a TOML parser path -that the reader would otherwise not need. - -The writer interns the ops of every chain in the chain registry, -then only the plane codecs that the plane records reference, in the -order of the built-in registry. Writers up to v1.1 interned all 16 -built-in plane codecs, about 2 KB per file, whether or not a plane -used them. Readers resolve `codec_table_idx` by position and check -the table hash against the bytes present, so they decode both -layouts. The table is a pure function of the chain registry and the -codec set; the explode/implode transcoder relies on this to rebuild -the same bytes from the codec list in the registry manifest. - -## Plane records - -Each tensor's record holds a fixed header plus a sequence of plane -records, one per terminal of the preprocessing chain. The plane -record carries: - -``` -codec_id u16 legacy CodecId enum -codec_table_idx u16 new: index into extension table -role u8 role bytes from PPG terminal -state_source u8 none, inline, shared, external -state_version u8 per-codec version of the state -state_info u16 interpretation depends on source -inline_state_bytes bytes present if source == inline -shape list per-plane shape -payload_hash u64 optional -plane_flags u8 bitfield -payload_len u64 -crc32 u32 optional -chunk_table Option<...> -external_state Option<...> -payload_bytes bytes -``` - -The `codec_id` field predates the `codec_table_idx`. New writers fill -both. Readers that see a valid `codec_table_idx` route through the -extension table; readers that see only the legacy `codec_id` fall -back to the in-tree codec table. This dual path will be removed once -no field-deployed writer emits the legacy form alone. - -The state source is the most subtle field. PTWM's compression -strategy depends on whether a codec needs per-plane state, a state -shared across planes within a tensor, a state shared across tensors -(the prelude case), or no state at all. The four-valued enum lets -the writer record exactly one of those four, which the reader uses -to find the bytes to feed back into the codec. - -## Compact single-tensor frame - -`Compressor.compress` writes one tensor. The container sections -above cost about 500 to 1,000 bytes for one tensor, which is more -than the payload of a small tensor. By default the compressor writes -a compact frame instead, with its own magic `\x89PTWS`: - -``` -magic [u8;5] \x89PTWS -flags u8 payload hash, plane CRC, shape, dtype name, - delta reference, stored, size; unknown bits - and plane CRC with payload hash rejected -dtype_code varint -input_format u8 -orig_size varint (if flagged) -shape varint ndim + varint dims (if flagged) -dtype_name varint length + utf8 (if flagged) -delta_ref [u8;32] blake3 of the reference (if flagged) -payload_hash u64 xxh64 of the raw bytes (if flagged) -body +index_len u64 (footer) +sentinel [u8; 8] \x89TWMEOF\x1a (footer) ``` -`orig_size` is stored only when the shape does not give it as -`product(shape) × element_size(dtype_code)`. The dtype name is stored -only when it differs from the canonical name of `dtype_code`. Varints -are minimal unsigned LEB128, and a reader rejects any field written -in a form the writer would not produce, so each field has one -encoding for the value it holds. The frame as a whole is not unique -for a tensor: the writer chooses the codecs, the chain, a stored or a -coded body, and plane state 0 or 2 (below), and the reader accepts -any valid choice. A reader rejects bytes after the end of the body. +The header is 22 bytes for files whose sections are each shorter than +128 bytes, and grows by one byte per seven bits of section length. The +body runs from the end of the registry to the start of the index, +which is `file_len − 16 − index_len`. The reader checks all offsets +with checked arithmetic and rejects a file whose body would have a +negative length. -The body is either the raw bytes (the stored flag, written when no -codec makes the tensor smaller) or a coded tensor: the chain in the -compact chain wire (CW2, below), then one record per terminal plane: -a varint `codec_id << 2 | state` (state 0 none, 1 inline with a -version byte and length-prefixed bytes, 2 inline with version 0 and -no bytes; states 0 and 2 decode alike), the payload length, an -optional CRC32 and the payload. The coded body ends with the last -payload. The payload hash and the plane CRCs check decoded bytes, not -the coded payload. The payload hash already covers every plane of the -one tensor, so the writer stores plane CRCs only when it does not -store the hash, and a reader rejects a frame with both. The frame does not store plane layouts, plane lengths or -the plane count, because the decoder derives them from the chain. - -The fixed cost is about 10 to 70 bytes, depending on the chain and -on which integrity fields are on. With the default settings (payload -hash on) it is at most about 50 bytes for a bf16 tensor, 55 bytes for -an fp32 tensor with the split chain, and 61 bytes for an fp32 tensor -with the IntDelta or PredictorXor chain that smooth data selects. -Plane CRCs instead of the hash add up to 8 bytes more for four -planes. -A tensor whose planes need shared state (a prelude) or chunked -encoding cannot use the frame; the compressor writes a full -container for it. The frame holds no extension table, so it can -reference built-in ops and codecs only. - -### Compact chain wire (CW2) - -`chain::wire_compact` stores a chain in about 15 to 20 bytes, against -50 to 70 for the registry wire form: +`meta_xxh64` protects the sections a reader parses before it touches +a record. xxh64 is enough here: it detects corruption, and the trust +decision does not rest on it, because a built-in is stored by wire id +and a third-party entry carries its own canonical id and attestation. + +### Extension Table + +``` +count varint at most 65,536 and at most et_len / 2 +entry × count: + tag u8 + 0x01 OpId wire id (varint) built-in transform + 0x02 CodecId wire id (varint) built-in plane codec + 0x03 name length varint (1..=255), then the built-in name + built-in without a wire id + 0x7F length varint, then one legacy entry of exactly that length + third-party entry +``` + +The reader expands tags 0x01 to 0x03 into the built-in entry of its own +version, so a PTWC file does not depend on the release that wrote it. +Six built-in entries take 13 bytes, against 779 in the legacy table. + +The reader rejects unknown tags, ids and names, duplicate entries, a +built-in stored with tag 0x7F (of this release or, by its label, of +another one), a tag 0x03 entry for a built-in that has a wire id, and +any entry with embedded WASM. Embedded WASM exists only in `\x89PTWX` +files; the PTWC writer refuses such entries. The writer orders entries +like the legacy writer: chain ops first, then the referenced plane +codecs. + +### Prelude + +The legacy prelude bytes, parsed within `prelude_len`. Shared state is +rare and the format is already compact enough. + +### Registry + +``` +count varint at most 0x8000 (chain ids 0..=0x7FFF) +chain × count: chain_len varint, then a CW2 chain +``` + +The chain id is the position. A node's `op_ref` is an Extension Table +index that must name a built-in transform. A registry chain may have a +bound Source (see [CW2](#compact-chain-wire-cw2)), which takes its +params from the record that uses the chain. Tensors that differ only in +shape then share one registry chain: 64 fp32 tensors with 64 distinct +shapes need two registry chains (48 bytes) in PTWC, against 64 chains +(3,872 bytes) in the legacy layout. + +### Tensor record + +Records follow each other in write order. A record does not store its +name: the index does. + +``` +tflags u8 bit 0 PAYLOAD_HASH, bit 1 INLINE_CHAIN, bit 2 DEPS, + bit 3 SHAPE, bit 4 DTYPE_NAME (needs SHAPE), + bit 5 RAW_META (excludes SHAPE); bits 6-7 zero +[chain_id] varint < registry count (not INLINE_CHAIN) +dtype_code varint at most 65,535 +input_format u8 +orig_size varint +[shape] varint ndim (at most 64), varint per dimension (SHAPE) +[dtype_name] varint length + UTF-8 (DTYPE_NAME) +[raw_meta] varint length + bytes (RAW_META) +[payload_hash] u64 xxh64 (seed 0) of the raw tensor (PAYLOAD_HASH) +[inline chain] varint length + compact Extension Table of the chain's + ops in first-use order, then varint length + CW2 chain + (INLINE_CHAIN) +[dependencies] varint count (at least 1); per dependency: ref_kind u8, + ref_flags u8 (bit 0: hash), varint length + ref bytes, + [blake3 32 B] (DEPS) +n_planes varint at most 255 +plane × n_planes +``` + +SHAPE stands for the CBOR shape map of the legacy record, built from +the shape and the canonical dtype name of `dtype_code`, or from the +stored dtype name under DTYPE_NAME. The reader rejects a DTYPE_NAME +equal to the canonical name. The writer uses SHAPE only when it +reproduces the record's metadata bytes exactly; other metadata goes in +RAW_META. + +The record must end exactly at the length the index gives. The +in-memory `chain_ref`, `flags` and `name` are rebuilt from the record +and the index, as the legacy parser would produce them. + +### Plane record + +``` +pflags u8 bit 0 CRC32, bit 1 CHUNKS, bits 2-3 state source + (0 none, 1 inline, 2 shared, 3 external), + bit 4 LAYOUT_ROWS, bit 5 STATE_EXTRA; bits 6-7 zero +role u8 PlaneRole +codec_ref varint Extension Table index of a plane codec +[state_version] u8 (STATE_EXTRA) +[state_info] varint at most 65,535 (STATE_EXTRA) +[inline_len] varint (state inline) +payload_len varint +[row_len] varint 1..=u32::MAX (LAYOUT_ROWS) +[crc32] u32 CRC32 of the decoded plane (CRC32) +[chunks] varint n (at most payload_len); n × (varint offset, + varint decoded_size), each at most u32::MAX (CHUNKS) +[external] ref_kind u8, ref_flags u8 (bit 0: hash), varint length + + ref bytes, state_format_version u8, [blake3 32 B] + (state external) +[inline state] inline_len bytes +payload payload_len bytes +``` + +STATE_EXTRA is set exactly when `state_version` or `state_info` is +non-zero. The plane does not store `codec_id`: the reader derives it +from the Extension Table entry, and the writer rejects a record whose +`codec_id` disagrees. Role and layout are stored so that parsing a +record never needs the chain. + +A typical plane takes 4 bytes plus its payload, or 8 with a CRC. + +### Tensor index + +``` +count varint at most index_len / 3 +entry × count, in body order: + name varint length (1..=65,535) + UTF-8 + rec_len varint, at least 3 +``` + +Record offsets are the running sum of `rec_len` from the start of the +body, and the sum must equal the body length exactly, so records cannot +overlap or leave gaps. Names must be unique. The reader sorts the +entries by name hash in memory; a lookup compares names, so a hash +collision cannot return the wrong tensor. `tensor_names` reads the +index alone, without touching the records. + +### Sizes + +Measured with `benchmarks/container_overhead.py`, fp32 tensors of shape +`(64,)` named `model.layers.NNN.attn.q.bias` (28 bytes), written with +`compress_ptwm_blob`. "Non-payload" excludes plane payloads and inline +and shared codec state. + +| Case | PTWC B/tensor | Legacy B/tensor | +|---|---|---| +| 64 tensors, same shape, hash and CRC on | 83 | 235 | +| 64 tensors, distinct shapes, hash and CRC on | 83 | 288 | +| 64 tensors, same shape, hash and CRC off | 59 | 211 | +| 1 tensor, hash and CRC on | 165 | 1,176 | + +Per tensor, the PTWC index entry takes 31 bytes (mostly the name) and +the record about 50 bytes with the payload hash and four plane CRCs, +26 without. CI holds these to the bounds in +`tests/ptwm/benchmarks/test_container_overhead.py` and +`crates/ptwm-core/tests/ptwc.rs`. + +## Compact chain wire (CW2) + +`chain::wire_compact` stores a chain in about 15 to 30 bytes, against +50 to 70 for the legacy chain wire. PTWC registry and inline chains and +the PTWS frame use it. ``` n_nodes, n_edges, n_terms varint each (n_nodes 1..=255, others <= 255) node: op_ref varint, param_len varint (<= 255), params n_spine varint (<= n_edges; only when n_edges > 0) -edge: src, src_out, dst, dst_in varint; etag u8 (bit0 role, - bit1 vendor); [compact role]; [varint length + vendor bytes] +edge: src, src_out, dst, dst_in varint; etag u8 (bit 0 role, + bit 1 vendor; other bits zero); [compact role]; + [varint length (1..=65,535) + vendor bytes] terminal: ref u8 (node << 4 | out, or 0xFF + node u8 + out u8); compact role ``` -- `op_ref` is the `OpId` wire id in the compact frame. +- `op_ref` is an Extension Table index in PTWC and the `OpId` wire id + in PTWS. - The first `n_spine` edges are `(i, 0) → (i + 1, 0)` with no role - override and no vendor bytes, and are not stored. -- A Source node with `param_len = 0` is a bound Source: its params - are those of the frame's shape and dtype (`u8` ndim, `u32` dims, - `u16` chain dtype). The writer binds a Source exactly when its - params equal these, so chains that differ only in shape encode to - the same bytes. + override and no vendor bytes, and are not stored. The writer always + stores the longest such run, and the reader rejects a stored first + edge that would extend it. +- A Source node with `param_len = 0` is a *bound Source*: its params + are those of the record's or frame's shape and dtype (`u8` ndim, + `u32` dims, `u16` chain dtype). The writer binds a Source exactly + when its params equal these, and the reader rejects explicit params + that equal them. Real Source params are never empty, so the marker + cannot collide with an explicit Source. A record without a shape + cannot use a chain with a bound Source. - A compact role is one byte for every role production chains use (scale formats, nibble kinds, mantissa and integer bytes up to 8, `IntN` up to 32 bits, raw, index, exponent byte, fp4 value, XOR - residual), and `0xFF` followed by the registry role encoding for - the rest. - -The frame follows the same versioning rule as the container: no -numeric version field, and a new magic for an incompatible layout. -Readers that predate the frame reject `\x89PTWS` as "not a .ptwm -container" rather than misparse it. `CompressionConfig(framing= -Framing.CONTAINER)` keeps the full container for consumers pinned to -such readers, and the CLI writes `.ptwm` files as full containers so -the random-access and store tools can open them. - -## Wire-format invariants - -The reader assumes the following without checking, since checks are -expensive on large files and a violated invariant would surface as a -codec error at decode time anyway: - -The plane shape always agrees with the role declared in the chain -terminal. The PPG validator enforces this at chain construction; -malformed chains never reach the wire format. - -Plane payloads are aligned to their codec's preferred alignment. The -writer pads with zeros if necessary; the reader trusts the alignment -and may panic on misaligned access. - -The tensor index is sorted by name hash. The reader uses a binary -search and would return wrong results on an unsorted index. - -The reader does check the magic, the extension-table hash, and the -EOF sentinel. The first two protect against accidental corruption; -the last protects against truncation. - -## Why blake3 and not sha256 - -Two reasons. Blake3 is faster on the workstation CPUs PTWM targets -for compression work, by about a factor of three on the inputs that -matter. The speed matters because the extension-table hash is -computed on every file write and verified on every file read; the -manifest plus binary hash is computed when verifying a signed -extension bundle. A slower hash would dominate the open path on -small files. - -Blake3 also avoids the cryptographic agility surface that PTWM has -no use for. The library is committed to a single hash function; -adding a second one to support sha256 callers would mean every plane -hash needs a discriminator. The choice may need to be revisited if -NIST publishes a new standard that supplants blake3 in policy -documents, but the migration cost is bounded by the number of -on-disk hash fields, which is small. + residual), and `0xFF` followed by the legacy role encoding for the + rest. An escaped role that has a one-byte form is rejected. + +## Compact single-tensor frame (`\x89PTWS`) + +`Compressor.compress` writes one tensor. A container's sections cost +about 120 to 140 bytes (PTWC) or 500 to 1,000 bytes (legacy) for one +tensor, which is more than the payload of a small tensor. By default +the compressor writes a compact frame instead: + +``` +magic [u8; 5] \x89PTWS +flags u8 bit 0 payload hash, bit 1 plane CRC, bit 2 shape, + bit 3 dtype name, bit 4 delta reference, + bit 5 stored, bit 6 size; bit 7 zero, and + plane CRC with payload hash is rejected +dtype_code varint at most 65,535 +input_format u8 +orig_size varint (size) +shape varint ndim (<= 64) + varint dims (shape) +dtype_name varint length + UTF-8 (dtype name) +delta_ref [u8; 32] blake3 of the reference (delta reference) +payload_hash u64 xxh64 of the raw bytes (payload hash) +body +``` + +`orig_size` is stored only when the shape does not give it as +`product(shape) × element_size(dtype_code)`. The dtype name is stored +only when it differs from the canonical name of `dtype_code`. A reader +rejects any field written in a form the writer would not produce, so +each field has one encoding for the value it holds. The frame as a +whole is not unique for a tensor: the writer chooses the codecs, the +chain, a stored or a coded body, and plane state 0 or 2 (below), and +the reader accepts any valid choice. A reader rejects bytes after the +end of the body. + +The body is either the raw bytes (the stored flag, written when no +codec makes the tensor smaller) or a coded tensor: a CW2 chain whose +`op_ref` is the `OpId` wire id and whose bound Source takes the +frame's shape and dtype, then one record per terminal plane: + +``` +head varint CodecId wire id << 2 | state + (0 none, 1 inline, 2 inline with version 0 and + no bytes; 3 rejected) +[state] version u8, varint length, bytes (state 1) +payload_len varint +[crc32] u32 CRC32 of the decoded plane (plane CRC) +payload payload_len bytes +``` + +States 0 and 2 decode alike. The payload hash already covers every +plane of the one tensor, so the writer stores plane CRCs only when it +does not store the hash. The frame does not store plane layouts, plane +lengths or the plane count, because the decoder derives them from the +chain. + +The fixed cost is about 20 to 70 bytes, depending on the chain and on +which integrity fields are on. With the default settings (payload hash +on) it is at most 43 bytes for a bf16 tensor, 55 bytes for an fp32 +tensor with the split chain, and 61 bytes for an fp32 tensor with the +IntDelta or PredictorXor chain that smooth data selects. Plane CRCs +instead of the hash add up to 8 bytes more for four planes. + +A tensor whose planes need shared state (a prelude) or chunked encoding +cannot use the frame; the compressor writes a full container for it, +in the configured `container_layout`. The frame holds no Extension +Table, so it can reference built-in ops and codecs only. + +`CompressionConfig(framing=Framing.CONTAINER)` writes a full container +instead of the frame, for the random-access, store and transcode tools. + +## Compatibility + +| File | Readers before v1.2 | v1.2 and later | +|---|---|---| +| `\x89PTWS` frame (default single-tensor output) | reject: "not a .ptwm container" | read | +| `\x89PTWC` container (default multi-tensor output) | reject: bad magic | read | +| `\x89PTWM` container | read if its built-in ids are of the reader's version | read, built-in ids of any release | + +Old readers fail cleanly on the new formats. A Rust reader before v1.2 +passes the EOF sentinel check and fails in `Header::read` with "bad +magic". A Python reader before v1.2 checks the 5-byte `\x89PTWM` +prefix and reports that the input is not a container. The transcoder +before v1.2 rejects the `"ptwm-manifest-c"` manifest of a PTWC +container. + +To write files that a v1.1 reader opens: + +- Python: `CompressionConfig(framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY)`, or + `container_layout="legacy"` on `compress_ptwm_blob`, + `write_ptwm_shard`, `compress_safetensors_file`, the LMDB and + WebDataset writers and `_rust.compress_model`. +- CLI: `ptwm compress --container-layout legacy`. +- Rust: `CompressorOptions { container_layout: ContainerLayout::Legacy, + .. }` or `ContainerWriter::new`. + +A legacy file carries the built-in ids of the writer's crate version +(see [Built-in entries](#built-in-entries)), and a v1.1 reader +recognises only the ids of version 1.1.0. This tree's crate version is +1.1.0, so its legacy output opens in v1.1. A build with another crate +version writes legacy files that v1.1 cannot open unless the writer +pins the built-in version with `ContainerWriter::set_builtin_version` +(Rust only). + +## Reader limits + +The readers treat every file as untrusted input. In both layouts: + +- every count is checked against the bytes that remain before a + vector is allocated, and every offset and length is added with + checked arithmetic; +- the chain's Source must declare exactly `orig_size` bytes in PTWC + and PTWS. A legacy Source may declare up to 8 × `orig_size`, so that + files written by the v1.1 `Format.BYTE` bug keep decoding; +- the plane lengths are derived symbolically from the chain before any + plane is decoded, and no plane may exceed `16 × orig_size + 4096` + bytes (`max_plane_len`), so a corrupt file cannot make a codec + allocate far more than the tensor; +- a decoded plane must have exactly its derived length in PTWC and + PTWS. In the legacy layout it may differ (older files carry Sources + that overstate or understate the raw size) but stays under the plane + bound. In every layout the inverse chain must give exactly + `orig_size` bytes; +- a file whose Extension Table holds only built-ins opens without + extension discovery or keyring access. + +## Why blake3 for the legacy table hash + +Blake3 is faster on the workstation CPUs PTWM targets than SHA-256, by +about a factor of three on the inputs that matter. The legacy +Extension Table hash is computed on every legacy write and checked on +every legacy open, and a slower hash would dominate the open path for +small files. PTWC uses xxh64 for its metadata check (see +[Header and footer](#header-and-footer)) because its table is a few +bytes and trust does not rest on the checksum. Blake3 stays the hash +for canonical ids, dependency references and extension bundles. + +Blake3 also avoids a hash-agility surface that PTWM has no use for. +Adding SHA-256 would mean every hash field needs a discriminator. The +choice may need revisiting if policy documents come to require another +hash, but the migration cost is bounded by the number of on-disk hash +fields, which is small. ## Sentinel public key -Built-in contributions use a canonical id derived from a sentinel -public key, `BUILTIN_PUBKEY = [0xBB; 32]`. The sentinel is treated as -trusted by construction; the verifier short-circuits before checking -the keyring when it sees an id derived from this key. The sentinel -bytes are recognisable, which helps when reading hex dumps of an -extension table during debugging. - -A real Ed25519 public key would never have all bytes equal to -0xBB, since the key derivation would require a chosen-private-key -attack on Ed25519 to construct one. The probability of accidentally -matching the sentinel from a real key is `2^-256`, which is the -collision bound the codebase relies on elsewhere. +Built-in contributions derive their canonical ids from a sentinel +public key, `BUILTIN_PUBKEY = [0xBB; 32]`. The reader recognises a +built-in by looking its id up in the table of derived built-in ids +(`builtin_registry`), after mapping ids of other releases (see +[Built-in entries](#built-in-entries)), and trusts it without a +keyring check. The sentinel bytes are recognisable, which helps when +reading hex dumps of an Extension Table. + +A real Ed25519 public key with all bytes equal to 0xBB would require a +chosen-private-key attack on Ed25519. The chance of a real key +matching the sentinel by accident is `2^-256`, the collision bound the +codebase relies on elsewhere. diff --git a/docs/site/content/20.concepts/5.container.md b/docs/site/content/20.concepts/5.container.md index 2e8bc5f..949bada 100644 --- a/docs/site/content/20.concepts/5.container.md +++ b/docs/site/content/20.concepts/5.container.md @@ -4,63 +4,108 @@ description: Multi-tensor container with random-access lookup, hash-verified pay --- The `.ptwm` file is a multi-tensor container. A single file holds many -tensors with a shared header, a tensor index, and per-plane codec -dispatch. +tensors with a shared Extension Table, a chain registry, a tensor +index, and per-plane codec dispatch. + +PTWM writes the container in one of two layouts: + +| Layout | Magic | Default | Readers | +|---|---|---|---| +| Compact (PTWC) | `\x89PTWC\r\n\x1a\n` | yes, since v1.2 | v1.2 and later | +| Legacy (PTWM) | `\x89PTWM\r\n\x1a\n` | no | v1.2 and later; v1.1 with a caveat | + +Both hold the same content and decode the same way. The compact layout +stores built-in codecs and operations by wire id, uses varints, and +keeps tensor names in the index. For 64 small fp32 tensors it costs +about 83 bytes per tensor besides the coded data, against 235 bytes for +the legacy layout. Choose the legacy layout only for readers older than +v1.2 (see [Compatibility](#compatibility)). + +Single-tensor output from `Compressor.compress` uses a third, smaller +format by default, the compact frame (`\x89PTWS`). See +[Single-tensor frame](#single-tensor-frame). ## Layout ```mermaid graph TD - subgraph Container [".ptwm File Container"] - M[Magic + Version
'PTWM' + byte] - P[Prelude
Shared State & Codec Menu] - I[Tensor Index
Name → Offset Maps] - - subgraph T1 [Tensor 1: Plane Dispatch Records] - R1_0[Plane 0: Codec, Shape, Hash] - R1_1[Plane 1: Codec, Shape, Hash] - end - - subgraph TN [Tensor N: Plane Dispatch Records] - RN_0[Plane 0: Codec, Shape, Hash] - RN_1[...] - end - - subgraph Payloads [Concatenated Plane Payloads] - PL1_0[Tensor 1 Plane 0 Bytes] - PL1_1[Tensor 1 Plane 1 Bytes] - PLN_0[...] + subgraph Container [".ptwm container"] + H[Header
magic, section lengths, checksum] + E[Extension Table
ops and codecs the file uses] + P[Prelude
shared codec state] + C[Chain registry
preprocessing chains] + subgraph Body [Tensor records, in write order] + T1[Tensor 1: dtype, shape, hash,
plane records + payloads] + TN[Tensor N: ...] end + I[Tensor index
name → record] + F[EOF sentinel] end - M --> P - P --> I - I --> T1 - I --> TN - T1 --> Payloads - TN --> Payloads + H --> E --> P --> C --> Body --> I --> F ``` -See `crates/ptwm-core/src/compressor.rs` for the exact wire format -fields and field widths. +Each tensor record holds its chain reference, dtype, shape and payload +hash, then one plane record per output of the chain. A plane record +names the codec of the plane, its codec state, an optional CRC32 and +chunk table, and is followed by the plane's payload bytes. + +The byte-level layout of both layouts, and of the single-tensor frame, +is in [`docs/architecture/wire-format.md`](https://github.com/khwstolle/ptwm/blob/master/docs/architecture/wire-format.md). ## Random access -The tensor index is small (a manifest of name → offset records). Opening -a `.ptwm` and fetching one tensor touches only: +The tensor index maps each name to its record. Opening a `.ptwm` and +fetching one tensor touches only: -1. The header + index (a few KB). -2. The plane records for that tensor (≤ 200 bytes per plane). -3. The plane payloads for that tensor. +1. The header, Extension Table, prelude, chain registry and index. For + a built-in-only file of 64 tensors these take about 2 KB in the + compact layout, most of it the names in the index. +2. The record of that tensor: about 50 bytes for an fp32 tensor in the + compact layout, 190 bytes in the legacy layout, plus its payloads. -Use [`TensorIndex`](/guides/random-access) to do this from Python. +Use [`TensorIndex`](/guides/random-access) to do this from Python. It +opens either layout and lists tensors in write order. ## Hash verification -An xxHash annotates each plane payload. The decoder verifies the hash -before returning bytes to the caller. A mismatch raises -`CompressionMethodNotSupportedError`, which lets callers distinguish -corruption from unsupported features. +Each tensor record can carry an xxh64 of the raw tensor bytes +(`emit_payload_hash`, 8 bytes per tensor), and each plane record a +CRC32 of the decoded plane (`emit_plane_crc`, 4 bytes per plane). The +decoder checks both before it returns bytes to the caller. A mismatch +raises `ptwm.Error`. + +## Single-tensor frame + +`Compressor.compress` writes a compact frame (`\x89PTWS`) by default: +one tensor with only the fields its decoder needs, at a fixed cost of +about 20 to 70 bytes. A tensor that needs shared codec state or chunked +planes falls back to a full container in the configured layout. +`Framing.CONTAINER` always writes a full container, which the +random-access, store and transcode tools need. + +## Compatibility + +Readers before v1.2 reject the compact container and the compact frame +with a "not a .ptwm container" or "bad magic" error; they do not +misread them. To write files that v1.1 opens, select the legacy layout +and the full container: + +```python +from ptwm import CompressionConfig, ContainerLayout, Framing + +config = CompressionConfig( + framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY, +) +``` + +The multi-tensor writers (`compress_ptwm_blob`, `write_ptwm_shard`, +`compress_safetensors_file`, the LMDB and WebDataset writers) take +`container_layout="legacy"`, and the CLI takes +`ptwm compress --container-layout legacy`. A legacy file opens in v1.1 +only when its built-in ids are those of version 1.1.0; see the +[release notes](/release-notes) for the details. ## Modes for safetensors input @@ -72,5 +117,6 @@ in one of two shapes; the `mode` argument selects which: | `"a"` (native) | A `.ptwm` directory with one `.ptwm` blob per shard | You control the loader and want the cleanest random-access surface. | | `"b"` (shell) | A `.safetensors` shell wrapping a `.ptwm` blob | You need the result to look like a regular `.safetensors` file to existing tooling. | -Both modes round-trip through PTWM. Mode `b` is what the -[safetensors integration](/guides/safetensors) emits by default. +Both modes round-trip through PTWM and both take `container_layout`. +Mode `b` is what the [safetensors integration](/guides/safetensors) +emits by default. diff --git a/docs/site/content/30.guides/01.cli.md b/docs/site/content/30.guides/01.cli.md index a35ed2a..c2f4506 100644 --- a/docs/site/content/30.guides/01.cli.md +++ b/docs/site/content/30.guides/01.cli.md @@ -37,6 +37,18 @@ ptwm decompress model.safetensors.ptwm # restores model.safetensors ``` +## Container layout + +New `.ptwm` files use the compact container layout (`\x89PTWC`), which +PTWM v1.2 and later read. For files that v1.1 opens, write the legacy +layout: + +```bash +ptwm compress model.safetensors --container-layout legacy +``` + +See the [release notes](/release-notes) for the details. + ## Delta compression Compress relative to a reference file. The output omits the reference — diff --git a/docs/site/content/30.guides/02.python-api.md b/docs/site/content/30.guides/02.python-api.md index 750c7d3..3552c70 100644 --- a/docs/site/content/30.guides/02.python-api.md +++ b/docs/site/content/30.guides/02.python-api.md @@ -65,9 +65,30 @@ assert torch.equal(tensor, restored) | `threads` | `None` (up to 16 cores) | Maximum threads for parallel encode | | `compression_threshold` | `0.95` | Skip compression if ratio exceeds this threshold | | `compression_chunk` | `262144` | Chunk size in bytes (256 KB) | -| `framing` | `Framing.COMPACT` | `COMPACT` writes a compact single-tensor frame (10 to 130 B of fixed overhead; raw bytes stored when no codec helps). `CONTAINER` writes a full `.ptwm` container, which readers before v1.2 and the random-access tools require. | +| `framing` | `Framing.COMPACT` | `COMPACT` writes a compact single-tensor frame (`\x89PTWS`, about 20 to 70 B of fixed overhead; raw bytes stored when no codec helps) and falls back to a full container for tensors that need shared codec state or chunking. `CONTAINER` always writes a full `.ptwm` container, which the random-access, store and transcode tools require. | +| `container_layout` | `ContainerLayout.COMPACT` | Layout of a full container (`CONTAINER`, or the `COMPACT` fallback). `COMPACT` writes `\x89PTWC` (about 120 to 140 B for one tensor). `LEGACY` writes `\x89PTWM` (about 0.5 to 1 KB for one tensor), the only layout readers before v1.2 open. Strings `"compact"` and `"legacy"` are accepted. | | `emit_payload_hash` | `True` | Store an xxh64 of the raw bytes (8 B per tensor), checked on decode | -| `emit_plane_crc` | `True` | Store a CRC32 per coded plane (4 B per plane), checked on decode | +| `emit_plane_crc` | `True` | Store a CRC32 per coded plane (4 B per plane), checked on decode. The compact frame stores plane CRCs only when `emit_payload_hash` is off, since the hash already covers every plane | + +## Files for readers before v1.2 + +The default outputs, the compact frame and the compact container, need +PTWM v1.2 or later to decode. Older readers reject them with an error; +they do not misread them. For files that v1.1 opens, write a full +container in the legacy layout: + +```python +from ptwm import CompressionConfig, Compressor, ContainerLayout, Framing + +legacy = Compressor(CompressionConfig( + framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY, +)) +``` + +The multi-tensor writers take the same choice as a keyword, +`container_layout="legacy"`. See the +[release notes](/release-notes) for the built-in id caveat. ## Full reference diff --git a/docs/site/content/70.release-notes.md b/docs/site/content/70.release-notes.md new file mode 100644 index 0000000..e0ab7fd --- /dev/null +++ b/docs/site/content/70.release-notes.md @@ -0,0 +1,85 @@ +--- +title: Release notes +description: Changes between PTWM releases that affect file compatibility. +--- + +## v1.2 (unreleased) + +### Default output needs PTWM v1.2 or later + +The default outputs changed format. Readers before v1.2 reject them +with an error ("not a .ptwm container" or "bad magic"); they do not +misread them. + +| Output | v1.1 | v1.2 default | +|---|---|---| +| `Compressor.compress` (one tensor) | legacy container, `\x89PTWM` | compact frame, `\x89PTWS` | +| multi-tensor writers and the CLI | legacy container, `\x89PTWM` | compact container, `\x89PTWC` | + +The multi-tensor writers are `compress_ptwm_blob`, `write_ptwm_shard`, +`compress_safetensors_file`, the LMDB and WebDataset writers, +`_rust.compress_model`, the Rust `compress_model` with +`CompressorOptions::default()`, and `ptwm compress`. + +v1.2 reads every format: `\x89PTWM`, `\x89PTWX`, `\x89PTWC` and +`\x89PTWS`. + +### Writing files for v1.1 readers + +Write a full container in the legacy layout: + +- Python, one tensor: `CompressionConfig(framing=Framing.CONTAINER, + container_layout=ContainerLayout.LEGACY)`. +- Python, many tensors: `container_layout="legacy"` (or + `ContainerLayout.LEGACY`) on the writers listed above. +- CLI: `ptwm compress --container-layout legacy`. +- Rust: `CompressorOptions { container_layout: ContainerLayout::Legacy, + .. }`, or `ContainerWriter::new`, which always writes the legacy + layout. + +A legacy file also records the ids of the built-in codecs and +operations it uses, and these ids are derived from the crate version of +the writer. A v1.1 reader recognises only the ids of version 1.1.0, so +it opens a legacy file only when the writer's crate version is 1.1.0. +v1.2 readers accept the built-in ids of any release. A writer with +another crate version can pin the ids with the Rust +`ContainerWriter::set_builtin_version("1.1.0")`; Python and the CLI +have no such option. + +### Smaller output + +Non-payload bytes (everything except coded data), measured with +`python -m benchmarks.container_overhead`: + +| Case | legacy layout | v1.2 default | +|---|---|---| +| one 4 KB fp32 tensor, `Compressor.compress` | 1,014 B (`\x89PTWM`) | 50 B (`\x89PTWS`) | +| one 4 KB fp32 tensor, `Framing.CONTAINER` | 1,014 B (`\x89PTWM`) | 134 B (`\x89PTWC`) | +| 64 small fp32 tensors, per tensor | 235 B (`\x89PTWM`) | 83 B (`\x89PTWC`) | + +The v1.1.0 release wrote about 2.7 KB of fixed bytes for one tensor, +most of it Extension Table entries for codecs the file did not use. +Legacy containers written by v1.2 hold only the codec entries their +planes reference. + +### Other changes + +- Writers reject duplicate tensor names and more than 32,768 chains, + instead of writing files whose extra entries no reader can reach. +- Readers check every count and length before they allocate, size each + plane from the chain before decoding it, and reject a chain whose + Source declares far more bytes than the tensor holds. +- `Format.BYTE` input: the chain's Source now declares exactly the raw + byte count. v1.1 could declare up to 8 times more for some dtypes; + v1.2 still decodes those files. +- `ptwm._rust.ContainerHandle` opens a container once for many + queries and borrows `bytes` and read-only `mmap` buffers without a + copy. `TensorIndex.names()` returns tensors in write order. +- The transcoder (`explode_ptwm` / `implode_ptwm`) writes a + `"ptwm-manifest-c"` version 2 manifest for compact containers. The + legacy `"ptwm-manifest"` version 1 manifest is unchanged apart from + optional fields, and `implode_ptwm` rejects manifests with unknown + formats, versions or fields. + +The byte layouts are in +[`docs/architecture/wire-format.md`](https://github.com/khwstolle/ptwm/blob/master/docs/architecture/wire-format.md). diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py index e39da58..3d72ce7 100644 --- a/tests/ptwm/benchmarks/test_container_overhead.py +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -1,9 +1,11 @@ -"""Size limits for single-tensor ``Compressor.compress`` output. - -The fixed bytes per output (header, tables, framing) used to be about 2.7 KB, -so small tensors came out larger than their input. These tests hold the -default (compact) framing to the limits below and keep the full container's -Extension Table trimmed. The sweep matches ``benchmarks/container_overhead``. +"""Size limits for PTWM output. + +The fixed bytes per single-tensor output (header, tables, framing) used to be +about 2.7 KB, so small tensors came out larger than their input. These tests +hold the default (compact) framing and the compact container (PTWC) to the +limits below, keep the legacy container's Extension Table trimmed, and bound +the per-tensor overhead of multi-tensor containers. The sweeps match +``benchmarks/container_overhead``. """ from __future__ import annotations @@ -15,10 +17,17 @@ from benchmarks.container_overhead import ( DTYPES, + FORMATS, SIZES, + ShapeFn, + codec_data_bytes, + compress_checked, + distinct_shapes, make_tensor, measure, + measure_multi, non_payload_bytes, + same_shape, tiled_rows, ) @@ -113,3 +122,71 @@ def test_container_framing_floor_stays_trimmed(dname: str, dtype: torch.dtype) - ContainerLayout.LEGACY, ) assert row.non_payload < 1100, row + + +def test_formats_write_the_expected_magic() -> None: + t = make_tensor(torch.float32, "normal", 4096) + labels = [ + measure(t, framing, "fp32", "normal", layout).framing + for framing, layout in FORMATS + ] + assert labels == ["ptws", "ptwc", "ptwm"] + + +@pytest.mark.parametrize(("dname", "dtype"), DTYPES) +@pytest.mark.parametrize("nbytes", [1024, 65536]) +def test_codec_data_is_counted_alike_in_every_format( + dname: str, dtype: torch.dtype, nbytes: int +) -> None: + # The three formats carry the same coded planes for the same tensor, so + # payload plus inline state must come out equal. A count that left inline + # state out for one format would differ here. + t = make_tensor(dtype, "normal", nbytes) + data = { + f"{framing.name}/{layout.name}": codec_data_bytes( + compress_checked(t, framing, layout) + ) + for framing, layout in FORMATS + } + assert len(set(data.values())) == 1, data + + +@pytest.mark.parametrize(("dname", "dtype"), DTYPES) +@pytest.mark.parametrize("nbytes", SIZES) +def test_compact_container_overhead_is_bounded( + dname: str, dtype: torch.dtype, nbytes: int +) -> None: + # A one-tensor PTWC file pays its header, Extension Table, registry, + # index and footer alone. Measured maxima over SIZES: fp32 141 B, + # bf16 117 B. + row = measure( + make_tensor(dtype, "normal", nbytes), + Framing.CONTAINER, + dname, + "normal", + ContainerLayout.COMPACT, + ) + assert row.framing == "ptwc", row + assert row.non_payload <= 160, row + + +@pytest.mark.parametrize( + ("shape_of", "hashes", "limit"), + [(same_shape, True, 95), (distinct_shapes, True, 100), (same_shape, False, 70)], + ids=["same-shape", "distinct-shapes", "no-hash-no-crc"], +) +def test_multi_tensor_overhead_per_tensor( + shape_of: ShapeFn, hashes: bool, limit: int +) -> None: + # 64 fp32 (64,) bias vectors (synthesis cases A and B). Measured per + # tensor: 83.1 B, 82.6 B and 59.1 B, against 235 B, 288 B and 211 B for + # the legacy layout. + compact = measure_multi( + "fp32", "float32", shape_of, 64, ContainerLayout.COMPACT, hashes + ) + legacy = measure_multi( + "fp32", "float32", shape_of, 64, ContainerLayout.LEGACY, hashes + ) + assert (compact.layout, legacy.layout) == ("ptwc", "ptwm") + assert compact.per_tensor <= limit, compact + assert compact.per_tensor < legacy.per_tensor / 2, (compact, legacy) From 6693641517d4ec58a8d10460e8869ce10844e326 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 03:23:01 +0000 Subject: [PATCH 25/34] fix(S10-docs): address review findings - CLI: `ptwm compress model.safetensors --container-layout legacy` now writes every tensor of the per-tensor shell as a full \x89PTWM container (Framing.CONTAINER) instead of a \x89PTWS frame. - Legacy writers stamp built-in ids and labels of LEGACY_BUILTIN_VERSION ("1.1.0") instead of the crate version, so legacy output stays readable by v1.1 after a version bump. The transcoder manifest omits builtin_version when it equals that constant. - benchmarks: codec_data_bytes read shared state from the wrong manifest key and never counted it; tests now cover inline and shared state. - Bound the Extension Table, capability blob and dependency allocations by the input; reword the reader-limits section to match. - Docs: correct the CLI guide, release notes, wire-format.md and 5.container.md (mode a is the safetensors default). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- benchmarks/container_overhead.py | 7 +- crates/ptwm-core/src/container.rs | 53 +++++++++++- crates/ptwm-core/src/extension/builtins.rs | 8 ++ crates/ptwm-core/src/extension/mod.rs | 6 +- crates/ptwm-core/src/extension/table.rs | 54 +++++++++++- crates/ptwm-core/src/tensor_record.rs | 84 ++++++++++++++++++- crates/ptwm-core/src/transcode/manifest.rs | 7 +- crates/ptwm-core/src/transcode/mod.rs | 12 ++- docs/architecture/wire-format.md | 32 +++---- docs/site/content/20.concepts/5.container.md | 8 +- docs/site/content/30.guides/01.cli.md | 10 ++- docs/site/content/70.release-notes.md | 26 +++--- python/ptwm/cli/compress.py | 22 +++-- .../benchmarks/test_container_overhead.py | 76 ++++++++++++++++- tests/ptwm/cli/test_container_layout_flag.py | 56 +++++++++++++ 15 files changed, 394 insertions(+), 67 deletions(-) diff --git a/benchmarks/container_overhead.py b/benchmarks/container_overhead.py index b458a4b..3761159 100644 --- a/benchmarks/container_overhead.py +++ b/benchmarks/container_overhead.py @@ -156,10 +156,9 @@ def codec_data_bytes(blob: bytes) -> int: # Payload plus inline state; the raw size for a stored frame. return _rust.decode_compact(blob)[6] registry_json, tensors, members = _rust.explode_ptwm(blob) - total = sum( - len(members[s["payload"]]) - for s in json.loads(registry_json).get("shared_state", []) - ) + # Both manifest forms keep the shared states under "registry". + shared = json.loads(registry_json)["registry"].get("shared_state", []) + total = sum(len(members[s["payload"]]) for s in shared) for _name, _key, tensor_json in tensors: for plane in json.loads(tensor_json)["planes"]: total += plane["comp_len"] diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 37d58d3..3b064c6 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -157,7 +157,7 @@ pub struct BuiltExtensionTable { /// [`crate::extension::builtin_entries`]. /// /// Built-in entries carry the ids and labels of `builtin_version` -/// ([`crate::extension::CURRENT_BUILTIN_VERSION`] for new files). The +/// ([`crate::extension::LEGACY_BUILTIN_VERSION`] for new legacy files). The /// returned `codec_idx_map` is keyed by the current canonical ids whatever /// the version. /// @@ -336,6 +336,19 @@ impl ContainerLayout { } } +/// Built-in version a new writer of `layout` stamps on built-in entries. +/// +/// The legacy layout uses [`crate::extension::LEGACY_BUILTIN_VERSION`], so +/// a v1.1 reader recognises its built-in ids whatever the crate version. +/// The compact layout stores built-ins without a version; its tables use +/// the crate version internally. +fn default_builtin_version(layout: ContainerLayout) -> String { + match layout { + ContainerLayout::Legacy => crate::extension::LEGACY_BUILTIN_VERSION.to_owned(), + ContainerLayout::Compact => crate::extension::CURRENT_BUILTIN_VERSION.to_owned(), + } +} + /// Serializes tensor records for one container writer's layout and tables. /// /// [`ContainerWriter::record_encoder`] returns it so that callers can @@ -452,6 +465,7 @@ impl ContainerWriter { } let header = Header::new_ptwm(); + let builtin_version = default_builtin_version(layout); Ok(Self { writer, @@ -459,7 +473,7 @@ impl ContainerWriter { compact_index: Vec::new(), encoder: RecordEncoder { layout, - builtin_version: crate::extension::CURRENT_BUILTIN_VERSION.to_owned(), + builtin_version: builtin_version.clone(), codec_ids: Vec::new(), registry_bound: Vec::new(), }, @@ -477,7 +491,7 @@ impl ContainerWriter { finalized: false, prelude_section_len: 0, names: HashSet::new(), - builtin_version: crate::extension::CURRENT_BUILTIN_VERSION.to_owned(), + builtin_version, raw_extension_table: None, table_set: false, }) @@ -668,7 +682,7 @@ impl ContainerWriter { } /// Write built-in Extension Table entries with the ids and labels of - /// `version` instead of the crate version. This applies to the table + /// `version` instead of [`crate::extension::LEGACY_BUILTIN_VERSION`]. This applies to the table /// built by [`set_chain_registry`](Self::set_chain_registry) and /// [`set_chain_registry_with_codecs`](Self::set_chain_registry_with_codecs) /// and to inline-chain tables. The transcoder uses it to reproduce files @@ -2480,6 +2494,37 @@ mod tests { ); } + /// A legacy writer stamps built-ins with the ids and labels of `1.1.0`, + /// which a v1.1 reader recognises, whatever the crate version. The + /// compact layout keeps the crate version internally. + #[test] + fn legacy_writer_stamps_legacy_builtin_version() { + use crate::extension::{BUILTIN_LABEL_PREFIX, BUILTIN_PUBKEY, LEGACY_BUILTIN_VERSION}; + assert_eq!(LEGACY_BUILTIN_VERSION, "1.1.0"); + let mut buf = Cursor::new(Vec::new()); + let mut w = ContainerWriter::new(&mut buf).unwrap(); + assert_eq!(w.builtin_version, "1.1.0"); + assert_eq!(w.encoder.builtin_version, "1.1.0"); + w.set_chain_registry(ChainRegistry { + chains: vec![(0, make_minimal_chain(1))], + }) + .unwrap(); + assert!(!w.extension_table.entries.is_empty()); + for e in &w.extension_table.entries { + let rest = e.human_label.strip_prefix(BUILTIN_LABEL_PREFIX).unwrap(); + let (name, ver) = rest.split_once('@').unwrap(); + assert_eq!(ver, "1.1.0", "{}", e.human_label); + assert_eq!( + e.canonical_id, + CanonicalId::derive(&BUILTIN_PUBKEY, name, "1.1.0") + ); + } + + let mut buf = Cursor::new(Vec::new()); + let w = ContainerWriter::with_layout(&mut buf, ContainerLayout::Compact).unwrap(); + assert_eq!(w.builtin_version, crate::extension::CURRENT_BUILTIN_VERSION); + } + /// Serialized Extension Table of a container written with `registry`. fn table_bytes_for(registry: &ChainRegistry, codecs: CodecSet<'_>) -> Vec { build_extension_table(registry, codecs, crate::extension::CURRENT_BUILTIN_VERSION) diff --git a/crates/ptwm-core/src/extension/builtins.rs b/crates/ptwm-core/src/extension/builtins.rs index e071006..1da2c35 100644 --- a/crates/ptwm-core/src/extension/builtins.rs +++ b/crates/ptwm-core/src/extension/builtins.rs @@ -18,6 +18,14 @@ use crate::transforms::op::OpId; /// Version that the built-in ids and labels of this build derive from. pub const CURRENT_BUILTIN_VERSION: &str = env!("CARGO_PKG_VERSION"); +/// Version whose built-in ids and labels the legacy (`\x89PTWM`) writers +/// stamp by default. It is fixed at the release whose reader the legacy +/// layout exists for: a v1.1 reader recognises only the ids derived from +/// `1.1.0`, so legacy output stays readable there whatever the crate +/// version of the writer. Readers of every release map these ids back to +/// the current ones (see [`normalize_builtin_ids`]). +pub const LEGACY_BUILTIN_VERSION: &str = "1.1.0"; + /// Prefix of every built-in label: `io.ptwm.builtin.@`. pub const BUILTIN_LABEL_PREFIX: &str = "io.ptwm.builtin."; diff --git a/crates/ptwm-core/src/extension/mod.rs b/crates/ptwm-core/src/extension/mod.rs index f7f30d1..860e449 100644 --- a/crates/ptwm-core/src/extension/mod.rs +++ b/crates/ptwm-core/src/extension/mod.rs @@ -17,9 +17,9 @@ pub mod table; pub use builder::ExtensionTableBuilder; pub use builtins::{ - BUILTIN_LABEL_PREFIX, BuiltinRef, BuiltinRegistry, CURRENT_BUILTIN_VERSION, builtin_entries, - builtin_entry, builtin_entry_at_version, builtin_registry, is_builtin, - is_valid_builtin_version, normalize_builtin_ids, + BUILTIN_LABEL_PREFIX, BuiltinRef, BuiltinRegistry, CURRENT_BUILTIN_VERSION, + LEGACY_BUILTIN_VERSION, builtin_entries, builtin_entry, builtin_entry_at_version, + builtin_registry, is_builtin, is_valid_builtin_version, normalize_builtin_ids, }; pub use capability::{Capability, CapabilityMap}; pub use dispatch::{BuiltinKind, dispatch_builtin}; diff --git a/crates/ptwm-core/src/extension/table.rs b/crates/ptwm-core/src/extension/table.rs index 8c22e87..24efee8 100644 --- a/crates/ptwm-core/src/extension/table.rs +++ b/crates/ptwm-core/src/extension/table.rs @@ -74,6 +74,9 @@ pub enum Attestation { PgpSignature(Vec), // raw signature bytes } +/// Entries [`ExtensionTable::read`] reserves before it parses any. +const PREALLOC_ENTRIES: usize = 64; + impl ExtensionTable { pub fn empty() -> Self { Self { @@ -96,7 +99,10 @@ impl ExtensionTable { "extension table count too large: {count} (max {MAX_TABLE_ENTRIES})" ))); } - let mut entries = Vec::with_capacity(count as usize); + // `r` does not expose the bytes that remain, so reserve at most + // `PREALLOC_ENTRIES` up front; the vector grows only as entries + // parse, so a forged count cannot allocate more than the input holds. + let mut entries = Vec::with_capacity((count as usize).min(PREALLOC_ENTRIES)); for _ in 0..count { entries.push(ExtensionTableEntry::read(r)?); } @@ -221,9 +227,18 @@ impl ExtensionTableEntry { "capability blob too large: {cap_len} bytes (max {MAX_CAPABILITIES_LEN})" ))); } - let mut cap_bytes = vec![0u8; cap_len as usize]; - r.read_exact(&mut cap_bytes) + // Read through `take` so the buffer grows with the bytes the input + // holds, not with the declared length. + let mut cap_bytes = Vec::new(); + r.by_ref() + .take(u64::from(cap_len)) + .read_to_end(&mut cap_bytes) .map_err(io_err("read capabilities"))?; + if cap_bytes.len() != cap_len as usize { + return Err(io_err("read capabilities")(std::io::Error::from( + std::io::ErrorKind::UnexpectedEof, + ))); + } let capabilities = CapabilityMap::from_cbor(&cap_bytes)?; let mut at_kind = [0u8; 1]; @@ -459,4 +474,37 @@ mod tests { "oversized cap_len should be rejected, got: {res:?}" ); } + + /// A capabilities length under the limit but past the end of the input + /// fails as a truncated read; the reader does not reserve the declared + /// length up front. + #[test] + fn truncated_capabilities_are_rejected() { + let mut table_buf = 1u32.to_le_bytes().to_vec(); + table_buf.extend_from_slice(&[0xAA; 32]); // canonical_id + table_buf.extend_from_slice(&0u16.to_le_bytes()); // human_label_len + table_buf.extend_from_slice(&1u16.to_le_bytes()); // kind = Transform + table_buf.extend_from_slice(&1u16.to_le_bytes()); // abi_version + table_buf.push(1); // lifecycle = Thread + table_buf.push(FLAVOR_WASM); // flavor_hints + table_buf.extend_from_slice(&MAX_CAPABILITIES_LEN.to_le_bytes()); + table_buf.extend_from_slice(&[0xA0; 7]); // 7 of the declared bytes + match ExtensionTable::from_bytes(&table_buf) { + Err(ExtensionError::TableParse(msg)) => { + assert!(msg.contains("read capabilities"), "{msg}"); + } + other => panic!("expected a truncated read, got {other:?}"), + } + } + + /// A count of `MAX_TABLE_ENTRIES` with no entries behind it fails on the + /// first entry. + #[test] + fn table_count_without_entries_is_rejected() { + let table_buf = MAX_TABLE_ENTRIES.to_le_bytes().to_vec(); + assert!(matches!( + ExtensionTable::from_bytes(&table_buf), + Err(ExtensionError::TableParse(_)) + )); + } } diff --git a/crates/ptwm-core/src/tensor_record.rs b/crates/ptwm-core/src/tensor_record.rs index 1ac5658..80b7bf7 100644 --- a/crates/ptwm-core/src/tensor_record.rs +++ b/crates/ptwm-core/src/tensor_record.rs @@ -9,7 +9,7 @@ use crate::chain::Chain; use crate::chain::wire::{read_chain, write_chain}; use crate::error::PtwmCoreError; use crate::extension::table::ExtensionTable; -use crate::extension::{CURRENT_BUILTIN_VERSION, ExtensionTableBuilder, normalize_builtin_ids}; +use crate::extension::{ExtensionTableBuilder, LEGACY_BUILTIN_VERSION, normalize_builtin_ids}; use crate::plane_record::{PlaneRecord, parse_plane_record, write_plane_record}; pub const TENSOR_FLAG_DEPENDENCIES: u8 = 0x01; @@ -90,12 +90,14 @@ pub(crate) fn checked_len>(n: usize, what: &str) -> Result) -> Result<(), PtwmCoreError> { - write_tensor_record_with_builtin_version(rec, out, CURRENT_BUILTIN_VERSION) + write_tensor_record_with_builtin_version(rec, out, LEGACY_BUILTIN_VERSION) } /// [`write_tensor_record`] with the inline chain's Extension Table written @@ -367,7 +369,9 @@ pub fn parse_tensor_record_with_builtin_version( need(buf, pos, 2, "dep_count")?; let dep_count = u16::from_le_bytes(buf[pos..pos + 2].try_into().unwrap()) as usize; pos += 2; - let mut deps = Vec::with_capacity(dep_count); + // Every dependency takes at least 4 bytes, so the reservation is + // bounded by the bytes that remain, not by the forged count. + let mut deps = Vec::with_capacity(dep_count.min((buf.len() - pos) / 4)); for _ in 0..dep_count { need(buf, pos, 4, "dependency header")?; let ref_kind_raw = buf[pos]; @@ -572,6 +576,78 @@ mod tests { assert_eq!(parsed.terminals.len(), 2); } + /// The inline chain table of a legacy record carries the ids and labels + /// of `1.1.0`, which a v1.1 reader recognises, whatever the crate + /// version of the writer. + #[test] + fn inline_chain_builtins_use_legacy_version() { + assert_eq!(LEGACY_BUILTIN_VERSION, "1.1.0"); + let rec = TensorRecord { + chain_ref: 0x8001, + dtype_code: 0x0010, + input_format: 0, + flags: 0, + orig_size: 2048, + name: "w".to_string(), + payload_hash: None, + inline_chain: Some(minimal_chain()), + dependencies: Vec::new(), + tensor_metadata: None, + terminals: vec![identity_plane(), identity_plane()], + }; + let mut out = Vec::new(); + write_tensor_record(&rec, &mut out).unwrap(); + let text = String::from_utf8_lossy(&out); + let labels: Vec<&str> = text + .match_indices("io.ptwm.builtin.") + .map(|(i, _)| &text[i..]) + .collect(); + assert!(!labels.is_empty()); + for label in labels { + let rest = &label["io.ptwm.builtin.".len()..]; + let (_, ver) = rest.split_once('@').unwrap(); + assert!(ver.starts_with("1.1.0"), "{label:?}"); + } + let (_, _, version) = parse_tensor_record_with_builtin_version(&out).unwrap(); + assert_eq!(version.as_deref(), Some("1.1.0")); + } + + /// A forged dependency count with no dependencies behind it fails as a + /// truncated record. + #[test] + fn forged_dependency_count_is_rejected() { + let rec = TensorRecord { + chain_ref: 0x0002, + dtype_code: 0x0002, + input_format: 0, + flags: 0, + orig_size: 4096, + name: "w".to_string(), + payload_hash: None, + inline_chain: None, + dependencies: vec![Dependency { + ref_kind: RefKind::LocalTensor, + expected_hash: None, + ref_bytes: b"x".to_vec(), + }], + tensor_metadata: None, + terminals: vec![identity_plane()], + }; + let mut out = Vec::new(); + write_tensor_record(&rec, &mut out).unwrap(); + // The dependency section starts with its u16 count, followed by the + // dependency's kind byte and flags byte. + let dep = [1u8, 0, RefKind::LocalTensor as u8, 0, 1, 0, b'x']; + let at = out + .windows(dep.len()) + .position(|w| w == dep) + .expect("dependency bytes"); + out[at..at + 2].copy_from_slice(&u16::MAX.to_le_bytes()); + out.truncate(at + 2); + let err = parse_tensor_record(&out).unwrap_err().to_string(); + assert!(err.contains("truncated"), "{err}"); + } + #[test] fn record_with_deps_and_hash_roundtrip() { let mut expected_hash = [0u8; 32]; diff --git a/crates/ptwm-core/src/transcode/manifest.rs b/crates/ptwm-core/src/transcode/manifest.rs index 08b20bc..223931f 100644 --- a/crates/ptwm-core/src/transcode/manifest.rs +++ b/crates/ptwm-core/src/transcode/manifest.rs @@ -43,9 +43,10 @@ pub struct RegistryManifest { pub struct ContainerFlags { pub embedded_wasm: bool, /// Release whose built-in Extension Table ids and labels the container - /// was written with, when it is not the release reading it. `implode` - /// writes built-in entries (file-level and inline-chain tables) with - /// this version. Absent for the crate version. + /// was written with. `implode` writes built-in entries (file-level and + /// inline-chain tables) with this version. Absent for + /// [`crate::extension::LEGACY_BUILTIN_VERSION`] (`1.1.0`), the default of + /// the legacy writer. #[serde(default, skip_serializing_if = "Option::is_none")] pub builtin_version: Option, } diff --git a/crates/ptwm-core/src/transcode/mod.rs b/crates/ptwm-core/src/transcode/mod.rs index 551e368..6f23f8d 100644 --- a/crates/ptwm-core/src/transcode/mod.rs +++ b/crates/ptwm-core/src/transcode/mod.rs @@ -43,7 +43,7 @@ use crate::error::PtwmCoreError; use crate::extension::compact_table::{ CompactEntry, TAG_CODEC, TAG_FULL, TAG_NAMED, TAG_OP, classify, read_compact_table, }; -use crate::extension::{CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionTable}; +use crate::extension::{CanonicalId, ExtensionTable, LEGACY_BUILTIN_VERSION}; use crate::header::FLAG_PTWX; use crate::layout::PlaneLayout; use crate::metadata::{decode_shape_metadata, encode_shape_metadata}; @@ -388,7 +388,7 @@ pub fn explode(blob: &[u8]) -> Result { if record.inline_chain.is_some() { let version = inline_builtin_version .as_deref() - .unwrap_or(CURRENT_BUILTIN_VERSION); + .unwrap_or(LEGACY_BUILTIN_VERSION); let mut reencoded = Vec::new(); write_tensor_record_with_builtin_version(&record, &mut reencoded, version)?; if blob.get(start..start + consumed) != Some(reencoded.as_slice()) { @@ -440,9 +440,7 @@ pub fn explode(blob: &[u8]) -> Result { // tables are always re-encoded, so their version wins; the file-level // table is rebuilt with it when that reproduces the bytes. let builtin_version = inline_builtin_version.or_else(|| reader.builtin_version.clone()); - let version = builtin_version - .as_deref() - .unwrap_or(CURRENT_BUILTIN_VERSION); + let version = builtin_version.as_deref().unwrap_or(LEGACY_BUILTIN_VERSION); // Check that implode can rebuild the Extension Table and the // shared-prelude section byte-exactly. When no table builder reproduces @@ -480,7 +478,7 @@ pub fn explode(blob: &[u8]) -> Result { manifest_version: MANIFEST_VERSION, container_flags: ContainerFlags { embedded_wasm, - builtin_version: builtin_version.filter(|v| v != CURRENT_BUILTIN_VERSION), + builtin_version: builtin_version.filter(|v| v != LEGACY_BUILTIN_VERSION), }, registry_hash_blake3: hex::encode(reader.header.extension_table_hash), registry: Registry { @@ -827,7 +825,7 @@ pub fn implode( .container_flags .builtin_version .as_deref() - .unwrap_or(CURRENT_BUILTIN_VERSION); + .unwrap_or(LEGACY_BUILTIN_VERSION); let prelude_entries = prelude_from_manifest(®istry.registry.shared_state, &mut members)?; let chain_registry = chains_from_manifest(®istry.registry.chains)?; diff --git a/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index 55db4cb..538a774 100644 --- a/docs/architecture/wire-format.md +++ b/docs/architecture/wire-format.md @@ -187,9 +187,13 @@ parser. A built-in entry has the canonical id `blake3(BUILTIN_PUBKEY ‖ 0x00 ‖ name ‖ 0x00 ‖ version)` and the label -`io.ptwm.builtin.@`, where `version` is the crate -version of the writer (`CURRENT_BUILTIN_VERSION`). The id therefore -changes with every release, although the entry means the same thing. +`io.ptwm.builtin.@`. Releases up to v1.1 wrote their own +crate version there, so the id changed with every release although the +entry means the same thing. From v1.2 the legacy writers write +`LEGACY_BUILTIN_VERSION` (`1.1.0`) whatever their crate version, so +v1.1 readers recognise their built-ins. The transcoder can write other +versions with `ContainerWriter::set_builtin_version` to reproduce files +of other releases. From v1.2 the reader handles this with `normalize_builtin_ids`: after the hash check, an entry whose label names a built-in and whose id @@ -670,25 +674,25 @@ To write files that a v1.1 reader opens: `container_layout="legacy"` on `compress_ptwm_blob`, `write_ptwm_shard`, `compress_safetensors_file`, the LMDB and WebDataset writers and `_rust.compress_model`. -- CLI: `ptwm compress --container-layout legacy`. +- CLI: `ptwm compress --container-layout legacy`. It applies to every + output, and for a `.safetensors` file without `--mode` it writes each + tensor as a full `\x89PTWM` container instead of a `\x89PTWS` frame. - Rust: `CompressorOptions { container_layout: ContainerLayout::Legacy, .. }` or `ContainerWriter::new`. -A legacy file carries the built-in ids of the writer's crate version -(see [Built-in entries](#built-in-entries)), and a v1.1 reader -recognises only the ids of version 1.1.0. This tree's crate version is -1.1.0, so its legacy output opens in v1.1. A build with another crate -version writes legacy files that v1.1 cannot open unless the writer -pins the built-in version with `ContainerWriter::set_builtin_version` -(Rust only). +A v1.1 reader recognises only the built-in ids of version 1.1.0. The +legacy writers stamp those ids whatever their own crate version (see +[Built-in entries](#built-in-entries)), so their output opens in v1.1. ## Reader limits The readers treat every file as untrusted input. In both layouts: -- every count is checked against the bytes that remain before a - vector is allocated, and every offset and length is added with - checked arithmetic; +- a vector sized from a count or length read from the file is bounded + by the bytes that remain or by a small constant: an 8-bit or 16-bit + field, or 64 Extension Table entries, after which the table grows + only as entries parse. Every offset and length is added with checked + arithmetic; - the chain's Source must declare exactly `orig_size` bytes in PTWC and PTWS. A legacy Source may declare up to 8 × `orig_size`, so that files written by the v1.1 `Format.BYTE` bug keep decoding; diff --git a/docs/site/content/20.concepts/5.container.md b/docs/site/content/20.concepts/5.container.md index 949bada..a43db15 100644 --- a/docs/site/content/20.concepts/5.container.md +++ b/docs/site/content/20.concepts/5.container.md @@ -103,8 +103,9 @@ config = CompressionConfig( The multi-tensor writers (`compress_ptwm_blob`, `write_ptwm_shard`, `compress_safetensors_file`, the LMDB and WebDataset writers) take `container_layout="legacy"`, and the CLI takes -`ptwm compress --container-layout legacy`. A legacy file opens in v1.1 -only when its built-in ids are those of version 1.1.0; see the +`ptwm compress --container-layout legacy` for every output, including +the per-tensor `.ptwm.safetensors` file. Legacy writers record the +built-in ids of version 1.1.0, which v1.1 readers recognise; see the [release notes](/release-notes) for the details. ## Modes for safetensors input @@ -118,5 +119,4 @@ in one of two shapes; the `mode` argument selects which: | `"b"` (shell) | A `.safetensors` shell wrapping a `.ptwm` blob | You need the result to look like a regular `.safetensors` file to existing tooling. | Both modes round-trip through PTWM and both take `container_layout`. -Mode `b` is what the [safetensors integration](/guides/safetensors) -emits by default. +Mode `a` is the default. diff --git a/docs/site/content/30.guides/01.cli.md b/docs/site/content/30.guides/01.cli.md index c2f4506..65699b5 100644 --- a/docs/site/content/30.guides/01.cli.md +++ b/docs/site/content/30.guides/01.cli.md @@ -40,13 +40,19 @@ ptwm decompress model.safetensors.ptwm ## Container layout New `.ptwm` files use the compact container layout (`\x89PTWC`), which -PTWM v1.2 and later read. For files that v1.1 opens, write the legacy -layout: +PTWM v1.2 and later read. Without `--mode`, a `.safetensors` file is +written as a `.ptwm.safetensors` file that holds one compact frame +(`\x89PTWS`) per tensor, which also needs v1.2. For files that v1.1 +opens, write the legacy layout: ```bash ptwm compress model.safetensors --container-layout legacy ``` +The flag applies to every output of `ptwm compress`. In the +`.ptwm.safetensors` output it makes every tensor a full `\x89PTWM` +container. + See the [release notes](/release-notes) for the details. ## Delta compression diff --git a/docs/site/content/70.release-notes.md b/docs/site/content/70.release-notes.md index e0ab7fd..e3369b6 100644 --- a/docs/site/content/70.release-notes.md +++ b/docs/site/content/70.release-notes.md @@ -14,12 +14,16 @@ misread them. | Output | v1.1 | v1.2 default | |---|---|---| | `Compressor.compress` (one tensor) | legacy container, `\x89PTWM` | compact frame, `\x89PTWS` | -| multi-tensor writers and the CLI | legacy container, `\x89PTWM` | compact container, `\x89PTWC` | +| `ptwm compress model.safetensors` (one blob per tensor in a safetensors file) | legacy container per tensor, `\x89PTWM` | compact frame per tensor, `\x89PTWS` | +| multi-tensor writers, other `ptwm compress` outputs | legacy container, `\x89PTWM` | compact container, `\x89PTWC` | The multi-tensor writers are `compress_ptwm_blob`, `write_ptwm_shard`, `compress_safetensors_file`, the LMDB and WebDataset writers, -`_rust.compress_model`, the Rust `compress_model` with -`CompressorOptions::default()`, and `ptwm compress`. +`_rust.compress_model` and the Rust `compress_model` with +`CompressorOptions::default()`. The other `ptwm compress` outputs are +`--mode a`, `--mode b`, the LMDB and WebDataset stores, +`--file_compression`, `--delta` and `.bin` inputs. A tensor that does +not fit the compact frame is written as a `\x89PTWC` container. v1.2 reads every format: `\x89PTWM`, `\x89PTWX`, `\x89PTWC` and `\x89PTWS`. @@ -32,19 +36,19 @@ Write a full container in the legacy layout: container_layout=ContainerLayout.LEGACY)`. - Python, many tensors: `container_layout="legacy"` (or `ContainerLayout.LEGACY`) on the writers listed above. -- CLI: `ptwm compress --container-layout legacy`. +- CLI: `ptwm compress --container-layout legacy`, with any `--mode`. + For a `.safetensors` file without `--mode`, this also writes every + tensor as a full `\x89PTWM` container instead of a `\x89PTWS` frame. - Rust: `CompressorOptions { container_layout: ContainerLayout::Legacy, .. }`, or `ContainerWriter::new`, which always writes the legacy layout. A legacy file also records the ids of the built-in codecs and -operations it uses, and these ids are derived from the crate version of -the writer. A v1.1 reader recognises only the ids of version 1.1.0, so -it opens a legacy file only when the writer's crate version is 1.1.0. -v1.2 readers accept the built-in ids of any release. A writer with -another crate version can pin the ids with the Rust -`ContainerWriter::set_builtin_version("1.1.0")`; Python and the CLI -have no such option. +operations it uses, and each id is derived from a release version. A +v1.1 reader recognises only the ids of version 1.1.0, so legacy +writers stamp the ids of 1.1.0 (`LEGACY_BUILTIN_VERSION`) whatever +their own crate version. v1.2 readers accept the built-in ids of any +release. ### Smaller output diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index c661065..e0fdb9e 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -27,7 +27,9 @@ _CONTAINER_LAYOUT_HELP = ( "Layout of written .ptwm containers: 'compact' (default; magic " "\\x89PTWC, readable by ptwm v1.2 and later) or 'legacy' (magic " - "\\x89PTWM, readable by v1.1 and later)." + "\\x89PTWM, readable by v1.1 and later). With 'legacy', per-tensor " + "safetensors output also uses full \\x89PTWM containers instead of " + "the v1.2 \\x89PTWS frame." ) @@ -383,10 +385,12 @@ def compress_safetensors_file( device: int | None = None, container_layout: str = "compact", ) -> None: - """Compress a safetensors file. + r"""Compress a safetensors file. - ``container_layout`` applies to tensors that do not fit the compact - single-tensor frame and are written as full containers. + With ``container_layout="compact"`` (the default) each tensor is a + ``\x89PTWS`` frame, or a ``\x89PTWC`` container when it does not fit + the frame. With ``container_layout="legacy"`` each tensor is a + ``\x89PTWM`` container, which v1.1 readers open. """ import torch # noqa: PLC0415 from safetensors import safe_open # noqa: PLC0415 @@ -397,7 +401,7 @@ def compress_safetensors_file( Compressor, Method, ) - from ptwm._config import Format # noqa: PLC0415 + from ptwm._config import Format, Framing, container_layout_name # noqa: PLC0415 from ptwm.utils._safetensors import ( # noqa: PLC0415 COMPRESSED_DTYPE, COMPRESSION_METHOD, @@ -433,6 +437,13 @@ def compress_safetensors_file( ) compressor_cache = {} + # The compact single-tensor frame is a v1.2 format, so legacy output + # writes every tensor as a full container. + framing = ( + Framing.CONTAINER + if container_layout_name(container_layout) == "legacy" + else Framing.COMPACT + ) with safe_open(filename, "pt", "cpu") as f: for name in f.keys(): # noqa: SIM118 @@ -451,6 +462,7 @@ def compress_safetensors_file( codec_menu=codec_menu, codec=codec, device=device, + framing=framing, container_layout=container_layout, ) ) diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py index 3d72ce7..b4f6ddf 100644 --- a/tests/ptwm/benchmarks/test_container_overhead.py +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -10,10 +10,14 @@ from __future__ import annotations +import json + import pytest import torch -from ptwm import CompressionConfig, Compressor, Decompressor, Format +from ptwm import CompressionConfig, Compressor, Decompressor, Format, _rust from ptwm._config import ContainerLayout, Framing +from ptwm._magic import container_kind +from ptwm.codecs import CodecId from benchmarks.container_overhead import ( DTYPES, @@ -139,8 +143,8 @@ def test_codec_data_is_counted_alike_in_every_format( dname: str, dtype: torch.dtype, nbytes: int ) -> None: # The three formats carry the same coded planes for the same tensor, so - # payload plus inline state must come out equal. A count that left inline - # state out for one format would differ here. + # the codec data must come out equal. These inputs carry no codec state; + # the two tests below cover inline and shared state. t = make_tensor(dtype, "normal", nbytes) data = { f"{framing.name}/{layout.name}": codec_data_bytes( @@ -151,6 +155,72 @@ def test_codec_data_is_counted_alike_in_every_format( assert len(set(data.values())) == 1, data +def _payload_only_bytes(blob: bytes) -> int: + """Sum of the plane payloads of a container, without any codec state.""" + _registry, tensors, _members = _rust.explode_ptwm(blob) + return sum( + plane["comp_len"] + for _name, _key, tensor_json in tensors + for plane in json.loads(tensor_json)["planes"] + ) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16]) +def test_inline_state_is_counted_alike_in_every_format(dtype: torch.dtype) -> None: + # ArithmeticO0 stores its frequency table as inline codec state, so the + # count must include it in every format. A count that left inline state + # out of the containers would come out smaller than the frame's. + t = torch.randn(16384, generator=torch.Generator().manual_seed(3)).to(dtype) + data = {} + for framing, layout in FORMATS: + blob = Compressor( + CompressionConfig( + input_format=Format.TORCH, + framing=framing, + container_layout=layout, + codec_menu=[CodecId.ArithmeticO0], + ) + ).compress(t) + back = Decompressor().decompress(blob) + assert torch.equal(back.view(torch.uint8), t.view(torch.uint8)) + data[container_kind(blob)] = codec_data_bytes(blob) + if framing == Framing.CONTAINER: + # The input really carries inline state. + assert codec_data_bytes(blob) > _payload_only_bytes(blob), layout + assert set(data) == {"ptws", "ptwc", "ptwm"}, data + assert len(set(data.values())) == 1, data + + +@pytest.mark.parametrize("layout", [ContainerLayout.LEGACY, ContainerLayout.COMPACT]) +def test_shared_state_is_counted(layout: ContainerLayout) -> None: + # A container with one shared (prelude) state, built through the + # transcoder: its bytes count as codec data in both container layouts. + # A single-tensor frame cannot hold shared state. + t = make_tensor(torch.float32, "normal", 4096) + blob = compress_checked(t, Framing.CONTAINER, layout) + registry_json, tensors, members = _rust.explode_ptwm(blob) + registry = json.loads(registry_json) + state = bytes(range(200)) + registry["registry"]["shared_state"].append( + { + "id": 0, + "codec": int(CodecId.Order1ScaleAC), + "state_format_version": 1, + "applies_to_mask": 1, + "hash_xxh64": f"{_rust.xxhash64(state):016x}", + "name": "s", + "payload": "shared.bin", + } + ) + with_state = _rust.implode_ptwm( + json.dumps(registry), + [tensor_json for _name, _key, tensor_json in tensors], + {**members, "shared.bin": state}, + ) + assert container_kind(with_state) == container_kind(blob) + assert codec_data_bytes(with_state) == codec_data_bytes(blob) + len(state) + + @pytest.mark.parametrize(("dname", "dtype"), DTYPES) @pytest.mark.parametrize("nbytes", SIZES) def test_compact_container_overhead_is_bounded( diff --git a/tests/ptwm/cli/test_container_layout_flag.py b/tests/ptwm/cli/test_container_layout_flag.py index 2836d0c..0184b62 100644 --- a/tests/ptwm/cli/test_container_layout_flag.py +++ b/tests/ptwm/cli/test_container_layout_flag.py @@ -260,3 +260,59 @@ def test_directory_compression_with_several_processes(tmp_path: Path) -> None: blob = (tmp_path / (src.name + ".ptwm")).read_bytes() assert container_kind(blob) == "ptwm" assert Decompressor().decompress(blob) == src.read_bytes() + + +def _builtin_label_versions(blob: bytes) -> set[bytes]: + """Versions of every ``io.ptwm.builtin.@`` label in ``blob``.""" + import re # noqa: PLC0415 + + return set(re.findall(rb"io\.ptwm\.builtin\.[a-z0-9_]+@([0-9.]+)", blob)) + + +@pytest.mark.parametrize("layout", [None, "compact", "legacy"]) +def test_per_tensor_shell_honours_the_layout( + tmp_path: Path, layout: str | None +) -> None: + """The default per-tensor path for a ``.safetensors`` file writes v1.1 + readable ``\\x89PTWM`` containers with ``--container-layout legacy``. It + used to write ``\\x89PTWS`` frames whatever the flag said.""" + from safetensors import safe_open # noqa: PLC0415 + + src = _safetensors_file(tmp_path) + argv = [str(src), "--force"] + if layout is not None: + argv += ["--container-layout", layout] + handle_compress(_parse(argv)) + out = tmp_path / "model.ptwm.safetensors" + whole = out.read_bytes() + + from ptwm.integrations._safetensors import ( # noqa: PLC0415 + decompress_safetensors_tensor, + ) + from safetensors.torch import load_file # noqa: PLC0415 + + want = load_file(str(src)) + kinds = set() + with safe_open(str(out), "pt") as f: + assert set(f.keys()) == set(want) + for name, tensor in want.items(): + stored = f.get_tensor(name) + if stored.dtype != torch.uint8: + # Stored raw: compression did not make it smaller. + assert torch.equal(stored, tensor), name + continue + kinds.add(container_kind(stored.numpy().tobytes())) + got = decompress_safetensors_tensor(stored) + if isinstance(got, torch.Tensor): + got = got.contiguous().view(torch.uint8).numpy().tobytes() + assert bytes(got) == tensor.contiguous().view(torch.uint8).numpy().tobytes() + assert kinds + if layout == "legacy": + assert kinds == {"ptwm"} + assert b"\x89PTWS" not in whole + assert b"\x89PTWC" not in whole + # A v1.1 reader recognises only the built-in ids of 1.1.0. + assert _builtin_label_versions(whole) == {b"1.1.0"} + else: + assert kinds <= {"ptws", "ptwc"} + assert b"\x89PTWM" not in whole From 29285a5ffeeaf7ef42a0dfa066766d454c6fe0c3 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:10:11 +0000 Subject: [PATCH 26/34] fix(core): stop file-declared sizes from aborting the reader zstd_decode_bounded reserved min(plane length, plane bound) bytes before decoding. Both come from the file, so a 53-byte frame with an edited shape asked for 32 GiB and the process aborted in handle_alloc_error. The first reservation is now also capped by what the payload can decode to (32 KiB per compressed byte) and made with try_reserve; the buffer grows only as data decodes. zstd_dict decodes with a dictionary through the same bounded stream reader instead of a bulk buffer of the declared size, and Order1ScaleAC grows its output in chunks with try_reserve. The legacy reader now checks at open that every record name hashes to its index entry and that names are unique. A bit flip in a record name used to rename the tensor silently, giving two tensors one name. The container reader names a \x89PTWS frame instead of reporting it as truncated (it has no EOF sentinel). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/codec_tagged.rs | 113 +++++++++++++++--- .../ptwm-core/src/codecs/order1_scale_ac.rs | 14 ++- crates/ptwm-core/src/codecs/zstd_dict.rs | 78 ++++++++---- crates/ptwm-core/src/container.rs | 111 ++++++++++++++++- tests/ptwm/core/test_compact_frame.py | 88 ++++++++++++++ 5 files changed, 361 insertions(+), 43 deletions(-) diff --git a/crates/ptwm-core/src/codec_tagged.rs b/crates/ptwm-core/src/codec_tagged.rs index 2978411..a764acc 100644 --- a/crates/ptwm-core/src/codec_tagged.rs +++ b/crates/ptwm-core/src/codec_tagged.rs @@ -214,31 +214,54 @@ pub fn zstd_decode(blob: &[u8]) -> Result, PtwmCoreError> { }) } -/// Decode a Zstd blob (one or more frames) to at most `max_len` bytes. +/// Largest number of bytes one compressed Zstd byte can decode to. /// -/// The output buffer starts at `size_hint.min(max_len)` bytes and the -/// decoder stops reading one byte past `max_len`, so a small frame that -/// expands to gigabytes fails after at most `max_len + 1` decoded bytes. -/// `size_hint` is the expected length; a plane that turns out longer is -/// still read, up to the bound. +/// A Zstd block decodes to at most 128 KiB, and the smallest block that +/// produces output (an RLE block: 3-byte header plus 1 byte) takes 4 bytes, +/// so no input decodes to more than 32 KiB per byte. #[cfg(feature = "codec-zstd")] -pub fn zstd_decode_bounded( - blob: &[u8], +const ZSTD_MAX_EXPANSION: usize = 128 * 1024 / 4; + +/// Initial output reservation for a Zstd payload of `blob_len` bytes. +/// +/// `size_hint` and `max_len` come from the file, so neither alone may size +/// an allocation: a 53-byte frame could otherwise ask for 32 GiB before any +/// byte is decoded. The reservation is also capped by what `blob_len` bytes +/// can decode to. The output still grows past it while real data arrives. +#[cfg(feature = "codec-zstd")] +fn zstd_initial_capacity(blob_len: usize, size_hint: usize, max_len: usize) -> usize { + size_hint + .min(max_len) + .min(blob_len.saturating_mul(ZSTD_MAX_EXPANSION)) +} + +/// Read at most `max_len` decoded bytes from `decoder` into a new buffer. +/// +/// The buffer starts at [`zstd_initial_capacity`] bytes, reserved with +/// `try_reserve` so a failed allocation is an error, not an abort. The read +/// stops one byte past `max_len`, so a small frame that expands to +/// gigabytes fails after at most `max_len + 1` decoded bytes. +#[cfg(feature = "codec-zstd")] +pub(crate) fn zstd_read_bounded( + mut decoder: R, + codec: &'static str, + blob_len: usize, size_hint: usize, max_len: usize, ) -> Result, PtwmCoreError> { use std::io::Read; - if blob.is_empty() { - return Ok(Vec::new()); - } let err = |e: std::io::Error| PtwmCoreError::CodecDecode { - codec: "zstd", + codec, msg: e.to_string(), }; - let mut decoder = zstd::stream::read::Decoder::with_buffer(blob).map_err(err)?; - let first = size_hint.min(max_len); - let mut out = Vec::with_capacity(first); - // Read the expected length into the exact-size buffer first, then probe + let first = zstd_initial_capacity(blob_len, size_hint, max_len); + let mut out = Vec::new(); + out.try_reserve_exact(first) + .map_err(|e| PtwmCoreError::CodecDecode { + codec, + msg: format!("cannot reserve {first} bytes for the decoded plane: {e}"), + })?; + // Read the expected length into the reserved buffer first, then probe // for more. In the usual case the probe reads nothing and the buffer is // never grown. (&mut decoder) @@ -254,13 +277,35 @@ pub fn zstd_decode_bounded( } if out.len() > max_len { return Err(PtwmCoreError::CodecDecode { - codec: "zstd", + codec, msg: format!("frame decodes to more than {max_len} bytes"), }); } Ok(out) } +/// Decode a Zstd blob (one or more frames) to at most `max_len` bytes. +/// +/// `size_hint` is the expected length. It only sizes the first reservation, +/// together with the payload length (see [`zstd_read_bounded`]); a plane +/// that turns out longer is still read, up to the bound. +#[cfg(feature = "codec-zstd")] +pub fn zstd_decode_bounded( + blob: &[u8], + size_hint: usize, + max_len: usize, +) -> Result, PtwmCoreError> { + if blob.is_empty() { + return Ok(Vec::new()); + } + let decoder = + zstd::stream::read::Decoder::with_buffer(blob).map_err(|e| PtwmCoreError::CodecDecode { + codec: "zstd", + msg: e.to_string(), + })?; + zstd_read_bounded(decoder, "zstd", blob.len(), size_hint, max_len) +} + #[cfg(test)] mod tests { use super::*; @@ -397,4 +442,38 @@ mod tests { let blob = zstd_encode(&[], 3).unwrap(); assert!(zstd_decode(&blob).unwrap().is_empty()); } + + /// A file-supplied size hint must not size the first allocation. A + /// 1 TiB hint on a small frame used to abort the process in + /// `Vec::with_capacity`. + #[cfg(feature = "codec-zstd")] + #[test] + fn zstd_bounded_ignores_huge_size_hint() { + let data = vec![0u8; 4096]; + let blob = zstd_encode(&data, 3).unwrap(); + let back = zstd_decode_bounded(&blob, 1 << 40, 1 << 41).unwrap(); + assert_eq!(back, data); + assert!(back.capacity() <= blob.len() * ZSTD_MAX_EXPANSION); + } + + #[cfg(feature = "codec-zstd")] + #[test] + fn zstd_bounded_rejects_output_past_max_len() { + let data = vec![7u8; 10_000]; + let blob = zstd_encode(&data, 3).unwrap(); + let err = zstd_decode_bounded(&blob, 1 << 40, 9_999).unwrap_err(); + assert!(err.to_string().contains("more than 9999 bytes"), "{err}"); + assert_eq!(zstd_decode_bounded(&blob, 0, 10_000).unwrap(), data); + } + + #[cfg(feature = "codec-zstd")] + #[test] + fn zstd_initial_capacity_is_capped_by_payload() { + assert_eq!( + zstd_initial_capacity(10, 1 << 40, 1 << 41), + 10 * ZSTD_MAX_EXPANSION + ); + assert_eq!(zstd_initial_capacity(1 << 20, 4096, 1 << 41), 4096); + assert_eq!(zstd_initial_capacity(1 << 20, 1 << 30, 100), 100); + } } diff --git a/crates/ptwm-core/src/codecs/order1_scale_ac.rs b/crates/ptwm-core/src/codecs/order1_scale_ac.rs index c8561ae..ffb09ee 100644 --- a/crates/ptwm-core/src/codecs/order1_scale_ac.rs +++ b/crates/ptwm-core/src/codecs/order1_scale_ac.rs @@ -620,10 +620,22 @@ impl PlaneCodec for Order1ScaleAC { let row_len = resolve_row_len_decode(decoded_len, layout)?; let mut dec = crate::range_coder::RangeDecoder::new(payload); - let mut out = Vec::with_capacity(decoded_len); + // `decoded_len` comes from the file, so it does not size one up-front + // allocation. The buffer grows in chunks with `try_reserve`, and a + // failed allocation is an error instead of an abort. + const RESERVE_CHUNK: usize = 1 << 20; + let mut out: Vec = Vec::new(); let mut row_pos = 0usize; let mut prev: Option = None; for _ in 0..decoded_len { + if out.len() == out.capacity() { + let want = (decoded_len - out.len()).min(out.len().max(RESERVE_CHUNK)); + out.try_reserve_exact(want) + .map_err(|e| PtwmCoreError::CodecDecode { + codec: "Order1ScaleAC", + msg: format!("cannot reserve {want} bytes for the decoded plane: {e}"), + })?; + } if row_pos == 0 { prev = None; } diff --git a/crates/ptwm-core/src/codecs/zstd_dict.rs b/crates/ptwm-core/src/codecs/zstd_dict.rs index f963139..41e9b21 100644 --- a/crates/ptwm-core/src/codecs/zstd_dict.rs +++ b/crates/ptwm-core/src/codecs/zstd_dict.rs @@ -148,12 +148,7 @@ impl PlaneCodec for ZstdDict { if state_bytes.is_empty() { return crate::codec_tagged::zstd_decode_bounded(payload, decoded_len, max_len); } - decode_with_dict(payload, state_bytes, decoded_len, max_len).map_err(|e| { - PtwmCoreError::CodecDecode { - codec: "zstd_dict", - msg: e.to_string(), - } - }) + decode_with_dict(payload, state_bytes, decoded_len, max_len) } #[cfg(not(feature = "codec-zstd"))] { @@ -185,29 +180,32 @@ fn encode_with_dict(plane: &[u8], dict: &[u8], level: i32) -> std::io::Result std::io::Result> { - let mut decompressor = zstd::bulk::Decompressor::with_dictionary(dict)?; - // The bulk Decompressor writes into a buffer of exactly the - // capacity we pass in. The dispatcher knows the decoded length - // from the propagated terminal descriptor and forwards it via the - // `decoded_len` argument. When that is zero (older callers that - // don't propagate it), fall back to the frame's pledged source - // size if present, then to an 8x heuristic. The capacity never - // exceeds `max_len`: a frame that needs more fails instead. - let capacity = if decoded_len > 0 { - decoded_len - } else { - zstd::bulk::Decompressor::upper_bound(payload) - .unwrap_or_else(|| payload.len().saturating_mul(8)) - } - .min(max_len); - decompressor.decompress(payload, capacity) +) -> Result, PtwmCoreError> { + let decoder = zstd::stream::read::Decoder::with_dictionary(payload, dict).map_err(|e| { + PtwmCoreError::CodecDecode { + codec: "zstd_dict", + msg: e.to_string(), + } + })?; + crate::codec_tagged::zstd_read_bounded( + decoder, + "zstd_dict", + payload.len(), + decoded_len, + max_len, + ) } #[cfg(all(test, feature = "codec-zstd"))] @@ -234,6 +232,40 @@ mod tests { assert_eq!(dec, data); } + /// A file-supplied plane length must not size the first allocation, + /// with or without a dictionary. + #[test] + fn huge_decoded_len_does_not_preallocate() { + let c = ZstdDict; + let pattern = b"the quick brown fox jumps over the lazy dog "; + let dict: Vec = pattern.repeat(8); + let data: Vec = pattern.repeat(64); + let enc = c.encode(&data, Some(&dict), &PlaneLayout::Flat).unwrap(); + let plain = c.encode(&data, None, &PlaneLayout::Flat).unwrap(); + for (state, payload) in [(&dict[..], &enc.payload), (&[][..], &plain.payload)] { + let version = if state.is_empty() { + 0 + } else { + STATE_FORMAT_VERSION + }; + let dec = c + .decode_bounded( + version, + state, + payload, + &PlaneLayout::Flat, + 1 << 40, + 1 << 41, + ) + .unwrap(); + assert_eq!(dec, data); + let err = c + .decode_bounded(version, state, payload, &PlaneLayout::Flat, 1 << 40, 100) + .unwrap_err(); + assert!(err.to_string().contains("more than 100 bytes"), "{err}"); + } + } + #[test] fn dictionary_round_trips() { // The encoder no longer echoes the dict back through diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 3b064c6..899c7b9 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -1372,6 +1372,15 @@ impl<'a> ContainerReader<'a> { /// Shared implementation for [`open`](Self::open) and /// [`open_partial`](Self::open_partial). fn open_with(buf: &'a [u8], skip_missing: bool) -> Result { + // A compact single-tensor frame has no EOF sentinel. Name it here so + // the sentinel check does not report a valid frame as truncated. + if buf.starts_with(&crate::compact::MAGIC_COMPACT) { + return Err(PtwmCoreError::InvalidContainer( + "this is a compact single-tensor frame (\\x89PTWS), not a container; \ + decode it with decode_compact" + .into(), + )); + } verify_eof_sentinel(buf)?; if sniff(buf) == Some(ContainerMagic::Ptwc) { return Self::open_compact(buf, skip_missing); @@ -1455,7 +1464,7 @@ impl<'a> ContainerReader<'a> { // point-in-time snapshot of the host's extension directory. let router = crate::flavor::PlaneCodecRouter::new(installed.clone()); - Ok(Self { + let reader = Self { buf, header, extension_table, @@ -1471,7 +1480,40 @@ impl<'a> ContainerReader<'a> { layout: ContainerLayout::Legacy, codec_ids: Vec::new(), registry_bound: Vec::new(), - }) + }; + reader.check_legacy_record_names()?; + Ok(reader) + } + + /// Check that every legacy record's name hashes to its index entry's + /// `name_hash` and that no two records share a name. + /// + /// The legacy index stores only the hash, and nothing else covers the + /// name bytes, so without this check a bit flip in a record name renames + /// the tensor: two records then report the same name and the original + /// name is not found. The compact index stores the names and rejects + /// duplicates in [`crate::ptwc::read_index`]. + fn check_legacy_record_names(&self) -> Result<(), PtwmCoreError> { + use crate::tensor_record::peek_tensor_record_name; + + let mut names: HashSet<&str> = HashSet::with_capacity(self.index.len()); + for e in &self.index { + let name = peek_tensor_record_name(self.entry_bytes(e)?)?; + if name_hash(name) != e.name_hash { + return Err(PtwmCoreError::InvalidContainer(format!( + "tensor record name {name:?} does not match its index entry \ + (name hash {:#018x}, index has {:#018x})", + name_hash(name), + e.name_hash + ))); + } + if !names.insert(name) { + return Err(PtwmCoreError::InvalidContainer(format!( + "duplicate tensor name {name:?}" + ))); + } + } + Ok(()) } /// [`open_with`](Self::open_with) for the compact layout. @@ -2645,6 +2687,71 @@ mod tests { w.finalize().unwrap(); } + /// A compact frame passed to the container reader is named as such, not + /// reported as a truncated file (it has no EOF sentinel). + #[test] + fn container_open_names_a_compact_frame() { + let mut frame = crate::compact::MAGIC_COMPACT.to_vec(); + frame.extend_from_slice(&[0u8; 40]); + for open in [ContainerReader::open, ContainerReader::open_partial] { + let Err(err) = open(&frame) else { + panic!("compact frame opened as a container") + }; + let err = err.to_string(); + assert!(err.contains("compact single-tensor frame"), "{err}"); + assert!(!err.contains("truncated"), "{err}"); + } + } + + /// A legacy record whose name does not hash to its index entry, or two + /// records with one name, are rejected when the container opens. A bit + /// flip in a record name once renamed the tensor silently. + #[test] + fn legacy_open_rejects_renamed_and_duplicate_records() { + let mut buf = Cursor::new(Vec::new()); + { + let mut w = ContainerWriter::new(&mut buf).unwrap(); + for name in ["layer.0.weight", "layer.1.weight"] { + w.append_tensor(&sample_tensor_record(name)).unwrap(); + } + w.finalize().unwrap(); + } + let good = buf.into_inner(); + ContainerReader::open(&good).unwrap(); + let find = |hay: &[u8], needle: &[u8]| { + let at: Vec = hay + .windows(needle.len()) + .enumerate() + .filter(|(_, w)| *w == needle) + .map(|(i, _)| i) + .collect(); + assert_eq!(at.len(), 1, "{needle:?} must occur once"); + at[0] + }; + + // One bit flip turns the '0' of "layer.0.weight" into '1'. + let mut renamed = good.clone(); + let at = find(&renamed, b"layer.0.weight"); + renamed[at + 6] ^= 1; + for open in [ContainerReader::open, ContainerReader::open_partial] { + let Err(err) = open(&renamed) else { + panic!("renamed record accepted") + }; + let err = err.to_string(); + assert!(err.contains("does not match its index entry"), "{err}"); + } + + // The same rename with the index hash patched to match. + let mut dup = renamed.clone(); + let at = find(&dup, &name_hash("layer.0.weight").to_le_bytes()); + dup[at..at + 8].copy_from_slice(&name_hash("layer.1.weight").to_le_bytes()); + let Err(err) = ContainerReader::open(&dup) else { + panic!("duplicate name accepted") + }; + let err = err.to_string(); + assert!(err.contains("duplicate tensor name"), "{err}"); + } + /// Two names with the same index hash: lookups compare the record name /// and pick the right entry, whatever their order in the hash run. #[test] diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index ac3ccfb..9dd2e17 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -2,6 +2,10 @@ from __future__ import annotations +import subprocess +import sys +import textwrap + import numpy as np import pytest import torch @@ -209,6 +213,90 @@ def test_corrupted_frame_is_rejected() -> None: Decompressor().decompress(bytes(blob[:-3])) +_SHAPE_BOMB = textwrap.dedent( + """ + import numpy as np + from ptwm import CompressionConfig, Compressor, Decompressor, Format, Method + from ptwm._exceptions import Error + + + def varint(n): + out = bytearray() + while True: + b, n = n & 0x7F, n >> 7 + if n: + out.append(b | 0x80) + else: + out.append(b) + return bytes(out) + + + def read_varint(b, i): + v = shift = 0 + while True: + c = b[i] + i += 1 + v |= (c & 0x7F) << shift + shift += 7 + if c < 0x80: + return v, i + + + cfg = CompressionConfig(input_format=Format.NUMPY, method=Method.ZSTD) + f = Compressor(cfg).compress(np.zeros((64, 64), dtype=np.uint8)) + assert f[:5] == b"\\x89PTWS" + _, i = read_varint(f, 6) # dtype code + i += 1 # input format + ndim, i = read_varint(f, i) + start = i + for _ in range(ndim): + _, i = read_varint(f, i) + # A 32 GiB plane declared by a frame of about 50 bytes. + bomb = f[:start] + varint(1 << 20) + varint(1 << 15) + f[i:] + try: + Decompressor().decompress(bomb) + except Error as e: + print("rejected:", e) + else: + raise SystemExit("decoded a 32 GiB shape from a 4 KiB payload") + """ +) + + +def test_declared_shape_does_not_size_allocation() -> None: + """A frame whose shape claims 32 GiB fails with an error, not an abort. + + The decoder once reserved the declared plane length before decoding, + so this frame killed the process with SIGABRT. It runs in a child + process so a regression fails this test instead of the whole session. + """ + r = subprocess.run( + [sys.executable, "-c", _SHAPE_BOMB], + capture_output=True, + text=True, + timeout=120, + check=False, + ) + assert r.returncode == 0, (r.returncode, r.stdout, r.stderr) + assert "rejected:" in r.stdout + + +def test_container_readers_name_a_compact_frame() -> None: + """Container APIs report a frame as a frame, not as a truncated file.""" + blob = Compressor(CompressionConfig(input_format=Format.NUMPY)).compress( + np.arange(64, dtype=np.float32) + ) + assert blob[:5] == COMPACT + for call in ( + lambda: _rust.decode_model(blob), + lambda: _rust.list_tensor_names(blob), + lambda: _rust.decode_tensor(blob, "_"), + ): + with pytest.raises(Error, match="compact single-tensor frame") as e: + call() + assert "truncated" not in str(e.value) + + def test_unknown_magic_is_rejected() -> None: with pytest.raises(HeaderParseError): Decompressor().decompress(b"\x89PTWQ" + bytes(20)) From 6b1eaadc53727c4cdaef2344a159a04f08f56598 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:10:24 +0000 Subject: [PATCH 27/34] fix(core): keep the v1.1 Extension Table in legacy containers The legacy writer interned only the codecs its planes used. Exploded LMDB and WebDataset stores of such a file carry "ptwm-manifest" v1 manifests whose codec_table_idx values index that short table. The v1.1 transcoder rebuilds the table with all 16 built-in codecs and copies the indices unchanged, so v1.1 decoded those stores with the wrong codecs ("bad huffman tag", rANS frequency errors). compress_model now writes the full v1.1.0 table for the legacy layout and trims it only in the compact layout. The legacy fixtures stay as they are; legacy_trimmed.ptwm and legacy_p4_byte_fp32.ptwm are now built by rewriting writer output with a trimmed table, and a new test checks that legacy writer output equals its own v1.1.0-layout rewrite. A Python test checks every plane's codec_table_idx in exploded legacy manifests against the table the v1.1 transcoder rebuilds. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/compressor.rs | 26 ++++-- crates/ptwm-core/src/container.rs | 11 ++- crates/ptwm-core/tests/legacy_fixtures.rs | 84 +++++++++++++++---- python/ptwm/_config.py | 6 +- .../benchmarks/test_container_overhead.py | 19 +++-- tests/ptwm/cli/test_container_layout_flag.py | 6 +- tests/ptwm/stores/test_compact_layout.py | 80 ++++++++++++++++++ .../ptwm/wire_format/test_extension_table.py | 63 +++++++++++--- 8 files changed, 241 insertions(+), 54 deletions(-) diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 1a2f0e0..35a7504 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -994,16 +994,26 @@ pub fn compress_model( let mut writer = ContainerWriter::with_layout(writer, opts.container_layout)?; writer.set_prelude(&prelude_entries)?; - // Intern only the codecs the planes reference: interning every built-in - // codec costs about 125 B per unused entry, which dominates small files. - let mut used_codecs: Vec = Vec::new(); - for cp in chosen_per_tensor.iter().flatten() { - let id = codec_canonical_id(cp.codec_id); - if !used_codecs.contains(&id) { - used_codecs.push(id); + if compact { + // Intern only the codecs the planes reference: interning every + // built-in codec costs about 125 B per unused entry, which dominates + // small files. + let mut used_codecs: Vec = Vec::new(); + for cp in chosen_per_tensor.iter().flatten() { + let id = codec_canonical_id(cp.codec_id); + if !used_codecs.contains(&id) { + used_codecs.push(id); + } } + writer.set_chain_registry_with_codecs(registry, &used_codecs)?; + } else { + // The legacy layout keeps the v1.1 table: every built-in codec, in + // built-in order. A v1.1 transcoder rebuilds exactly this table when + // it implodes an exploded store, and copies each plane's + // `codec_table_idx` unchanged, so a trimmed table would make those + // indices name the wrong codecs. + writer.set_chain_registry(registry)?; } - writer.set_chain_registry_with_codecs(registry, &used_codecs)?; // Pre-compute the chunked-path payloads in parallel before the sequential // write loop. Otherwise encode_plane_chunked runs inside the writer's diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 899c7b9..00f5811 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -519,8 +519,9 @@ impl ContainerWriter { /// second table-building pass. Any previously set extension table is /// replaced. /// - /// This is the legacy (v1.1 and earlier) table layout. Writers that know - /// the set of codecs in use should call + /// This is the legacy (v1.1 and earlier) table layout, and the one legacy + /// writers use. Compact-layout writers that know the set of codecs in use + /// should call /// [`set_chain_registry_with_codecs`](Self::set_chain_registry_with_codecs) /// instead, which omits the unreferenced codec entries (about 125 B each). /// @@ -543,6 +544,12 @@ impl ContainerWriter { /// Returns `InvalidContainer` if a codec in `codecs` has no built-in /// Extension Table entry. /// + /// Use it for the compact layout. A v1.1 transcoder cannot rebuild a + /// trimmed legacy table: its implode always interns every built-in + /// codec and keeps each plane's `codec_table_idx`, so an exploded store + /// of such a file decodes with the wrong codecs there. Legacy files meant + /// for v1.1 readers use [`set_chain_registry`](Self::set_chain_registry). + /// /// Must be called before any `append_tensor` call. pub fn set_chain_registry_with_codecs( &mut self, diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs index 35fd954..3b06a27 100644 --- a/crates/ptwm-core/tests/legacy_fixtures.rs +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -8,9 +8,9 @@ //! | `legacy_v1_1_0.ptwm` | Multi-tensor container with the v1.1.0 Extension Table layout (every built-in codec interned, `ContainerWriter::set_chain_registry`). Holds bf16, fp16, fp32, fp8 and int8 tensors, two scale tensors that share an `Order1ScaleAC` prelude state, a chunked bf16 tensor, an fp8 tensor with an inline chain, and a tensor with a dependency. Built by [`rewrite_as_v1_1`] from current `compress_model` output; built-in ids and labels are derived for version `1.1.0`. Byte-identical to what this generator writes when built against the v1.1.0 release, whose `compress_model` emits the same layout. | //! | `legacy_v1_0_0.ptwm` | The v1.1.0 fixture with every built-in id and label re-derived for version `1.0.0` and the Extension Table hash fixed up. It is not v1.0.0 writer output: for these inputs the v1.0.0 `compress_model` writes no shared `Order1ScaleAC` prelude state. It covers v1.0.0 ids on a file that has the v1.1.0 features. | //! | `legacy_v1_0_0_native.ptwm` | Written once by the v1.0.0 release (commit `a17d403`): this generator, built against that tree, run as `rewrite_as_v1_1(&[main, chunked], INLINE_TENSOR)`. Real v1.0.0 writer output, with v1.0.0 built-in ids. `regenerate_fixtures` does not rewrite it. | -//! | `legacy_trimmed.ptwm` | `compress_model` output of the current writer, whose Extension Table holds only the codecs the planes use. Same tensors as the v1.1.0 file minus the two scale tensors and the chunked tensor. | +//! | `legacy_trimmed.ptwm` | `compress_model` output of the current writer rewritten with an Extension Table that holds only the codecs the planes use (`set_chain_registry_with_codecs`, the layout pre-release v1.2 legacy writers used). Same tensors as the v1.1.0 file minus the two scale tensors and the chunked tensor. | //! | `legacy_p4_byte_fp32_v1_1_0.ptwm` | An fp32 tensor from `Format.BYTE` input, with the chain built the way the v1.1 Python path builds it: the Source node declares `[n_bytes]` fp32 elements, four times the real size. v1.1.0 layout (full Extension Table), built-in ids for `1.1.0`. This is the layout of P4 files written by v1.1.x; it is byte-identical to the v1.1.0 `compress_model` output for the same input. | -//! | `legacy_p4_byte_fp32.ptwm` | The same tensor written by the current writer (trimmed Extension Table). | +//! | `legacy_p4_byte_fp32.ptwm` | The same tensor with a trimmed Extension Table, as pre-release v1.2 legacy writers wrote it. | //! //! The tests check that every fixture decodes to the expected bytes and that //! `implode(explode(x)) == x`. The expected bytes are stored as @@ -125,7 +125,7 @@ const EXPECTED: &[(&str, &[(&str, usize, u64)])] = &[ /// blake3 of `compress_model` output for [`tensors`] with the options of /// [`main_options`] (see [`main_container`]). Printed by `regenerate_fixtures`. const LEGACY_WRITER_GOLDEN_BLAKE3: &str = - "b74c3fd1d81d5bbf4703fb55fb6427781d9aa6d2a92a44127cf4dc91ee4b9791"; + "ec513f5c09f76d8f220a69dd3c47b96dec0ca7e3fcd1a075f601d3688b3ffc9a"; // ── Inputs ─────────────────────────────────────────────────────────────────── @@ -485,10 +485,11 @@ fn trimmed_tensors() -> Vec { .collect() } -/// The current legacy writer's output for [`trimmed_tensors`]. +/// The current legacy writer's output for [`trimmed_tensors`], rewritten +/// with a trimmed Extension Table. fn trimmed_container() -> Vec { let data = trimmed_tensors(); - compress(&inputs(&data), main_options(None)) + rewrite_as_trimmed(&[&compress(&inputs(&data), main_options(None))]) } // ── v1.1.0 layout rewrite ──────────────────────────────────────────────────── @@ -515,6 +516,18 @@ fn records_in_file_order(reader: &ContainerReader<'_>) -> Vec { /// entries are taken from the parts unchanged (ids must not clash). The /// record of `inline_tensor` gets its chain stored inline. fn rewrite_as_v1_1(parts: &[&[u8]], inline_tensor: &str) -> Vec { + rewrite_tables(parts, inline_tensor, false) +} + +/// Like [`rewrite_as_v1_1`], but the Extension Table holds only the codecs +/// the planes use (`set_chain_registry_with_codecs`). The legacy writer +/// produced this layout before it went back to the v1.1.0 table; readers +/// still accept it. +fn rewrite_as_trimmed(parts: &[&[u8]]) -> Vec { + rewrite_tables(parts, "", true) +} + +fn rewrite_tables(parts: &[&[u8]], inline_tensor: &str, trimmed: bool) -> Vec { let readers: Vec> = parts .iter() .map(|b| ContainerReader::open(b).expect("open part")) @@ -554,19 +567,34 @@ fn rewrite_as_v1_1(parts: &[&[u8]], inline_tensor: &str) -> Vec { } } + let registry = ChainRegistry { + chains: chains + .into_iter() + .enumerate() + .map(|(i, c)| (i as u16, c)) + .collect(), + }; + let mut used_codecs: Vec = Vec::new(); + for (rec, part) in &records { + let table = &readers[*part].extension_table; + for plane in &rec.terminals { + let id = table.entries[plane.codec_table_idx as usize].canonical_id; + if !used_codecs.contains(&id) { + used_codecs.push(id); + } + } + } let mut cursor = Cursor::new(Vec::new()); { let mut writer = ContainerWriter::new(&mut cursor).unwrap(); writer.set_prelude(&prelude).unwrap(); - writer - .set_chain_registry(ChainRegistry { - chains: chains - .into_iter() - .enumerate() - .map(|(i, c)| (i as u16, c)) - .collect(), - }) - .unwrap(); + if trimmed { + writer + .set_chain_registry_with_codecs(registry, &used_codecs) + .unwrap(); + } else { + writer.set_chain_registry(registry).unwrap(); + } for (mut rec, part) in records { let table = &readers[part].extension_table; for plane in &mut rec.terminals { @@ -972,7 +1000,7 @@ fn rederive_builtins_rejects_length_change() { rederive_builtins(&mut blob, "1.1.0", "1.10.0"); } -/// Current writer output with a trimmed table and shared `Order1ScaleAC` +/// Current writer output with the v1.1.0 table layout and shared `Order1ScaleAC` /// prelude state (the default for models with scale planes) decodes to the /// inputs and transcodes byte-exactly. No fixture file: the bytes come from /// [`main_container`], whose hash is pinned by @@ -982,7 +1010,7 @@ fn current_writer_shared_state_transcode_round_trip() { let blob = main_container(); let reader = ContainerReader::open(&blob).unwrap(); assert!(!reader.prelude.is_empty(), "shared prelude state"); - assert!(plane_codec_count(&reader) < 16, "trimmed table"); + assert_eq!(plane_codec_count(&reader), 16, "v1.1.0 table layout"); assert!( records_in_file_order(&reader) .iter() @@ -1013,10 +1041,30 @@ fn legacy_writer_output_is_golden() { ); assert!( trimmed_container() == load(TRIMMED), - "compress_model output != {TRIMMED}" + "trimmed rewrite of compress_model output != {TRIMMED}" + ); + assert!( + rewrite_as_trimmed(&[&p4_container()]) == load(P4), + "trimmed rewrite of the P4 container != {P4}" ); } +/// The legacy writer interns every built-in codec in built-in order, as +/// v1.1.0 did. A v1.1 transcoder rebuilds this table when it implodes an +/// exploded store and keeps each plane's `codec_table_idx`; a trimmed table +/// would make those indices name the wrong codecs. +#[test] +fn legacy_writer_keeps_v1_1_0_table_layout() { + for blob in [main_container(), p4_container()] { + let reader = ContainerReader::open(&blob).unwrap(); + assert_eq!(plane_codec_count(&reader), 16); + assert!( + rewrite_as_v1_1(&[&blob], "") == blob, + "legacy writer output differs from its v1.1.0-layout rewrite" + ); + } +} + /// The v1.1.0 table layout (`set_chain_registry`) must keep producing the /// same bytes. #[test] @@ -1038,7 +1086,7 @@ fn regenerate_fixtures() { assert!(rederive_builtins(&mut v1_0_0, "1.1.0", "1.0.0") > 0); let trimmed = trimmed_container(); let golden = main_container(); - let p4 = p4_container(); + let p4 = rewrite_as_trimmed(&[&p4_container()]); let p4_v1_1_0 = p4_v1_1_0_container(); let dir = fixture_dir(); diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 44cc6ce..85b4edf 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -85,9 +85,9 @@ class ContainerLayout(Enum): built-in extensions by id, integers as varints and tensor names only in the index; it needs v1.2 or later to read. ``LEGACY`` (magic ``\x89PTWM``) is the fixed-width layout that v1.1 and later read. Its - output is byte-identical to the current pre-PTWC writer, whose Extension - Table holds only the referenced codecs; the v1.1.0 writer interned all 16 - built-in codecs, so its bytes differ. + Extension Table interns all 16 built-in codecs in built-in order, as the + v1.1.0 writer does, so a v1.1 transcoder rebuilds exploded LMDB and + WebDataset stores of it correctly. Strings coerce case-insensitively: ``ContainerLayout("compact")`` and ``ContainerLayout("LEGACY")`` work, and the bindings take the lower-case diff --git a/tests/ptwm/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py index b4f6ddf..d38246b 100644 --- a/tests/ptwm/benchmarks/test_container_overhead.py +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -3,8 +3,9 @@ The fixed bytes per single-tensor output (header, tables, framing) used to be about 2.7 KB, so small tensors came out larger than their input. These tests hold the default (compact) framing and the compact container (PTWC) to the -limits below, keep the legacy container's Extension Table trimmed, and bound -the per-tensor overhead of multi-tensor containers. The sweeps match +limits below, check that the legacy container keeps the v1.1.0 Extension +Table, and bound the per-tensor overhead of multi-tensor containers. The +sweeps match ``benchmarks/container_overhead``. """ @@ -114,10 +115,11 @@ def test_tiled_rows_ratio_is_at_least_60() -> None: @pytest.mark.parametrize(("dname", "dtype"), DTYPES) -def test_container_framing_floor_stays_trimmed(dname: str, dtype: torch.dtype) -> None: - # v1.1 wrote about 2.5 KB of non-payload bytes per single-tensor - # container, 2.1 KB of it unreferenced codec entries. The bound is for - # the legacy layout, so the layout is pinned. +def test_legacy_container_keeps_the_v1_1_table(dname: str, dtype: torch.dtype) -> None: + # The legacy layout interns all 16 built-in codecs (about 2.1 KB), as + # v1.1.0 did: the v1.1 transcoder rebuilds that table for exploded + # stores. A trimmed legacy table (about 1 KB in total) broke them. The + # upper bound keeps the rest of the legacy overhead from growing. row = measure( make_tensor(dtype, "normal", 4096), Framing.CONTAINER, @@ -125,7 +127,8 @@ def test_container_framing_floor_stays_trimmed(dname: str, dtype: torch.dtype) - "normal", ContainerLayout.LEGACY, ) - assert row.non_payload < 1100, row + assert row.framing == "ptwm", row + assert 2000 < row.non_payload < 3200, row def test_formats_write_the_expected_magic() -> None: @@ -249,7 +252,7 @@ def test_multi_tensor_overhead_per_tensor( shape_of: ShapeFn, hashes: bool, limit: int ) -> None: # 64 fp32 (64,) bias vectors (synthesis cases A and B). Measured per - # tensor: 83.1 B, 82.6 B and 59.1 B, against 235 B, 288 B and 211 B for + # tensor: 83.1 B, 82.6 B and 59.1 B, against 262 B, 317 B and 238 B for # the legacy layout. compact = measure_multi( "fp32", "float32", shape_of, 64, ContainerLayout.COMPACT, hashes diff --git a/tests/ptwm/cli/test_container_layout_flag.py b/tests/ptwm/cli/test_container_layout_flag.py index 0184b62..c293158 100644 --- a/tests/ptwm/cli/test_container_layout_flag.py +++ b/tests/ptwm/cli/test_container_layout_flag.py @@ -41,7 +41,9 @@ def _safetensors_file(tmp_path: Path) -> Path: path = tmp_path / "model.safetensors" save_file( { - "a.weight": torch.randn(64, 32, generator=g).to(torch.bfloat16), + # Large enough to compress in the legacy layout too, whose + # Extension Table of all 16 built-in codecs takes about 2 KB. + "a.weight": torch.randn(256, 128, generator=g).to(torch.bfloat16), "b.weight": torch.randn(128, generator=g), }, str(path), @@ -166,7 +168,7 @@ def test_webdataset_store_takes_the_layout(tmp_path: Path, layout: str | None) - with WebDatasetWeightStore.open(out) as store: assert manifest_container_layout(store._registry_json) == (layout or "compact") - want = torch.randn(64, 32, generator=torch.Generator().manual_seed(0)) + want = torch.randn(256, 128, generator=torch.Generator().manual_seed(0)) assert torch.equal(store.get_tensor("a.weight"), want.to(torch.bfloat16)) diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py index 8a0035e..2e1b01b 100644 --- a/tests/ptwm/stores/test_compact_layout.py +++ b/tests/ptwm/stores/test_compact_layout.py @@ -322,3 +322,83 @@ def test_new_writes_default_to_the_compact_layout(tmp_path, store_tensors): write_webdataset(wd, store_tensors) index = json.loads((wd / INDEX_NAME).read_text()) assert index["registry"]["ptwm_format"] == "ptwm-manifest-c" + + +# ── v1.1 transcoder compatibility ───────────────────────────────────────────── + +# Plane codecs in the order the v1.1.0 transcoder interns them after the chain +# ops when it implodes a "ptwm-manifest" v1 manifest +# (ContainerWriter::set_chain_registry). +V1_1_PLANE_CODECS = ( + "identity", + "huffman", + "rans", + "zstd", + "zstd_dict", + "fpc", + "tans", + "order1_scale_ac", + "arithmetic_o0", + "arithmetic_o0_adaptive", + "arithmetic_o1", + "per_group_codebook", + "context_mixing_lite", + "huff_llm_5bit", + "order1_arithmetic", + "neural_predictor", +) + + +def _assert_v1_1_codec_indices(registry: dict, tensor_manifests: list[dict]) -> int: + """Check every plane's ``codec_table_idx`` against the v1.1 table. + + The v1.1 implode rebuilds the Extension Table as the chain ops followed + by every built-in plane codec, and copies ``codec_table_idx`` from the + manifest unchanged. Returns the number of planes checked. + """ + reg = registry["registry"] + n_ops = len(reg["ops"]) + names = {c["id"]: c["codec"] for c in reg["codecs"]} + n = 0 + for tm in tensor_manifests: + for plane in tm["planes"]: + want = n_ops + V1_1_PLANE_CODECS.index(names[plane["codec"]]) + assert plane["codec_table_idx"] == want, (tm["name"], plane) + n += 1 + return n + + +@pytest.mark.parametrize("method_hint", [2, 3, 4]) +def test_legacy_manifest_indices_match_the_v1_1_table(store_tensors, method_hint): + """A v1.1 implode of a legacy manifest resolves each plane to its codec. + + The legacy writer once interned only the codecs its planes used, so the + manifest's indices pointed at other codecs in the table v1.1 rebuilds, + and v1.1 failed to decode LMDB-exploded and WebDataset stores. + """ + blob = compress_ptwm_blob( + store_tensors, method_hint=method_hint, container_layout="legacy" + ) + model = explode_model(blob) + registry = json.loads(model.registry_json) + assert registry["ptwm_format"] == "ptwm-manifest" + manifests = [json.loads(j) for _, _, j in model.tensors] + assert _assert_v1_1_codec_indices(registry, manifests) > 0 + + +def test_legacy_webdataset_indices_match_the_v1_1_table(tmp_path, store_tensors): + wd = tmp_path / "m.wds" + write_webdataset(wd, store_tensors, container_layout="legacy") + for shard in sorted(wd.glob("weights-*-of-*.tar")): + with tarfile.open(shard, "r") as tar: + f = tar.extractfile(SHARDMETA_MEMBER) + assert f is not None + registry = json.loads(f.read()) + manifests = [] + for m in tar.getmembers(): + if m.name.endswith(".json") and m.name != SHARDMETA_MEMBER: + g = tar.extractfile(m) + assert g is not None + manifests.append(json.loads(g.read())) + assert manifests + assert _assert_v1_1_codec_indices(registry, manifests) > 0 diff --git a/tests/ptwm/wire_format/test_extension_table.py b/tests/ptwm/wire_format/test_extension_table.py index 306dacd..679eb63 100644 --- a/tests/ptwm/wire_format/test_extension_table.py +++ b/tests/ptwm/wire_format/test_extension_table.py @@ -51,19 +51,56 @@ def test_extension_table_is_present_for_huffman_chain(): ) -def test_extension_table_holds_only_referenced_codecs(): - """The writer interns the chain ops plus the codecs the planes use. The - v1.1 layout interned all 16 built-in plane codecs, about 2 KB.""" +V1_1_PLANE_CODECS = ( + "identity", + "huffman", + "rans", + "zstd", + "zstd_dict", + "fpc", + "tans", + "order1_scale_ac", + "arithmetic_o0", + "arithmetic_o0_adaptive", + "arithmetic_o1", + "per_group_codebook", + "context_mixing_lite", + "huff_llm_5bit", + "order1_arithmetic", + "neural_predictor", +) + + +def test_legacy_extension_table_holds_every_builtin_codec(): + """The legacy writer interns the chain ops plus all 16 built-in plane + codecs, as v1.1.0 did, about 2 KB. The v1.1 transcoder rebuilds that + table for exploded stores and keeps each plane's codec index, so a + table with only the referenced codecs decoded there with the wrong + codecs. The compact layout keeps only the referenced codecs.""" tensor = torch.full((32,), 0.5, dtype=torch.bfloat16) - blob = Compressor( - CompressionConfig( - input_format=Format.TORCH, - framing=Framing.CONTAINER, - container_layout=ContainerLayout.LEGACY, - ) - ).compress(tensor) + + def compress(layout: ContainerLayout) -> bytes: + return Compressor( + CompressionConfig( + input_format=Format.TORCH, + framing=Framing.CONTAINER, + container_layout=layout, + ) + ).compress(tensor) + + blob = compress(ContainerLayout.LEGACY) assert blob[:5] == b"\x89PTWM" table_off = int.from_bytes(blob[11:19], "little") - count = int.from_bytes(blob[table_off : table_off + 4], "little") - assert count < 16, f"extension table has {count} entries" - assert len(blob) < 1024, f"single-tensor container is {len(blob)} bytes" + table_len = int.from_bytes(blob[19:27], "little") + table = blob[table_off : table_off + table_len] + count = int.from_bytes(table[:4], "little") + assert count > len(V1_1_PLANE_CODECS), f"extension table has {count} entries" + positions = [ + table.index(f"io.ptwm.builtin.{name}@1.1.0".encode()) + for name in V1_1_PLANE_CODECS + ] + assert positions == sorted(positions), "codecs out of built-in order" + + compact = compress(ContainerLayout.COMPACT) + assert compact[:5] == b"\x89PTWC" + assert len(compact) < 1024, f"single-tensor container is {len(compact)} bytes" From df0775335c8a4cc5184197de6e26ca582e24de7d Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:10:41 +0000 Subject: [PATCH 28/34] fix(python): map numpy dtype instances to the right DType DType.from_dtype compared np.dtype instances with `in`, which uses numpy's coercing __eq__: np.dtype("float64") == None and np.dtype("int64") == int are both true. float64 arrays mapped to NONE (code 0) and int64 arrays to INT32, so Compressor(Format.NUMPY) wrote frames that Decompressor could not rebuild. np.dtype inputs are now compared as dtypes, then by name. Compressor.compress raises InvalidDTypeError for a TORCH or NUMPY input whose dtype has no code, instead of writing code 0. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- python/ptwm/core/_compressor.py | 7 ++++++- python/ptwm/utils/_torch.py | 20 ++++++++++++++++++-- tests/ptwm/core/test_compact_frame.py | 21 ++++++++++++++++++++- tests/ptwm/utils/test_torch.py | 27 +++++++++++++++++++++++++++ 4 files changed, 71 insertions(+), 4 deletions(-) diff --git a/python/ptwm/core/_compressor.py b/python/ptwm/core/_compressor.py index a0d3fe1..f267faf 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -5,7 +5,7 @@ from .. import _rust from .._config import CompressionConfig, Format, Framing, Method -from .._exceptions import CompressionMethodNotSupportedError +from .._exceptions import CompressionMethodNotSupportedError, InvalidDTypeError from ..delta import encode as delta_encode from ..preprocessing._chains import ( ClassifierRole, @@ -122,6 +122,11 @@ def compress( shape: tuple[int, ...] | None = None else: dtype_code = DType.from_dtype(data.dtype).code + if dtype_code == DType.NONE.code: + # Code 0 has no dtype mapping, so the frame could not be + # decoded back to a tensor or array. + msg = f"unsupported dtype {data.dtype} for {self.config.input_format}" + raise InvalidDTypeError(msg) shape = tuple(data.shape) ba_bytes = _to_raw_bytes(data) diff --git a/python/ptwm/utils/_torch.py b/python/ptwm/utils/_torch.py index 422b1d8..d94671e 100644 --- a/python/ptwm/utils/_torch.py +++ b/python/ptwm/utils/_torch.py @@ -213,6 +213,17 @@ def _from_dtype_cached(cls, dtype: Any) -> "DType": @classmethod def _from_dtype_uncached(cls, dtype: Any) -> "DType": + if isinstance(dtype, np.dtype): + # Compare numpy dtypes as dtypes. ``np.dtype.__eq__`` coerces the + # other operand, so ``np.dtype("float64") == None`` and + # ``np.dtype("int64") == int`` are both true, and the generic test + # below would return NONE and INT32 for them. + for member in cls: + if member.numpy_dtype is not None and dtype == np.dtype( + member.numpy_dtype + ): + return member + return cls._from_name(dtype) for member in cls: if dtype in ( member.torch_dtype, @@ -221,8 +232,13 @@ def _from_dtype_uncached(cls, dtype: Any) -> "DType": member.dtype_str, ): return member - # Fallback: numpy 2.x / ml_dtypes BF16/FP8/FP4 expose .name matching - # the canonical dtype_str (e.g. np.dtype('bfloat16').name == "bfloat16"). + return cls._from_name(dtype) + + @classmethod + def _from_name(cls, dtype: Any) -> "DType": + """Match ``dtype.name`` against the canonical names, else NONE.""" + # numpy 2.x / ml_dtypes BF16/FP8/FP4 expose .name matching the + # canonical dtype_str (e.g. np.dtype('bfloat16').name == "bfloat16"). name = getattr(dtype, "name", None) if isinstance(name, str): name_lower = name.lower() diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index 9dd2e17..78f97a7 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -19,7 +19,7 @@ Framing, _rust, ) -from ptwm._exceptions import Error, HeaderParseError +from ptwm._exceptions import Error, HeaderParseError, InvalidDTypeError COMPACT = b"\x89PTWS" # A full container has the default (compact) layout unless the config pins @@ -118,6 +118,25 @@ def test_numpy_round_trip() -> None: assert np.array_equal(back, a) +@pytest.mark.parametrize("dtype", [np.float64, np.int64, np.float32, np.int32]) +@pytest.mark.parametrize("framing", [Framing.COMPACT, Framing.CONTAINER]) +def test_numpy_round_trip_keeps_64_bit_dtypes(dtype, framing: Framing) -> None: + """float64 and int64 arrays once stored dtype code 0 and INT32.""" + a = np.arange(6, dtype=dtype).reshape(2, 3) + cfg = CompressionConfig(input_format=Format.NUMPY, framing=framing) + out = Decompressor().decompress(Compressor(cfg).compress(a)) + assert isinstance(out, np.ndarray) + assert out.dtype == a.dtype + np.testing.assert_array_equal(out, a) + + +def test_numpy_dtype_without_code_is_rejected() -> None: + """An array whose dtype has no code fails before anything is written.""" + a = np.zeros(4, dtype="datetime64[s]") + with pytest.raises(InvalidDTypeError, match="unsupported dtype"): + Compressor(CompressionConfig(input_format=Format.NUMPY)).compress(a) + + def test_byte_round_trip() -> None: data = bytes(range(256)) * 9 blob = Compressor(CompressionConfig(bytearray_dtype="float32")).compress(data) diff --git a/tests/ptwm/utils/test_torch.py b/tests/ptwm/utils/test_torch.py index 48f7f32..4bc18f1 100644 --- a/tests/ptwm/utils/test_torch.py +++ b/tests/ptwm/utils/test_torch.py @@ -341,3 +341,30 @@ def __init__(self, name: str) -> None: assert DType.from_dtype(_FakeNpDtype("float8_e5m2")) is DType.FLOAT8_E5M2 # Unknown name still falls through. assert DType.from_dtype(_FakeNpDtype("not_a_dtype")) is DType.NONE + + +@pytest.mark.parametrize( + ("name", "member"), + [ + ("float64", DType.FLOAT64), + ("int64", DType.INT64), + ("float32", DType.FLOAT32), + ("float16", DType.FLOAT16), + ("int32", DType.INT32), + ("int16", DType.INT16), + ("int8", DType.INT8), + ("uint8", DType.UINT8), + ("uint64", DType.UINT64), + ("complex64", DType.COMPLEX64), + ("complex128", DType.COMPLEX128), + ("bool", DType.BOOL), + ], +) +def test_from_dtype_maps_numpy_dtype_instances(name: str, member: DType) -> None: + """``np.dtype`` instances map like their scalar types. + + ``np.dtype("float64") == None`` and ``np.dtype("int64") == int`` are + true in numpy, so these once mapped to NONE and INT32. + """ + assert DType.from_dtype(np.dtype(name)) is member + assert DType.from_dtype(np.dtype(name).type) is member From 7b03138b091914bff4be21c1a443c8fb92e6e4c6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:10:41 +0000 Subject: [PATCH 29/34] fix(cli): record compressed tensors in safetensors files without metadata set_compressed_tensors_metadata wrote the ptwm_compressed_vectors entry only into a non-empty metadata dict. A source saved without metadata (f.metadata() is None) lost the entry, so readers and `ptwm decompress` returned each compressed tensor as its raw uint8 frame. The function now creates the dict when it is None and returns it, and decompress writes no metadata when none is left. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- python/ptwm/cli/compress.py | 2 +- python/ptwm/cli/decompress.py | 2 +- python/ptwm/utils/_safetensors.py | 17 +++++++-- tests/ptwm/integrations/test_safetensors.py | 42 +++++++++++++++++++++ 4 files changed, 57 insertions(+), 6 deletions(-) diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index e0fdb9e..664b37d 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -482,7 +482,7 @@ def compress_safetensors_file( metadata = f.metadata() - set_compressed_tensors_metadata(compressed_tensor_infos, metadata) + metadata = set_compressed_tensors_metadata(compressed_tensor_infos, metadata) save_file(tensors, str(compressed_path), metadata) if delete and not hf_cache: diff --git a/python/ptwm/cli/decompress.py b/python/ptwm/cli/decompress.py index e432303..80c2d87 100644 --- a/python/ptwm/cli/decompress.py +++ b/python/ptwm/cli/decompress.py @@ -195,7 +195,7 @@ def decompress_safetensors_file( if metadata: metadata.pop("ptwm_compressed_vectors", None) - save_file(tensors, str(decompressed_path), metadata) + save_file(tensors, str(decompressed_path), metadata or None) if delete and not hf_cache: full_path.unlink() diff --git a/python/ptwm/utils/_safetensors.py b/python/ptwm/utils/_safetensors.py index 2274ecc..ff9bf60 100644 --- a/python/ptwm/utils/_safetensors.py +++ b/python/ptwm/utils/_safetensors.py @@ -37,11 +37,20 @@ def build_compressed_tensor_info( def set_compressed_tensors_metadata( - compressed_tensor_infos: dict[str, CompressedTensorInfo], metadata: dict[str, str] -) -> None: - """Set file-level metadata on all compressed tensors.""" - if metadata: + compressed_tensor_infos: dict[str, CompressedTensorInfo], + metadata: dict[str, str] | None, +) -> dict[str, str]: + """Record the compressed tensors in the file-level metadata. + + Updates ``metadata`` in place and returns it. ``None`` (a safetensors + file saved without metadata) gives a new dict. Without this entry, + readers return each compressed tensor as its raw uint8 frame. + """ + if metadata is None: + metadata = {} + if compressed_tensor_infos or metadata: metadata[METADATA_KEY] = json.dumps(compressed_tensor_infos) + return metadata def get_compressed_tensors_metadata( diff --git a/tests/ptwm/integrations/test_safetensors.py b/tests/ptwm/integrations/test_safetensors.py index ad347bd..0a4a56c 100644 --- a/tests/ptwm/integrations/test_safetensors.py +++ b/tests/ptwm/integrations/test_safetensors.py @@ -65,3 +65,45 @@ def test_decompress_safetensors_tensor_success() -> None: assert isinstance(decompressed, torch.Tensor) assert decompressed.dtype == torch.float32 assert torch.equal(decompressed, original_tensor) + + +def test_safetensors_without_metadata_round_trips(tmp_path) -> None: + """A source file saved without metadata still records its compressed tensors. + + The ``ptwm_compressed_vectors`` entry was written only into an existing + metadata dict, so readers returned each compressed tensor as raw uint8. + """ + from ptwm.cli.compress import compress_safetensors_file + from ptwm.cli.decompress import decompress_safetensors_file + from ptwm.integrations._safetensors import SafeOpen + from ptwm.utils._safetensors import METADATA_KEY + + g = torch.Generator().manual_seed(0) + tensors = { + "big": torch.randn(256, 256, generator=g).to(torch.bfloat16), + "small": torch.randn(64, 32, generator=g).to(torch.bfloat16), + } + src = tmp_path / "m.safetensors" + save_file(tensors, str(src)) + with safe_open(str(src), framework="pt", device="cpu") as f: + assert f.metadata() is None + + compress_safetensors_file(str(src), force=True) + compressed = tmp_path / "m.ptwm.safetensors" + with safe_open(str(compressed), framework="pt", device="cpu") as f: + meta = f.metadata() + assert meta is not None + assert METADATA_KEY in meta + + with SafeOpen(str(compressed), "pt") as f: + for name, want in tensors.items(): + got = f.get_tensor(name) + assert got.dtype == want.dtype + assert torch.equal(got, want) + + src.unlink() + decompress_safetensors_file(str(compressed), force=True) + with safe_open(str(src), framework="pt", device="cpu") as f: + assert f.metadata() is None + for name, want in tensors.items(): + assert torch.equal(f.get_tensor(name), want) From 4682092baddb37a305b82e28d20ca641bc3ef824 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:10:41 +0000 Subject: [PATCH 30/34] fix(python): coerce CompressionConfig.framing like container_layout framing="compact" (or "COMPACT", or 1) stayed a raw value that never equalled Framing.COMPACT, so the compressor silently wrote a full container, and an unknown name was not rejected. __post_init__ now converts framing to the enum. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- python/ptwm/_config.py | 5 +++++ tests/ptwm/core/test_compact_frame.py | 26 ++++++++++++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 85b4edf..d8d6899 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -190,6 +190,11 @@ def __post_init__(self) -> None: object.__setattr__( self, "container_layout", ContainerLayout(self.container_layout) ) + # The same for framing ("compact" / "container", or the int value). + # Without it, a string never equals Framing.COMPACT and the output is + # silently a full container. + if not isinstance(self.framing, Framing): + object.__setattr__(self, "framing", Framing(self.framing)) @classmethod def from_resolved_policy( diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py index 78f97a7..1c1e5c1 100644 --- a/tests/ptwm/core/test_compact_frame.py +++ b/tests/ptwm/core/test_compact_frame.py @@ -47,6 +47,32 @@ def test_framing_accepts_names() -> None: assert Framing("COMPACT") is Framing.COMPACT +@pytest.mark.parametrize( + ("framing", "magic"), + [ + ("compact", COMPACT), + ("COMPACT", COMPACT), + (1, COMPACT), + ("container", CONTAINER), + (2, CONTAINER), + ], +) +def test_config_coerces_framing(framing, magic: bytes) -> None: + """A framing name or value selects the same output as the enum. + + The config once kept the raw value, which never equalled + ``Framing.COMPACT``, so ``framing="compact"`` wrote a full container. + """ + cfg = _torch_cfg(framing=framing) + assert isinstance(cfg.framing, Framing) + assert Compressor(cfg).compress(torch.randn(64))[:5] == magic + + +def test_config_rejects_unknown_framing() -> None: + with pytest.raises(ValueError, match="bogus"): + _torch_cfg(framing="bogus") + + @pytest.mark.parametrize( ("layout", "magic"), [ From ce312e0d922a0231be1960da17b165479c2eb8f6 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:11:04 +0000 Subject: [PATCH 31/34] docs: correct compatibility errors, CW2 rules and size figures - Compatibility: readers before v1.2 reject a \x89PTWS frame with "not a .ptwm container" (Python Decompressor) or "missing EOF sentinel (file may be truncated)" (Rust and _rust container functions); "bad magic" applies only to \x89PTWC. Same fix in the release notes. - CW2: add the terminal-ref rules (one byte only for node <= 14 and out <= 15, 0xF0-0xFE reserved, no needless escape), the compact role byte table and the 255 bound on edge indices. - Legacy containers keep the v1.1.0 Extension Table; explain why the v1.1 transcoder needs it and update the legacy size figures. - PTWC one-tensor cost is about 85-120 B for bf16 and 125-140 B for fp32; the PTWS default bound for smooth bf16 data is 49 B. - Reader limits: declared plane lengths no longer size allocations. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- docs/architecture/wire-format.md | 124 +++++++++++++------ docs/site/content/20.concepts/5.container.md | 2 +- docs/site/content/30.guides/02.python-api.md | 2 +- docs/site/content/70.release-notes.md | 33 +++-- 4 files changed, 106 insertions(+), 55 deletions(-) diff --git a/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index 538a774..c3319fb 100644 --- a/docs/architecture/wire-format.md +++ b/docs/architecture/wire-format.md @@ -164,19 +164,20 @@ entry × count: ``` The table lists every contribution the file needs at decode time: -first the ops of every registry chain, then the plane codecs that the -plane records reference, in built-in registry order. Chain nodes and -plane records refer to entries by position. A built-in entry costs -about 130 bytes (six entries take 779 bytes in a one-tensor fp32 -file), which is why the compact container stores built-ins by wire id -instead. - -Writers up to the v1.1.0 release interned all 16 built-in plane -codecs, about 2 KB per file, whether or not a plane used them. Readers resolve -`codec_table_idx` by position and check the table hash against the -bytes present, so they read both forms. The transcoder rebuilds the -table from the chain registry and the codec list and keeps whichever -form reproduces the recorded hash. +first the ops of every registry chain, then the plane codecs, in +built-in registry order. Chain nodes and plane records refer to entries +by position. A built-in entry costs about 130 bytes, which is why the +compact container stores built-ins by wire id instead. + +The legacy writers intern all 16 built-in plane codecs, about 2 KB per +file, whether or not a plane uses them. This is the v1.1.0 layout, and +the v1.1 transcoder depends on it (see [Compatibility](#compatibility)). +`ContainerWriter::set_chain_registry_with_codecs` writes a table with +only the referenced codecs. Readers resolve `codec_table_idx` by +position and check the table hash against the bytes present, so they +read both forms. The transcoder rebuilds the table from the chain +registry and the codec list and keeps whichever form reproduces the +recorded hash. The flavor hints are a bitfield rather than an enum because a contribution can ship in several flavours and the reader picks one by @@ -536,10 +537,10 @@ and shared codec state. | Case | PTWC B/tensor | Legacy B/tensor | |---|---|---| -| 64 tensors, same shape, hash and CRC on | 83 | 235 | -| 64 tensors, distinct shapes, hash and CRC on | 83 | 288 | -| 64 tensors, same shape, hash and CRC off | 59 | 211 | -| 1 tensor, hash and CRC on | 165 | 1,176 | +| 64 tensors, same shape, hash and CRC on | 83 | 262 | +| 64 tensors, distinct shapes, hash and CRC on | 83 | 317 | +| 64 tensors, same shape, hash and CRC off | 59 | 238 | +| 1 tensor, hash and CRC on | 165 | 3,029 | Per tensor, the PTWC index entry takes 31 bytes (mostly the name) and the record about 50 bytes with the payload hash and four plane CRCs, @@ -557,11 +558,10 @@ the PTWS frame use it. n_nodes, n_edges, n_terms varint each (n_nodes 1..=255, others <= 255) node: op_ref varint, param_len varint (<= 255), params n_spine varint (<= n_edges; only when n_edges > 0) -edge: src, src_out, dst, dst_in varint; etag u8 (bit 0 role, - bit 1 vendor; other bits zero); [compact role]; +edge: src, src_out, dst, dst_in varint each (<= 255); etag u8 + (bit 0 role, bit 1 vendor; other bits zero); [compact role]; [varint length (1..=65,535) + vendor bytes] -terminal: ref u8 (node << 4 | out, or 0xFF + node u8 + out u8); - compact role +terminal: ref u8; compact role ``` - `op_ref` is an Extension Table index in PTWC and the `OpId` wire id @@ -577,17 +577,36 @@ terminal: ref u8 (node << 4 | out, or 0xFF + node u8 + out u8); that equal them. Real Source params are never empty, so the marker cannot collide with an explicit Source. A record without a shape cannot use a chain with a bound Source. -- A compact role is one byte for every role production chains use - (scale formats, nibble kinds, mantissa and integer bytes up to 8, - `IntN` up to 32 bits, raw, index, exponent byte, fp4 value, XOR - residual), and `0xFF` followed by the legacy role encoding for the - rest. An escaped role that has a one-byte form is rejected. +- A terminal `ref` is one byte, `node << 4 | out`, when `node <= 14` + and `out <= 15`. Otherwise it is `0xFF` followed by `node` u8 and + `out` u8. The reader rejects bytes `0xF0` to `0xFE` (reserved) and an + escape for a pair that fits the one-byte form. +- A compact role is one byte for every role production chains use, and + `0xFF` followed by the legacy role encoding (`Role::write`) for the + rest. The reader rejects every other byte and an escaped role that + has a one-byte form: + + | byte | role | + |---|---| + | `0x00` to `0x04` | `Scale { format }`, format tag 0 E8M0, 1 E4M3, 2 E5M2, 3 F32, 4 F16 | + | `0x08` to `0x0C` | `GlobalScale { format }`, `0x08` + format tag | + | `0x10` to `0x12` | `Nibble { kind }`, kind 0 Exponent, 1 SignMantissa, 2 Value | + | `0x18` | `Value { Fp4E2m1 }` | + | `0x19` | `ExponentByte` | + | `0x1A` | `Raw` | + | `0x1B` | `Index` | + | `0x1C` | `Residual { Xor }` | + | `0x40 + 8 × (of − 1) + index` | `MantissaByte { index, of }`, `index < of <= 8` | + | `0x80 + 8 × (of − 1) + index` | `IntegerByte { index, of }`, `index < of <= 8` | + | `0xC0 + bits − 1` | `Value { IntN { bits } }`, `1 <= bits <= 32` | + | `0xFF` | escape, then the legacy role bytes | ## Compact single-tensor frame (`\x89PTWS`) -`Compressor.compress` writes one tensor. A container's sections cost -about 120 to 140 bytes (PTWC) or 500 to 1,000 bytes (legacy) for one -tensor, which is more than the payload of a small tensor. By default +`Compressor.compress` writes one tensor. For one tensor, a container's +sections cost about 85 to 120 bytes for bf16 and 125 to 140 bytes for +fp32 (PTWC), or 2.5 to 2.9 KB (legacy, which interns every built-in +codec). That is more than the payload of a small tensor. By default the compressor writes a compact frame instead: ``` @@ -639,10 +658,11 @@ chain. The fixed cost is about 20 to 70 bytes, depending on the chain and on which integrity fields are on. With the default settings (payload hash -on) it is at most 43 bytes for a bf16 tensor, 55 bytes for an fp32 -tensor with the split chain, and 61 bytes for an fp32 tensor with the -IntDelta or PredictorXor chain that smooth data selects. Plane CRCs -instead of the hash add up to 8 bytes more for four planes. +on) it is at most 43 bytes for a bf16 tensor and 55 bytes for an fp32 +tensor with the split chain, and 49 bytes for a bf16 tensor and 61 +bytes for an fp32 tensor with the IntDelta or PredictorXor chain that +smooth data selects. Plane CRCs instead of the hash add up to 8 bytes +more for four planes. A tensor whose planes need shared state (a prelude) or chunked encoding cannot use the frame; the compressor writes a full container for it, @@ -656,16 +676,22 @@ instead of the frame, for the random-access, store and transcode tools. | File | Readers before v1.2 | v1.2 and later | |---|---|---| -| `\x89PTWS` frame (default single-tensor output) | reject: "not a .ptwm container" | read | -| `\x89PTWC` container (default multi-tensor output) | reject: bad magic | read | +| `\x89PTWS` frame (default single-tensor output) | reject: Python `Decompressor` "not a .ptwm container"; Rust and `_rust` container functions "missing EOF sentinel (file may be truncated)" | read | +| `\x89PTWC` container (default multi-tensor output) | reject: Python `Decompressor` "not a .ptwm container"; Rust and `_rust` container functions "bad magic" | read | | `\x89PTWM` container | read if its built-in ids are of the reader's version | read, built-in ids of any release | -Old readers fail cleanly on the new formats. A Rust reader before v1.2 -passes the EOF sentinel check and fails in `Header::read` with "bad -magic". A Python reader before v1.2 checks the 5-byte `\x89PTWM` -prefix and reports that the input is not a container. The transcoder -before v1.2 rejects the `"ptwm-manifest-c"` manifest of a PTWC -container. +Old readers fail cleanly on the new formats; none misreads them. The +Python `Decompressor` before v1.2 checks the 5-byte `\x89PTWM` prefix +and reports that the input is not a container. The Rust reader before +v1.2 (`ContainerReader::open`, and through it `_rust.decode_model`, +`_rust.decode_tensor`, `_rust.list_tensor_names` and the transcoder) +checks the EOF sentinel first. A PTWC container ends with the +sentinel, so the reader goes on to `Header::read` and fails with "bad +magic". A PTWS frame has no sentinel, so the reader fails with +"missing EOF sentinel (file may be truncated)", although the file is +intact. From v1.2 the container functions name a PTWS frame as such +and point to `decode_compact`. The transcoder before v1.2 rejects the +`"ptwm-manifest-c"` manifest of a PTWC container. To write files that a v1.1 reader opens: @@ -684,6 +710,17 @@ A v1.1 reader recognises only the built-in ids of version 1.1.0. The legacy writers stamp those ids whatever their own crate version (see [Built-in entries](#built-in-entries)), so their output opens in v1.1. +The legacy writers also keep the v1.1.0 Extension Table layout: the +chain ops, then all 16 built-in plane codecs in built-in order. The +LMDB (`layout="exploded"`) and WebDataset stores hold exploded +`"ptwm-manifest"` version 1 manifests, and a v1.1 reader rebuilds each +tensor with the v1.1 transcoder. That transcoder always interns all 16 +codecs and copies each plane's `codec_table_idx` from the manifest +unchanged, so the indices must refer to that table. A legacy container +whose table holds only the referenced codecs (written with +`ContainerWriter::set_chain_registry_with_codecs`) still opens in v1.1, +but its exploded stores decode there with the wrong codecs. + ## Reader limits The readers treat every file as untrusted input. In both layouts: @@ -700,6 +737,11 @@ The readers treat every file as untrusted input. In both layouts: plane is decoded, and no plane may exceed `16 × orig_size + 4096` bytes (`max_plane_len`), so a corrupt file cannot make a codec allocate far more than the tensor; +- `orig_size` and the shape come from the file too, so the declared + plane length does not size an allocation by itself. The Zstd decoders + reserve at most what the payload can decode to (32 KiB per compressed + byte) and grow the buffer as data arrives; reservations use + `try_reserve`, so a failed allocation is an error, not an abort; - a decoded plane must have exactly its derived length in PTWC and PTWS. In the legacy layout it may differ (older files carry Sources that overstate or understate the raw size) but stays under the plane diff --git a/docs/site/content/20.concepts/5.container.md b/docs/site/content/20.concepts/5.container.md index a43db15..3d95c72 100644 --- a/docs/site/content/20.concepts/5.container.md +++ b/docs/site/content/20.concepts/5.container.md @@ -17,7 +17,7 @@ PTWM writes the container in one of two layouts: Both hold the same content and decode the same way. The compact layout stores built-in codecs and operations by wire id, uses varints, and keeps tensor names in the index. For 64 small fp32 tensors it costs -about 83 bytes per tensor besides the coded data, against 235 bytes for +about 83 bytes per tensor besides the coded data, against 262 bytes for the legacy layout. Choose the legacy layout only for readers older than v1.2 (see [Compatibility](#compatibility)). diff --git a/docs/site/content/30.guides/02.python-api.md b/docs/site/content/30.guides/02.python-api.md index 3552c70..ff1e740 100644 --- a/docs/site/content/30.guides/02.python-api.md +++ b/docs/site/content/30.guides/02.python-api.md @@ -66,7 +66,7 @@ assert torch.equal(tensor, restored) | `compression_threshold` | `0.95` | Skip compression if ratio exceeds this threshold | | `compression_chunk` | `262144` | Chunk size in bytes (256 KB) | | `framing` | `Framing.COMPACT` | `COMPACT` writes a compact single-tensor frame (`\x89PTWS`, about 20 to 70 B of fixed overhead; raw bytes stored when no codec helps) and falls back to a full container for tensors that need shared codec state or chunking. `CONTAINER` always writes a full `.ptwm` container, which the random-access, store and transcode tools require. | -| `container_layout` | `ContainerLayout.COMPACT` | Layout of a full container (`CONTAINER`, or the `COMPACT` fallback). `COMPACT` writes `\x89PTWC` (about 120 to 140 B for one tensor). `LEGACY` writes `\x89PTWM` (about 0.5 to 1 KB for one tensor), the only layout readers before v1.2 open. Strings `"compact"` and `"legacy"` are accepted. | +| `container_layout` | `ContainerLayout.COMPACT` | Layout of a full container (`CONTAINER`, or the `COMPACT` fallback). `COMPACT` writes `\x89PTWC` (about 85 to 120 B for one bf16 tensor, 125 to 140 B for one fp32 tensor). `LEGACY` writes `\x89PTWM` (about 2.5 to 2.9 KB for one tensor, most of it the v1.1 table of all 16 built-in codecs), the only layout readers before v1.2 open. Strings `"compact"` and `"legacy"` are accepted. | | `emit_payload_hash` | `True` | Store an xxh64 of the raw bytes (8 B per tensor), checked on decode | | `emit_plane_crc` | `True` | Store a CRC32 per coded plane (4 B per plane), checked on decode. The compact frame stores plane CRCs only when `emit_payload_hash` is off, since the hash already covers every plane | diff --git a/docs/site/content/70.release-notes.md b/docs/site/content/70.release-notes.md index e3369b6..e7e6030 100644 --- a/docs/site/content/70.release-notes.md +++ b/docs/site/content/70.release-notes.md @@ -8,8 +8,11 @@ description: Changes between PTWM releases that affect file compatibility. ### Default output needs PTWM v1.2 or later The default outputs changed format. Readers before v1.2 reject them -with an error ("not a .ptwm container" or "bad magic"); they do not -misread them. +with an error; they do not misread them. The Python `Decompressor` +reports "not a .ptwm container". The Rust reader and the `_rust` +container functions report "bad magic" for a `\x89PTWC` container and +"missing EOF sentinel (file may be truncated)" for a `\x89PTWS` frame, +although the frame is intact. | Output | v1.1 | v1.2 default | |---|---|---| @@ -57,14 +60,16 @@ Non-payload bytes (everything except coded data), measured with | Case | legacy layout | v1.2 default | |---|---|---| -| one 4 KB fp32 tensor, `Compressor.compress` | 1,014 B (`\x89PTWM`) | 50 B (`\x89PTWS`) | -| one 4 KB fp32 tensor, `Framing.CONTAINER` | 1,014 B (`\x89PTWM`) | 134 B (`\x89PTWC`) | -| 64 small fp32 tensors, per tensor | 235 B (`\x89PTWM`) | 83 B (`\x89PTWC`) | +| one 4 KB fp32 tensor, `Compressor.compress` | 2,867 B (`\x89PTWM`) | 50 B (`\x89PTWS`) | +| one 4 KB fp32 tensor, `Framing.CONTAINER` | 2,867 B (`\x89PTWM`) | 134 B (`\x89PTWC`) | +| 64 small fp32 tensors, per tensor | 262 B (`\x89PTWM`) | 83 B (`\x89PTWC`) | -The v1.1.0 release wrote about 2.7 KB of fixed bytes for one tensor, -most of it Extension Table entries for codecs the file did not use. -Legacy containers written by v1.2 hold only the codec entries their -planes reference. +Most of the legacy bytes for one tensor are Extension Table entries for +all 16 built-in codecs, whether the file uses them or not. Legacy +containers written by v1.2 keep that v1.1.0 table: the v1.1 transcoder +rebuilds it when it reads an exploded LMDB or WebDataset store, and +keeps each plane's codec index from the manifest, so a shorter table +would make v1.1 decode those stores with the wrong codecs. ### Other changes @@ -81,9 +86,13 @@ planes reference. copy. `TensorIndex.names()` returns tensors in write order. - The transcoder (`explode_ptwm` / `implode_ptwm`) writes a `"ptwm-manifest-c"` version 2 manifest for compact containers. The - legacy `"ptwm-manifest"` version 1 manifest is unchanged apart from - optional fields, and `implode_ptwm` rejects manifests with unknown - formats, versions or fields. + legacy `"ptwm-manifest"` version 1 manifest keeps its v1.1 form and + gains two optional fields, `extension_table_b64` and + `builtin_version`, and `implode_ptwm` rejects manifests with unknown + formats, versions or fields. v1.1 ignores the optional fields, so it + rebuilds a manifest correctly only when they are absent. Legacy files + that v1.2 writes produce no such fields; legacy files with another + Extension Table or built-in version may. The byte layouts are in [`docs/architecture/wire-format.md`](https://github.com/khwstolle/ptwm/blob/master/docs/architecture/wire-format.md). From a5a1548953ff4484931b0ad9e277e60e8fecb3bc Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 05:17:18 +0000 Subject: [PATCH 32/34] style(tests): clear clippy lints in test code Drop struct updates that set no field and Copy-type clones, so clippy --all-targets -D warnings passes. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/compressor.rs | 3 --- crates/ptwm-core/src/container.rs | 2 -- crates/ptwm-core/tests/ptwc.rs | 6 +++--- 3 files changed, 3 insertions(+), 8 deletions(-) diff --git a/crates/ptwm-core/src/compressor.rs b/crates/ptwm-core/src/compressor.rs index 35a7504..c8a5735 100644 --- a/crates/ptwm-core/src/compressor.rs +++ b/crates/ptwm-core/src/compressor.rs @@ -1412,7 +1412,6 @@ mod tests { chunk_size: None, // Parses legacy records, so the layout is pinned. container_layout: ContainerLayout::Legacy, - ..Default::default() }, ) .expect("compress succeeds"); @@ -1818,7 +1817,6 @@ mod tests { chunk_size: None, // Parses legacy records, so the layout is pinned. container_layout: ContainerLayout::Legacy, - ..Default::default() }, ) .unwrap(); @@ -1903,7 +1901,6 @@ mod tests { chunk_size: None, // Parses legacy records, so the layout is pinned. container_layout: ContainerLayout::Legacy, - ..Default::default() }, ) .unwrap(); diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 00f5811..6321279 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -3446,7 +3446,6 @@ mod tests { chunk_size: None, // Reads the legacy index and records, so the layout is pinned. container_layout: ContainerLayout::Legacy, - ..Default::default() }, ) .expect("compress succeeds"); @@ -3915,7 +3914,6 @@ label = "test-author" chunk_size, // Parses legacy records, so the layout is pinned. container_layout: ContainerLayout::Legacy, - ..Default::default() }, ) .unwrap(); diff --git a/crates/ptwm-core/tests/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs index 1d543de..7a1f0be 100644 --- a/crates/ptwm-core/tests/ptwc.rs +++ b/crates/ptwm-core/tests/ptwc.rs @@ -1457,18 +1457,18 @@ fn read_extension_table_checks_the_legacy_region_and_hash() { }; // offset + length overflows u64 / usize. - let mut h = header.clone(); + let mut h = header; h.extension_table_offset = u64::MAX - 4; let err = read_extension_table(&patched(&h)).unwrap_err().to_string(); assert!(err.contains("exceeds the buffer"), "{err}"); - let mut h = header.clone(); + let mut h = header; h.extension_table_length = u64::MAX; let err = read_extension_table(&patched(&h)).unwrap_err().to_string(); assert!(err.contains("exceeds the buffer"), "{err}"); // A region that runs past the end of the file. - let mut h = header.clone(); + let mut h = header; h.extension_table_length = legacy.len() as u64; assert!(read_extension_table(&patched(&h)).is_err()); From c84174c0da5a9d67151989a3e7d38da27d00a97b Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 15:36:51 +0000 Subject: [PATCH 33/34] fix(core): bound raw-table import and rANS/tANS raw planes - encode_registry_for_table checks every registry op against the table before encoding and rejects tables beyond the u16 index space. A full imported table that lacked an op used to panic in the table builder. - set_compact_tables uses a checked u16 index. - Rans and Tans decode_bounded check the output length, which also covers the raw-fallback tag that carries no length prefix. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- crates/ptwm-core/src/codecs/rans.rs | 28 +++++++++++++-- crates/ptwm-core/src/codecs/tans.rs | 28 +++++++++++++-- crates/ptwm-core/src/container.rs | 54 ++++++++++++++++++++++++++++- 3 files changed, 103 insertions(+), 7 deletions(-) diff --git a/crates/ptwm-core/src/codecs/rans.rs b/crates/ptwm-core/src/codecs/rans.rs index ab02b1e..210974f 100644 --- a/crates/ptwm-core/src/codecs/rans.rs +++ b/crates/ptwm-core/src/codecs/rans.rs @@ -4,7 +4,7 @@ //! payload. Delegates to the `crate::codec_tagged::rans_{encode,decode}` //! helpers, which already apply the raw-fallback-on-expansion guard. -use crate::codec::{CodecId, Encoded, PlaneCodec}; +use crate::codec::{CodecId, Encoded, PlaneCodec, check_decoded_len}; use crate::codec_tagged::{rans_decode, rans_encode}; use crate::codecs::length_prefix; use crate::error::PtwmCoreError; @@ -152,16 +152,21 @@ impl PlaneCodec for Rans { decoded_len: usize, max_len: usize, ) -> Result, PtwmCoreError> { + // The prefix check rejects an oversized coded plane before it is + // decoded. The check on the output also covers the raw-fallback tag, + // which carries no length prefix. if let Some((&TAG_RANS, rest)) = payload.split_first() { length_prefix::check_within(rest, max_len, "rans")?; } - self.decode( + let out = self.decode( state_format_version, state_bytes, payload, layout, decoded_len, - ) + )?; + check_decoded_len(self.id(), out.len(), max_len)?; + Ok(out) } } @@ -293,4 +298,21 @@ mod tests { .unwrap(); assert!(dec == data); } + + #[test] + fn decode_bounded_rejects_an_oversized_raw_fallback_plane() { + // A plane stored raw (tag 0) carries no length prefix; the bound + // must still apply to its output. + let payload = vec![0u8; 100_001]; + let err = Rans + .decode_bounded(0, &[], &payload, &PlaneLayout::Flat, 4096, 4096) + .unwrap_err() + .to_string(); + assert!(err.contains("bound is 4096"), "{err}"); + // Within the bound the raw plane decodes. + let ok = Rans + .decode_bounded(0, &[], &payload[..4097], &PlaneLayout::Flat, 4096, 4096) + .unwrap(); + assert_eq!(ok.len(), 4096); + } } diff --git a/crates/ptwm-core/src/codecs/tans.rs b/crates/ptwm-core/src/codecs/tans.rs index 2c2246d..ab5071a 100644 --- a/crates/ptwm-core/src/codecs/tans.rs +++ b/crates/ptwm-core/src/codecs/tans.rs @@ -5,7 +5,7 @@ //! `TAG_RAW` / length-prefix framing the rANS codec uses, so a plane that //! tANS can't shrink is stored verbatim and decode stays self-describing. -use crate::codec::{CodecId, Encoded, PlaneCodec}; +use crate::codec::{CodecId, Encoded, PlaneCodec, check_decoded_len}; use crate::codecs::length_prefix; use crate::entropy::tans; use crate::error::PtwmCoreError; @@ -145,16 +145,21 @@ impl PlaneCodec for Tans { decoded_len: usize, max_len: usize, ) -> Result, PtwmCoreError> { + // The prefix check rejects an oversized coded plane before it is + // decoded. The check on the output also covers the raw-fallback tag, + // which carries no length prefix. if let Some((&TAG_TANS, rest)) = payload.split_first() { length_prefix::check_within(rest, max_len, "tans")?; } - self.decode( + let out = self.decode( state_format_version, state_bytes, payload, layout, decoded_len, - ) + )?; + check_decoded_len(self.id(), out.len(), max_len)?; + Ok(out) } } @@ -281,4 +286,21 @@ mod tests { .unwrap(); assert!(dec == data); } + + #[test] + fn decode_bounded_rejects_an_oversized_raw_fallback_plane() { + // A plane stored raw (tag 0) carries no length prefix; the bound + // must still apply to its output. + let payload = vec![0u8; 100_001]; + let err = Tans + .decode_bounded(0, &[], &payload, &PlaneLayout::Flat, 4096, 4096) + .unwrap_err() + .to_string(); + assert!(err.contains("bound is 4096"), "{err}"); + // Within the bound the raw plane decodes. + let ok = Tans + .decode_bounded(0, &[], &payload[..4097], &PlaneLayout::Flat, 4096, 4096) + .unwrap(); + assert_eq!(ok.len(), 4096); + } } diff --git a/crates/ptwm-core/src/container.rs b/crates/ptwm-core/src/container.rs index 6321279..632f9c1 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -258,9 +258,34 @@ pub fn encode_registry_for_table( return Err(bad("embedded WASM entries are not supported".into())); } + // Every entry must be addressable by a u16 table index. + if table.entries.len() > usize::from(u16::MAX) + 1 { + return Err(bad(format!( + "{} entries; a u16 index addresses at most 65536", + table.entries.len() + ))); + } + let mut normalised = table.clone(); normalize_builtin_ids(&mut normalised); let n_entries = normalised.entries.len(); + + // Check every registry op against the table before encoding. Encoding + // interns a missing op, which grows the table and can overflow the u16 + // index space of a full table. + let present: std::collections::HashSet = + normalised.entries.iter().map(|e| e.canonical_id).collect(); + for (_, chain) in ®istry.chains { + for node in &chain.nodes { + if !present.contains(&node.op.canonical_id()) { + return Err(bad(format!( + "chain op {:?} is missing from the table", + node.op + ))); + } + } + } + let mut codec_idx_map = HashMap::new(); for (i, e) in normalised.entries.iter().enumerate() { if let Some(BuiltinKind::Codec(_)) = dispatch_builtin(&e.canonical_id) { @@ -657,7 +682,12 @@ impl ContainerWriter { if let Some(crate::extension::BuiltinKind::Codec(_)) = crate::extension::dispatch_builtin(&e.canonical_id) { - codec_idx_map.entry(e.canonical_id).or_insert(i as u16); + let idx = u16::try_from(i).map_err(|_| { + PtwmCoreError::InvalidContainer(format!( + "extension table entry {i} is past the u16 index range" + )) + })?; + codec_idx_map.entry(e.canonical_id).or_insert(idx); } } self.extension_table = table; @@ -2625,6 +2655,28 @@ mod tests { assert!(encode_registry_for_table(b"\xff\xff\xff\xff", ®istry).is_err()); } + /// A full raw table (65,536 third-party entries) whose registry needs an + /// op the table lacks must fail with an error. Encoding would intern the + /// op past the u16 index space, which used to panic. + #[test] + fn full_raw_extension_table_missing_an_op_is_an_error() { + let registry = ChainRegistry { + chains: vec![(0, v3_passthrough_chain(4))], + }; + let base = sample_extension_table_entry(); + let entries: Vec<_> = (0u32..1 << 16) + .map(|i| { + let mut e = base.clone(); + e.canonical_id = + crate::extension::CanonicalId::derive(&[0x22; 32], &format!("bar{i}"), "0.1.0"); + e + }) + .collect(); + let bytes = ExtensionTable { entries }.to_bytes().unwrap(); + let e = encode_registry_for_table(&bytes, ®istry).unwrap_err(); + assert!(e.to_string().contains("missing from the table"), "{e}"); + } + /// A raw table is written verbatim; the registry indices resolve /// against it, here a table whose ops come after its codecs. #[test] From 17ab3249011fd214eeb2954fd6ba4ae0bd714926 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 28 Oct 2026 15:36:51 +0000 Subject: [PATCH 34/34] fix(cli): guard the multi-shard mode-b output against silent overwrite The overwrite prompt checked only the single-shard file. A multi-shard run replaces the directory without the last suffix, which was deleted with no prompt and no --force. A regular file in that place raised NotADirectoryError after the model was compressed. The guard now covers both destinations, and the replace step removes a file or a directory. Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01H7NivSsfkoNxDZvZv1SnNm --- python/ptwm/cli/compress.py | 29 +++++---- tests/ptwm/cli/test_container_layout_flag.py | 65 ++++++++++++++++++++ 2 files changed, 83 insertions(+), 11 deletions(-) diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index 664b37d..47aaa0a 100644 --- a/python/ptwm/cli/compress.py +++ b/python/ptwm/cli/compress.py @@ -1135,20 +1135,21 @@ def _compress_safetensors_integration(args, mode: str) -> int: # noqa: ANN001 ) out_dir = out_path out_dir.mkdir(parents=True, exist_ok=True) - is_existing_non_empty = any(out_dir.iterdir()) + existing = [out_dir] if any(out_dir.iterdir()) else [] else: out_path = ( Path(args.out) if args.out else input_path.with_suffix(".ptwm.safetensors") ) # The shards are written to a staging directory first and moved - # into place afterwards (see `_finish_mode_b`). + # into place afterwards (see `_finish_mode_b`). One shard becomes + # `out_path`; several become the directory `out_path` without its + # last suffix. Either may be replaced, so the guard covers both. out_dir = out_path.parent / (out_path.stem + "__staging") - is_existing_non_empty = out_path.exists() + existing = [p for p in (out_path, out_path.with_suffix("")) if p.exists()] - if not args.force and is_existing_non_empty: - user_input = ( - input(f"{out_path} already exists; overwrite (y/n)? ").strip().lower() - ) + if not args.force and existing: + shown = " and ".join(str(p) for p in existing) + user_input = input(f"{shown} already exists; overwrite (y/n)? ").strip().lower() if user_input not in ("yes", "y"): return 0 @@ -1198,18 +1199,24 @@ def _finish_mode_b(staging: Path, out_path: Path) -> Path: """ shards = sorted(staging.glob("*.safetensors")) if len(shards) == 1: - if out_path.is_dir(): - shutil.rmtree(out_path) + _remove_path(out_path) shards[0].replace(out_path) shutil.rmtree(staging) return out_path final_dir = out_path.with_suffix("") - if final_dir.exists(): - shutil.rmtree(final_dir) + _remove_path(final_dir) staging.rename(final_dir) return final_dir +def _remove_path(path: Path) -> None: + """Remove ``path``, a file or a directory, if it exists.""" + if path.is_dir() and not path.is_symlink(): + shutil.rmtree(path) + elif path.exists() or path.is_symlink(): + path.unlink() + + def _resolve_compress_mode(args) -> str | None: # noqa: ANN001 """Pick the effective compress mode for a .safetensors input. diff --git a/tests/ptwm/cli/test_container_layout_flag.py b/tests/ptwm/cli/test_container_layout_flag.py index c293158..f5a29a6 100644 --- a/tests/ptwm/cli/test_container_layout_flag.py +++ b/tests/ptwm/cli/test_container_layout_flag.py @@ -318,3 +318,68 @@ def test_per_tensor_shell_honours_the_layout( else: assert kinds <= {"ptws", "ptwc"} assert b"\x89PTWM" not in whole + + +def _fake_two_shard_compress(input_path, out_dir, **_kwargs): # noqa: ANN001, ANN003 + """Stand-in for the integration writer that emits two mode-b shards.""" + out_dir = Path(out_dir) + for i in range(2): + (out_dir / f"model-{i:05d}-of-00002.safetensors").write_bytes(b"shard") + (out_dir / "model.safetensors.index.json").write_text("{}") + + +def test_mode_b_guard_covers_the_multi_shard_directory( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """Without --force, an existing multi-shard output directory is not + replaced unless the user agrees.""" + from ptwm import integrations # noqa: PLC0415 + + src = _safetensors_file(tmp_path) + out = tmp_path / "m.ptwm.safetensors" + previous = tmp_path / "m.ptwm" + previous.mkdir() + (previous / "keep.safetensors").write_bytes(b"old") + + calls: list[str] = [] + monkeypatch.setattr( + integrations, + "compress_safetensors_file", + lambda *a, **k: calls.append("compress") or _fake_two_shard_compress(*a, **k), + ) + prompts: list[str] = [] + monkeypatch.setattr("builtins.input", lambda msg: prompts.append(msg) or "n") + handle_compress(_parse([str(src), "--mode", "b", "--out", str(out)])) + assert len(prompts) == 1 + assert str(previous) in prompts[0] + assert calls == [] + assert (previous / "keep.safetensors").read_bytes() == b"old" + + # Agreeing replaces the directory with the new shards. + monkeypatch.setattr("builtins.input", lambda _msg: "y") + handle_compress(_parse([str(src), "--mode", "b", "--out", str(out)])) + assert calls == ["compress"] + assert not (previous / "keep.safetensors").exists() + assert sorted(p.name for p in previous.iterdir()) == [ + "model-00000-of-00002.safetensors", + "model-00001-of-00002.safetensors", + "model.safetensors.index.json", + ] + + +def test_mode_b_replaces_a_file_in_the_multi_shard_place( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """A regular file where the shard directory goes is replaced, not a + NotADirectoryError after the model was compressed.""" + from ptwm import integrations # noqa: PLC0415 + + src = _safetensors_file(tmp_path) + (tmp_path / "out").write_bytes(b"a file") + monkeypatch.setattr( + integrations, "compress_safetensors_file", _fake_two_shard_compress + ) + argv = [str(src), "--mode", "b", "--out", str(tmp_path / "out.safetensors")] + handle_compress(_parse([*argv, "--force"])) + assert (tmp_path / "out").is_dir() + assert not (tmp_path / "out__staging").exists()