feat(pco): Phase 2 — Classic/None/I64 tANS decode

- PcoBin: bin record (weight, lower, offsetBits)
- PcoTansDecoder: spread_state_symbols + tANS node table;
  decodePage writes raw U64 latents to MemorySegment
- PcoEncoding.Decoder: reads chunk meta bits, builds tANS
  table per chunk, decodes pages; converts U64→I64 via
  sign-bit flip (from_latent_ordered); Phase 3+ ptypes and
  delta=Consecutive throw clear "not yet implemented"

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: dfa1

core/src/main/java/io/github/dfa1/vortex/encoding/PcoBin.java

@@ -0,0 +1,9 @@
+package io.github.dfa1.vortex.encoding;
+
+/// One bin in a pco latent variable: a numerical range [lower, lower + 2^offsetBits).
+///
+/// {@code weight} is the bin's count in the tANS table (sum of weights == table size).
+/// {@code lower} is the raw unsigned lower bound (U64 for 64-bit latents).
+/// {@code offsetBits} is the log2 of the range size (0 = single value).
+record PcoBin(int weight, long lower, int offsetBits) {
+}

core/src/main/java/io/github/dfa1/vortex/encoding/PcoEncoding.java

@@ -1,12 +1,18 @@
 package io.github.dfa1.vortex.encoding;
 
 import com.google.protobuf.InvalidProtocolBufferException;
+import io.github.dfa1.vortex.core.ArrayStats;
 import io.github.dfa1.vortex.core.DType;
+import io.github.dfa1.vortex.core.PType;
 import io.github.dfa1.vortex.core.VortexException;
 import io.github.dfa1.vortex.core.array.Array;
+import io.github.dfa1.vortex.core.array.LongArray;
 import io.github.dfa1.vortex.proto.EncodingProtos;
 
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.ValueLayout;
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 
 /// Decoder for {@code vortex.pco} (pcodec numerical compression).
 ///
@@ -19,20 +25,26 @@
 ///
 /// <p>Wire format (pcodec layer, per chunk/page):
 /// <ul>
-///   <li>Chunk meta: mode nibble + extra mode bits + delta nibble + extra delta bits +
-///       per-latent: ans_size_log (4b), bin_count (15b), per-bin {weight-1, lower, offset_bits}</li>
-///   <li>Page: initial latent state (delta state_n + 4 tANS state indices) → byte align →
-///       per 256-batch: tANS-decoded bin indices + offset bits</li>
+///   <li>Chunk meta: [4b mode][extra mode bits][4b delta][extra delta bits]
+///       [per-latent: 4b ans_size_log, 15b n_bins, per-bin {weight-1, lower, offset_bits}]
+///       [0–7b alignment]</li>
+///   <li>Page: [4 × ans_size_log b initial states][0–7b alignment]
+///       [per 256-batch: ANS bits for all k, then offset bits for all k]</li>
 ///   <li>All bit packing little-endian (LSB first)</li>
 /// </ul>
 ///
