From 94f8ac6bbcc6d8d3079a12f0f000d41dae272cf5 Mon Sep 17 00:00:00 2001 From: "kiloconnect[bot]" <240665456+kiloconnect[bot]@users.noreply.github.com> Date: Sat, 18 Jul 2026 12:22:46 +0000 Subject: [PATCH] Fix wavelet tree get_segment crash for alphabet_size=1 When alphabet_size == 1, the wavelet tree has no internal nodes (root_ == npos). The get_segment_impl method unconditionally called copy_segment_content(root_, ...) which accessed nodes_[npos], causing an out-of-bounds access since nodes_ is empty. rank_impl and select_impl already handled root_ == npos correctly (the traversal loop simply does not execute), but get_segment_impl had no guard for this degenerate case. Fix: add an early return in get_segment_impl when root_ == npos. In this case all symbols are identical, so the result is the single symbol (inverse_permutation_.front()) repeated (end - begin) times. Also add edge-case test coverage: - AlphabetSizeOne: rank, select, and get_segment for single-symbol alphabets (both Standard and Huffman builds) - AlphabetSizeZero: empty alphabet edge case - EmptyData: non-empty alphabet with no data - SymbolWithZeroOccurrences: symbols present in the alphabet but absent from the data - AllSameSymbolLargeAlphabet: one symbol repeated with a large alphabet where most symbols have zero occurrences --- include/pixie/wavelet_tree/index.h | 5 ++ src/tests/wavelet_tree_tests.cpp | 113 +++++++++++++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/include/pixie/wavelet_tree/index.h b/include/pixie/wavelet_tree/index.h index 4c6a8fc..94af9f2 100644 --- a/include/pixie/wavelet_tree/index.h +++ b/include/pixie/wavelet_tree/index.h @@ -382,6 +382,11 @@ class WaveletTreeIndex : public WaveletTreeBase> { if (alphabet_size_ == 0 || data_size_ == 0 || begin >= end) [[unlikely]] { return {}; } + if (root_ == npos) [[unlikely]] { + std::vector result(end - begin, + inverse_permutation_.front()); + return result; + } auto length = static_cast(end - begin); std::vector result(2 * length); copy_segment_content(root_, begin, end, diff --git a/src/tests/wavelet_tree_tests.cpp b/src/tests/wavelet_tree_tests.cpp index 03e900e..58940f9 100644 --- a/src/tests/wavelet_tree_tests.cpp +++ b/src/tests/wavelet_tree_tests.cpp @@ -145,6 +145,119 @@ TEST(WaveletTreeTest, SmokeSegment) { } } +TEST(WaveletTreeTest, AlphabetSizeOne) { + // alphabet_size == 1: tree has no internal nodes (root_ == npos). + // This previously crashed in get_segment_impl. + const std::vector data = {0, 0, 0, 0, 0}; + size_t data_size = 5, alphabet_size = 1; + + for (auto build_type : {pixie::WaveletTreeBuildType::Standard, + pixie::WaveletTreeBuildType::Huffman}) { + WaveletTree wavelet_tree(alphabet_size, data, build_type); + + // rank: all positions have symbol 0 + for (size_t pos = 0; pos <= data_size; pos++) { + EXPECT_EQ(wavelet_tree.rank(0, pos), pos); + } + // rank of out-of-range symbol returns 0 + EXPECT_EQ(wavelet_tree.rank(1, 3), 0); + + // select: the k-th occurrence of symbol 0 is at position k-1 + for (size_t rank = 1; rank <= data_size; rank++) { + EXPECT_EQ(wavelet_tree.select(0, rank), rank - 1); + } + // out-of-range rank returns data_size + EXPECT_EQ(wavelet_tree.select(0, data_size + 1), data_size); + + // get_segment: all elements are 0 + for (size_t begin = 0; begin <= data_size; begin++) { + for (size_t end = begin; end <= data_size; end++) { + auto segment = wavelet_tree.get_segment(begin, end); + EXPECT_EQ(segment.size(), end - begin); + for (size_t i = 0; i < end - begin; i++) { + EXPECT_EQ(segment[i], 0); + } + } + } + } +} + +TEST(WaveletTreeTest, AlphabetSizeZero) { + size_t alphabet_size = 0; + const std::vector data = {}; + WaveletTree wavelet_tree(alphabet_size, data); + + EXPECT_EQ(wavelet_tree.size(), 0); + EXPECT_TRUE(wavelet_tree.empty()); + EXPECT_EQ(wavelet_tree.rank(0, 0), 0); + EXPECT_EQ(wavelet_tree.select(0, 1), 0); + EXPECT_EQ(wavelet_tree.get_segment(0, 0).size(), 0); +} + +TEST(WaveletTreeTest, EmptyData) { + size_t alphabet_size = 4; + const std::vector data = {}; + WaveletTree wavelet_tree(alphabet_size, data); + + EXPECT_EQ(wavelet_tree.size(), 0); + EXPECT_TRUE(wavelet_tree.empty()); + for (uint64_t symb = 0; symb < alphabet_size; symb++) { + EXPECT_EQ(wavelet_tree.rank(symb, 0), 0); + EXPECT_EQ(wavelet_tree.select(symb, 1), 0); + } +} + +TEST(WaveletTreeTest, SymbolWithZeroOccurrences) { + // alphabet_size = 5, but symbol 4 never appears in the data. + const std::vector data = {0, 1, 2, 3, 0, 1, 2, 3}; + size_t data_size = 8, alphabet_size = 5; + + WaveletTree wavelet_tree(alphabet_size, data); + + // rank of absent symbol is 0 at every position + for (size_t pos = 0; pos <= data_size; pos++) { + EXPECT_EQ(wavelet_tree.rank(4, pos), 0); + } + + // select of absent symbol returns data_size + EXPECT_EQ(wavelet_tree.select(4, 1), data_size); + EXPECT_EQ(wavelet_tree.select(4, 2), data_size); + + // get_segment still works correctly + for (size_t begin = 0; begin <= data_size; begin++) { + for (size_t end = begin; end <= data_size; end++) { + auto segment = wavelet_tree.get_segment(begin, end); + EXPECT_EQ(segment.size(), end - begin); + for (size_t i = 0; i < end - begin; i++) { + EXPECT_EQ(segment[i], data[begin + i]); + } + } + } +} + +TEST(WaveletTreeTest, AllSameSymbolLargeAlphabet) { + // All data is symbol 0, but alphabet_size is large. + // Most symbols have 0 occurrences. + const std::vector data = {0, 0, 0, 0, 0, 0, 0, 0}; + size_t data_size = 8, alphabet_size = 256; + + for (auto build_type : {pixie::WaveletTreeBuildType::Standard, + pixie::WaveletTreeBuildType::Huffman}) { + WaveletTree wavelet_tree(alphabet_size, data, build_type); + + EXPECT_EQ(wavelet_tree.rank(0, data_size), data_size); + for (uint64_t symb = 1; symb < alphabet_size; symb++) { + EXPECT_EQ(wavelet_tree.rank(symb, data_size), 0); + EXPECT_EQ(wavelet_tree.select(symb, 1), data_size); + } + + auto segment = wavelet_tree.get_segment(0, data_size); + for (auto s : segment) { + EXPECT_EQ(s, 0); + } + } +} + TEST(WaveletTreeTest, SerializationSmoke) { size_t data_size = 4096, alphabet_size = 100;