JPEG XL: faster EPF, RGBA output, 2D DCTs and LZ77 reads

The edge-preserving filter computes its sigma factors per 8x8 block and its tap
offsets per row instead of per pixel. The RGBA output divides by alpha only for
premultiplied images. 2D DCTs transform the columns in place instead of
transposing squares. LZ77 reads past a copy take a small, inlinable path, and
the gradient and zero predictor channels of single leaf trees decode in their
own small loops. The output is identical.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Lfv41PqHEGbU7SLg5onpG1

Author: matmen

src/codecs/jxl/color.ts

@@ -7,6 +7,9 @@ import { ImageHeader, TF_PQ, TF_HLG, PRI_BT2100, PRI_P3, EXTRA_ALPHA, EXTRA_SPOT
 import { toSRGB, hdrToLinear, ToneMapping, REFERENCE_WHITE } from '../hdr.ts';
 import type { Planes } from './frame.ts';
 
+// a sample (0-1) as a byte
+const byte = (v: number): number => v <= 0 ? 0 : v >= 1 ? 255 : Math.round(v * 255);
+
 // linear light to sRGB like libjxl (TF_SRGB::EncodedFromDisplay, lib/jxl/cms/transfer_functions-inl.h): a rational
 // polynomial of the square root instead of the power, within 5e-7
 const [P0, P1, P2, P3, P4] = [-5.135152395e-04, 5.287254571e-03, 3.903842876e-01, 1.474205315e+00, 7.352629620e-01].map(Math.fround);
@@ -119,16 +122,23 @@ export class ColorOutput {
             out = toSRGB(rgb, this.toneMapping);
         } else {
             out = new Uint8Array(4 * n);
-            const byte = (v: number): number => v <= 0 ? 0 : v >= 1 ? 255 : Math.round(v * 255);
-            for (let i = 0; i < n; i++) {
-                const a = premultiplied ? Math.max(alpha![i], 2 ** -26) : 1;
-                out[4 * i] = byte(r[i] / a);
-                out[4 * i + 1] = byte(g[i] / a);
-                out[4 * i + 2] = byte(b[i] / a);
-                out[4 * i + 3] = 255;
+            if (premultiplied) {
+                for (let i = 0, o = 0; i < n; i++, o += 4) {
+                    const a = Math.max(alpha![i], 2 ** -26);
+                    out[o] = byte(r[i] / a);
+                    out[o + 1] = byte(g[i] / a);
+                    out[o + 2] = byte(b[i] / a);
+                }
+            } else {
+                for (let i = 0, o = 0; i < n; i++, o += 4) {
+                    out[o] = byte(r[i]);
+                    out[o + 1] = byte(g[i]);
+                    out[o + 2] = byte(b[i]);
+                }
             }
         }
-        if (alpha) for (let i = 0; i < n; i++) out[4 * i + 3] = alpha[i] <= 0 ? 0 : alpha[i] >= 1 ? 255 : Math.round(alpha[i] * 255);
+        if (alpha) for (let i = 0, o = 3; i < n; i++, o += 4) out[o] = byte(alpha[i]);
+        else for (let o = 3; o < out.length; o += 4) out[o] = 255;
         return out;
     }
 }

src/codecs/jxl/dct.ts

@@ -97,7 +97,7 @@ export function dct(io: Float64Array, o: number, n: number, scratch: Float64Arra
     }
 }
 
-const work = { row: new Float64Array(256), scratch: new Float64Array(256), col: new Float64Array(256 * 2) };
+const work = { row: new Float64Array(256), scratch: new Float64Array(256) };
 
 const SEC8 = secHalf(8);
 
@@ -205,40 +205,7 @@ export function dct2d(b: Block, inverse: boolean): void {
         return;
     }
     for (let y = 0; y < h; y++) rowDct(y, w);