-/// <p>Phase 1: skeleton only — parses metadata, validates header, dispatches on PType.
-/// Phase 2 adds Classic/None decode for I64; later phases extend to all ptypes and modes.
+/// <p>Supported (Phase 2): Classic mode, None delta, non-null, I64.
+/// Other modes/deltas/ptypes throw with a clear "not yet implemented" message.
 public final class PcoEncoding implements Encoding {
 
     static final byte PCO_FORMAT_MAJOR = 0x04;
     static final byte PCO_FORMAT_MINOR = 0x01;
 
+    // bits needed to encode offset_bits field per latent type
+    static final int BITS_TO_ENCODE_OFFSET_BITS_64 = 7; // log2(64) + 1
+    static final int BITS_TO_ENCODE_OFFSET_BITS_32 = 6; // log2(32) + 1
+    static final int BITS_TO_ENCODE_OFFSET_BITS_16 = 5; // log2(16) + 1
+
     @Override
     public EncodingId encodingId() {
         return EncodingId.VORTEX_PCO;
@@ -58,12 +70,92 @@ static EncodeResult encode(DType dtype, Object data) {
 
     static final class Decoder {
 
+        private static final ValueLayout.OfLong LE_LONG =
+                ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
+
         static Array decode(DecodeContext ctx) {
             EncodingProtos.PcoMetadata meta = parseMeta(ctx);
             validateHeader(meta);
-            throw new VortexException(EncodingId.VORTEX_PCO,
-                    "pco decode not yet implemented — Phase 2 pending (chunks="
-                            + meta.getChunksCount() + ")");
+
+            DType dtype = ctx.dtype();
+            if (!(dtype instanceof DType.Primitive dt)) {
+                throw new VortexException(EncodingId.VORTEX_PCO,
+                        "pco decode requires Primitive dtype, got: " + dtype);
+            }
+            PType ptype = dt.ptype();
+            if (ptype != PType.I64) {
+                throw new VortexException(EncodingId.VORTEX_PCO,
+                        "pco decode Phase 2: only I64 supported, got: " + ptype);
+            }
+
+            long n = ctx.rowCount();
+            MemorySegment out = ctx.arena().allocate(n * Long.BYTES);
+
+            int nChunks = meta.getChunksCount();
+            int bufIdx = 0;
+            long outByteOffset = 0L;
+
+            for (int c = 0; c < nChunks; c++) {
+                EncodingProtos.PcoChunkInfo chunkInfo = meta.getChunks(c);
+                MemorySegment chunkMetaBuf = ctx.buffer(bufIdx++);
+
+                PcoChunkMeta chunkMeta = readChunkMeta(chunkMetaBuf);
+                PcoTansDecoder tans = PcoTansDecoder.build(chunkMeta.ansSizeLog(), chunkMeta.bins());
+
+                int nPages = chunkInfo.getPagesCount();
+                for (int p = 0; p < nPages; p++) {
+                    int pageN = chunkInfo.getPages(p).getNValues();
+                    MemorySegment pageBuf = ctx.buffer(bufIdx++);
+
+                    LeBitReader pageReader = new LeBitReader(pageBuf);
+                    int[] stateIdxs = new int[PcoTansDecoder.ANS_INTERLEAVING];
+                    for (int i = 0; i < PcoTansDecoder.ANS_INTERLEAVING; i++) {
+                        stateIdxs[i] = (int) pageReader.readBits(chunkMeta.ansSizeLog());
+                    }
+                    pageReader.alignToByte();
+
+                    tans.decodePage(pageReader, stateIdxs, pageN, out, outByteOffset);
+                    outByteOffset += (long) pageN * Long.BYTES;
+                }
+            }
+
+            // Convert U64 latents → I64: flip sign bit (from_latent_ordered for signed types)
+            for (long i = 0; i < n; i++) {
+                long byteOff = i * Long.BYTES;
+                out.set(LE_LONG, byteOff, out.get(LE_LONG, byteOff) ^ Long.MIN_VALUE);
+            }
+
+            return new LongArray(dtype, n, out, ArrayStats.empty());
+        }
+
+        private static PcoChunkMeta readChunkMeta(MemorySegment buf) {
+            LeBitReader r = new LeBitReader(buf);
+
+            int modeNibble = (int) r.readBits(4);
+            if (modeNibble != 0) {
+                throw new VortexException(EncodingId.VORTEX_PCO,
+                        "pco mode " + modeNibble + " not yet implemented (only Classic=0)");
+            }
+            int deltaNibble = (int) r.readBits(4);
+            if (deltaNibble != 0) {
+                throw new VortexException(EncodingId.VORTEX_PCO,
+                        "pco delta " + deltaNibble + " not yet implemented (only None=0)");
+            }
+
+            // One primary latent variable for Classic + None delta.
+            int ansSizeLog = (int) r.readBits(4);
+            int nBins = (int) r.readBits(15);
+
+            PcoBin[] bins = new PcoBin[nBins];
+            for (int b = 0; b < nBins; b++) {
+                int weight = (int) r.readBits(ansSizeLog) + 1;
+                long lower = r.readBits(64);  // dtype_size = 64 for I64/U64
+                int offsetBits = (int) r.readBits(BITS_TO_ENCODE_OFFSET_BITS_64);
+                bins[b] = new PcoBin(weight, lower, offsetBits);
+            }
+            r.alignToByte(); // drain padding at end of chunk meta
+
+            return new PcoChunkMeta(ansSizeLog, bins);
         }
 
         private static EncodingProtos.PcoMetadata parseMeta(DecodeContext ctx) {
@@ -74,7 +166,8 @@ private static EncodingProtos.PcoMetadata parseMeta(DecodeContext ctx) {
             try {
                 return EncodingProtos.PcoMetadata.parseFrom(raw.duplicate());
             } catch (InvalidProtocolBufferException e) {
-                throw new VortexException(EncodingId.VORTEX_PCO, "invalid PcoMetadata: " + e.getMessage());
+                throw new VortexException(EncodingId.VORTEX_PCO,
+                        "invalid PcoMetadata: " + e.getMessage());
             }
         }
 
@@ -92,4 +185,7 @@ private static void validateHeader(EncodingProtos.PcoMetadata meta) {
             }
         }
     }
+
+    private record PcoChunkMeta(int ansSizeLog, PcoBin[] bins) {
+    }
 }

core/src/main/java/io/github/dfa1/vortex/encoding/PcoTansDecoder.java

@@ -0,0 +1,125 @@
+package io.github.dfa1.vortex.encoding;
+
+import java.lang.foreign.MemorySegment;
+import java.lang.foreign.ValueLayout;
+import java.nio.ByteOrder;
+
+/// 4-way interleaved tANS decoder for one pco latent variable.
+///
+/// Build via {@link #build(int, PcoBin[])}; then call {@link #decodePage} once per page.
+/// Port of {@code pco/src/ans/spec.rs} (spread) and {@code pco/src/ans/decoding.rs} (nodes).
+final class PcoTansDecoder {
+
+    private static final ValueLayout.OfLong LE_LONG =
+            ValueLayout.JAVA_LONG_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
+
+    static final int BATCH_N = 256;
+    static final int ANS_INTERLEAVING = 4;
+
+    // All arrays indexed by state index in [0, tableSize).
+    private final int[] nextStateIdxBase; // = (symbolXs[sym] << bitsToRead) - tableSize
+    private final int[] bitsToRead;       // bits consumed from bit stream per ANS step
+    private final int[] nodeOffsetBits;   // offset bits for this bin (stored per-state for cache locality)
+    private final long[] stateLowers;     // bin.lower for each state
+
+    private PcoTansDecoder(int[] nextStateIdxBase, int[] bitsToRead,
+                           int[] nodeOffsetBits, long[] stateLowers) {
+        this.nextStateIdxBase = nextStateIdxBase;
+        this.bitsToRead = bitsToRead;
+        this.nodeOffsetBits = nodeOffsetBits;
+        this.stateLowers = stateLowers;
+    }
+
+    /// Build the decode table from chunk latent-var metadata.
+    ///
+    /// Port of {@code Spec::from_weights} + {@code Decoder::new} from pcodec.
+    static PcoTansDecoder build(int ansSizeLog, PcoBin[] bins) {
+        if (bins.length == 0) {
+            // Degenerate: no bins → 1-state table, all offsets zero.
+            return new PcoTansDecoder(new int[]{0}, new int[]{0}, new int[]{0}, new long[]{0L});
+        }
+
+        int tableSize = 1 << ansSizeLog;
+        int[] weights = new int[bins.length];
+        for (int i = 0; i < bins.length; i++) {
+            weights[i] = bins[i].weight();
+        }
+
+        int[] stateSymbols = spreadStateSymbols(ansSizeLog, weights, tableSize);
+
+        int[] symbolXs = weights.clone();
+        int[] nextStateIdxBase = new int[tableSize];
+        int[] bitsToRead = new int[tableSize];
+        int[] nodeOffsetBits = new int[tableSize];
+        long[] stateLowers = new long[tableSize];
+
+        for (int s = 0; s < tableSize; s++) {
+            int sym = stateSymbols[s];
+            int xs = symbolXs[sym];
+            int btr = Integer.numberOfLeadingZeros(xs) - Integer.numberOfLeadingZeros(tableSize);
+            int nextBase = xs << btr;
+            nextStateIdxBase[s] = nextBase - tableSize;
+            bitsToRead[s] = btr;
+            nodeOffsetBits[s] = sym < bins.length ? bins[sym].offsetBits() : 0;
+            stateLowers[s] = sym < bins.length ? bins[sym].lower() : 0L;
+            symbolXs[sym]++;
+        }
+
+        return new PcoTansDecoder(nextStateIdxBase, bitsToRead, nodeOffsetBits, stateLowers);
+    }
+
+    /// Port of {@code Spec::spread_state_symbols} from pcodec.
+    ///
+    /// Spreads symbols across the table with a stride of ~3/5 * tableSize (odd).
+    static int[] spreadStateSymbols(int ansSizeLog, int[] weights, int tableSize) {
+        int[] stateSymbols = new int[tableSize];
+        int stride = (3 * tableSize) / 5;
+        if (stride % 2 == 0) {
+            stride++;
+        }
+        int modMask = tableSize - 1;
+        int step = 0;
+        for (int sym = 0; sym < weights.length; sym++) {
+            for (int k = 0; k < weights[sym]; k++) {
+                stateSymbols[(stride * step) & modMask] = sym;
+                step++;
+            }
+        }
+        return stateSymbols;
+    }
+
+    /// Decode {@code n} raw latent values (U64) from {@code reader} into {@code out}.
+    ///
+    /// Caller must have already read 4 initial ANS state indices and called
+    /// {@link LeBitReader#alignToByte()} before this call.
+    /// {@code ansStateIdxs} is modified in place and not valid after return.
+    void decodePage(LeBitReader reader, int[] ansStateIdxs, int n,
+                    MemorySegment out, long outByteOffset) {
+        int remaining = n;
+        long pos = outByteOffset;
+        long[] batchLowers = new long[BATCH_N];
+        int[] batchOffsetBits = new int[BATCH_N];
+
+        while (remaining > 0) {
+            int batchN = Math.min(remaining, BATCH_N);
+
+            // Phase 1 — read all ANS bin indices for this batch (sequential bit stream).
+            for (int i = 0; i < batchN; i++) {
+                int si = ansStateIdxs[i % ANS_INTERLEAVING];
+                batchLowers[i] = stateLowers[si];
+                batchOffsetBits[i] = nodeOffsetBits[si];
+                long ansVal = reader.readBits(bitsToRead[si]);
+                ansStateIdxs[i % ANS_INTERLEAVING] = nextStateIdxBase[si] + (int) ansVal;
+            }
+
+            // Phase 2 — read all offsets and reconstruct latents.
+            for (int i = 0; i < batchN; i++) {
+                long offset = reader.readBits(batchOffsetBits[i]);
+                out.set(LE_LONG, pos, batchLowers[i] + offset);
+                pos += Long.BYTES;
+            }
+
+            remaining -= batchN;
+        }
+    }
+}