Add kernel interface, generic impl, and test suite (#848)

src/openzl/codecs/common/bitstream/ff_bitstream.h

@@ -31,6 +31,10 @@ ZL_INLINE ZL_Report ZS_BitCStreamFF_finish(ZS_BitCStreamFF* bits);
 ZL_INLINE void
 ZS_BitCStreamFF_write(ZS_BitCStreamFF* bits, size_t value, size_t nbBits);
 ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits);
+ZL_INLINE uint8_t* ZS_BitCStreamFF_reserveAlignedBits(
+        ZS_BitCStreamFF* bits,
+        size_t nbBits);
+ZL_INLINE void ZS_BitCStreamFF_commitReservedBits(ZS_BitCStreamFF* bits);
 
 ZL_INLINE void ZS_BitCStreamFF_writeExpGolomb(
         ZS_BitCStreamFF* bits,
@@ -52,6 +56,7 @@ typedef struct {
     uint8_t const* ptr;
     uint8_t const* limit;
     uint8_t const* end;
+    uint8_t const* begin;
 } ZS_BitDStreamFF;
 
 ZL_INLINE ZS_BitDStreamFF
@@ -62,6 +67,9 @@ ZL_INLINE size_t
 ZS_BitDStreamFF_peek(ZS_BitDStreamFF const* bits, size_t nbBits);
 ZL_INLINE void ZS_BitDStreamFF_skip(ZS_BitDStreamFF* bits, size_t nbBits);
 ZL_INLINE void ZS_BitDStreamFF_reload(ZS_BitDStreamFF* bits);
+ZL_INLINE uint8_t const* ZS_BitDStreamFF_popAlignedBits(
+        ZS_BitDStreamFF* bits,
+        size_t nbBits);
 
 ZL_INLINE uint32_t
 ZS_BitDStreamFF_readExpGolomb(ZS_BitDStreamFF* bits, size_t order)
@@ -80,11 +88,13 @@ ZS_BitDStreamFF_readExpGolomb(ZS_BitDStreamFF* bits, size_t order)
 
 ZL_INLINE ZS_BitCStreamFF ZS_BitCStreamFF_init(uint8_t* dst, size_t dstCapacity)
 {
+    const size_t limit =
+            dstCapacity < sizeof(size_t) ? 0 : dstCapacity - sizeof(size_t) + 1;
     return (ZS_BitCStreamFF){
         .container = 0,
         .nbBits    = 0,
         .ptr       = dst,
-        .limit     = dst + dstCapacity - sizeof(size_t),
+        .limit     = dst + limit,
         .end       = dst + dstCapacity,
         .begin     = dst,
     };
@@ -94,7 +104,7 @@ ZL_INLINE ZL_Report ZS_BitCStreamFF_finish(ZS_BitCStreamFF* bits)
 {
     ZL_RESULT_DECLARE_SCOPE_REPORT((ZL_OperationContext*)NULL);
     size_t bytesToWrite = (bits->nbBits + 7) / 8;
-    if (bits->end < bits->ptr + bytesToWrite)
+    if ((size_t)(bits->end - bits->ptr) < bytesToWrite)
         ZL_ERR(internalBuffer_tooSmall);
     if (bytesToWrite) {
         ZL_ASSERT_EQ(
@@ -116,7 +126,7 @@ ZS_BitCStreamFF_write(ZS_BitCStreamFF* bits, size_t value, size_t nbBits)
 
 ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits)
 {
-    if (bits->ptr > bits->limit) {
+    if (bits->ptr >= bits->limit) {
         return;
     }
     size_t const nbBytes = bits->nbBits >> 3;
@@ -126,6 +136,62 @@ ZL_INLINE void ZS_BitCStreamFF_flush(ZS_BitCStreamFF* bits)
     bits->container >>= (nbBytes << 3);
 }
 
+ZL_INLINE uint8_t* ZS_BitCStreamFF_reserveAlignedBits(
+        ZS_BitCStreamFF* bits,
+        size_t nbBits)
+{
+    if (nbBits > SIZE_MAX - 7) {
+        return NULL;
+    }
+
+    size_t const bufferedBytes = (bits->nbBits + 7) / 8;
+    if (bufferedBytes > (size_t)(bits->end - bits->ptr)) {
+        return NULL;
+    }
+
+    size_t const nbBytes   = (nbBits + 7) / 8;
+    size_t const available = (size_t)(bits->end - bits->ptr) - bufferedBytes;
+    if (nbBytes > available) {
+        return NULL;
+    }
+
+    ZL_ASSERT_LE(bufferedBytes, sizeof(size_t));
+    if (bufferedBytes != 0) {
+        ZL_writeLE64_N(bits->ptr, bits->container, bufferedBytes);
+        bits->ptr += bufferedBytes;
+        bits->container = 0;
+        bits->nbBits    = 0;
+    }
+    ZL_ASSERT_EQ(bits->nbBits, 0);
+
+    uint8_t* const out = bits->ptr;
+    bits->ptr          = out + (nbBits / 8);
+    bits->nbBits       = nbBits & 7;
+    bits->container    = 0;
+    return out;
+}
+
+ZL_INLINE void ZS_BitCStreamFF_commitReservedBits(ZS_BitCStreamFF* bits)
+{
+    if (bits->nbBits == 0) {
+        bits->container = 0;
+        return;
+    }
+    bits->container = bits->ptr[0] & (((size_t)1 << bits->nbBits) - 1);
+}
+
+// Little-endian load of the first @p nbBytes (< sizeof(size_t)) bytes at @p src
+// into a container word. Used for streams shorter than a full word, where every
+// byte lives in the container.
+ZL_INLINE size_t ZS_BitDStreamFF_loadPartial(uint8_t const* src, size_t nbBytes)
+{
+    size_t container = 0;
+    for (size_t i = 0; i < nbBytes; ++i) {
+        container |= (size_t)src[i] << (i << 3);
+    }
+    return container;
+}
+
 ZL_INLINE ZS_BitDStreamFF
 ZS_BitDStreamFF_init(uint8_t const* src, size_t srcSize)
 {
@@ -136,19 +202,19 @@ ZS_BitDStreamFF_init(uint8_t const* src, size_t srcSize)
             .ptr        = src,
             .limit      = src + srcSize - sizeof(size_t) + 1,
             .end        = src + srcSize,
+            .begin      = src,
         };
         return bits;
     } else {
-        ZS_BitDStreamFF bits = {
-            .container  = 0,
-            .nbBitsRead = (sizeof(size_t) - srcSize) * 8,
-            .ptr        = src + srcSize,
-            .limit      = src,
-            .end        = src + srcSize,
+        uint8_t const* const end = srcSize > 0 ? src + srcSize : src;
+        ZS_BitDStreamFF bits     = {
+                .container  = ZS_BitDStreamFF_loadPartial(src, srcSize),
+                .nbBitsRead = (sizeof(size_t) - srcSize) * 8,
+                .ptr        = end,
+                .limit      = src,
+                .end        = end,
+                .begin      = src,
         };
-        for (size_t i = 0; i < srcSize; ++i) {
-            bits.container |= (size_t)src[i] << (i << 3);
-        }
         return bits;
     }
 }
@@ -159,7 +225,13 @@ ZL_INLINE ZL_Report ZS_BitDStreamFF_finish(ZS_BitDStreamFF const* bits)
     if (bits->nbBitsRead > ZS_BITSTREAM_READ_MAX_BITS) {
         ZL_ERR(GENERIC);
     }
-    return ZL_returnSuccess();
+    size_t bytesRead = (size_t)(bits->ptr - bits->begin);
+    bytesRead += bits->nbBitsRead >> 3;
+    bytesRead += (bits->nbBitsRead & 7) != 0;
+    if ((size_t)(bits->end - bits->begin) < sizeof(size_t)) {
+        bytesRead -= sizeof(size_t);
+    }
+    return ZL_returnValue(bytesRead);
 }
 
 ZL_INLINE size_t ZS_BitDStreamFF_read(ZS_BitDStreamFF* bits, size_t nbBits)
@@ -189,22 +261,93 @@ ZL_INLINE void ZS_BitDStreamFF_skip(ZS_BitDStreamFF* bits, size_t nbBits)
 
 ZL_INLINE void ZS_BitDStreamFF_reload(ZS_BitDStreamFF* bits)
 {
-    bits->ptr += bits->nbBitsRead >> 3;
-    if (ZL_LIKELY(bits->ptr < bits->limit)) {
-        size_t const next = ZL_readLEST(bits->ptr);
+    uint8_t const* const ptr = bits->ptr + (bits->nbBitsRead >> 3);
+    if (ZL_LIKELY(ptr < bits->limit)) {
+        bits->ptr = ptr;
         bits->nbBitsRead &= 7;
-        bits->container = next >> bits->nbBitsRead;
+        bits->container = ZL_readLEST(ptr) >> bits->nbBitsRead;
         return;
     }
 
-    if (bits->ptr >= bits->end)
+    // Past the end: leave ptr and nbBitsRead untouched. This keeps the consumed
+    // bit count -- (ptr - begin) * 8 + nbBitsRead -- exact for finish() and
+    // popAlignedBits(), while an over-read still leaves nbBitsRead too large
+    // for finish() to accept.
+    if (ptr >= bits->end)
         return;
 
-    uint8_t const* const limit = bits->limit - 1;
-    size_t const skippedBits   = (size_t)((bits->ptr - limit) << 3);
-    size_t const next          = ZL_readLEST(limit);
+    uint8_t const* const tail = bits->limit - 1;
+    size_t const skippedBits  = (size_t)((ptr - tail) << 3);
+    bits->ptr                 = ptr;
     bits->nbBitsRead &= 7;
-    bits->container = next >> (bits->nbBitsRead + skippedBits);
+    bits->container = ZL_readLEST(tail) >> (bits->nbBitsRead + skippedBits);
+}
+
+ZL_INLINE uint8_t const* ZS_BitDStreamFF_popAlignedBits(
+        ZS_BitDStreamFF* bits,
+        size_t nbBits)
+{
+    if (nbBits > SIZE_MAX - 7) {
+        return NULL;
+    }
+
+    // check if the stream is already in an invalid state
+    if (bits->nbBitsRead > ZS_BITSTREAM_READ_MAX_BITS) {
+        return NULL;
+    }
+
+    size_t const streamSize = (size_t)(bits->end - bits->begin);
+    if (streamSize > SIZE_MAX / 8) {
+        return NULL;
+    }
+
+    size_t consumedBits =
+            (size_t)(bits->ptr - bits->begin) * 8 + bits->nbBitsRead;
+    if (streamSize < sizeof(size_t)) {
+        consumedBits -= sizeof(size_t) * 8;
+    }
+
+    size_t const bitSize = streamSize * 8;
+    if (consumedBits > bitSize || consumedBits > SIZE_MAX - 7) {
+        return NULL;
+    }
+
+    size_t const alignedBits = ((consumedBits + 7) / 8) * 8;
+    if (nbBits > bitSize - alignedBits) {
+        return NULL;
+    }
+
+    uint8_t const* const out = bits->begin + alignedBits / 8;
+
+    // Reposition the stream just past the aligned region we handed out.
+    size_t const bitPos      = alignedBits + nbBits;
+    uint8_t const* const pos = bits->begin + bitPos / 8;
+    size_t const bitShift    = bitPos & 7;
+
+    if (streamSize >= sizeof(size_t)) {
+        // Refill the 64-bit window at `pos`, clamping the load back from the
+        // end when fewer than 8 bytes remain; the matching shift exposes the
+        // same bit either way. A shift of the full word width (the whole
+        // stream consumed) leaves an empty window.
+        uint8_t const* const tail = bits->end - sizeof(size_t);
+        uint8_t const* const load = pos <= tail ? pos : tail;
+        size_t const shift        = bitPos - (size_t)(load - bits->begin) * 8;
+        bits->container =
+                shift < sizeof(size_t) * 8 ? ZL_readLEST(load) >> shift : 0;
+        bits->nbBitsRead = bitShift;
+        bits->ptr        = pos;
+    } else {
+        // Short stream: every byte already lives in the container.
+        size_t const remaining = (size_t)(bits->end - pos);
+        bits->container        = 0;
+        for (size_t i = 0; i < remaining; ++i) {
+            bits->container |= (size_t)pos[i] << (i << 3);
+        }
+        bits->container >>= bitShift;
+        bits->nbBitsRead = (sizeof(size_t) - remaining) * 8 + bitShift;
+        bits->ptr        = bits->end;
+    }
+    return out;
 }
 
 ZL_END_C_DECLS

src/openzl/codecs/pivco_huffman/README.md

@@ -0,0 +1,22 @@
+# PivCo-Huffman
+
+This is an implementation of PivCo-Huffman coding from Marcin Zukowski.
+The code is based on ideas from:
+
+- [The paper](https://arxiv.org/abs/2606.05765)
+- [The repo](https://github.com/MarcinZukowski/pivco-huffman)
+- [Ryg's blog](https://fgiesen.wordpress.com/2026/06/21/pivco-huffman-merge-operations/)
+
+OpenZL re-implements PivCo-Huffman instead of using the upstream repo for several reasons:
+
+1. OpenZL needs full control over the wire-format.
+2. OpenZL needs to be hardened against decoding corrupted data.
+3. OpenZL needs to be able to dynamically dispatch to kernels based on the CPUID,
+e.g. to detect AVX512 support at runtime.
+
+The goal is to work with upstream to factor out the primitives so that OpenZL can
+import and use upstream's primitives, as those are the most complex part of PivCo,
+and are independent of the wire format.
+
+In the meantime, any meaningful improvements will be upstreamed, so there is a
+single reference implementation for PivCo Huffman.

src/openzl/codecs/pivco_huffman/arch/common_pivco_arch.h

@@ -0,0 +1,20 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+#ifndef OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_ARCH_H
+#define OPENZL_CODECS_PIVCO_HUFFMAN_ARCH_COMMON_PIVCO_ARCH_H
+
+#include "openzl/shared/portability.h"
+
+ZL_BEGIN_C_DECLS
+
+/**
+ * Slop bytes that every bitmap and rank buffer must reserve past its logical
+ * length. The encode and decode kernels may over-read and over-write up to this
+ * many trailing bytes where noted in the docs as `SLOP`, so callers pad their
+ * buffers by this amount. It is sized to the widest fixed-width group the SIMD
+ * kernels process.
+ */
+#define ZL_PIVCO_HUFFMAN_SLOP 64
+
+ZL_END_C_DECLS
+
+#endif