diff --git a/.github/workflows/fuzz.yml b/.github/workflows/fuzz.yml index 28fc5fe..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 11-target sweep runs on the weekly cron and on - # manual dispatch. Per-job whole-minute round-up × 11 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"]') }} + 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/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 64d828c..07d1e42 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) | 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 new file mode 100644 index 0000000..3761159 --- /dev/null +++ b/benchmarks/container_overhead.py @@ -0,0 +1,327 @@ +# ruff: noqa: T201 +r"""Fixed overhead of PTWM output: single-tensor and multi-tensor. + +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 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 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: + 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 + + +@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: + """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 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) + # 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"] + 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 data (see the module docstring).""" + return len(blob) - codec_data_bytes(blob) + + +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, container_layout=layout + ) + ).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, + 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=container_kind(blob) or "?", + dtype=dtype, + kind=kind, + raw=len(raw), + out=len(blob), + non_payload=non_payload_bytes(blob), + zstd19=len(z), + ) + + +def sweep( + formats: tuple[tuple[Framing, ContainerLayout], ...] = FORMATS, +) -> list[Row]: + rows = [] + 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, 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"{'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:<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{'format':<7} {'rows':>5} {'raw B':>9} {'out B':>7} {'ratio':>7} {'zstd19 B':>9}" + ) + for framing, layout in FORMATS: + for rows in (1, 8, 64): + r = measure(tiled_rows(rows), framing, "bf16", "tiled", layout) + print( + 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/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/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/runtime.rs b/crates/ptwm-core/src/chain/runtime.rs index c8b9fca..98d8d5f 100644 --- a/crates/ptwm-core/src/chain/runtime.rs +++ b/crates/ptwm-core/src/chain/runtime.rs @@ -270,65 +270,178 @@ 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 is the product of two cases that existing files can +/// combine: +/// +/// - A Source that overstates the raw size by up to 8×. Only files written +/// by the Python byte-input path before the P4 fix do this, and their +/// chains are 1-D byte passthroughs or byte splits, whose planes are no +/// larger than the Source. +/// - Ops that expand a plane. MxFp4 deinterleave at most doubles it. The +/// spherical and alpha-stable scale planes can reach 4× their input (one +/// 8-byte scale per row of a single 16-bit element), but their builders +/// accept only 2-D shapes, which the Python path always sizes exactly. +/// +/// So no chain a writer emits multiplies the two, and 16× leaves 2× of +/// headroom over the largest case (8× overstated Source, or a 4× scale +/// plane). The constant 4096 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) +} + +/// 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) +} - // 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 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 +1475,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/chain/wire.rs b/crates/ptwm-core/src/chain/wire.rs index 258a454..583f041 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. @@ -351,19 +340,12 @@ 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 (mut 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(), - )); - } + // 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)) } @@ -452,6 +434,101 @@ 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()); + } + + /// 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").unwrap(); + 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/chain/wire_compact.rs b/crates/ptwm-core/src/chain/wire_compact.rs new file mode 100644 index 0000000..f98cdc5 --- /dev/null +++ b/crates/ptwm-core/src/chain/wire_compact.rs @@ -0,0 +1,1025 @@ +//! 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 and the `\x89PTWC` compact container use it. 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]>, + 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> { + 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 { + 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; + } + } + 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]>, + 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; + 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 { + 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 && binding == SourceBinding::Tensor(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 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], + 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; + + #[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(); + 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/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..a764acc 100644 --- a/crates/ptwm-core/src/codec_tagged.rs +++ b/crates/ptwm-core/src/codec_tagged.rs @@ -214,6 +214,98 @@ pub fn zstd_decode(blob: &[u8]) -> Result, PtwmCoreError> { }) } +/// Largest number of bytes one compressed Zstd byte can decode to. +/// +/// 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")] +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; + let err = |e: std::io::Error| PtwmCoreError::CodecDecode { + codec, + msg: e.to_string(), + }; + 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) + .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, + 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::*; @@ -350,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/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/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/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/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/rans.rs b/crates/ptwm-core/src/codecs/rans.rs index 03205ab..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; @@ -142,6 +142,32 @@ 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> { + // 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")?; + } + let out = self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + )?; + check_decoded_len(self.id(), out.len(), max_len)?; + Ok(out) + } } #[cfg(test)] @@ -272,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 0302582..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; @@ -135,6 +135,32 @@ 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> { + // 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")?; + } + let out = self.decode( + state_format_version, + state_bytes, + payload, + layout, + decoded_len, + )?; + check_decoded_len(self.id(), out.len(), max_len)?; + Ok(out) + } } #[cfg(test)] @@ -260,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/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..41e9b21 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,19 @@ 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) } #[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(), )) @@ -160,22 +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. - let capacity = if decoded_len > 0 { - decoded_len - } else { - zstd::bulk::Decompressor::upper_bound(payload) - .unwrap_or_else(|| payload.len().saturating_mul(8)) - }; - decompressor.decompress(payload, capacity) +fn decode_with_dict( + payload: &[u8], + dict: &[u8], + decoded_len: usize, + max_len: usize, +) -> 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"))] @@ -202,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/compact.rs b/crates/ptwm-core/src/compact.rs new file mode 100644 index 0000000..a582008 --- /dev/null +++ b/crates/ptwm-core/src/compact.rs @@ -0,0 +1,1157 @@ +//! 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, +//! 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) +//! [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) +//! 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 +//! 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: +//! 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 +//! [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. 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 +//! 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::runtime::{ + InverseContext, chain_source_descriptor, inverse_chain, max_plane_len, + propagate_terminal_descriptors_bounded, +}; +use crate::chain::validate::validate_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, canonical_dtype_name}; +use crate::error::PtwmCoreError; +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; + +/// 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. 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; +/// 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; +/// 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_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 { + 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`. + 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 +} + +/// 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 +/// [`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 orig_size = raw.len() as u64; + + 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); + // Shape metadata is recorded only when both shape and dtype name are + // known, matching the full container's tensor metadata. + 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 { + 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()); + } + bound_source = bound_source_params(shape, input.dtype_code); + } + 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()); + } + + // 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 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 (chain, planes). `Ok(None)` when a plane needs +/// shared or external state. +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() { + return Ok(None); + } + let chain: &Chain = &input.candidate_chains[plan.pass1[0].chain_idx]; + let planes = &plan.chosen_per_tensor[0]; + + let mut body = Vec::with_capacity( + 64 + planes + .iter() + .map(|p| p.payload_bytes.len() + p.encoded_state_bytes.len() + 16) + .sum::(), + ); + write_chain_compact(chain, bound_source, op_wire_ref, &mut body)?; + + 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), + StateSource::Inline + if cp.encoded_state_version == 0 && cp.encoded_state_bytes.is_empty() => + { + write_varint(&mut body, codec | STATE_INLINE_EMPTY) + } + StateSource::Inline => { + 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 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)) +} + +/// 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.take(MAGIC_COMPACT.len())?; + let flags = r.u8()?; + if flags & !KNOWN_FLAGS != 0 { + return Err(r.err(format_args!( + "unknown flag bits {:#04x}", + 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 { + Some(r.varint()?) + } else { + None + }; + + let mut shape = None; + let mut dtype_name = None; + if flags & FLAG_SHAPE != 0 { + 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()?); + } + 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(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")); + } + 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 { + Some(r.array::<32>()?) + } else { + None + }; + let payload_hash = if flags & FLAG_PAYLOAD_HASH != 0 { + 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 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.take(r.remaining())?; + if rest.len() as u64 != orig_size { + return Err(r.err(format_args!( + "stored body is {} bytes, header says {orig_size}", + rest.len() + ))); + } + 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, + h.flags & FLAG_PLANE_CRC != 0, + orig_size, + bound.as_deref(), + &mut payload_len, + )? + }; + r.finish("body")?; + + if let Some(expected) = h.payload_hash + && xxh64(&raw, 0) != expected + { + return Err(r.err("payload hash mismatch")); + } + + Ok(CompactTensor { + 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 Cursor<'_>, + has_crc: bool, + orig_size: u64, + bound_source: Option<&[u8]>, + payload_len_out: &mut u64, +) -> Result, PtwmCoreError> { + 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 + // 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 { + return Err(r.err(format_args!( + "chain Source declares {} bytes, header says {orig_size}", + source.length_bytes + ))); + } + validate_chain(&chain, 0)?; + let descriptors = propagate_terminal_descriptors_bounded(&chain, max_plane_len(orig_size))?; + + let n_planes = descriptors.len(); + let mut terminal_planes: Vec = Vec::with_capacity(n_planes); + 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) + .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 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(r.err(format_args!("unknown state code {s}"))), + }; + let payload_len = r.len_varint()?; + let crc = if has_crc { Some(r.u32_le()?) } 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(|| 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(|_| 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. + let bytes = codec.decode_bounded( + state_version, + state, + payload, + &layout, + decoded_len, + decoded_len, + )?; + if bytes.len() != decoded_len { + return Err(r.err(format_args!( + "plane decoded to {} bytes, chain says {decoded_len}", + bytes.len() + ))); + } + if let Some(expected) = crc + && crc32fast::hash(&bytes) != expected + { + return Err(r.err("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(r.err(format_args!( + "decoded {} bytes, header says {orig_size}", + source.bytes.len() + ))); + } + Ok(source.bytes.as_ref().to_vec()) +} + +#[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(); + + 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, + ..Default::default() + } + } + + 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 declares"), "{err}"); + 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 + // 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}"); + } + } + + /// 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), (false, true)] { + 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, 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}, 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(); + 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_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); + 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, 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])) + .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 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(), + ); + } + 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/compressor.rs b/crates/ptwm-core/src/compressor.rs index 32e43a6..c8a5735 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; @@ -12,11 +13,11 @@ 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; -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}; @@ -38,7 +40,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"), @@ -87,7 +89,11 @@ 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 compact container layout +/// (`\x89PTWC`, readable by v1.2 and later). Set `container_layout` to +/// [`ContainerLayout::Legacy`] for output that v1.1 and later read. +#[derive(Default)] pub struct CompressorOptions { /// Header hint to readers about the dominant codec. pub method_hint: u16, @@ -105,6 +111,9 @@ 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. Defaults to + /// [`ContainerLayout::Compact`]. + pub container_layout: ContainerLayout, } // --------------------------------------------------------------------------- @@ -139,7 +148,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 +297,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 +469,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)` @@ -518,9 +527,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) { @@ -616,46 +659,73 @@ 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 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 => { - 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(); } - (ChainRegistry { chains: registry }, tensor_to_idx) } // --------------------------------------------------------------------------- // 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 +802,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 +916,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,18 +926,94 @@ 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 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 (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)?; - writer.set_chain_registry(registry)?; + 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)?; + } // Pre-compute the chunked-path payloads in parallel before the sequential // write loop. Otherwise encode_plane_chunked runs inside the writer's @@ -908,6 +1056,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]; @@ -981,11 +1130,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, @@ -1006,26 +1150,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)?; @@ -1272,6 +1410,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, }, ) .expect("compress succeeds"); @@ -1336,6 +1476,100 @@ 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, + ..Default::default() + }; + 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 @@ -1373,6 +1607,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"); @@ -1492,6 +1727,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .unwrap(); @@ -1579,6 +1815,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, }, ) .unwrap(); @@ -1661,6 +1899,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, }, ) .unwrap(); @@ -1682,6 +1922,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 dc84825..632f9c1 100644 --- a/crates/ptwm-core/src/container.rs +++ b/crates/ptwm-core/src/container.rs @@ -7,21 +7,36 @@ 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}; +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::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}; +use crate::tensor_record_compact::{ + RecordContext, parse_compact_record, plane_codec_ids, record_binding, write_compact_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 +67,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 +96,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!( @@ -98,6 +119,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) @@ -105,9 +132,302 @@ 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`]. +#[derive(Debug)] +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`]. +/// +/// Built-in entries carry the ids and labels of `builtin_version` +/// ([`crate::extension::LEGACY_BUILTIN_VERSION`] for new legacy 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, builtin_entries, builtin_registry}; + + let mut builder = ExtensionTableBuilder::with_builtin_version(builtin_version)?; + 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_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_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 canonical_id in codec_ids { + let keep = match codecs { + CodecSet::AllBuiltins => true, + CodecSet::Only(wanted) => wanted.contains(&canonical_id), + }; + if keep { + 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); + } + } + + Ok(BuiltExtensionTable { + table: ExtensionTable { + entries: builder.finish(), + }, + registry_bytes, + codec_idx_map, + }) +} + +/// 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())); + } + + // 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) { + 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, + }) +} + +/// 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. +/// +/// `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 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. + #[default] + 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, + } + } +} + +/// 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 +/// 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, @@ -132,27 +452,56 @@ 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, + /// 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 { - /// 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(); + let builtin_version = default_builtin_version(layout); Ok(Self { writer, + layout, + compact_index: Vec::new(), + encoder: RecordEncoder { + layout, + builtin_version: builtin_version.clone(), + codec_ids: Vec::new(), + registry_bound: Vec::new(), + }, header, header_position, tensors_start_pos: 0, @@ -166,6 +515,10 @@ impl ContainerWriter { prelude_written: false, finalized: false, prelude_section_len: 0, + names: HashSet::new(), + builtin_version, + raw_extension_table: None, + table_set: false, }) } @@ -191,49 +544,158 @@ impl ContainerWriter { /// second table-building pass. Any previously set extension table is /// replaced. /// + /// 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). + /// /// 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. + /// + /// 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, + 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)?; + 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 = registry; + self.codec_idx_map = built.codec_idx_map; + self.refresh_encoder(); + Ok(()) + } - // 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) + /// 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)) { - let canonical_id = entry.canonical_id; - let idx = builder.intern(entry); - codec_idx_map.insert(canonical_id, idx); + 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 + } - let entries = builder.finish(); - self.extension_table = ExtensionTable { entries }; - self.chain_registry_bytes = registry_bytes; + /// 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) + { + 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; + self.raw_extension_table = None; self.chain_registry = registry; self.codec_idx_map = codec_idx_map; + self.table_set = true; + self.refresh_encoder(); Ok(()) } @@ -251,6 +713,76 @@ 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; + self.refresh_encoder(); + } + + /// Write built-in Extension Table entries with the ids and labels of + /// `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 + /// 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 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(()) + } + + /// 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(), + )); + } + 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()); + self.chain_registry_bytes = built.registry_bytes; + self.chain_registry = registry; + self.codec_idx_map = built.codec_idx_map; + self.table_set = true; + self.refresh_encoder(); + Ok(()) } /// Write the extension table, prelude section, and chain registry section, @@ -261,13 +793,19 @@ 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). - 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) @@ -331,11 +869,34 @@ 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)?; + // 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 + .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 { @@ -349,7 +910,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 @@ -360,15 +921,31 @@ impl ContainerWriter { name_hash: name_hash(&record.name), tensor_offset, tensor_len, + name: None, }); + 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 - /// (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, @@ -377,6 +954,30 @@ impl ContainerWriter { if self.finalized { 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; @@ -392,7 +993,9 @@ 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(()) } @@ -407,6 +1010,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 @@ -550,6 +1164,171 @@ 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)) +} + +/// 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> { @@ -585,6 +1364,20 @@ 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. + /// 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> { @@ -616,128 +1409,37 @@ 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); + } if buf.len() < HEADER_LEN { return Err(PtwmCoreError::InvalidContainer( "file shorter than header".into(), )); } let header = Header::from_bytes(&buf[..HEADER_LEN])?; - - // 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 { - 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}")))? - }; - - // 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(); - - // 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 }); + if header.is_ptwx() != (header.magic == MAGIC_PTWX) { + return Err(PtwmCoreError::InvalidContainer( + "the PTWX header flag disagrees with the magic".into(), + )); } - // ── 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. - { - // 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(), - }; + 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); - // 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()) @@ -757,13 +1459,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() @@ -776,16 +1489,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 @@ -793,7 +1501,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, @@ -805,27 +1513,219 @@ impl<'a> ContainerReader<'a> { tensors_offset, chain_registry_offset, missing_extensions: needs, + builtin_version, + 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. + 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) + } + + /// 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(); + 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. 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 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) + }) } /// Retrieve the shape + dtype_name embedded in this tensor's metadata, @@ -833,70 +1733,46 @@ 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 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; - 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 - /// 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()); } - 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() + // 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 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)) }) @@ -908,11 +1784,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) } @@ -924,13 +1796,22 @@ 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}; 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( @@ -939,13 +1820,54 @@ 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 + // 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)?; + 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!( + "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` / @@ -957,11 +1879,12 @@ impl<'a> ContainerReader<'a> { )) })?; validate_chain(chain, n_dependencies)?; + let plane_bound = max_plane_len(rec.orig_size); - // 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)?; + // 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", @@ -1054,7 +1977,21 @@ 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) + .filter(|&t| !strict || t == descriptor.length_bytes) + .ok_or_else(|| { + PtwmCoreError::InvalidContainer(format!( + "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; let decoded_chunks: Vec> = chunks .par_iter() .enumerate() @@ -1088,12 +2025,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( @@ -1166,16 +2106,49 @@ impl<'a> ContainerReader<'a> { )); } }; - // Decoded length comes from the propagated terminal descriptor. - let decoded_len = descriptor.length_bytes as usize; - codec.decode( + // 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()) + })?; + // 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. + 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, - )? + 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. + 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 { @@ -1199,6 +2172,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. @@ -1222,8 +2202,10 @@ 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; fn sample_tensor_record(name: &str) -> TensorRecord { TensorRecord { @@ -1525,6 +2507,417 @@ 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 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:?}"); + // 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 { + 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() + ); + } + + /// 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) + .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 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] + 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()); + 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(); + } + + /// 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] + 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; + + let patch = |len: u64| { + let mut h = header; + 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).expect_err("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)).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)).expect_err("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)).expect_err("must fail"); + assert!(err.to_string().contains("shared prelude"), "{err}"); + } + #[test] fn chain_registry_roundtrip() { use crate::extension::{ExtensionTable, ExtensionTableBuilder}; @@ -1854,6 +3247,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -1907,6 +3301,7 @@ mod tests { forced_codec: Some(CodecId::Rans), allow_codec_ids: None, chunk_size: Some(256), + ..Default::default() }, ) .expect("compress succeeds"); @@ -1957,6 +3352,7 @@ mod tests { forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -2100,6 +3496,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, }, ) .expect("compress succeeds"); @@ -2501,6 +3899,7 @@ label = "test-author" forced_codec: None, allow_codec_ids: None, chunk_size: None, + ..Default::default() }, ) .expect("compress succeeds"); @@ -2514,4 +3913,219 @@ 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 { + 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 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, + // Parses legacy records, so the layout is pinned. + container_layout: ContainerLayout::Legacy, + }, + ) + .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 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 + // 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/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/builder.rs b/crates/ptwm-core/src/extension/builder.rs index 8fef2f3..7ac7e77 100644 --- a/crates/ptwm-core/src/extension/builder.rs +++ b/crates/ptwm-core/src/extension/builder.rs @@ -2,14 +2,28 @@ use std::collections::HashMap; -use super::{CanonicalId, ExtensionTableEntry}; +use super::{ + CURRENT_BUILTIN_VERSION, CanonicalId, ExtensionError, ExtensionTableEntry, + builtin_entry_at_version, builtin_registry, is_valid_builtin_version, +}; /// 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 { + by_id: HashMap::new(), + entries: Vec::new(), + builtin_version: CURRENT_BUILTIN_VERSION.to_owned(), + } + } } impl ExtensionTableBuilder { @@ -17,6 +31,71 @@ 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. + /// + /// 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 + /// `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 +144,53 @@ 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").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)); + 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 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![ + 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..1da2c35 100644 --- a/crates/ptwm-core/src/extension/builtins.rs +++ b/crates/ptwm-core/src/extension/builtins.rs @@ -1,24 +1,212 @@ -//! 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"); + +/// 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."; + +/// 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) + } + + /// 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 { + 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 +218,281 @@ 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 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 { + 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()) +} + +/// 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.@`, and `` passes +/// [`is_valid_builtin_version`]; +/// - `` 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; + }; + if !is_valid_builtin_version(ver) { + 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")); + } + + /// 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"); + 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/compact_table.rs b/crates/ptwm-core/src/extension/compact_table.rs new file mode 100644 index 0000000..ec43ddb --- /dev/null +++ b/crates/ptwm-core/src/extension/compact_table.rs @@ -0,0 +1,412 @@ +//! 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/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/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 d6a2be5..860e449 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; @@ -15,7 +16,11 @@ 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, + 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}; pub use error::ExtensionError; @@ -32,9 +37,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/extension/table.rs b/crates/ptwm-core/src/extension/table.rs index ad99cf9..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)?); } @@ -114,6 +120,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]; @@ -123,6 +142,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"))?; @@ -186,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]; @@ -424,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/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 41c504c..cb2661e 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; @@ -38,15 +39,18 @@ 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; pub mod types; +pub mod varint; pub use dtype::{Dtype, PreprocessingModes}; pub use error::PtwmCoreError; 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/ptwc.rs b/crates/ptwm-core/src/ptwc.rs new file mode 100644 index 0000000..81ce002 --- /dev/null +++ b/crates/ptwm-core/src/ptwc.rs @@ -0,0 +1,448 @@ +//! 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 (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 +//! 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. +//! +//! 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. + +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 +} + +#[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.rs b/crates/ptwm-core/src/tensor_record.rs index d28e697..80b7bf7 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::{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; @@ -79,7 +79,48 @@ 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::() + )) + }) +} + +/// Append the serialized record to `out`. Built-in entries of an inline +/// chain's Extension Table carry the ids and labels of +/// [`LEGACY_BUILTIN_VERSION`]. +/// +/// On error, `out` is truncated back to its length before the call, so a +/// failed write never leaves a partial record behind. pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), PtwmCoreError> { + write_tensor_record_with_builtin_version(rec, out, LEGACY_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, builtin_version); + if result.is_err() { + out.truncate(start); + } + 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; if has_inline && rec.inline_chain.is_none() { @@ -105,9 +146,21 @@ 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 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 = 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) + @@ -144,7 +197,7 @@ pub fn write_tensor_record(rec: &TensorRecord, out: &mut Vec) -> Result<(), // 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 { @@ -159,7 +212,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 +221,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 +231,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,7 +244,38 @@ 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> { + 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(), @@ -257,6 +341,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 @@ -265,10 +350,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}")) })?; @@ -283,7 +369,9 @@ pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCore 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]; @@ -354,6 +442,7 @@ pub fn parse_tensor_record(buf: &[u8]) -> Result<(TensorRecord, usize), PtwmCore terminals, }, pos, + inline_version, )) } @@ -487,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]; @@ -586,4 +747,95 @@ 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"); + + // 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] + 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-core/src/tensor_record_compact.rs b/crates/ptwm-core/src/tensor_record_compact.rs new file mode 100644 index 0000000..ce30821 --- /dev/null +++ b/crates/ptwm-core/src/tensor_record_compact.rs @@ -0,0 +1,1267 @@ +//! 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 `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}; +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, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::chain::{ChainEdge, ChainNode, TerminalRef}; + use crate::types::Role; + + const CODEC_IDS: &[Option] = &[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 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()); + 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/manifest.rs b/crates/ptwm-core/src/transcode/manifest.rs index 6122099..223931f 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,11 +17,18 @@ 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 /// 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 +39,88 @@ 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. `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, } #[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, +} + +/// 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 { /// Stable wire op id (`OpId::as_u16`). pub id: u16, @@ -49,6 +129,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 +141,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 +155,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 +164,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 +172,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 +183,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 +195,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 +219,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 +231,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,15 +239,19 @@ pub struct DependencyJson { } #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(deny_unknown_fields)] pub struct PlaneJson { pub index: u32, /// Human-readable plane role label. 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`. @@ -184,6 +277,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 +286,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 a8bd0a7..6f23f8d 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. @@ -23,19 +35,28 @@ 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::{ + 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::{CanonicalId, ExtensionTable, LEGACY_BUILTIN_VERSION}; 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, write_tensor_record_with_builtin_version, }; use crate::transforms::op::OpId; use crate::types::PlaneRole; use crate::types::role::Role; +use crate::varint::write_varint; use manifest::*; @@ -229,12 +250,82 @@ 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, -/// 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 == ContainerLayout::Compact { + return explode_compact(&reader); + } let embedded_wasm = reader.header.flags & FLAG_PTWX != 0; if embedded_wasm { return Err(err( @@ -245,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. @@ -282,6 +353,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 @@ -291,7 +363,43 @@ 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, 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. + 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(_) => {} + } + } + // `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(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()) { + 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()); } @@ -328,16 +436,57 @@ 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(LEGACY_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 != LEGACY_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 = @@ -350,6 +499,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) { @@ -493,30 +783,157 @@ 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], 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 { - return Err(err(format!( - "unexpected ptwm_format {:?}", - registry.ptwm_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()); + 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}")))?; 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(LEGACY_BUILTIN_VERSION); + + 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()); + { + 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. 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). + 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 { + 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)?; + } + writer.finalize()?; + } + Ok(cursor.into_inner()) +} - // 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 { +/// [`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 @@ -535,35 +952,32 @@ pub fn implode( state_bytes, }); } + Ok(prelude_entries) +} - // Rebuild the chain registry (order preserved). +/// Rebuild the chain registry (order preserved). +fn chains_from_manifest(chains: &[ChainEntry]) -> Result { let mut chain_registry = ChainRegistry::new(); - for c in ®istry.registry.chains { + 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. +/// 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, &mut members)?); + records.push(record_from_manifest(&tm, members)?); } - - // Re-drive the deterministic container writer. - let mut cursor = Cursor::new(Vec::::new()); - { - let mut writer = ContainerWriter::new(&mut cursor)?; - writer.set_prelude(&prelude_entries)?; - writer.set_chain_registry(chain_registry)?; - for record in &records { - writer.append_tensor(record)?; - } - writer.finalize()?; - } - Ok(cursor.into_inner()) + Ok(records) } fn record_from_manifest( diff --git a/crates/ptwm-core/src/transcode/tests.rs b/crates/ptwm-core/src/transcode/tests.rs index f4ed981..3a24661 100644 --- a/crates/ptwm-core/src/transcode/tests.rs +++ b/crates/ptwm-core/src/transcode/tests.rs @@ -286,3 +286,405 @@ 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, + std::slice::from_ref(&b.manifest_json), + 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" + ); +} + +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 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())], + }; + 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(version).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(); + } + 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" + ); + 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); +} + +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-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/op.rs b/crates/ptwm-core/src/transforms/op.rs index d49dc9b..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, @@ -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/src/transforms/source.rs b/crates/ptwm-core/src/transforms/source.rs index 2abde55..6aa5a5d 100644 --- a/crates/ptwm-core/src/transforms/source.rs +++ b/crates/ptwm-core/src/transforms/source.rs @@ -48,11 +48,60 @@ 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; + } + 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()); } - n_elements * (element_width_for(dtype_code).bits_per_element() as u64 / 8) + out.extend_from_slice(&chain_dtype.to_le_bytes()); + Some(out) } /// Mandatory entry node of every PPG chain. Takes no inputs (the runtime @@ -89,17 +138,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 +174,7 @@ impl Source { residual_of: None, is_nibble_packed: is_nibble_packed_dtype(self.dtype_code), vendor_bytes: vec![], - } + }) } } @@ -129,7 +195,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> { @@ -210,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(); @@ -362,4 +459,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/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/crates/ptwm-core/src/types/role.rs b/crates/ptwm-core/src/types/role.rs index 8c957ed..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, @@ -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), @@ -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, @@ -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), @@ -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/src/varint.rs b/crates/ptwm-core/src/varint.rs new file mode 100644 index 0000000..accac41 --- /dev/null +++ b/crates/ptwm-core/src/varint.rs @@ -0,0 +1,264 @@ +//! 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`, 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; + +/// 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/crates/ptwm-core/tests/decode_bomb.rs b/crates/ptwm-core/tests/decode_bomb.rs new file mode 100644 index 0000000..8a24bca --- /dev/null +++ b/crates/ptwm-core/tests/decode_bomb.rs @@ -0,0 +1,247 @@ +//! 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::{ContainerLayout, 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, + ..Default::default() + } +} + +/// 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); +} + +/// 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(); + 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() + .decode_tensor("t") + .unwrap(), + raw + ); + buf +} + +#[test] +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(); + 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_err() + .to_string(); + assert_bounded(err, t0); + } +} 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 0000000..c50d05c Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32.ptwm differ 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 0000000..00ff06f Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/legacy_p4_byte_fp32_v1_1_0.ptwm differ 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 0000000..7f63ad3 Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/legacy_trimmed.ptwm differ 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 0000000..9c14457 Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/legacy_v1_0_0.ptwm differ 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 0000000..24b6f48 Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/legacy_v1_0_0_native.ptwm differ 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 0000000..883fd1e Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/legacy_v1_1_0.ptwm differ 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 0000000..e432e61 Binary files /dev/null and b/crates/ptwm-core/tests/fixtures/production_chains.bin differ diff --git a/crates/ptwm-core/tests/legacy_fixtures.rs b/crates/ptwm-core/tests/legacy_fixtures.rs new file mode 100644 index 0000000..3b06a27 --- /dev/null +++ b/crates/ptwm-core/tests/legacy_fixtures.rs @@ -0,0 +1,1412 @@ +//! 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 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 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 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 +//! `(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 +//! +//! ```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`, +//! `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; +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, 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}; +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 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`. +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), + ], + ), + ( + 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, + &[ + ("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)]), + (P4_V1_1_0, &[("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 = + "ec513f5c09f76d8f220a69dd3c47b96dec0ca7e3fcd1a075f601d3688b3ffc9a"; + +// ── 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`), +/// pinned to the legacy layout that these fixtures freeze. +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, + container_layout: ContainerLayout::Legacy, + } +} + +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`], rewritten +/// with a trimmed Extension Table. +fn trimmed_container() -> Vec { + let data = trimmed_tensors(); + rewrite_as_trimmed(&[&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 { + 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")) + .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 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(); + 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 { + 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. +/// +/// 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(), + "rederive_builtins only maps between versions of equal length" + ); + if from == to { + return 0; + } + // 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; + 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; + } + + 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. +/// +/// 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(&[n / 4])], + dtype_code: 1, + // Format.BYTE + input_format: 1, + raw_bytes: &bytes, + shape: None, + dtype_name: None, + delta_reference_blake3: 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. +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 { + 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 +} + +/// `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) { + 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 + .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 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) { + 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() + .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, V1_0_0_NATIVE, TRIMMED, P4, P4_V1_1_0] + .iter() + .map(|f| load(f).len() as u64) + .sum(); + assert!(total < 96 * 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. + assert_eq!(plane_codec_count(&reader), 16); +} + +#[test] +fn v1_0_0_fixture_decodes() { + assert_decodes(V1_0_0); +} + +#[test] +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() { + 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(); + 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] +fn v1_0_0_native_fixture_decodes() { + assert_decodes(V1_0_0_NATIVE); +} + +#[test] +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 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 +/// [`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_eq!(plane_codec_count(&reader), 16, "v1.1.0 table layout"); + 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 +/// 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), + "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] +fn v1_1_0_layout_writer_output_is_golden() { + assert!(v1_1_0_container() == load(V1_1_0)); +} + +/// 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() { + 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 = rewrite_as_trimmed(&[&p4_container()]); + let p4_v1_1_0 = p4_v1_1_0_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), + (P4_V1_1_0, &p4_v1_1_0), + ]; + 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), + ("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() { + 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()); + } +} + +// ── 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/ptwc.rs b/crates/ptwm-core/tests/ptwc.rs new file mode 100644 index 0000000..7a1f0be --- /dev/null +++ b/crates/ptwm-core/tests/ptwc.rs @@ -0,0 +1,1517 @@ +//! 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::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}; + +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" + ); + assert_transcodes(&compact, fixture); + } +} + +#[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}"); + // 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 ─────────────────────────────────── + +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); + 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"); + if chunk.is_none() { + assert!(!reader.prelude.is_empty(), "no shared state"); + } + } + } + assert_eq!(decoded[0], decoded[1]); + } + } +} + +/// New multi-tensor writes default to the compact layout; the legacy +/// layout is still available on request. +#[test] +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_eq!(legacy[..5], *b"\x89PTWM"); + let from_default = compress(&data, CompressorOptions::default()); + assert_eq!(from_default[..5], *b"\x89PTWC"); +} + +#[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); + } + assert_transcodes(&compact, "bound Source"); + + // 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() + ); + } +} + +/// 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"); + assert_transcodes(&blob, "BYTE input"); + } +} + +#[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"); + assert_transcodes(&compact, "raw metadata"); + } +} + +#[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"); + assert_transcodes(&compact, "external state"); + // 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 +/// 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 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; + 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); + } +} + +// ── 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}"); +} + +// ── 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; + 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; + 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; + 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_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 + .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:?}"); + let decoded: Vec = reader + .decode_model() + .unwrap() + .into_iter() + .map(|(name, _)| name) + .collect(); + assert_eq!(decoded, names, "{layout:?}"); + } +} 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/handle.rs b/crates/ptwm-py/src/handle.rs new file mode 100644 index 0000000..0663c53 --- /dev/null +++ b/crates/ptwm-py/src/handle.rs @@ -0,0 +1,327 @@ +//! `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 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, PyMemoryView}; + +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 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]>), +} + +/// 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()) +} + +/// 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 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, 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)) + } 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: &Bound<'_, PyAny>, 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 write 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 a064bb2..0f55366 100644 --- a/crates/ptwm-py/src/inspect.rs +++ b/crates/ptwm-py/src/inspect.rs @@ -5,8 +5,7 @@ use std::io::Read; use pyo3::prelude::*; -use ptwm_core::extension::ExtensionTable; -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()) @@ -16,10 +15,17 @@ fn to_pyerr(e: E) -> PyErr { /// /// 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. +/// +/// 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 = @@ -27,31 +33,7 @@ pub fn list_extension_table_entries(path: String) -> PyResult buf.len() { - return Err(to_pyerr(format!( - "'{}': extension table region [{}..{}] exceeds file size {}", - path, - off, - off + len, - buf.len() - ))); - } - let table_bytes = &buf[off..off + len]; - let table = ExtensionTable::from_bytes(table_bytes).map_err(to_pyerr)?; + let table = read_extension_table(&buf).map_err(|e| to_pyerr(format!("'{path}': {e}")))?; Ok(table .entries .iter() diff --git a/crates/ptwm-py/src/lib.rs b/crates/ptwm-py/src/lib.rs index 7a56c71..c7fd17c 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; @@ -45,6 +47,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)?)?; @@ -58,11 +62,13 @@ 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)?; ext::register(m)?; inspect::register(m)?; + handle::register(m)?; delta_scheme::register(m)?; hardware::register(m)?; plane_codec_cuda::register(m)?; @@ -311,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 v1.1 and later). +/// /// 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="compact"))] fn compress_model<'py>( py: Python<'py>, records: Vec<( @@ -335,7 +344,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 {}", @@ -435,6 +452,7 @@ fn compress_model<'py>( forced_codec, allow_codec_ids, chunk_size: streaming_chunk, + container_layout, }; let cursor = std::io::Cursor::new(&mut out); @@ -446,41 +464,152 @@ fn compress_model<'py>( Ok(pyo3::types::PyBytes::new(py, &out)) } -/// Decode a single tensor from a container blob. +/// Compress one tensor into a compact single-tensor frame. /// -/// `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. +/// 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, + ..Default::default() + }; + 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, + )) +} + +// The functions below open a temporary `ContainerHandle` per call (see +// `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, 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 /// 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: &Bound<'_, PyAny>, 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 @@ -492,30 +621,13 @@ fn decode_tensor<'py>( #[pyo3(signature = (blob, skip_missing = false))] fn decode_model<'py>( py: Python<'py>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: &Bound<'_, PyAny>, 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`. @@ -523,162 +635,67 @@ fn decode_model<'py>( #[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_info( py: Python<'_>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: &Bound<'_, PyAny>, name: &str, 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(); - let info = py - .allow_threads( - || -> Result, String, u16, u8)>, PtwmCoreError> { - let reader = if skip_missing { - ContainerReader::open_partial(&data)? - } 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)?; - 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: &Bound<'_, PyAny>, name: &str, skip_missing: bool, ) -> PyResult>> { - use ptwm_core::tensor_record::{RefKind, parse_tensor_record}; - - 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 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)?; - 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>, + blob: &Bound<'_, PyAny>, 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) + 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`. #[pyfunction] +#[pyo3(signature = (blob, name, skip_missing = false))] fn decode_tensor_shape( py: Python<'_>, - blob: &Bound<'_, pyo3::types::PyBytes>, + blob: &Bound<'_, PyAny>, 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)?; - 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. +/// 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>, -) -> PyResult> { - use ptwm_core::tensor_record::parse_tensor_record; - - let data: Vec = blob.as_bytes().to_vec(); - let summaries = py - .allow_threads( - || -> Result, PtwmCoreError> { - let reader = ContainerReader::open(&data)?; - 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)?; - 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) + blob: &Bound<'_, PyAny>, + skip_missing: bool, +) -> PyResult> { + ContainerHandle::open(py, blob, skip_missing)?.plane_summary_rows(py) } /// Explode a `.ptwm` container into its canonical PTWM-JSON projection. @@ -800,13 +817,12 @@ 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 (mut 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(); + // 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()); @@ -819,6 +835,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/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/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/docs/architecture/wire-format.md b/docs/architecture/wire-format.md index 91bb620..c3319fb 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,221 +29,755 @@ 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, 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 +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.@`. 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 +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`) -## Top-level layout +`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. -The container is a flat byte stream with the following sections in -order. All offsets and lengths are stored in little-endian. +``` +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. + +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. + +### 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. - -## 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. - -## 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. +index_len u64 (footer) +sentinel [u8; 8] \x89TWMEOF\x1a (footer) +``` + +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. + +`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 | 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, +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 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; compact role +``` + +- `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. 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 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. 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: + +``` +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 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, +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: 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; 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: + +- 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`. 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 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: + +- 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; +- 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; +- `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 + 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..3d95c72 100644 --- a/docs/site/content/20.concepts/5.container.md +++ b/docs/site/content/20.concepts/5.container.md @@ -4,63 +4,109 @@ 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 262 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` 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 @@ -72,5 +118,5 @@ 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 `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 a35ed2a..65699b5 100644 --- a/docs/site/content/30.guides/01.cli.md +++ b/docs/site/content/30.guides/01.cli.md @@ -37,6 +37,24 @@ 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. 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 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 80166af..ff1e740 100644 --- a/docs/site/content/30.guides/02.python-api.md +++ b/docs/site/content/30.guides/02.python-api.md @@ -65,6 +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 (`\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 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 | + +## 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..e7e6030 --- /dev/null +++ b/docs/site/content/70.release-notes.md @@ -0,0 +1,98 @@ +--- +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; 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 | +|---|---|---| +| `Compressor.compress` (one tensor) | legacy container, `\x89PTWM` | compact frame, `\x89PTWS` | +| `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` 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`. + +### 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`, 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 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 + +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` | 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`) | + +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 + +- 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 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). 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..a94df8a 100644 --- a/fuzz/Cargo.toml +++ b/fuzz/Cargo.toml @@ -95,3 +95,27 @@ 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 + +[[bin]] +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/chain_wire_compact/bf16_const b/fuzz/corpus/chain_wire_compact/bf16_const new file mode 100644 index 0000000..2e2d811 Binary files /dev/null and b/fuzz/corpus/chain_wire_compact/bf16_const differ diff --git a/fuzz/corpus/chain_wire_compact/bf16_randn b/fuzz/corpus/chain_wire_compact/bf16_randn new file mode 100644 index 0000000..618ad33 Binary files /dev/null and b/fuzz/corpus/chain_wire_compact/bf16_randn differ diff --git a/fuzz/corpus/chain_wire_compact/bf16_smooth b/fuzz/corpus/chain_wire_compact/bf16_smooth new file mode 100644 index 0000000..d35a461 Binary files /dev/null and b/fuzz/corpus/chain_wire_compact/bf16_smooth differ diff --git a/fuzz/corpus/chain_wire_compact/byte_bf16 b/fuzz/corpus/chain_wire_compact/byte_bf16 new file mode 100644 index 0000000..b5a9946 Binary files /dev/null and b/fuzz/corpus/chain_wire_compact/byte_bf16 differ diff --git a/fuzz/corpus/chain_wire_compact/fp32_const b/fuzz/corpus/chain_wire_compact/fp32_const new file mode 100644 index 0000000..a247021 Binary files /dev/null and b/fuzz/corpus/chain_wire_compact/fp32_const differ diff --git a/fuzz/corpus/chain_wire_compact/fp32_smooth b/fuzz/corpus/chain_wire_compact/fp32_smooth new file mode 100644 index 0000000..ced2007 Binary files /dev/null and b/fuzz/corpus/chain_wire_compact/fp32_smooth differ diff --git a/fuzz/corpus/compact_decode/bf16_const_bare b/fuzz/corpus/compact_decode/bf16_const_bare new file mode 100644 index 0000000..8b6cd35 Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_const_bare differ diff --git a/fuzz/corpus/compact_decode/bf16_const_crc b/fuzz/corpus/compact_decode/bf16_const_crc new file mode 100644 index 0000000..02a8160 Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_const_crc differ diff --git a/fuzz/corpus/compact_decode/bf16_const_hash b/fuzz/corpus/compact_decode/bf16_const_hash new file mode 100644 index 0000000..a07363a Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_const_hash differ diff --git a/fuzz/corpus/compact_decode/bf16_randn_bare b/fuzz/corpus/compact_decode/bf16_randn_bare new file mode 100644 index 0000000..888b28f Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_randn_bare differ diff --git a/fuzz/corpus/compact_decode/bf16_randn_crc b/fuzz/corpus/compact_decode/bf16_randn_crc new file mode 100644 index 0000000..0507d10 Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_randn_crc differ diff --git a/fuzz/corpus/compact_decode/bf16_randn_hash b/fuzz/corpus/compact_decode/bf16_randn_hash new file mode 100644 index 0000000..8808745 Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_randn_hash differ diff --git a/fuzz/corpus/compact_decode/bf16_smooth_bare b/fuzz/corpus/compact_decode/bf16_smooth_bare new file mode 100644 index 0000000..0e4404d Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_smooth_bare differ diff --git a/fuzz/corpus/compact_decode/bf16_smooth_crc b/fuzz/corpus/compact_decode/bf16_smooth_crc new file mode 100644 index 0000000..3966e30 Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_smooth_crc differ diff --git a/fuzz/corpus/compact_decode/bf16_smooth_hash b/fuzz/corpus/compact_decode/bf16_smooth_hash new file mode 100644 index 0000000..d91b847 Binary files /dev/null and b/fuzz/corpus/compact_decode/bf16_smooth_hash differ diff --git a/fuzz/corpus/compact_decode/byte_bf16_hash b/fuzz/corpus/compact_decode/byte_bf16_hash new file mode 100644 index 0000000..1c61d57 Binary files /dev/null and b/fuzz/corpus/compact_decode/byte_bf16_hash differ diff --git a/fuzz/corpus/compact_decode/fp16_randn_bare b/fuzz/corpus/compact_decode/fp16_randn_bare new file mode 100644 index 0000000..e2ab027 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp16_randn_bare differ diff --git a/fuzz/corpus/compact_decode/fp16_randn_crc b/fuzz/corpus/compact_decode/fp16_randn_crc new file mode 100644 index 0000000..e2ab027 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp16_randn_crc differ diff --git a/fuzz/corpus/compact_decode/fp16_randn_hash b/fuzz/corpus/compact_decode/fp16_randn_hash new file mode 100644 index 0000000..7cf6f74 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp16_randn_hash differ diff --git a/fuzz/corpus/compact_decode/fp32_const_bare b/fuzz/corpus/compact_decode/fp32_const_bare new file mode 100644 index 0000000..48c43e2 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_const_bare differ diff --git a/fuzz/corpus/compact_decode/fp32_const_crc b/fuzz/corpus/compact_decode/fp32_const_crc new file mode 100644 index 0000000..d6683c9 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_const_crc differ diff --git a/fuzz/corpus/compact_decode/fp32_const_hash b/fuzz/corpus/compact_decode/fp32_const_hash new file mode 100644 index 0000000..d5408c4 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_const_hash differ diff --git a/fuzz/corpus/compact_decode/fp32_randn_bare b/fuzz/corpus/compact_decode/fp32_randn_bare new file mode 100644 index 0000000..70356b3 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_randn_bare differ diff --git a/fuzz/corpus/compact_decode/fp32_randn_crc b/fuzz/corpus/compact_decode/fp32_randn_crc new file mode 100644 index 0000000..fb7510f Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_randn_crc differ diff --git a/fuzz/corpus/compact_decode/fp32_randn_hash b/fuzz/corpus/compact_decode/fp32_randn_hash new file mode 100644 index 0000000..65c1b15 Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_randn_hash differ diff --git a/fuzz/corpus/compact_decode/fp32_smooth_bare b/fuzz/corpus/compact_decode/fp32_smooth_bare new file mode 100644 index 0000000..5ef04dc Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_smooth_bare differ diff --git a/fuzz/corpus/compact_decode/fp32_smooth_crc b/fuzz/corpus/compact_decode/fp32_smooth_crc new file mode 100644 index 0000000..684f7fa Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_smooth_crc differ diff --git a/fuzz/corpus/compact_decode/fp32_smooth_hash b/fuzz/corpus/compact_decode/fp32_smooth_hash new file mode 100644 index 0000000..ead5f2d Binary files /dev/null and b/fuzz/corpus/compact_decode/fp32_smooth_hash differ diff --git a/fuzz/corpus/compact_decode/fp8_e4m3_bare b/fuzz/corpus/compact_decode/fp8_e4m3_bare new file mode 100644 index 0000000..54f284e Binary files /dev/null and b/fuzz/corpus/compact_decode/fp8_e4m3_bare differ diff --git a/fuzz/corpus/compact_decode/fp8_e4m3_crc b/fuzz/corpus/compact_decode/fp8_e4m3_crc new file mode 100644 index 0000000..54f284e Binary files /dev/null and b/fuzz/corpus/compact_decode/fp8_e4m3_crc differ diff --git a/fuzz/corpus/compact_decode/fp8_e4m3_hash b/fuzz/corpus/compact_decode/fp8_e4m3_hash new file mode 100644 index 0000000..0dc87ce Binary files /dev/null and b/fuzz/corpus/compact_decode/fp8_e4m3_hash differ 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 diff --git a/fuzz/corpus/container_open/ptwc_bound_source b/fuzz/corpus/container_open/ptwc_bound_source new file mode 100644 index 0000000..172f6cb Binary files /dev/null and b/fuzz/corpus/container_open/ptwc_bound_source differ diff --git a/fuzz/corpus/container_open/ptwc_trimmed b/fuzz/corpus/container_open/ptwc_trimmed new file mode 100644 index 0000000..9368bc9 Binary files /dev/null and b/fuzz/corpus/container_open/ptwc_trimmed differ diff --git a/fuzz/corpus/extension_table_compact/bound_source b/fuzz/corpus/extension_table_compact/bound_source new file mode 100644 index 0000000..3b37837 Binary files /dev/null and b/fuzz/corpus/extension_table_compact/bound_source differ diff --git a/fuzz/corpus/extension_table_compact/trimmed b/fuzz/corpus/extension_table_compact/trimmed new file mode 100644 index 0000000..d767621 Binary files /dev/null and b/fuzz/corpus/extension_table_compact/trimmed differ 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 0000000..4a27e84 Binary files /dev/null and b/fuzz/corpus/parse_wire_compact/record_3 differ diff --git a/fuzz/fuzz_targets/chain_wire_compact.rs b/fuzz/fuzz_targets/chain_wire_compact.rs new file mode 100644 index 0000000..ebd0b6d --- /dev/null +++ b/fuzz/fuzz_targets/chain_wire_compact.rs @@ -0,0 +1,33 @@ +#![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. +// 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; +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 new file mode 100644 index 0000000..74a5ee1 --- /dev/null +++ b/fuzz/fuzz_targets/compact_decode.rs @@ -0,0 +1,54 @@ +#![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. 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}; +use ptwm_core::varint::read_varint; + +/// 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; + } + 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. + 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()); + buf.extend_from_slice(&MAGIC_COMPACT); + buf.extend_from_slice(data); + 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/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"); + } +}); diff --git a/python/ptwm/__init__.py b/python/ptwm/__init__.py index 19f343e..5f774fc 100644 --- a/python/ptwm/__init__.py +++ b/python/ptwm/__init__.py @@ -7,17 +7,20 @@ Public entry points: :class:`Compressor`, :class:`Decompressor`, :class:`CompressionConfig`, :class:`DecompressionConfig`, -:class:`Method`, :class:`Format`. 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, Lossy, Method, ) @@ -39,9 +42,11 @@ "CompressionConfig", "CompressionMethodNotSupportedError", "Compressor", + "ContainerLayout", "DecompressionConfig", "Decompressor", "Format", + "Framing", "Lossy", "Method", "Error", diff --git a/python/ptwm/_config.py b/python/ptwm/_config.py index 3b824be..d8d6899 100644 --- a/python/ptwm/_config.py +++ b/python/ptwm/_config.py @@ -42,6 +42,85 @@ def _missing_(cls, value): return None +class Framing(Enum): + r"""Byte layout of a single-tensor :meth:`Compressor.compress` output. + + ``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. Set + ``framing=Framing.CONTAINER`` and + ``container_layout=ContainerLayout.LEGACY`` for output that older releases + read. + """ + + 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 ContainerLayout(Enum): + r"""On-disk layout of a multi-tensor ``.ptwm`` container. + + ``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 v1.1 and later read. Its + 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 + :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``. + + ``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: + msg = f"unknown container_layout {layout!r} (expected 'compact' or 'legacy')" + raise ValueError(msg) from None + + class Lossy(Enum): NONE = 0 INTEGER = 1 @@ -87,6 +166,35 @@ class CompressionConfig: """ codec: str | None = None device: int | None = None + framing: Framing = Framing.COMPACT + """Output layout; see :class:`Framing`.""" + container_layout: ContainerLayout = ContainerLayout.COMPACT + """Layout of a full container written with ``Framing.CONTAINER`` (or as + 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 + """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. + """ + + 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) + ) + # 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/python/ptwm/_magic.py b/python/ptwm/_magic.py new file mode 100644 index 0000000..58477b1 --- /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 (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 + :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: + r"""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 c4d8820..9eaec13 100644 --- a/python/ptwm/_rust/__init__.py +++ b/python/ptwm/_rust/__init__.py @@ -10,13 +10,22 @@ 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 a read-only +``mmap`` are borrowed without a copy; every other buffer is copied once.""" + __all__ = [ + "ContainerHandle", "blake3_hash", "builtin_canonical_id", "codec_id_registry", + "compress_compact", "compress_model", + "container_layout", + "decode_compact", "decode_model", "decode_tensor", "decode_tensor_delta_hash", @@ -188,8 +197,10 @@ def compress_model( method_hint: int = 3, streaming_chunk: int | None = None, codec_menu: list[int] | None = None, + *, + container_layout: ContainerLayout | str = "compact", ) -> 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, @@ -202,7 +213,14 @@ 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``: ``"compact"`` (default) writes ``\x89PTWC`` + (readable by v1.2 and later), ``"legacy"`` writes ``\x89PTWM`` (readable + 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``). """ + layout = container_layout_name(container_layout) return bytes( _ext.compress_model( records, @@ -212,12 +230,74 @@ def compress_model( method_hint, streaming_chunk, codec_menu, + container_layout=layout, ) ) @translate_errors -def decode_tensor(blob: bytes, name: str, skip_missing: bool = False) -> bytes: +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: Buffer, name: str, skip_missing: bool = False) -> bytes: """Decode a single tensor from a ``.ptwm`` container blob. Returns raw decompressed bytes. @@ -226,10 +306,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. """ @@ -237,14 +317,16 @@ 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: 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) + return _ext.decode_tensor_shape(blob, name, skip_missing) @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``. @@ -255,31 +337,133 @@ 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.""" +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. + """ return list(_ext.list_tensor_names(blob, skip_missing)) @translate_errors def list_plane_summaries( - blob: bytes, + blob: Buffer, 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 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: Buffer) -> 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`` 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. + + Errors from the binding surface as :class:`ptwm.Error`. + """ + + __slots__ = ("_h",) + + @translate_errors + def __init__(self, blob: Buffer, 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 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 + 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 ------------------- @@ -287,13 +471,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). """ @@ -309,5 +499,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/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/README.md b/python/ptwm/cli/README.md index 94eb07d..1222d5a 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 +v1.1 and later read. diff --git a/python/ptwm/cli/compress.py b/python/ptwm/cli/compress.py index c40daf5..47aaa0a 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 @@ -22,6 +24,27 @@ } +_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). With 'legacy', per-tensor " + "safetensors output also uses full \\x89PTWM containers instead of " + "the v1.2 \\x89PTWS frame." +) + + +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,13 +154,19 @@ 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, DecompressionConfig, Decompressor, + Framing, Method, ) @@ -162,6 +191,10 @@ 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 (they read both layouts). + framing=Framing.CONTAINER, + container_layout=container_layout, is_streaming=is_streaming, streaming_chunk=chunk_size, method=Method(method), @@ -246,13 +279,18 @@ 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, DecompressionConfig, Decompressor, + Framing, Method, ) @@ -284,7 +322,9 @@ 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, streaming_chunk=chunk_size, delta_compressed_type="file", @@ -298,21 +338,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) @@ -344,8 +383,15 @@ 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.""" + r"""Compress a safetensors file. + + 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 from safetensors.torch import save_file # noqa: PLC0415 @@ -355,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, @@ -391,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 @@ -409,6 +462,8 @@ def compress_safetensors_file( codec_menu=codec_menu, codec=codec, device=device, + framing=framing, + container_layout=container_layout, ) ) compressor = compressor_cache[dtype_str] @@ -427,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: @@ -467,6 +522,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 @@ -563,6 +619,7 @@ def compress_path( codec_menu, codec, device, + container_layout, ) else: compression_func = compress_safetensors_file @@ -576,31 +633,47 @@ def compress_path( codec_menu, codec, device, + container_layout, ) 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}", @@ -695,6 +768,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) @@ -736,11 +810,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 @@ -903,7 +978,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( @@ -955,6 +1033,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 @@ -986,6 +1065,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 @@ -1002,11 +1082,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 { @@ -1046,22 +1135,29 @@ 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`). 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") - out_dir.mkdir(parents=True, exist_ok=True) - 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 + 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 @@ -1078,14 +1174,49 @@ 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()) + 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: + _remove_path(out_path) + shards[0].replace(out_path) + shutil.rmtree(staging) + return out_path + final_dir = out_path.with_suffix("") + _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. @@ -1156,6 +1287,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: @@ -1175,6 +1307,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( @@ -1187,6 +1320,7 @@ def handle_compress(args): codec_menu=codec_menu, codec=args.codec, device=args.device, + container_layout=_container_layout_arg(args), ) else: compress_file( @@ -1204,6 +1338,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( @@ -1227,6 +1362,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/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/core/_compressor.py b/python/ptwm/core/_compressor.py index 51915cf..f267faf 100644 --- a/python/ptwm/core/_compressor.py +++ b/python/ptwm/core/_compressor.py @@ -4,12 +4,12 @@ import torch from .. import _rust -from .._config import CompressionConfig, Format, Method -from .._exceptions import CompressionMethodNotSupportedError +from .._config import CompressionConfig, Format, Framing, Method +from .._exceptions import CompressionMethodNotSupportedError, InvalidDTypeError 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,33 @@ 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 + 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) + 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, ) @@ -174,7 +150,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}" @@ -191,13 +170,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: @@ -205,6 +187,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,9 +222,10 @@ 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, + container_layout=self.config.container_layout.value, ) diff --git a/python/ptwm/core/_decompressor.py b/python/ptwm/core/_decompressor.py index 474312f..0c8bd9e 100644 --- a/python/ptwm/core/_decompressor.py +++ b/python/ptwm/core/_decompressor.py @@ -7,6 +7,7 @@ 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 @@ -35,11 +36,14 @@ 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 - come from ``decode_tensor``; shape/dtype/input_format (when present) - come from the CBOR shape-metadata map. Format recovery: + 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 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 @@ -73,19 +77,32 @@ 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): + if mv_data.nbytes < len(MAGIC_PTWS): msg = "too short to be a .ptwm container" raise HeaderParseError(msg) - if bytes(mv_data[: len(magic)]) != magic: + + blob = data if isinstance(data, bytes) else bytes(mv_data) + kind = container_kind(blob) + info: tuple[list[int], str, int, int] | None + if kind == "ptws": + 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 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) - blob = bytes(mv_data) - skip = self.config.skip_missing - raw = _rust.decode_tensor(blob, "_", skip) - expected_hash = _rust.decode_tensor_delta_hash(blob, "_", skip) - if expected_hash is not None: if self.config.delta_second_data is None: msg = ( @@ -109,7 +126,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 @@ -157,15 +173,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/integrations/_compress_safetensors.py b/python/ptwm/integrations/_compress_safetensors.py index 6454649..ed6fee6 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, @@ -26,6 +27,7 @@ from ..preprocessing._chains import ( CHAIN_BYTE_PASSTHROUGH_VALUE, ClassifierRole, + build_chain_blobs, chains_for, ) from ..preprocessing._explorer import ExploreOptions, explore_chains @@ -108,7 +110,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 +173,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( @@ -188,8 +191,9 @@ def compress_safetensors_file( use_user_cache: bool = True, fast: bool = False, fast_chunk: int = 512 * 1024, + container_layout: ContainerLayout | str = "compact", ) -> AuditLog: - """Compress a safetensors file into a sharded ``.ptwm`` directory. + r"""Compress a safetensors file into a sharded ``.ptwm`` directory. Parameters ---------- @@ -209,12 +213,20 @@ 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: ``"compact"`` (default; + ``\x89PTWC``, v1.2 and later) or ``"legacy"`` (``\x89PTWM``, readable + 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. 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) @@ -337,6 +349,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 @@ -371,6 +384,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: @@ -386,6 +400,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/preprocessing/_chains.py b/python/ptwm/preprocessing/_chains.py index 0bdb34d..58a1ead 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,33 @@ 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 +# Largest number of row counts _byte_passthrough_shape tries. +_PASSTHROUGH_ROW_SEARCH = 1 << 16 + + +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) # --------------------------------------------------------------------------- @@ -731,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=[ @@ -787,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=[ @@ -962,6 +997,127 @@ 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 _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], + *, + 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``. 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. + """ + 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) + 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 = _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" + raise RuntimeError(msg) + blobs.append(chain.to_bytes()) + return blobs + + # --------------------------------------------------------------------------- # Production chains table # --------------------------------------------------------------------------- @@ -984,11 +1140,18 @@ def _builder(shape: list[int]) -> Chain | None: # 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, @@ -1073,5 +1236,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/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/sharding/_writer.py b/python/ptwm/sharding/_writer.py index 1d2e030..b827274 100644 --- a/python/ptwm/sharding/_writer.py +++ b/python/ptwm/sharding/_writer.py @@ -7,9 +7,10 @@ from typing import Any from .. import _rust +from .._config import ContainerLayout, container_layout_name from ..preprocessing._chains import ( - CHAIN_BYTE_PASSTHROUGH_VALUE, ClassifierRole, + build_chain_blobs, chains_for, ) from ..utils import DType @@ -17,14 +18,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( @@ -35,15 +41,21 @@ def compress_ptwm_blob( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, + container_layout: ContainerLayout | str = "compact", ) -> 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``: ``"compact"`` (default; ``\x89PTWC``, v1.2 and + 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. """ + layout = container_layout_name(container_layout) records = [] auto_chains: list[list[bytes]] = [] for name, raw, dtype_in, shape in tensors: @@ -59,7 +71,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 @@ -72,6 +84,7 @@ def compress_ptwm_blob( emit_plane_crc, method_hint, streaming_chunk, + container_layout=layout, ) @@ -84,10 +97,13 @@ def write_ptwm_shard( emit_payload_hash: bool = True, emit_plane_crc: bool = True, streaming_chunk: int | None = None, + container_layout: ContainerLayout | str = "compact", ) -> 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, @@ -96,5 +112,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/__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 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..b9ead75 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,35 +77,45 @@ 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 = "compact", ) -> Path: - """Write ``tensors`` into an LMDB environment at ``path``. - - ``layout`` is ``"blob"`` (default) or ``"exploded"``. Returns ``path``. + r"""Write ``tensors`` into an LMDB environment at ``path``. + + ``layout`` is ``"blob"`` (default) or ``"exploded"``. ``container_layout`` + (``"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` + 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')" raise ValueError(msg) + container_layout = container_layout_name(container_layout) 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) 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: @@ -213,6 +224,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", {})) @@ -359,13 +372,17 @@ def put_tensor( *, delta_reference: str | None = None, method_hint: int = 3, + container_layout: ContainerLayout | str | None = None, ) -> 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; ``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" @@ -391,8 +408,13 @@ def put_tensor( else: stored, store_shape, store_dtype = raw, shape, dtype + 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 + [(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..fe6dcf2 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, container_layout_name from ..sharding import compress_ptwm_blob, plan_shards from ._common import ( ExplodedModel, @@ -70,18 +71,28 @@ 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 = "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`` (``"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 + created. """ + container_layout = container_layout_name(container_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/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/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/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/benchmarks/test_container_overhead.py b/tests/ptwm/benchmarks/test_container_overhead.py new file mode 100644 index 0000000..d38246b --- /dev/null +++ b/tests/ptwm/benchmarks/test_container_overhead.py @@ -0,0 +1,265 @@ +"""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, 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``. +""" + +from __future__ import annotations + +import json + +import pytest +import torch +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, + FORMATS, + SIZES, + ShapeFn, + codec_data_bytes, + compress_checked, + distinct_shapes, + make_tensor, + measure, + measure_multi, + non_payload_bytes, + same_shape, + 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 + +# 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: + 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 + assert row.non_payload <= COMPACT_OVERHEAD[dname], 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 + assert row.non_payload <= COMPACT_OVERHEAD[dname], row + 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}" + + +@pytest.mark.parametrize(("dname", "dtype"), DTYPES) +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, + dname, + "normal", + ContainerLayout.LEGACY, + ) + assert row.framing == "ptwm", row + assert 2000 < row.non_payload < 3200, 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 + # 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( + compress_checked(t, framing, layout) + ) + for framing, layout in FORMATS + } + 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( + 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 262 B, 317 B and 238 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) 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..f5a29a6 --- /dev/null +++ b/tests/ptwm/cli/test_container_layout_flag.py @@ -0,0 +1,385 @@ +"""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( + { + # 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), + ) + 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"] + + +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(256, 128, 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() + + +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 + + +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() diff --git a/tests/ptwm/core/test_compact_frame.py b/tests/ptwm/core/test_compact_frame.py new file mode 100644 index 0000000..1c1e5c1 --- /dev/null +++ b/tests/ptwm/core/test_compact_frame.py @@ -0,0 +1,376 @@ +"""Compact single-tensor frame: selection, fallbacks and round trips.""" + +from __future__ import annotations + +import subprocess +import sys +import textwrap + +import numpy as np +import pytest +import torch +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + DecompressionConfig, + Decompressor, + Format, + Framing, + _rust, +) +from ptwm._exceptions import Error, HeaderParseError, InvalidDTypeError + +COMPACT = b"\x89PTWS" +# 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: + 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 + + +@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"), + [ + (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) + 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) + + +@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) + 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)) + + +@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) + + +@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) + 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, 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) + ).compress(t) + 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: + 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])) + + +_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)) + 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) + + +@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/core/test_container_handle.py b/tests/ptwm/core/test_container_handle.py new file mode 100644 index 0000000..bf6dbf6 --- /dev/null +++ b/tests/ptwm/core/test_container_handle.py @@ -0,0 +1,435 @@ +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 numpy as np +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: + # 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") + 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_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. + 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/core/test_zstd.py b/tests/ptwm/core/test_zstd.py index f0d5dbc..a3ad9d2 100644 --- a/tests/ptwm/core/test_zstd.py +++ b/tests/ptwm/core/test_zstd.py @@ -15,8 +15,10 @@ from ptwm import ( CompressionConfig, Compressor, + ContainerLayout, Decompressor, Format, + Framing, Method, ) @@ -69,12 +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") + config = CompressionConfig( + 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/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/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) diff --git a/tests/ptwm/preprocessing/test_chains.py b/tests/ptwm/preprocessing/test_chains.py index d2a0b3d..51a8c77 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, @@ -203,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 == [ @@ -269,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): 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..8e289af --- /dev/null +++ b/tests/ptwm/preprocessing/test_production_chain_fixture.py @@ -0,0 +1,96 @@ +"""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)) + # 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) + ): + 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 diff --git a/tests/ptwm/stores/test_compact_layout.py b/tests/ptwm/stores/test_compact_layout.py new file mode 100644 index 0000000..2e1b01b --- /dev/null +++ b/tests/ptwm/stores/test_compact_layout.py @@ -0,0 +1,404 @@ +"""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 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.sharding._writer import _default_chain_bytes +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_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") + + +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()[:MAGIC_LEN] == 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, 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): + 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, 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]) +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[:MAGIC_LEN] == 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())[:MAGIC_LEN] == 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,), container_layout="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")[:MAGIC_LEN] == MAGIC["compact"] + store.put_tensor( + "d", + raw_of(step), + torch.float32, + (8,), + delta_reference="v", + container_layout="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) + 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")[: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")[: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) + 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 ──────────────────────────────────────────────────────────────── + + +@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 + + +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" + + +# ── 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" + + +# ── 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/stores/test_lmdb.py b/tests/ptwm/stores/test_lmdb.py index 201eec3..d239c0d 100644 --- a/tests/ptwm/stores/test_lmdb.py +++ b/tests/ptwm/stores/test_lmdb.py @@ -7,13 +7,19 @@ import pytest import torch from ptwm import _rust +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 -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" @@ -21,11 +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()) - assert blob[:5] == b"\x89PTWM" + 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 0f3e6f5..2b99b3e 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 @@ -92,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/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/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 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..fdb907a --- /dev/null +++ b/tests/ptwm/wire_format/test_cross_version_builtins.py @@ -0,0 +1,177 @@ +"""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) + + +@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. 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] + 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, + container_layout=layout, + ) + + current = compress(blob) + assert current.startswith(magic) + 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(" 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 diff --git a/tests/ptwm/wire_format/test_extension_table.py b/tests/ptwm/wire_format/test_extension_table.py index fdde250..679eb63 100644 --- a/tests/ptwm/wire_format/test_extension_table.py +++ b/tests/ptwm/wire_format/test_extension_table.py @@ -1,14 +1,31 @@ -"""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, Method +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + 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, + container_layout=ContainerLayout.LEGACY, + ) ).compress(tensor) + assert blob[:5] == b"\x89PTWM" # Header layout (little-endian): # magic: [u8; 9] bytes 0.. 9 @@ -32,3 +49,58 @@ def test_extension_table_is_present_for_huffman_chain(): assert count >= 1, ( "expected at least one Extension Table entry referenced by a Huffman chain" ) + + +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) + + 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") + 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" 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..29e4a38 --- /dev/null +++ b/tests/ptwm/wire_format/test_legacy_hardening.py @@ -0,0 +1,219 @@ +"""Regression tests for reader hardening on legacy ``\\x89PTWM`` containers.""" + +from __future__ import annotations + +import struct + +import pytest +import torch +from ptwm import ( + CompressionConfig, + Compressor, + ContainerLayout, + 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, + container_layout=ContainerLayout.LEGACY, + ) + ).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_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 + + 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(" 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("