-    // (the square blocks of the size of the shorter side transposed, then the columns transformed as rows)
-    const size = Math.min(w, h);
-    const transposeSquares = (): void => {
-        for (let by = 0; by < h; by += size) {
-            for (let bx = 0; bx < w; bx += size) {
-                for (let dy = 0; dy < size; dy++) {
-                    for (let dx = dy + 1; dx < size; dx++) {
-                        const p = at(bx + dx, by + dy), q = at(bx + dy, by + dx), t = buf[p];
-                        buf[p] = buf[q];
-                        buf[q] = t;
-                    }
-                }
-            }
-        }
-    };
-    transposeSquares();
-    if (size === h) {
-        for (let y = 0; y < h; y++) for (let x0 = 0; x0 < w; x0 += h) rowDct(y, h, x0);
-    } else {
-        // (rows of w values, h / w of them make a column)
-        const col = work.col, scratch = work.scratch;
-        for (let y = 0; y < w; y++) {
-            for (let k = 0; k < h / size; k++) {
-                const p = at(0, y + k * size);
-                for (let x = 0; x < w; x++) col[k * w + x] = buf[p + x];
-            }
-            dct(col, 0, h, scratch, 0, inverse);
-            for (let k = 0; k < h / size; k++) {
-                const p = at(0, y + k * size);
-                for (let x = 0; x < w; x++) buf[p + x] = col[k * w + x];
-            }
-        }
-    }
-    transposeSquares();
+    for (let x = 0; x < w; x++) colDct(x, h);
 }
 
 // the DCT scale factors of the LF of large varblocks (for c up to 32 of a 256 point DCT)

src/codecs/jxl/entropy.ts

@@ -454,6 +454,8 @@ export class EntropyDecoder {
         const n = (config.splitExponent - inToken + ((token - config.split) >>> inToken)) & 31;
         const low = token & ((1 << config.lsb) - 1);
         const high = (token >>> config.lsb) & ((1 << config.msb) - 1) | 1 << config.msb;
+        // (in integers while it fits in 31 bits)
+        if (n + config.msb + config.lsb < 31) return ((high << n | r.bits(n)) << config.lsb) | low;
         return ((high * 2 ** n + r.bits(n)) * 2 ** config.lsb + low) >>> 0;
     }
 
