Carry JPEG tokens and WebP bits at their natural width

CHANGELOG.md

@@ -73,6 +73,20 @@ in the README).
   - The native rasterizer's serial point and segment passes no longer
     materialize an `0..n` index vector to read back in order (4 bytes per mark),
     and the per-point density quantizer keeps one temporary instead of three.
+  - The JPEG token pipeline carries every field at its natural width. It held
+    ~20 parallel `int64` arrays over the nonzero coefficients — positions that
+    top out at 62, categories at 12, symbols that are a byte by definition —
+    while the coefficients themselves fit `int16`, because an orthonormal 8×8
+    DCT of level-shifted 8-bit samples is bounded by 1024 and quantizers are
+    ≥ 1. The RGB→YCbCr transform also promoted the whole frame to interleaved
+    float before splitting it into planes. A 3200×2400 photographic encode now
+    peaks at **200 MB instead of 731 MB** and is ~3% faster; the chart-shaped
+    export in the memory suite drops 418 MB → 205 MB.
+  - The lossless WebP packer scatters bits in bounded entry blocks instead of
+    one masked pass per bit position over the whole token stream (five machine
+    words per entry, per position), and the header now rides in the token
+    buffer rather than being concatenated onto the front of it. A 3200×2400
+    export peaks at 135 MB instead of 216 MB at unchanged speed.
 - `memory_report()` reports `capacity_bytes` per column and
   `canonical_capacity_bytes` per store, and builds `resident_array_bytes` from
   the capacity total. A streamed column's `values` is a prefix view of its

python/xy/_jpeg.py

@@ -128,11 +128,16 @@ def _zigzag_order() -> np.ndarray:
 
 
 def _huff_lookup(bits: bytes, values: bytes) -> tuple[np.ndarray, np.ndarray]:
-    """Canonical symbol→(code, length) arrays from a BITS/HUFFVAL spec."""
+    """Canonical symbol→(code, length) arrays from a BITS/HUFFVAL spec.
+
+    Codes are at most 16 bits and lengths at most 16, so the lookup tables are
+    int32/uint8 rather than int64: every token gathers one entry from each, and
+    the gathered arrays are the widest thing the entropy stage carries.
+    """
     if len(values) != sum(bits):
         raise AssertionError("Huffman spec mismatch: HUFFVAL count != sum(BITS)")
-    codes = np.zeros(256, dtype=np.int64)
-    lens = np.zeros(256, dtype=np.int64)
+    codes = np.zeros(256, dtype=np.int32)
+    lens = np.zeros(256, dtype=np.uint8)
     code = 0
     k = 0
     for length in range(1, 17):
@@ -175,7 +180,8 @@ def _dct_matrix() -> np.ndarray:
 
 
 def _bit_size(magnitude: np.ndarray) -> np.ndarray:
-    return np.searchsorted(_POW2, magnitude, side="right").astype(np.int64)
+    # Categories top out at 12, so the result is a byte per token, not 8.
+    return np.searchsorted(_POW2, magnitude, side="right").astype(np.uint8)
 
 
 def _scaled_quant(base: np.ndarray, quality: int) -> np.ndarray:
