From 445d7d8f5e685abd89371bd937dfbbf189b5bd22 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Tue, 8 Sep 2026 21:35:24 +0300 Subject: [PATCH 1/3] Faster histogram collection and improved partitioning --- include/pixie/detail/byte_histogram.h | 51 ++++++++++ include/pixie/detail/wavelet_partition.h | 124 +++++++++++++++++++++++ include/pixie/file_archive.h | 12 ++- include/pixie/wavelet_tree/index.h | 89 ++++++++++++---- 4 files changed, 252 insertions(+), 24 deletions(-) create mode 100644 include/pixie/detail/byte_histogram.h create mode 100644 include/pixie/detail/wavelet_partition.h diff --git a/include/pixie/detail/byte_histogram.h b/include/pixie/detail/byte_histogram.h new file mode 100644 index 0000000..a6007a0 --- /dev/null +++ b/include/pixie/detail/byte_histogram.h @@ -0,0 +1,51 @@ +#pragma once + +/** + * @file byte_histogram.h + * @brief Dependency-reduced histogram construction for byte streams. + */ + +#include +#include +#include +#include + +namespace pixie::detail { + +/** @brief Incremental 256-bin byte histogram with four update lanes. */ +class ByteHistogram { + public: + /** @brief Add one contiguous chunk to the histogram. */ + void add(std::span bytes) { + std::size_t index = 0; + for (; index + kLaneCount <= bytes.size(); index += kLaneCount) { + ++partial_[0][std::to_integer(bytes[index])]; + ++partial_[1][std::to_integer(bytes[index + 1])]; + ++partial_[2][std::to_integer(bytes[index + 2])]; + ++partial_[3][std::to_integer(bytes[index + 3])]; + } + for (; index < bytes.size(); ++index) { + ++partial_[index & (kLaneCount - 1)] + [std::to_integer(bytes[index])]; + } + } + + /** @brief Return the combined symbol counts. */ + std::array counts() const { + std::array result{}; + for (std::size_t symbol = 0; symbol < result.size(); ++symbol) { + for (const auto& lane : partial_) { + result[symbol] += lane[symbol]; + } + } + return result; + } + + private: + static constexpr std::size_t kLaneCount = 4; + static_assert((kLaneCount & (kLaneCount - 1)) == 0); + + alignas(64) std::array, kLaneCount> partial_{}; +}; + +} // namespace pixie::detail diff --git a/include/pixie/detail/wavelet_partition.h b/include/pixie/detail/wavelet_partition.h new file mode 100644 index 0000000..efeba0b --- /dev/null +++ b/include/pixie/detail/wavelet_partition.h @@ -0,0 +1,124 @@ +#pragma once + +/** + * @file wavelet_partition.h + * @brief Bulk direction-mask and stable-partition primitives. + */ + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#if defined(PIXIE_AVX2_SUPPORT) || defined(PIXIE_AVX512_SUPPORT) +#include +#endif + +namespace pixie::detail { + +template +std::uint64_t wavelet_direction_mask(std::span ranks, + Symbol middle) { + std::uint64_t mask = 0; + std::size_t offset = 0; + + if constexpr (std::same_as) { +#if defined(PIXIE_AVX512_SUPPORT) + if (ranks.size() == 64) { + const __m512i values = _mm512_loadu_si512(ranks.data()); + const __m512i threshold = _mm512_set1_epi8(static_cast(middle)); + return _mm512_cmpge_epu8_mask(values, threshold); + } +#endif +#if defined(PIXIE_AVX2_SUPPORT) + const __m256i threshold = _mm256_set1_epi8(static_cast(middle)); + for (; offset + 32 <= ranks.size(); offset += 32) { + const __m256i values = _mm256_loadu_si256( + reinterpret_cast(ranks.data() + offset)); + const __m256i greater_or_equal = + _mm256_cmpeq_epi8(_mm256_max_epu8(values, threshold), values); + mask |= static_cast(static_cast( + _mm256_movemask_epi8(greater_or_equal))) + << offset; + } +#endif + } + + for (; offset < ranks.size(); ++offset) { + mask |= static_cast(ranks[offset] >= middle) << offset; + } + return mask; +} + +/** + * @brief Build packed node directions and stably partition ranks. + * @param input Ranks in the node's original subsequence order. + * @param middle First rank belonging to the right child. + * @param output Destination split into left then right subsequences. + * @param expected_left Number of ranks expected in the left child. + * @param write_left Whether the left subsequence is needed by a child node. + * @param write_right Whether the right subsequence is needed by a child node. + * @param directions Packed destination for one direction bit per input rank. + */ +template +void partition_wavelet_ranks(std::span input, + Symbol middle, + std::span output, + std::size_t expected_left, + bool write_left, + bool write_right, + PackedBitBuilder& directions) { + const bool output_required = write_left || write_right; + if ((output_required && input.size() != output.size()) || + expected_left > input.size()) { + throw std::invalid_argument("Invalid wavelet partition buffers"); + } + + std::size_t left = 0; + std::size_t right = expected_left; + for (std::size_t offset = 0; offset < input.size(); offset += 64) { + const std::size_t width = std::min(64, input.size() - offset); + const std::span block = input.subspan(offset, width); + const std::uint64_t right_mask = wavelet_direction_mask(block, middle); + directions.write_bits(right_mask, width); + + const std::uint64_t valid_mask = + width == 64 ? std::numeric_limits::max() + : (std::uint64_t{1} << width) - 1; + std::uint64_t left_mask = (~right_mask) & valid_mask; + if (write_left) { + while (left_mask != 0) { + const unsigned index = std::countr_zero(left_mask); + output[left++] = block[index]; + left_mask &= left_mask - 1; + } + } else { + left += std::popcount(left_mask); + } + + std::uint64_t remaining_right = right_mask & valid_mask; + if (write_right) { + while (remaining_right != 0) { + const unsigned index = std::countr_zero(remaining_right); + output[right++] = block[index]; + remaining_right &= remaining_right - 1; + } + } else { + right += std::popcount(remaining_right); + } + } + + if (left != expected_left || right != input.size()) { + throw std::invalid_argument( + "Wavelet partition does not match the supplied symbol counts"); + } +} + +} // namespace pixie::detail diff --git a/include/pixie/file_archive.h b/include/pixie/file_archive.h index fa7de6e..195af37 100644 --- a/include/pixie/file_archive.h +++ b/include/pixie/file_archive.h @@ -5,6 +5,7 @@ * @brief Self-contained byte-oriented file archives with line extraction. */ +#include #include #include #include @@ -378,7 +379,9 @@ class FileArchiveIndex : public FileArchiveBase>, * receives `(const FileArchiveSourceMetadata&, consumer)` and must pass the * same immutable content to `consumer` as byte spans on both calls. The * first pass derives metadata and symbol counts; the second constructs the - * tree without retaining complete source contents. + * tree without retaining original source buffers. Wavelet construction + * temporarily owns up to two byte buffers, each as large as the logical + * content. * @throws std::invalid_argument for invalid metadata or changed content. */ template @@ -546,8 +549,7 @@ class FileArchiveIndex : public FileArchiveBase>, [](const auto& left, const auto& right) { return left.path < right.path; }); - std::array - symbol_counts{}; + detail::ByteHistogram histogram; std::vector content_hashes; content_hashes.reserve(sources.size()); std::string paths; @@ -589,9 +591,9 @@ class FileArchiveIndex : public FileArchiveBase>, } logical_size_ += chunk.size(); utf8.Consume(chunk); + histogram.add(chunk); for (const std::byte byte : chunk) { const std::uint8_t value = std::to_integer(byte); - ++symbol_counts[value]; newlines += value == '\n' ? 1U : 0U; has_content = true; last_byte = value; @@ -613,6 +615,8 @@ class FileArchiveIndex : public FileArchiveBase>, paths_ = MakeStorage( std::as_bytes(std::span(paths.data(), paths.size()))); + const auto symbol_counts = histogram.counts(); + tree_.emplace( file_archive_detail::kByteAlphabetSize, symbol_counts, [&](auto&& emit) { diff --git a/include/pixie/wavelet_tree/index.h b/include/pixie/wavelet_tree/index.h index a45da64..c2baedf 100644 --- a/include/pixie/wavelet_tree/index.h +++ b/include/pixie/wavelet_tree/index.h @@ -1,6 +1,7 @@ #pragma once #include +#include #include #include #include @@ -44,6 +45,7 @@ class WaveletTreeIndex node_index_t left_child = npos; node_index_t right_child = npos; std::size_t middle; + std::size_t left_size = 0; PackedBitBuilder stream; explicit PreWaveletNode(std::size_t middle) : middle(middle) {} }; @@ -278,6 +280,7 @@ class WaveletTreeIndex middle = begin + (middle == npos ? (end - begin) / 2 : middle); nodes.emplace_back(middle); + nodes[result].left_size = prefix_sum[middle] - prefix_sum[begin]; nodes[result].stream.reserve_bits(prefix_sum[end] - prefix_sum[begin]); nodes[result].parent = parent; nodes[result].left_child = @@ -288,6 +291,64 @@ class WaveletTreeIndex return result; } + void build_node_streams(node_index_t node, + std::span input, + std::span output, + std::vector& nodes) + requires(std::same_as) + { + PreWaveletNode& current = nodes[node]; + detail::partition_wavelet_ranks( + std::span(input), static_cast(current.middle), + output, current.left_size, current.left_child != npos, + current.right_child != npos, current.stream); + + if (current.left_child != npos) { + build_node_streams(current.left_child, output.first(current.left_size), + input.first(current.left_size), nodes); + } + if (current.right_child != npos) { + build_node_streams(current.right_child, output.subspan(current.left_size), + input.subspan(current.left_size), nodes); + } + } + + template + void build_bit_streams(std::span symbol_counts, + ForEachSymbol& for_each_symbol, + std::vector& nodes) + requires(std::same_as) + { + std::vector actual_counts(alphabet_size_); + std::vector ranks; + if (root_ != npos) { + ranks.reserve(data_size_); + } + for_each_symbol([&](Symbol symbol) { + const std::size_t original = checked_symbol_index(symbol, alphabet_size_); + if (actual_counts[original] == std::numeric_limits::max()) { + throw std::length_error("Wavelet-tree symbol count is too large"); + } + ++actual_counts[original]; + if (root_ != npos) { + ranks.push_back(static_cast(permutation_[original])); + } + }); + if (!std::ranges::equal(actual_counts, symbol_counts)) { + throw std::invalid_argument( + "Wavelet-tree emitted symbols do not match their counts"); + } + if (root_ == npos) { + return; + } + + const PreWaveletNode& root = nodes[root_]; + const bool root_needs_output = + root.left_child != npos || root.right_child != npos; + std::vector scratch(root_needs_output ? data_size_ : 0); + build_node_streams(root_, ranks, scratch, nodes); + } + /** * @brief Recursively copies segment of the original data corresponding to the * node @@ -461,25 +522,7 @@ class WaveletTreeIndex nodes); } - std::vector actual_counts(alphabet_size_); - for_each_symbol([&](Symbol symbol) { - const std::size_t original = checked_symbol_index(symbol, alphabet_size_); - if (actual_counts[original] == std::numeric_limits::max()) { - throw std::length_error("Wavelet-tree symbol count is too large"); - } - ++actual_counts[original]; - const std::size_t permuted = permutation_[original]; - for (node_index_t current = root_; current != npos;) { - auto& node = nodes[current]; - const bool go_right = permuted >= node.middle; - node.stream.write_bit(go_right); - current = go_right ? node.right_child : node.left_child; - } - }); - if (!std::ranges::equal(actual_counts, symbol_counts)) { - throw std::invalid_argument( - "Wavelet-tree emitted symbols do not match their counts"); - } + build_bit_streams(symbol_counts, for_each_symbol, nodes); nodes_.reserve(nodes.size()); for (auto& node : nodes) { @@ -494,6 +537,9 @@ class WaveletTreeIndex /** * @brief Construct from a contiguous sequence of typed symbols. + * @details Construction temporarily owns up to two buffers of one `Symbol` + * per input symbol. The buffers are released before the permanent node + * indexes are materialized. * @param alphabet_size Dense alphabet size; every symbol must be smaller. * @param data Input symbols retained only for the duration of construction. * @param build_type Standard or Huffman-shaped construction. @@ -525,7 +571,10 @@ class WaveletTreeIndex * @details @p for_each_symbol is invoked exactly once with a consumer that * accepts one `Symbol`. Emitted symbols must exactly match @p symbol_counts; * this permits callers to scan a replayable source once for counts and once - * for construction without materializing the sequence. + * for construction without materializing the sequence themselves. + * Construction temporarily owns up to two buffers of one `Symbol` per + * emitted symbol. The buffers are released before the permanent node indexes + * are materialized. * @param alphabet_size Dense alphabet size. * @param symbol_counts Count for every symbol in alphabet order. * @param for_each_symbol Callable accepting the construction consumer. From 87e3dce176d58f3ed9a73f380e1b0828cde21bc7 Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Wed, 9 Sep 2026 01:59:15 +0300 Subject: [PATCH 2/3] Precomputed partition table --- include/pixie/detail/wavelet_partition.h | 338 +++++++++++++++++++++-- src/tests/wavelet_tree_tests.cpp | 41 +++ 2 files changed, 352 insertions(+), 27 deletions(-) diff --git a/include/pixie/detail/wavelet_partition.h b/include/pixie/detail/wavelet_partition.h index efeba0b..c0a58bb 100644 --- a/include/pixie/detail/wavelet_partition.h +++ b/include/pixie/detail/wavelet_partition.h @@ -9,10 +9,12 @@ #include #include +#include #include #include #include #include +#include #include #include #include @@ -23,6 +25,171 @@ namespace pixie::detail { +#if defined(PIXIE_AVX2_SUPPORT) + +consteval auto make_wavelet_compaction_table() { + std::array, 256> table{}; + for (std::size_t mask = 0; mask < table.size(); ++mask) { + auto& shuffle = table[mask]; + shuffle.fill(0x80); + std::size_t left = 0; + std::size_t right = 8; + for (std::size_t index = 0; index < 8; ++index) { + if ((mask & (std::size_t{1} << index)) == 0) { + shuffle[left++] = static_cast(index); + } else { + shuffle[right++] = static_cast(index); + } + } + } + return table; +} + +alignas(64) inline constexpr auto kWaveletCompactionTable = + make_wavelet_compaction_table(); + +inline void store_compacted_bytes(std::uint8_t* destination, + std::uint64_t packed, + unsigned count) { + switch (count) { + case 8: + std::memcpy(destination, &packed, 8); + return; + case 7: { + const std::uint32_t suffix = static_cast(packed >> 24); + std::memcpy(destination, &packed, 4); + std::memcpy(destination + 3, &suffix, 4); + return; + } + case 6: { + const std::uint16_t suffix = static_cast(packed >> 32); + std::memcpy(destination, &packed, 4); + std::memcpy(destination + 4, &suffix, 2); + return; + } + case 5: + std::memcpy(destination, &packed, 4); + destination[4] = static_cast(packed >> 32); + return; + case 4: + std::memcpy(destination, &packed, 4); + return; + case 3: { + const std::uint16_t prefix = static_cast(packed); + std::memcpy(destination, &prefix, 2); + destination[2] = static_cast(packed >> 16); + return; + } + case 2: { + const std::uint16_t prefix = static_cast(packed); + std::memcpy(destination, &prefix, 2); + return; + } + case 1: + destination[0] = static_cast(packed); + return; + default: + return; + } +} + +// Full-width stores avoid a variable-size store in the hot loop. The logical +// cursor advances by only the valid byte count, so the next store replaces the +// preceding store's unused suffix. The allocation boundary still uses an exact +// store. +class OverlappingByteWriter { + public: + OverlappingByteWriter(std::uint8_t* destination, std::uint8_t* end) + : destination_(destination), end_(end) {} + + void append(std::uint64_t packed, unsigned count) { + if (count == 0) { + return; + } + if (end_ - destination_ >= 8) { + std::memcpy(destination_, &packed, 8); + } else { + store_compacted_bytes(destination_, packed, count); + } + destination_ += count; + } + + void append(std::uint8_t value) { append(value, 1); } + + private: + std::uint8_t* destination_; + std::uint8_t* end_; +}; + +// A left-partition store can overlap the start of its adjacent right partition +// by at most seven bytes. Retain enough of the right stream to restore it once +// partitioning is complete. +class RightPartitionPrefix { + public: + void append(std::uint64_t packed, unsigned count) { + if (size_ == bytes_.size() || count == 0) { + return; + } + const std::size_t copied = + std::min(count, bytes_.size() - size_); + std::memcpy(bytes_.data() + size_, &packed, copied); + size_ += copied; + } + + void append(std::uint8_t value) { append(value, 1); } + + void restore(std::uint8_t* destination) const { + std::memcpy(destination, bytes_.data(), size_); + } + + private: + std::array bytes_{}; + std::size_t size_ = 0; +}; + +template +std::size_t compact_wavelet_block_avx2(std::span input, + std::uint64_t right_mask, + OverlappingByteWriter& left_output, + OverlappingByteWriter& right_output, + RightPartitionPrefix& right_prefix, + std::size_t& left, + std::size_t& right) { + std::size_t offset = 0; + for (; offset + 8 <= input.size(); offset += 8) { + const auto mask = static_cast(right_mask >> offset); + const unsigned right_count = std::popcount(mask); + const unsigned left_count = 8 - right_count; + + if constexpr (WriteLeft || WriteRight) { + const __m128i values = _mm_loadl_epi64( + reinterpret_cast(input.data() + offset)); + const __m128i shuffle = _mm_load_si128(reinterpret_cast( + kWaveletCompactionTable[mask].data())); + const __m128i compacted = _mm_shuffle_epi8(values, shuffle); + if constexpr (WriteLeft) { + const std::uint64_t packed = + static_cast(_mm_cvtsi128_si64(compacted)); + left_output.append(packed, left_count); + } + if constexpr (WriteRight) { + const std::uint64_t packed = static_cast( + _mm_cvtsi128_si64(_mm_srli_si128(compacted, 8))); + if constexpr (WriteLeft) { + right_prefix.append(packed, right_count); + } + right_output.append(packed, right_count); + } + } + + left += left_count; + right += right_count; + } + return offset; +} + +#endif + template std::uint64_t wavelet_direction_mask(std::span ranks, Symbol middle) { @@ -57,6 +224,139 @@ std::uint64_t wavelet_direction_mask(std::span ranks, return mask; } +template +void partition_wavelet_ranks_scalar(std::span input, + Symbol middle, + std::span output, + PackedBitBuilder& directions, + std::size_t& left, + std::size_t& right) { + for (std::size_t offset = 0; offset < input.size(); offset += 64) { + const std::size_t width = std::min(64, input.size() - offset); + const std::span block = input.subspan(offset, width); + const std::uint64_t right_mask = wavelet_direction_mask(block, middle); + directions.write_bits(right_mask, width); + + const std::uint64_t valid_mask = + width == 64 ? std::numeric_limits::max() + : (std::uint64_t{1} << width) - 1; + std::uint64_t left_mask = (~right_mask) & valid_mask; + if constexpr (WriteLeft) { + while (left_mask != 0) { + const unsigned index = std::countr_zero(left_mask); + output[left++] = block[index]; + left_mask &= left_mask - 1; + } + } else { + left += std::popcount(left_mask); + } + + std::uint64_t remaining_right = right_mask & valid_mask; + if constexpr (WriteRight) { + while (remaining_right != 0) { + const unsigned index = std::countr_zero(remaining_right); + output[right++] = block[index]; + remaining_right &= remaining_right - 1; + } + } else { + right += std::popcount(remaining_right); + } + } +} + +#if defined(PIXIE_AVX2_SUPPORT) + +template +void partition_wavelet_bytes_avx2(std::span input, + std::uint8_t middle, + std::span output, + PackedBitBuilder& directions, + std::size_t& left, + std::size_t& right) { + std::uint8_t* const output_end = + output.empty() ? nullptr : output.data() + output.size(); + OverlappingByteWriter left_output( + WriteLeft && !output.empty() ? output.data() + left : nullptr, + output_end); + OverlappingByteWriter right_output( + WriteRight && !output.empty() ? output.data() + right : nullptr, + output_end); + RightPartitionPrefix right_prefix; + const std::size_t right_begin = right; + + for (std::size_t offset = 0; offset < input.size(); offset += 64) { + const std::size_t width = std::min(64, input.size() - offset); + const std::span block = input.subspan(offset, width); + const std::uint64_t right_mask = wavelet_direction_mask(block, middle); + directions.write_bits(right_mask, width); + + const std::size_t compacted = + compact_wavelet_block_avx2( + block, right_mask, left_output, right_output, right_prefix, left, + right); + const std::uint64_t valid_mask = + width == 64 ? std::numeric_limits::max() + : (std::uint64_t{1} << width) - 1; + const std::uint64_t remaining_mask = + compacted == 64 + ? 0 + : valid_mask & + (std::numeric_limits::max() << compacted); + std::uint64_t left_mask = (~right_mask) & remaining_mask; + if constexpr (WriteLeft) { + while (left_mask != 0) { + const unsigned index = std::countr_zero(left_mask); + left_output.append(block[index]); + ++left; + left_mask &= left_mask - 1; + } + } else { + left += std::popcount(left_mask); + } + + std::uint64_t remaining_right = right_mask & remaining_mask; + if constexpr (WriteRight) { + while (remaining_right != 0) { + const unsigned index = std::countr_zero(remaining_right); + if constexpr (WriteLeft) { + right_prefix.append(block[index]); + } + right_output.append(block[index]); + ++right; + remaining_right &= remaining_right - 1; + } + } else { + right += std::popcount(remaining_right); + } + } + + if constexpr (WriteLeft && WriteRight) { + if (!output.empty()) { + right_prefix.restore(output.data() + right_begin); + } + } +} + +#endif + +template +void partition_wavelet_ranks_impl(std::span input, + Symbol middle, + std::span output, + PackedBitBuilder& directions, + std::size_t& left, + std::size_t& right) { +#if defined(PIXIE_AVX2_SUPPORT) + if constexpr (std::same_as) { + partition_wavelet_bytes_avx2( + input, middle, output, directions, left, right); + return; + } +#endif + partition_wavelet_ranks_scalar( + input, middle, output, directions, left, right); +} + /** * @brief Build packed node directions and stably partition ranks. * @param input Ranks in the node's original subsequence order. @@ -83,36 +383,20 @@ void partition_wavelet_ranks(std::span input, std::size_t left = 0; std::size_t right = expected_left; - for (std::size_t offset = 0; offset < input.size(); offset += 64) { - const std::size_t width = std::min(64, input.size() - offset); - const std::span block = input.subspan(offset, width); - const std::uint64_t right_mask = wavelet_direction_mask(block, middle); - directions.write_bits(right_mask, width); - - const std::uint64_t valid_mask = - width == 64 ? std::numeric_limits::max() - : (std::uint64_t{1} << width) - 1; - std::uint64_t left_mask = (~right_mask) & valid_mask; - if (write_left) { - while (left_mask != 0) { - const unsigned index = std::countr_zero(left_mask); - output[left++] = block[index]; - left_mask &= left_mask - 1; - } - } else { - left += std::popcount(left_mask); - } - - std::uint64_t remaining_right = right_mask & valid_mask; + if (write_left) { if (write_right) { - while (remaining_right != 0) { - const unsigned index = std::countr_zero(remaining_right); - output[right++] = block[index]; - remaining_right &= remaining_right - 1; - } + partition_wavelet_ranks_impl(input, middle, output, + directions, left, right); } else { - right += std::popcount(remaining_right); + partition_wavelet_ranks_impl(input, middle, output, + directions, left, right); } + } else if (write_right) { + partition_wavelet_ranks_impl(input, middle, output, directions, + left, right); + } else { + partition_wavelet_ranks_impl(input, middle, output, + directions, left, right); } if (left != expected_left || right != input.size()) { diff --git a/src/tests/wavelet_tree_tests.cpp b/src/tests/wavelet_tree_tests.cpp index 52d1cb1..0f741be 100644 --- a/src/tests/wavelet_tree_tests.cpp +++ b/src/tests/wavelet_tree_tests.cpp @@ -219,6 +219,47 @@ TEST(WaveletTreeTest, TypedByteSymbolsRoundTripWithoutWidening) { EXPECT_EQ(wrong_symbol_reader.position(), 0u); } +TEST(WaveletTreeTest, ByteBuildHandlesSkewAndPackedBlockTails) { + std::vector data; + const std::array frequencies = {129, 65, 33, 17, 9, + 5, 3, 2, 2}; + for (std::size_t symbol = 0; symbol < frequencies.size(); ++symbol) { + data.insert(data.end(), frequencies[symbol], + static_cast(symbol)); + } + std::mt19937_64 rng(42); + std::shuffle(data.begin(), data.end(), rng); + + for (const auto build_type : {pixie::WaveletTreeBuildType::Standard, + pixie::WaveletTreeBuildType::Huffman}) { + const pixie::WaveletTree tree(256, data, build_type); + EXPECT_EQ(tree.get_segment(0, data.size()), data); + for (std::size_t symbol = 0; symbol < frequencies.size(); ++symbol) { + EXPECT_EQ(tree.rank(static_cast(symbol), data.size()), + frequencies[symbol]); + EXPECT_LT( + tree.select(static_cast(symbol), frequencies[symbol]), + data.size()); + } + } + + for (const std::size_t size : {7u, 8u, 9u, 15u, 16u, 17u, 63u, 64u, 65u}) { + for (std::size_t left_size = 0; left_size <= size; ++left_size) { + SCOPED_TRACE(testing::Message() + << "size=" << size << ", left_size=" << left_size); + std::vector boundary_data(size); + for (std::size_t index = 0; index < size; ++index) { + boundary_data[index] = static_cast( + index < left_size ? index % 2 : 2 + index % 2); + } + std::shuffle(boundary_data.begin(), boundary_data.end(), rng); + const pixie::WaveletTree tree( + 4, boundary_data, pixie::WaveletTreeBuildType::Standard); + EXPECT_EQ(tree.get_segment(0, boundary_data.size()), boundary_data); + } + } +} + TEST(WaveletTreeTest, BuildsFromCountsAndOneStreamedPass) { const std::vector data = {3, 0, 1, 3, 2, 1, 0}; const std::array counts = {2, 2, 1, 2}; From 68b691f409001a5a05b103a673ef581326cb929c Mon Sep 17 00:00:00 2001 From: Sanya239 Date: Thu, 10 Sep 2026 18:45:53 +0300 Subject: [PATCH 3/3] Speed up Huffman wavelet tree construction --- include/pixie/detail/huffman_build_table.h | 412 +++++++++++++++++ include/pixie/detail/wavelet_partition.h | 489 +++++++++++++++++++++ include/pixie/file_archive.h | 5 +- include/pixie/wavelet_tree/index.h | 275 +++++++++++- src/tests/wavelet_tree_tests.cpp | 72 +++ 5 files changed, 1242 insertions(+), 11 deletions(-) create mode 100644 include/pixie/detail/huffman_build_table.h diff --git a/include/pixie/detail/huffman_build_table.h b/include/pixie/detail/huffman_build_table.h new file mode 100644 index 0000000..d2b3fc3 --- /dev/null +++ b/include/pixie/detail/huffman_build_table.h @@ -0,0 +1,412 @@ +#pragma once + +/** + * @file huffman_build_table.h + * @brief Length-limited byte Huffman tree construction for wavelet indexes. + * + * This file is derived from PivCo's `huffman_table.c` and has been modified + * for Pixie's header-only C++ representation. Wire-format, flat-subtree, + * entropy-coding, and decoder tables were removed. The two-queue length + * builder, length limiter, fused canonical tree shaping, and in-order rank + * assignment are retained. Both projects are distributed under Apache-2.0. + */ + +#include +#include +#include +#include +#include +#include +#include + +namespace pixie::detail { + +inline constexpr std::size_t kByteAlphabetSize = 256; +inline constexpr std::size_t kMaximumHuffmanNodes = 2 * kByteAlphabetSize - 1; +inline constexpr std::size_t kMaximumHuffmanCodeLength = 11; + +struct HuffmanBuildNode { + std::int16_t symbol = -1; + std::int16_t left = -1; + std::int16_t right = -1; +}; + +enum class HuffmanBuildNodeType : std::uint8_t { + kInternal, + kBothLeaves, + kLeftLeaf, + kLeaf, +}; + +struct HuffmanBuildTable { + std::array tree{}; + std::array split_rank{}; + std::array node_type{}; + std::array symbol_to_rank{}; + std::int16_t root = -1; + std::size_t node_count = 0; + std::size_t symbol_count = 0; +}; + +namespace huffman_build_detail { + +struct FrequencyLeaf { + std::size_t frequency; + std::uint16_t symbol; +}; + +// Stable LSD radix sort. The input is seeded in symbol order, so equal +// frequencies retain PivCo's (frequency, symbol) tie discipline. +inline void sort_leaves_by_frequency( + std::span leaves, + std::array& temporary) { + std::size_t maximum = 0; + for (const FrequencyLeaf leaf : leaves) { + maximum = std::max(maximum, leaf.frequency); + } + std::size_t byte_count = 0; + while (maximum != 0) { + ++byte_count; + maximum >>= 8; + } + + FrequencyLeaf* source = leaves.data(); + FrequencyLeaf* destination = temporary.data(); + for (std::size_t byte = 0; byte < byte_count; ++byte) { + const std::size_t shift = byte * 8; + std::array counts{}; + for (std::size_t index = 0; index < leaves.size(); ++index) { + ++counts[(source[index].frequency >> shift) & 0xff]; + } + std::size_t prefix = 0; + for (std::size_t& count : counts) { + const std::size_t current = count; + count = prefix; + prefix += current; + } + for (std::size_t index = 0; index < leaves.size(); ++index) { + const std::size_t bucket = (source[index].frequency >> shift) & 0xff; + destination[counts[bucket]++] = source[index]; + } + std::swap(source, destination); + } + if (source != leaves.data()) { + std::copy_n(source, leaves.size(), leaves.data()); + } +} + +inline std::array build_code_lengths( + std::span frequencies, + std::span used_symbols) { + std::array lengths{}; + if (used_symbols.size() < 2) { + if (!used_symbols.empty()) { + lengths[used_symbols.front()] = 1; + } + return lengths; + } + + std::array leaves{}; + for (std::size_t index = 0; index < used_symbols.size(); ++index) { + const std::uint16_t symbol = used_symbols[index]; + leaves[index] = {frequencies[symbol], symbol}; + } + std::array sort_temporary{}; + sort_leaves_by_frequency(std::span(leaves).first(used_symbols.size()), + sort_temporary); + + std::array node_frequency{}; + std::array parent{}; + const std::size_t leaf_count = used_symbols.size(); + for (std::size_t index = 0; index < leaf_count; ++index) { + node_frequency[index] = leaves[index].frequency; + } + + std::size_t next_leaf = 0; + std::size_t internal_head = leaf_count; + std::size_t next_internal = leaf_count; + const auto take_minimum = [&]() { + if (next_leaf < leaf_count && + (internal_head == next_internal || + node_frequency[next_leaf] <= node_frequency[internal_head])) { + return next_leaf++; + } + return internal_head++; + }; + for (std::size_t remaining = leaf_count; remaining > 1; --remaining) { + const std::size_t left = take_minimum(); + const std::size_t right = take_minimum(); + node_frequency[next_internal] = + node_frequency[left] + node_frequency[right]; + parent[left] = static_cast(next_internal); + parent[right] = static_cast(next_internal); + ++next_internal; + } + + const std::size_t root = next_internal - 1; + std::array depth{}; + for (std::size_t index = root; index-- > 0;) { + depth[index] = static_cast(depth[parent[index]] + 1); + } + for (std::size_t index = 0; index < leaf_count; ++index) { + lengths[leaves[index].symbol] = depth[index]; + } + return lengths; +} + +// PivCo's DEFLATE-style limiter caps tree height so byte-rank construction has +// a small, predictable recursion and scratch bound. +inline void limit_code_lengths( + std::array& lengths) { + std::array length_counts{}; + std::size_t maximum = 0; + for (const std::uint8_t length : lengths) { + if (length != 0) { + ++length_counts[length]; + maximum = std::max(maximum, static_cast(length)); + } + } + if (maximum <= kMaximumHuffmanCodeLength) { + return; + } + + for (std::size_t length = maximum; length > kMaximumHuffmanCodeLength; + --length) { + length_counts[kMaximumHuffmanCodeLength] += length_counts[length]; + length_counts[length] = 0; + } + + std::size_t kraft = 0; + for (std::size_t length = 1; length <= kMaximumHuffmanCodeLength; ++length) { + kraft += length_counts[length] << (kMaximumHuffmanCodeLength - length); + } + constexpr std::size_t kKraftTarget = std::size_t{1} + << kMaximumHuffmanCodeLength; + while (kraft > kKraftTarget) { + std::size_t best = kMaximumHuffmanCodeLength - 1; + while (best > 0 && length_counts[best] == 0) { + --best; + } + if (best == 0) { + break; + } + --length_counts[best]; + ++length_counts[best + 1]; + kraft -= std::size_t{1} << (kMaximumHuffmanCodeLength - best - 1); + } + while (kraft < kKraftTarget && + length_counts[kMaximumHuffmanCodeLength] != 0) { + bool shortened = false; + for (std::size_t length = kMaximumHuffmanCodeLength - 1; length > 0; + --length) { + const std::size_t delta = + (std::size_t{1} << (kMaximumHuffmanCodeLength - length)) - 1; + if (kraft + delta <= kKraftTarget) { + --length_counts[kMaximumHuffmanCodeLength]; + ++length_counts[length]; + kraft += delta; + shortened = true; + break; + } + } + if (!shortened) { + break; + } + } + + struct LengthSymbol { + std::uint8_t length; + std::uint8_t symbol; + }; + std::array ordered{}; + std::size_t count = 0; + for (std::size_t symbol = 0; symbol < lengths.size(); ++symbol) { + if (lengths[symbol] != 0) { + ordered[count++] = { + std::min(lengths[symbol], kMaximumHuffmanCodeLength), + static_cast(symbol)}; + } + } + for (std::size_t index = 1; index < count; ++index) { + const LengthSymbol current = ordered[index]; + std::size_t position = index; + while (position != 0 && ordered[position - 1].length > current.length) { + ordered[position] = ordered[position - 1]; + --position; + } + ordered[position] = current; + } + std::size_t ordered_index = 0; + for (std::size_t length = 1; length <= kMaximumHuffmanCodeLength; ++length) { + for (std::size_t index = 0; index < length_counts[length]; ++index) { + lengths[ordered[ordered_index++].symbol] = + static_cast(length); + } + } +} + +struct FusedChunk { + std::uint16_t suffix_bits; + std::uint16_t depth; + std::uint16_t symbol_count; + std::uint16_t root_code; + std::size_t symbol_begin; +}; + +inline std::uint16_t assign_inorder_ranks(HuffmanBuildTable& table, + std::int16_t node_id, + std::uint16_t rank) { + const HuffmanBuildNode& node = table.tree[node_id]; + if (node.symbol >= 0) { + table.symbol_to_rank[node.symbol] = static_cast(rank); + return static_cast(rank + 1); + } + rank = assign_inorder_ranks(table, node.left, rank); + table.split_rank[node_id] = static_cast(rank - 1); + return assign_inorder_ranks(table, node.right, rank); +} + +inline void build_fused_tree( + const std::array& lengths, + HuffmanBuildTable& table) { + std::array length_counts{}; + std::size_t maximum_length = 0; + for (const std::uint8_t length : lengths) { + if (length != 0) { + ++length_counts[length]; + maximum_length = + std::max(maximum_length, static_cast(length)); + } + } + + std::array ordered_symbols{}; + std::array length_begin{}; + std::array cursor{}; + std::size_t accumulated = 0; + for (std::size_t length = 1; length <= maximum_length; ++length) { + length_begin[length] = accumulated; + cursor[length] = accumulated; + accumulated += length_counts[length]; + } + length_begin[maximum_length + 1] = accumulated; + for (std::size_t symbol = 0; symbol < lengths.size(); ++symbol) { + const std::uint8_t length = lengths[symbol]; + if (length != 0) { + ordered_symbols[cursor[length]++] = static_cast(symbol); + } + } + + std::array chunks{}; + std::size_t chunk_count = 0; + for (std::size_t length = 1; length <= maximum_length; ++length) { + std::size_t symbol = length_begin[length]; + for (std::size_t pair = 0; pair < length_counts[length] / 2; ++pair) { + chunks[chunk_count++] = {1, static_cast(length - 1), 2, 0, + symbol}; + symbol += 2; + } + if ((length_counts[length] & 1U) != 0) { + chunks[chunk_count++] = {0, static_cast(length), 1, 0, + symbol}; + } + } + for (std::size_t index = 1; index < chunk_count; ++index) { + const FusedChunk current = chunks[index]; + std::size_t position = index; + while (position != 0 && chunks[position - 1].depth > current.depth) { + chunks[position] = chunks[position - 1]; + --position; + } + chunks[position] = current; + } + + std::uint32_t code = 0; + std::size_t previous_depth = 0; + for (std::size_t index = 0; index < chunk_count; ++index) { + FusedChunk& chunk = chunks[index]; + code <<= chunk.depth - previous_depth; + chunk.root_code = static_cast(code); + ++code; + previous_depth = chunk.depth; + } + + table.root = 0; + table.node_count = 1; + for (std::size_t index = 0; index < chunk_count; ++index) { + const FusedChunk& chunk = chunks[index]; + std::int16_t node_id = table.root; + for (std::size_t bit = chunk.depth; bit-- > 0;) { + const bool right = ((chunk.root_code >> bit) & 1U) != 0; + std::int16_t& child = + right ? table.tree[node_id].right : table.tree[node_id].left; + if (child < 0) { + child = static_cast(table.node_count++); + } + node_id = child; + } + if (chunk.suffix_bits == 1) { + HuffmanBuildNode& node = table.tree[node_id]; + node.left = static_cast(table.node_count++); + table.tree[node.left].symbol = ordered_symbols[chunk.symbol_begin]; + node.right = static_cast(table.node_count++); + table.tree[node.right].symbol = ordered_symbols[chunk.symbol_begin + 1]; + } else { + table.tree[node_id].symbol = ordered_symbols[chunk.symbol_begin]; + } + } + + assign_inorder_ranks(table, table.root, 0); + for (std::size_t index = 0; index < table.node_count; ++index) { + const HuffmanBuildNode& node = table.tree[index]; + if (node.symbol >= 0) { + table.node_type[index] = HuffmanBuildNodeType::kLeaf; + continue; + } + const bool left_leaf = table.tree[node.left].symbol >= 0; + const bool right_leaf = table.tree[node.right].symbol >= 0; + if (left_leaf && right_leaf) { + table.node_type[index] = HuffmanBuildNodeType::kBothLeaves; + } else if (left_leaf) { + table.node_type[index] = HuffmanBuildNodeType::kLeftLeaf; + } else { + table.node_type[index] = HuffmanBuildNodeType::kInternal; + } + } +} + +} // namespace huffman_build_detail + +/** + * @brief Build PivCo's fused, non-flat byte Huffman tree. + * @param frequencies Frequencies for a dense alphabet of at most 256 symbols. + * @return Tree, node dispatch classes, symbol ranks, and split ranks. + */ +inline HuffmanBuildTable build_huffman_table( + std::span frequencies) { + HuffmanBuildTable table; + std::array used_symbols{}; + for (std::size_t symbol = 0; symbol < frequencies.size(); ++symbol) { + if (frequencies[symbol] != 0) { + used_symbols[table.symbol_count++] = static_cast(symbol); + } + } + if (table.symbol_count == 0) { + return table; + } + if (table.symbol_count == 1) { + table.root = 0; + table.node_count = 1; + table.tree[0].symbol = used_symbols[0]; + table.node_type[0] = HuffmanBuildNodeType::kLeaf; + table.symbol_to_rank[used_symbols[0]] = 0; + return table; + } + + auto lengths = huffman_build_detail::build_code_lengths( + frequencies, std::span(used_symbols).first(table.symbol_count)); + huffman_build_detail::limit_code_lengths(lengths); + huffman_build_detail::build_fused_tree(lengths, table); + return table; +} + +} // namespace pixie::detail diff --git a/include/pixie/detail/wavelet_partition.h b/include/pixie/detail/wavelet_partition.h index c0a58bb..800ed8f 100644 --- a/include/pixie/detail/wavelet_partition.h +++ b/include/pixie/detail/wavelet_partition.h @@ -23,6 +23,11 @@ #include #endif +#if defined(__aarch64__) && defined(__ARM_NEON) +#define PIXIE_WAVELET_NEON_SUPPORT +#include +#endif + namespace pixie::detail { #if defined(PIXIE_AVX2_SUPPORT) @@ -48,6 +53,61 @@ consteval auto make_wavelet_compaction_table() { alignas(64) inline constexpr auto kWaveletCompactionTable = make_wavelet_compaction_table(); +struct WaveletCompaction16Tables { + std::array, 256> low{}; + std::array, 256> high{}; +}; + +// Adapted from PivCo's p16rev partition tables. One shuffle places the left +// ranks forward at the front and the right ranks in reverse at the back. A +// fixed reverse shuffle then recovers the stable right partition. +consteval WaveletCompaction16Tables make_wavelet_compaction_16_tables() { + WaveletCompaction16Tables tables; + for (std::size_t mask = 0; mask < 256; ++mask) { + std::size_t left = 0; + std::size_t right = 15; + for (std::size_t index = 0; index < 8; ++index) { + if ((mask & (std::size_t{1} << index)) == 0) { + tables.low[mask][left++] = static_cast(index); + } else { + tables.low[mask][right--] = static_cast(index); + } + } + + left = 8; + right = 15; + for (std::size_t index = 0; index < 8; ++index) { + if ((mask & (std::size_t{1} << index)) == 0) { + tables.high[mask][left++] = static_cast(index + 8); + } else { + tables.high[mask][right--] = static_cast(index + 8); + } + } + } + return tables; +} + +alignas(64) inline constexpr auto kWaveletCompaction16Tables = + make_wavelet_compaction_16_tables(); + +inline std::uint16_t wavelet_direction_mask_16(__m128i ranks, + __m128i first_right) { + const __m128i right = + _mm_cmpeq_epi8(_mm_min_epu8(ranks, first_right), first_right); + return static_cast(_mm_movemask_epi8(right)); +} + +inline __m128i compact_wavelet_16(__m128i ranks, std::uint16_t right_mask) { + const std::uint8_t low_mask = static_cast(right_mask); + const std::uint8_t high_mask = static_cast(right_mask >> 8); + const unsigned low_right = std::popcount(low_mask); + const __m128i low = _mm_load_si128(reinterpret_cast( + kWaveletCompaction16Tables.low[low_mask].data())); + const __m128i high = _mm_loadu_si128(reinterpret_cast( + kWaveletCompaction16Tables.high[high_mask].data() + low_right)); + return _mm_shuffle_epi8(ranks, _mm_or_si128(low, high)); +} + inline void store_compacted_bytes(std::uint8_t* destination, std::uint64_t packed, unsigned count) { @@ -190,6 +250,66 @@ std::size_t compact_wavelet_block_avx2(std::span input, #endif +#if defined(PIXIE_WAVELET_NEON_SUPPORT) + +struct NeonWaveletCompactionTables { + std::array, 256> low{}; + std::array, 256> high{}; +}; + +consteval NeonWaveletCompactionTables make_neon_wavelet_compaction_tables() { + NeonWaveletCompactionTables tables; + for (std::size_t mask = 0; mask < 256; ++mask) { + std::size_t left = 0; + std::size_t right = 15; + for (std::size_t index = 0; index < 8; ++index) { + if ((mask & (std::size_t{1} << index)) == 0) { + tables.low[mask][left++] = static_cast(index); + } else { + tables.low[mask][right--] = static_cast(index); + } + } + left = 8; + right = 15; + for (std::size_t index = 0; index < 8; ++index) { + if ((mask & (std::size_t{1} << index)) == 0) { + tables.high[mask][left++] = static_cast(index + 8); + } else { + tables.high[mask][right--] = static_cast(index + 8); + } + } + } + return tables; +} + +alignas(64) inline constexpr auto kNeonWaveletCompactionTables = + make_neon_wavelet_compaction_tables(); + +inline std::uint16_t wavelet_direction_mask_16_neon(uint8x16_t ranks, + uint8x16_t first_right) { + static constexpr std::array kBitWeights = { + 1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128}; + const uint8x16_t bits = + vandq_u8(vcgeq_u8(ranks, first_right), vld1q_u8(kBitWeights.data())); + const std::uint8_t low = vaddv_u8(vget_low_u8(bits)); + const std::uint8_t high = vaddv_u8(vget_high_u8(bits)); + return static_cast(low | (std::uint16_t{high} << 8)); +} + +inline uint8x16_t compact_wavelet_16_neon(uint8x16_t ranks, + std::uint16_t right_mask) { + const std::uint8_t low_mask = static_cast(right_mask); + const std::uint8_t high_mask = static_cast(right_mask >> 8); + const unsigned low_right = std::popcount(low_mask); + const uint8x16_t low = + vld1q_u8(kNeonWaveletCompactionTables.low[low_mask].data()); + const uint8x16_t high = + vld1q_u8(kNeonWaveletCompactionTables.high[high_mask].data() + low_right); + return vqtbl1q_u8(ranks, vorrq_u8(low, high)); +} + +#endif + template std::uint64_t wavelet_direction_mask(std::span ranks, Symbol middle) { @@ -405,4 +525,373 @@ void partition_wavelet_ranks(std::span input, } } +/** Convert byte symbols to in-order leaf ranks. */ +inline void map_wavelet_byte_ranks( + std::span symbols, + std::span ranks, + const std::array& symbol_to_rank, + const std::array& symbol_to_high_rank) { + if (symbols.size() != ranks.size()) { + throw std::invalid_argument("Invalid wavelet rank mapping buffers"); + } + std::size_t offset = 0; +#if defined(PIXIE_WAVELET_NEON_SUPPORT) + if (symbols.size() >= 20) { + uint8x16x4_t table0; + uint8x16x4_t table1; + uint8x16x4_t table2; + uint8x16x4_t table3; + for (std::size_t lane = 0; lane < 4; ++lane) { + table0.val[lane] = vld1q_u8(symbol_to_rank.data() + lane * 16); + table1.val[lane] = vld1q_u8(symbol_to_rank.data() + 64 + lane * 16); + table2.val[lane] = vld1q_u8(symbol_to_rank.data() + 128 + lane * 16); + table3.val[lane] = vld1q_u8(symbol_to_rank.data() + 192 + lane * 16); + } + const uint8x16_t offset64 = vdupq_n_u8(64); + const uint8x16_t offset128 = vdupq_n_u8(128); + const uint8x16_t offset192 = vdupq_n_u8(192); + for (; offset + 20 <= symbols.size(); offset += 20) { + const uint8x16_t input = vld1q_u8(symbols.data() + offset); + std::uint32_t tail; + std::memcpy(&tail, symbols.data() + offset + 16, sizeof(tail)); + uint8x16_t mapped = vqtbl4q_u8(table0, input); + const unsigned rank0 = symbol_to_rank[static_cast(tail)]; + mapped = vqtbx4q_u8(mapped, table1, vsubq_u8(input, offset64)); + const unsigned rank1 = + symbol_to_rank[static_cast(tail >> 8)]; + mapped = vqtbx4q_u8(mapped, table2, vsubq_u8(input, offset128)); + const unsigned rank2 = + symbol_to_rank[static_cast(tail >> 16)]; + mapped = vqtbx4q_u8(mapped, table3, vsubq_u8(input, offset192)); + const unsigned rank3 = + symbol_to_rank[static_cast(tail >> 24)]; + vst1q_u8(ranks.data() + offset, mapped); + const std::uint32_t mapped_tail = + rank0 | (rank1 << 8) | (rank2 << 16) | (rank3 << 24); + std::memcpy(ranks.data() + offset + 16, &mapped_tail, + sizeof(mapped_tail)); + } + } + (void)symbol_to_high_rank; +#else + if constexpr (std::endian::native == std::endian::little) { + for (; offset + 16 <= symbols.size(); offset += 16) { + std::uint64_t low_symbols; + std::uint64_t high_symbols; + std::memcpy(&low_symbols, symbols.data() + offset, sizeof(low_symbols)); + std::memcpy(&high_symbols, symbols.data() + offset + 8, + sizeof(high_symbols)); + + std::array pairs; + for (std::size_t pair = 0; pair < 4; ++pair) { + pairs[pair] = + symbol_to_rank[static_cast(low_symbols)] + + symbol_to_high_rank[static_cast(low_symbols >> 8)]; + low_symbols >>= 16; + pairs[pair + 4] = + symbol_to_rank[static_cast(high_symbols)] + + symbol_to_high_rank[static_cast(high_symbols >> 8)]; + high_symbols >>= 16; + } + std::memcpy(ranks.data() + offset, pairs.data(), 16); + } + } +#endif + for (; offset < symbols.size(); ++offset) { + ranks[offset] = symbol_to_rank[symbols[offset]]; + } +} + +/** Convert one byte block to in-order leaf ranks in place. */ +inline void map_wavelet_byte_ranks( + std::span symbols, + const std::array& symbol_to_rank, + const std::array& symbol_to_high_rank) { + map_wavelet_byte_ranks(std::span(symbols), symbols, + symbol_to_rank, symbol_to_high_rank); +} + +template +std::size_t partition_wavelet_byte_block_scalar( + std::span ranks, + std::uint8_t first_right, + std::span right_output, + PackedBitBuilder& directions) { + std::size_t left = 0; + std::size_t right = 0; + for (std::size_t offset = 0; offset < ranks.size(); offset += 64) { + const std::size_t width = std::min(64, ranks.size() - offset); + std::uint64_t right_mask = 0; + for (std::size_t index = 0; index < width; ++index) { + const std::uint8_t rank = ranks[offset + index]; + const bool goes_right = rank >= first_right; + right_mask |= static_cast(goes_right) << index; + if (goes_right) { + if constexpr (WriteRight) { + right_output[right] = rank; + } + ++right; + } else { + if constexpr (WriteLeft) { + ranks[left] = rank; + } + ++left; + } + } + directions.write_bits(right_mask, width); + } + return left; +} + +#if defined(PIXIE_WAVELET_NEON_SUPPORT) + +// PivCo's p16rev idea expressed with AArch64 TBL. The two halves are loaded +// before either in-place store, and the template removes dead leaf scatters. +template +std::size_t partition_wavelet_byte_block_neon( + std::span ranks, + std::uint8_t first_right, + std::span right_output, + PackedBitBuilder& directions) { + static constexpr std::array kReverse = { + 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0}; + const uint8x16_t reverse = vld1q_u8(kReverse.data()); + const uint8x16_t threshold = vdupq_n_u8(first_right); + std::size_t left = 0; + std::size_t right = 0; + std::size_t offset = 0; + + for (; offset + 32 <= ranks.size(); offset += 32) { + const uint8x16_t low = vld1q_u8(ranks.data() + offset); + const uint8x16_t high = vld1q_u8(ranks.data() + offset + 16); + const std::uint16_t low_mask = + wavelet_direction_mask_16_neon(low, threshold); + const std::uint16_t high_mask = + wavelet_direction_mask_16_neon(high, threshold); + const std::uint32_t right_mask = + low_mask | (static_cast(high_mask) << 16); + const unsigned low_right = std::popcount(low_mask); + const unsigned low_left = 16 - low_right; + const unsigned right_count = std::popcount(right_mask); + directions.write_bits(right_mask, 32); + + if constexpr (WriteLeft || WriteRight) { + const uint8x16_t compacted_low = compact_wavelet_16_neon(low, low_mask); + const uint8x16_t compacted_high = + compact_wavelet_16_neon(high, high_mask); + if constexpr (WriteLeft) { + vst1q_u8(ranks.data() + left, compacted_low); + vst1q_u8(ranks.data() + left + low_left, compacted_high); + } + if constexpr (WriteRight) { + vst1q_u8(right_output.data() + right, + vqtbl1q_u8(compacted_low, reverse)); + vst1q_u8(right_output.data() + right + low_right, + vqtbl1q_u8(compacted_high, reverse)); + } + } + left += 32 - right_count; + right += right_count; + } + for (; offset + 16 <= ranks.size(); offset += 16) { + const uint8x16_t values = vld1q_u8(ranks.data() + offset); + const std::uint16_t right_mask = + wavelet_direction_mask_16_neon(values, threshold); + const unsigned right_count = std::popcount(right_mask); + directions.write_bits(right_mask, 16); + if constexpr (WriteLeft || WriteRight) { + const uint8x16_t compacted = compact_wavelet_16_neon(values, right_mask); + if constexpr (WriteLeft) { + vst1q_u8(ranks.data() + left, compacted); + } + if constexpr (WriteRight) { + vst1q_u8(right_output.data() + right, vqtbl1q_u8(compacted, reverse)); + } + } + left += 16 - right_count; + right += right_count; + } + for (; offset < ranks.size(); ++offset) { + const std::uint8_t rank = ranks[offset]; + const bool goes_right = rank >= first_right; + directions.write_bit(goes_right); + if (goes_right) { + if constexpr (WriteRight) { + right_output[right] = rank; + } + ++right; + } else { + if constexpr (WriteLeft) { + ranks[left] = rank; + } + ++left; + } + } + return left; +} + +#endif + +#if defined(PIXIE_AVX2_SUPPORT) + +// PivCo's non-flat p16rev kernel adapted to append direction masks directly to +// a Pixie node stream. This deliberately keeps the full/right/none compile-time +// specializations so leaf children do not cause dead scatter traffic. +template +std::size_t partition_wavelet_byte_block_x86( + std::span ranks, + std::uint8_t first_right, + std::span right_output, + PackedBitBuilder& directions) { + static constexpr std::array kReverse = { + 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0}; + const __m128i reverse = + _mm_loadu_si128(reinterpret_cast(kReverse.data())); + const __m128i threshold = _mm_set1_epi8(static_cast(first_right)); + std::size_t left = 0; + std::size_t right = 0; + std::size_t offset = 0; + +#if defined(PIXIE_AVX512_SUPPORT) && defined(__AVX512VBMI2__) + const __m512i threshold_512 = + _mm512_set1_epi8(static_cast(first_right - 1)); + for (; offset + 64 <= ranks.size(); offset += 64) { + const __m512i values = _mm512_loadu_si512(ranks.data() + offset); + const __mmask64 right_mask = _mm512_cmpgt_epu8_mask(values, threshold_512); + const unsigned right_count = std::popcount(right_mask); + directions.write_bits(right_mask, 64); + if constexpr (WriteLeft) { + _mm512_mask_compressstoreu_epi8(ranks.data() + left, ~right_mask, values); + } + if constexpr (WriteRight) { + _mm512_mask_compressstoreu_epi8(right_output.data() + right, right_mask, + values); + } + left += 64 - right_count; + right += right_count; + } +#endif + + for (; offset + 32 <= ranks.size(); offset += 32) { + const __m128i low = _mm_loadu_si128( + reinterpret_cast(ranks.data() + offset)); + const __m128i high = _mm_loadu_si128( + reinterpret_cast(ranks.data() + offset + 16)); + const std::uint16_t low_mask = wavelet_direction_mask_16(low, threshold); + const std::uint16_t high_mask = wavelet_direction_mask_16(high, threshold); + const std::uint32_t right_mask = + low_mask | (static_cast(high_mask) << 16); + const unsigned low_right = std::popcount(low_mask); + const unsigned low_left = 16 - low_right; + const unsigned right_count = std::popcount(right_mask); + directions.write_bits(right_mask, 32); + + if constexpr (WriteLeft || WriteRight) { + const __m128i compacted_low = compact_wavelet_16(low, low_mask); + const __m128i compacted_high = compact_wavelet_16(high, high_mask); + if constexpr (WriteLeft) { + _mm_storeu_si128(reinterpret_cast<__m128i*>(ranks.data() + left), + compacted_low); + _mm_storeu_si128( + reinterpret_cast<__m128i*>(ranks.data() + left + low_left), + compacted_high); + } + if constexpr (WriteRight) { + _mm_storeu_si128( + reinterpret_cast<__m128i*>(right_output.data() + right), + _mm_shuffle_epi8(compacted_low, reverse)); + _mm_storeu_si128( + reinterpret_cast<__m128i*>(right_output.data() + right + low_right), + _mm_shuffle_epi8(compacted_high, reverse)); + } + } + left += 32 - right_count; + right += right_count; + } + + for (; offset + 16 <= ranks.size(); offset += 16) { + const __m128i values = _mm_loadu_si128( + reinterpret_cast(ranks.data() + offset)); + const std::uint16_t right_mask = + wavelet_direction_mask_16(values, threshold); + const unsigned right_count = std::popcount(right_mask); + directions.write_bits(right_mask, 16); + if constexpr (WriteLeft || WriteRight) { + const __m128i compacted = compact_wavelet_16(values, right_mask); + if constexpr (WriteLeft) { + _mm_storeu_si128(reinterpret_cast<__m128i*>(ranks.data() + left), + compacted); + } + if constexpr (WriteRight) { + _mm_storeu_si128( + reinterpret_cast<__m128i*>(right_output.data() + right), + _mm_shuffle_epi8(compacted, reverse)); + } + } + left += 16 - right_count; + right += right_count; + } + + for (; offset < ranks.size(); ++offset) { + const std::uint8_t rank = ranks[offset]; + const bool goes_right = rank >= first_right; + directions.write_bit(goes_right); + if (goes_right) { + if constexpr (WriteRight) { + right_output[right] = rank; + } + ++right; + } else { + if constexpr (WriteLeft) { + ranks[left] = rank; + } + ++left; + } + } + return left; +} + +#endif + +/** + * @brief Emit a cache-sized byte block's directions and live child ranks. + * @details The left child is compacted in place and the right child is written + * to disjoint scratch. Node streams are appended bit-contiguously across + * blocks; PivCo's wire framing and per-block bitmap padding are not retained. + */ +inline std::size_t partition_wavelet_byte_block( + std::span ranks, + std::uint8_t first_right, + std::span right_output, + bool write_left, + bool write_right, + PackedBitBuilder& directions) { + if (write_right && right_output.size() < ranks.size()) { + throw std::invalid_argument("Invalid wavelet block partition buffers"); + } + + const auto partition = [&]() { +#if defined(PIXIE_AVX2_SUPPORT) + return partition_wavelet_byte_block_x86( + ranks, first_right, right_output, directions); +#elif defined(PIXIE_WAVELET_NEON_SUPPORT) + return partition_wavelet_byte_block_neon( + ranks, first_right, right_output, directions); +#else + return partition_wavelet_byte_block_scalar( + ranks, first_right, right_output, directions); +#endif + }; + if (write_left) { + return write_right ? partition.template operator()() + : partition.template operator()(); + } + return write_right ? partition.template operator()() + : partition.template operator()(); +} + } // namespace pixie::detail + +#if defined(PIXIE_WAVELET_NEON_SUPPORT) +#undef PIXIE_WAVELET_NEON_SUPPORT +#endif diff --git a/include/pixie/file_archive.h b/include/pixie/file_archive.h index 195af37..40d90aa 100644 --- a/include/pixie/file_archive.h +++ b/include/pixie/file_archive.h @@ -630,12 +630,15 @@ class FileArchiveIndex : public FileArchiveBase>, "File-archive source changed between build passes"); } content_size += chunk.size(); + const auto symbols = std::span( + reinterpret_cast(chunk.data()), + chunk.size()); for (const std::byte byte : chunk) { const std::uint8_t value = std::to_integer(byte); content_hash ^= value; content_hash *= 1099511628211ULL; - emit(value); } + emit(symbols); }); if (content_size != records_[index].content_size || content_hash != content_hashes[index]) { diff --git a/include/pixie/wavelet_tree/index.h b/include/pixie/wavelet_tree/index.h index c2baedf..d73436d 100644 --- a/include/pixie/wavelet_tree/index.h +++ b/include/pixie/wavelet_tree/index.h @@ -1,5 +1,7 @@ #pragma once +#include +#include #include #include #include @@ -39,6 +41,11 @@ class WaveletTreeIndex 'P', 'X', 'W', 'A', 'V', 'E', 'T', '\0'}; static constexpr std::uint32_t kSerializationVersion = 5; static constexpr std::size_t kSerializationHeaderBytes = 24; +#if defined(__APPLE__) && defined(__aarch64__) + static constexpr std::size_t kByteConstructionBlockSize = 16 * 1024; +#else + static constexpr std::size_t kByteConstructionBlockSize = 32 * 1024; +#endif struct PreWaveletNode { node_index_t parent = npos; @@ -220,13 +227,20 @@ class WaveletTreeIndex return; } if (validation == DeserializationValidation::kFull) { + bool has_mapped_symbol = false; for (std::size_t symbol = symbol_begin; symbol < symbol_end; ++symbol) { - if (leaves_[symbol] != node) { + if (leaves_[symbol] == node) { + has_mapped_symbol = true; + } else if (leaves_[symbol] != npos) { throw std::invalid_argument( "Serialized wavelet-tree leaf map disagrees with topology"); } } + if (expected_size != 0 && !has_mapped_symbol) { + throw std::invalid_argument( + "Serialized wavelet-tree leaf is missing from its map"); + } } }; validate_branch(metadata.left_child, current.symbol_begin, @@ -291,6 +305,82 @@ class WaveletTreeIndex return result; } + static std::size_t count_huffman_subtree( + const detail::HuffmanBuildTable& table, + std::int16_t source_node, + std::span symbol_counts, + std::array& subtree_counts) { + const detail::HuffmanBuildNode& node = table.tree[source_node]; + if (node.symbol >= 0) { + return subtree_counts[source_node] = symbol_counts[node.symbol]; + } + const std::size_t left = + count_huffman_subtree(table, node.left, symbol_counts, subtree_counts); + const std::size_t right = + count_huffman_subtree(table, node.right, symbol_counts, subtree_counts); + return subtree_counts[source_node] = left + right; + } + + node_index_t materialize_huffman_node( + const detail::HuffmanBuildTable& table, + std::int16_t source_node, + node_index_t parent, + const std::array& + subtree_counts, + std::vector& nodes) + requires(std::same_as && + std::same_as) + { + const detail::HuffmanBuildNode& source = table.tree[source_node]; + if (source.symbol >= 0) { + leaves_[table.symbol_to_rank[source.symbol]] = parent; + return npos; + } + + const node_index_t result = nodes.size(); + const std::size_t middle = + static_cast(table.split_rank[source_node]) + 1; + nodes.emplace_back(middle); + nodes[result].parent = parent; + nodes[result].left_size = subtree_counts[source.left]; + nodes[result].stream.reserve_bits(subtree_counts[source_node]); + nodes[result].left_child = materialize_huffman_node( + table, source.left, result, subtree_counts, nodes); + nodes[result].right_child = materialize_huffman_node( + table, source.right, result, subtree_counts, nodes); + return result; + } + + void build_huffman_byte_topology(std::span symbol_counts, + std::vector& nodes) + requires(std::same_as && + std::same_as) + { + const detail::HuffmanBuildTable table = + detail::build_huffman_table(symbol_counts); + permutation_.resize(alphabet_size_); + inverse_permutation_.resize(alphabet_size_); + + std::size_t unused_rank = table.symbol_count; + for (std::size_t symbol = 0; symbol < alphabet_size_; ++symbol) { + const std::size_t rank = symbol_counts[symbol] == 0 + ? unused_rank++ + : table.symbol_to_rank[symbol]; + permutation_[symbol] = rank; + inverse_permutation_[rank] = symbol; + } + if (table.symbol_count < 2) { + root_ = npos; + return; + } + + std::array subtree_counts{}; + count_huffman_subtree(table, table.root, symbol_counts, subtree_counts); + nodes.reserve(table.symbol_count - 1); + root_ = materialize_huffman_node(table, table.root, npos, subtree_counts, + nodes); + } + void build_node_streams(node_index_t node, std::span input, std::span output, @@ -313,10 +403,117 @@ class WaveletTreeIndex } } + // PivCo-style block walk: keep the left subsequence in place, put the right + // subsequence in scratch, and omit scatter for leaf children. Separate node + // builders make every block append bit-contiguously without wire padding. + void build_byte_block_streams(node_index_t node, + std::span ranks, + std::span scratch, + std::vector& nodes) + requires(std::same_as && + std::same_as) + { + PreWaveletNode& current = nodes[node]; + const bool write_left = current.left_child != npos; + const bool write_right = current.right_child != npos; + const std::size_t left_size = detail::partition_wavelet_byte_block( + ranks, static_cast(current.middle), scratch, write_left, + write_right, current.stream); + const std::size_t right_size = ranks.size() - left_size; + + if (write_left && left_size != 0) { + build_byte_block_streams(current.left_child, ranks.first(left_size), + scratch.subspan(right_size, left_size), nodes); + } + if (write_right && right_size != 0) { + build_byte_block_streams(current.right_child, scratch.first(right_size), + ranks.first(right_size), nodes); + } + } + template - void build_bit_streams(std::span symbol_counts, - ForEachSymbol& for_each_symbol, - std::vector& nodes) + void build_byte_bit_streams(std::span symbol_counts, + ForEachSymbol& for_each_symbol, + std::vector& nodes) + requires(std::same_as && + std::same_as) + { + std::array symbol_to_rank{}; + std::array symbol_to_high_rank{}; + for (std::size_t symbol = 0; symbol < alphabet_size_; ++symbol) { + const auto rank = static_cast(permutation_[symbol]); + symbol_to_rank[symbol] = rank; + symbol_to_high_rank[symbol] = static_cast(rank) << 8; + } + + const std::size_t block_capacity = std::min( + kByteConstructionBlockSize, std::max(data_size_, 1)); + std::vector block; + std::vector scratch(block_capacity); + block.reserve(block_capacity); + detail::ByteHistogram actual_histogram; + const auto encode_block = [&] { + if (block.empty()) { + return; + } + if (root_ != npos) { + build_byte_block_streams(root_, block, + std::span(scratch).first(block.size()), nodes); + } + block.clear(); + }; + + std::size_t emitted_size = 0; + const auto append_symbols = [&](std::span symbols) { + if (emitted_size > data_size_ || + symbols.size() > data_size_ - emitted_size) { + throw std::invalid_argument( + "Wavelet-tree emitted symbols do not match their counts"); + } + emitted_size += symbols.size(); + actual_histogram.add(std::as_bytes(symbols)); + while (!symbols.empty()) { + const std::size_t copied = + std::min(block_capacity - block.size(), symbols.size()); + const std::size_t destination = block.size(); + block.resize(destination + copied); + detail::map_wavelet_byte_ranks( + symbols.first(copied), + std::span(block).subspan(destination, copied), symbol_to_rank, + symbol_to_high_rank); + symbols = symbols.subspan(copied); + if (block.size() == block_capacity) { + encode_block(); + } + } + }; + for_each_symbol([&](auto&& emitted) { + using Emitted = std::remove_cvref_t; + if constexpr (std::same_as) { + append_symbols(std::span(&emitted, 1)); + } else { + append_symbols(std::span(emitted)); + } + }); + encode_block(); + + const auto actual_counts = actual_histogram.counts(); + if (emitted_size != data_size_ || + !std::ranges::equal(actual_counts.begin(), + actual_counts.begin() + alphabet_size_, + symbol_counts.begin(), symbol_counts.end()) || + std::ranges::any_of(actual_counts.begin() + alphabet_size_, + actual_counts.end(), + [](std::size_t count) { return count != 0; })) { + throw std::invalid_argument( + "Wavelet-tree emitted symbols do not match their counts"); + } + } + + template + void build_generic_bit_streams(std::span symbol_counts, + ForEachSymbol& for_each_symbol, + std::vector& nodes) requires(std::same_as) { std::vector actual_counts(alphabet_size_); @@ -349,6 +546,19 @@ class WaveletTreeIndex build_node_streams(root_, ranks, scratch, nodes); } + template + void build_bit_streams(std::span symbol_counts, + ForEachSymbol& for_each_symbol, + std::vector& nodes) + requires(std::same_as) + { + if constexpr (std::same_as) { + build_byte_bit_streams(symbol_counts, for_each_symbol, nodes); + } else { + build_generic_bit_streams(symbol_counts, for_each_symbol, nodes); + } + } + /** * @brief Recursively copies segment of the original data corresponding to the * node @@ -444,6 +654,20 @@ class WaveletTreeIndex leaves_.assign(alphabet_size_, npos); std::vector nodes; + if constexpr (std::same_as) { + if (build_type == WaveletTreeBuildType::Huffman) { + if (alphabet_size_ != 0) { + build_huffman_byte_topology(symbol_counts, nodes); + } + build_bit_streams(symbol_counts, for_each_symbol, nodes); + nodes_.reserve(nodes.size()); + for (auto& node : nodes) { + nodes_.emplace_back(std::move(node)); + } + return; + } + } + std::vector nodes_structure; if (alphabet_size_ != 0) { nodes.reserve(alphabet_size_); @@ -552,15 +776,33 @@ class WaveletTreeIndex requires(std::same_as) { validate_alphabet_size(alphabet_size); - std::vector counts(alphabet_size); - for (const Symbol symbol : data) { - ++counts[checked_symbol_index(symbol, alphabet_size)]; + std::vector counts; + if constexpr (std::same_as) { + detail::ByteHistogram histogram; + histogram.add(std::as_bytes(data)); + const auto byte_counts = histogram.counts(); + if (std::ranges::any_of(byte_counts.begin() + alphabet_size, + byte_counts.end(), + [](std::size_t count) { return count != 0; })) { + throw std::invalid_argument( + "Wavelet-tree symbol is outside the alphabet"); + } + counts.assign(byte_counts.begin(), byte_counts.begin() + alphabet_size); + } else { + counts.assign(alphabet_size, 0); + for (const Symbol symbol : data) { + ++counts[checked_symbol_index(symbol, alphabet_size)]; + } } build_from_counts( alphabet_size, counts, [&](auto&& emit) { - for (const Symbol symbol : data) { - emit(symbol); + if constexpr (requires { emit(data); }) { + emit(data); + } else { + for (const Symbol symbol : data) { + emit(symbol); + } } }, build_type); @@ -607,6 +849,12 @@ class WaveletTreeIndex return 0; } symbol_index = permutation_[symbol_index]; + if (root_ == npos) [[unlikely]] { + return data_size_ != 0 && symbol_index == 0 ? pos : 0; + } + if (leaves_[symbol_index] == npos) [[unlikely]] { + return 0; + } for (node_index_t current = root_; current != npos;) { const WaveletNode& node = nodes_[current]; if (symbol_index < node.middle) { @@ -630,10 +878,17 @@ class WaveletTreeIndex */ size_t select_impl(Symbol symbol, size_t rank) const { std::size_t symbol_index = static_cast(symbol); - if (symbol_index >= alphabet_size_ || data_size_ == 0) [[unlikely]] { + if (symbol_index >= alphabet_size_ || data_size_ == 0 || rank == 0) + [[unlikely]] { return data_size_; } symbol_index = permutation_[symbol_index]; + if (root_ == npos) [[unlikely]] { + return symbol_index == 0 && rank <= data_size_ ? rank - 1 : data_size_; + } + if (leaves_[symbol_index] == npos) [[unlikely]] { + return data_size_; + } node_index_t current = leaves_[symbol_index]; for (; current != npos; current = nodes_[current].parent) { const WaveletNode& node = nodes_[current]; diff --git a/src/tests/wavelet_tree_tests.cpp b/src/tests/wavelet_tree_tests.cpp index 0f741be..d2a1340 100644 --- a/src/tests/wavelet_tree_tests.cpp +++ b/src/tests/wavelet_tree_tests.cpp @@ -260,6 +260,69 @@ TEST(WaveletTreeTest, ByteBuildHandlesSkewAndPackedBlockTails) { } } +TEST(WaveletTreeTest, HuffmanByteBuildStitchesBlocksAndMarksAbsentSymbols) { + constexpr std::size_t kBlockSize = 32 * 1024; + std::vector data(2 * kBlockSize + 137); + const std::array used = {0, 10, 255}; + std::array, 3> positions; + for (std::size_t index = 0; index < data.size(); ++index) { + const std::size_t choice = (index * 17 + index / 11) % used.size(); + data[index] = used[choice]; + positions[choice].push_back(index); + } + + const pixie::WaveletTree tree( + 256, data, pixie::WaveletTreeBuildType::Huffman); + EXPECT_EQ(tree.get_segment(0, data.size()), data); + for (std::size_t choice = 0; choice < used.size(); ++choice) { + EXPECT_EQ(tree.rank(used[choice], data.size()), positions[choice].size()); + EXPECT_EQ(tree.select(used[choice], positions[choice].size()), + positions[choice].back()); + } + EXPECT_EQ(tree.rank(42, data.size()), 0u); + EXPECT_EQ(tree.select(42, 1), data.size()); + + pixie::VectorOutputSink output; + pixie::BinaryWriter writer(output); + tree.serialize(writer); + writer.finish(); + const std::vector artifact = output.take(); + pixie::BinaryReader reader(artifact); + const auto restored = pixie::WaveletTreeView::deserialize( + reader, pixie::DeserializationValidation::kFull); + EXPECT_TRUE(reader.empty()); + EXPECT_EQ(restored.get_segment(kBlockSize - 3, kBlockSize + 3), + std::vector(data.begin() + kBlockSize - 3, + data.begin() + kBlockSize + 3)); + EXPECT_EQ(restored.rank(42, restored.size()), 0u); +} + +TEST(WaveletTreeTest, HuffmanByteBuildLimitsDeepCodes) { + constexpr std::size_t kUsedSymbols = 20; + std::array frequencies{}; + frequencies[0] = 1; + frequencies[1] = 1; + for (std::size_t symbol = 2; symbol < frequencies.size(); ++symbol) { + frequencies[symbol] = frequencies[symbol - 1] + frequencies[symbol - 2]; + } + + std::vector data; + for (std::size_t symbol = 0; symbol < frequencies.size(); ++symbol) { + data.insert(data.end(), frequencies[symbol], + static_cast(symbol)); + } + std::mt19937_64 random(831); + std::shuffle(data.begin(), data.end(), random); + + const pixie::WaveletTree tree( + 256, data, pixie::WaveletTreeBuildType::Huffman); + EXPECT_EQ(tree.get_segment(0, data.size()), data); + for (std::size_t symbol = 0; symbol < frequencies.size(); ++symbol) { + EXPECT_EQ(tree.rank(static_cast(symbol), data.size()), + frequencies[symbol]); + } +} + TEST(WaveletTreeTest, BuildsFromCountsAndOneStreamedPass) { const std::vector data = {3, 0, 1, 3, 2, 1, 0}; const std::array counts = {2, 2, 1, 2}; @@ -276,6 +339,15 @@ TEST(WaveletTreeTest, BuildsFromCountsAndOneStreamedPass) { EXPECT_EQ(passes, 1u); EXPECT_EQ(tree.get_segment(0, data.size()), data); + const pixie::WaveletTree batched_tree( + 4, counts, + [&](auto&& emit) { + emit(std::span(data).first(3)); + emit(std::span(data).subspan(3)); + }, + pixie::WaveletTreeBuildType::Huffman); + EXPECT_EQ(batched_tree.get_segment(0, data.size()), data); + const std::array wrong_counts = {2, 2, 2, 1}; EXPECT_THROW((pixie::WaveletTree( 4, wrong_counts,