@@ -464,28 +466,33 @@ export class EntropyDecoder {
 
     readClustered(r: BitReader, cluster: number, distMultiplier = 0): number {
         if (!this.lz77) return this.uint(r, this.configs[cluster], this.symbol(r, cluster));
+        // (small, so that it's inlined: the rest of a copy, else readLz77)
+        if (this.toCopy === 0) return this.readLz77(r, cluster, distMultiplier);
+        const window = this.window as Uint32Array, v = window[this.copyPos++ & WINDOW_MASK];
+        this.toCopy--;
+        window[this.decoded++ & WINDOW_MASK] = v;
+        return v;
+    }
+
+    // a value or the start of a copy, with LZ77
+    private readLz77(r: BitReader, cluster: number, distMultiplier: number): number {
         const window = this.window ??= new Uint32Array(1 << 20);
         let v: number;
-        if (this.toCopy > 0) {
+        const token = this.symbol(r, cluster);
+        if (token >= this.minSymbol) {
+            if (this.decoded === 0) throw new Error('JPEG XL: LZ77 copy before any value');
+            const distCluster = this.clusters[this.clusters.length - 1];
+            this.toCopy = this.uint(r, this.lenConfig as IntegerConfig, token - this.minSymbol) + this.minLength;
+            let distance = this.uint(r, this.configs[distCluster], this.symbol(r, distCluster));
+            if (distMultiplier) {
+                if (distance < 120) distance = Math.max(SPECIAL_DISTANCES[distance * 2] + distMultiplier * SPECIAL_DISTANCES[distance * 2 + 1] - 1, 0);
+                else distance -= 120;
+            }
+            distance = Math.min(Math.min(distance, WINDOW_MASK) + 1, this.decoded);
+            this.copyPos = this.decoded - distance;
             v = window[this.copyPos++ & WINDOW_MASK];
             this.toCopy--;
-        } else {
-            const token = this.symbol(r, cluster);
-            if (token >= this.minSymbol) {
-                if (this.decoded === 0) throw new Error('JPEG XL: LZ77 copy before any value');
-                const distCluster = this.clusters[this.clusters.length - 1];
-                this.toCopy = this.uint(r, this.lenConfig as IntegerConfig, token - this.minSymbol) + this.minLength;
-                let distance = this.uint(r, this.configs[distCluster], this.symbol(r, distCluster));
-                if (distMultiplier) {
-                    if (distance < 120) distance = Math.max(SPECIAL_DISTANCES[distance * 2] + distMultiplier * SPECIAL_DISTANCES[distance * 2 + 1] - 1, 0);
-                    else distance -= 120;
-                }
-                distance = Math.min(Math.min(distance, WINDOW_MASK) + 1, this.decoded);
-                this.copyPos = this.decoded - distance;
-                v = window[this.copyPos++ & WINDOW_MASK];
-                this.toCopy--;
-            } else v = this.uint(r, this.configs[cluster], token);
-        }
+        } else v = this.uint(r, this.configs[cluster], token);
         window[this.decoded++ & WINDOW_MASK] = v;
         return v;
     }

src/codecs/jxl/filters.ts

@@ -62,43 +62,74 @@ function outputRow(cur: Float32Array[], out: Float32Array[], y: number, w: numbe
     // (rows 2 away from the edges: the taps without mirroring)
     const innerRow = rows[0] === (y - 2) * w && rows[4] === (y + 2) * w;
     const row = y * w;
-    for (let x = 0; x < w; x++) {
-        const i = row + x;
-        const sg = sigma ? sigma[sigmaRow + (x >> 3)] : sigmaModular;
+    // per tap pair: the offsets of +u in the sums and in the row
+    const sumOff = TAPS.subarray(0, n), srcOff = TAPS.subarray(8, 8 + n);
+    for (let k = 0; k < n; k++) {
+        sumOff[k] = u[2 * k + 1] * pw + u[2 * k];
+        srcOff[k] = u[2 * k + 1] * w + u[2 * k];
+    }
+    for (let x0 = 0; x0 < w; x0 += 8) {
+        const x1 = Math.min(x0 + 8, w);
+        const sg = sigma ? sigma[sigmaRow + (x0 >> 3)] : sigmaModular;
         if (sg < 0.3) {
-            o0[i] = i0[i];
-            o1[i] = i1[i];
-            o2[i] = i2[i];
+            o0.set(i0.subarray(row + x0, row + x1), row + x0);
+            o1.set(i1.subarray(row + x0, row + x1), row + x0);
+            o2.set(i2.subarray(row + x0, row + x1), row + x0);
             continue;
         }
-        const sm = yBorder || (x & 7) === 0 || (x & 7) === 7 ? stepMul * borderMul : stepMul;
-        const negInvSigma = 6.6 * (Math.SQRT1_2 - 1) / sg * sm;
-        let sumW = 1, sum0 = i0[i], sum1 = i1[i], sum2 = i2[i];
-        const p = p0 + x, inner = innerRow && x >= 2 && x < w - 2;
-        for (let k = 0; k < n; k++) {
-            const ux = u[2 * k], uy = u[2 * k + 1], sum = sums[k];
-            let weight = 1 + sum[p] * negInvSigma;
-            if (weight > 0) {
-                const t = inner ? i + uy * w + ux : rows[2 + uy] + xs[x + 5 + ux];
-                sumW += weight;
-                sum0 += weight * i0[t];
-                sum1 += weight * i1[t];
-                sum2 += weight * i2[t];
-            }
-            weight = 1 + sum[p - uy * pw - ux] * negInvSigma;
-            if (weight > 0) {
-                const t = inner ? i - uy * w - ux : rows[2 - uy] + xs[x + 5 - ux];
-                sumW += weight;
-                sum0 += weight * i0[t];
-                sum1 += weight * i1[t];
-                sum2 += weight * i2[t];
+        // (the block's columns 0 and 7, and its rows 0 and 7, are weighted by borderMul)
+        const borderInvSigma = 6.6 * (Math.SQRT1_2 - 1) / sg * (stepMul * borderMul), innerInvSigma = 6.6 * (Math.SQRT1_2 - 1) / sg * stepMul;
+        for (let x = x0; x < x1; x++) {
+            const i = row + x;
+            const negInvSigma = yBorder || (x & 7) === 0 || (x & 7) === 7 ? borderInvSigma : innerInvSigma;
+            let sumW = 1, sum0 = i0[i], sum1 = i1[i], sum2 = i2[i];
+            const p = p0 + x;
+            if (innerRow && x >= 2 && x < w - 2) {
+                for (let k = 0; k < n; k++) {
+                    const sum = sums[k], so = sumOff[k], to = srcOff[k];
+                    let weight = 1 + sum[p] * negInvSigma;
+                    if (weight > 0) {
+                        sumW += weight;
+                        sum0 += weight * i0[i + to];
+                        sum1 += weight * i1[i + to];
+                        sum2 += weight * i2[i + to];
+                    }
+                    weight = 1 + sum[p - so] * negInvSigma;
+                    if (weight > 0) {
+                        sumW += weight;
+                        sum0 += weight * i0[i - to];
+                        sum1 += weight * i1[i - to];
+                        sum2 += weight * i2[i - to];
+                    }
+                }
+            } else {
+                for (let k = 0; k < n; k++) {
+                    const ux = u[2 * k], uy = u[2 * k + 1], sum = sums[k];
+                    let weight = 1 + sum[p] * negInvSigma;
+                    if (weight > 0) {
+                        const t = rows[2 + uy] + xs[x + 5 + ux];
+                        sumW += weight;
+                        sum0 += weight * i0[t];
+                        sum1 += weight * i1[t];
+                        sum2 += weight * i2[t];
+                    }
+                    weight = 1 + sum[p - uy * pw - ux] * negInvSigma;
+                    if (weight > 0) {
+                        const t = rows[2 - uy] + xs[x + 5 - ux];
+                        sumW += weight;
+                        sum0 += weight * i0[t];
+                        sum1 += weight * i1[t];
+                        sum2 += weight * i2[t];
+                    }
+                }
             }
+            o0[i] = sum0 / sumW;
+            o1[i] = sum1 / sumW;
+            o2[i] = sum2 / sumW;
         }
-        o0[i] = sum0 / sumW;
-        o1[i] = sum1 / sumW;
-        o2[i] = sum2 / sumW;
     }
 }
+const TAPS = new Int32Array(16);
 
 /**
  * The edge-preserving filter on 3 w x h channels

src/codecs/jxl/modular.ts

@@ -647,6 +647,23 @@ export function decodeModular(r: BitReader, image: ModularImage, stream: number,
     if (undo) image.undo();
 }
 
+// the channels of a tree with a single leaf (small functions, so that the entropy decoding is inlined): no prediction
+function decodeZero(r: BitReader, dec: EntropyDecoder, cluster: number, data: Int32Array, mult: number, offset: number, distMultiplier: number): void {
+    for (let i = 0; i < data.length; i++) data[i] = (Math.imul(unpackSigned(dec.readClustered(r, cluster, distMultiplier)), mult) + offset) | 0;
+}
+
+// the gradient predictor (with the first row and column from the left and top)
+function decodeGradient(r: BitReader, dec: EntropyDecoder, cluster: number, data: Int32Array, w: number, h: number, distMultiplier: number): void {
+    for (let x = 0; x < w; x++) data[x] = (unpackSigned(dec.readClustered(r, cluster, distMultiplier)) + (x ? data[x - 1] : 0)) | 0;
+    for (let y = 1, i = w; y < h; y++) {
+        data[i] = (unpackSigned(dec.readClustered(r, cluster, distMultiplier)) + data[i - w]) | 0;
+        i++;
+        for (let x = 1; x < w; x++, i++) {
+            data[i] = (unpackSigned(dec.readClustered(r, cluster, distMultiplier)) + clampedGradient(data[i - w], data[i - 1], data[i - w - 1])) | 0;
+        }
+    }
+}
+
 function decodeChannel(r: BitReader, dec: EntropyDecoder, globalTree: Tree, image: ModularImage, chan: number, stream: number, distMultiplier: number): void {
     const c = image.channels[chan], { w, h } = c, data = c.data;
     const t = specialize(globalTree, chan, stream);
@@ -670,16 +687,11 @@ function decodeChannel(r: BitReader, dec: EntropyDecoder, globalTree: Tree, imag
                 data.fill((Math.imul(unpackSigned(single), mult) + offset) | 0);
                 return;
             }
-            for (let i = 0; i < w * h; i++) data[i] = (Math.imul(unpackSigned(dec.readClustered(r, cluster, distMultiplier)), mult) + offset) | 0;
+            decodeZero(r, dec, cluster, data, mult, offset, distMultiplier);
             return;
         }
         if (pred === 5 && mult === 1 && offset === 0) {
-            for (let y = 0, i = 0; y < h; y++) {
-                for (let x = 0; x < w; x++, i++) {
-                    const left = x ? data[i - 1] : y ? data[i - w] : 0, top = y ? data[i - w] : left, topleft = x && y ? data[i - w - 1] : left;
-                    data[i] = (unpackSigned(dec.readClustered(r, cluster, distMultiplier)) + clampedGradient(top, left, topleft)) | 0;
-                }
-            }
+            decodeGradient(r, dec, cluster, data, w, h, distMultiplier);
             return;
         }
     }