diff --git a/CMakeLists.txt b/CMakeLists.txt index 6101a98..72cf718 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -256,6 +256,15 @@ if (PIXIE_TESTS) PRIVATE ${sdsl_lite_SOURCE_DIR}/include) endif () + add_executable(pivco_huffman_tests + src/tests/pivco_huffman_tests.cpp) + target_include_directories(pivco_huffman_tests + PUBLIC include) + target_link_libraries(pivco_huffman_tests + gtest + gtest_main + ${PIXIE_DIAGNOSTICS_LIBS}) + set(PIXIE_TEST_TARGETS bit_algorithms_unittests rank_select_unittests @@ -267,7 +276,8 @@ if (PIXIE_TESTS) storage_tests excess_positions_tests excess_record_lows_tests - rmq_tests) + rmq_tests + pivco_huffman_tests) foreach (test_target IN LISTS PIXIE_TEST_TARGETS) gtest_discover_tests(${test_target} DISCOVERY_MODE PRE_TEST @@ -399,6 +409,24 @@ if (PIXIE_BENCHMARKS) benchmark benchmark_main ${PIXIE_DIAGNOSTICS_LIBS}) + + add_executable(pivco_huffman_benchmarks + src/benchmarks/pivco_huffman_benchmarks.cpp) + target_include_directories(pivco_huffman_benchmarks + PUBLIC include) + target_link_libraries(pivco_huffman_benchmarks + benchmark + benchmark_main + ${PIXIE_DIAGNOSTICS_LIBS}) + + add_executable(pivco_huffman_file_benchmarks + src/benchmarks/pivco_huffman_file_benchmarks.cpp) + target_include_directories(pivco_huffman_file_benchmarks + PUBLIC include) + target_link_libraries(pivco_huffman_file_benchmarks + benchmark + benchmark_main + ${PIXIE_DIAGNOSTICS_LIBS}) endif () # --------------------------------------------------------------------------- diff --git a/include/pixie/huffman.h b/include/pixie/huffman.h new file mode 100644 index 0000000..68acab5 --- /dev/null +++ b/include/pixie/huffman.h @@ -0,0 +1,89 @@ +#pragma once + +/** + * @file huffman.h + * @brief Common CRTP contract for Huffman entropy codecs. + * + * PivCo-Huffman ("Pivot-Coded Huffman") reuses the wavelet-tree "tree of + * bitmaps" layout to turn sequential, bit-by-bit Huffman-tree traversals into + * vectorizable operations. A codec encodes a byte sequence into a compressed + * stream by building a Huffman-shaped tree of per-node bitmaps, and decodes it + * back by traversing that tree. Concrete implementations live under + * ``. + * + * @see Marcin Zukowski, "PivCo-Huffman", v1.0 (2026). + */ + +#include +#include +#include +#include + +namespace pixie { + +/** + * @brief CRTP facade for Huffman entropy codecs. + * + * The contract is the source of truth for observable codec semantics. Each + * concrete implementation inherits `HuffmanBase` and supplies the + * required `*_impl()` extension points. There is no virtual dispatch: the + * facade delegates statically through CRTP, mirroring the other Pixie families. + * + * Range and ownership conventions: + * - Symbol sequences are zero-based byte streams (`symbol_type`). + * - Compressed streams are byte-oriented views (`std::byte`). + * - The caller keeps any non-owning view alive for the codec lifetime. + * + * @tparam Impl Concrete codec type implementing the `*_impl()` contract. + * + * @see `` for the available concrete + * implementations. + */ +template +class HuffmanBase { + public: + /** @brief Symbol type handled by the codec: one byte per symbol. */ + using symbol_type = std::uint8_t; + + /** + * @brief Number of symbols in the original (uncompressed) stream. + * @return Logical uncompressed symbol count. + */ + std::size_t uncompressed_size() const { + return impl().uncompressed_size_impl(); + } + + /** + * @brief Number of bytes in the compressed representation. + * @return Compressed stream size in bytes. + */ + std::size_t compressed_size() const { return impl().compressed_size_impl(); } + + /** + * @brief Check whether the codec holds no data. + * @return `true` when `uncompressed_size() == 0`. + */ + bool empty() const { return uncompressed_size() == 0; } + + /** + * @brief Read-only view of the compressed byte stream. + * @return Span over the serialized representation. + * + * @note The returned view is invalidated by codec destruction. + */ + std::span compressed_data() const { + return impl().compressed_data_impl(); + } + + /** + * @brief Reconstruct the original symbol sequence. + * @return Decoded symbols of length `uncompressed_size()`. + */ + std::vector decode() const { return impl().decode_impl(); } + + private: + /** @brief Return this facade as its concrete CRTP implementation. */ + const Impl& impl() const { return static_cast(*this); } +}; + +} // namespace pixie diff --git a/include/pixie/huffman/implementations.h b/include/pixie/huffman/implementations.h new file mode 100644 index 0000000..e5455cf --- /dev/null +++ b/include/pixie/huffman/implementations.h @@ -0,0 +1,12 @@ +#pragma once + +/** + * @file implementations.h + * @brief All Huffman codec implementations provided by Pixie. + * + * - `PivCoHuffman`: simple scalar reference codec storing a Huffman-shaped + * tree of per-node routing bitmaps as packed 64-bit words. + */ + +#include +#include diff --git a/include/pixie/huffman/pivco_huffman.h b/include/pixie/huffman/pivco_huffman.h new file mode 100644 index 0000000..27bfb39 --- /dev/null +++ b/include/pixie/huffman/pivco_huffman.h @@ -0,0 +1,1333 @@ +#pragma once + +/** + * @file pivco_huffman.h + * @brief Simple scalar PivCo-Huffman codec with canonical Huffman codes. + * + * Stores a Huffman-shaped tree of per-node routing bitmaps (the "tree of + * bitmaps" layout shared with wavelet trees). Encoding walks each input symbol + * from the root, appending one direction bit to every internal node on its + * path. Decoding merges child symbol streams bottom-up from leaves to root, + * using the node bitmap as a selector. + * + * The tree uses canonical Huffman code lengths (<= 15 bits, enforced with a + * zlib-style overflow correction): the wire format stores one file-level set + * of 256 per-symbol lengths (128 bytes as 4-bit nibbles), and the tree + * structure is reconstructed once at load time and reused by every block. + * Same-length symbols are reshaped into flat subtrees as described by the + * PivCo-Huffman paper. + * + * This implementation uses packed `std::uint64_t` bitmaps, flat subtrees, + * contiguous arenas, and bottom-up decoding. The merge/partition primitives are + * vectorized in `pivco_simd.h` (AVX-512 VBMI2, AVX2, and scalar fallbacks); + * selective ANS remains future work. + * + * Range and ownership conventions follow `HuffmanBase`: symbols are bytes, + * the compressed stream is a byte view, and the codec owns its serialized form. + */ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace pixie { + +/** + * @brief Simple scalar PivCo-Huffman codec with canonical Huffman codes. + * @details Implements `HuffmanBase` by building a Huffman tree over the + * byte alphabet, canonicalizing it (grouping same-length codes), + * storing one routing bitmap per internal node as packed 64-bit + * words, and decoding by bottom-up merging of child symbol streams. + * The serialized form is one file-level set of 256 code lengths + * followed by block-local packed node bitmaps; the shared tree is + * reconstructed from the lengths at load time. + */ +class PivCoHuffman : public HuffmanBase { + public: + /** @brief Symbol type handled by the codec: one byte per symbol. */ + using symbol_type = std::uint8_t; + + /** + * @brief Build a codec by encoding @p input. + * @param input Symbol sequence to compress. + */ + explicit PivCoHuffman(std::span input) { build(input); } + + /** + * @brief Load a codec from a previously serialized compressed stream. + * @param compressed Compressed byte view produced by another instance. + */ + explicit PivCoHuffman(std::span compressed) { + deserialize(compressed); + } + + // --- CRTP extension points ---------------------------------------------- + + /** @brief Size of the original symbol stream. */ + std::size_t uncompressed_size_impl() const { return uncompressed_size_; } + + /** @brief Size of the compressed byte stream. */ + std::size_t compressed_size_impl() const { return compressed_.size(); } + + /** @brief Read-only view of the compressed byte stream. */ + std::span compressed_data_impl() const { + return compressed_; + } + + /** @brief Reconstruct the original symbol sequence by decoding each block. */ + std::vector decode_impl() const { + if (uncompressed_size_ == 0) { + return std::vector(); + } + // SIMD merge kernels may touch one vector past their logical output. Keep + // the same slack as the internal workspace; resize it away without copying + // after the last block has decoded. + std::vector result(uncompressed_size_ + kWorkspaceSlack); + std::size_t pos = blocks_offset_; + std::size_t output_offset = 0; + while (output_offset < uncompressed_size_) { + deserialize_block_payload(pos); + decode_block_into(std::span(result).subspan( + output_offset, block_uncompressed_size_)); + output_offset += block_uncompressed_size_; + } + result.resize(uncompressed_size_); + return result; + } + + private: + /// @brief Sentinel node index meaning "no node". + static constexpr std::size_t kNpos = static_cast(-1); + + /// @brief Byte-alphabet size (256 symbols). + static constexpr std::size_t kAlphabet = 256; + + /// @brief Bits per word used by the packed bitmap. + static constexpr std::size_t kWordBits = 64; + + /// @brief Maximum supported Huffman code length (matches common decoders). + static constexpr std::size_t kMaxCodeLength = 15; + + /// @brief Maximum depth of a flat subtree (table size = 2^depth ≤ 256). + /// @details A flat subtree replaces `2^d - 1` internal nodes (each with a + /// 1-bit-per-symbol bitmap and a merge pass) with a single node + /// storing `d` bits per symbol and a lookup table. This turns `d` + /// decode merge passes into a single table-lookup pass. + static constexpr std::uint8_t kMaxFlatDepth = 8; + + /// @brief Maximum nodes in a full binary tree over the byte alphabet. + static constexpr std::size_t kMaxTreeNodes = 2 * kAlphabet - 1; + + /// @brief Default block size for block-based processing (64 KiB). + /// @details All blocks share one file-level Huffman tree but carry their own + /// routing bitmaps. This keeps the decode workspace (2 × block_size + /// = 128 KiB) within the per-core L2 while avoiding per-block tree + /// construction and serialization overhead. + static constexpr std::size_t kBlockSize = 64 * 1024; + + /// @brief Extra bytes appended to the ping-pong workspace beyond `2 * block`. + /// @details Vectorized merge/partition kernels read and write a full 8- or + /// 64-byte vector on each iteration even when the final group is + /// shorter; this slack makes those over-reads/over-writes stay + /// within the allocated buffer rather than spilling past the end of + /// a child region that abuts the workspace boundary. + static constexpr std::size_t kWorkspaceSlack = 64; + + /** + * @brief Lightweight descriptor for a node's routing bitmap. + * @details The packed 64-bit words live in a single shared arena + * (`arena_`); this struct only records where the node's words begin + * (`offset`, in words) and how many symbols they describe + * (`count`). Splitting descriptor from storage lets one contiguous + * allocation back every internal node, eliminating per-node heap + * traffic during decode and improving cache locality. + */ + struct Bitmap { + std::size_t offset = 0; + std::size_t count = 0; ///< Number of symbols (not bits). + + /** @brief Number of 64-bit words backing the bitmap. + * @param flat_depth 0 for normal nodes (1 bit/symbol); >0 for flat + * nodes (`flat_depth` bits/symbol). */ + std::size_t word_count(std::uint8_t flat_depth = 0) const { + std::size_t bits_per = flat_depth > 0 ? flat_depth : 1; + return (count * bits_per + kWordBits - 1) / kWordBits; + } + }; + + /** + * @brief A node in the Huffman tree of bitmaps. + * @details Internal nodes carry a routing bitmap with one bit per symbol + * that passes through the node (0 = left, 1 = right), in input + * order. Leaves carry only their symbol. + */ + struct Node { + std::size_t left = kNpos; + std::size_t right = kNpos; + std::uint8_t symbol = 0; + bool is_leaf = false; + Bitmap bits; + /// @brief 0 = normal node (1 bit/symbol); >0 = flat subtree of this depth + /// (stores `flat_depth` bits/symbol + a lookup table). + std::uint8_t flat_depth = 0; + /// @brief Whether this flat table is the balanced bit-reversal mapping. + bool flat_bitreverse = false; + /// @brief Offset of this flat node's decode table in `flat_tables_`. + std::size_t flat_table_offset = 0; + }; + + /// @brief A symbol with its frequency weight, used for Huffman construction. + struct Leaf { + std::size_t weight; + std::uint8_t symbol; + }; + + /** @brief Compact root-to-leaf code used by the scalar encoder. + * @details Codes are stored MSB-first in the low `length` bits. The + * 15-bit length limit keeps every code in one 16-bit value and + * avoids the per-symbol dynamic path vectors used previously. */ + struct SymbolCode { + std::uint16_t bits = 0; + std::uint8_t length = 0; + }; + + using DirectionTables = + std::array, kMaxCodeLength>; + using FrequencyTable = std::array; + + /// @brief Total uncompressed size across all blocks (set during + /// build/deserialize, read by `uncompressed_size_impl()`). + std::size_t uncompressed_size_ = 0; + /// @brief Per-block serialized size from the wire (fixed; last block may + /// be smaller). Used by decode to iterate blocks. + std::size_t block_size_ = kBlockSize; + /// @brief Serialized compressed stream (header + per-block data). + std::vector compressed_; + /// @brief Byte offset of the first block payload in `compressed_`. + std::size_t blocks_offset_ = 0; + + // --- shared tree and current-block state -------------------------------- + // The topology and flat tables are file-level. Bitmap counts, offsets, + // and words are replaced for each block during const decoding. + mutable std::size_t root_ = kNpos; + mutable std::vector nodes_; + /// @brief One contiguous allocation backing every internal node's bitmap + /// words. Node `i` owns `nodes_[i].bits.word_count()` words starting + /// at `arena_[nodes_[i].bits.offset]`. + mutable std::vector arena_; + /// @brief Reusable decode/encode scratch: two ping-pong halves of size + /// `block_uncompressed_size_`. Reused across blocks; `mutable` + /// because `decode_impl()` is const. + mutable std::vector workspace_; + /// @brief Uncompressed size of the current block being processed. + mutable std::size_t block_uncompressed_size_ = 0; + mutable std::array code_lengths_{}; + /// @brief Contiguous storage for all flat-subtree decode tables. + mutable std::vector flat_tables_; + + // --- construction -------------------------------------------------------- + + /** @brief Build and serialize a shared tree and all block payloads. + * @details Counts each block once, combines those counts to build one + * file-level Huffman tree, then reuses its topology and direction + * tables while encoding every block. The wire format is: + * - total_uncompressed_size (8 bytes) + * - block_size (8 bytes) + * - 256 file-level code lengths as 128 bytes of nibbles + * - Per block: block_uncompressed_size (8 bytes) followed by + * per-internal-node bitmaps. */ + void build(std::span input) { + uncompressed_size_ = input.size(); + compressed_.clear(); + write(uncompressed_size_); + write(block_size_); + blocks_offset_ = compressed_.size(); + if (input.empty()) { + return; + } + + const std::size_t block_count = + (input.size() + block_size_ - 1) / block_size_; + std::vector block_frequencies; + block_frequencies.reserve(block_count); + FrequencyTable file_frequencies{}; + + std::size_t offset = 0; + while (offset < input.size()) { + const std::size_t this_block = + std::min(input.size() - offset, block_size_); + block_frequencies.push_back( + count_frequencies(input.subspan(offset, this_block))); + const FrequencyTable& frequencies = block_frequencies.back(); + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + file_frequencies[symbol] += frequencies[symbol]; + } + offset += this_block; + } + + compute_code_lengths(file_frequencies, code_lengths_); + build_noncanonical_tree(code_lengths_); + detect_flat_subtrees(); + serialize_code_lengths(); + blocks_offset_ = compressed_.size(); + + std::array codes{}; + DirectionTables directions{}; + if (root_ != kNpos && !nodes_[root_].is_leaf) { + assign_codes(root_, 0, 0, codes); + directions = build_direction_tables(codes); + } + + offset = 0; + std::size_t block_index = 0; + while (offset < input.size()) { + const std::size_t this_block = + std::min(input.size() - offset, block_size_); + build_block_payload(input.subspan(offset, this_block), + block_frequencies[block_index], directions); + serialize_block_payload(); + offset += this_block; + ++block_index; + } + } + + /** @brief Build one block's bitmaps using the file-level tree. */ + void build_block_payload(std::span input, + const FrequencyTable& freq, + const DirectionTables& directions) { + block_uncompressed_size_ = input.size(); + if (block_uncompressed_size_ == 0) { + return; + } + + // Single-symbol input: root is a leaf, no internal nodes, no arena. + if (root_ == kNpos || nodes_[root_].is_leaf) { + return; + } + + // Assign exact per-node weights (subtree frequency sums) and arena + // offsets in a single post-order walk, then allocate the arena once. + std::size_t arena_words = 0; + assign_weights_and_offsets(root_, freq, arena_words); + arena_.assign(arena_words, 0); + + // Fill node bitmaps via top-down in-place partition into the reusable + // workspace. This mirrors decode: + // the root reads its input from workspace half 0, writes its bitmap, + // and partitions symbols into half 1 for its children. Each level + // ping-pongs between the two halves — no per-node allocations. + workspace_.resize(block_uncompressed_size_ * 2 + kWorkspaceSlack); + std::copy(input.begin(), input.end(), workspace_.data()); + encode_partition(root_, 0, 0, directions, freq); + } + + /** @brief Count a block's byte frequencies with independent accumulators. + * @details Eight partial histograms break the dependency chain caused by + * repeatedly incrementing a hot symbol. The block is at most + * 64 KiB, so 32-bit partial counters cannot overflow. */ + static FrequencyTable count_frequencies(std::span input) { + constexpr std::size_t kLanes = 8; + std::array, kLanes> partial{}; + std::size_t i = 0; + for (; i + kLanes <= input.size(); i += kLanes) { + partial[0][input[i]]++; + partial[1][input[i + 1]]++; + partial[2][input[i + 2]]++; + partial[3][input[i + 3]]++; + partial[4][input[i + 4]]++; + partial[5][input[i + 5]]++; + partial[6][input[i + 6]]++; + partial[7][input[i + 7]]++; + } + for (; i < input.size(); ++i) { + partial[0][input[i]]++; + } + + FrequencyTable frequencies{}; + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + for (std::size_t lane = 0; lane < kLanes; ++lane) { + frequencies[symbol] += partial[lane][symbol]; + } + } + return frequencies; + } + + /** @brief Compute Huffman code lengths, limited to ≤ kMaxCodeLength. + * @details Two-phase approach (same as zlib/zstd): + * 1. Sort present symbols by frequency and build a standard + * Huffman tree with the linear two-queue algorithm, then + * extract per-leaf depths by walking parent pointers. + * 2. If any depth exceeds kMaxCodeLength, apply the zlib-style + * "overflow fixup": collapse all lengths > L into L, then + * rebalance the Kraft sum by borrowing from shorter codes. + * + * This is O(n log n) for sorting and O(n + L) after that, with fixed-size + * storage for the byte alphabet. */ + void compute_code_lengths(const std::array& freq, + std::array& lengths) { + lengths.fill(0); + + std::array leaves{}; + std::size_t leaf_count = 0; + for (std::size_t s = 0; s < kAlphabet; s++) { + if (freq[s] > 0) { + leaves[leaf_count++] = {freq[s], static_cast(s)}; + } + } + if (leaf_count == 0) { + return; + } + if (leaf_count == 1) { + lengths[leaves[0].symbol] = 1; + return; + } + + std::sort(leaves.begin(), leaves.begin() + leaf_count, + [](const Leaf& a, const Leaf& b) { + if (a.weight != b.weight) { + return a.weight < b.weight; + } + return a.symbol < b.symbol; + }); + + // Build the Huffman tree with the linear two-queue algorithm. Sorted + // leaves form one queue; newly-created internal nodes form the other. + // Internal weights are produced in nondecreasing order, so selecting the + // smaller queue front replaces all binary-heap maintenance. + struct TmpNode { + std::size_t weight = 0; + std::size_t parent = kNpos; + }; + std::array tree{}; + for (std::size_t i = 0; i < leaf_count; ++i) { + tree[i].weight = leaves[i].weight; + } + + std::size_t leaf_head = 0; + std::size_t internal_head = leaf_count; + std::size_t next_internal = leaf_count; + auto take_smallest = [&]() { + const bool take_leaf = + leaf_head < leaf_count && + (internal_head == next_internal || + tree[leaf_head].weight <= tree[internal_head].weight); + return take_leaf ? leaf_head++ : internal_head++; + }; + + const std::size_t node_count = 2 * leaf_count - 1; + while (next_internal < node_count) { + const std::size_t left = take_smallest(); + const std::size_t right = take_smallest(); + tree[next_internal].weight = tree[left].weight + tree[right].weight; + tree[left].parent = next_internal; + tree[right].parent = next_internal; + ++next_internal; + } + + // Extract code lengths by walking from each leaf to root. + std::uint8_t max_len = 0; + for (std::size_t i = 0; i < leaf_count; ++i) { + std::uint8_t depth = 0; + std::size_t cur = i; + while (tree[cur].parent != kNpos) { + cur = tree[cur].parent; + depth++; + } + lengths[leaves[i].symbol] = depth; + if (depth > max_len) { + max_len = depth; + } + } + + if (max_len <= kMaxCodeLength) { + return; + } + + // Length limiting via zlib-style overflow fixup. + limit_code_lengths(lengths, + std::span(leaves.data(), leaf_count)); + } + + /** @brief Limit Huffman code lengths to ≤ kMaxCodeLength. + * @details Collapses all lengths > L into L, rebalances the length + * histogram, then assigns its longest codes to the least frequent + * symbols. */ + void limit_code_lengths(std::array& lengths, + std::span leaves) { + // Count codes at each length and track how many exceeded the limit. This + // is the same overflow definition used by zlib's gen_bitlen(). + std::array bl_count{}; + int overflow = 0; + for (const Leaf& leaf : leaves) { + std::uint8_t len = lengths[leaf.symbol]; + if (len > kMaxCodeLength) { + len = kMaxCodeLength; + ++overflow; + } + ++bl_count[len]; + } + + // Move one leaf down from the deepest available shorter level, creating + // two children there, and remove one overlong code from level L. + while (overflow > 0) { + int bits = static_cast(kMaxCodeLength) - 1; + while (bits > 0 && bl_count[bits] == 0) { + --bits; + } + --bl_count[bits]; + bl_count[bits + 1] += 2; + --bl_count[kMaxCodeLength]; + overflow -= 2; + } + + // Leaves arrive sorted by ascending frequency. Assign the longest codes + // first so the fixup preserves the minimum weighted path length for its + // corrected length histogram. + lengths.fill(0); + std::size_t leaf_index = 0; + for (int len = static_cast(kMaxCodeLength); len >= 1; --len) { + for (int count = 0; count < bl_count[len]; ++count) { + lengths[leaves[leaf_index++].symbol] = static_cast(len); + } + } + } + + // --- bit helpers --------------------------------------------------------- + // Small functions for reading/writing individual bits and multi-bit + // values from/to packed `std::uint64_t` word arrays. Used by both the + // normal (1-bit-per-symbol) and flat (d-bits-per-symbol) code paths. + + /** @brief Reverse the low @p width bits of an at-most-eight-bit value. */ + static std::uint8_t reverse_low_bits(std::uint8_t value, std::uint8_t width) { + value = static_cast(((value & 0x55u) << 1) | + ((value & 0xAAu) >> 1)); + value = static_cast(((value & 0x33u) << 2) | + ((value & 0xCCu) >> 2)); + value = static_cast((value << 4) | (value >> 4)); + return value >> (8 - width); + } + + // --- flat subtree detection ---------------------------------------------- + + /** @brief Sentinel value meaning "subtree has non-uniform leaf depths". */ + static constexpr std::uint8_t kNotUniform = 0xFF; + + /** @brief Detect maximal flat subtrees in two linear tree passes. + * @details The post-order pass computes uniform leaf depths without + * allocating per node. A top-down pass then selects only maximal + * eligible subtrees and appends their tables to one contiguous + * arena. This follows the paper's flat-subtree optimization while + * avoiding temporary child tables that a bottom-up marker would + * immediately discard when their parent is also flat. */ + void detect_flat_subtrees() const { + flat_tables_.clear(); + if (root_ != kNpos) { + std::array uniform_depths{}; + compute_uniform_depths(root_, uniform_depths); + mark_flat_topdown(root_, uniform_depths); + } + } + + /** @brief Record each subtree's uniform leaf depth in post-order. */ + std::uint8_t compute_uniform_depths( + std::size_t idx, + std::array& depths) const { + const Node& n = nodes_[idx]; + if (n.is_leaf) { + return depths[idx] = 0; + } + + const std::uint8_t left_depth = compute_uniform_depths(n.left, depths); + const std::uint8_t right_depth = compute_uniform_depths(n.right, depths); + + if (left_depth == kNotUniform || right_depth == kNotUniform || + left_depth != right_depth) { + return depths[idx] = kNotUniform; + } + return depths[idx] = static_cast(left_depth + 1); + } + + /** @brief Mark maximal flat subtrees and append their lookup tables. */ + void mark_flat_topdown( + std::size_t idx, + const std::array& depths) const { + Node& n = nodes_[idx]; + if (n.is_leaf) { + return; + } + + const std::uint8_t depth = depths[idx]; + if (depth >= 2 && depth <= kMaxFlatDepth) { + n.flat_depth = depth; + n.flat_table_offset = flat_tables_.size(); + flat_tables_.resize(flat_tables_.size() + (std::size_t{1} << depth)); + n.flat_bitreverse = build_flat_table( + idx, depth, 0, flat_tables_.data() + n.flat_table_offset); + return; + } + + mark_flat_topdown(n.left, depths); + mark_flat_topdown(n.right, depths); + } + + /** @brief Populate the flat lookup table by walking the subtree. + * @param idx Current node. + * @param flat_depth Total depth of the flat subtree. + * @param code Suffix accumulated so far (MSB-first: left=0, + * right=1). + * @param table Output table of size `2^total_depth`. + * @return True when the completed table is the balanced bit-reversal + * permutation used by the specialized SIMD kernels. */ + bool build_flat_table(std::size_t idx, + std::uint8_t flat_depth, + std::uint32_t code, + symbol_type* table) const { + const Node& n = nodes_[idx]; + if (n.is_leaf) { + table[reverse_low_bits(static_cast(code), flat_depth)] = + n.symbol; + return n.symbol == code; + } + const bool left_bitreverse = + build_flat_table(n.left, flat_depth, code << 1, table); + const bool right_bitreverse = + build_flat_table(n.right, flat_depth, (code << 1) | 1, table); + return left_bitreverse && right_bitreverse; + } + + /** @brief Total frequency weight of a subtree (sum of leaf frequencies). */ + std::size_t subtree_weight( + std::size_t idx, + const std::array& freq) const { + if (idx == kNpos) { + return 0; + } + const Node& n = nodes_[idx]; + if (n.is_leaf) { + return freq[n.symbol]; + } + return subtree_weight(n.left, freq) + subtree_weight(n.right, freq); + } + + using SymbolsByLength = + std::array, kMaxCodeLength + 1>; + using LengthCounts = std::array; + + /** @brief Rebuild a Huffman tree from code lengths, reshaped to maximize + * flat-subtree opportunities. + * @details Instead of building the canonical tree bit-by-bit (which scatters + * same-length symbols across the tree), this builder groups + * same-length symbols into power-of-two flat subtrees, then + * assembles them into a tree whose root-to-leaf path lengths still + * equal the original code lengths. This is the paper's "non- + * canonical subtree" optimization (Section 2.2.4): same average + * code lengths, different (better) tree shape. */ + void build_noncanonical_tree( + const std::array& lengths) const { + nodes_.clear(); + nodes_.reserve(2 * kAlphabet); + + SymbolsByLength by_length{}; + LengthCounts length_counts{}; + std::size_t symbol_count = 0; + std::uint8_t only_symbol = 0; + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + const std::uint8_t length = lengths[symbol]; + if (length == 0) { + continue; + } + by_length[length][length_counts[length]++] = + static_cast(symbol); + only_symbol = static_cast(symbol); + ++symbol_count; + } + + // Special case: single distinct symbol → root is a leaf. + if (symbol_count == 1) { + nodes_.push_back(Node{}); + root_ = 0; + nodes_[0].is_leaf = true; + nodes_[0].symbol = only_symbol; + return; + } + + // Build the reshaped tree using the power-of-two grouping strategy. + build_reshaped_tree(by_length, length_counts); + } + + /** @brief Build a non-canonical tree that groups same-length symbols into + * power-of-two flat subtrees. + * @details Uses a recursive slot-filling approach. At each depth, we have + * `slots` open positions. We fill them with leaves (grouped into + * power-of-2 balanced subtrees for flat-node detection) and + * internal nodes (which recurse deeper). The split is determined + * by the Kraft equality: `leaves_at_depth + 2*internal_at_depth = + * slots`, and `internal_at_depth` creates the slots for the next + * depth. */ + /** @brief Build a balanced binary subtree of depth @p d with @p count + * leaves (count must be 2^d). + * @param symbols Array of symbol values (count entries). + * @param count Number of symbols (must be 2^d). + * @param d Depth of the subtree (0 = single leaf). + * @returns Node index of the subtree root. */ + std::size_t build_balanced_subtree(const std::uint8_t* symbols, + std::size_t count, + std::uint8_t d) const { + if (d == 0) { + nodes_.push_back(Node{}); + std::size_t idx = nodes_.size() - 1; + nodes_[idx].is_leaf = true; + nodes_[idx].symbol = symbols[0]; + return idx; + } + std::size_t half = count / 2; + std::size_t left = build_balanced_subtree(symbols, half, d - 1); + std::size_t right = build_balanced_subtree(symbols + half, half, d - 1); + nodes_.push_back(Node{}); + std::size_t idx = nodes_.size() - 1; + nodes_[idx].left = left; + nodes_[idx].right = right; + return idx; + } + + void build_reshaped_tree(const SymbolsByLength& by_length, + const LengthCounts& length_counts) const { + // Per-depth consumption cursor. + LengthCounts cursor{}; + root_ = build_slots(by_length, length_counts, cursor, 1, 2); + } + + /** @brief Fill `slots` open positions at `depth` with leaves and internal + * nodes, returning the root of the resulting balanced subtree. + * @param by_length Symbols grouped by code length. + * @param cursor Per-depth consumption cursor (how many symbols used). + * @param depth Current tree depth (1 = children of root). + * @param slots Number of open positions (must be a power of 2). */ + std::size_t build_slots(const SymbolsByLength& by_length, + const LengthCounts& length_counts, + LengthCounts& cursor, + std::uint8_t depth, + std::size_t slots) const { + if (slots == 0) { + return kNpos; + } + + // How many leaves to place at this depth. + std::size_t available = length_counts[depth] - cursor[depth]; + std::size_t leaves_here = std::min(available, slots); + std::size_t internal = slots - leaves_here; + + // Build children left-to-right: first leaves (grouped), then internal. + std::array children{}; + std::size_t child_count = 0; + + // Group leaves into power-of-2 balanced subtrees. + std::size_t leaf_off = cursor[depth]; + std::size_t remaining = leaves_here; + while (remaining > 0) { + std::uint8_t d = 0; + std::size_t subset = 1; + while (subset * 2 <= remaining && d < kMaxFlatDepth) { + subset *= 2; + d++; + } + children[child_count++] = + build_balanced_subtree(by_length[depth].data() + leaf_off, subset, d); + leaf_off += subset; + remaining -= subset; + } + cursor[depth] += leaves_here; + + // Internal nodes: each recurses with 2 slots at depth+1. + for (std::size_t i = 0; i < internal; i++) { + children[child_count++] = + build_slots(by_length, length_counts, cursor, depth + 1, 2); + } + + // Assemble children into a balanced binary tree. + return assemble_balanced( + std::span(children.data(), child_count)); + } + + /** @brief Assemble a list of subtree roots into a balanced binary tree. + * @details Recursively pairs children left-right, creating internal nodes. + */ + std::size_t assemble_balanced(std::span children) const { + if (children.empty()) { + return kNpos; + } + if (children.size() == 1) { + return children[0]; + } + // Pair up: left half vs right half. + std::size_t mid = children.size() / 2; + std::size_t left = assemble_balanced(children.first(mid)); + std::size_t right = assemble_balanced(children.subspan(mid)); + nodes_.push_back(Node{}); + std::size_t idx = nodes_.size() - 1; + nodes_[idx].left = left; + nodes_[idx].right = right; + return idx; + } + + /** @brief Post-order assignment of exact per-node weights and arena offsets. + * @details The weight of an internal node is the sum of leaf frequencies in + * its subtree (exact, not a parent-derived upper bound). Assigns + * `bits.count` and `bits.offset` for each internal node and + * accumulates the total arena word count. Returns the subtree + * weight so each parent can sum its two children. */ + std::size_t assign_weights_and_offsets( + std::size_t idx, + const std::array& freq, + std::size_t& arena_words) { + if (idx == kNpos) { + return 0; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return freq[n.symbol]; + } + + if (n.flat_depth > 0) { + // Flat node: compute weight from subtree, assign arena for its + // d-bits-per-symbol bitmap, but DON'T recurse into children (their + // bitmaps aren't stored — the flat node replaces them). + std::size_t weight = subtree_weight(idx, freq); + n.bits.count = weight; + n.bits.offset = arena_words; + arena_words += n.bits.word_count(n.flat_depth); + return weight; + } + + std::size_t left_weight = + assign_weights_and_offsets(n.left, freq, arena_words); + std::size_t right_weight = + assign_weights_and_offsets(n.right, freq, arena_words); + n.bits.count = left_weight + right_weight; + n.bits.offset = arena_words; + arena_words += n.bits.word_count(0); + return n.bits.count; + } + + /** @brief Exact weight of a child node, used to place partitions. + * @details Internal nodes carry their precomputed `bits.count`; leaves + * carry their symbol frequency (no bitmap stored). */ + std::size_t child_weight( + std::size_t idx, + const std::array& freq) const { + const Node& n = nodes_[idx]; + return n.is_leaf ? freq[n.symbol] : n.bits.count; + } + + /** @brief Assign compact MSB-first codes with a depth-first tree walk. */ + void assign_codes(std::size_t idx, + std::uint16_t bits, + std::uint8_t length, + std::array& codes) const { + if (nodes_[idx].is_leaf) { + codes[nodes_[idx].symbol] = {bits, length}; + return; + } + assign_codes(nodes_[idx].left, static_cast(bits << 1), + length + 1, codes); + assign_codes(nodes_[idx].right, static_cast((bits << 1) | 1), + length + 1, codes); + } + + /** @brief Build the per-depth direction tables shared by all tree nodes. + * @details A symbol can occur in only one node at a given depth, so one + * 256-byte table per depth replaces rebuilding the same direction + * information separately for every internal node. */ + static DirectionTables build_direction_tables( + const std::array& codes) { + DirectionTables directions{}; + for (std::size_t symbol = 0; symbol < kAlphabet; ++symbol) { + const SymbolCode code = codes[symbol]; + for (std::uint8_t depth = 0; depth < code.length; ++depth) { + directions[depth][symbol] = static_cast( + (code.bits >> (code.length - depth - 1)) & 1u); + } + } + return directions; + } + + /** @brief Pack one flat node's fixed-width wire codes. + * @details Inverts the node's small decode table once, then packs directly + * from the symbol-to-wire-code lookup. Widths dividing a byte get + * dedicated loops; other widths track word position without a + * division per symbol. */ + static void encode_flat_values(std::uint64_t* bits, + const symbol_type* src, + std::size_t count, + std::uint8_t flat_depth, + const symbol_type* table, + bool known_bitreverse) { +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + if (known_bitreverse) { + pivco_flat_encode_neon(bits, src, nullptr, count, flat_depth, true); + return; + } +#endif + + // Three zero pad bytes make overlapping dword gathers at symbol 255 safe + // in the AVX2 flat encoder. + std::array wire_code{}; + const std::size_t alphabet_size = std::size_t{1} << flat_depth; + for (std::size_t code = 0; code < alphabet_size; ++code) { + wire_code[table[code]] = static_cast(code); + } + +#if defined(__AVX2__) && defined(__SSE4_1__) && !defined(__AVX512VBMI2__) + if (flat_depth == 8 && + pivco_flat_encode_d8_bitreverse(bits, src, wire_code.data(), count, + known_bitreverse)) { + return; + } +#if defined(__BMI2__) + if (flat_depth >= 2 && flat_depth <= 7 && + pivco_flat_encode_bitreverse(bits, src, wire_code.data(), count, + flat_depth, known_bitreverse)) { + return; + } + if (flat_depth == 3) { + pivco_flat_encode_d3(bits, src, wire_code.data(), count); + return; + } +#endif +#endif + +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + pivco_flat_encode_neon(bits, src, wire_code.data(), count, flat_depth, + false); + return; +#endif + + auto* bytes = reinterpret_cast(bits); + if (flat_depth == 8) { + for (std::size_t i = 0; i < count; ++i) { + bytes[i] = wire_code[src[i]]; + } + return; + } + if (flat_depth == 4) { + std::size_t i = 0; + for (; i + 2 <= count; i += 2) { + bytes[i / 2] = static_cast(wire_code[src[i]] | + (wire_code[src[i + 1]] << 4)); + } + if (i < count) { + bytes[i / 2] = wire_code[src[i]]; + } + return; + } + if (flat_depth == 2) { + std::size_t i = 0; + for (; i + 4 <= count; i += 4) { + bytes[i / 4] = static_cast( + wire_code[src[i]] | (wire_code[src[i + 1]] << 2) | + (wire_code[src[i + 2]] << 4) | (wire_code[src[i + 3]] << 6)); + } + if (i < count) { + std::uint8_t packed = 0; + for (std::size_t lane = 0; i + lane < count; ++lane) { + packed |= + static_cast(wire_code[src[i + lane]] << (2 * lane)); + } + bytes[i / 4] = packed; + } + return; + } + + std::size_t word_index = 0; + std::size_t bit_offset = 0; + for (std::size_t i = 0; i < count; ++i) { + const std::uint64_t packed = wire_code[src[i]]; + bits[word_index] |= packed << bit_offset; + if (bit_offset + flat_depth > kWordBits) { + bits[word_index + 1] |= packed >> (kWordBits - bit_offset); + } + bit_offset += flat_depth; + if (bit_offset >= kWordBits) { + bit_offset -= kWordBits; + ++word_index; + } + } + } + + /** @brief Top-down in-place partition encoding into the reusable workspace. + * @details The dual of decode: at each internal node, read the incoming + * symbol stream from workspace half `depth`, fill the node's + * arena bitmap with direction bits, and partition symbols into + * half `depth+1` (left at `out_base`, right at + * `out_base + left_weight`). Leaves return immediately. No + * per-node allocations — the workspace is pre-sized and + * ping-ponged by depth, exactly mirroring decode. + * @param idx Current node index. + * @param depth Current tree depth (selects workspace half + path index). + * @param out_base Offset within the current half where this node's stream + * begins. + * @param directions Per-depth symbol direction lookup tables. + * @param freq Symbol frequencies (sizing child partitions). */ + void encode_partition(std::size_t idx, + std::size_t depth, + std::size_t out_base, + const DirectionTables& directions, + const std::array& freq) { + if (nodes_[idx].is_leaf) { + return; + } + + Node& n = nodes_[idx]; + if (n.bits.count == 0) { + return; + } + + // Flat node: write d-bit suffix per symbol, no partition, no recursion. + if (n.flat_depth > 0) { + std::uint64_t* bits = arena_.data() + n.bits.offset; + const symbol_type* src = workspace_half(depth) + out_base; + if (n.flat_depth == 1) { + pivco_partition_encode_none(bits, src, directions[depth].data(), + n.bits.count); + } else { + const symbol_type* table = flat_tables_.data() + n.flat_table_offset; + encode_flat_values(bits, src, n.bits.count, n.flat_depth, table, + n.flat_bitreverse); + } + return; + } + + // Normal node: 1-bit-per-symbol bitmap + partition into children. + const std::size_t left_weight = child_weight(n.left, freq); + const std::size_t weight = n.bits.count; + + const symbol_type* src = workspace_half(depth) + out_base; + symbol_type* left_dst = workspace_half(depth + 1) + out_base; + symbol_type* right_dst = workspace_half(depth + 1) + out_base + left_weight; + + std::uint64_t* bits = arena_.data() + n.bits.offset; + const std::uint8_t* dir = directions[depth].data(); + const bool left_is_leaf = nodes_[n.left].is_leaf; + const bool right_is_leaf = nodes_[n.right].is_leaf; + if (left_is_leaf && right_is_leaf) { + pivco_partition_encode_cst_cst(bits, src, nodes_[n.right].symbol, weight); + } else if (right_is_leaf) { +#if ((defined(__AVX2__) && defined(__SSE4_1__)) || \ + defined(PIXIE_PIVCO_NEON_SUPPORT)) && \ + !defined(__AVX512VBMI2__) + pivco_partition_encode_vec_cst( + left_dst, bits, src, nodes_[n.right].symbol, weight, left_weight); +#else + pivco_partition_encode_left(left_dst, bits, src, dir, weight, + left_weight); +#endif + } else if (left_is_leaf) { +#if ((defined(__AVX2__) && defined(__SSE4_1__)) || \ + defined(PIXIE_PIVCO_NEON_SUPPORT)) && \ + !defined(__AVX512VBMI2__) + pivco_partition_encode_cst_vec(right_dst, bits, src, + nodes_[n.left].symbol, weight, + weight - left_weight); +#else + pivco_partition_encode_right(right_dst, bits, src, dir, weight, + weight - left_weight); +#endif + } else { + pivco_partition_encode(left_dst, right_dst, bits, src, dir, weight, + left_weight, weight - left_weight); + } + + encode_partition(n.left, depth + 1, out_base, directions, freq); + encode_partition(n.right, depth + 1, out_base + left_weight, directions, + freq); + } + + // --- decode -------------------------------------------------------------- + + /** @brief Decode the current block directly into its final output slice. + * @details Internal levels still ping-pong through the reusable two-block + * workspace, but the root writes into @p output. This avoids a + * temporary block vector and the subsequent append/copy. */ + void decode_block_into(std::span output) const { + if (root_ == kNpos || output.empty()) { + return; + } + if (nodes_[root_].is_leaf) { + std::fill(output.begin(), output.end(), nodes_[root_].symbol); + return; + } + workspace_.resize(output.size() * 2 + kWorkspaceSlack); + decode_node(root_, output.size(), 0, 0, output.data()); + } + + /** @brief Recursively decode a node's output stream bottom-up into the + * reusable workspace. + * @param weight Number of symbols this node must produce. + * @param depth Tree depth. The root writes directly to the caller's + * output; deeper levels select a workspace half. + * @param out_base Offset within the selected half where this node's output + * begins. Children are placed at `out_base` (left) and + * `out_base + left_weight` (right) in the opposite half, + * then merged into `[out_base, out_base + weight)` here. + * @param block_output Final output address used by the root node. */ + void decode_node(std::size_t idx, + std::size_t weight, + std::size_t depth, + std::size_t out_base, + symbol_type* block_output) const { + if (weight == 0) { + return; + } + const Node& n = nodes_[idx]; + symbol_type* dst = + depth == 0 ? block_output : workspace_half(depth) + out_base; + if (n.is_leaf) { + std::fill(dst, dst + weight, n.symbol); + return; + } + + // Flat node: read d-bit suffixes, lookup symbols, no merge. + if (n.flat_depth > 0) { + const std::uint64_t* bits = arena_.data() + n.bits.offset; + const symbol_type* table = flat_tables_.data() + n.flat_table_offset; + pivco_flat_decode(dst, bits, table, weight, n.flat_depth, + n.flat_bitreverse); + return; + } + + // Normal node: decode non-leaf children, then select the paper's matching + // MVV/MCV/MVC/MCC merge primitive. Internal children already carry their + // exact stream size, so no bitmap popcount pass is needed. + const Node& left = nodes_[n.left]; + const Node& right = nodes_[n.right]; + const std::uint64_t* bits = arena_.data() + n.bits.offset; + if (left.is_leaf && right.is_leaf) { + pivco_merge_decode_cst_cst(dst, left.symbol, right.symbol, bits, weight); + return; + } + + const std::size_t left_weight = + left.is_leaf ? weight - right.bits.count : left.bits.count; + const std::size_t right_weight = weight - left_weight; + if (!left.is_leaf) { + decode_node(n.left, left_weight, depth + 1, out_base, block_output); + } + if (!right.is_leaf) { + decode_node(n.right, right_weight, depth + 1, out_base + left_weight, + block_output); + } + + const symbol_type* src = workspace_half(depth + 1) + out_base; + const symbol_type* left_out = src; + const symbol_type* right_out = src + left_weight; + if (left.is_leaf) { + pivco_merge_decode_cst_vec(dst, left.symbol, right_out, bits, weight); + } else if (right.is_leaf) { + pivco_merge_decode_vec_cst(dst, left_out, right.symbol, bits, weight); + } else { + pivco_merge_decode(dst, left_out, right_out, bits, weight); + } + } + + /** @brief Pointer to the workspace half that holds depth-@p depth outputs. */ + symbol_type* workspace_half(std::size_t depth) const { + return workspace_.data() + (depth % 2) * block_uncompressed_size_; + } + + // --- serialization ------------------------------------------------------- + + /** @brief Append the shared file-level code lengths as packed nibbles. */ + void serialize_code_lengths() { + for (std::size_t i = 0; i < kAlphabet; i += 2) { + std::uint8_t byte = static_cast( + code_lengths_[i] | (code_lengths_[i + 1] << 4)); + write(byte); + } + } + + /** @brief Append one block's serialized payload to `compressed_`. + * @details Per-block wire format: + * - block_uncompressed_size (8 bytes) + * - For each internal node (pre-order): bits.count (8 bytes) + + * packed words. Leaves and topology are implied by the shared + * file-level code lengths. */ + void serialize_block_payload() { + write(block_uncompressed_size_); + if (block_uncompressed_size_ == 0) { + return; + } + + // Write internal node bitmaps in pre-order (matches deserialize order). + serialize_node(root_); + } + + /** @brief Serialize a node's bitmap and recurse (pre-order). + * @details Flat nodes are terminal: their children's bitmaps are not + * serialized (the flat node replaces them). */ + void serialize_node(std::size_t idx) { + if (idx == kNpos) { + return; + } + const auto& n = nodes_[idx]; + if (n.is_leaf) { + return; + } + write(n.bits.count); + write_words(n.bits.offset, n.bits.word_count(n.flat_depth)); + if (n.flat_depth > 0) { + return; // flat: children not serialized + } + serialize_node(n.left); + serialize_node(n.right); + } + + /** @brief Load the serialized stream and reconstruct its shared tree. + * @details Copies @p data into `compressed_` and reads the top-level + * header and file-level code lengths. Per-block bitmap payloads + * are loaded on demand during `decode_impl()`. */ + void deserialize(std::span data) { + compressed_.assign(data.begin(), data.end()); + nodes_.clear(); + arena_.clear(); + root_ = kNpos; + if (data.empty()) { + uncompressed_size_ = 0; + blocks_offset_ = 0; + return; + } + std::size_t pos = 0; + uncompressed_size_ = read(compressed_, pos); + block_size_ = read(compressed_, pos); + if (uncompressed_size_ == 0) { + blocks_offset_ = pos; + return; + } + + for (std::size_t i = 0; i < kAlphabet; i += 2) { + const std::uint8_t byte = read(compressed_, pos); + code_lengths_[i] = byte & 0x0F; + code_lengths_[i + 1] = (byte >> 4) & 0x0F; + } + build_noncanonical_tree(code_lengths_); + detect_flat_subtrees(); + blocks_offset_ = pos; + } + + /** @brief Load one block's arena from `compressed_` at @p pos. + * @details Reuses the file-level tree, reads the block's uncompressed size, + * and loads its arena words via a two-pass wire walk. Advances + * @p pos past the block payload. */ + void deserialize_block_payload(std::size_t& pos) const { + block_uncompressed_size_ = read(compressed_, pos); + if (block_uncompressed_size_ == 0) { + return; + } + + // Single-symbol input: root is a leaf, no internal nodes, no arena. + if (root_ == kNpos || nodes_[root_].is_leaf) { + return; + } + + // Pass 1: read each internal node's `count` (and skip its word bytes on + // the wire), assign contiguous arena offsets, and size the arena exactly. + const std::size_t bitmap_start = pos; + std::size_t arena_words = 0; + scan_node_counts(root_, compressed_, pos, arena_words); + arena_.assign(arena_words, 0); + + // Pass 2: rewind to the bitmap section and read each node's packed words + // into its arena slot. + pos = bitmap_start; + read_node_words(root_, compressed_, pos); + } + + /** @brief Pass 1: read each internal node's `count` from the wire, assign its + * arena offset, and advance past its word bytes without storing + * them. Sums the total arena word count so the caller allocates once. + * @details Pre-order traversal matches the serialization order. Flat nodes + * are terminal (children not visited). */ + void scan_node_counts(std::size_t idx, + std::span data, + std::size_t& pos, + std::size_t& arena_words) const { + if (idx == kNpos) { + return; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return; + } + n.bits.count = read(data, pos); + n.bits.offset = arena_words; + arena_words += n.bits.word_count(n.flat_depth); + pos += n.bits.word_count(n.flat_depth) * sizeof(std::uint64_t); + if (n.flat_depth > 0) { + return; // flat: children not on the wire + } + scan_node_counts(n.left, data, pos, arena_words); + scan_node_counts(n.right, data, pos, arena_words); + } + + /** @brief Pass 2: read each internal node's packed words from the wire into + * its arena slot (pre-order, matching serialization). Flat nodes are + * terminal. */ + void read_node_words(std::size_t idx, + std::span data, + std::size_t& pos) const { + if (idx == kNpos) { + return; + } + Node& n = nodes_[idx]; + if (n.is_leaf) { + return; + } + pos += sizeof(std::size_t); // count, already read in pass 1 + const std::size_t words = n.bits.word_count(n.flat_depth); + const std::size_t bytes = words * sizeof(std::uint64_t); + if (bytes > 0) { + std::memcpy(arena_.data() + n.bits.offset, data.data() + pos, bytes); + } + pos += bytes; + if (n.flat_depth > 0) { + return; // flat: children not on the wire + } + read_node_words(n.left, data, pos); + read_node_words(n.right, data, pos); + } + + /** @brief Append a little-endian, fixed-width value to the byte buffer. */ + template + void write(const T& value) { + const std::size_t old = compressed_.size(); + compressed_.resize(old + sizeof(T)); + std::memcpy(compressed_.data() + old, &value, sizeof(T)); + } + + /** @brief Append a node's arena-backed words directly to the byte buffer. */ + void write_words(std::size_t offset, std::size_t word_count) { + const std::size_t bytes = word_count * sizeof(std::uint64_t); + const std::size_t old = compressed_.size(); + compressed_.resize(old + bytes); + if (bytes > 0) { + std::memcpy(compressed_.data() + old, arena_.data() + offset, bytes); + } + } + + /** @brief Read a fixed-width value from @p data at @p pos. */ + template + static T read(std::span data, std::size_t& pos) { + T value; + std::memcpy(&value, data.data() + pos, sizeof(T)); + pos += sizeof(T); + return value; + } +}; + +} // namespace pixie diff --git a/include/pixie/huffman/pivco_simd.h b/include/pixie/huffman/pivco_simd.h new file mode 100644 index 0000000..bafe099 --- /dev/null +++ b/include/pixie/huffman/pivco_simd.h @@ -0,0 +1,2566 @@ +#pragma once + +/** + * @file pivco_simd.h + * @brief SIMD kernels for PivCo-Huffman encode/decode primitives. + * + * Implements the two core PivCo-Huffman primitives from the paper with + * architecture-specific kernels and a portable scalar fallback: + * + * - `pivco_merge_decode` (bottom-up decode, `merge_vec_vec`): interleaves two + * dense child symbol streams under a 1-bit-per-symbol routing bitmap into a + * single dense output. A set bitmap bit selects the right child, a clear bit + * the left child. + * - `pivco_partition_encode` (top-down encode partition): the inverse — given + * per-symbol direction bits, splits a dense input into left/right child + * streams and writes the routing bitmap. + * - `pivco_flat_decode` (bottom-up `merge_flat_D`): unpacks fixed-width leaf + * indices and translates them through a flat-subtree symbol table. + * + * Bitmap convention (matches `pivco_huffman.h`): bits are packed LSB-first + * into 64-bit words; bit 0 of word 0 is the first symbol. + * + * Backends, selected by compiler feature macros (`-march`): + * - AVX-512 VBMI + VBMI2: byte-granular `vpermb`, `vpmultishiftqb`, + * `vpexpandb`, and `vpcompressb` operating on 64 symbols per iteration. + * These extensions implement vector direction lookup, partition, and the + * full `merge_flat_D` family without changing the wire format. + * - AVX2 + SSE4.1: 8-byte-at-a-time `_mm_shuffle_epi8` with 256-entry lookup + * tables, mirroring the paper's NEON `vqtbl1q_u8` technique. + * - AArch64 NEON: paper-style `vqtbl*`/`vqtbx*` table lookup, compact and + * expand tables, and 16-symbol fixed-width unpacking. + * - Scalar: the original bit-by-bit / symbol-by-symbol loops. + * + * Every SIMD backend produces byte-identical output to the scalar fallback + * (the operation is a deterministic permutation). The existing round-trip tests + * under release and AddressSanitizer serve as the differential oracle for every + * backend the test matrix builds and runs; the AVX-512 backend is verified only + * when built on VBMI/VBMI2-capable hardware, since there is no runtime + * dispatch. + */ + +#include +#include +#include +#include +#include + +#if defined(__AVX2__) || defined(__AVX512VBMI2__) +#include +#endif + +#if defined(__aarch64__) && (defined(__ARM_NEON) || defined(__ARM_NEON__)) +#include +#define PIXIE_PIVCO_NEON_SUPPORT 1 +#endif + +namespace pixie::pivco_simd_detail { + +/// @brief Bitmap bit-ordering convention used throughout this header. +/// @details Every routing bitmap packs one bit per symbol, LSB-first within +/// each +/// 64-bit word: bit `k` is the direction bit of symbol `k`, where a +/// set bit selects the right child and a clear bit the left child. All +/// mask builders (`build_dir_mask8`, the AVX-512 scalar loop), LUT +/// tables, merge kernels, and partition tails MUST agree on this +/// convention — flipping it in one backend only would make the +/// serialized bitmap backend-dependent and silently break +/// cross-`-march` encode/decode. + +// --------------------------------------------------------------------------- +// Scalar helpers (always compiled; used as tail + fallback). +// --------------------------------------------------------------------------- + +/// @brief Build an 8-bit routing mask for symbols @p s8 (one byte per symbol) +/// under the LSB-first convention (bit k = dir[s8[k]]). +/// @details Shared source of truth for the 8-symbol direction mask; the AVX2 +/// partition kernel and (transitively) the AVX-512 tail both build +/// their per-group mask here. +inline std::uint8_t build_dir_mask8(const std::uint8_t* s8, + const std::uint8_t* dir) noexcept { + return static_cast((dir[s8[0]] << 0) | (dir[s8[1]] << 1) | + (dir[s8[2]] << 2) | (dir[s8[3]] << 3) | + (dir[s8[4]] << 4) | (dir[s8[5]] << 5) | + (dir[s8[6]] << 6) | (dir[s8[7]] << 7)); +} + +/// @brief Scalar merge tail: process the remaining `< 8` symbols bit-by-bit. +template +inline void merge_scalar_tail(std::uint8_t* dst, + std::size_t& out_i, + const std::uint8_t* left, + std::uint8_t left_symbol, + std::size_t& c_left, + const std::uint8_t* right, + std::uint8_t right_symbol, + std::size_t& c_right, + const std::uint64_t* bits, + std::size_t i, + std::size_t count) noexcept { + for (; i < count; ++i) { + const std::uint8_t bit = + static_cast((bits[i / 64] >> (i % 64)) & 1u); + if (bit) { + if constexpr (kRightConstant) { + dst[out_i++] = right_symbol; + } else { + dst[out_i++] = right[c_right]; + } + ++c_right; + } else { + if constexpr (kLeftConstant) { + dst[out_i++] = left_symbol; + } else { + dst[out_i++] = left[c_left]; + } + ++c_left; + } + } +} + +/// @brief Write only the routing bitmap for a scalar partition tail. +inline void partition_bitmap_tail(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t start_i, + std::size_t weight) noexcept { + std::size_t word_idx = start_i / 64; + std::size_t bit_pos = start_i % 64; + for (std::size_t i = start_i; i < weight; ++i) { + if (dir[src[i]]) { + bits[word_idx] |= (1ull << bit_pos); + } + if (++bit_pos == 64) { + bit_pos = 0; + ++word_idx; + } + } +} + +/// @brief Scalar partition tail: route the remaining symbols (fewer than the +/// vector width) one at a time, writing the bitmap bit and appending to +/// the left/right child stream. +/// @details Shared by the AVX2 kernel tail, the AVX-512 kernel tail, and the +/// scalar fallback, so the LSB-first `|=`-into-pre-zeroed-arena +/// convention lives in exactly one place. @p bits is the bitmap base +/// (word array); @p start_bit is the absolute bit offset where the +/// tail begins (must be a multiple of the full-group width so it is +/// already byte-aligned within the word packing). +template +inline void partition_scalar_tail(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t& c_left, + std::size_t& c_right, + std::size_t start_i, + std::size_t weight) noexcept { + std::size_t i = start_i; + std::size_t word_idx = i / 64; + std::size_t bit_pos = i % 64; + for (; i < weight; ++i) { + const std::uint8_t s = src[i]; + if (dir[s]) { + bits[word_idx] |= (1ull << bit_pos); + if constexpr (kWriteRight) { + right_dst[c_right] = s; + } + ++c_right; + } else { + if constexpr (kWriteLeft) { + left_dst[c_left] = s; + } + ++c_left; + } + if (++bit_pos == 64) { + bit_pos = 0; + ++word_idx; + } + } +} + +// =========================================================================== +// AArch64 NEON backend (table lookup and compaction, 8/16 symbols per batch). +// =========================================================================== +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + +/// @brief 256 eight-byte table controls for one NEON `vqtbl1_u8` operation. +struct NeonShufTable { + alignas(16) std::uint8_t entry[256][8]; +}; + +/// @brief Paper `expand_tab`: select left/right dense elements for each mask. +inline const NeonShufTable& merge_lut_neon() { + static const NeonShufTable lut = [] { + NeonShufTable table{}; + for (int mask = 0; mask < 256; ++mask) { + int left = 0; + int right = 0; + for (int lane = 0; lane < 8; ++lane) { + if ((mask >> lane) & 1) { + table.entry[mask][lane] = static_cast(8 + right++); + } else { + table.entry[mask][lane] = static_cast(left++); + } + } + } + return table; + }(); + return lut; +} + +/// @brief Paper `compress_tab` for elements routed to the right child. +inline const NeonShufTable& compress_right_lut_neon() { + static const NeonShufTable lut = [] { + NeonShufTable table{}; + for (int mask = 0; mask < 256; ++mask) { + int output = 0; + for (int lane = 0; lane < 8; ++lane) { + if ((mask >> lane) & 1) { + table.entry[mask][output++] = static_cast(lane); + } + } + for (; output < 8; ++output) { + table.entry[mask][output] = 0xff; + } + } + return table; + }(); + return lut; +} + +/// @brief Paper `compress_tab` for elements routed to the left child. +inline const NeonShufTable& compress_left_lut_neon() { + static const NeonShufTable lut = [] { + NeonShufTable table{}; + for (int mask = 0; mask < 256; ++mask) { + int output = 0; + for (int lane = 0; lane < 8; ++lane) { + if (((mask >> lane) & 1) == 0) { + table.entry[mask][output++] = static_cast(lane); + } + } + for (; output < 8; ++output) { + table.entry[mask][output] = 0xff; + } + } + return table; + }(); + return lut; +} + +/// @brief Four 64-byte NEON lookup banks, enough for any byte alphabet. +struct ByteLookupNeon { + uint8x16x4_t bank0; + uint8x16x4_t bank1; + uint8x16x4_t bank2; + uint8x16x4_t bank3; +}; + +/// @brief Return four zero vectors in the layout expected by `vqtbl4*`. +inline uint8x16x4_t zero_lookup_bank_neon() noexcept { + const uint8x16_t zero = vdupq_n_u8(0); + uint8x16x4_t bank; + bank.val[0] = zero; + bank.val[1] = zero; + bank.val[2] = zero; + bank.val[3] = zero; + return bank; +} + +/// @brief Load one complete 64-byte `vqtbl4*` bank. +inline uint8x16x4_t load_lookup_bank_neon(const std::uint8_t* table) noexcept { + uint8x16x4_t bank; + bank.val[0] = vld1q_u8(table); + bank.val[1] = vld1q_u8(table + 16); + bank.val[2] = vld1q_u8(table + 32); + bank.val[3] = vld1q_u8(table + 48); + return bank; +} + +/// @brief Load exactly `2^D` table bytes without reading past small tables. +template +inline ByteLookupNeon load_byte_lookup_neon( + const std::uint8_t* table) noexcept { + static_assert(kDepth >= 1 && kDepth <= 8); + ByteLookupNeon lookup{zero_lookup_bank_neon(), zero_lookup_bank_neon(), + zero_lookup_bank_neon(), zero_lookup_bank_neon()}; + if constexpr (kDepth <= 3) { + std::uint64_t entries = 0; + std::memcpy(&entries, table, std::size_t{1} << kDepth); + lookup.bank0.val[0] = vcombine_u8(vcreate_u8(entries), vdup_n_u8(0)); + } else if constexpr (kDepth == 4) { + lookup.bank0.val[0] = vld1q_u8(table); + } else if constexpr (kDepth == 5) { + lookup.bank0.val[0] = vld1q_u8(table); + lookup.bank0.val[1] = vld1q_u8(table + 16); + } else { + lookup.bank0 = load_lookup_bank_neon(table); + if constexpr (kDepth >= 7) { + lookup.bank1 = load_lookup_bank_neon(table + 64); + } + if constexpr (kDepth == 8) { + lookup.bank2 = load_lookup_bank_neon(table + 128); + lookup.bank3 = load_lookup_bank_neon(table + 192); + } + } + return lookup; +} + +/// @brief Translate sixteen byte indices through a preloaded `2^D` table. +template +inline uint8x16_t lookup_bytes16_neon(const ByteLookupNeon& lookup, + uint8x16_t indices) noexcept { + uint8x16_t result = vqtbl4q_u8(lookup.bank0, indices); + if constexpr (kDepth >= 7) { + result = + vqtbx4q_u8(result, lookup.bank1, vsubq_u8(indices, vdupq_n_u8(64))); + } + if constexpr (kDepth == 8) { + result = + vqtbx4q_u8(result, lookup.bank2, vsubq_u8(indices, vdupq_n_u8(128))); + result = + vqtbx4q_u8(result, lookup.bank3, vsubq_u8(indices, vdupq_n_u8(192))); + } + return result; +} + +/// @brief Translate eight byte indices through a preloaded `2^D` table. +template +inline uint8x8_t lookup_bytes8_neon(const ByteLookupNeon& lookup, + uint8x8_t indices) noexcept { + uint8x8_t result = vqtbl4_u8(lookup.bank0, indices); + if constexpr (kDepth >= 7) { + result = vqtbx4_u8(result, lookup.bank1, vsub_u8(indices, vdup_n_u8(64))); + } + if constexpr (kDepth == 8) { + result = vqtbx4_u8(result, lookup.bank2, vsub_u8(indices, vdup_n_u8(128))); + result = vqtbx4_u8(result, lookup.bank3, vsub_u8(indices, vdup_n_u8(192))); + } + return result; +} + +/// @brief Reverse the low D bits independently in sixteen byte lanes. +template +inline uint8x16_t reverse_low_bits16_neon(uint8x16_t values) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + const uint8x16_t reversed = vrbitq_u8(values); + if constexpr (kDepth == 8) { + return reversed; + } else { + return vshrq_n_u8(reversed, 8 - kDepth); + } +} + +/// @brief Scalar bit reverse for the final fewer-than-16 encode symbols. +template +inline std::uint8_t reverse_low_bits_scalar_neon(std::uint8_t value) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + value = static_cast((value >> 4) | (value << 4)); + value = static_cast(((value & 0xccu) >> 2) | + ((value & 0x33u) << 2)); + value = static_cast(((value & 0xaau) >> 1) | + ((value & 0x55u) << 1)); + return static_cast(value >> (8 - kDepth)); +} + +/// @brief Translate flat codes through either a table or a known bit reverse. +template +inline uint8x16_t translate_flat_codes16_neon(const ByteLookupNeon* lookup, + uint8x16_t codes) noexcept { + if constexpr (kBitReverse) { + return reverse_low_bits16_neon(codes); + } else { + return lookup_bytes16_neon(*lookup, codes); + } +} + +/// @brief Eight-lane form of `translate_flat_codes16_neon`. +template +inline uint8x8_t translate_flat_codes8_neon(const ByteLookupNeon* lookup, + uint8x8_t codes) noexcept { + if constexpr (kBitReverse) { + return vget_low_u8( + reverse_low_bits16_neon(vcombine_u8(codes, codes))); + } else { + return lookup_bytes8_neon(*lookup, codes); + } +} + +/// @brief Convert eight lanes containing 0/1 into an LSB-first byte mask. +inline std::uint8_t lane_bits_to_mask_neon(uint8x8_t lane_bits) noexcept { + alignas(8) static constexpr std::uint8_t kBitWeights[8] = {1, 2, 4, 8, + 16, 32, 64, 128}; + return vaddv_u8(vmul_u8(lane_bits, vld1_u8(kBitWeights))); +} + +/// @brief Store only the valid prefix of an eight-byte compaction result. +inline void store_masked_neon(std::uint8_t* dst, + uint8x8_t values, + std::size_t count) noexcept { + if (count >= 8) { + vst1_u8(dst, values); + } else if (count != 0) { + alignas(8) std::uint8_t temporary[8]; + vst1_u8(temporary, values); + std::memcpy(dst, temporary, count); + } +} + +/// @brief Compact one eight-symbol partition group into its requested child +/// streams. +template +inline void partition_group8_neon(std::uint8_t* left_dst, + std::uint8_t* right_dst, + uint8x8_t values, + std::uint8_t routing, + std::size_t left_weight, + std::size_t right_weight, + std::size_t& left_count, + std::size_t& right_count) noexcept { + const uint8x16_t table = vcombine_u8(values, vdup_n_u8(0)); + if constexpr (kWriteRight) { + const uint8x8_t controls = + vld1_u8(compress_right_lut_neon().entry[routing]); + store_masked_neon(right_dst + right_count, vqtbl1_u8(table, controls), + right_weight - right_count); + } + if constexpr (kWriteLeft) { + const uint8x8_t controls = vld1_u8(compress_left_lut_neon().entry[routing]); + store_masked_neon(left_dst + left_count, vqtbl1_u8(table, controls), + left_weight - left_count); + } + const std::size_t routed_right = + std::popcount(static_cast(routing)); + right_count += routed_right; + left_count += 8 - routed_right; +} + +/// @brief Compact the non-constant side of one eight-symbol leaf partition. +template +inline void partition_one_constant_group8_neon(std::uint8_t* output, + uint8x8_t values, + uint8x8_t equal_lanes, + std::size_t output_weight, + std::size_t& output_count, + std::uint8_t& routing) noexcept { + const std::uint8_t equal = lane_bits_to_mask_neon(vshr_n_u8(equal_lanes, 7)); + routing = kLeftConstant ? static_cast(~equal) : equal; + const auto& lut = + kLeftConstant ? compress_right_lut_neon() : compress_left_lut_neon(); + const uint8x8_t packed = + vqtbl1_u8(vcombine_u8(values, vdup_n_u8(0)), vld1_u8(lut.entry[routing])); + store_masked_neon(output + output_count, packed, + output_weight - output_count); + const std::size_t routed_right = + std::popcount(static_cast(routing)); + output_count += kLeftConstant ? routed_right : 8 - routed_right; +} + +/// @brief Merge two constant leaves without a shuffle-table load. +inline void merge_decode_cst_cst_neon(std::uint8_t* dst, + std::uint8_t left_symbol, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + alignas(16) static constexpr std::uint8_t kBitPositions[16] = { + 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128}; + const uint8x16_t positions = vld1q_u8(kBitPositions); + const uint8x16_t left = vdupq_n_u8(left_symbol); + const uint8x16_t right = vdupq_n_u8(right_symbol); + const auto* bit_bytes = reinterpret_cast(bits); + std::size_t i = 0; + for (; i + 16 <= count; i += 16) { + const uint8x16_t routing = vcombine_u8(vdup_n_u8(bit_bytes[i / 8]), + vdup_n_u8(bit_bytes[i / 8 + 1])); + vst1q_u8(dst + i, vbslq_u8(vtstq_u8(routing, positions), right, left)); + } + for (; i < count; ++i) { + dst[i] = + ((bits[i / 64] >> (i % 64)) & 1u) != 0 ? right_symbol : left_symbol; + } +} + +/// @brief Paper-style NEON dense/constant merge in eight-symbol batches. +template +inline void merge_decode_neon(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + if constexpr (kLeftConstant && kRightConstant) { + merge_decode_cst_cst_neon(dst, left_symbol, right_symbol, bits, count); + return; + } + + const auto& merge_table = merge_lut_neon(); + const auto* bit_bytes = reinterpret_cast(bits); + std::size_t output = 0; + std::size_t left_count = 0; + std::size_t right_count = 0; + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const std::uint8_t routing = bit_bytes[i / 8]; + uint8x8_t left_values; + if constexpr (kLeftConstant) { + left_values = vdup_n_u8(left_symbol); + } else { + left_values = vld1_u8(left + left_count); + } + uint8x8_t right_values; + if constexpr (kRightConstant) { + right_values = vdup_n_u8(right_symbol); + } else { + right_values = vld1_u8(right + right_count); + } + const uint8x8_t controls = vld1_u8(merge_table.entry[routing]); + vst1_u8(dst + output, + vqtbl1_u8(vcombine_u8(left_values, right_values), controls)); + output += 8; + const std::size_t routed_right = + std::popcount(static_cast(routing)); + right_count += routed_right; + left_count += 8 - routed_right; + } + merge_scalar_tail( + dst, output, left, left_symbol, left_count, right, right_symbol, + right_count, bits, i, count); +} + +/// @brief Paper-style direction lookup plus NEON `compress_tab` partition. +template +inline void partition_encode_neon(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + const ByteLookupNeon direction_lookup = load_byte_lookup_neon<8>(dir); + auto* bit_bytes = reinterpret_cast(bits); + std::size_t left_count = 0; + std::size_t right_count = 0; + std::size_t i = 0; + for (; i + 16 <= weight; i += 16) { + const uint8x16_t values = vld1q_u8(src + i); + const uint8x16_t directions = + lookup_bytes16_neon<8>(direction_lookup, values); + const std::uint8_t low_routing = + lane_bits_to_mask_neon(vget_low_u8(directions)); + const std::uint8_t high_routing = + lane_bits_to_mask_neon(vget_high_u8(directions)); + bit_bytes[i / 8] = low_routing; + bit_bytes[i / 8 + 1] = high_routing; + partition_group8_neon( + left_dst, right_dst, vget_low_u8(values), low_routing, left_weight, + right_weight, left_count, right_count); + partition_group8_neon( + left_dst, right_dst, vget_high_u8(values), high_routing, left_weight, + right_weight, left_count, right_count); + } + for (; i + 8 <= weight; i += 8) { + const uint8x8_t values = vld1_u8(src + i); + const std::uint8_t routing = + lane_bits_to_mask_neon(lookup_bytes8_neon<8>(direction_lookup, values)); + bit_bytes[i / 8] = routing; + partition_group8_neon( + left_dst, right_dst, values, routing, left_weight, right_weight, + left_count, right_count); + } + + if constexpr (kWriteLeft || kWriteRight) { + partition_scalar_tail(left_dst, right_dst, bits, + src, dir, left_count, + right_count, i, weight); + } else { + partition_bitmap_tail(bits, src, dir, i, weight); + } +} + +/// @brief NEON partition when one child is a constant leaf. +template +inline void partition_encode_one_constant_neon( + std::uint8_t* output, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t constant_symbol, + std::size_t weight, + std::size_t output_weight) noexcept { + const uint8x8_t constant = vdup_n_u8(constant_symbol); + auto* bit_bytes = reinterpret_cast(bits); + std::size_t output_count = 0; + std::size_t i = 0; + for (; i + 16 <= weight; i += 16) { + const uint8x16_t values = vld1q_u8(src + i); + const uint8x16_t equal = vceqq_u8(values, vdupq_n_u8(constant_symbol)); + std::uint8_t low_routing = 0; + std::uint8_t high_routing = 0; + partition_one_constant_group8_neon( + output, vget_low_u8(values), vget_low_u8(equal), output_weight, + output_count, low_routing); + partition_one_constant_group8_neon( + output, vget_high_u8(values), vget_high_u8(equal), output_weight, + output_count, high_routing); + bit_bytes[i / 8] = low_routing; + bit_bytes[i / 8 + 1] = high_routing; + } + for (; i + 8 <= weight; i += 8) { + const uint8x8_t values = vld1_u8(src + i); + std::uint8_t routing = 0; + partition_one_constant_group8_neon( + output, values, vceq_u8(values, constant), output_weight, output_count, + routing); + bit_bytes[i / 8] = routing; + } + + for (; i < weight; ++i) { + const bool is_constant = src[i] == constant_symbol; + const bool goes_right = kLeftConstant ? !is_constant : is_constant; + if (goes_right) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + if (kLeftConstant ? goes_right : !goes_right) { + output[output_count++] = src[i]; + } + } +} + +/// @brief Build a two-leaf routing bitmap sixteen symbols at a time. +inline void partition_encode_cst_cst_neon(std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight) noexcept { + const uint8x16_t right = vdupq_n_u8(right_symbol); + auto* bit_bytes = reinterpret_cast(bits); + std::size_t i = 0; + for (; i + 16 <= weight; i += 16) { + const uint8x16_t equal = vceqq_u8(vld1q_u8(src + i), right); + bit_bytes[i / 8] = lane_bits_to_mask_neon(vshr_n_u8(vget_low_u8(equal), 7)); + bit_bytes[i / 8 + 1] = + lane_bits_to_mask_neon(vshr_n_u8(vget_high_u8(equal), 7)); + } + for (; i < weight; ++i) { + if (src[i] == right_symbol) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + } +} + +/// @brief Load exactly the bytes containing sixteen packed D-bit codes. +template +inline uint8x16_t load_packed_codes16_neon( + const std::uint8_t* packed) noexcept { + static_assert(kDepth >= 3 && kDepth <= 8); + constexpr std::size_t kPackedBytes = 2 * kDepth; + std::uint64_t low = 0; + std::uint64_t high = 0; + constexpr std::size_t kLowBytes = kPackedBytes < 8 ? kPackedBytes : 8; + std::memcpy(&low, packed, kLowBytes); + if constexpr (kPackedBytes > 8) { + std::memcpy(&high, packed + 8, kPackedBytes - 8); + } + return vcombine_u8(vcreate_u8(low), vcreate_u8(high)); +} + +/// @brief Unpack sixteen tightly packed D-bit codes into sixteen byte lanes. +template +inline uint8x16_t unpack_codes16_neon(const std::uint8_t* packed) noexcept { + static_assert(kDepth >= 3 && kDepth <= 7); + alignas(16) static constexpr auto kLowIndices = [] { + std::array controls{}; + for (std::size_t lane = 0; lane < controls.size(); ++lane) { + controls[lane] = static_cast((lane * kDepth) / 8); + } + return controls; + }(); + alignas(16) static constexpr auto kHighIndices = [] { + std::array controls{}; + for (std::size_t lane = 0; lane < controls.size(); ++lane) { + controls[lane] = static_cast((lane * kDepth) / 8 + 1); + } + return controls; + }(); + alignas(16) static constexpr auto kRightShifts = [] { + std::array shifts{}; + for (std::size_t lane = 0; lane < shifts.size(); ++lane) { + shifts[lane] = -static_cast((lane * kDepth) % 8); + } + return shifts; + }(); + alignas(16) static constexpr auto kLeftShifts = [] { + std::array shifts{}; + for (std::size_t lane = 0; lane < shifts.size(); ++lane) { + const std::uint8_t offset = (lane * kDepth) % 8; + shifts[lane] = static_cast(offset == 0 ? 8 : 8 - offset); + } + return shifts; + }(); + + const uint8x16_t bytes = load_packed_codes16_neon(packed); + const uint8x16_t low = vqtbl1q_u8(bytes, vld1q_u8(kLowIndices.data())); + const uint8x16_t high = vqtbl1q_u8(bytes, vld1q_u8(kHighIndices.data())); + const uint8x16_t codes = + vorrq_u8(vshlq_u8(low, vld1q_s8(kRightShifts.data())), + vshlq_u8(high, vld1q_s8(kLeftShifts.data()))); + return vandq_u8(codes, vdupq_n_u8((std::uint8_t{1} << kDepth) - 1)); +} + +/// @brief Decode one flat subtree with NEON table lookup or bit reversal. +template +inline void flat_decode_neon(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + ByteLookupNeon lookup; + const ByteLookupNeon* lookup_ptr = nullptr; + if constexpr (!kBitReverse) { + lookup = load_byte_lookup_neon(table); + lookup_ptr = &lookup; + } + const auto* packed = reinterpret_cast(bits); + std::size_t i = 0; + if constexpr (kDepth == 2) { + const uint8x8_t mask = vdup_n_u8(0x03); + for (; i + 32 <= count; i += 32) { + const uint8x8_t bytes = vld1_u8(packed + i / 4); + uint8x8x4_t symbols; + symbols.val[0] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vand_u8(bytes, mask)); + symbols.val[1] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vand_u8(vshr_n_u8(bytes, 2), mask)); + symbols.val[2] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vand_u8(vshr_n_u8(bytes, 4), mask)); + symbols.val[3] = translate_flat_codes8_neon<2, kBitReverse>( + lookup_ptr, vshr_n_u8(bytes, 6)); + vst4_u8(dst + i, symbols); + } + } else if constexpr (kDepth == 4) { + const uint8x8_t mask = vdup_n_u8(0x0f); + for (; i + 16 <= count; i += 16) { + const uint8x8_t bytes = vld1_u8(packed + i / 2); + uint8x8x2_t symbols; + symbols.val[0] = translate_flat_codes8_neon<4, kBitReverse>( + lookup_ptr, vand_u8(bytes, mask)); + symbols.val[1] = translate_flat_codes8_neon<4, kBitReverse>( + lookup_ptr, vshr_n_u8(bytes, 4)); + vst2_u8(dst + i, symbols); + } + } else if constexpr (kDepth == 8) { + for (; i + 16 <= count; i += 16) { + const uint8x16_t codes = vld1q_u8(packed + i); + const uint8x16_t symbols = + translate_flat_codes16_neon<8, kBitReverse>(lookup_ptr, codes); + vst1q_u8(dst + i, symbols); + } + } else { + for (; i + 16 <= count; i += 16) { + const auto* group = packed + (i * kDepth) / 8; + const uint8x16_t codes = unpack_codes16_neon(group); + const uint8x16_t symbols = + translate_flat_codes16_neon(lookup_ptr, codes); + vst1q_u8(dst + i, symbols); + } + } + + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + for (; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t code = bits[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + code |= bits[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[code & kCodeMask]; + } +} + +/// @brief Densely pack eight byte codes into the low `8 * D` bits. +template +inline std::uint64_t pack_codes8_neon(std::uint64_t byte_codes) noexcept { + static_assert(kDepth >= 3 && kDepth <= 8); + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + std::uint64_t dense = 0; + for (std::size_t lane = 0; lane < 8; ++lane) { + dense |= ((byte_codes >> (lane * 8)) & kCodeMask) << (lane * kDepth); + } + return dense; +} + +/// @brief Encode flat wire codes with NEON lookup/bit-reverse and fixed pack. +template +inline void flat_encode_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + ByteLookupNeon lookup; + const ByteLookupNeon* lookup_ptr = nullptr; + if constexpr (!kBitReverse) { + lookup = load_byte_lookup_neon<8>(wire_code); + lookup_ptr = &lookup; + } + auto* packed = reinterpret_cast(bits); + std::size_t i = 0; + if constexpr (kDepth == 8) { + for (; i + 16 <= count; i += 16) { + const uint8x16_t symbols = vld1q_u8(src + i); + const uint8x16_t codes = + translate_flat_codes16_neon<8, kBitReverse>(lookup_ptr, symbols); + vst1q_u8(packed + i, codes); + } + for (; i < count; ++i) { + if constexpr (kBitReverse) { + packed[i] = reverse_low_bits_scalar_neon<8>(src[i]); + } else { + packed[i] = wire_code[src[i]]; + } + } + return; + } + + for (; i + 16 <= count; i += 16) { + uint8x16_t codes; + if constexpr (kBitReverse) { + codes = reverse_low_bits16_neon(vld1q_u8(src + i)); + } else { + codes = lookup_bytes16_neon<8>(*lookup_ptr, vld1q_u8(src + i)); + } + if constexpr (kDepth == 2) { + const uint8x16_t even = vuzp1q_u8(codes, codes); + const uint8x16_t odd = vuzp2q_u8(codes, codes); + const uint8x8_t pairs = vget_low_u8( + vorrq_u8(even, vshlq_n_u8(odd, static_cast(kDepth)))); + const uint8x8_t packed_codes = + vorr_u8(vuzp1_u8(pairs, pairs), vshl_n_u8(vuzp2_u8(pairs, pairs), 4)); + const std::uint32_t dense = + vget_lane_u32(vreinterpret_u32_u8(packed_codes), 0); + std::memcpy(packed + i / 4, &dense, sizeof(dense)); + } else if constexpr (kDepth == 4) { + const uint8x16_t even = vuzp1q_u8(codes, codes); + const uint8x16_t odd = vuzp2q_u8(codes, codes); + vst1_u8(packed + i / 2, vget_low_u8(vorrq_u8(even, vshlq_n_u8(odd, 4)))); + } else { + const uint64x2_t byte_codes = vreinterpretq_u64_u8(codes); + const std::uint64_t low = + pack_codes8_neon(vgetq_lane_u64(byte_codes, 0)); + const std::uint64_t high = + pack_codes8_neon(vgetq_lane_u64(byte_codes, 1)); + auto* output = packed + (i * kDepth) / 8; + std::memcpy(output, &low, kDepth); + std::memcpy(output + kDepth, &high, kDepth); + } + } + + std::size_t word_index = (i * kDepth) / 64; + std::size_t bit_offset = (i * kDepth) % 64; + for (; i < count; ++i) { + std::uint64_t code; + if constexpr (kBitReverse) { + code = reverse_low_bits_scalar_neon(src[i]); + } else { + code = wire_code[src[i]]; + } + bits[word_index] |= code << bit_offset; + if (bit_offset + kDepth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + bit_offset += kDepth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +/// @brief Select table lookup or the preclassified balanced-table fast path. +template +inline void flat_encode_maybe_bitreverse_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + bool known_bitreverse) noexcept { + if (known_bitreverse) { + flat_encode_neon(bits, src, wire_code, count); + } else { + flat_encode_neon(bits, src, wire_code, count); + } +} + +/// @brief Runtime depth dispatch for NEON flat encoding. +inline void flat_encode_dispatch_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse) noexcept { + switch (depth) { + case 2: + flat_encode_maybe_bitreverse_neon<2>(bits, src, wire_code, count, + known_bitreverse); + return; + case 3: + flat_encode_maybe_bitreverse_neon<3>(bits, src, wire_code, count, + known_bitreverse); + return; + case 4: + flat_encode_maybe_bitreverse_neon<4>(bits, src, wire_code, count, + known_bitreverse); + return; + case 5: + flat_encode_maybe_bitreverse_neon<5>(bits, src, wire_code, count, + known_bitreverse); + return; + case 6: + flat_encode_maybe_bitreverse_neon<6>(bits, src, wire_code, count, + known_bitreverse); + return; + case 7: + flat_encode_maybe_bitreverse_neon<7>(bits, src, wire_code, count, + known_bitreverse); + return; + case 8: + flat_encode_maybe_bitreverse_neon<8>(bits, src, wire_code, count, + known_bitreverse); + return; + default: + return; + } +} + +#endif // PIXIE_PIVCO_NEON_SUPPORT + +// =========================================================================== +// AVX2 backend (SSE4.1 byte-shuffle lookup tables; AVX2 implies SSE4.1). +// =========================================================================== +#if defined(__AVX2__) && defined(__SSE4_1__) + +/// @brief 256-entry × 16-byte shuffle-control table, 16-byte aligned per entry. +/// @details Stored as plain bytes (not `__m128i`) to avoid dropping the +/// `__may_alias__` attribute that `std::array<__m128i>` would warn +/// about; entries are loaded as aligned vectors at use time. +struct ShufTable { + alignas(16) std::uint8_t entry[256][16]; +}; + +/// @brief Load shuffle entry @p m as an aligned 128-bit vector. +inline __m128i load_shuf(const ShufTable& t, std::uint8_t m) noexcept { + return _mm_load_si128( + reinterpret_cast(&t.entry[m])); // NOLINT +} + +/// @brief 256-entry shuffle table for the 8-byte merge. +/// @details For an 8-bit mask `m` (bit i = 1 -> right), entry `m` is a +/// 16-byte shuffle control that, applied to a 128-bit vector holding +/// `left[0..7]` in its low qword and `right[0..7]` in its high qword, +/// produces the merged 8-byte output. Mirrors the paper's NEON +/// `expand_tab`. +inline const ShufTable& merge_lut() { + static const ShufTable lut = [] { + ShufTable t{}; + for (int m = 0; m < 256; ++m) { + int nl = 0; + int nr = 0; + for (int i = 0; i < 8; ++i) { + if ((m >> i) & 1) { + t.entry[m][i] = static_cast(8 + nr); + ++nr; + } else { + t.entry[m][i] = static_cast(nl); + ++nl; + } + } + for (int i = 8; i < 16; ++i) { + t.entry[m][i] = 0x80; // zeroed lanes (unused, stored as full 16 bytes) + } + } + return t; + }(); + return lut; +} + +/// @brief 256-entry shuffle tables for the 8-byte compress (partition). +/// @details `compress_right[m]` gathers the elements at the 1-bit positions of +/// mask `m` (in order) to the front of the 128-bit result; +/// `compress_left[m]` does the same for the 0-bit positions. Trailing +/// lanes are zeroed. Mirrors the paper's NEON `compress_tab`. +inline const ShufTable& compress_right_lut() { + static const ShufTable lut = [] { + ShufTable t{}; + for (int m = 0; m < 256; ++m) { + int p = 0; + for (int i = 0; i < 8; ++i) { + if ((m >> i) & 1) { + t.entry[m][p++] = static_cast(i); + } + } + for (int i = p; i < 16; ++i) { + t.entry[m][i] = 0x80; + } + } + return t; + }(); + return lut; +} + +inline const ShufTable& compress_left_lut() { + static const ShufTable lut = [] { + ShufTable t{}; + for (int m = 0; m < 256; ++m) { + int p = 0; + for (int i = 0; i < 8; ++i) { + if (((m >> i) & 1) == 0) { + t.entry[m][p++] = static_cast(i); + } + } + for (int i = p; i < 16; ++i) { + t.entry[m][i] = 0x80; + } + } + return t; + }(); + return lut; +} + +/// @brief Merge two constant leaves 32 symbols at a time with AVX2. +/// @details Four routing bytes are widened to four qwords, each byte is +/// broadcast over its eight output lanes, and a fixed bit-position +/// vector expands the bitmap without a shuffle-table lookup. +inline void merge_decode_cst_cst_avx2(std::uint8_t* dst, + std::uint8_t left_symbol, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + const auto* bit_bytes = reinterpret_cast(bits); + const __m256i broadcast_control = + _mm256_setr_epi8(0, 0, 0, 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8, 0, 0, 0, + 0, 0, 0, 0, 0, 8, 8, 8, 8, 8, 8, 8, 8); + const __m256i bit_positions = _mm256_setr_epi8( + 1, 2, 4, 8, 16, 32, 64, static_cast(0x80), 1, 2, 4, 8, 16, 32, 64, + static_cast(0x80), 1, 2, 4, 8, 16, 32, 64, static_cast(0x80), + 1, 2, 4, 8, 16, 32, 64, static_cast(0x80)); + const __m256i zero = _mm256_setzero_si256(); + const __m256i left = _mm256_set1_epi8(static_cast(left_symbol)); + const __m256i delta = + _mm256_set1_epi8(static_cast(left_symbol ^ right_symbol)); + + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + std::uint32_t routing = 0; + std::memcpy(&routing, bit_bytes + i / 8, sizeof(routing)); + const __m128i routing_bytes = _mm_cvtsi32_si128(static_cast(routing)); + const __m256i routing_qwords = _mm256_cvtepu8_epi64(routing_bytes); + const __m256i broadcast = + _mm256_shuffle_epi8(routing_qwords, broadcast_control); + const __m256i clear_bits = + _mm256_cmpeq_epi8(_mm256_and_si256(broadcast, bit_positions), zero); + const __m256i symbols = + _mm256_xor_si256(left, _mm256_andnot_si256(clear_bits, delta)); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(dst + i), symbols); + } + for (; i < count; ++i) { + dst[i] = + ((bits[i / 64] >> (i % 64)) & 1u) != 0 ? right_symbol : left_symbol; + } +} + +template +inline void merge_decode_avx2(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + if constexpr (kLeftConstant && kRightConstant) { + merge_decode_cst_cst_avx2(dst, left_symbol, right_symbol, bits, count); + return; + } + + const auto& mlut = merge_lut(); + const auto* bit_bytes = reinterpret_cast(bits); + std::size_t out_i = 0; + std::size_t c_left = 0; + std::size_t c_right = 0; + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + // AVX2 is little-endian on every supported x86 target. Eight-symbol + // groups are byte-aligned in the LSB-first bitmap, so loading the routing + // byte directly avoids a word-index/shift sequence in every merge group. + const std::uint8_t m8 = bit_bytes[i / 8]; + __m128i left_values; + if constexpr (kLeftConstant) { + left_values = _mm_set1_epi8(static_cast(left_symbol)); + } else { + left_values = + _mm_loadl_epi64(reinterpret_cast(left + c_left)); + } + __m128i right_values; + if constexpr (kRightConstant) { + right_values = _mm_set1_epi8(static_cast(right_symbol)); + } else { + right_values = + _mm_loadl_epi64(reinterpret_cast(right + c_right)); + } + const __m128i combined = _mm_unpacklo_epi64(left_values, right_values); + const __m128i out = _mm_shuffle_epi8(combined, load_shuf(mlut, m8)); + _mm_storel_epi64(reinterpret_cast<__m128i*>(dst + out_i), out); + out_i += 8; + const int nr = std::popcount(static_cast(m8)); + c_right += static_cast(nr); + c_left += static_cast(8 - nr); + } + merge_scalar_tail( + dst, out_i, left, left_symbol, c_left, right, right_symbol, c_right, bits, + i, count); +} + +/// @brief Store @p v (low @p n bytes valid) at @p dst without writing past +/// @p dst + @p n. +/// @details The codec partitions into a ping-pong workspace where a node's left +/// and right child regions are contiguous, and a node's right region +/// is immediately followed by its sibling's region (already written by +/// the parent, consumed later in pre-order). A full-width vector store +/// would clobber the start of the adjacent region, so when the valid +/// byte count is less than the vector width the store is masked to +/// exactly +/// @p n bytes. +inline void store_masked_epi64(std::uint8_t* dst, __m128i v, int n) noexcept { + if (n >= 8) { + _mm_storel_epi64(reinterpret_cast<__m128i*>(dst), v); + } else if (n > 0) { + alignas(8) std::uint8_t tmp[8]; + _mm_storel_epi64(reinterpret_cast<__m128i*>(tmp), v); + std::memcpy(dst, tmp, static_cast(n)); + } +} + +template +inline void partition_encode_avx2(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + std::uint8_t* bits_bytes = reinterpret_cast(bits); + + const std::size_t full = weight / 8; + std::size_t c_left = 0; + std::size_t c_right = 0; + std::size_t g = 0; + for (; g < full; ++g) { + const std::uint8_t* s8 = src + g * 8; + const std::uint8_t m8 = build_dir_mask8(s8, dir); + // Byte-aligned write: full groups always start on a byte boundary in the + // LSB-first bitmap, so byte g holds exactly symbols [8g, 8g+8). + bits_bytes[g] = m8; + if constexpr (kWriteLeft || kWriteRight) { + const int nr = std::popcount(static_cast(m8)); + const int nl = 8 - nr; + const __m128i data = + _mm_loadl_epi64(reinterpret_cast(s8)); + if constexpr (kWriteRight) { + const __m128i right_values = + _mm_shuffle_epi8(data, load_shuf(compress_right_lut(), m8)); + store_masked_epi64(right_dst + c_right, right_values, + static_cast(right_weight - c_right)); + } + if constexpr (kWriteLeft) { + const __m128i left_values = + _mm_shuffle_epi8(data, load_shuf(compress_left_lut(), m8)); + store_masked_epi64(left_dst + c_left, left_values, + static_cast(left_weight - c_left)); + } + c_right += static_cast(nr); + c_left += static_cast(nl); + } + } + + // Scalar tail for the trailing < 8 symbols (byte-unaligned bitmap bits). + if constexpr (kWriteLeft || kWriteRight) { + partition_scalar_tail( + left_dst, right_dst, bits, src, dir, c_left, c_right, full * 8, weight); + } else { + partition_bitmap_tail(bits, src, dir, full * 8, weight); + } +} + +/// @brief AVX2 partition when one child is a constant leaf. +/// @tparam kLeftConstant True when the left child is the constant; the dense +/// right stream is written. False writes the dense left +/// stream and treats the right child as constant. +template +inline void partition_encode_one_constant_avx2( + std::uint8_t* output, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t constant_symbol, + std::size_t weight, + std::size_t output_weight) noexcept { + auto* bit_bytes = reinterpret_cast(bits); + const __m128i constant = _mm_set1_epi8(static_cast(constant_symbol)); + std::size_t output_count = 0; + const std::size_t full = weight / 8; + std::size_t group = 0; + for (; group < full; ++group) { + const std::uint8_t* s8 = src + group * 8; + const __m128i data = _mm_loadl_epi64(reinterpret_cast(s8)); + const std::uint8_t equal_mask = static_cast( + _mm_movemask_epi8(_mm_cmpeq_epi8(data, constant))); + const std::uint8_t routing = + kLeftConstant ? static_cast(~equal_mask) : equal_mask; + bit_bytes[group] = routing; + + const auto& lut = + kLeftConstant ? compress_right_lut() : compress_left_lut(); + const __m128i packed = _mm_shuffle_epi8(data, load_shuf(lut, routing)); + store_masked_epi64(output + output_count, packed, + static_cast(output_weight - output_count)); + const std::size_t right_count = + std::popcount(static_cast(routing)); + output_count += kLeftConstant ? right_count : 8 - right_count; + } + + std::size_t word_index = (group * 8) / 64; + std::size_t bit_position = (group * 8) % 64; + for (std::size_t i = group * 8; i < weight; ++i) { + const bool is_constant = src[i] == constant_symbol; + const bool goes_right = kLeftConstant ? !is_constant : is_constant; + if (goes_right) { + bits[word_index] |= std::uint64_t{1} << bit_position; + } + if (kLeftConstant ? goes_right : !goes_right) { + output[output_count++] = src[i]; + } + if (++bit_position == 64) { + bit_position = 0; + ++word_index; + } + } +} + +/// @brief Decode a 2-bit flat subtree, 32 symbols per iteration. +inline void flat_decode_d2_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + alignas(16) std::uint8_t table_lanes[16]{}; + std::memcpy(table_lanes, table, 4); + const __m128i lookup = + _mm_load_si128(reinterpret_cast(table_lanes)); + const __m128i mask = _mm_set1_epi8(3); + + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + const __m128i data = + _mm_loadl_epi64(reinterpret_cast(packed + i / 4)); + const __m128i code0 = _mm_and_si128(data, mask); + const __m128i code1 = _mm_and_si128(_mm_srli_epi16(data, 2), mask); + const __m128i code2 = _mm_and_si128(_mm_srli_epi16(data, 4), mask); + const __m128i code3 = _mm_and_si128(_mm_srli_epi16(data, 6), mask); + const __m128i code01 = _mm_unpacklo_epi8(code0, code1); + const __m128i code23 = _mm_unpacklo_epi8(code2, code3); + const __m128i indices0 = _mm_unpacklo_epi16(code01, code23); + const __m128i indices1 = _mm_unpackhi_epi16(code01, code23); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), + _mm_shuffle_epi8(lookup, indices0)); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), + _mm_shuffle_epi8(lookup, indices1)); + } + for (; i < count; ++i) { + dst[i] = table[(packed[i / 4] >> (2 * (i % 4))) & 3u]; + } +} + +/// @brief Decode a 4-bit flat subtree, 32 symbols per iteration. +inline void flat_decode_d4_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + const __m128i lookup = + _mm_loadu_si128(reinterpret_cast(table)); + const __m128i mask = _mm_set1_epi8(0x0f); + + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + const __m128i data = + _mm_loadu_si128(reinterpret_cast(packed + i / 2)); + const __m128i low = _mm_and_si128(data, mask); + const __m128i high = _mm_and_si128(_mm_srli_epi16(data, 4), mask); + const __m128i low_symbols = _mm_shuffle_epi8(lookup, low); + const __m128i high_symbols = _mm_shuffle_epi8(lookup, high); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), + _mm_unpacklo_epi8(low_symbols, high_symbols)); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i + 16), + _mm_unpackhi_epi8(low_symbols, high_symbols)); + } + for (; i < count; ++i) { + dst[i] = table[(packed[i / 2] >> (4 * (i % 2))) & 0x0fu]; + } +} + +/// @brief Reverse the low @p width bits of one byte. +inline std::uint8_t reverse_low_bits_avx2(std::uint8_t value, + std::uint8_t width) noexcept { + value = static_cast(((value & 0x55u) << 1) | + ((value & 0xaau) >> 1)); + value = static_cast(((value & 0x33u) << 2) | + ((value & 0xccu) >> 2)); + value = static_cast((value << 4) | (value >> 4)); + return value >> (8 - width); +} + +/// @brief Decode a known full-alphabet depth-8 bit-reversal table. +/// @details A balanced tree built from symbols 0..255 stores +/// `table[reverse8(code)] = code`. AVX2 can apply that permutation +/// with two nibble `pshufb` operations, avoiding the arbitrary byte +/// lookup that AVX2 cannot otherwise vectorize efficiently. +inline void flat_decode_d8_bitreverse_unchecked_avx2( + std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + alignas(16) static constexpr std::uint8_t kReverseNibble[16] = { + 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15}; + const __m128i reverse128 = + _mm_load_si128(reinterpret_cast(kReverseNibble)); + const __m256i reverse = _mm256_broadcastsi128_si256(reverse128); + const __m256i nibble_mask = _mm256_set1_epi8(0x0f); + std::size_t i = 0; + for (; i + 32 <= count; i += 32) { + const __m256i codes = + _mm256_loadu_si256(reinterpret_cast(packed + i)); + const __m256i low = _mm256_and_si256(codes, nibble_mask); + const __m256i high = + _mm256_and_si256(_mm256_srli_epi16(codes, 4), nibble_mask); + const __m256i reversed_low = _mm256_shuffle_epi8(reverse, low); + const __m256i reversed_high = _mm256_shuffle_epi8(reverse, high); + const __m256i symbols = + _mm256_or_si256(_mm256_slli_epi16(reversed_low, 4), reversed_high); + _mm256_storeu_si256(reinterpret_cast<__m256i*>(dst + i), symbols); + } + for (; i < count; ++i) { + dst[i] = table[packed[i]]; + } +} + +/// @brief Decode the common full-alphabet depth-8 bit-reversal table. +/// @returns True when @p table has the reshaped tree's bit-reversal layout and +/// the output was produced; false for an arbitrary 256-byte table. +inline bool flat_decode_d8_bitreverse_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + for (std::size_t code = 0; code < 256; ++code) { + if (table[code] != + reverse_low_bits_avx2(static_cast(code), 8)) { + return false; + } + } + flat_decode_d8_bitreverse_unchecked_avx2(dst, packed, table, count); + return true; +} + +/// @brief Expand eight tightly packed D-bit codes into eight byte lanes. +/// @details Eight codes occupy exactly D bytes, so every eight-code group is +/// byte-aligned. BMI2 `pdep` inserts the padding bits between codes in +/// one instruction. The constexpr fallback is fully unrolled for +/// AVX2 CPUs that do not also expose BMI2. +template +inline std::uint64_t expand_codes8_avx2(const std::uint8_t* packed) noexcept { + static_assert(kDepth >= 3 && kDepth <= 8); + std::uint64_t dense = 0; + std::memcpy(&dense, packed, kDepth); + if constexpr (kDepth == 8) { + return dense; + } + + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; +#if defined(__BMI2__) + constexpr std::uint64_t kByteLaneMask = + kCodeMask * UINT64_C(0x0101010101010101); + return _pdep_u64(dense, kByteLaneMask); +#else + std::uint64_t expanded = 0; + for (std::size_t lane = 0; lane < 8; ++lane) { + expanded |= ((dense >> (lane * kDepth)) & kCodeMask) << (lane * 8); + } + return expanded; +#endif +} + +/// @brief Reverse the low D bits independently in 16 byte lanes. +template +inline __m128i reverse_low_bits16_avx2(__m128i values) noexcept { + static_assert(kDepth >= 2 && kDepth <= 7); + alignas(16) static constexpr std::uint8_t kReverseNibble[16] = { + 0, 8, 4, 12, 2, 10, 6, 14, 1, 9, 5, 13, 3, 11, 7, 15}; + const __m128i reverse = + _mm_load_si128(reinterpret_cast(kReverseNibble)); + const __m128i nibble_mask = _mm_set1_epi8(0x0f); + const __m128i low = _mm_and_si128(values, nibble_mask); + const __m128i high = _mm_and_si128(_mm_srli_epi16(values, 4), nibble_mask); + const __m128i reversed = + _mm_or_si128(_mm_slli_epi16(_mm_shuffle_epi8(reverse, low), 4), + _mm_shuffle_epi8(reverse, high)); + constexpr std::uint8_t kCodeMask = (std::uint8_t{1} << kDepth) - 1; + return _mm_and_si128(_mm_srli_epi16(reversed, 8 - kDepth), + _mm_set1_epi8(static_cast(kCodeMask))); +} + +/// @brief Decode a known balanced low-alphabet table as bit reversal. +template +inline void flat_decode_bitreverse_unchecked_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 6 && kDepth <= 7); + std::size_t i = 0; + for (; i + 16 <= count; i += 16) { + const std::size_t byte_offset = (i * kDepth) / 8; + const std::uint64_t codes0 = + expand_codes8_avx2(packed + byte_offset); + const std::uint64_t codes1 = + expand_codes8_avx2(packed + byte_offset + kDepth); + const __m128i codes = _mm_set_epi64x(static_cast(codes1), + static_cast(codes0)); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), + reverse_low_bits16_avx2(codes)); + } + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + const auto* words = reinterpret_cast(packed); + for (; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t code = words[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + code |= words[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[code & kCodeMask]; + } +} + +/// @brief Decode a balanced low-alphabet flat table as bit reversal. +/// @returns True when the table has the expected permutation and was decoded. +template +inline bool flat_decode_bitreverse_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 6 && kDepth <= 7); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + for (std::size_t code = 0; code < kTableSize; ++code) { + if (table[code] != + reverse_low_bits_avx2(static_cast(code), kDepth)) { + return false; + } + } + flat_decode_bitreverse_unchecked_avx2(dst, packed, table, count); + return true; +} + +/// @brief Lookup byte indices in a power-of-two collection of 16-byte banks. +/// @details AVX2 has no arbitrary byte permutation across 256 entries. Each +/// leaf performs one `pshufb` using the low nibble; balanced +/// `blendv` levels select the bank with the remaining high bits. +/// The recursion is compile-time and emits straight-line SIMD code. +template +inline __m128i lookup_banked_bytes_avx2(const std::uint8_t* table, + __m128i low_nibbles, + __m128i codes) noexcept { + static_assert(kBankCount > 0 && std::has_single_bit(kBankCount)); + if constexpr (kBankCount == 1) { + const __m128i bank = _mm_loadu_si128( + reinterpret_cast(table + kFirstBank * 16)); + return _mm_shuffle_epi8(bank, low_nibbles); + } else { + constexpr std::size_t kHalf = kBankCount / 2; + constexpr int kSelectBit = 3 + std::countr_zero(kBankCount); + const __m128i low = + lookup_banked_bytes_avx2(table, low_nibbles, codes); + const __m128i high = lookup_banked_bytes_avx2( + table, low_nibbles, codes); + const __m128i select = _mm_slli_epi16(codes, 7 - kSelectBit); + return _mm_blendv_epi8(low, high, select); + } +} + +/// @brief Decode flat depths 3 and 5..7 in 16-symbol AVX2 batches. +/// @details The existing depth-2/depth-4 kernels unpack 32 symbols per batch +/// and remain faster for those widths. Other depths use BMI2 to +/// expand two groups of eight packed codes, followed by a banked +/// `pshufb` lookup. This avoids the former per-symbol shift, branch, +/// and dependent table load loop. +template +inline void flat_decode_banked_avx2(std::uint8_t* dst, + const std::uint8_t* packed, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth == 3 || (kDepth >= 5 && kDepth <= 7)); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + constexpr std::size_t kBankCount = kTableSize <= 16 ? 1 : kTableSize / 16; + const __m128i nibble_mask = _mm_set1_epi8(0x0f); + + // A depth-3 table has only eight valid bytes, while the banked lookup loads + // a complete vector. Pad that one small table locally; depths >=5 already + // contain one or more complete 16-byte banks. + alignas(16) std::uint8_t padded_depth3[16]{}; + const std::uint8_t* lookup_table = table; + if constexpr (kDepth == 3) { + std::memcpy(padded_depth3, table, kTableSize); + lookup_table = padded_depth3; + } + + std::size_t i = 0; + for (; i + 16 <= count; i += 16) { + const std::size_t byte_offset = (i * kDepth) / 8; + const std::uint64_t codes0 = + expand_codes8_avx2(packed + byte_offset); + const std::uint64_t codes1 = + expand_codes8_avx2(packed + byte_offset + kDepth); + const __m128i codes = _mm_set_epi64x(static_cast(codes1), + static_cast(codes0)); + const __m128i low_nibbles = _mm_and_si128(codes, nibble_mask); + const __m128i symbols = lookup_banked_bytes_avx2<0, kBankCount>( + lookup_table, low_nibbles, codes); + _mm_storeu_si128(reinterpret_cast<__m128i*>(dst + i), symbols); + } + + const auto* words = reinterpret_cast(packed); + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + for (; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t code = words[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + code |= words[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[code & kCodeMask]; + } +} + +#if defined(__BMI2__) +/// @brief Encode one depth-3 flat subtree with AVX2 gather and BMI2 packing. +/// @details AVX2 has no byte gather, so eight input symbols are widened to +/// dword indices and looked up together. The gathered codes are +/// narrowed back to bytes, then `pext` packs their low three bits. +/// Wider flat depths were benchmarked and deliberately retain their +/// faster scalar encoders. The caller provides three readable pad +/// bytes after the 256-entry code table because dword gathers overlap +/// adjacent byte entries. +inline void flat_encode_d3_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + constexpr std::uint8_t kDepth = 3; + constexpr std::uint64_t kByteLaneMask = UINT64_C(0x0707070707070707); + auto* packed = reinterpret_cast(bits); + const __m256i byte_mask = _mm256_set1_epi32(0xff); + const __m128i zero = _mm_setzero_si128(); + + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const __m128i symbols = + _mm_loadl_epi64(reinterpret_cast(src + i)); + const __m256i indices = _mm256_cvtepu8_epi32(symbols); + const __m256i gathered = _mm256_and_si256( + _mm256_i32gather_epi32(reinterpret_cast(wire_code), indices, + 1), + byte_mask); + const __m128i gathered_low = _mm256_castsi256_si128(gathered); + const __m128i gathered_high = _mm256_extracti128_si256(gathered, 1); + const __m128i codes16 = _mm_packus_epi32(gathered_low, gathered_high); + const __m128i codes8 = _mm_packus_epi16(codes16, zero); + const std::uint64_t codes = + static_cast(_mm_cvtsi128_si64(codes8)); + const std::uint64_t dense = _pext_u64(codes, kByteLaneMask); + std::memcpy(packed + (i * kDepth) / 8, &dense, kDepth); + } + + std::size_t word_index = (i * kDepth) / 64; + std::size_t bit_offset = (i * kDepth) % 64; + for (; i < count; ++i) { + const std::uint64_t code = wire_code[src[i]]; + bits[word_index] |= code << bit_offset; + if (bit_offset + kDepth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + bit_offset += kDepth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +/// @brief Encode a known balanced low-alphabet table as bit reversal. +template +inline void flat_encode_bitreverse_unchecked_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 7); + constexpr std::uint64_t kCodeMask = (std::uint64_t{1} << kDepth) - 1; + constexpr std::uint64_t kByteLaneMask = + kCodeMask * UINT64_C(0x0101010101010101); + auto* packed = reinterpret_cast(bits); + std::size_t i = 0; + for (; i + 8 <= count; i += 8) { + const __m128i symbols = + _mm_loadl_epi64(reinterpret_cast(src + i)); + const __m128i codes = reverse_low_bits16_avx2(symbols); + const std::uint64_t byte_codes = + static_cast(_mm_cvtsi128_si64(codes)); + const std::uint64_t dense = _pext_u64(byte_codes, kByteLaneMask); + std::memcpy(packed + (i * kDepth) / 8, &dense, kDepth); + } + + std::size_t word_index = (i * kDepth) / 64; + std::size_t bit_offset = (i * kDepth) % 64; + for (; i < count; ++i) { + const std::uint64_t code = wire_code[src[i]]; + bits[word_index] |= code << bit_offset; + if (bit_offset + kDepth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + bit_offset += kDepth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +/// @brief Encode a balanced low-alphabet table as vectorized bit reversal. +/// @returns True when @p wire_code has the expected mapping and was encoded. +template +inline bool flat_encode_bitreverse_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 7); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + for (std::size_t symbol = 0; symbol < kTableSize; ++symbol) { + if (wire_code[symbol] != + reverse_low_bits_avx2(static_cast(symbol), kDepth)) { + return false; + } + } + flat_encode_bitreverse_unchecked_avx2(bits, src, wire_code, count); + return true; +} + +/// @brief Select checked or preclassified balanced-table encoding. +template +inline bool flat_encode_bitreverse_maybe_avx2(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + bool known_bitreverse) noexcept { + if (known_bitreverse) { + flat_encode_bitreverse_unchecked_avx2(bits, src, wire_code, count); + return true; + } + return flat_encode_bitreverse_avx2(bits, src, wire_code, count); +} +#endif + +#endif // AVX2 + SSE4.1 + +// =========================================================================== +// AVX-512 VBMI2 backend (byte-granular expand/compress, 64 bytes/iteration). +// =========================================================================== +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + +#if defined(__AVX512VBMI__) + +/// @brief Four-bank byte lookup used for 64-, 128-, and 256-entry tables. +/// @details `vpermb` addresses one 64-byte bank. The high two index bits select +/// the bank after all required bank permutations run in parallel. +template +struct ByteLookupAvx512 { + __m512i bank0; + __m512i bank1; + __m512i bank2; + __m512i bank3; +}; + +/// @brief Load a `2^kDepth` byte table without reading beyond its allocation. +template +inline ByteLookupAvx512 load_byte_lookup_avx512( + const std::uint8_t* table) noexcept { + static_assert(kDepth >= 1 && kDepth <= 8); + constexpr std::size_t kTableSize = std::size_t{1} << kDepth; + constexpr __mmask64 kFirstBankMask = + kTableSize >= 64 ? ~__mmask64{0} + : (__mmask64{1} << kTableSize) - __mmask64{1}; + const __m512i zero = _mm512_setzero_si512(); + ByteLookupAvx512 lookup{ + _mm512_maskz_loadu_epi8(kFirstBankMask, table), zero, zero, zero}; + if constexpr (kDepth >= 7) { + lookup.bank1 = _mm512_loadu_si512(table + 64); + } + if constexpr (kDepth == 8) { + lookup.bank2 = _mm512_loadu_si512(table + 128); + lookup.bank3 = _mm512_loadu_si512(table + 192); + } + return lookup; +} + +/// @brief Translate 64 byte indices through a preloaded `2^kDepth` table. +template +inline __m512i lookup_bytes_avx512( + __m512i indices, + const ByteLookupAvx512& lookup) noexcept { + const __m512i low_indices = _mm512_and_si512(indices, _mm512_set1_epi8(0x3f)); + __m512i result = _mm512_permutexvar_epi8(low_indices, lookup.bank0); + if constexpr (kDepth >= 7) { + const __m512i banks = + _mm512_and_si512(_mm512_srli_epi16(indices, 6), _mm512_set1_epi8(3)); + const __mmask64 bank1 = _mm512_cmpeq_epi8_mask(banks, _mm512_set1_epi8(1)); + result = _mm512_mask_mov_epi8( + result, bank1, _mm512_permutexvar_epi8(low_indices, lookup.bank1)); + if constexpr (kDepth == 8) { + const __mmask64 bank2 = + _mm512_cmpeq_epi8_mask(banks, _mm512_set1_epi8(2)); + const __mmask64 bank3 = + _mm512_cmpeq_epi8_mask(banks, _mm512_set1_epi8(3)); + result = _mm512_mask_mov_epi8( + result, bank2, _mm512_permutexvar_epi8(low_indices, lookup.bank2)); + result = _mm512_mask_mov_epi8( + result, bank3, _mm512_permutexvar_epi8(low_indices, lookup.bank3)); + } + } + return result; +} + +/// @brief Decode one scalar tail after full 64-symbol flat-D groups. +template +inline void flat_decode_tail_avx512(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t start, + std::size_t count) noexcept { + constexpr std::uint64_t kMask = (std::uint64_t{1} << kDepth) - 1; + for (std::size_t i = start; i < count; ++i) { + const std::size_t bit_position = i * kDepth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + std::uint64_t packed = bits[word_index] >> bit_offset; + if (bit_offset + kDepth > 64) { + packed |= bits[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[packed & kMask]; + } +} + +/// @brief AVX-512 implementation of the paper's `merge_flat_D` primitive. +/// @details Each 64-symbol group occupies exactly @p kDepth 64-bit words. The +/// words are permuted into eight 8-code windows, aligned with vector +/// shifts, unpacked by `vpmultishiftqb`, and translated by `vpermb`. +template +inline void flat_decode_avx512(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count) noexcept { + static_assert(kDepth >= 2 && kDepth <= 8); + const ByteLookupAvx512 lookup = + load_byte_lookup_avx512(table); + + alignas(64) static constexpr auto kWordIndices = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = (lane * 8 * kDepth) / 64; + } + return values; + }(); + alignas(64) static constexpr auto kNextWordIndices = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = ((lane * 8 * kDepth) / 64 + 1) % 8; + } + return values; + }(); + alignas(64) static constexpr auto kRightShifts = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = (lane * 8 * kDepth) % 64; + } + return values; + }(); + alignas(64) static constexpr auto kLeftShifts = [] { + std::array values{}; + for (std::size_t lane = 0; lane < values.size(); ++lane) { + values[lane] = 64 - ((lane * 8 * kDepth) % 64); + } + return values; + }(); + alignas(64) static constexpr auto kMultishiftControl = [] { + std::array values{}; + for (std::size_t lane = 0; lane < 8; ++lane) { + for (std::size_t code = 0; code < 8; ++code) { + values[lane * 8 + code] = static_cast(code * kDepth); + } + } + return values; + }(); + + const __m512i word_indices = _mm512_load_si512(kWordIndices.data()); + const __m512i next_word_indices = _mm512_load_si512(kNextWordIndices.data()); + const __m512i right_shifts = _mm512_load_si512(kRightShifts.data()); + const __m512i left_shifts = _mm512_load_si512(kLeftShifts.data()); + const __m512i multishift_control = + _mm512_load_si512(kMultishiftControl.data()); + const __m512i code_mask = + _mm512_set1_epi8(static_cast((1u << kDepth) - 1u)); + constexpr __mmask8 kPackedWordMask = + static_cast<__mmask8>((1u << kDepth) - 1u); + + std::size_t i = 0; + for (; i + 64 <= count; i += 64) { + __m512i indices; + if constexpr (kDepth == 8) { + // Eight-bit codes are already byte indices; only lookup is required. + indices = + _mm512_loadu_si512(reinterpret_cast(bits) + i); + } else { + const __m512i packed = + _mm512_maskz_loadu_epi64(kPackedWordMask, bits + (i / 64) * kDepth); + const __m512i current_words = + _mm512_permutexvar_epi64(word_indices, packed); + const __m512i next_words = + _mm512_permutexvar_epi64(next_word_indices, packed); + const __m512i windows = + _mm512_or_si512(_mm512_srlv_epi64(current_words, right_shifts), + _mm512_sllv_epi64(next_words, left_shifts)); + indices = _mm512_and_si512( + _mm512_multishift_epi64_epi8(multishift_control, windows), code_mask); + } + _mm512_storeu_si512(dst + i, lookup_bytes_avx512(indices, lookup)); + } + flat_decode_tail_avx512(dst, bits, table, i, count); +} + +#endif // AVX-512 VBMI + +template +inline void merge_decode_avx512(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + std::size_t out_i = 0; + std::size_t c_left = 0; + std::size_t c_right = 0; + const std::size_t full = count / 64; + for (std::size_t w = 0; w < full; ++w) { + const __mmask64 km = bits[w]; + __m512i left_values; + if constexpr (kLeftConstant) { + left_values = _mm512_set1_epi8(static_cast(left_symbol)); + } else { + left_values = _mm512_loadu_si512(left + c_left); + } + __m512i right_values; + if constexpr (kRightConstant) { + right_values = _mm512_set1_epi8(static_cast(right_symbol)); + } else { + right_values = _mm512_loadu_si512(right + c_right); + } + // Fill clear-bit positions with left symbols, then set-bit positions with + // right symbols. maskz zeroes the right positions first; the second expand + // keeps the left values there. + __m512i d; + if constexpr (kLeftConstant) { + d = left_values; + } else { + d = _mm512_maskz_expand_epi8(~km, left_values); + } + if constexpr (kRightConstant) { + d = _mm512_mask_mov_epi8(d, km, right_values); + } else { + d = _mm512_mask_expand_epi8(d, km, right_values); + } + _mm512_storeu_si512(dst + out_i, d); + const std::size_t nr = std::popcount(bits[w]); + out_i += 64; + c_right += nr; + c_left += 64 - nr; + } + // Tail handled by the AVX2 path if available, otherwise scalar. +#if defined(__AVX2__) && defined(__SSE4_1__) + const std::size_t done = full * 64; + const std::uint8_t* tail_left = kLeftConstant ? left : left + c_left; + const std::uint8_t* tail_right = kRightConstant ? right : right + c_right; + merge_decode_avx2( + dst + done, tail_left, left_symbol, tail_right, right_symbol, bits + full, + count - done); + (void)out_i; +#else + merge_scalar_tail( + dst, out_i, left, left_symbol, c_left, right, right_symbol, c_right, bits, + full * 64, count); +#endif +} + +template +inline void partition_encode_avx512(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + // Build the routing mask 64 symbols at a time, store the word, then compress + // the 64 input symbols into left/right streams with vpcompressb. + const std::size_t full = weight / 64; + std::size_t c_left = 0; + std::size_t c_right = 0; +#if defined(__AVX512VBMI__) + const ByteLookupAvx512<8> direction_lookup = load_byte_lookup_avx512<8>(dir); +#endif + std::size_t g = 0; + for (; g < full; ++g) { + const std::uint8_t* s64 = src + g * 64; + __m512i data; + __mmask64 km = 0; +#if defined(__AVX512VBMI__) + data = _mm512_loadu_si512(s64); + const __m512i directions = lookup_bytes_avx512(data, direction_lookup); + km = _mm512_cmpneq_epi8_mask(directions, _mm512_setzero_si512()); +#else + for (int k = 0; k < 64; ++k) { + km |= static_cast<__mmask64>(dir[s64[k]]) << k; + } +#endif + bits[g] = km; + if constexpr (kWriteLeft || kWriteRight) { +#if !defined(__AVX512VBMI__) + data = _mm512_loadu_si512(s64); +#endif + const std::size_t nr = std::popcount(km); + const std::size_t nl = 64 - nr; + if constexpr (kWriteRight) { + const __m512i right_packed = _mm512_maskz_compress_epi8(km, data); + const std::size_t remaining = right_weight - c_right; + if (remaining >= 64) { + _mm512_storeu_si512(right_dst + c_right, right_packed); + } else { + const __mmask64 store_mask = (1ull << remaining) - 1; + _mm512_mask_storeu_epi8(right_dst + c_right, store_mask, + right_packed); + } + } + if constexpr (kWriteLeft) { + const __m512i left_packed = _mm512_maskz_compress_epi8(~km, data); + const std::size_t remaining = left_weight - c_left; + if (remaining >= 64) { + _mm512_storeu_si512(left_dst + c_left, left_packed); + } else { + const __mmask64 store_mask = (1ull << remaining) - 1; + _mm512_mask_storeu_epi8(left_dst + c_left, store_mask, left_packed); + } + } + c_right += nr; + c_left += nl; + } + } + // Tail for the trailing < 64 symbols. `__AVX512VBMI2__` implies `__AVX2__` + // on every real target, so delegate to the AVX2 kernel (which itself defers + // the final < 8 symbols to the shared `partition_scalar_tail`). There is no + // reachable `#else` branch, so no scalar tail is duplicated here. + const std::size_t done_words = g; + const std::size_t done_syms = g * 64; + std::uint8_t* tail_left = kWriteLeft ? left_dst + c_left : left_dst; + std::uint8_t* tail_right = kWriteRight ? right_dst + c_right : right_dst; + partition_encode_avx2( + tail_left, tail_right, bits + done_words, src + done_syms, dir, + weight - done_syms, left_weight - c_left, right_weight - c_right); +} + +#endif // AVX-512 VBMI2 + +// =========================================================================== +// Scalar fallback (no SIMD). +// =========================================================================== + +template +inline void merge_decode_scalar(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + std::size_t out_i = 0; + std::size_t c_left = 0; + std::size_t c_right = 0; + const std::size_t words = (count + 63) / 64; + for (std::size_t w = 0; w < words; ++w) { + std::uint64_t word = bits[w]; + std::size_t limit = 64; + if (w + 1 == words) { + limit = count - w * 64; + } + for (std::size_t i = 0; i < limit; ++i) { + if (word & 1ull) { + if constexpr (kRightConstant) { + dst[out_i++] = right_symbol; + } else { + dst[out_i++] = right[c_right]; + } + ++c_right; + } else { + if constexpr (kLeftConstant) { + dst[out_i++] = left_symbol; + } else { + dst[out_i++] = left[c_left]; + } + ++c_left; + } + word >>= 1; + } + } +} + +template +inline void partition_encode_scalar(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight) noexcept { + std::size_t c_left = 0; + std::size_t c_right = 0; + partition_scalar_tail( + left_dst, right_dst, bits, src, dir, c_left, c_right, 0, weight); +} + +/// @brief Portable fixed-width unpack and lookup for flat subtrees. +inline void flat_decode_scalar(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count, + std::uint8_t depth) noexcept { + if (depth == 8) { + const auto* bytes = reinterpret_cast(bits); + for (std::size_t i = 0; i < count; ++i) { + dst[i] = table[bytes[i]]; + } + return; + } + + const std::uint64_t mask = (std::uint64_t{1} << depth) - 1; + std::size_t word_index = 0; + std::size_t bit_offset = 0; + for (std::size_t i = 0; i < count; ++i) { + std::uint64_t packed = bits[word_index] >> bit_offset; + if (bit_offset + depth > 64) { + packed |= bits[word_index + 1] << (64 - bit_offset); + } + dst[i] = table[packed & mask]; + bit_offset += depth; + if (bit_offset >= 64) { + bit_offset -= 64; + ++word_index; + } + } +} + +template +inline void merge_decode_dispatch(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t left_symbol, + const std::uint8_t* right, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + merge_decode_avx512( + dst, left, left_symbol, right, right_symbol, bits, count); +#elif defined(__AVX2__) && defined(__SSE4_1__) + merge_decode_avx2( + dst, left, left_symbol, right, right_symbol, bits, count); +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + merge_decode_neon( + dst, left, left_symbol, right, right_symbol, bits, count); +#else + merge_decode_scalar( + dst, left, left_symbol, right, right_symbol, bits, count); +#endif +} + +template +inline void partition_encode_dispatch(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + partition_encode_avx512( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#elif defined(__AVX2__) && defined(__SSE4_1__) + partition_encode_avx2( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + partition_encode_neon( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +#else + partition_encode_scalar(left_dst, right_dst, bits, + src, dir, weight); + (void)left_weight; + (void)right_weight; +#endif +} + +/// @brief Dispatch the paper's merge-flat-D primitive family. +inline void flat_decode_dispatch(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse) noexcept { +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + if (depth == 1) { + merge_decode_avx512(dst, nullptr, table[0], nullptr, table[1], + bits, count); + return; + } +#elif defined(__AVX2__) && defined(__SSE4_1__) + if (depth == 1) { + merge_decode_avx2(dst, nullptr, table[0], nullptr, table[1], + bits, count); + return; + } +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + if (depth == 1) { + merge_decode_neon(dst, nullptr, table[0], nullptr, table[1], + bits, count); + return; + } +#endif + +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) && defined(__AVX512VBMI__) + switch (depth) { + case 2: + flat_decode_avx512<2>(dst, bits, table, count); + return; + case 3: + flat_decode_avx512<3>(dst, bits, table, count); + return; + case 4: + flat_decode_avx512<4>(dst, bits, table, count); + return; + case 5: + flat_decode_avx512<5>(dst, bits, table, count); + return; + case 6: + flat_decode_avx512<6>(dst, bits, table, count); + return; + case 7: + flat_decode_avx512<7>(dst, bits, table, count); + return; + case 8: + flat_decode_avx512<8>(dst, bits, table, count); + return; + default: + break; + } +#endif + +#if defined(__AVX2__) && defined(__SSE4_1__) + const auto* packed = reinterpret_cast(bits); + switch (depth) { + case 2: + flat_decode_d2_avx2(dst, packed, table, count); + return; + case 3: + flat_decode_banked_avx2<3>(dst, packed, table, count); + return; + case 4: + flat_decode_d4_avx2(dst, packed, table, count); + return; + case 5: + flat_decode_banked_avx2<5>(dst, packed, table, count); + return; + case 6: + if (known_bitreverse) { + flat_decode_bitreverse_unchecked_avx2<6>(dst, packed, table, count); + return; + } + if (flat_decode_bitreverse_avx2<6>(dst, packed, table, count)) { + return; + } + flat_decode_banked_avx2<6>(dst, packed, table, count); + return; + case 7: + if (known_bitreverse) { + flat_decode_bitreverse_unchecked_avx2<7>(dst, packed, table, count); + return; + } + if (flat_decode_bitreverse_avx2<7>(dst, packed, table, count)) { + return; + } + flat_decode_banked_avx2<7>(dst, packed, table, count); + return; + case 8: + if (known_bitreverse) { + flat_decode_d8_bitreverse_unchecked_avx2(dst, packed, table, count); + return; + } + if (flat_decode_d8_bitreverse_avx2(dst, packed, table, count)) { + return; + } + break; + default: + break; + } +#endif + +#if defined(PIXIE_PIVCO_NEON_SUPPORT) + switch (depth) { + case 2: + flat_decode_neon<2>(dst, bits, table, count); + return; + case 3: + flat_decode_neon<3>(dst, bits, table, count); + return; + case 4: + flat_decode_neon<4>(dst, bits, table, count); + return; + case 5: + flat_decode_neon<5>(dst, bits, table, count); + return; + case 6: + flat_decode_neon<6>(dst, bits, table, count); + return; + case 7: + if (known_bitreverse) { + flat_decode_neon<7, true>(dst, bits, table, count); + } else { + flat_decode_neon<7, false>(dst, bits, table, count); + } + return; + case 8: + if (known_bitreverse) { + flat_decode_neon<8, true>(dst, bits, table, count); + } else { + flat_decode_neon<8, false>(dst, bits, table, count); + } + return; + default: + break; + } +#endif + (void)known_bitreverse; + flat_decode_scalar(dst, bits, table, count, depth); +} + +} // namespace pixie::pivco_simd_detail + +// --------------------------------------------------------------------------- +// Public dispatch entry points. +// --------------------------------------------------------------------------- + +namespace pixie { + +/// @brief Bottom-up decode merge: interleave @p left / @p right child streams +/// under @p bits into @p dst, producing @p count symbols. +/// @param dst Output buffer of (at least) @p count bytes. +/// @param left Left child output stream (consumed left_weight symbols). +/// @param right Right child output stream (consumed right_weight symbols). +/// @param bits Packed 1-bit-per-symbol routing bitmap, LSB-first per word. +/// Bit value 1 selects @p right, 0 selects @p left. +/// @param count Number of symbols (== bitmap bit count). +/// @details The caller must guarantee a few bytes of slack past the end of +/// @p dst, @p left, and @p right for vectorized over-reads/over-writes +/// (the codec's workspace is padded accordingly). +inline void pivco_merge_decode(std::uint8_t* dst, + const std::uint8_t* left, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch(dst, left, 0, right, 0, + bits, count); +} + +/// @brief Merge two constant-symbol leaves under a routing bitmap. +inline void pivco_merge_decode_cst_cst(std::uint8_t* dst, + std::uint8_t left_symbol, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch( + dst, nullptr, left_symbol, nullptr, right_symbol, bits, count); +} + +/// @brief Merge a constant left leaf with a dense right child stream. +inline void pivco_merge_decode_cst_vec(std::uint8_t* dst, + std::uint8_t left_symbol, + const std::uint8_t* right, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch( + dst, nullptr, left_symbol, right, 0, bits, count); +} + +/// @brief Merge a dense left child stream with a constant right leaf. +inline void pivco_merge_decode_vec_cst(std::uint8_t* dst, + const std::uint8_t* left, + std::uint8_t right_symbol, + const std::uint64_t* bits, + std::size_t count) noexcept { + pivco_simd_detail::merge_decode_dispatch( + dst, left, 0, nullptr, right_symbol, bits, count); +} + +/// @brief Top-down encode partition: split @p src into left/right child streams +/// and write the routing bitmap, using per-symbol direction bits. +/// @param left_dst Output for symbols whose direction bit is 0 (exactly +/// @p left_weight symbols). +/// @param right_dst Output for symbols whose direction bit is 1. Must equal +/// @p left_dst + @p left_weight: the two streams are stored +/// contiguously, and each is immediately followed by a +/// later-consumed sibling region, so both stores are masked +/// to never spill past their own region end. +/// @param bits Destination bitmap (pre-zeroed), LSB-first per word. +/// @param src Input symbols to partition. +/// @param dir Per-symbol direction bit table (0 or 1), indexed by +/// symbol. +/// @param weight Number of input symbols (== bitmap bit count). +/// @param left_weight Number of symbols routed left (== bitmap zero-count). +/// @param right_weight Number of symbols routed right (== bitmap one-count). +inline void pivco_partition_encode(std::uint8_t* left_dst, + std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + left_dst, right_dst, bits, src, dir, weight, left_weight, right_weight); +} + +/// @brief Build a routing bitmap and materialize only the left child stream. +inline void pivco_partition_encode_left(std::uint8_t* left_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t left_weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + left_dst, nullptr, bits, src, dir, weight, left_weight, 0); +} + +/// @brief Build a routing bitmap and materialize only the right child stream. +inline void pivco_partition_encode_right(std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + nullptr, right_dst, bits, src, dir, weight, 0, right_weight); +} + +/// @brief Build only the routing bitmap when both children are leaves. +inline void pivco_partition_encode_none(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* dir, + std::size_t weight) noexcept { + pivco_simd_detail::partition_encode_dispatch( + nullptr, nullptr, bits, src, dir, weight, 0, 0); +} + +/// @brief Build a bitmap for a node whose two children are constant leaves. +/// @details A set bit means @p src equals @p right_symbol. Unlike the generic +/// direction-table path, SIMD backends compare complete vectors and +/// pack the equality results directly into routing bytes. +inline void pivco_partition_encode_cst_cst(std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight) noexcept { + std::size_t i = 0; +#if defined(__AVX512VBMI2__) && defined(__AVX512BW__) + const __m512i right = _mm512_set1_epi8(static_cast(right_symbol)); + for (; i + 64 <= weight; i += 64) { + const __m512i symbols = _mm512_loadu_si512(src + i); + bits[i / 64] = _mm512_cmpeq_epi8_mask(symbols, right); + } +#elif defined(__AVX2__) + const __m256i right = _mm256_set1_epi8(static_cast(right_symbol)); + auto* bit_bytes = reinterpret_cast(bits); + for (; i + 32 <= weight; i += 32) { + const __m256i symbols = + _mm256_loadu_si256(reinterpret_cast(src + i)); + const std::uint32_t routing = static_cast( + _mm256_movemask_epi8(_mm256_cmpeq_epi8(symbols, right))); + std::memcpy(bit_bytes + i / 8, &routing, sizeof(routing)); + } +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) + pivco_simd_detail::partition_encode_cst_cst_neon(bits, src, right_symbol, + weight); + return; +#endif + for (; i < weight; ++i) { + if (src[i] == right_symbol) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + } +} + +#if defined(__AVX2__) && defined(__SSE4_1__) +/// @brief Partition a constant left leaf and materialize the right stream. +inline void pivco_partition_encode_cst_vec(std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t left_symbol, + std::size_t weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_avx2( + right_dst, bits, src, left_symbol, weight, right_weight); +} + +/// @brief Partition a dense left stream and a constant right leaf. +inline void pivco_partition_encode_vec_cst(std::uint8_t* left_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight, + std::size_t left_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_avx2( + left_dst, bits, src, right_symbol, weight, left_weight); +} +#elif defined(PIXIE_PIVCO_NEON_SUPPORT) +/// @brief Partition a constant left leaf and materialize the right stream. +inline void pivco_partition_encode_cst_vec(std::uint8_t* right_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t left_symbol, + std::size_t weight, + std::size_t right_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_neon( + right_dst, bits, src, left_symbol, weight, right_weight); +} + +/// @brief Partition a dense left stream and a constant right leaf. +inline void pivco_partition_encode_vec_cst(std::uint8_t* left_dst, + std::uint64_t* bits, + const std::uint8_t* src, + std::uint8_t right_symbol, + std::size_t weight, + std::size_t left_weight) noexcept { + pivco_simd_detail::partition_encode_one_constant_neon( + left_dst, bits, src, right_symbol, weight, left_weight); +} +#endif + +#if defined(PIXIE_PIVCO_NEON_SUPPORT) +/// @brief Pack a depth-2..8 flat subtree through the AArch64 NEON backend. +inline void pivco_flat_encode_neon(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse) noexcept { + pivco_simd_detail::flat_encode_dispatch_neon(bits, src, wire_code, count, + depth, known_bitreverse); +} +#endif + +#if defined(__AVX2__) && defined(__SSE4_1__) && defined(__BMI2__) +/// @brief Encode a balanced depth-2..7 table as AVX2 bit reversal. +/// @param known_bitreverse Skip the table scan when construction already +/// classified the mapping. +inline bool pivco_flat_encode_bitreverse( + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse = false) noexcept { + switch (depth) { + case 2: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<2>( + bits, src, wire_code, count, known_bitreverse); + case 3: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<3>( + bits, src, wire_code, count, known_bitreverse); + case 4: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<4>( + bits, src, wire_code, count, known_bitreverse); + case 5: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<5>( + bits, src, wire_code, count, known_bitreverse); + case 6: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<6>( + bits, src, wire_code, count, known_bitreverse); + case 7: + return pivco_simd_detail::flat_encode_bitreverse_maybe_avx2<7>( + bits, src, wire_code, count, known_bitreverse); + default: + return false; + } +} + +/// @brief Pack a depth-3 flat subtree through the AVX2/BMI2 kernel. +/// @pre @p wire_code has 259 readable bytes. +inline void pivco_flat_encode_d3(std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count) noexcept { + pivco_simd_detail::flat_encode_d3_avx2(bits, src, wire_code, count); +} +#endif + +#if defined(__AVX2__) && defined(__SSE4_1__) +/// @brief Encode the common full-alphabet bit-reversal table with AVX2. +/// @returns True when @p wire_code is the expected bit-reversal permutation. +/// @param known_bitreverse Skip the table scan when construction already +/// classified the mapping. +inline bool pivco_flat_encode_d8_bitreverse( + std::uint64_t* bits, + const std::uint8_t* src, + const std::uint8_t* wire_code, + std::size_t count, + bool known_bitreverse = false) noexcept { + if (known_bitreverse) { + pivco_simd_detail::flat_decode_d8_bitreverse_unchecked_avx2( + reinterpret_cast(bits), src, wire_code, count); + return true; + } + return pivco_simd_detail::flat_decode_d8_bitreverse_avx2( + reinterpret_cast(bits), src, wire_code, count); +} +#endif + +/// @brief Unpack and translate a flat subtree's fixed-width wire codes. +/// @param known_bitreverse Skip table classification when the caller already +/// established the balanced mapping during construction. +inline void pivco_flat_decode(std::uint8_t* dst, + const std::uint64_t* bits, + const std::uint8_t* table, + std::size_t count, + std::uint8_t depth, + bool known_bitreverse = false) noexcept { + pivco_simd_detail::flat_decode_dispatch(dst, bits, table, count, depth, + known_bitreverse); +} + +} // namespace pixie diff --git a/src/benchmarks/pivco_huffman_benchmarks.cpp b/src/benchmarks/pivco_huffman_benchmarks.cpp new file mode 100644 index 0000000..297b5c7 --- /dev/null +++ b/src/benchmarks/pivco_huffman_benchmarks.cpp @@ -0,0 +1,222 @@ +#include +#include + +#include +#include +#include +#include + +using pixie::PivCoHuffman; + +// --------------------------------------------------------------------------- +// Configuration: +// PIVCO_BENCH_SEED - random seed (default 42) +// +// Input sizes are swept by Google Benchmark's range mechanism, not an env +// var: 1 KiB -> 1 MiB -> 1 GiB (multiplier 1024). This exercises the codec +// from cache-resident to gigabyte scale in a single run. +// +// To run a subset of datasets or sizes, use Google Benchmark's filter flag, +// e.g. `--benchmark_filter="BM_Encode/Uniform256"`. +// --------------------------------------------------------------------------- + +namespace { + +/// @brief RNG seed, overridable via `PIVCO_BENCH_SEED`. +std::uint64_t bench_seed() { + if (const char* env = std::getenv("PIVCO_BENCH_SEED")) { + try { + return std::stoull(env); + } catch (...) { + } + } + return 42; +} + +/// @brief Signature of a dataset generator: produces @p n bytes from @p rng. +using DatasetMaker = std::vector (*)(std::size_t, + std::mt19937_64&); + +// --- built-in datasets ----------------------------------------------------- + +/// @brief Uniformly distributed bytes over the full 256-symbol alphabet. +std::vector make_uniform256(std::size_t n, std::mt19937_64& rng) { + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng()); + } + return data; +} + +/// @brief Binary alphabet {0, 1}: shortest possible codes (depth 1). +std::vector make_binary(std::size_t n, std::mt19937_64& rng) { + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng() & 1u); + } + return data; +} + +/// @brief Four-symbol alphabet {0,1,2,3}: code depth 2. +std::vector make_low_entropy_4(std::size_t n, + std::mt19937_64& rng) { + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng() & 3u); + } + return data; +} + +/// @brief Uniform distribution over a power-of-two alphabet. +/// @details Depths 3--7 exercise every flat width between the four-symbol and +/// full-byte datasets. +template +std::vector make_uniform_pow2(std::size_t n, + std::mt19937_64& rng) { + static_assert(kDepth >= 3 && kDepth <= 7); + constexpr std::uint8_t kMask = (std::uint8_t{1} << kDepth) - 1; + std::vector data(n); + for (auto& b : data) { + b = static_cast(rng()) & kMask; + } + return data; +} + +/// @brief English-text-like distribution: ~17% space, ~80% lowercase letters, +/// ~2% digits, ~1% punctuation. +std::vector make_text_like(std::size_t n, std::mt19937_64& rng) { + std::vector data(n); + std::uniform_int_distribution pick(0, 99); + std::uniform_int_distribution letter('a', 'z'); + std::uniform_int_distribution digit('0', '9'); + std::uniform_int_distribution punct('!', '/'); + for (auto& b : data) { + const int r = pick(rng); + if (r < 17) { + b = ' '; + } else if (r < 97) { + b = static_cast(letter(rng)); + } else if (r < 99) { + b = static_cast(digit(rng)); + } else { + b = static_cast(punct(rng)); + } + } + return data; +} + +/// @brief Highly skewed: 99% one symbol, 1% random over 256. Exercises a deep, +/// lopsided Huffman tree. +std::vector make_skewed99(std::size_t n, std::mt19937_64& rng) { + std::vector data(n, 0); + std::uniform_int_distribution coin(0, 99); + for (auto& b : data) { + if (coin(rng) < 1) { + b = static_cast(rng()); + } + } + return data; +} + +/// @brief Degenerate single-symbol input (root is a leaf): best-case path. +std::vector make_single_symbol(std::size_t n, + std::mt19937_64& rng) { + (void)rng; + return std::vector(n, 123); +} + +/// @brief Compression efficiency in bits per input byte (8 = no compression). +void report_ratio(benchmark::State& state, + std::size_t compressed_bytes, + std::size_t input_bytes) { + state.counters["bpb"] = static_cast(compressed_bytes) * 8.0 / + static_cast(input_bytes); +} + +/// @brief Common range of input sizes: 1 KiB, 1 MiB, 1 GiB. +/// @details Applied to every registered benchmark. The 1 GiB tier exercises +/// gigabyte-scale inputs; at that size Google Benchmark runs a single +/// iteration per case. +constexpr std::int64_t kSizeLo = 1ull << 10; // 1 KiB +constexpr std::int64_t kSizeHi = 1ull << 30; // 1 GiB +constexpr std::int64_t kSizeMult = 1024; + +// --- benchmarks ------------------------------------------------------------ + +/// @brief Measure full encoding: build Huffman tree, fill node bitmaps, and +/// serialize the compressed stream. The compressed size is read from +/// the first timed iteration (no separate probe build), which keeps +/// gigabyte-scale setup cost to a single encode. +void BM_Encode(benchmark::State& state, DatasetMaker make) { + const std::size_t size = static_cast(state.range(0)); + std::mt19937_64 rng(bench_seed()); + const std::vector data = make(size, rng); + + std::size_t compressed = 0; + for (auto _ : state) { + PivCoHuffman codec(data); + compressed = codec.compressed_size(); + benchmark::DoNotOptimize(codec.compressed_data().data()); + benchmark::ClobberMemory(); + } + if (compressed > 0) { + report_ratio(state, compressed, size); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +/// @brief Measure decoding: reconstruct the original byte stream from the +/// in-memory tree. Encoding is done once, outside the timed loop. +void BM_Decode(benchmark::State& state, DatasetMaker make) { + const std::size_t size = static_cast(state.range(0)); + std::mt19937_64 rng(bench_seed()); + const std::vector data = make(size, rng); + + const PivCoHuffman codec(data); + report_ratio(state, codec.compressed_size(), size); + + for (auto _ : state) { + std::vector out = codec.decode(); + benchmark::DoNotOptimize(out.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +} // namespace + +// Register encode + decode for every built-in dataset, each swept over the +// full size range. Each registration is named `BM_//`, so +// `--benchmark_filter` selects subsets freely. +#define PIXIE_PIVCO_REGISTER(Op, Dataset, Maker) \ + BENCHMARK_CAPTURE(BM_##Op, Dataset, Maker) \ + ->Unit(benchmark::kMillisecond) \ + ->UseRealTime() \ + ->Repetitions(5) \ + ->ReportAggregatesOnly(true) \ + ->RangeMultiplier(kSizeMult) \ + ->Range(kSizeLo, kSizeHi); + +PIXIE_PIVCO_REGISTER(Encode, Uniform256, make_uniform256) +PIXIE_PIVCO_REGISTER(Encode, Binary, make_binary) +PIXIE_PIVCO_REGISTER(Encode, LowEntropy4, make_low_entropy_4) +PIXIE_PIVCO_REGISTER(Encode, Uniform8, make_uniform_pow2<3>) +PIXIE_PIVCO_REGISTER(Encode, Uniform16, make_uniform_pow2<4>) +PIXIE_PIVCO_REGISTER(Encode, Uniform32, make_uniform_pow2<5>) +PIXIE_PIVCO_REGISTER(Encode, Uniform64, make_uniform_pow2<6>) +PIXIE_PIVCO_REGISTER(Encode, Uniform128, make_uniform_pow2<7>) +PIXIE_PIVCO_REGISTER(Encode, TextLike, make_text_like) +PIXIE_PIVCO_REGISTER(Encode, Skewed99, make_skewed99) +PIXIE_PIVCO_REGISTER(Encode, SingleSymbol, make_single_symbol) + +PIXIE_PIVCO_REGISTER(Decode, Uniform256, make_uniform256) +PIXIE_PIVCO_REGISTER(Decode, Binary, make_binary) +PIXIE_PIVCO_REGISTER(Decode, LowEntropy4, make_low_entropy_4) +PIXIE_PIVCO_REGISTER(Decode, Uniform8, make_uniform_pow2<3>) +PIXIE_PIVCO_REGISTER(Decode, Uniform16, make_uniform_pow2<4>) +PIXIE_PIVCO_REGISTER(Decode, Uniform32, make_uniform_pow2<5>) +PIXIE_PIVCO_REGISTER(Decode, Uniform64, make_uniform_pow2<6>) +PIXIE_PIVCO_REGISTER(Decode, Uniform128, make_uniform_pow2<7>) +PIXIE_PIVCO_REGISTER(Decode, TextLike, make_text_like) +PIXIE_PIVCO_REGISTER(Decode, Skewed99, make_skewed99) +PIXIE_PIVCO_REGISTER(Decode, SingleSymbol, make_single_symbol) diff --git a/src/benchmarks/pivco_huffman_file_benchmarks.cpp b/src/benchmarks/pivco_huffman_file_benchmarks.cpp new file mode 100644 index 0000000..1795103 --- /dev/null +++ b/src/benchmarks/pivco_huffman_file_benchmarks.cpp @@ -0,0 +1,135 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +using pixie::PivCoHuffman; + +// --------------------------------------------------------------------------- +// File-based PivCo-Huffman benchmark. +// +// Reads a real-world text file (default: prose_pride.txt, "Pride and +// Prejudice" from Project Gutenberg) and measures encode/decode throughput +// on it. This matches the article's evaluation methodology, which uses real +// datasets rather than synthetic random data. +// +// Configuration: +// PIVCO_BENCH_FILE - path to the input file (default: prose_pride.txt) +// --------------------------------------------------------------------------- + +namespace { + +/// @brief Read the entire contents of @p path into a byte vector. +std::vector read_file(const std::string& path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + std::fprintf(stderr, "Cannot open '%s'\n", path.c_str()); + std::exit(1); + } + const std::streamsize size = file.tellg(); + file.seekg(0); + std::vector data(static_cast(size)); + if (size > 0) { + file.read(reinterpret_cast(data.data()), size); + } + return data; +} + +/// @brief Input file path, overridable via `PIVCO_BENCH_FILE`. +const std::string& bench_file() { + static const std::string path = [] { + if (const char* env = std::getenv("PIVCO_BENCH_FILE")) { + return std::string(env); + } + return std::string("prose_pride.txt"); + }(); + return path; +} + +/// @brief Compression efficiency in bits per input byte. +void report_ratio(benchmark::State& state, + std::size_t compressed_bytes, + std::size_t input_bytes) { + state.counters["bpb"] = static_cast(compressed_bytes) * 8.0 / + static_cast(input_bytes); +} + +/// @brief Compute Shannon entropy of the data in bits per byte, for reference. +double entropy(std::span data) { + if (data.empty()) { + return 0.0; + } + std::array freq{}; + for (auto b : data) { + freq[b]++; + } + double h = 0.0; + const double n = static_cast(data.size()); + for (std::size_t i = 0; i < 256; i++) { + if (freq[i] > 0) { + const double p = static_cast(freq[i]) / n; + h -= p * std::log2(p); + } + } + return h; +} + +// --- benchmarks ------------------------------------------------------------ + +/// @brief Measure full encoding of the file. +void BM_EncodeFile(benchmark::State& state) { + const std::vector data = read_file(bench_file()); + const std::size_t size = data.size(); + + // Report entropy once as a reference counter. + state.counters["entropy_bpb"] = entropy(data); + + std::size_t compressed = 0; + for (auto _ : state) { + PivCoHuffman codec(data); + compressed = codec.compressed_size(); + benchmark::DoNotOptimize(codec.compressed_data().data()); + benchmark::ClobberMemory(); + } + if (compressed > 0) { + report_ratio(state, compressed, size); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +/// @brief Measure decoding of the file. Encoding is done once, outside timing. +void BM_DecodeFile(benchmark::State& state) { + const std::vector data = read_file(bench_file()); + const std::size_t size = data.size(); + + state.counters["entropy_bpb"] = entropy(data); + + const PivCoHuffman codec(data); + report_ratio(state, codec.compressed_size(), size); + + for (auto _ : state) { + std::vector out = codec.decode(); + benchmark::DoNotOptimize(out.data()); + benchmark::ClobberMemory(); + } + state.SetBytesProcessed(static_cast(state.iterations() * size)); +} + +} // namespace + +BENCHMARK(BM_EncodeFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Repetitions(10) + ->ReportAggregatesOnly(true); + +BENCHMARK(BM_DecodeFile) + ->Unit(benchmark::kMillisecond) + ->UseRealTime() + ->Repetitions(10) + ->ReportAggregatesOnly(true); diff --git a/src/tests/pivco_huffman_tests.cpp b/src/tests/pivco_huffman_tests.cpp new file mode 100644 index 0000000..965ac4c --- /dev/null +++ b/src/tests/pivco_huffman_tests.cpp @@ -0,0 +1,402 @@ +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +using pixie::PivCoHuffman; + +static_assert( + std::is_base_of_v, PivCoHuffman>); +static_assert(std::is_same_v); + +namespace { +// Round-trip through the in-memory tree of the same codec instance. +std::vector decode_same(const PivCoHuffman& codec) { + return codec.decode(); +} + +// Round-trip through the serialized byte stream: simulate writing the +// compressed form to disk and loading it into a fresh codec instance. +std::vector decode_via_bytes(const PivCoHuffman& codec) { + const auto bytes = codec.compressed_data(); + PivCoHuffman copy(bytes); + return copy.decode(); +} +} // namespace + +TEST(PivCoHuffmanSmoke, EmptyInputRoundTrips) { + const std::vector data; + const PivCoHuffman codec(data); + EXPECT_TRUE(codec.empty()); + EXPECT_EQ(codec.uncompressed_size(), 0u); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, SingleSymbolRoundTrips) { + const std::vector data(1000, 42); + const PivCoHuffman codec(data); + EXPECT_EQ(codec.uncompressed_size(), data.size()); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, TwoSymbolRoundTrips) { + std::vector data; + for (std::size_t i = 0; i < 500; i++) { + data.push_back(static_cast(i % 2)); + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, KnownInputRoundTrips) { + const std::vector data{3, 2, 0, 3, 1, 1, 2, 0, 3, 1}; + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, CompressedDataIsNonEmpty) { + const std::vector data{3, 2, 0, 3, 1, 1, 2, 0, 3, 1}; + const PivCoHuffman codec(data); + EXPECT_GT(codec.compressed_size(), 0u); + EXPECT_EQ(codec.compressed_data().size(), codec.compressed_size()); +} + +TEST(PivCoHuffmanSmoke, FullAlphabetRoundTrips) { + std::vector data(256); + for (std::size_t i = 0; i < 256; i++) { + data[i] = static_cast(i); + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, FlatPowerOfTwoAlphabetsRoundTrip) { + std::mt19937 rng(1871); + for (std::uint8_t depth = 1; depth <= 8; ++depth) { + const std::size_t alphabet_size = std::size_t{1} << depth; + std::vector data; + data.reserve(alphabet_size * 16); + for (std::size_t repetition = 0; repetition < 16; ++repetition) { + for (std::size_t symbol = 0; symbol < alphabet_size; ++symbol) { + data.push_back(static_cast(symbol)); + } + } + std::shuffle(data.begin(), data.end(), rng); + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data) << "flat depth " << +depth; + EXPECT_EQ(decode_via_bytes(codec), data) << "flat depth " << +depth; + } +} + +TEST(PivCoSimd, FlatKernelsDecodeFullGroupsAndTails) { + std::mt19937 rng(4189); + for (std::uint8_t depth = 1; depth <= 8; ++depth) { + const std::size_t alphabet_size = std::size_t{1} << depth; + const std::size_t count = 129 + depth; + std::vector table(alphabet_size); + for (std::size_t i = 0; i < alphabet_size; ++i) { + table[i] = static_cast(i); + } + std::shuffle(table.begin(), table.end(), rng); + + std::vector bits((count * depth + 63) / 64, 0); + std::vector expected(count); + for (std::size_t i = 0; i < count; ++i) { + const std::uint64_t code = rng() % alphabet_size; + const std::size_t bit_position = i * depth; + const std::size_t word_index = bit_position / 64; + const std::size_t bit_offset = bit_position % 64; + bits[word_index] |= code << bit_offset; + if (bit_offset + depth > 64) { + bits[word_index + 1] |= code >> (64 - bit_offset); + } + expected[i] = table[code]; + } + + std::vector decoded(count); + pixie::pivco_flat_decode(decoded.data(), bits.data(), table.data(), count, + depth); + EXPECT_EQ(decoded, expected) << "flat depth " << +depth; + } +} + +TEST(PivCoSimd, MergeVariantsMatchReferenceAcrossFullGroupsAndTail) { + constexpr std::size_t kCount = 263; + constexpr std::size_t kPadding = 64; + std::vector bits((kCount + 63) / 64, 0); + std::vector left(kCount + kPadding, 0); + std::vector right(kCount + kPadding, 0); + std::vector expected(kCount); + std::size_t left_count = 0; + std::size_t right_count = 0; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((i * 37 + i / 5 + 11) % 13) < 6; + if (goes_right) { + bits[i / 64] |= std::uint64_t{1} << (i % 64); + right[right_count] = static_cast(0x80u + right_count % 97); + expected[i] = right[right_count++]; + } else { + left[left_count] = static_cast(left_count % 113); + expected[i] = left[left_count++]; + } + } + + std::vector decoded(kCount + kPadding, 0); + pixie::pivco_merge_decode(decoded.data(), left.data(), right.data(), + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); + + constexpr std::uint8_t kLeftSymbol = 17; + constexpr std::uint8_t kRightSymbol = 231; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((bits[i / 64] >> (i % 64)) & 1u) != 0; + expected[i] = goes_right ? kRightSymbol : kLeftSymbol; + } + std::fill(decoded.begin(), decoded.end(), 0); + pixie::pivco_merge_decode_cst_cst(decoded.data(), kLeftSymbol, kRightSymbol, + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); + + right_count = 0; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((bits[i / 64] >> (i % 64)) & 1u) != 0; + expected[i] = goes_right ? right[right_count++] : kLeftSymbol; + } + std::fill(decoded.begin(), decoded.end(), 0); + pixie::pivco_merge_decode_cst_vec(decoded.data(), kLeftSymbol, right.data(), + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); + + left_count = 0; + for (std::size_t i = 0; i < kCount; ++i) { + const bool goes_right = ((bits[i / 64] >> (i % 64)) & 1u) != 0; + expected[i] = goes_right ? kRightSymbol : left[left_count++]; + } + std::fill(decoded.begin(), decoded.end(), 0); + pixie::pivco_merge_decode_vec_cst(decoded.data(), left.data(), kRightSymbol, + bits.data(), kCount); + EXPECT_TRUE(std::equal(expected.begin(), expected.end(), decoded.begin())); +} + +TEST(PivCoSimd, PartitionVariantsBuildTheSameDirectionMask) { + constexpr std::size_t kCount = 257; + std::array directions{}; + for (std::size_t symbol = 0; symbol < directions.size(); ++symbol) { + directions[symbol] = + static_cast(((symbol * 73) ^ (symbol >> 2)) & 1u); + } + + std::mt19937 rng(7127); + std::vector input(kCount); + std::vector expected_left; + std::vector expected_right; + std::vector expected_bits((kCount + 63) / 64, 0); + for (std::size_t i = 0; i < input.size(); ++i) { + input[i] = static_cast(rng()); + if (directions[input[i]]) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + expected_right.push_back(input[i]); + } else { + expected_left.push_back(input[i]); + } + } + + std::vector left(expected_left.size()); + std::vector right(expected_right.size()); + std::vector full_bits(expected_bits.size(), 0); + pixie::pivco_partition_encode(left.data(), right.data(), full_bits.data(), + input.data(), directions.data(), input.size(), + left.size(), right.size()); + EXPECT_EQ(left, expected_left); + EXPECT_EQ(right, expected_right); + EXPECT_EQ(full_bits, expected_bits); + + std::vector left_bits(expected_bits.size(), 0); + std::fill(left.begin(), left.end(), 0); + pixie::pivco_partition_encode_left(left.data(), left_bits.data(), + input.data(), directions.data(), + input.size(), left.size()); + EXPECT_EQ(left, expected_left); + EXPECT_EQ(left_bits, expected_bits); + + std::vector right_bits(expected_bits.size(), 0); + std::fill(right.begin(), right.end(), 0); + pixie::pivco_partition_encode_right(right.data(), right_bits.data(), + input.data(), directions.data(), + input.size(), right.size()); + EXPECT_EQ(right, expected_right); + EXPECT_EQ(right_bits, expected_bits); + + std::vector bitmap_only(expected_bits.size(), 0); + pixie::pivco_partition_encode_none(bitmap_only.data(), input.data(), + directions.data(), input.size()); + EXPECT_EQ(bitmap_only, expected_bits); +} + +TEST(PivCoSimd, ConstantLeafBitmapMatchesReferenceAcrossTail) { + constexpr std::size_t kCount = 267; + constexpr std::uint8_t kRightSymbol = 193; + std::vector input(kCount); + std::vector expected_bits((kCount + 63) / 64, 0); + for (std::size_t i = 0; i < kCount; ++i) { + const bool is_right = ((i * 19 + i / 3) % 11) < 5; + input[i] = is_right ? kRightSymbol : static_cast(i % 127); + if (is_right) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + } + } + + std::vector bits(expected_bits.size(), 0); + pixie::pivco_partition_encode_cst_cst(bits.data(), input.data(), kRightSymbol, + input.size()); + EXPECT_EQ(bits, expected_bits); +} + +#if (defined(__AVX2__) && defined(__SSE4_1__)) || \ + defined(PIXIE_PIVCO_NEON_SUPPORT) +TEST(PivCoSimd, OneConstantPartitionsMatchReferenceAcrossTail) { + constexpr std::size_t kCount = 269; + constexpr std::uint8_t kConstant = 241; + std::vector input(kCount); + std::vector expected_dense; + std::vector expected_bits((kCount + 63) / 64, 0); + for (std::size_t i = 0; i < kCount; ++i) { + const bool is_constant = ((i * 23 + i / 7) % 17) < 8; + input[i] = is_constant ? kConstant : static_cast(i % 191); + if (!is_constant) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + expected_dense.push_back(input[i]); + } + } + + std::vector dense(expected_dense.size()); + std::vector bits(expected_bits.size(), 0); + pixie::pivco_partition_encode_cst_vec(dense.data(), bits.data(), input.data(), + kConstant, input.size(), dense.size()); + EXPECT_EQ(dense, expected_dense); + EXPECT_EQ(bits, expected_bits); + + std::fill(expected_bits.begin(), expected_bits.end(), 0); + expected_dense.clear(); + for (std::size_t i = 0; i < kCount; ++i) { + const bool is_constant = input[i] == kConstant; + if (is_constant) { + expected_bits[i / 64] |= std::uint64_t{1} << (i % 64); + } else { + expected_dense.push_back(input[i]); + } + } + std::fill(dense.begin(), dense.end(), 0); + std::fill(bits.begin(), bits.end(), 0); + pixie::pivco_partition_encode_vec_cst(dense.data(), bits.data(), input.data(), + kConstant, input.size(), dense.size()); + EXPECT_EQ(dense, expected_dense); + EXPECT_EQ(bits, expected_bits); +} +#endif + +TEST(PivCoHuffmanSmoke, RandomUniformRoundTrips) { + std::mt19937 rng(239); + std::uniform_int_distribution dist(0, 255); + std::vector data(10000); + for (auto& b : data) { + b = static_cast(dist(rng)); + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, SkewedInputRoundTrips) { + // Most symbols are one value, a few are outliers: exercises a deep, lopsided + // Huffman tree without any single-symbol shortcut. + std::mt19937 rng(7); + std::uniform_int_distribution coin(0, 99); + std::vector data(10000, 7); + for (std::size_t i = 0; i < data.size(); i++) { + if (coin(rng) < 5) { + data[i] = static_cast(rng() % 256); + } + } + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, LengthLimitedTreeRoundTrips) { + // Fibonacci weights create a maximally deep unconstrained Huffman tree and + // exercise the 15-bit length correction. + std::vector data; + std::size_t previous = 1; + std::size_t current = 1; + for (std::uint8_t symbol = 0; symbol < 17; ++symbol) { + const std::size_t frequency = symbol < 2 ? 1 : previous + current; + if (symbol >= 2) { + previous = current; + current = frequency; + } + data.insert(data.end(), frequency, symbol); + } + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, MultipleBlocksRoundTrip) { + std::mt19937 rng(9127); + std::uniform_int_distribution dist(0, 255); + std::vector data(3 * 64 * 1024 + 137); + for (auto& byte : data) { + byte = static_cast(dist(rng)); + } + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data); + EXPECT_EQ(decode_via_bytes(codec), data); +} + +TEST(PivCoHuffmanSmoke, SharedTreeHandlesBlockLocalMissingSymbols) { + constexpr std::size_t kBlockSize = 64 * 1024; + std::vector data; + data.reserve(3 * kBlockSize); + data.insert(data.end(), kBlockSize, 17); + data.insert(data.end(), kBlockSize, 93); + data.insert(data.end(), kBlockSize, 241); + + const PivCoHuffman codec(data); + EXPECT_EQ(codec.decode(), data); + EXPECT_EQ(codec.decode(), data); + + const PivCoHuffman copy(codec.compressed_data()); + EXPECT_EQ(copy.decode(), data); + EXPECT_EQ(copy.decode(), data); +} + +TEST(PivCoHuffmanSmoke, RandomizedAlphabetAndBlockSizesRoundTrip) { + std::mt19937 rng(3319); + for (std::size_t test_case = 0; test_case < 32; ++test_case) { + const std::size_t alphabet_size = 1 + rng() % 256; + const std::size_t size = 1 + rng() % (2 * 64 * 1024); + std::vector data(size); + for (auto& byte : data) { + byte = static_cast(rng() % alphabet_size); + } + + const PivCoHuffman codec(data); + EXPECT_EQ(decode_same(codec), data) << "test case " << test_case; + EXPECT_EQ(decode_via_bytes(codec), data) << "test case " << test_case; + } +}