@@ -192,69 +198,108 @@ def _component_tokens(
     Returns parallel arrays (block, seq, table, symbol, amplitude, amp_bits);
     `seq` orders tokens within a block (0 = DC, 255 = EOB) so a single sort on
     (block, component, seq) later interleaves the MCU stream.
+
+    Every field is carried at its natural width rather than as an int64. A
+    photographic image mints millions of tokens and this pipeline holds ~20
+    parallel arrays at once, so width is the whole cost: positions and runs are
+    at most 62, categories at most 12, symbols are a byte by definition, and
+    amplitudes fit int16 because an orthonormal 8x8 DCT of level-shifted 8-bit
+    samples is bounded by sqrt(64) * 128 = 1024 before quantization (and
+    quantizers are >= 1). Only the block index needs more than two bytes.
     """
     n = coef.shape[0]
 
-    # DC is coded differentially along the component's block sequence.
-    diff = np.diff(coef[:, 0], prepend=np.int64(0))
+    # DC is coded differentially along the component's block sequence. These are
+    # per-block, not per-token, so they stay small regardless of width.
+    diff = np.diff(coef[:, 0], prepend=np.int16(0))
     dsize = _bit_size(np.abs(diff))
     # T.81 amplitude coding: negatives are sent as v + 2**size - 1.
-    dampl = np.where(diff < 0, diff + (np.int64(1) << dsize) - 1, diff)
+    dampl = np.where(diff < 0, diff + (np.int16(1) << dsize) - 1, diff).astype(np.int16)
     dc_tok = (
-        np.arange(n, dtype=np.int64),
-        np.zeros(n, dtype=np.int64),
-        np.full(n, dc_tbl, dtype=np.int64),
+        np.arange(n, dtype=np.int32),
+        np.zeros(n, dtype=np.uint8),
+        np.full(n, dc_tbl, dtype=np.uint8),
         dsize,
         dampl,
         dsize,
     )
 
     ac = coef[:, 1:]
     blk, pos = np.nonzero(ac)  # row-major: block-ascending, position-ascending
-    last = np.full(n, -1, dtype=np.int64)
+    last = np.full(n, -1, dtype=np.int16)
     if blk.size:
         last[blk] = pos  # row-major order → last write per block wins
         val = ac[blk, pos]
+        # np.nonzero hands back two intp vectors; narrow both before anything
+        # else is derived from them, so the int64 pair is transient rather than
+        # the base width of the whole pipeline.
+        blk = blk.astype(np.int32)
+        pos = pos.astype(np.uint8)
         first = np.empty(blk.size, dtype=bool)
         first[0] = True
         np.not_equal(blk[1:], blk[:-1], out=first[1:])
-        prev = np.where(first, np.int64(-1), np.concatenate((pos[:1] * 0 - 1, pos[:-1])))
-        run = pos - prev - 1
+        # Previous nonzero position within the block, -1 at a block's first.
+        # Built in place: the shift-and-mask form allocated a concatenate and a
+        # where, and an unsigned `pos[:1] * 0 - 1` would wrap to 255.
+        prev = np.empty(pos.size, dtype=np.int8)
+        prev[0] = -1
+        prev[1:] = pos[:-1]
+        np.copyto(prev, np.int8(-1), where=first)
+        run = (pos - prev - np.int16(1)).astype(np.uint8)
+        del prev
         zrl = run >> 4  # each 16 zeros of run becomes a ZRL (0xF0) token
         asize = _bit_size(np.abs(val))
         sym = ((run & 15) << 4) | asize
-        aampl = np.where(val < 0, val + (np.int64(1) << asize) - 1, val)
+        del run
+        aampl = np.where(val < 0, val + (np.int16(1) << asize) - 1, val).astype(np.int16)
+        del val
         # Expand each nonzero into its ZRL prefix + the coefficient token,
         # numbering tokens within their block via a segmented cumsum: `start`
         # is the global exclusive cumsum of token counts and `base` forward-
         # fills each block's opening value, so `start - base` restarts at 0.
-        tot = zrl + 1
-        cum = np.cumsum(tot)
+        tot = zrl + np.uint8(1)
+        cum = np.cumsum(tot, dtype=np.int32)
         start = cum - tot
-        base = np.maximum.accumulate(np.where(first, start, 0))
-        rep = np.repeat(np.arange(blk.size, dtype=np.int64), tot)
-        j = np.arange(cum[-1], dtype=np.int64) - np.repeat(start, tot)
+        ntok_ac = int(cum[-1])
+        del cum
+        base = np.maximum.accumulate(np.where(first, start, np.int32(0)))
+        del first
+        rep = np.repeat(np.arange(blk.size, dtype=np.int32), tot)
+        # `j` only ever counts within one nonzero's ZRL run, so it is <= 3.
+        j = (np.arange(ntok_ac, dtype=np.int32) - np.repeat(start, tot)).astype(np.uint8)
+        del tot
         is_zrl = j < zrl[rep]
+        seq = (start - base)[rep]
+        del start, base
+        seq += 1
+        seq += j
         ac_tok = (
-            blk[rep].astype(np.int64),
-            1 + (start - base)[rep] + j,
-            np.full(rep.size, ac_tbl, dtype=np.int64),
-            np.where(is_zrl, np.int64(0xF0), sym[rep]),
-            np.where(is_zrl, np.int64(0), aampl[rep]),
-            np.where(is_zrl, np.int64(0), asize[rep]),
+            blk[rep],
+            seq.astype(np.uint8),
+            np.full(rep.size, ac_tbl, dtype=np.uint8),
+            np.where(is_zrl, np.uint8(0xF0), sym[rep]),
+            np.where(is_zrl, np.int16(0), aampl[rep]),
+            np.where(is_zrl, np.uint8(0), asize[rep]),
         )
     else:
-        ac_tok = tuple(np.empty(0, dtype=np.int64) for _ in range(6))
+        ac_tok = (
+            np.empty(0, dtype=np.int32),
+            np.empty(0, dtype=np.uint8),
+            np.empty(0, dtype=np.uint8),
+            np.empty(0, dtype=np.uint8),
+            np.empty(0, dtype=np.int16),
+            np.empty(0, dtype=np.uint8),
+        )
 
     # EOB unless the block's final zigzag coefficient (AC position 62) is set.
-    eob = np.flatnonzero(last != 62).astype(np.int64)
+    eob = np.flatnonzero(last != 62).astype(np.int32)
     eob_tok = (
         eob,
-        np.full(eob.size, 255, dtype=np.int64),
-        np.full(eob.size, ac_tbl, dtype=np.int64),
-        np.zeros(eob.size, dtype=np.int64),
-        np.zeros(eob.size, dtype=np.int64),
-        np.zeros(eob.size, dtype=np.int64),
+        np.full(eob.size, 255, dtype=np.uint8),
+        np.full(eob.size, ac_tbl, dtype=np.uint8),
+        np.zeros(eob.size, dtype=np.uint8),
+        np.zeros(eob.size, dtype=np.int16),
+        np.zeros(eob.size, dtype=np.uint8),
     )
     merged = [np.concatenate(parts) for parts in zip(dc_tok, ac_tok, eob_tok, strict=True)]
     return merged[0], merged[1], merged[2], merged[3], merged[4], merged[5]
@@ -267,6 +312,11 @@ def _component_tokens(
 #: leftover bits into the next one, which keeps the byte stream identical.
 _ENTROPY_BIT_CHUNK = 1 << 18
 
+#: Pixels per band of the RGB→YCbCr transform. The transform is elementwise, so
+#: banding it is bit-identical; this bounds the interleaved f32 promotion that
+#: would otherwise be three floats per pixel of the whole frame.
+_YCBCR_PIXEL_CHUNK = 1 << 18
+
 
 def _pack_entropy(chunk: np.ndarray, nbits: np.ndarray) -> bytes:
     """MSB-first bit packing of (value, bit-count) chunks, with 1-padding to a
@@ -278,7 +328,7 @@ def _pack_entropy(chunk: np.ndarray, nbits: np.ndarray) -> bytes:
     local, so it applies per group.
     """
     total = int(nbits.sum())
-    cum = np.cumsum(nbits)
+    cum = np.cumsum(nbits, dtype=np.int64)
     # Token index at each bit-budget boundary: a group is [t0, t1) tokens, and a
     # single token never spans groups (its bits stay contiguous, as before).
     edges = np.searchsorted(
@@ -293,7 +343,9 @@ def _pack_entropy(chunk: np.ndarray, nbits: np.ndarray) -> bytes:
         values = chunk[t0:t1]
         widths = nbits[t0:t1]
         group_total = int(widths.sum())
-        start = np.cumsum(widths) - widths
+        # Pinned to int64: bit widths are uint8, and an unsigned accumulator
+        # would make `arange(int64) - start` promote to float64 under NEP 50.
+        start = np.cumsum(widths, dtype=np.int64) - widths
         idx = np.repeat(np.arange(values.size, dtype=np.int64), widths)
         offset = np.arange(group_total, dtype=np.int64) - np.repeat(start, widths)
         bits = ((values[idx] >> (widths[idx] - 1 - offset)) & 1).astype(np.uint8)
@@ -395,17 +447,26 @@ def encode(rgba: np.ndarray, *, quality: int = 90) -> bytes:
 
     # Alpha is ignored: the caller has already composited onto an opaque
     # background, so only the RGB planes carry information.
-    rgb = rgba[..., :3].astype(np.float32)
-    r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
+    #
     # JFIF BT.601 full-range transform. The −128 DCT level shift cancels the
     # +128 chroma offset, so Y is shifted here and Cb/Cr are left centered.
-    y = 0.299 * r + 0.587 * g + 0.114 * b - 128.0
-    cb = -0.168736 * r - 0.331264 * g + 0.5 * b
-    cr = 0.5 * r - 0.418688 * g - 0.081312 * b
-    # The interleaved f32 source is dead once the planes exist, and it is three
-    # floats per pixel — hand it back before the per-component pipeline runs.
-    # (`r`/`g`/`b` are views into it, so all four references have to go.)
-    del rgb, r, g, b
+    #
+    # Done in row bands: the transform is elementwise, so a band produces
+    # bit-identical f32 output, but promoting the whole frame to interleaved
+    # float first costs three floats per pixel on top of the three planes —
+    # 92 MB of transient on a 3200x2400 export, more than the planes themselves.
+    y = np.empty((h, w), dtype=np.float32)
+    cb = np.empty((h, w), dtype=np.float32)
+    cr = np.empty((h, w), dtype=np.float32)
+    band = max(1, _YCBCR_PIXEL_CHUNK // w)
+    for r0 in range(0, h, band):
+        r1 = min(r0 + band, h)
+        rgb = rgba[r0:r1, :, :3].astype(np.float32)
+        r, g, b = rgb[..., 0], rgb[..., 1], rgb[..., 2]
+        y[r0:r1] = 0.299 * r + 0.587 * g + 0.114 * b - 128.0
+        cb[r0:r1] = -0.168736 * r - 0.331264 * g + 0.5 * b
+        cr[r0:r1] = 0.5 * r - 0.418688 * g - 0.081312 * b
+        del rgb, r, g, b
 
     qy = _scaled_quant(_QUANT_LUMA, quality)
     qc = _scaled_quant(_QUANT_CHROMA, quality)
@@ -450,7 +511,10 @@ def encode(rgba: np.ndarray, *, quality: int = 90) -> bytes:
         scaled += bias
         del bias
         np.trunc(scaled, out=scaled)
-        quant = scaled.astype(np.int64)
+        # int16 is exact here and halves every array the token stage derives:
+        # an orthonormal DCT of level-shifted 8-bit samples is bounded by 1024
+        # in magnitude and quantizers are >= 1, so |coefficient| <= 1024.
+        quant = scaled.astype(np.int16)
         del scaled
         # Exact-math AC magnitudes cap at 1020 (category 10); clamp is one-LSB
         # insurance against float rounding ever minting category 11, which the
@@ -461,7 +525,9 @@ def encode(rgba: np.ndarray, *, quality: int = 90) -> bytes:
         del quant
         # 4:4:4 → one block per component per MCU, so the MCU-interleaved
         # order Y, Cb, Cr is a stable sort on (block, component, in-block seq).
-        keys.append(((toks[0] * 3 + comp) << 8) | toks[1])
+        # The key is widened deliberately: the maximum 65535x65535 frame has 67M
+        # blocks, and `block * 3 << 8` leaves int32 well before that.
+        keys.append(((toks[0].astype(np.int64) * 3 + comp) << 8) | toks[1])
         # Only the four emitted fields outlive the key: the block and in-block
         # sequence numbers are folded into it, and on a photographic image each
         # dropped field is an int64 per token.

python/xy/_webp.py

@@ -208,18 +208,32 @@ def _write_prefix_code(put, hist: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     return emit_len, emit_code
 
 
+#: Entries per bit-scatter pass. Each pass builds a boolean mask and three
+#: compacted gathers over the entries it covers, so an unbounded pass costs
+#: ~5 machine words per entry *per bit position* -- 216 MB of transient on a
+#: 3200x2400 export. Bounding it changes nothing about which bits are written
+#: or where, because every entry scatters to absolute positions in `bits`.
+_PACK_ENTRY_CHUNK = 1 << 17
+
+
 def _pack_lsb(values: np.ndarray, nbits: np.ndarray) -> bytes:
     """Pack (value, nbits) pairs into an LSB-first byte stream.
 
     Expands to one uint8 per bit and lets packbits fold them: one vectorized
-    pass per bit *position* (<= ~40) instead of a Python loop per symbol.
+    pass per bit *position* (<= ~40) instead of a Python loop per symbol. The
+    passes run over bounded entry blocks, and each block only iterates as far
+    as its own widest entry.
     """
-    ends = np.cumsum(nbits, dtype=np.int64)
-    starts = ends - nbits
-    bits = np.zeros(int(ends[-1]), np.uint8)
-    for k in range(int(nbits.max())):
-        m = nbits > k
-        bits[starts[m] + k] = (values[m] >> np.uint64(k)) & np.uint64(1)
+    starts = np.cumsum(nbits, dtype=np.int64)
+    total = int(starts[-1]) if starts.size else 0
+    starts -= nbits  # exclusive scan, in place: the ends array is dead here
+    bits = np.zeros(total, np.uint8)
+    for i0 in range(0, values.size, _PACK_ENTRY_CHUNK):
+        i1 = min(i0 + _PACK_ENTRY_CHUNK, values.size)
+        vals, widths, offs = values[i0:i1], nbits[i0:i1], starts[i0:i1]
+        for k in range(int(widths.max())):
+            m = widths > k
+            bits[offs[m] + k] = (vals[m] >> np.uint64(k)) & np.uint64(1)
     return np.packbits(bits, bitorder="little").tobytes()
 
 
@@ -293,10 +307,18 @@ def put(v: int, n: int) -> None:
     # --- token stream as (value, nbits) entries: two per literal (green+red,
     # blue+alpha packed pairwise, <=30 bits each) and one per reference
     # (length code + extra bits + distance code, <=40 bits).
+    # The header rides in the same buffer rather than being concatenated onto
+    # the front of it: `put` has finished by now, so its length is known, and
+    # joining afterwards would hold a second copy of the whole token stream.
+    nhead = len(head_vals)
     per_seg = 2 + nref
-    offsets = np.cumsum(per_seg) - per_seg
-    ev = np.zeros(int(per_seg.sum()), np.uint64)
+    offsets = np.cumsum(per_seg)
+    offsets -= per_seg
+    offsets += nhead
+    ev = np.zeros(nhead + int(per_seg.sum()), np.uint64)
     eb = np.zeros(ev.size, np.uint8)
+    ev[:nhead] = head_vals
+    eb[:nhead] = head_bits
     gi, ri, bi, ai = lit[:, 1], lit[:, 0], lit[:, 2], lit[:, 3]
     ev[offsets] = g_code[gi] | (r_code[ri] << g_len[gi].astype(np.uint64))
     eb[offsets] = g_len[gi] + r_len[ri]
@@ -312,10 +334,7 @@ def put(v: int, n: int) -> None:
         )
         eb[pos] = g_len[ref_sym] + _LP_EBITS[run] + d_len[1]
 
-    payload = _pack_lsb(
-        np.concatenate([np.asarray(head_vals, np.uint64), ev]),
-        np.concatenate([np.asarray(head_bits, np.uint8), eb]),
-    )
+    payload = _pack_lsb(ev, eb)
     chunk = b"VP8L" + struct.pack("<I", len(payload)) + payload
     if len(payload) & 1:
         chunk += b"\x00"  # RIFF chunks are even-aligned; pad byte is unsized

python/xy/export.py

@@ -921,15 +921,33 @@ def _export_dimensions(
     return w, h
 
 
+#: Pixels per band of the alpha flatten. The composite is elementwise, so the
+#: band size cannot change a byte; it bounds the six uint16 full-frame
+#: temporaries the one-shot form held (117 MB on a 3200x2400 export).
+_FLATTEN_PIXEL_CHUNK = 1 << 18
+
+
 def _flatten_alpha(rgba: "Any") -> "Any":
-    """Composite leftover alpha over white — the JPEG determinism backstop."""
+    """Composite leftover alpha over white — the JPEG determinism backstop.
+
+    Returns `(h, w, 3)`: both callers hand the result straight to the JPEG
+    encoder, which ignores alpha, so a fourth opaque plane is a quarter of a
+    frame carried for nothing. Composited in row bands — `rgb * a + 255 *
+    (255 - a) + 127` peaks at 65152 for any 8-bit input, so the uint16
+    intermediates are exact, and banding an elementwise expression is
+    bit-identical.
+    """
     import numpy as np
 
-    alpha = rgba[..., 3:4].astype(np.uint16)
-    rgb = (rgba[..., :3].astype(np.uint16) * alpha + 255 * (255 - alpha) + 127) // 255
-    out = np.empty_like(rgba)
-    out[..., :3] = rgb.astype(np.uint8)
-    out[..., 3] = 255
+    h, w = rgba.shape[0], rgba.shape[1]
+    out = np.empty((h, w, 3), dtype=np.uint8)
+    band = max(1, _FLATTEN_PIXEL_CHUNK // max(w, 1))
+    for r0 in range(0, h, band):
+        r1 = min(r0 + band, h)
+        chunk = rgba[r0:r1]
+        alpha = chunk[..., 3:4].astype(np.uint16)
+        rgb = (chunk[..., :3].astype(np.uint16) * alpha + 255 * (255 - alpha) + 127) // 255
+        out[r0:r1] = rgb.astype(np.uint8)
     return out