diff --git a/BUCK b/BUCK index dbf79a2617ab..b80438676ab7 100644 --- a/BUCK +++ b/BUCK @@ -381,8 +381,7 @@ cpp_library_wrapper(name="rocksdb_lib", srcs=[ "utilities/transactions/write_prepared_txn_db.cc", "utilities/transactions/write_unprepared_txn.cc", "utilities/transactions/write_unprepared_txn_db.cc", - "utilities/trie_index/bitvector.cc", - "utilities/trie_index/louds_trie.cc", + "utilities/trie_index/trie.cc", "utilities/trie_index/trie_index_factory.cc", "utilities/ttl/db_ttl_impl.cc", "utilities/types_util.cc", diff --git a/CMakeLists.txt b/CMakeLists.txt index 7e9e5f5bb032..de185db42451 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1107,8 +1107,7 @@ set(SOURCES utilities/transactions/write_prepared_txn_db.cc utilities/transactions/write_unprepared_txn.cc utilities/transactions/write_unprepared_txn_db.cc - utilities/trie_index/bitvector.cc - utilities/trie_index/louds_trie.cc + utilities/trie_index/trie.cc utilities/trie_index/trie_index_factory.cc utilities/types_util.cc utilities/ttl/db_ttl_impl.cc diff --git a/db_stress_tool/db_stress_tool.cc b/db_stress_tool/db_stress_tool.cc index 38fae20bb56e..8d13ccdf686d 100644 --- a/db_stress_tool/db_stress_tool.cc +++ b/db_stress_tool/db_stress_tool.cc @@ -567,15 +567,6 @@ int db_stress_tool(int argc, char** argv) { FLAGS_atomic_flush = true; } - // Trie UDI uses zero-copy pointers into block data, which is - // incompatible with mmap_read. - if (FLAGS_use_trie_index && FLAGS_mmap_read) { - return ReturnValidationError( - "use_trie_index is incompatible with mmap_read. The trie index uses " - "zero-copy pointers into block data which is unsafe with mmap'd " - "reads."); - } - // TrieIndexFactory requires plain BytewiseComparator, but timestamps use // BytewiseComparator.u64ts. if (FLAGS_use_trie_index && FLAGS_user_timestamp_size > 0) { diff --git a/docs/components/index_factory/index.md b/docs/components/index_factory/index.md index 892fea75b267..71203a1d1565 100644 --- a/docs/components/index_factory/index.md +++ b/docs/components/index_factory/index.md @@ -14,7 +14,7 @@ of, the built-in index selected by `BlockBasedTableOptions::index_type`. | Mode / read selection options | [`include/rocksdb/table.h`](../../../include/rocksdb/table.h), [`include/rocksdb/options.h`](../../../include/rocksdb/options.h) | | Built-in index adapter | `table/block_based/builtin_index_factory.{h,cc}` | | Reader wrapper | `table/block_based/user_defined_index_wrapper.h` | -| Reference implementation | `utilities/trie_index/` (LOUDS trie, registered as `trie_index`) | +| Reference implementation | `utilities/trie_index/` (registered as `trie_index`) | ## SPI shape @@ -40,7 +40,7 @@ tiers an SST is written with and which one reads use by default. | Mode | Standard index written | Custom index written | Default read target | Missing/unusable custom block | |------|------------------------|----------------------|---------------------|-------------------------------| | `kStandardOnly` | full | no (factory ignored entirely) | standard | n/a | -| `kStandardDefault` | full | yes | standard | warn, fall back to standard | +| `kStandardDefault` | full | yes | standard | log, fall back to standard | | `kStandardRequired` | full | yes | standard | hard error at open | | `kCustomDefault` | full | yes | custom | hard error at open | | `kCustomOnly` | stub only | yes | custom | hard error at open | @@ -129,12 +129,14 @@ baking a degraded reader into a long-lived `Rep`. Custom factories cannot be combined with user-defined timestamps unless `kStandardOnly` ignores the factory. -The trie format stores data-block offsets and sizes as 32-bit integers. A block -with either field above `UINT32_MAX` makes construction return `NotSupported`, -including in `kStandardDefault`, where trie construction is optional. This can -fail flushes, compactions, or `SstFileWriter::Finish()`. Total SST size above -4 GiB alone is not the rejection condition. File-size targets do not split L0 -outputs or `SstFileWriter` outputs; bound those outputs or use `kStandardOnly`. +`BlockBasedTableFactory::ValidateOptions()` calls the factory's +`ValidateOptions()` in every mode except `kStandardOnly`. `DB::Open()` and +`CreateColumnFamily()` run that validation; `SetOptions()` and `SstFileWriter` +do not, and the table builder does not repeat it. A factory that cannot build +with some option should therefore also fail `NewBuilder()`. The built-in trie +rejects comparators other than `BytewiseComparator` in both places, so +`DB::Open()` and `CreateColumnFamily()` fail up front, and the other paths fail +at the first table build. ## Cache and memory accounting @@ -150,7 +152,7 @@ outputs or `SstFileWriter` outputs; bound those outputs or use `kStandardOnly`. allocations, excluding the serialized block. A reader whose estimate includes the block can override `MemoryUsageIncludesIndexBlock()` to return true; the wrapper then subtracts the block size from that estimate. `TrieIndexReader` - reports just its auxiliary lookup tables. + reads the block in place, so it reports only its own object. ## Parallel compression @@ -172,9 +174,11 @@ Recommended path: `kStandardDefault` -> `kStandardRequired` -> `kCustomDefault` -> `kCustomOnly`, compacting between steps so existing SSTs pick up the new layout. -Rolling back from `kCustomDefault` or below is immediate: those files always -carry a full standard index, so lowering `index_mode` sends reads back to it -without waiting for compaction. +Rolling back from `kCustomDefault` or below needs no compaction: those files +always carry a full standard index. Index routing is fixed when a table reader +opens, though. Lowering `index_mode` through `SetOptions()` reaches only SSTs +opened afterwards; SSTs already open keep their routing until they are +reopened, e.g. by reopening the DB. `kCustomOnly` is effectively a one-way door. Those SSTs have no usable standard index, so before dropping the factory you must move to `kCustomDefault` and @@ -215,12 +219,16 @@ compact every file, with the factory still attached. `ValidateOptions()` (and, for callers such as `SstFileWriter` that never reach it, the table builder) rejects: -- a custom factory together with user-defined timestamps; +- a custom factory together with user-defined timestamps, except in + `kStandardOnly`; - `kCustomDefault` / `kCustomOnly` with `kTwoLevelIndexSearch` or `partition_filters`; - `kCustomOnly` with `format_version < 6`; - `kCustomDefault` / `kCustomOnly` with no factory. +`ValidateOptions()` alone also rejects anything the factory's own +`ValidateOptions()` rejects, except in `kStandardOnly`. + ## Observability - `SST_USER_DEFINED_INDEX_LOAD_FAIL_COUNT` counts SSTs that could not use their diff --git a/include/rocksdb/index_factory.h b/include/rocksdb/index_factory.h index e87be1bd9741..87c608b4d60a 100644 --- a/include/rocksdb/index_factory.h +++ b/include/rocksdb/index_factory.h @@ -352,6 +352,13 @@ struct IndexFactoryOptions { // // Extends Customizable for string-based construction (CreateFromString), // options serialization, and Name()-based identification. +// +// A factory can override ValidateOptions() to reject options it does not +// support. BlockBasedTableFactory calls it in every index mode except +// kStandardOnly, which ignores the factory. It runs when DB::Open() and +// CreateColumnFamily() validate the options, but not for SetOptions() or +// SstFileWriter, so a factory that cannot build with an option should also +// fail NewBuilder(). // --------------------------------------------------------------------------- class IndexFactory : public Customizable { public: @@ -359,7 +366,7 @@ class IndexFactory : public Customizable { static const char* Type() { return "UserDefinedIndexFactory"; } - // Create an IndexFactory from a string identifier (e.g., "trie"). + // Create an IndexFactory from a string identifier (e.g., "trie_index"). static Status CreateFromString(const ConfigOptions& config_options, const std::string& value, std::shared_ptr* factory); diff --git a/include/rocksdb/table.h b/include/rocksdb/table.h index ad7015bf2988..12725d86ca81 100644 --- a/include/rocksdb/table.h +++ b/include/rocksdb/table.h @@ -660,11 +660,15 @@ struct BlockBasedTableOptions { // // Rollback: // From kCustomDefault: switch to kStandardDefault or kStandardOnly. - // The standard index is fully populated, so SSTs are immediately - // readable. kStandardOnly silently ignores any factory pointer that - // remains set in the config. + // The standard index is fully populated, so no rewrite is needed. + // kStandardOnly silently ignores any factory pointer that remains set + // in the config. // From kCustomOnly: switch to kCustomDefault and compact to rewrite // all SSTs with both indexes before downgrading further. + // An SST's index routing is fixed when its table reader is opened. A + // change made through SetOptions() applies to SSTs opened afterwards; + // SSTs already open keep their routing until they are reopened, e.g. by + // reopening the DB. // // Backup/restore and other Options-serialization reopen paths: // user_defined_index_factory (shared_ptr) does not survive Options diff --git a/src.mk b/src.mk index e4b5ce5fab41..70200c8c3952 100644 --- a/src.mk +++ b/src.mk @@ -352,8 +352,7 @@ LIB_SOURCES = \ utilities/transactions/write_unprepared_txn.cc \ utilities/transactions/write_unprepared_txn_db.cc \ utilities/ttl/db_ttl_impl.cc \ - utilities/trie_index/bitvector.cc \ - utilities/trie_index/louds_trie.cc \ + utilities/trie_index/trie.cc \ utilities/trie_index/trie_index_factory.cc \ utilities/types_util.cc \ utilities/wal_filter.cc \ diff --git a/table/block_based/block_based_table_factory.cc b/table/block_based/block_based_table_factory.cc index 364ec34aa85c..fc7f9bfb9b17 100644 --- a/table/block_based/block_based_table_factory.cc +++ b/table/block_based/block_based_table_factory.cc @@ -421,7 +421,14 @@ static struct BlockBasedTableTypeInfo { OptionTypeInfo::AsCustomSharedPtr( offsetof(struct BlockBasedTableOptions, user_defined_index_factory), - OptionVerificationType::kByNameAllowFromNull)}, + OptionVerificationType::kByNameAllowFromNull) + // BlockBasedTableFactory::ValidateOptions() validates the factory + // only in the index modes that use it. + .SetValidateFunc( + [](const DBOptions& /*db_opts*/, + const ColumnFamilyOptions& /*cf_opts*/, + const std::string& /*name*/, + const void* /*addr*/) { return Status::OK(); })}, {"whole_key_filtering", {offsetof(struct BlockBasedTableOptions, whole_key_filtering), OptionType::kBoolean, OptionVerificationType::kNormal}}, @@ -826,6 +833,18 @@ Status BlockBasedTableFactory::ValidateOptions( "user_defined_index_factory is incompatible with user-defined " "timestamps"); } + // kStandardOnly ignores the factory, so only the other modes are bound by + // its requirements. The user_defined_index_factory option's validate func + // skips the generic Configurable check so that this one decides. + if (table_options_.GetEffectiveIndexMode() != + BlockBasedTableOptions::IndexMode::kStandardOnly && + table_options_.user_defined_index_factory != nullptr) { + Status s = table_options_.user_defined_index_factory->ValidateOptions( + db_opts, cf_opts); + if (!s.ok()) { + return s; + } + } if (table_options_.cache_index_and_filter_blocks && table_options_.no_block_cache) { return Status::InvalidArgument( diff --git a/tools/db_bench_tool.cc b/tools/db_bench_tool.cc index 96ad546b5a7e..ccc7b013f3fe 100644 --- a/tools/db_bench_tool.cc +++ b/tools/db_bench_tool.cc @@ -777,9 +777,7 @@ DEFINE_bool(index_with_first_key, false, "Include first key in the index"); DEFINE_bool(use_trie_index, false, "Use trie-based user-defined index (UDI) for block-based tables. " - "Builds a LOUDS-encoded succinct trie from separator keys, " - "providing space reduction compared to the standard index. " - "Requires BytewiseComparator."); + "Builds a trie from separator keys. Requires BytewiseComparator."); DEFINE_int32(index_mode, static_cast( diff --git a/tools/db_crashtest.py b/tools/db_crashtest.py index 3cade51b40dd..982fc32b9dd9 100644 --- a/tools/db_crashtest.py +++ b/tools/db_crashtest.py @@ -1643,12 +1643,7 @@ def finalize_and_sanitize(src_params): dest_params["open_read_fault_one_in"] = 0 dest_params["sync_fault_injection"] = 0 - # UDI is compatible with all operation types and iteration directions - # except mmap_read. if dest_params.get("use_trie_index") == 1: - # Trie UDI uses zero-copy pointers into block data, which is - # incompatible with mmap_read. - dest_params["mmap_read"] = 0 index_mode = dest_params.get("index_mode") if index_mode in (2, 3): # kCustomDefault/kCustomOnly reject partitioned index @@ -2291,6 +2286,17 @@ def gen_cmd_params(args, unknown_params=None): if v is not None: params[k] = v + # With the default 16 KiB blocks an SST has only a few dozen separators. + # Smaller blocks give the trie more buckets and nodes, and user keys whose + # versions span several blocks. kStandardOnly (index_mode 0) never builds + # the trie, and an explicit --block_size wins. + if ( + getattr(args, "block_size", None) is None + and params.get("use_trie_index") == 1 + and params.get("index_mode") != 0 + ): + params["block_size"] = lambda: random.choice([1024, 4096, 16384]) + # Reserve a small fraction of ordinary crash-test invocations for the # file-scoped WAL recovery state machine. Do this after merging command # line arguments so automatic selection cannot override an explicit opt-out diff --git a/tools/db_crashtest_test.py b/tools/db_crashtest_test.py index 5b75e5693976..5315bdad4e38 100644 --- a/tools/db_crashtest_test.py +++ b/tools/db_crashtest_test.py @@ -961,6 +961,33 @@ def test_finalize_keeps_parallel_compression_with_trie_index(self): self.assertEqual(1, finalized["use_trie_index"]) self.assertEqual(4, finalized["compression_parallel_threads"]) + def test_trie_block_size_draw(self): + db_crashtest = self.load_db_crashtest() + + def block_sizes(**overrides): + # auto_readahead_size=0 keeps the unrelated 1 KiB block shrink off. + args = self.build_mode_args( + test_type="blackbox", use_trie_index=1, auto_readahead_size=0, + **overrides + ) + params = db_crashtest.gen_cmd_params(args) + params["db"] = self.test_tmpdir + return { + db_crashtest.finalize_and_sanitize(dict(params))["block_size"] + for _ in range(60) + } + + # Smaller blocks are drawn for the modes that build the trie. + drawn = block_sizes(index_mode=1) + self.assertTrue(drawn <= {1024, 4096, 16384}, drawn) + self.assertGreater(len(drawn), 1) + # An explicit --block_size is kept. + self.assertEqual({8192}, block_sizes(index_mode=1, block_size=8192)) + # kStandardOnly never builds the trie, so it keeps the default. + self.assertEqual( + {db_crashtest.default_params["block_size"]}, block_sizes(index_mode=0) + ) + def test_finalize_disables_test_batches_snapshots_when_disable_wal(self): db_crashtest = self.load_db_crashtest() params = self.build_params( @@ -1295,7 +1322,7 @@ def test_finalize_trie_index_mode_rules(self): finalized = db_crashtest.finalize_and_sanitize(params) self.assertEqual(index_mode, finalized["index_mode"]) - self.assertEqual(0, finalized["mmap_read"]) + self.assertEqual(1, finalized["mmap_read"]) if index_mode in (2, 3): self.assertIn(finalized["index_type"], [0, 3]) self.assertEqual(0, finalized["partition_filters"]) diff --git a/unreleased_history/behavior_changes/trie_index_format.md b/unreleased_history/behavior_changes/trie_index_format.md new file mode 100644 index 000000000000..7fa463279676 --- /dev/null +++ b/unreleased_history/behavior_changes/trie_index_format.md @@ -0,0 +1 @@ +The experimental trie index (`TrieIndexFactory`) has a new on-disk format that is smaller, faster to search and supports data blocks past 4 GiB. Trie blocks from earlier releases fail with `NotSupported`: in `kStandardDefault` their SSTs fall back to the standard index, so compact there to rewrite them. SSTs written in `kCustomOnly` have no standard index, so rewrite them in another mode before upgrading. Earlier releases cannot read the new trie blocks. The empty user key no longer fails flushes and compactions as a trie separator, and `DB::Open()` and `CreateColumnFamily()` now reject the trie with a non-bytewise comparator outside `kStandardOnly`. diff --git a/utilities/trie_index/bitvector.cc b/utilities/trie_index/bitvector.cc deleted file mode 100644 index d3d383f85137..000000000000 --- a/utilities/trie_index/bitvector.cc +++ /dev/null @@ -1,537 +0,0 @@ -// Copyright (c) Meta Platforms, Inc. and affiliates. -// This source code is licensed under both the GPLv2 (found in the -// COPYING file in the root directory) and Apache 2.0 License -// (found in the LICENSE.Apache file in the root directory). - -#include "utilities/trie_index/bitvector.h" - -#include -#include -#include - -#include "port/lang.h" -#include "util/coding.h" - -namespace ROCKSDB_NAMESPACE { -namespace trie_index { - -// ============================================================================ -// Bitvector serialization format: -// -// [num_bits : uint64_t (fixed 8 bytes)] -// [num_ones : uint64_t (fixed 8 bytes)] -// [words : num_words * 8 bytes, where num_words = ceil(num_bits/64)] -// [rank_lut : num_rank_samples * 4 bytes (uint32_t), padded to 8] -// [select1_hints : num_select1_hints * 4 bytes (uint32_t), padded to 8] -// [select0_hints : num_select0_hints * 4 bytes (uint32_t), padded to 8] -// -// num_rank_samples = num_bits / kBitsPerRankSample + 1 -// (The +1 is for the sentinel entry at the end that stores total popcount.) -// -// num_select1_hints = ceil(num_ones / kOnesPerSelectHint) if num_ones > 0, -// else 0. Analogously for num_select0_hints with (num_bits - num_ones). -// -// Using uint32_t for rank LUT entries halves the LUT memory overhead compared -// to uint64_t. See bitvector.h for why this is safe for trie index bitvectors. -// ============================================================================ - -size_t Bitvector::SerializedSize() const { - // Header: 2 uint64_t (num_bits, num_ones). - // Words: num_words * 8 bytes. - // Rank LUT: num_rank_samples * 4 bytes, padded to 8. - // FindNthOneBit hints: num_select1_hints * 4 bytes, padded to 8. - // FindNthZeroBit hints: num_select0_hints * 4 bytes, padded to 8. - size_t rank_bytes = num_rank_samples_ * sizeof(uint32_t); - size_t rank_padded = (rank_bytes + 7) & ~size_t(7); - size_t s1_bytes = num_select1_hints_ * sizeof(uint32_t); - size_t s1_padded = (s1_bytes + 7) & ~size_t(7); - size_t s0_bytes = num_select0_hints_ * sizeof(uint32_t); - size_t s0_padded = (s0_bytes + 7) & ~size_t(7); - return 2 * sizeof(uint64_t) + num_words_ * sizeof(uint64_t) + rank_padded + - s1_padded + s0_padded; -} - -void Bitvector::EncodeTo(std::string* output) const { - size_t old_size = output->size(); - size_t needed = SerializedSize(); - output->resize(old_size + needed); - char* dst = &(*output)[old_size]; - - // Write header: num_bits, num_ones - memcpy(dst, &num_bits_, sizeof(uint64_t)); - dst += sizeof(uint64_t); - memcpy(dst, &num_ones_, sizeof(uint64_t)); - dst += sizeof(uint64_t); - - // Write words - if (num_words_ > 0) { - memcpy(dst, words_, num_words_ * sizeof(uint64_t)); - dst += num_words_ * sizeof(uint64_t); - } - - // Helper lambda to write a uint32_t array with 8-byte aligned padding. - auto write_u32_array = [&](const uint32_t* arr, uint64_t count) { - if (count > 0) { - size_t bytes = count * sizeof(uint32_t); - memcpy(dst, arr, bytes); - size_t padded = (bytes + 7) & ~size_t(7); - if (padded > bytes) { - memset(dst + bytes, 0, padded - bytes); - } - dst += padded; - } - }; - - // Write rank LUT, select1 hints, select0 hints. - write_u32_array(rank_lut_, num_rank_samples_); - write_u32_array(select1_hints_, num_select1_hints_); - write_u32_array(select0_hints_, num_select0_hints_); -} - -Status Bitvector::InitFromData(const char* data, size_t data_size, - size_t* bytes_consumed) { - const char* start = data; - - // Read header - if (data_size < 2 * sizeof(uint64_t)) { - return Status::Corruption("Bitvector: insufficient data for header"); - } - memcpy(&num_bits_, data, sizeof(uint64_t)); - data += sizeof(uint64_t); - memcpy(&num_ones_, data, sizeof(uint64_t)); - data += sizeof(uint64_t); - data_size -= 2 * sizeof(uint64_t); - - // Validate header fields from untrusted data. - if (num_ones_ > num_bits_) { - return Status::Corruption("Bitvector: num_ones exceeds num_bits"); - } - // The rank LUT uses uint32_t entries, so num_bits must fit in uint32_t. - // This also prevents integer overflow in derived size computations below. - if (num_bits_ > UINT32_MAX) { - return Status::Corruption( - "Bitvector: num_bits exceeds uint32_t range for rank LUT"); - } - - // Compute derived sizes - num_words_ = (num_bits_ + 63) / 64; - num_rank_samples_ = num_bits_ / kBitsPerRankSample + 1; - // Select hints: ceil(count / kOnesPerSelectHint) entries for ones/zeros. - // If there are 0 ones (or 0 zeros), there are 0 hints. - num_select1_hints_ = - (num_ones_ > 0) ? (num_ones_ - 1) / kOnesPerSelectHint + 1 : 0; - uint64_t num_zeros = num_bits_ - num_ones_; - num_select0_hints_ = - (num_zeros > 0) ? (num_zeros - 1) / kOnesPerSelectHint + 1 : 0; - - // Read words - size_t words_bytes = num_words_ * sizeof(uint64_t); - if (data_size < words_bytes) { - return Status::Corruption("Bitvector: insufficient data for words"); - } - // The data pointer must be 8-byte aligned for safe uint64_t access. - // This is guaranteed by the serialization format (header is 16 bytes, - // and the overall buffer comes from aligned block cache allocations). - if (reinterpret_cast(data) % alignof(uint64_t) != 0) { - return Status::Corruption("Bitvector: words data not 8-byte aligned"); - } - words_ = reinterpret_cast(data); - data += words_bytes; - data_size -= words_bytes; - - // Helper lambda: read a padded uint32_t array from the stream. - auto read_u32_array = [&](const uint32_t** out, uint64_t count, - const char* name) -> Status { - if (count == 0) { - *out = nullptr; - return Status::OK(); - } - size_t bytes = count * sizeof(uint32_t); - size_t padded = (bytes + 7) & ~size_t(7); - if (data_size < padded) { - return Status::Corruption( - std::string("Bitvector: insufficient data for ") + name); - } - if (reinterpret_cast(data) % alignof(uint32_t) != 0) { - return Status::Corruption(std::string("Bitvector: ") + name + - " not 4-byte aligned"); - } - *out = reinterpret_cast(data); - data += padded; - data_size -= padded; - return Status::OK(); - }; - - Status s; - s = read_u32_array(&rank_lut_, num_rank_samples_, "rank LUT"); - if (!s.ok()) { - return s; - } - s = read_u32_array(&select1_hints_, num_select1_hints_, "select1 hints"); - if (!s.ok()) { - return s; - } - s = read_u32_array(&select0_hints_, num_select0_hints_, "select0 hints"); - if (!s.ok()) { - return s; - } - - // Validate select hint values. Each hint is a rank LUT index used to - // narrow the FindNthOneBit/FindNthZeroBit search range. An out-of-bounds - // hint would cause rank_lut_[hint] to read past the LUT array. - for (uint64_t i = 0; i < num_select1_hints_; i++) { - if (select1_hints_[i] >= num_rank_samples_) { - return Status::Corruption( - "Bitvector: select1 hint value exceeds rank LUT size"); - } - } - for (uint64_t i = 0; i < num_select0_hints_; i++) { - if (select0_hints_[i] >= num_rank_samples_) { - return Status::Corruption( - "Bitvector: select0 hint value exceeds rank LUT size"); - } - } - - *bytes_consumed = static_cast(data - start); - - return Status::OK(); -} - -void Bitvector::BuildFrom(const BitvectorBuilder& builder) { - num_bits_ = builder.NumBits(); - num_words_ = (num_bits_ + 63) / 64; - num_rank_samples_ = num_bits_ / kBitsPerRankSample + 1; - - // We need to build rank LUT first to know num_ones_, then compute hint - // counts. Two-pass: first allocate for words + rank LUT, build rank LUT - // to get num_ones_, then resize to include select hints. - size_t words_bytes = num_words_ * sizeof(uint64_t); - size_t rank_bytes = num_rank_samples_ * sizeof(uint32_t); - - // Phase 1: allocate words + rank LUT. - owned_data_.resize(words_bytes + rank_bytes, '\0'); - if (num_words_ > 0) { - memcpy(&owned_data_[0], builder.Words().data(), words_bytes); - } - words_ = reinterpret_cast(owned_data_.data()); - rank_lut_ = - reinterpret_cast(owned_data_.data() + words_bytes); - assert(reinterpret_cast(rank_lut_) % alignof(uint32_t) == 0); - - BuildRankLUT(); - - // Phase 2: compute hint counts and reallocate to include them. - num_select1_hints_ = - (num_ones_ > 0) ? (num_ones_ - 1) / kOnesPerSelectHint + 1 : 0; - uint64_t num_zeros = num_bits_ - num_ones_; - num_select0_hints_ = - (num_zeros > 0) ? (num_zeros - 1) / kOnesPerSelectHint + 1 : 0; - - size_t s1_bytes = num_select1_hints_ * sizeof(uint32_t); - size_t s0_bytes = num_select0_hints_ * sizeof(uint32_t); - owned_data_.resize(words_bytes + rank_bytes + s1_bytes + s0_bytes, '\0'); - - // Recompute all pointers after resize (string may have reallocated). - RecomputePointers(); - assert(reinterpret_cast(rank_lut_) % alignof(uint32_t) == 0); - - BuildSelectHints(); -} - -void Bitvector::BuildRankLUT() { - // The rank LUT is mutable during construction, so we const_cast here. - // After BuildRankLUT() completes, rank_lut_ is treated as read-only. - uint32_t* lut = const_cast(rank_lut_); - - // Guard: uint32_t can hold values up to ~4 billion. The maximum cumulative - // rank is num_ones_ (at most num_bits_), so asserting num_bits_ fits in - // uint32_t guarantees all rank LUT entries fit. - assert(num_bits_ <= UINT32_MAX); - - uint64_t cumulative = 0; - for (uint64_t s = 0; s < num_rank_samples_; s++) { - lut[s] = static_cast(cumulative); - // Sum popcount for the kWordsPerRankSample words in this sample interval. - uint64_t word_start = s * kWordsPerRankSample; - uint64_t word_end = std::min(word_start + kWordsPerRankSample, num_words_); - for (uint64_t w = word_start; w < word_end; w++) { - cumulative += Popcount(words_[w]); - } - } - // Total number of 1-bits is the cumulative count after all words. - num_ones_ = cumulative; -} - -void Bitvector::BuildSelectHints() { - // Build select1 hints: for each k, select1_hints_[k] is the rank LUT - // index of the sample interval containing the (k * kOnesPerSelectHint)-th - // 1-bit. - if (num_select1_hints_ > 0) { - uint32_t* hints = const_cast(select1_hints_); - uint64_t sample_idx = 0; - for (uint64_t k = 0; k < num_select1_hints_; k++) { - uint64_t target = k * kOnesPerSelectHint; - // Advance sample_idx until rank_lut_[sample_idx+1] > target. - while (sample_idx + 1 < num_rank_samples_ && - rank_lut_[sample_idx + 1] <= target) { - sample_idx++; - } - hints[k] = static_cast(sample_idx); - } - } - - // Build select0 hints analogously using zero counts. - if (num_select0_hints_ > 0) { - uint32_t* hints = const_cast(select0_hints_); - uint64_t sample_idx = 0; - for (uint64_t k = 0; k < num_select0_hints_; k++) { - uint64_t target = k * kOnesPerSelectHint; - while (sample_idx + 1 < num_rank_samples_ && - ((sample_idx + 1) * kBitsPerRankSample - - rank_lut_[sample_idx + 1]) <= target) { - sample_idx++; - } - hints[k] = static_cast(sample_idx); - } - } -} - -// FindNthOneBit is now inline in bitvector.h for hot-path performance. - -uint64_t Bitvector::FindNthZeroBit(uint64_t i) const { - uint64_t num_zeros = num_bits_ - num_ones_; - if (i >= num_zeros) { - return num_bits_; - } - - // Use select0 hints to narrow the search range. - uint64_t lo; - uint64_t hi; - if (num_select0_hints_ > 0) { - uint64_t hint_idx = i / kOnesPerSelectHint; - lo = select0_hints_[hint_idx]; - hi = (hint_idx + 1 < num_select0_hints_) ? select0_hints_[hint_idx + 1] - : num_rank_samples_ - 1; - } else { - lo = 0; - hi = num_rank_samples_ - 1; - } - - // Linear scan within the narrowed range. - while (lo < hi) { - uint64_t zeros_next = (lo + 1) * kBitsPerRankSample - rank_lut_[lo + 1]; - if (zeros_next <= i) { - lo++; - } else { - break; - } - } - - uint64_t zeros_at_lo = lo * kBitsPerRankSample - rank_lut_[lo]; - uint64_t remaining = i - zeros_at_lo; - uint64_t word_start = lo * kWordsPerRankSample; - uint64_t word_end = std::min(word_start + kWordsPerRankSample, num_words_); - - for (uint64_t w = word_start; w < word_end; w++) { - // For the last word, we may have fewer than 64 valid bits. We need to - // mask off the padding bits so they don't count as zeros. - uint64_t valid_bits = 64; - if (w == num_words_ - 1 && num_bits_ % 64 != 0) { - valid_bits = num_bits_ % 64; - } - uint64_t zeros_in_word = valid_bits - Popcount(words_[w]); - if (remaining < zeros_in_word) { - // The target 0-bit is in this word. Invert to get 1-bits at zero - // positions, mask off padding, then find the n-th set bit in the word. - uint64_t word = ~words_[w]; - if (valid_bits < 64) { - word &= (uint64_t(1) << valid_bits) - 1; - } - return w * 64 + FindNthSetBitInWord(word, remaining); - } - remaining -= zeros_in_word; - } - - assert(false); - return num_bits_; -} - -uint64_t Bitvector::DistanceToNextSetBit(uint64_t pos) const { - assert(pos < num_bits_); - assert(GetBit(pos)); // pos must be a set bit. - - // Find the next set bit after pos (exclusive). - uint64_t next = NextSetBit(pos + 1); - return next - pos; -} - -// ============================================================================ -// EliasFano implementation -// ============================================================================ - -void EliasFano::BuildFrom(const uint64_t* values, uint64_t count, - uint64_t universe) { - count_ = count; - universe_ = universe; - - if (count == 0) { - low_bits_ = 0; - low_mask_ = 0; - low_words_ = nullptr; - num_low_words_ = 0; - // Build an empty high bitvector so it serializes/deserializes correctly. - // Without this, high_bv_ stays default-constructed (num_rank_samples_=0), - // but InitFromData computes num_rank_samples_=1 for 0-bit bitvectors, - // causing a size mismatch on deserialization. - BitvectorBuilder empty_builder; - high_bv_.BuildFrom(empty_builder); - return; - } - - // Compute low_bits = floor(log2(universe / count)). - // When universe <= count, low_bits = 0 (all bits are high). - // FloorLog2 returns floor(log2(x)) for x >= 1. - if (universe > count) { - low_bits_ = static_cast(FloorLog2(universe / count)); - } else { - low_bits_ = 0; - } - low_mask_ = - (low_bits_ < 64) ? ((uint64_t(1) << low_bits_) - 1) : ~uint64_t(0); - - // Build high-bits bitvector. - // The high part of each value is (value >> low_bits_). The bitvector - // has a 1-bit at position high[i] + i for each element i, with 0-bits - // filling the gaps. Total length = max_high + count. - uint64_t max_high = - (count > 0 && low_bits_ < 64) ? (values[count - 1] >> low_bits_) : 0; - uint64_t high_len = max_high + count; - - BitvectorBuilder high_builder; - high_builder.Reserve(high_len + 1); - - uint64_t prev_high = 0; - for (uint64_t i = 0; i < count; i++) { - assert(i == 0 || values[i] >= values[i - 1]); // Must be monotone. - uint64_t high = (low_bits_ < 64) ? (values[i] >> low_bits_) : 0; - // Append (high - prev_high) zeros followed by a 1. - assert(high >= prev_high); - high_builder.AppendMultiple(false, high - prev_high); - high_builder.Append(true); - prev_high = high; - } - - high_bv_.BuildFrom(high_builder); - - // Build packed low-bits array. - uint64_t total_low_bits = count * low_bits_; - num_low_words_ = (total_low_bits + 63) / 64; - owned_low_data_.resize(num_low_words_ * sizeof(uint64_t), '\0'); - uint64_t* low_buf = reinterpret_cast(&owned_low_data_[0]); - low_words_ = low_buf; - - // Pack low bits: value i's low bits go at bit position i * low_bits_. - for (uint64_t i = 0; i < count; i++) { - uint64_t low = values[i] & low_mask_; - uint64_t bit_pos = i * low_bits_; - uint64_t word_idx = bit_pos / 64; - uint64_t bit_idx = bit_pos % 64; - low_buf[word_idx] |= low << bit_idx; - // Handle word boundary crossing. - if (bit_idx > 0 && bit_idx + low_bits_ > 64) { - low_buf[word_idx + 1] |= low >> (64 - bit_idx); - } - } -} - -size_t EliasFano::SerializedSize() const { - // 3 uint64_t header + high bitvector + low words (8-byte aligned). - size_t low_bytes = num_low_words_ * sizeof(uint64_t); - return 3 * sizeof(uint64_t) + high_bv_.SerializedSize() + low_bytes; -} - -void EliasFano::EncodeTo(std::string* output) const { - size_t old_size = output->size(); - // Header: count, universe, low_bits - PutFixed64(output, count_); - PutFixed64(output, universe_); - PutFixed64(output, low_bits_); - - // High bitvector - high_bv_.EncodeTo(output); - - // Low words - size_t low_bytes = num_low_words_ * sizeof(uint64_t); - if (low_bytes > 0) { - size_t cur = output->size(); - output->resize(cur + low_bytes); - memcpy(&(*output)[cur], low_words_, low_bytes); - } - (void)old_size; -} - -Status EliasFano::InitFromData(const char* data, size_t data_size, - size_t* bytes_consumed) { - const char* start = data; - - // Read header: count, universe, low_bits. - if (data_size < 3 * sizeof(uint64_t)) { - return Status::Corruption("EliasFano: insufficient data for header"); - } - memcpy(&count_, data, sizeof(uint64_t)); - data += sizeof(uint64_t); - memcpy(&universe_, data, sizeof(uint64_t)); - data += sizeof(uint64_t); - memcpy(&low_bits_, data, sizeof(uint64_t)); - data += sizeof(uint64_t); - data_size -= 3 * sizeof(uint64_t); - - // Validate header fields from untrusted data. - if (low_bits_ > 63) { - return Status::Corruption("EliasFano: low_bits exceeds 63"); - } - // Upper-bound on count_: a sequence with > 2^30 elements would require - // tens of GB of data. This prevents integer overflow in subsequent - // count_ * low_bits_ and count_ * sizeof(...) arithmetic. Block checksums - // are the first line of defense; this is defense-in-depth for the - // deserialization path. - static constexpr uint64_t kMaxReasonableCount = uint64_t(1) << 30; - if (count_ > kMaxReasonableCount) { - return Status::Corruption("EliasFano: count exceeds reasonable limit"); - } - low_mask_ = - (low_bits_ < 64) ? ((uint64_t(1) << low_bits_) - 1) : ~uint64_t(0); - - // Read high bitvector. - size_t consumed = 0; - Status s = high_bv_.InitFromData(data, data_size, &consumed); - if (!s.ok()) { - return s; - } - data += consumed; - data_size -= consumed; - - // Read low words. The multiplication is safe from overflow because - // count_ <= 2^30 and low_bits_ <= 63, so count_ * low_bits_ <= 2^36. - uint64_t total_low_bits = count_ * low_bits_; - num_low_words_ = (total_low_bits + 63) / 64; - size_t low_bytes = num_low_words_ * sizeof(uint64_t); - if (data_size < low_bytes) { - return Status::Corruption("EliasFano: insufficient data for low words"); - } - if (low_bytes > 0) { - if (reinterpret_cast(data) % alignof(uint64_t) != 0) { - return Status::Corruption("EliasFano: low words not 8-byte aligned"); - } - low_words_ = reinterpret_cast(data); - } else { - low_words_ = nullptr; - } - data += low_bytes; - - *bytes_consumed = static_cast(data - start); - return Status::OK(); -} - -} // namespace trie_index -} // namespace ROCKSDB_NAMESPACE diff --git a/utilities/trie_index/bitvector.h b/utilities/trie_index/bitvector.h deleted file mode 100644 index 0372de67cd78..000000000000 --- a/utilities/trie_index/bitvector.h +++ /dev/null @@ -1,681 +0,0 @@ -// Copyright (c) Meta Platforms, Inc. and affiliates. -// This source code is licensed under both the GPLv2 (found in the -// COPYING file in the root directory) and Apache 2.0 License -// (found in the LICENSE.Apache file in the root directory). -// -// ***************************************************************** -// EXPERIMENTAL - subject to change while under development -// ***************************************************************** - -#pragma once - -#include -#include -#include -#include -#include -#include - -#include "rocksdb/slice.h" -#include "rocksdb/status.h" -#include "util/math.h" - -namespace ROCKSDB_NAMESPACE { -namespace trie_index { - -// ============================================================================ -// Bit-vector with O(1) rank and select operations. -// -// This is the foundational data structure for the LOUDS-based succinct trie. -// It stores a sequence of bits and supports: -// - rank1(pos): Count of 1-bits in positions [0, pos) -// - select1(i): Position of the i-th 1-bit (0-indexed) -// - rank0(pos): Count of 0-bits in positions [0, pos) -// - select0(i): Position of the i-th 0-bit (0-indexed) -// -// Implementation uses a two-level lookup table for rank (sampling every -// kBitsPerRankSample bits) and select hints + linear scan for select. Both -// rank and select operations are O(1). Select uses precomputed hint arrays -// that narrow the search to O(1) rank samples, then a word-level popcount -// scan within that interval. -// -// Memory layout (serialized): -// [num_bits: uint64_t] -// [num_ones: uint64_t] -// [raw bits: ceil(num_bits/64) uint64_t words, 64-bit aligned] -// [rank LUT: (num_bits/kBitsPerRankSample + 1) uint32_t entries, padded to 8] -// [select1 hints: uint32_t entries, padded to 8] -// [select0 hints: uint32_t entries, padded to 8] -// -// The rank LUT uses uint32_t entries (not uint64_t), which halves the LUT -// memory overhead. This is safe because the maximum cumulative popcount -// equals num_bits, and uint32_t can hold values up to ~4 billion. This -// limits individual bitvectors to ~4 billion bits, which is far beyond -// any realistic trie index: the largest bitvector (d_labels_) uses 256 -// bits per dense trie node, so the limit corresponds to ~16 million dense -// nodes. A typical SST file has 16K-64K data blocks, producing a trie -// with at most a few hundred thousand nodes. An assertion in BuildRankLUT -// guards against overflow. -// -// The rank LUT stores cumulative popcount at every kBitsPerRankSample-bit -// boundary. For positions between boundaries, we compute the remaining -// popcount using hardware popcount on the intermediate words. -// ============================================================================ - -// Number of bits per rank lookup table entry. Must be a power of 2 and a -// multiple of 64. 256 gives a good balance between LUT size overhead (~12.5% -// of bitvector size with uint32_t entries) and the number of popcounts needed -// per rank query (at most 4 words). Benchmarking showed 256 is ~20% faster -// than 512 for trie Seek due to fewer popcount iterations per Rank1 call. -inline constexpr uint64_t kBitsPerRankSample = 256; -inline constexpr uint64_t kWordsPerRankSample = kBitsPerRankSample / 64; - -// Number of 1-bits (or 0-bits) between select hint entries. Each hint stores -// the rank LUT index where the k-th group of ones/zeros begins, narrowing the -// Select search to a linear scan of 1-2 rank samples, making it O(1). -// 256 matches kBitsPerRankSample so hints map directly to rank LUT entries. -inline constexpr uint64_t kOnesPerSelectHint = 256; - -// Portable popcount using RocksDB's BitsSetToOne, which handles MSVC, -// GCC, and Clang with hardware POPCNT when available. -inline uint64_t Popcount(uint64_t x) { - return static_cast(BitsSetToOne(x)); -} - -// Count trailing zeros. Returns 64 if x == 0. Uses RocksDB's -// CountTrailingZeroBits for portability (MSVC + GCC/Clang). -inline uint64_t Ctz(uint64_t x) { - return x == 0 ? 64 : static_cast(CountTrailingZeroBits(x)); -} - -// Count trailing zeros for a known-nonzero value. Avoids the zero-check -// branch when the caller has already verified x != 0. -inline uint64_t CtzNonZero(uint64_t x) { - assert(x != 0); - return static_cast(CountTrailingZeroBits(x)); -} - -// Select the i-th set bit (0-indexed) within a 64-bit word. -// Precondition: i < Popcount(word). -#if defined(__BMI2__) && defined(__x86_64__) -// Hardware-accelerated using BMI2 PDEP: deposits bit i among set bits of word, -// then counts trailing zeros to find its position. Single PDEP + CTZ. -inline uint64_t FindNthSetBitInWord(uint64_t word, uint64_t i) { - return __builtin_ctzll(_pdep_u64(1ULL << i, word)); -} -#else -// Popcount-based binary search: O(log 64) = 6 steps (5 popcounts + 1 bit test). -// Matches the SuRF reference implementation's select64_popcount_search(). -inline uint64_t FindNthSetBitInWord(uint64_t word, uint64_t i) { - // Binary search: narrow down which 32-bit half, then 16, 8, 4, 2, 1. - uint64_t pos = 0; - uint64_t pc; - - pc = Popcount(word & 0xFFFFFFFFULL); - if (i >= pc) { - word >>= 32; - pos += 32; - i -= pc; - } - pc = Popcount(word & 0xFFFFULL); - if (i >= pc) { - word >>= 16; - pos += 16; - i -= pc; - } - pc = Popcount(word & 0xFFULL); - if (i >= pc) { - word >>= 8; - pos += 8; - i -= pc; - } - pc = Popcount(word & 0xFULL); - if (i >= pc) { - word >>= 4; - pos += 4; - i -= pc; - } - pc = Popcount(word & 0x3ULL); - if (i >= pc) { - word >>= 2; - pos += 2; - i -= pc; - } - pc = word & 1; - if (i >= pc) { - pos += 1; - } - return pos; -} -#endif - -// ============================================================================ -// BitvectorBuilder: Builds a bitvector incrementally, bit by bit. -// ============================================================================ -class BitvectorBuilder { - public: - BitvectorBuilder() : num_bits_(0) {} - - // Append a single bit (0 or 1). - void Append(bool bit) { - if (num_bits_ % 64 == 0) { - words_.push_back(0); - } - if (bit) { - words_.back() |= (uint64_t(1) << (num_bits_ % 64)); - } - num_bits_++; - } - - // Append a full 64-bit word of `nbits` bits (1..64). The bits are taken - // from the low `nbits` positions of `word` in LSB-first order. This is - // used by the LOUDS-Dense builder to emit 256-bit label bitmaps as 4 - // word-level operations instead of 256 individual Append() calls. - // Precondition: num_bits_ must be 64-bit aligned (i.e., num_bits_ % 64 == 0). - void AppendWord(uint64_t word, uint64_t nbits) { - assert(nbits > 0 && nbits <= 64); - assert(num_bits_ % 64 == 0); // Must be word-aligned. - words_.push_back(word); - num_bits_ += nbits; - } - - // Append `count` copies of `bit`. Optimized to operate at word granularity - // when possible, which is significantly faster than the bit-by-bit loop for - // large counts (e.g., appending 256 zeros for an empty dense node). - void AppendMultiple(bool bit, uint64_t count) { - if (count == 0) { - return; - } - - // Fill partial word at the end of the current buffer. - uint64_t partial = num_bits_ % 64; - if (partial > 0) { - uint64_t fill = std::min(count, 64 - partial); - if (bit) { - // Set bits [partial, partial + fill) in the last word. - uint64_t mask = - ((fill == 64) ? ~uint64_t(0) : ((uint64_t(1) << fill) - 1)) - << partial; - words_.back() |= mask; - } - // For bit==0, no action needed (words are zero-initialized). - num_bits_ += fill; - count -= fill; - } - - // Append full 64-bit words. - uint64_t full_word = bit ? ~uint64_t(0) : uint64_t(0); - while (count >= 64) { - words_.push_back(full_word); - num_bits_ += 64; - count -= 64; - } - - // Append remaining bits (< 64). - if (count > 0) { - if (bit) { - words_.push_back((uint64_t(1) << count) - 1); - } else { - words_.push_back(0); - } - num_bits_ += count; - } - } - - uint64_t NumBits() const { return num_bits_; } - - // Pre-allocate capacity for at least `num_bits` bits. Avoids repeated - // reallocations when the final size is known or can be estimated. - void Reserve(uint64_t num_bits) { words_.reserve((num_bits + 63) / 64); } - - // Access a specific bit. For testing/debugging only. - bool GetBit(uint64_t pos) const { - assert(pos < num_bits_); - return (words_[pos / 64] >> (pos % 64)) & 1; - } - - // Return the underlying word array. - const std::vector& Words() const { return words_; } - - private: - std::vector words_; - uint64_t num_bits_; -}; - -// ============================================================================ -// Bitvector: Immutable bitvector with O(1) rank and O(1) select. -// -// Created from serialized data (e.g., read from an SST file meta-block) or -// from a BitvectorBuilder. The bitvector does NOT own the underlying memory -// when created from a Slice -- the caller must ensure the data outlives this -// object. -// -// Select acceleration: select hints store the rank LUT index where every -// kOnesPerSelectHint-th 1-bit (or 0-bit) lives. This narrows the search -// in FindNthOneBit/FindNthZeroBit to a linear scan of O(1) rank samples, -// making Select O(1). The hint arrays add ~12.5% space overhead per bitvector. -// ============================================================================ -class Bitvector { - public: - Bitvector() - : words_(nullptr), - rank_lut_(nullptr), - select1_hints_(nullptr), - select0_hints_(nullptr), - num_bits_(0), - num_ones_(0), - num_words_(0), - num_rank_samples_(0), - num_select1_hints_(0), - num_select0_hints_(0) {} - - // Bitvector contains raw pointers (words_, rank_lut_, select hints) that - // may point into owned_data_ or into external memory (InitFromData). - // Copying would create dangling pointers, so copy is deleted. Move is safe - // only when the pointers point into owned_data_ (BuildFrom case) because - // std::string's move constructor preserves the buffer address for - // SSO-exceeding strings. For the InitFromData case, the pointers reference - // external memory and are unaffected by moving owned_data_ (which is empty). - ~Bitvector() = default; - - Bitvector(const Bitvector&) = delete; - Bitvector& operator=(const Bitvector&) = delete; - // Move constructor delegates to default ctor + move assignment to avoid - // duplicating the field-copy + recompute + zero-other logic. - Bitvector(Bitvector&& other) noexcept : Bitvector() { - *this = std::move(other); - } - Bitvector& operator=(Bitvector&& other) noexcept { - if (this != &other) { - words_ = other.words_; - rank_lut_ = other.rank_lut_; - select1_hints_ = other.select1_hints_; - select0_hints_ = other.select0_hints_; - num_bits_ = other.num_bits_; - num_ones_ = other.num_ones_; - num_words_ = other.num_words_; - num_rank_samples_ = other.num_rank_samples_; - num_select1_hints_ = other.num_select1_hints_; - num_select0_hints_ = other.num_select0_hints_; - owned_data_ = std::move(other.owned_data_); - // If this bitvector owns its data, the pointers must be re-seated into - // our owned_data_ buffer. std::string move may or may not preserve the - // buffer address (SSO optimization), so always re-seat. - if (!owned_data_.empty()) { - RecomputePointers(); - } - other.words_ = nullptr; - other.rank_lut_ = nullptr; - other.select1_hints_ = nullptr; - other.select0_hints_ = nullptr; - other.num_bits_ = 0; - other.num_ones_ = 0; - other.num_words_ = 0; - other.num_rank_samples_ = 0; - other.num_select1_hints_ = 0; - other.num_select0_hints_ = 0; - } - return *this; - } - - // Initialize from serialized data. The data pointer must remain valid for - // the lifetime of this object. On success, sets `*bytes_consumed` to the - // number of bytes read from the input. Returns Status::OK() on success, - // or Status::Corruption() if the data is malformed. - Status InitFromData(const char* data, size_t data_size, - size_t* bytes_consumed); - - // Append serialized representation to `output`. - void EncodeTo(std::string* output) const; - - // Build from a BitvectorBuilder. This allocates owned memory. - void BuildFrom(const BitvectorBuilder& builder); - - // ---- Core Operations ---- - - // Get the bit at position `pos`. - inline bool GetBit(uint64_t pos) const { - assert(pos < num_bits_); - return (words_[pos / 64] >> (pos % 64)) & 1; - } - - // rank1(pos): Count of 1-bits in positions [0, pos). - // pos can be in [0, num_bits_]. - inline uint64_t Rank1(uint64_t pos) const { - assert(pos <= num_bits_); - // Which rank sample does `pos` fall into? - uint64_t sample_idx = pos / kBitsPerRankSample; - uint64_t sample_rank = rank_lut_[sample_idx]; - // Count remaining 1-bits in words between the sample boundary and `pos`. - // Unrolled: kWordsPerRankSample is 4, so count is 0-3. The switch - // eliminates loop control overhead and guarantees no branch mispredictions - // from the loop counter. - uint64_t word_start = sample_idx * kWordsPerRankSample; - uint64_t word_end = pos / 64; - switch (word_end - word_start) { - case 3: - sample_rank += Popcount(words_[word_start + 2]); - [[fallthrough]]; - case 2: - sample_rank += Popcount(words_[word_start + 1]); - [[fallthrough]]; - case 1: - sample_rank += Popcount(words_[word_start]); - [[fallthrough]]; - case 0: - break; - } - // Count bits within the final partial word [0, pos % 64). - uint64_t remaining = pos % 64; - if (remaining > 0) { - uint64_t mask = (uint64_t(1) << remaining) - 1; - sample_rank += Popcount(words_[word_end] & mask); - } - return sample_rank; - } - - // rank0(pos): Count of 0-bits in positions [0, pos). - inline uint64_t Rank0(uint64_t pos) const { return pos - Rank1(pos); } - - // select1(i): Position of the i-th 1-bit (0-indexed). - // Returns num_bits_ if there are fewer than (i+1) 1-bits. - // Inlined for hot-path performance -- the hint lookup + linear scan is - // branch-predictor friendly and avoids function call overhead. - inline uint64_t FindNthOneBit(uint64_t i) const { - if (i >= num_ones_) { - return num_bits_; - } - // Use select hints to narrow the search range. - uint64_t lo; - uint64_t hi; - if (num_select1_hints_ > 0) { - uint64_t hint_idx = i / kOnesPerSelectHint; - lo = select1_hints_[hint_idx]; - hi = (hint_idx + 1 < num_select1_hints_) ? select1_hints_[hint_idx + 1] - : num_rank_samples_ - 1; - } else { - lo = 0; - hi = num_rank_samples_ - 1; - } - // Linear scan within the narrowed range. - while (lo < hi && rank_lut_[lo + 1] <= i) { - lo++; - } - // Scan words in the sample interval. - uint64_t remaining = i - rank_lut_[lo]; - uint64_t word_start = lo * kWordsPerRankSample; - uint64_t word_end = std::min(word_start + kWordsPerRankSample, num_words_); - for (uint64_t w = word_start; w < word_end; w++) { - uint64_t pc = Popcount(words_[w]); - if (remaining < pc) { - return w * 64 + FindNthSetBitInWord(words_[w], remaining); - } - remaining -= pc; - } - assert(false); - return num_bits_; - } - - // select0(i): Position of the i-th 0-bit (0-indexed). - // Returns num_bits_ if there are fewer than (i+1) 0-bits. - uint64_t FindNthZeroBit(uint64_t i) const; - - // Find the next set bit at or after position `pos`. - // Returns num_bits_ if no set bit is found. - // Inlined for hot-path performance -- called on every dense Seek, Advance, - // and DescendToLeftmostLeaf. - inline uint64_t NextSetBit(uint64_t pos) const { - if (pos >= num_bits_) { - return num_bits_; - } - uint64_t word_idx = pos / 64; - uint64_t bit_idx = pos % 64; - uint64_t word = words_[word_idx] >> bit_idx; - if (word != 0) { - uint64_t result = pos + CtzNonZero(word); - return (result < num_bits_) ? result : num_bits_; - } - for (uint64_t w = word_idx + 1; w < num_words_; w++) { - if (words_[w] != 0) { - uint64_t result = w * 64 + CtzNonZero(words_[w]); - return (result < num_bits_) ? result : num_bits_; - } - } - return num_bits_; - } - - // Find the last set bit strictly before position `pos`. - // Returns num_bits_ (sentinel) if no set bit exists before pos. - // Inlined for hot-path performance during reverse iteration. - inline uint64_t PrevSetBit(uint64_t pos) const { - if (pos == 0 || num_bits_ == 0) { - return num_bits_; - } - if (pos > num_bits_) { - pos = num_bits_; - } - pos--; - uint64_t word_idx = pos / 64; - uint64_t bit_idx = pos % 64; - uint64_t mask = - (bit_idx == 63) ? ~uint64_t(0) : ((uint64_t(1) << (bit_idx + 1)) - 1); - uint64_t word = words_[word_idx] & mask; - if (word != 0) { - return word_idx * 64 + static_cast(FloorLog2(word)); - } - for (int64_t w = static_cast(word_idx) - 1; w >= 0; w--) { - if (words_[w] != 0) { - return static_cast(w) * 64 + - static_cast(FloorLog2(words_[w])); - } - } - return num_bits_; - } - - // Find the distance from `pos` to the next set bit (exclusive). - // Returns the distance in bits. Used by sparse level to compute node size. - // pos must point to a set bit. - uint64_t DistanceToNextSetBit(uint64_t pos) const; - - // ---- Accessors ---- - uint64_t NumBits() const { return num_bits_; } - uint64_t NumOnes() const { return num_ones_; } - uint64_t NumZeros() const { return num_bits_ - num_ones_; } - - // Size in bytes of the serialized representation. - size_t SerializedSize() const; - - private: - // Build the rank LUT and select hints from current words. - // Used by BuildFrom(). - void BuildRankLUT(); - - // Build select hint arrays from the rank LUT. Called after BuildRankLUT(). - void BuildSelectHints(); - - // Re-compute pointers into owned_data_ after a move operation. - // Re-seats pointers into owned_data_ buffer. - // Layout: [words][rank_lut][select1_hints][select0_hints] - void RecomputePointers() { - const char* base = owned_data_.data(); - size_t words_bytes = num_words_ * sizeof(uint64_t); - size_t rank_bytes = num_rank_samples_ * sizeof(uint32_t); - size_t select1_bytes = num_select1_hints_ * sizeof(uint32_t); - - words_ = reinterpret_cast(base); - rank_lut_ = reinterpret_cast(base + words_bytes); - select1_hints_ = - reinterpret_cast(base + words_bytes + rank_bytes); - select0_hints_ = reinterpret_cast( - base + words_bytes + rank_bytes + select1_bytes); - } - - // Pointer to the raw bit words. May point into external data (InitFromData) - // or into owned_data_ (BuildFrom). - const uint64_t* words_; - - // Pointer to the rank lookup table (uint32_t entries to halve LUT size). - // Same ownership semantics as words_. - const uint32_t* rank_lut_; - - // Select hint arrays: select1_hints_[k] is the rank LUT index where the - // (k * kOnesPerSelectHint)-th 1-bit lives. select0_hints_ is analogous - // for 0-bits. These narrow Select to a linear scan of 1-2 rank samples. - const uint32_t* select1_hints_; - const uint32_t* select0_hints_; - - uint64_t num_bits_; - uint64_t num_ones_; - uint64_t num_words_; - uint64_t num_rank_samples_; - uint64_t num_select1_hints_; - uint64_t num_select0_hints_; - - // Owned storage when built from BitvectorBuilder (not from serialized data). - std::string owned_data_; -}; - -// ============================================================================ -// EliasFano: Compressed representation of a monotonically non-decreasing -// sequence of uint64_t values. Uses the Elias-Fano encoding which achieves -// near-optimal space (within 2 bits per element of the information-theoretic -// minimum) while supporting O(1) random access. -// -// Encoding: -// Given n values in [0, U], each value v is split into: -// - high part: v >> low_bits (unary-coded in a bitvector) -// - low part: v & ((1 << low_bits) - 1) (packed in a bit array) -// where low_bits = floor(log2(U / n)) when n > 0 and U > n, or 0 when -// n == 0 or U <= n. -// -// The high-bits bitvector has ones at position high[i] + i for each -// element i. The i-th value is recovered as: -// value = (FindNthOneBit(high_bv, i) - i) << low_bits | packed_low[i] -// -// Memory layout (serialized): -// [count : uint64_t] Number of elements -// [universe : uint64_t] Maximum value + 1 -// [low_bits : uint64_t] Number of low bits per element -// [high_bv : Bitvector] Unary-coded high parts (with rank/select) -// [low_words : ceil(count * low_bits / 64) * 8 bytes] Packed low parts -// -// Space for 32K offsets with U = 4GB: low_bits = 17, total ~76 KB vs -// 256 KB for raw uint64_t array (3.4x compression). -// ============================================================================ -class EliasFano { - public: - EliasFano() - : count_(0), - universe_(0), - low_bits_(0), - low_mask_(0), - low_words_(nullptr), - num_low_words_(0) {} - - ~EliasFano() = default; - - EliasFano(const EliasFano&) = delete; - EliasFano& operator=(const EliasFano&) = delete; - - // Custom move operations to re-seat low_words_ after moving owned_low_data_. - // The default move would leave low_words_ pointing into the moved-from - // object's owned_low_data_ buffer. For SSO-sized strings (owned_low_data_ - // with <= ~22 bytes, e.g. count_=1 low_bits_=8 -> 8 bytes), std::string - // move copies the SSO buffer to a new address rather than transferring a - // heap pointer, leaving low_words_ dangling. For heap-allocated strings - // move transfers the pointer (address preserved), but we re-seat - // unconditionally for correctness since SSO thresholds are implementation- - // defined. This mirrors Bitvector's move pattern (RecomputePointers). - EliasFano(EliasFano&& other) noexcept : EliasFano() { - *this = std::move(other); - } - EliasFano& operator=(EliasFano&& other) noexcept { - if (this != &other) { - count_ = other.count_; - universe_ = other.universe_; - low_bits_ = other.low_bits_; - low_mask_ = other.low_mask_; - num_low_words_ = other.num_low_words_; - high_bv_ = std::move(other.high_bv_); - owned_low_data_ = std::move(other.owned_low_data_); - // Re-seat low_words_ into our owned buffer. When owned_low_data_ is - // non-empty, low_words_ must point into it (BuildFrom path). When - // empty, low_words_ points to external memory (InitFromData path) - // or is nullptr (empty sequence). - if (!owned_low_data_.empty()) { - low_words_ = reinterpret_cast(owned_low_data_.data()); - } else { - low_words_ = other.low_words_; - } - other.count_ = 0; - other.universe_ = 0; - other.low_bits_ = 0; - other.low_mask_ = 0; - other.low_words_ = nullptr; - other.num_low_words_ = 0; - } - return *this; - } - - // Build from a sorted sequence of uint64_t values. - // Values must be monotonically non-decreasing. - void BuildFrom(const uint64_t* values, uint64_t count, uint64_t universe); - - // Initialize from serialized data. Returns Status::OK() on success. - Status InitFromData(const char* data, size_t data_size, - size_t* bytes_consumed); - - // Append serialized representation to `output`. - void EncodeTo(std::string* output) const; - - // Access the i-th value (0-indexed). i must be < count_. - inline uint64_t Access(uint64_t i) const { - assert(i < count_); - uint64_t high = high_bv_.FindNthOneBit(i) - i; - uint64_t low = 0; - if (low_bits_ > 0) { - // Extract low_bits_ bits starting at bit position i * low_bits_. - uint64_t bit_pos = i * low_bits_; - uint64_t word_idx = bit_pos / 64; - uint64_t bit_idx = bit_pos % 64; - low = (low_words_[word_idx] >> bit_idx) & low_mask_; - // Handle the case where the low bits span two words. - // Guard bit_idx > 0 to prevent UB: when bit_idx == 0 the shift - // (64 - bit_idx) would be 64 which overflows uint64_t. This is - // unreachable in practice (low_bits_ <= 63), but the guard makes - // it explicit for static analyzers. - if (bit_idx > 0 && bit_idx + low_bits_ > 64) { - uint64_t remaining = bit_idx + low_bits_ - 64; - low |= (low_words_[word_idx + 1] & ((uint64_t(1) << remaining) - 1)) - << (64 - bit_idx); - } - } - return (high << low_bits_) | low; - } - - uint64_t Count() const { return count_; } - uint64_t Universe() const { return universe_; } - - // Size in bytes of the serialized representation. - size_t SerializedSize() const; - - private: - uint64_t count_; // Number of elements. - uint64_t universe_; // Upper bound (max value + 1). - uint64_t low_bits_; // Number of low bits per element. - uint64_t low_mask_; // (1 << low_bits) - 1, precomputed for Access. - - // High bits: unary-coded bitvector with rank/select support. - Bitvector high_bv_; - - // Low bits: packed array of low_bits_-bit values. - // Points into external data (InitFromData) or into owned_low_data_. - const uint64_t* low_words_; - uint64_t num_low_words_; - - // Owned storage for low bits (when built from BuildFrom). - std::string owned_low_data_; -}; - -} // namespace trie_index -} // namespace ROCKSDB_NAMESPACE diff --git a/utilities/trie_index/louds_trie.cc b/utilities/trie_index/louds_trie.cc deleted file mode 100644 index fb141b7bbcbf..000000000000 --- a/utilities/trie_index/louds_trie.cc +++ /dev/null @@ -1,2794 +0,0 @@ -// Copyright (c) Meta Platforms, Inc. and affiliates. -// This source code is licensed under both the GPLv2 (found in the -// COPYING file in the root directory) and Apache 2.0 License -// (found in the LICENSE.Apache file in the root directory). - -#include "utilities/trie_index/louds_trie.h" - -#include -#include -#include -#include -#include - -#include "util/coding.h" - -namespace ROCKSDB_NAMESPACE { -namespace trie_index { - -static constexpr uint32_t kTrieFormatVersion = 1; -static constexpr uint32_t kTrieMagic = 0x54524945; // "TRIE" in ASCII - -// Header flags. The flags field is a 4-byte bitmask stored after -// dense_child_count in the serialized header. -static constexpr uint32_t kFlagSeqnoEncoding = 1u << 0; - -// ============================================================================ -// LoudsTrieBuilder implementation -// ============================================================================ - -LoudsTrieBuilder::LoudsTrieBuilder() - : cutoff_level_(0), - max_depth_(0), - dense_leaf_count_(0), - dense_node_count_(0), - dense_child_count_(0), - has_seqno_encoding_(false) {} - -void LoudsTrieBuilder::AddKey(const Slice& key, const TrieBlockHandle& handle) { - keys_.emplace_back(key.data(), key.size()); - handles_.push_back(handle); -} - -void LoudsTrieBuilder::AddKeyWithSeqno(const Slice& key, - const TrieBlockHandle& handle, - uint64_t seqno, uint32_t block_count) { - keys_.emplace_back(key.data(), key.size()); - handles_.push_back(handle); - seqnos_.push_back(seqno); - block_counts_.push_back(block_count); -} - -void LoudsTrieBuilder::AddOverflowBlock(const TrieBlockHandle& handle, - uint64_t seqno) { - // Seqno may be 0 when bottommost compaction zeroes all sequence numbers. - // In that case, every block in the same-key run has seqno=0. The reader's - // post-seek correction handles this correctly: the primary leaf's seqno=0 - // triggers the "never advance" guard (leaf_seqno != 0 check), so the seek - // returns the primary block. Next() iterates overflow blocks by index, not - // seqno, so all blocks are still visited in order. - overflow_handles_.push_back(handle); - overflow_seqnos_.push_back(seqno); -} - -// Compute the optimal cutoff level between dense and sparse encoding. -// Merges the node-per-level and label-per-level computations into a single -// pass over the key set, avoiding the redundant LCP computation that would -// occur if they were computed separately. -uint32_t LoudsTrieBuilder::ComputeCutoffLevel() const { - if (keys_.empty()) { - return 0; - } - - std::vector nodes_per_level(max_depth_ + 1, 0); - std::vector labels_per_level(max_depth_ + 1, 0); - nodes_per_level[0] = 1; - - for (size_t i = 0; i < keys_.size(); i++) { - // Compute LCP with previous key (once per key pair). - uint32_t lcp = 0; - if (i > 0) { - const auto& prev = keys_[i - 1]; - const auto& curr = keys_[i]; - uint32_t min_len = - static_cast(std::min(prev.size(), curr.size())); - while (lcp < min_len && prev[lcp] == curr[lcp]) { - lcp++; - } - } - - uint32_t key_len = static_cast(keys_[i].size()); - - // New nodes: levels (lcp+1) through key_len (same as before). - for (uint32_t l = lcp + 1; l <= key_len && l <= max_depth_; l++) { - nodes_per_level[l]++; - } - - // Labels: a label exists at level l if the key has a character at l - // and this label is distinct from the previous key's label (lcp <= l). - // For the first key (i==0), all levels contribute a label. - // Labels: count distinct labels per level. For sorted keys, a label at - // level l is "new" only if l >= lcp (the shared prefix with the previous - // key). For the first key, all levels contribute a new label. - uint32_t label_start = (i == 0) ? 0 : lcp; - for (uint32_t l = label_start; l < key_len && l <= max_depth_; l++) { - labels_per_level[l]++; - } - } - - // Find first level where sparse encoding is cheaper than dense. - for (uint32_t l = 0; l <= max_depth_; l++) { - uint64_t n = nodes_per_level[l]; - uint64_t labels = labels_per_level[l]; - if (n == 0) { - return l; - } - - // Dense: 256 bits for d_labels_ + 1 bit for d_is_prefix_key_ per node, - // plus 1 bit for d_has_child_ per label. - uint64_t dense_cost = n * 257 + labels; - // Sparse: 8 bits (label byte) + 1 bit (s_has_child_) + 1 bit (s_louds_) - // per label, plus 1 bit for s_is_prefix_key_ per node. - uint64_t sparse_cost = labels * 10 + n; - - if (sparse_cost < dense_cost) { - return l; - } - } - return max_depth_ + 1; -} - -void LoudsTrieBuilder::Finish() { - if (keys_.empty()) { - cutoff_level_ = 0; - max_depth_ = 0; - dense_leaf_count_ = 0; - dense_node_count_ = 0; - dense_child_count_ = 0; - SerializeAll(); - return; - } - - // Compute max depth. - max_depth_ = 0; - for (const auto& key : keys_) { - max_depth_ = std::max(max_depth_, static_cast(key.size())); - } - - // Determine cutoff level. - cutoff_level_ = ComputeCutoffLevel(); - cutoff_level_ = std::min(cutoff_level_, max_depth_ + 1); - - // ========================================================================= - // Build per-level trie data directly from sorted keys (streaming approach). - // - // Instead of constructing an explicit trie with per-node heap-allocated - // children vectors and then BFS-encoding it, we infer the trie structure - // directly from the sorted key sequence using LCP (longest common prefix) - // analysis and build flat per-level arrays. This is the approach used by - // the SuRF reference implementation (Zhang et al., SIGMOD 2018). - // - // Memory advantage: flat per-level vectors avoid the O(total_nodes) per-node - // allocation overhead. For N keys with average depth D: - // - Explicit tree: ~48 bytes per node x N*D nodes = ~48*N*D bytes - // - Per-level: ~10 bytes per label + ~17 bytes per node, all in flat - // vectors with a single allocation per level. - // - // Algorithm: For each key, compute LCP with the previous key. The LCP - // determines which levels share the trie path (existing nodes) and where - // new branches start. At each level l >= lcp, the key contributes a new - // label to the current (last) node at that level. - // - // Three key mechanisms handle deferred information: - // 1. Deferred internal marking: when a label is first added, we mark it - // as a leaf (has_child=false). If a subsequent key continues the path - // through this label, we retroactively mark it as internal. - // 2. Handle migration: when a leaf becomes internal (i.e., a prefix key - // emerges), the handle moves from the label's leaf_handle to the - // child node's prefix_handle. - // 3. Lazy node creation: terminal nodes at depth K (where a key of - // length K ends) are created on-demand when retroactive marking - // needs them. - // ========================================================================= - - // Per-level data: flat arrays for labels, has_child, and handle tracking. - // All labels within a level are concatenated; node boundaries are tracked - // by node_label_start (the index where each node's labels begin). - struct PerLevelData { - // ---- Per-label data (all nodes concatenated) ---- - std::vector labels; - // Using uint8_t instead of bool to avoid std::vector proxy issues. - std::vector has_child; - // Handle index for leaf labels (-1 if internal). When has_child[i] is - // false, leaf_handle[i] is the key index whose handle should be emitted - // for this leaf child. - std::vector leaf_handle; - - // ---- Per-node data ---- - std::vector is_prefix; - // Handle index for prefix key nodes (-1 if not a prefix key). - std::vector prefix_handle; - // Index in labels[] where each node's labels start. Used to derive - // louds bits and to iterate over nodes during LOUDS encoding. - std::vector node_label_start; - - uint64_t node_count() const { - return static_cast(is_prefix.size()); - } - - // Start a new node. Call this before adding labels to the node. - void StartNode() { - node_label_start.push_back(static_cast(labels.size())); - is_prefix.push_back(false); - prefix_handle.push_back(-1); - } - - // Add a label to the current (last) node. - void AddLabel(uint8_t label, bool is_internal, int64_t handle) { - labels.push_back(label); - has_child.push_back(is_internal); - leaf_handle.push_back(is_internal ? -1 : handle); - } - }; - - std::vector levels(max_depth_ + 1); - - // Initialize root node at level 0. - levels[0].StartNode(); - - for (size_t ki = 0; ki < keys_.size(); ki++) { - const auto& key = keys_[ki]; - uint32_t K = static_cast(key.size()); - - // Compute LCP with previous key. - uint32_t lcp = 0; - if (ki > 0) { - const auto& prev = keys_[ki - 1]; - uint32_t min_len = - static_cast(std::min(prev.size(), key.size())); - while (lcp < min_len && prev[lcp] == key[lcp]) { - lcp++; - } - } - - // Step 1: Retroactive internal marking. - // - // If lcp > 0, the label at level lcp-1 (the deepest shared label) may - // need to transition from leaf to internal. This happens when: - // - The previous key's path ended at or before level lcp (the label - // was originally a leaf), AND - // - The current key continues deeper through the same child. - // - // Since keys are sorted and lcp is the shared prefix length, the label - // at level lcp-1 is always the LAST label in the current (last) node at - // that level. We only need to mark this one label. - if (lcp > 0) { - auto& pl = levels[lcp - 1]; - if (!pl.has_child.empty() && !pl.has_child.back()) { - // Leaf -> internal transition. - pl.has_child.back() = true; - - // Handle migration: move the leaf's handle to a NEW child node at - // level lcp as a prefix key. The child node is always new because - // this label just transitioned from leaf to internal -- no child - // node existed before for this branch. - // - // The leaf_handle must be valid (>= 0) because the label was marked - // as a leaf in Step 2, which always sets the handle for leaf labels. - assert(pl.leaf_handle.back() >= 0); - int64_t h = pl.leaf_handle.back(); - pl.leaf_handle.back() = -1; - - if (lcp <= max_depth_) { - levels[lcp].StartNode(); - levels[lcp].is_prefix.back() = true; - levels[lcp].prefix_handle.back() = h; - } - } - } - - // Step 2: Add labels for levels lcp..K-1. - // - // At level lcp: add a label to the existing (last) node. - // At levels > lcp: create a new node and add its first label. - for (uint32_t l = lcp; l < K; l++) { - uint8_t label = static_cast(key[l]); - bool is_internal = (l + 1 < K); // Key continues to next level. - - if (l == lcp) { - // Adding to existing node. Verify it exists. - assert(levels[l].node_count() > 0); - } else { - // New child node at this level. - levels[l].StartNode(); - } - - // Associate handle with the last label if this is a leaf (key ends). - int64_t handle = is_internal ? -1 : static_cast(ki); - levels[l].AddLabel(label, is_internal, handle); - } - } - - // ========================================================================= - // Phase 2: Build LOUDS bitvectors and reorder handles from per-level data. - // - // The per-level data is already in BFS order (nodes at each level are - // created left-to-right as sorted keys are processed). We iterate by - // level and by node within each level: - // - For each node: emit prefix key handle (if any), then for each label: - // if leaf -> emit handle; if internal -> skip (child visited at next - // level). - // - Build dense/sparse bitvectors from labels and has_child flags. - // ========================================================================= - - d_labels_ = BitvectorBuilder(); - d_has_child_ = BitvectorBuilder(); - d_is_prefix_key_ = BitvectorBuilder(); - dense_node_count_ = 0; - dense_child_count_ = 0; - - s_labels_.clear(); - s_has_child_ = BitvectorBuilder(); - s_louds_ = BitvectorBuilder(); - s_is_prefix_key_ = BitvectorBuilder(); - - dense_leaf_count_ = 0; - - std::vector bfs_ordered_handles; - bfs_ordered_handles.reserve(keys_.size()); - // BFS-ordered seqno side-table data (same reordering as handles). - // Only used when has_seqno_encoding_ is true. - std::vector bfs_ordered_seqnos; - std::vector bfs_ordered_block_counts; - // Overflow arrays must also be BFS-reordered. Without this, the - // overflow_base_ prefix sum (computed from BFS-ordered block_counts) - // would index into key-sorted overflow arrays, mapping overflow blocks - // to the wrong leaves when BFS order differs from key-sorted order - // (which happens whenever separator keys have different lengths). - std::vector key_sorted_overflow_base; - std::vector bfs_ordered_overflow_handles; - std::vector bfs_ordered_overflow_seqnos; - if (has_seqno_encoding_) { - bfs_ordered_seqnos.reserve(keys_.size()); - bfs_ordered_block_counts.reserve(keys_.size()); - if (!overflow_handles_.empty()) { - key_sorted_overflow_base.resize(keys_.size()); - uint32_t sum = 0; - for (size_t i = 0; i < keys_.size(); i++) { - key_sorted_overflow_base[i] = sum; - assert(block_counts_[i] >= 1); - sum += block_counts_[i] - 1; - } - assert(sum == static_cast(overflow_handles_.size())); - bfs_ordered_overflow_handles.reserve(overflow_handles_.size()); - bfs_ordered_overflow_seqnos.reserve(overflow_seqnos_.size()); - } - } - - // Emit BFS-ordered data for a leaf with key index ki: primary handle, - // seqno side-table fields, and overflow blocks (if any). - auto emit_leaf = [&](size_t ki) { - bfs_ordered_handles.push_back(handles_[ki]); - if (has_seqno_encoding_) { - bfs_ordered_seqnos.push_back(seqnos_[ki]); - bfs_ordered_block_counts.push_back(block_counts_[ki]); - if (!key_sorted_overflow_base.empty() && block_counts_[ki] > 1) { - uint32_t base = key_sorted_overflow_base[ki]; - uint32_t count = block_counts_[ki] - 1; - for (uint32_t j = 0; j < count; j++) { - bfs_ordered_overflow_handles.push_back(overflow_handles_[base + j]); - bfs_ordered_overflow_seqnos.push_back(overflow_seqnos_[base + j]); - } - } - } - }; - - for (uint32_t level = 0; level <= max_depth_; level++) { - const auto& ld = levels[level]; - if (ld.node_count() == 0) { - continue; - } - - bool is_dense = (level < cutoff_level_); - - for (uint64_t ni = 0; ni < ld.node_count(); ni++) { - // Determine label range for this node. - uint64_t label_start = ld.node_label_start[ni]; - uint64_t label_end = (ni + 1 < ld.node_count()) - ? ld.node_label_start[ni + 1] - : static_cast(ld.labels.size()); - - // ---- Handle reordering: emit prefix key handle ---- - if (ld.is_prefix[ni] && ld.prefix_handle[ni] >= 0) { - emit_leaf(static_cast(ld.prefix_handle[ni])); - } - - // Skip pure leaf nodes (no labels) -- they are accounted for by - // has_child=0 in their parent. They don't produce LOUDS entries. - if (label_start == label_end) { - if (ld.is_prefix[ni]) { - // Prefix key node with no children at THIS level (but it was - // already counted above). This can happen for nodes created by - // lazy creation that never gained children. - if (is_dense) { - dense_leaf_count_++; - } - } - continue; - } - - if (is_dense) { - // ---- Dense encoding ---- - dense_node_count_++; - d_is_prefix_key_.Append(ld.is_prefix[ni]); - - // Build 256-bit label bitmap as 4 x 64-bit words. - uint64_t bitmap[4] = {}; - for (uint64_t li = label_start; li < label_end; li++) { - uint8_t label = ld.labels[li]; - bitmap[label / 64] |= uint64_t(1) << (label % 64); - } - for (int w = 0; w < 4; w++) { - d_labels_.AppendWord(bitmap[w], 64); - } - - // Emit has_child bits, count leaves, emit leaf handles. - for (uint64_t li = label_start; li < label_end; li++) { - bool is_internal = ld.has_child[li]; - d_has_child_.Append(is_internal); - - if (!is_internal) { - if (ld.leaf_handle[li] >= 0) { - emit_leaf(static_cast(ld.leaf_handle[li])); - } - dense_leaf_count_++; - } else if (level == cutoff_level_ - 1) { - dense_child_count_++; - } - } - - if (ld.is_prefix[ni]) { - dense_leaf_count_++; - } - } else { - // ---- Sparse encoding ---- - s_is_prefix_key_.Append(ld.is_prefix[ni]); - bool first_label = true; - for (uint64_t li = label_start; li < label_end; li++) { - bool is_internal = ld.has_child[li]; - - s_labels_.push_back(ld.labels[li]); - s_has_child_.Append(is_internal); - s_louds_.Append(first_label); - first_label = false; - - if (!is_internal) { - if (ld.leaf_handle[li] >= 0) { - emit_leaf(static_cast(ld.leaf_handle[li])); - } - } - } - } - } - } - - // When cutoff_level_ = 0, all nodes are sparse and the root is the sole - // "root sparse node" (not a child of any dense level). - if (cutoff_level_ == 0) { - dense_child_count_ = 1; - } - - assert(bfs_ordered_handles.size() == keys_.size()); - handles_ = std::move(bfs_ordered_handles); - if (has_seqno_encoding_) { - assert(bfs_ordered_seqnos.size() == keys_.size()); - assert(bfs_ordered_block_counts.size() == keys_.size()); - seqnos_ = std::move(bfs_ordered_seqnos); - block_counts_ = std::move(bfs_ordered_block_counts); - if (!overflow_handles_.empty()) { - assert(bfs_ordered_overflow_handles.size() == overflow_handles_.size()); - assert(bfs_ordered_overflow_seqnos.size() == overflow_seqnos_.size()); - overflow_handles_ = std::move(bfs_ordered_overflow_handles); - overflow_seqnos_ = std::move(bfs_ordered_overflow_seqnos); - } - } - - SerializeAll(); -} - -void LoudsTrieBuilder::SerializeAll() { - serialized_data_.clear(); - - PutFixed32(&serialized_data_, kTrieMagic); - PutFixed32(&serialized_data_, kTrieFormatVersion); - PutFixed64(&serialized_data_, keys_.size()); - PutFixed32(&serialized_data_, cutoff_level_); - PutFixed32(&serialized_data_, max_depth_); - PutFixed64(&serialized_data_, dense_leaf_count_); - PutFixed64(&serialized_data_, dense_node_count_); - PutFixed64(&serialized_data_, dense_child_count_); - - // Flags field. Contains feature bits that inform the reader about - // encoding decisions made during building. - uint32_t flags = 0; - if (has_seqno_encoding_) { - flags |= kFlagSeqnoEncoding; - } - PutFixed32(&serialized_data_, flags); - // 4 bytes of reserved padding to maintain 8-byte alignment (56-byte header). - PutFixed32(&serialized_data_, 0); - - // Dense section. - { - Bitvector bv; - bv.BuildFrom(d_labels_); - bv.EncodeTo(&serialized_data_); - } - { - Bitvector bv; - bv.BuildFrom(d_has_child_); - bv.EncodeTo(&serialized_data_); - } - { - Bitvector bv; - bv.BuildFrom(d_is_prefix_key_); - bv.EncodeTo(&serialized_data_); - } - - // Sparse section. - uint64_t sl_size = s_labels_.size(); - PutFixed64(&serialized_data_, sl_size); - if (sl_size > 0) { - serialized_data_.append(reinterpret_cast(s_labels_.data()), - sl_size); - size_t padding = (8 - (sl_size % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - Bitvector bv_s_has_child; - bv_s_has_child.BuildFrom(s_has_child_); - bv_s_has_child.EncodeTo(&serialized_data_); - - Bitvector bv_s_louds; - bv_s_louds.BuildFrom(s_louds_); - bv_s_louds.EncodeTo(&serialized_data_); - - { - Bitvector bv; - bv.BuildFrom(s_is_prefix_key_); - bv.EncodeTo(&serialized_data_); - } - - // Child position lookup tables for Select-free sparse traversal. - // Compute and serialize: s_child_start_pos_[k] and s_child_end_pos_[k] - // for each internal label (has_child=1). - uint64_t num_internal = bv_s_has_child.NumOnes(); - { - PutFixed64(&serialized_data_, num_internal); - - if (num_internal > 0 && sl_size <= UINT32_MAX) { - for (uint64_t k = 0; k < num_internal; k++) { - uint64_t child_node = dense_child_count_ + k; - uint64_t child_start = bv_s_louds.FindNthOneBit(child_node); - PutFixed32(&serialized_data_, static_cast(child_start)); - } - // Pad to 8-byte alignment - size_t bytes_written = num_internal * sizeof(uint32_t); - size_t padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - - for (uint64_t k = 0; k < num_internal; k++) { - uint64_t child_node = dense_child_count_ + k; - uint64_t child_start = bv_s_louds.FindNthOneBit(child_node); - uint64_t child_end = bv_s_louds.NextSetBit(child_start + 1); - if (child_end > sl_size) { - child_end = sl_size; - } - PutFixed32(&serialized_data_, static_cast(child_end)); - } - // Pad to 8-byte alignment - bytes_written = num_internal * sizeof(uint32_t); - padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - } - - // Path compression: detect and serialize fanout-1 chains. - // - // A chain starts at the child of internal label k and consists of >= 2 - // consecutive fanout-1 nodes (all internal except possibly the last, - // which may be a leaf). For each chain, we store: - // - suffix bytes (the label at each chain node) - // - chain length - // - the child_idx at the end of the chain (UINT32_MAX if ends at leaf) - // - // This lets Seek() skip entire chains with a single memcmp instead of - // traversing level by level with Rank1 at each step. - { - // Build child start/end arrays in memory for chain detection. - // These are the same values we just serialized above. - std::vector child_starts(num_internal); - std::vector child_ends(num_internal); - if (num_internal > 0 && sl_size <= UINT32_MAX) { - for (uint64_t k = 0; k < num_internal; k++) { - uint64_t child_node = dense_child_count_ + k; - uint64_t cs = bv_s_louds.FindNthOneBit(child_node); - child_starts[k] = static_cast(cs); - uint64_t ce = bv_s_louds.NextSetBit(cs + 1); - if (ce > sl_size) { - ce = sl_size; - } - child_ends[k] = static_cast(ce); - } - } - - // For each internal label k, detect if its child starts a chain. - std::string chain_suffix_data; - // Offsets into chain_suffix_data for each internal label. - // UINT32_MAX means no chain. - std::vector chain_offsets(num_internal, UINT32_MAX); - std::vector chain_lens(num_internal, 0); - // child_idx at chain end (UINT32_MAX = leaf). - std::vector chain_end_child_idx(num_internal, UINT32_MAX); - - // Build s_is_prefix_key bitvector for prefix key checks during chain - // detection. Chains must not contain prefix key nodes because the chain - // skip logic bypasses prefix key checks. - Bitvector bv_s_is_prefix_key; - bv_s_is_prefix_key.BuildFrom(s_is_prefix_key_); - - for (uint64_t k = 0; k < num_internal; k++) { - uint32_t cs = child_starts[k]; - uint32_t ce = child_ends[k]; - - // Child must be fanout-1 (single label). - if (ce - cs != 1) { - continue; - } - - // Check if the child's label is internal (has_child = 1). - if (!bv_s_has_child.GetBit(cs)) { - continue; - } - - // Check if the child node is a prefix key -- if so, cannot skip it. - { - uint64_t child_sparse_node = bv_s_louds.Rank1(cs + 1) - 1; - if (child_sparse_node < bv_s_is_prefix_key.NumBits() && - bv_s_is_prefix_key.GetBit(child_sparse_node)) { - continue; - } - } - - // Found a fanout-1 internal node. Follow the chain. - std::vector suffix; - suffix.push_back(s_labels_[cs]); - - // Get the child_idx of this chain node. - uint64_t cur_child_idx = bv_s_has_child.Rank1(cs + 1) - 1; - uint32_t last_child_idx = - UINT32_MAX; // Will be set to chain end's child. - - while (true) { - // cur_child_idx is the index of the current chain node's internal - // label. - if (cur_child_idx >= num_internal) { - break; - } - - uint32_t next_cs = child_starts[cur_child_idx]; - uint32_t next_ce = child_ends[cur_child_idx]; - - if (next_ce - next_cs != 1) { - // Child has fanout > 1. Chain ends here at an internal node. - last_child_idx = static_cast(cur_child_idx); - break; - } - - // Check if next node is a prefix key -- stop chain here. - { - uint64_t next_sparse_node = bv_s_louds.Rank1(next_cs + 1) - 1; - if (next_sparse_node < bv_s_is_prefix_key.NumBits() && - bv_s_is_prefix_key.GetBit(next_sparse_node)) { - // End chain before this prefix key node. - last_child_idx = static_cast(cur_child_idx); - break; - } - } - - // Next child is also fanout-1. Check if internal or leaf. - suffix.push_back(s_labels_[next_cs]); - - if (!bv_s_has_child.GetBit(next_cs)) { - // Chain ends at a leaf. - last_child_idx = UINT32_MAX; - break; - } - - // Continue chaining. - cur_child_idx = bv_s_has_child.Rank1(next_cs + 1) - 1; - } - - // Only store chains of meaningful length. Short chains don't save - // enough Rank1 calls to justify the metadata overhead (10 bytes per - // chain + suffix bytes). - static constexpr size_t kMinChainLength = 8; - if (suffix.size() >= kMinChainLength && suffix.size() <= UINT16_MAX) { - chain_offsets[k] = static_cast(chain_suffix_data.size()); - chain_lens[k] = static_cast(suffix.size()); - chain_end_child_idx[k] = last_child_idx; - chain_suffix_data.append(reinterpret_cast(suffix.data()), - suffix.size()); - } - } - - // Chain cost/benefit filter: only emit chains when they provide a net - // speed benefit. - // - // Key insight: chains help when a Seek is *likely* to hit a chain. - // For tries with few chains relative to keys (e.g., numeric keys with - // a single long shared-prefix chain), every Seek benefits. For tries - // with many chains relative to keys (e.g., random hex with thousands - // of short chains), each chain is rarely hit and the per-level bitmap - // check overhead outweighs the occasional hit. - // - // Metric: emit chains only when num_chains <= num_keys. When there are - // more chains than keys, most chains won't be hit by any Seek, making - // the bitmap overhead a net loss. Additionally, apply a space budget - // (10% of base trie size) to prevent excessive metadata. - { - static constexpr double kChainBudgetPct = 0.10; - static constexpr size_t kPerChainMetaBytes = 10; // 4 + 2 + 4 - - uint64_t candidate_count = 0; - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0) { - candidate_count++; - } - } - - uint64_t num_keys = handles_.size(); - bool too_many_chains = (candidate_count > num_keys); - - if (candidate_count == 0 || too_many_chains) { - for (uint64_t k = 0; k < num_internal; k++) { - chain_lens[k] = 0; - } - } else { - // Apply space budget to prevent excessive metadata even when - // the chain count is reasonable. - size_t base_trie_size = serialized_data_.size(); - size_t budget = static_cast(base_trie_size * kChainBudgetPct); - size_t bitmap_fixed_cost = - (num_internal + 7) / 8 + (num_internal / 512 + 1) * 8; - - if (budget <= bitmap_fixed_cost) { - for (uint64_t k = 0; k < num_internal; k++) { - chain_lens[k] = 0; - } - } else { - size_t available = budget - bitmap_fixed_cost; - size_t total_cost = 0; - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0) { - total_cost += kPerChainMetaBytes + chain_lens[k]; - } - } - - if (total_cost > available) { - // Over budget. Keep longest chains first. - struct ChainCandidate { - uint64_t idx; - size_t cost; - uint16_t length; - }; - std::vector candidates; - candidates.reserve(candidate_count); - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0) { - candidates.push_back( - {k, kPerChainMetaBytes + chain_lens[k], chain_lens[k]}); - } - } - std::sort(candidates.begin(), candidates.end(), - [](const ChainCandidate& a, const ChainCandidate& b) { - return a.length > b.length; - }); - - std::vector keep(num_internal, false); - size_t used = 0; - for (const auto& c : candidates) { - if (used + c.cost <= available) { - keep[c.idx] = true; - used += c.cost; - } - } - - std::string new_suffix_data; - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0 && !keep[k]) { - chain_lens[k] = 0; - chain_offsets[k] = UINT32_MAX; - chain_end_child_idx[k] = UINT32_MAX; - } else if (chain_lens[k] > 0) { - uint32_t old_off = chain_offsets[k]; - chain_offsets[k] = - static_cast(new_suffix_data.size()); - new_suffix_data.append(chain_suffix_data.data() + old_off, - chain_lens[k]); - } - } - chain_suffix_data = std::move(new_suffix_data); - } - } - } - } - - // Count actual chains and build bitmap + compact arrays. - uint64_t num_chains = 0; - BitvectorBuilder chain_bitmap_builder; - for (uint64_t k = 0; k < num_internal; k++) { - bool has_chain = (chain_lens[k] > 0); - chain_bitmap_builder.Append(has_chain); - if (has_chain) { - num_chains++; - } - } - - // Serialize: num_chains, then bitmap + compact arrays if any. - PutFixed64(&serialized_data_, num_chains); - - if (num_chains > 0) { - // Write chain bitmap (1 bit per internal label). - Bitvector chain_bitmap_bv; - chain_bitmap_bv.BuildFrom(chain_bitmap_builder); - chain_bitmap_bv.EncodeTo(&serialized_data_); - - // Write compact chain_offsets (uint32_t per chain). - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0) { - PutFixed32(&serialized_data_, chain_offsets[k]); - } - } - size_t bytes_written = num_chains * sizeof(uint32_t); - size_t padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - - // Write compact chain_lens (uint16_t per chain). - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0) { - serialized_data_.append(reinterpret_cast(&chain_lens[k]), - sizeof(uint16_t)); - } - } - bytes_written = num_chains * sizeof(uint16_t); - padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - - // Write compact chain_end_child_idx (uint32_t per chain). - for (uint64_t k = 0; k < num_internal; k++) { - if (chain_lens[k] > 0) { - PutFixed32(&serialized_data_, chain_end_child_idx[k]); - } - } - bytes_written = num_chains * sizeof(uint32_t); - padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - - // Write suffix data blob. - uint64_t chain_data_size = chain_suffix_data.size(); - PutFixed64(&serialized_data_, chain_data_size); - serialized_data_.append(chain_suffix_data); - padding = (8 - (chain_data_size % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - } - - // Block handles: packed uint32_t arrays for offsets and sizes. - // BFS leaf order does not match key-sorted order for keys of different - // lengths, so offsets are NOT monotone and cannot use Elias-Fano. - { - if (!handles_.empty()) { - size_t n = handles_.size(); - // Write offsets array (uint32_t, padded to 8-byte alignment). - for (size_t i = 0; i < n; i++) { - assert(handles_[i].offset <= UINT32_MAX); - PutFixed32(&serialized_data_, - static_cast(handles_[i].offset)); - } - size_t bytes_written = n * sizeof(uint32_t); - size_t padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - - // Write sizes array (uint32_t, padded to 8-byte alignment). - for (size_t i = 0; i < n; i++) { - assert(handles_[i].size <= UINT32_MAX); - PutFixed32(&serialized_data_, static_cast(handles_[i].size)); - } - bytes_written = n * sizeof(uint32_t); - padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - } - - // ========================================================================= - // Seqno side-table: serialized after handle arrays when kFlagSeqnoEncoding - // is set. Contains per-leaf seqno/block_count data in BFS leaf order, plus - // overflow block handles/seqnos for same-user-key runs. - // - // The trie always stores user-key-only separators. When same-user-key - // boundaries exist, duplicate separator keys are de-duplicated (only the - // first occurrence goes in the trie). Overflow blocks (the 2nd, 3rd, ... - // blocks in a same-key run) are stored here in the side-table. After a - // trie Seek lands on a leaf, the iterator uses the seqno data to determine - // whether to advance through overflow blocks. - // - // Format: - // num_overflow_blocks: uint32_t - // padding: uint32_t (0, for 8-byte alignment) - // leaf_seqnos: uint64_t[num_keys] (BFS order) - // leaf_block_counts: uint32_t[num_keys] (BFS order, padded to 8-byte) - // overflow_offsets: uint32_t[num_overflow] (padded to 8-byte) - // overflow_sizes: uint32_t[num_overflow] (padded to 8-byte) - // overflow_seqnos: uint64_t[num_overflow] - // ========================================================================= - if (has_seqno_encoding_) { - uint32_t num_overflow = static_cast(overflow_handles_.size()); - PutFixed32(&serialized_data_, num_overflow); - PutFixed32(&serialized_data_, 0); // padding for 8-byte alignment - - size_t n = handles_.size(); - - // leaf_seqnos: uint64_t per leaf (BFS order) - for (size_t i = 0; i < n; i++) { - PutFixed64(&serialized_data_, seqnos_[i]); - } - - // leaf_block_counts: uint32_t per leaf (BFS order, padded) - for (size_t i = 0; i < n; i++) { - PutFixed32(&serialized_data_, block_counts_[i]); - } - { - size_t bytes_written = n * sizeof(uint32_t); - size_t padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - - // overflow_offsets: uint32_t per overflow (padded) - for (uint32_t i = 0; i < num_overflow; i++) { - assert(overflow_handles_[i].offset <= UINT32_MAX); - PutFixed32(&serialized_data_, - static_cast(overflow_handles_[i].offset)); - } - if (num_overflow > 0) { - size_t bytes_written = num_overflow * sizeof(uint32_t); - size_t padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - - // overflow_sizes: uint32_t per overflow (padded) - for (uint32_t i = 0; i < num_overflow; i++) { - assert(overflow_handles_[i].size <= UINT32_MAX); - PutFixed32(&serialized_data_, - static_cast(overflow_handles_[i].size)); - } - if (num_overflow > 0) { - size_t bytes_written = num_overflow * sizeof(uint32_t); - size_t padding = (8 - (bytes_written % 8)) % 8; - serialized_data_.append(padding, '\0'); - } - - // overflow_seqnos: uint64_t per overflow (naturally 8-byte aligned) - for (uint32_t i = 0; i < num_overflow; i++) { - PutFixed64(&serialized_data_, overflow_seqnos_[i]); - } - } -} - -// ============================================================================ -// LoudsTrie implementation -// ============================================================================ - -LoudsTrie::LoudsTrie() - : num_keys_(0), - cutoff_level_(0), - max_depth_(0), - has_seqno_encoding_(false), - dense_leaf_count_(0), - dense_node_count_(0), - dense_child_count_(0), - s_labels_data_(nullptr), - s_labels_size_(0), - s_chain_suffix_data_(nullptr), - s_chain_suffix_size_(0), - handle_offsets_(nullptr), - handle_sizes_(nullptr), - leaf_seqnos_(nullptr), - leaf_block_counts_(nullptr), - overflow_offsets_(nullptr), - overflow_sizes_(nullptr), - overflow_seqnos_(nullptr), - num_overflow_blocks_(0) {} - -Status LoudsTrie::InitFromData(const Slice& data) { - const char* p = data.data(); - size_t remaining = data.size(); - - // The trie data contains bitvectors with uint64_t arrays and handle arrays - // with uint32_t entries, all accessed via reinterpret_cast pointers that - // require proper alignment. Block buffers from heap/cache allocations are - // typically aligned, but mmap'd data or other sources may not be. If the - // data is not 8-byte aligned, copy it into an owned aligned buffer. - // std::string::data() returns memory from new[]/malloc, which is aligned - // to at least alignof(max_align_t) (>= 8 on all supported platforms). - if (reinterpret_cast(p) % alignof(uint64_t) != 0) { - aligned_copy_.assign(p, remaining); - p = aligned_copy_.data(); - } - - // Header: magic(4) + version(4) + num_keys(8) + cutoff_level(4) + - // max_depth(4) + dense_leaf_count(8) + dense_node_count(8) + - // dense_child_count(8) + flags(4) + reserved(4) = 56. - if (remaining < 56) { - return Status::Corruption("Trie index: data too short for header"); - } - - uint32_t magic; - memcpy(&magic, p, 4); - p += 4; - remaining -= 4; - if (magic != kTrieMagic) { - return Status::Corruption("Trie index: bad magic number"); - } - - uint32_t version; - memcpy(&version, p, 4); - p += 4; - remaining -= 4; - if (version != kTrieFormatVersion) { - return Status::Corruption("Trie index: unsupported format version"); - } - - memcpy(&num_keys_, p, 8); - p += 8; - remaining -= 8; - - // Validate num_keys_ early, before any arithmetic that multiplies it by - // sizeof(uint32_t) or sizeof(uint64_t). Without this check, a crafted - // num_keys_ near SIZE_MAX / 4 causes silent size_t overflow in the handle - // array and seqno side-table parsing below, leading to buffer overreads. - // A trie with > 2^30 leaves is unrealistic (would require terabytes of - // data blocks). - static constexpr uint64_t kMaxReasonableKeys = uint64_t(1) << 30; - if (num_keys_ > kMaxReasonableKeys) { - return Status::Corruption("Trie index: num_keys exceeds reasonable limit"); - } - - memcpy(&cutoff_level_, p, 4); - p += 4; - remaining -= 4; - memcpy(&max_depth_, p, 4); - p += 4; - remaining -= 4; - - // Validate max_depth_ from untrusted data. The iterator allocates a - // key_buf_ of size MaxDepth()+1; if max_depth_ == UINT32_MAX, the +1 - // overflows uint32_t to 0, causing a zero-length allocation and subsequent - // buffer overflow. A key longer than 64 KB is unrealistic for a block - // index separator (RocksDB keys are typically < 1 KB). - static constexpr uint32_t kMaxReasonableDepth = 65536; - if (max_depth_ > kMaxReasonableDepth) { - return Status::Corruption("Trie index: max_depth exceeds reasonable limit"); - } - - memcpy(&dense_leaf_count_, p, 8); - p += 8; - remaining -= 8; - memcpy(&dense_node_count_, p, 8); - p += 8; - remaining -= 8; - memcpy(&dense_child_count_, p, 8); - p += 8; - remaining -= 8; - - // Read flags field. - uint32_t flags; - memcpy(&flags, p, 4); - p += 4; - remaining -= 4; - has_seqno_encoding_ = (flags & kFlagSeqnoEncoding) != 0; - // Skip 4-byte reserved padding. - p += 4; - remaining -= 4; - - // Dense bitvectors. - size_t consumed; - Status s; - - s = d_labels_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - s = d_has_child_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - s = d_is_prefix_key_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - // Validate dense section counts against bitvector sizes. These counts - // were read from the header (untrusted data) and will be used for leaf - // ordinal computation during traversal. Inconsistent values would cause - // incorrect rank arithmetic leading to wrong leaf indices -> wrong block - // handles. - if (dense_node_count_ > 0) { - // Each dense node uses 256 bits in d_labels_. - if (d_labels_.NumBits() != dense_node_count_ * 256) { - return Status::Corruption( - "Trie index: d_labels size inconsistent with dense_node_count"); - } - // d_has_child_ has one bit per set bit in d_labels_ (one per label). - if (d_has_child_.NumBits() != d_labels_.NumOnes()) { - return Status::Corruption( - "Trie index: d_has_child size inconsistent with d_labels"); - } - // d_is_prefix_key_ has one bit per dense node. - if (d_is_prefix_key_.NumBits() != dense_node_count_) { - return Status::Corruption( - "Trie index: d_is_prefix_key size inconsistent with " - "dense_node_count"); - } - } - if (dense_leaf_count_ > num_keys_) { - return Status::Corruption("Trie index: dense_leaf_count exceeds num_keys"); - } - - // Sparse section. - if (remaining < 8) { - return Status::Corruption("Trie index: truncated sparse labels size"); - } - memcpy(&s_labels_size_, p, 8); - p += 8; - remaining -= 8; - - if (s_labels_size_ > 0) { - if (remaining < s_labels_size_) { - return Status::Corruption("Trie index: truncated sparse labels"); - } - s_labels_data_ = reinterpret_cast(p); - p += s_labels_size_; - remaining -= s_labels_size_; - size_t padding = (8 - (s_labels_size_ % 8)) % 8; - if (remaining < padding) { - return Status::Corruption("Trie index: truncated sparse label padding"); - } - p += padding; - remaining -= padding; - } - - s = s_has_child_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - s = s_louds_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - s = s_is_prefix_key_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - // Validate sparse bitvector sizes match s_labels_size_. Each sparse - // label has exactly one bit in s_has_child_ and one bit in s_louds_. - // A mismatch means the serialized data is inconsistent, which would cause - // GetBit/Rank1 OOB reads during traversal. - if (s_labels_size_ > 0) { - if (s_has_child_.NumBits() != s_labels_size_) { - return Status::Corruption( - "Trie index: s_has_child size inconsistent with s_labels_size"); - } - if (s_louds_.NumBits() != s_labels_size_) { - return Status::Corruption( - "Trie index: s_louds size inconsistent with s_labels_size"); - } - } - - // Child position lookup tables for Select-free sparse traversal. - uint64_t num_internal = 0; - { - if (remaining < 8) { - return Status::Corruption( - "Trie index: truncated child position table count"); - } - memcpy(&num_internal, p, 8); - p += 8; - remaining -= 8; - - // Guard against integer overflow in subsequent size computations. - // Same limit as kMaxReasonableKeys used for num_keys_ validation. - static constexpr uint64_t kMaxReasonableInternal = 1ULL << 30; - if (num_internal > kMaxReasonableInternal) { - return Status::Corruption( - "Trie index: num_internal exceeds reasonable limit"); - } - - if (num_internal > 0) { - // Read s_child_start_pos_ - size_t start_bytes = num_internal * sizeof(uint32_t); - size_t start_padded = (start_bytes + 7) & ~size_t(7); - if (remaining < start_padded) { - return Status::Corruption( - "Trie index: truncated child start position table"); - } - s_child_start_pos_.resize(num_internal); - memcpy(s_child_start_pos_.data(), p, start_bytes); - p += start_padded; - remaining -= start_padded; - - // Read s_child_end_pos_ - size_t end_bytes = num_internal * sizeof(uint32_t); - size_t end_padded = (end_bytes + 7) & ~size_t(7); - if (remaining < end_padded) { - return Status::Corruption( - "Trie index: truncated child end position table"); - } - s_child_end_pos_.resize(num_internal); - memcpy(s_child_end_pos_.data(), p, end_bytes); - p += end_padded; - remaining -= end_padded; - - // Validate that child position values are within the sparse label array. - // A crafted input could have positions beyond s_labels_size_, causing - // OOB reads during traversal when labels are accessed at these positions. - for (uint64_t k = 0; k < num_internal; k++) { - if (s_child_start_pos_[k] >= s_labels_size_ || - s_child_end_pos_[k] > s_labels_size_ || - s_child_end_pos_[k] < s_child_start_pos_[k]) { - return Status::Corruption( - "Trie index: child position out of range for sparse labels"); - } - } - } - } - - // Path compression: chain metadata for fanout-1 chains. - // Format: num_chains (uint64), then if > 0: bitmap + compact arrays. - { - if (remaining < 8) { - return Status::Corruption("Trie index: truncated chain count"); - } - uint64_t num_chains; - memcpy(&num_chains, p, 8); - p += 8; - remaining -= 8; - - if (num_chains > 0) { - // Read chain bitmap. - s = s_chain_bitmap_.InitFromData(p, remaining, &consumed); - if (!s.ok()) { - return s; - } - p += consumed; - remaining -= consumed; - - // Read compact chain_offsets (uint32_t per chain). - size_t offsets_bytes = num_chains * sizeof(uint32_t); - size_t offsets_padded = (offsets_bytes + 7) & ~size_t(7); - if (remaining < offsets_padded) { - return Status::Corruption("Trie index: truncated chain offsets"); - } - s_chain_suffix_offsets_.resize(num_chains); - memcpy(s_chain_suffix_offsets_.data(), p, offsets_bytes); - p += offsets_padded; - remaining -= offsets_padded; - - // Read compact chain_lens (uint16_t per chain). - size_t lens_bytes = num_chains * sizeof(uint16_t); - size_t lens_padded = (lens_bytes + 7) & ~size_t(7); - if (remaining < lens_padded) { - return Status::Corruption("Trie index: truncated chain lengths"); - } - s_chain_lens_.resize(num_chains); - memcpy(s_chain_lens_.data(), p, lens_bytes); - p += lens_padded; - remaining -= lens_padded; - - // Read compact chain_end_child_idx (uint32_t per chain). - size_t end_bytes = num_chains * sizeof(uint32_t); - size_t end_padded = (end_bytes + 7) & ~size_t(7); - if (remaining < end_padded) { - return Status::Corruption("Trie index: truncated chain end indices"); - } - s_chain_end_child_idx_.resize(num_chains); - memcpy(s_chain_end_child_idx_.data(), p, end_bytes); - p += end_padded; - remaining -= end_padded; - - // Validate chain end child indices. Each value is either - // UINT32_MAX (chain ends at a leaf) or an index into the child - // position tables (< num_internal). An out-of-bounds index would - // cause s_child_start_pos_[idx] OOB reads during chain traversal. - for (size_t ci = 0; ci < num_chains; ci++) { - uint32_t idx = s_chain_end_child_idx_[ci]; - if (idx != UINT32_MAX && idx >= num_internal) { - return Status::Corruption( - "Trie index: chain end child index out of range"); - } - } - - // Read suffix data blob. - if (remaining < 8) { - return Status::Corruption("Trie index: truncated chain suffix size"); - } - memcpy(&s_chain_suffix_size_, p, 8); - p += 8; - remaining -= 8; - - size_t suffix_padded = (s_chain_suffix_size_ + 7) & ~size_t(7); - if (remaining < suffix_padded) { - return Status::Corruption("Trie index: truncated chain suffix data"); - } - s_chain_suffix_data_ = reinterpret_cast(p); - p += suffix_padded; - remaining -= suffix_padded; - - // Validate that each chain's suffix range fits within the suffix blob. - // A crafted input could have offsets beyond the blob, causing OOB reads - // during SeekImpl when chain suffixes are compared via memcmp. - for (size_t ci = 0; ci < num_chains; ci++) { - uint64_t end = static_cast(s_chain_suffix_offsets_[ci]) + - s_chain_lens_[ci]; - if (end > s_chain_suffix_size_) { - return Status::Corruption( - "Trie index: chain suffix offset + length exceeds suffix data"); - } - } - } - } - - // Block handles: packed uint32_t arrays for offsets and sizes. - if (num_keys_ > 0) { - size_t arr_bytes = num_keys_ * sizeof(uint32_t); - size_t arr_padded = (arr_bytes + 7) & ~size_t(7); - - // Read offsets array. - if (remaining < arr_padded) { - return Status::Corruption("Trie index: truncated handle offsets"); - } - if (reinterpret_cast(p) % alignof(uint32_t) != 0) { - return Status::Corruption("Trie index: handle offsets not aligned"); - } - handle_offsets_ = reinterpret_cast(p); - p += arr_padded; - remaining -= arr_padded; - - // Read sizes array. - if (remaining < arr_padded) { - return Status::Corruption("Trie index: truncated handle sizes"); - } - if (reinterpret_cast(p) % alignof(uint32_t) != 0) { - return Status::Corruption("Trie index: handle sizes not aligned"); - } - handle_sizes_ = reinterpret_cast(p); - p += arr_padded; - remaining -= arr_padded; - } - - // ========================================================================= - // Seqno side-table: deserialized when kFlagSeqnoEncoding is set. - // See SerializeAll() for the format description. - // ========================================================================= - if (has_seqno_encoding_ && num_keys_ > 0) { - // num_keys_ was already validated against kMaxReasonableKeys above. - - // Read num_overflow_blocks (uint32_t) + padding (uint32_t). - if (remaining < 8) { - return Status::Corruption( - "Trie index: truncated seqno side-table header"); - } - memcpy(&num_overflow_blocks_, p, 4); - p += 8; // skip num_overflow + padding - remaining -= 8; - - // Read leaf_seqnos: uint64_t[num_keys_] - { - size_t bytes = num_keys_ * sizeof(uint64_t); - if (remaining < bytes) { - return Status::Corruption("Trie index: truncated leaf seqnos"); - } - if (reinterpret_cast(p) % alignof(uint64_t) != 0) { - return Status::Corruption("Trie index: leaf seqnos not aligned"); - } - leaf_seqnos_ = reinterpret_cast(p); - p += bytes; - remaining -= bytes; - } - - // Read leaf_block_counts: uint32_t[num_keys_] (padded to 8-byte) - { - size_t bytes = num_keys_ * sizeof(uint32_t); - size_t padded = (bytes + 7) & ~size_t(7); - if (remaining < padded) { - return Status::Corruption("Trie index: truncated leaf block counts"); - } - if (reinterpret_cast(p) % alignof(uint32_t) != 0) { - return Status::Corruption("Trie index: leaf block counts not aligned"); - } - leaf_block_counts_ = reinterpret_cast(p); - p += padded; - remaining -= padded; - } - - if (num_overflow_blocks_ > 0) { - // Sanity-check num_overflow_blocks_ to prevent size_t overflow in - // subsequent arithmetic, same rationale as the num_keys_ check above. - if (num_overflow_blocks_ > kMaxReasonableKeys) { - return Status::Corruption( - "Trie index: num_overflow_blocks exceeds reasonable limit"); - } - - // Read overflow_offsets: uint32_t[num_overflow_blocks_] (padded) - { - size_t bytes = num_overflow_blocks_ * sizeof(uint32_t); - size_t padded = (bytes + 7) & ~size_t(7); - if (remaining < padded) { - return Status::Corruption("Trie index: truncated overflow offsets"); - } - if (reinterpret_cast(p) % alignof(uint32_t) != 0) { - return Status::Corruption("Trie index: overflow offsets not aligned"); - } - overflow_offsets_ = reinterpret_cast(p); - p += padded; - remaining -= padded; - } - - // Read overflow_sizes: uint32_t[num_overflow_blocks_] (padded) - { - size_t bytes = num_overflow_blocks_ * sizeof(uint32_t); - size_t padded = (bytes + 7) & ~size_t(7); - if (remaining < padded) { - return Status::Corruption("Trie index: truncated overflow sizes"); - } - if (reinterpret_cast(p) % alignof(uint32_t) != 0) { - return Status::Corruption("Trie index: overflow sizes not aligned"); - } - overflow_sizes_ = reinterpret_cast(p); - p += padded; - remaining -= padded; - } - - // Read overflow_seqnos: uint64_t[num_overflow_blocks_] - // This is the last section in the side-table, so we intentionally do - // not advance p/remaining after reading (no more data to parse). - { - size_t bytes = num_overflow_blocks_ * sizeof(uint64_t); - if (remaining < bytes) { - return Status::Corruption("Trie index: truncated overflow seqnos"); - } - if (reinterpret_cast(p) % alignof(uint64_t) != 0) { - return Status::Corruption("Trie index: overflow seqnos not aligned"); - } - overflow_seqnos_ = reinterpret_cast(p); - } - } - - // Compute overflow_base_ prefix sum for O(1) access to overflow arrays. - // overflow_base_[i] = sum of (block_count-1) for leaves [0, i). - overflow_base_.resize(num_keys_); - uint32_t running_sum = 0; - for (uint64_t i = 0; i < num_keys_; i++) { - if (leaf_block_counts_[i] == 0) { - return Status::Corruption( - "Trie index: leaf block count is zero for leaf " + - std::to_string(i)); - } - overflow_base_[i] = running_sum; - running_sum += leaf_block_counts_[i] - 1; - } - // Verify the prefix sum matches the declared overflow count. - if (running_sum != num_overflow_blocks_) { - return Status::Corruption( - "Trie index: overflow count mismatch with leaf block counts"); - } - - return Status::OK(); - } - - return Status::OK(); -} - -TrieBlockHandle LoudsTrie::GetHandle(uint64_t leaf_index) const { - assert(leaf_index < num_keys_); - return TrieBlockHandle{handle_offsets_[leaf_index], - handle_sizes_[leaf_index]}; -} - -// ============================================================================ -// LoudsTrieIterator implementation -// -// The iterator maintains a stack of positions (path_) from the root to the -// current leaf. Each entry records the position in the bitvector/label array -// at that depth, enabling backtracking for Next() and key reconstruction. -// -// Dense positions: pos = node_num * 256 + label_byte. The node_num is the -// global node number across all dense levels (concatenated). The label_byte -// is pos % 256. -// -// Sparse positions: pos = index into s_labels_ array. -// -// Leaf ordering: BFS leaf order matches sorted key order. For each node in -// BFS order, first its prefix key leaf (if any), then its child leaves -// (has_child = 0) in label order. -// ============================================================================ - -LoudsTrieIterator::LoudsTrieIterator(const LoudsTrie* trie) - : has_chains_(trie->HasChains()), - trie_(trie), - valid_(false), - leaf_index_(0), - key_len_(0), - key_cap_(0), - is_at_prefix_key_(false) { - // Allocate key buffer once, sized to the longest possible trie path. - // MaxDepth() is the length of the longest key, so no traversal will - // exceed this depth. - key_cap_ = trie_->MaxDepth() + 1; - key_buf_ = std::make_unique(key_cap_); -} - -// --- Dense helpers --- - -bool LoudsTrieIterator::DenseSeekLabel(uint64_t node_num, uint8_t target_byte, - uint64_t* out_pos) { - uint64_t pos = node_num * 256 + target_byte; - if (pos < trie_->d_labels_.NumBits() && trie_->d_labels_.GetBit(pos)) { - *out_pos = pos; - return true; - } - uint64_t node_end = (node_num + 1) * 256; - uint64_t next = trie_->d_labels_.NextSetBit(pos + 1); - if (next < node_end && next < trie_->d_labels_.NumBits()) { - *out_pos = next; - return false; - } - *out_pos = node_end; - return false; -} - -uint64_t LoudsTrieIterator::DenseChildNodeNumFromRank( - uint64_t label_rank) const { - // In the concatenated dense model, node 0 is the root (no parent). - // Each internal child (d_has_child_ bit = 1) creates a new node numbered - // 1, 2, 3, ... in BFS order. The child's global node number = - // rank1(d_has_child_, label_rank + 1). - return trie_->d_has_child_.Rank1(label_rank + 1); -} - -uint64_t LoudsTrieIterator::DenseLeafIndexFromRank(uint64_t pos, - uint64_t label_rank) const { - uint64_t has_child_rank = trie_->d_has_child_.Rank1(label_rank + 1); - return DenseLeafIndexFromRankAndHasChildRank(pos, label_rank, has_child_rank); -} - -uint64_t LoudsTrieIterator::DenseLeafIndexFromRankAndHasChildRank( - uint64_t pos, uint64_t label_rank, uint64_t has_child_rank) const { - // Leaf ordinal for a dense leaf at position `pos` in d_labels_. - // - // BFS leaf order: for each node 0..N, prefix key first, then non-internal - // children. So leaf index = - // (prefix keys at nodes [0, node_num]) + - // (non-internal labels at positions [0, label_rank]) - 1 (0-indexed). - // - // When there are no prefix keys, the prefix_keys term is zero and we - // can skip the d_is_prefix_key_.Rank1 call entirely. - uint64_t leaf_labels = (label_rank + 1) - has_child_rank; - if (trie_->d_is_prefix_key_.NumOnes() == 0) { - return leaf_labels - 1; - } - uint64_t node_num = pos / 256; - uint64_t prefix_keys = trie_->d_is_prefix_key_.Rank1(node_num + 1); - return prefix_keys + leaf_labels - 1; -} - -uint64_t LoudsTrieIterator::DensePrefixKeyLeafIndex(uint64_t node_num) const { - // Leaf ordinal for a prefix key at dense node N. The prefix key comes - // after all leaves at nodes [0, N-1] and before leaf labels at node N. - // Index = (prefix keys before N) + (non-internal labels before N). - uint64_t labels_before_node = trie_->d_labels_.Rank1(node_num * 256); - uint64_t internal_before = trie_->d_has_child_.Rank1(labels_before_node); - uint64_t leaf_labels_before = labels_before_node - internal_before; - uint64_t prefix_keys_before = trie_->d_is_prefix_key_.Rank1(node_num); - return prefix_keys_before + leaf_labels_before; -} - -// --- Sparse helpers --- - -bool LoudsTrieIterator::SparseSeekLabel(uint64_t node_start_pos, - uint64_t node_end_pos, - uint8_t target_byte, - uint64_t* out_pos) { - // Labels within a sparse node are stored in sorted order. Find the first - // label >= target_byte. - // - // Strategy: use linear scan for small nodes (<=16 labels) and binary - // search for larger nodes. Profiling shows that the vast majority of - // sparse nodes have small fanout (often 1-10 children), where linear - // scan is faster than std::lower_bound due to sequential memory access - // and predictable branches. The threshold of 16 was chosen because - // 16 bytes fits in a single cache line and binary search gains an - // advantage only around log2(16)=4 comparisons with ~50% branch - // misprediction rate. - static constexpr uint64_t kLinearScanThreshold = 16; - - const uint8_t* base = trie_->s_labels_data_; - const uint8_t* begin = base + node_start_pos; - const uint8_t* end = base + node_end_pos; - uint64_t size = node_end_pos - node_start_pos; - - const uint8_t* it; - if (size <= kLinearScanThreshold) { - // Linear scan: fastest for small nodes. - it = begin; - while (it < end && *it < target_byte) { - ++it; - } - } else { - it = std::lower_bound(begin, end, target_byte); - } - - if (it == end) { - *out_pos = node_end_pos; - return false; - } - *out_pos = static_cast(it - base); - return (*it == target_byte); -} - -uint64_t LoudsTrieIterator::SparseChildNodeNum(uint64_t pos) const { - // The first dense_child_count_ sparse nodes are children of the last - // dense level. Additional sparse internal children add nodes after that. - // child_node = dense_child_count_ + rank1(s_has_child_, pos+1) - 1 - return trie_->dense_child_count_ + trie_->s_has_child_.Rank1(pos + 1) - 1; -} - -uint64_t LoudsTrieIterator::SparseLeafIndex(uint64_t pos) const { - uint64_t has_child_rank = trie_->s_has_child_.Rank1(pos + 1); - return SparseLeafIndexFromHasChildRank(pos, has_child_rank); -} - -uint64_t LoudsTrieIterator::SparseLeafIndexFromHasChildRank( - uint64_t pos, uint64_t has_child_rank) const { - // Leaf ordinal for a sparse leaf at position pos. - // = dense_leaf_count + prefix_keys_at_nodes[0..N] + - // non-internal_labels[0..pos] - 1 - // - // When there are no prefix keys (the common case), the prefix_keys term - // is zero and we can skip both SparseNodeNum (1 Rank1 on s_louds_) and - // s_is_prefix_key_.Rank1, leaving only the precomputed has_child rank. - uint64_t leaf_labels = (pos + 1) - has_child_rank; - - if (trie_->s_is_prefix_key_.NumOnes() == 0) { - return trie_->dense_leaf_count_ + leaf_labels - 1; - } - uint64_t sparse_node = SparseNodeNum(pos); - uint64_t prefix_keys = trie_->s_is_prefix_key_.Rank1(sparse_node + 1); - return trie_->dense_leaf_count_ + prefix_keys + leaf_labels - 1; -} - -uint64_t LoudsTrieIterator::SparsePrefixKeyLeafIndex( - uint64_t sparse_node_num) const { - // Leaf ordinal for a sparse prefix key. Same logic as - // DensePrefixKeyLeafIndex but offset by dense_leaf_count_. - uint64_t start_pos = SparseNodeStartPos(sparse_node_num); - uint64_t prefix_keys_before = trie_->s_is_prefix_key_.Rank1(sparse_node_num); - uint64_t internal_before = trie_->s_has_child_.Rank1(start_pos); - uint64_t leaf_labels_before = start_pos - internal_before; - return trie_->dense_leaf_count_ + prefix_keys_before + leaf_labels_before; -} - -uint64_t LoudsTrieIterator::SparseNodeNum(uint64_t pos) const { - return trie_->s_louds_.Rank1(pos + 1) - 1; -} - -uint64_t LoudsTrieIterator::SparseNodeStartPos(uint64_t sparse_node_num) const { - if (sparse_node_num == 0) { - return 0; - } - - // Use the precomputed lookup table for nodes that are children of sparse - // internal labels. These are the most common case during traversal. - if (sparse_node_num >= trie_->dense_child_count_ && - !trie_->s_child_start_pos_.empty()) { - uint64_t internal_idx = sparse_node_num - trie_->dense_child_count_; - if (internal_idx < trie_->s_child_start_pos_.size()) { - return trie_->s_child_start_pos_[internal_idx]; - } - } - - // Fallback to FindNthOneBit for: - // - Root sparse nodes (children of dense nodes) - // - Very large tries where lookup table wasn't built - return trie_->s_louds_.FindNthOneBit(sparse_node_num); -} - -uint64_t LoudsTrieIterator::SparseNodeEndPos(uint64_t start_pos) const { - uint64_t next = trie_->s_louds_.NextSetBit(start_pos + 1); - if (next >= trie_->s_louds_.NumBits()) { - return trie_->s_labels_size_; - } - return next; -} - -// Helper: descend from (in_dense, node_num) to leftmost leaf, pushing -// path entries and building key_buf_. Returns true if a leaf was found. -bool LoudsTrieIterator::DescendToLeftmostLeaf(bool in_dense, - uint64_t node_num) { - while (true) { - if (in_dense) { - // Check prefix key first. - if (trie_->d_is_prefix_key_.NumBits() > 0 && - node_num < trie_->d_is_prefix_key_.NumBits() && - trie_->d_is_prefix_key_.GetBit(node_num)) { - is_at_prefix_key_ = true; - leaf_index_ = DensePrefixKeyLeafIndex(node_num); - valid_ = true; - return true; - } - - uint64_t base = node_num * 256; - if (base >= trie_->d_labels_.NumBits()) { - valid_ = false; - return false; - } - uint64_t first = trie_->d_labels_.NextSetBit(base); - if (first >= base + 256 || first >= trie_->d_labels_.NumBits()) { - valid_ = false; - return false; - } - - path_.push_back(LevelPos::MakeDense(first)); - AppendKeySlot() = static_cast(first % 256); - - uint64_t label_rank = trie_->d_labels_.Rank1(first + 1) - 1; - if (!trie_->d_has_child_.GetBit(label_rank)) { - leaf_index_ = DenseLeafIndexFromRank(first, label_rank); - valid_ = true; - return true; - } - - uint64_t child = DenseChildNodeNumFromRank(label_rank); - if (child < trie_->dense_node_count_) { - node_num = child; - in_dense = true; - } else { - node_num = child - trie_->dense_node_count_; - in_dense = false; - } - } else { - // Check prefix key first. - if (trie_->s_is_prefix_key_.NumBits() > 0 && - node_num < trie_->s_is_prefix_key_.NumBits() && - trie_->s_is_prefix_key_.GetBit(node_num)) { - is_at_prefix_key_ = true; - leaf_index_ = SparsePrefixKeyLeafIndex(node_num); - valid_ = true; - return true; - } - - uint64_t start = SparseNodeStartPos(node_num); - if (start >= trie_->s_labels_size_) { - valid_ = false; - return false; - } - - path_.push_back(LevelPos::MakeSparse(start)); - AppendKeySlot() = static_cast(trie_->s_labels_data_[start]); - - if (!trie_->s_has_child_.GetBit(start)) { - leaf_index_ = SparseLeafIndex(start); - valid_ = true; - return true; - } - - node_num = SparseChildNodeNum(start); - in_dense = false; - } - } -} - -bool LoudsTrieIterator::SeekToFirst() { - valid_ = false; - leaf_index_ = 0; - key_len_ = 0; - path_.clear(); - is_at_prefix_key_ = false; - - if (trie_->NumKeys() == 0) { - return false; - } - - // Descend directly from root to the leftmost leaf. - // Compared to Seek(""), this skips the SeekImpl target-consumption loop - // (a no-op for empty target) and the redundant prefix key check that - // SeekImpl performs at the root before calling DescendToLeftmostLeaf. - // DescendToLeftmostLeaf itself handles prefix keys at every node. - bool in_dense = (trie_->cutoff_level_ > 0); - return DescendToLeftmostLeaf(in_dense, /*node_num=*/0); -} - -// Main Seek implementation. -// Uses SuRF-style Select-free traversal for sparse regions: instead of -// tracking node_num and calling FindNthOneBit to find node boundaries, we -// track (start_pos, end_pos) directly and use only Rank1 + array lookup. -template -bool LoudsTrieIterator::SeekImpl(const Slice& target) { - valid_ = false; - leaf_index_ = 0; - key_len_ = 0; - path_.clear(); - is_at_prefix_key_ = false; - - if (trie_->NumKeys() == 0) { - return false; - } - - bool in_dense = (trie_->cutoff_level_ > 0); - uint64_t node_num = 0; - - // SuRF-style: For sparse traversal, track (start_pos, end_pos) directly - // instead of node_num. This eliminates FindNthOneBit calls entirely. - uint64_t sparse_start = 0; - uint64_t sparse_end = 0; - bool have_sparse_bounds = false; - - assert(target.size() <= UINT32_MAX); - for (uint32_t depth = 0; depth < static_cast(target.size()); - depth++) { - uint8_t target_byte = static_cast(target[depth]); - - if (in_dense) { - uint64_t pos; - bool exact = DenseSeekLabel(node_num, target_byte, &pos); - uint64_t node_end = (node_num + 1) * 256; - - if (pos >= node_end) { - // No label >= target_byte. Backtrack. - return Advance(); - } - - path_.push_back(LevelPos::MakeDense(pos)); - AppendKeySlot() = static_cast(pos % 256); - - // Cache label_rank: avoids redundant Rank1(d_labels_) in both - // has_child check and DenseChildNodeNumFromRank/DenseLeafIndexFromRank. - uint64_t label_rank = trie_->d_labels_.Rank1(pos + 1) - 1; - bool is_internal = trie_->d_has_child_.GetBit(label_rank); - // Compute has_child_rank once, reuse for both leaf index and child - // node number. DenseChildNodeNumFromRank(lr) = Rank1(lr + 1) and - // DenseLeafIndexFromRankAndHasChildRank also needs Rank1(lr + 1). - uint64_t has_child_rank = trie_->d_has_child_.Rank1(label_rank + 1); - - if (!exact) { - // Landed on label > target_byte. Go to leftmost leaf in subtree. - if (!is_internal) { - leaf_index_ = DenseLeafIndexFromRankAndHasChildRank(pos, label_rank, - has_child_rank); - valid_ = true; - return true; - } - uint64_t child = has_child_rank; // = DenseChildNodeNumFromRank(lr) - bool child_dense = (child < trie_->dense_node_count_); - uint64_t cn = child_dense ? child : child - trie_->dense_node_count_; - return DescendToLeftmostLeaf(child_dense, cn); - } - - if (!is_internal) { - // This label is a leaf. Check if target is fully consumed. - if (depth == static_cast(target.size()) - 1) { - // Exact match: trie key == target. - leaf_index_ = DenseLeafIndexFromRankAndHasChildRank(pos, label_rank, - has_child_rank); - valid_ = true; - return true; - } - // Target has more bytes, so trie key is a proper prefix of target. - // Trie key < target. Advance to the NEXT leaf. - leaf_index_ = DenseLeafIndexFromRankAndHasChildRank(pos, label_rank, - has_child_rank); - valid_ = true; - return Advance(); - } - - uint64_t child = has_child_rank; // = DenseChildNodeNumFromRank(lr) - if (child < trie_->dense_node_count_) { - node_num = child; - in_dense = true; - } else { - node_num = child - trie_->dense_node_count_; - in_dense = false; - have_sparse_bounds = false; // Will compute on first sparse access - } - } else { - // SuRF-style sparse traversal: Use (start, end) positions directly. - // No FindNthOneBit calls - only Rank1 + array lookup. - uint64_t start; - uint64_t end; - if (have_sparse_bounds) { - start = sparse_start; - end = sparse_end; - } else { - // First entry into sparse region from dense, or from root. - // Need to compute bounds using FindNthOneBit (only once per - // dense->sparse). - start = SparseNodeStartPos(node_num); - end = SparseNodeEndPos(start); - } - - // Fast path for fanout-1 nodes: the most common case for tries built - // from keys with long common prefixes (e.g., zero-padded numeric keys). - // Avoids the SparseSeekLabel function call and reduces branch logic to - // a single byte comparison. - if (start + 1 == end) { - uint8_t label = trie_->s_labels_data_[start]; - - if (label == target_byte) { - // Exact match (the overwhelmingly common case for fanout-1). - path_.push_back(LevelPos::MakeSparse(start)); - AppendKeySlot() = static_cast(label); - - bool is_internal = trie_->s_has_child_.GetBit(start); - uint64_t has_child_rank = trie_->s_has_child_.Rank1(start + 1); - if (is_internal) { - // Internal node: use rank for child lookup. - if (!trie_->s_child_start_pos_.empty()) { - uint64_t child_idx = has_child_rank - 1; - if (child_idx < trie_->s_child_start_pos_.size()) { - // ---- Path compression: chain skip ---- - // Check if this child starts a fanout-1 chain. If so, compare - // the remaining target bytes against the chain suffix with a - // single memcmp instead of traversing level by level. - // Guarded by if constexpr: when kHasChains=false, the - // compiler eliminates this entire block from the generated - // code. - if constexpr (kHasChains) { - if (child_idx < trie_->s_chain_bitmap_.NumBits() && - trie_->s_chain_bitmap_.GetBit(child_idx)) { - uint64_t chain_idx = - trie_->s_chain_bitmap_.Rank1(child_idx + 1) - 1; - uint16_t chain_len = trie_->s_chain_lens_[chain_idx]; - const uint8_t* suffix = - trie_->s_chain_suffix_data_ + - trie_->s_chain_suffix_offsets_[chain_idx]; - uint32_t target_remaining = - static_cast(target.size()) - depth - 1; - - if (target_remaining >= chain_len) { - // Target has enough bytes to cover the full chain. - const uint8_t* target_bytes = - reinterpret_cast(target.data()) + - depth + 1; - int cmp = memcmp(target_bytes, suffix, chain_len); - if (cmp == 0) { - // Full chain match! Push all chain nodes onto path_. - // Walk the chain using child position tables to get - // each node's label position for path reconstruction. - uint64_t cur_idx = child_idx; - for (uint16_t ci = 0; ci < chain_len; ci++) { - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - // Move to next chain node (last node handled - // below). - if (ci + 1 < chain_len) { - cur_idx = trie_->s_has_child_.Rank1(cs + 1) - 1; - } - } - // Chain fully matched. Advance depth past the chain. - depth += chain_len; - - // Set up for the node after the chain. - uint32_t end_child_idx = - trie_->s_chain_end_child_idx_[chain_idx]; - if (end_child_idx == UINT32_MAX) { - // Chain ends at a leaf. Check if target is - // consumed. - uint32_t last_cs = trie_->s_child_start_pos_[cur_idx]; - uint64_t last_hcr = - trie_->s_has_child_.Rank1(last_cs + 1); - if (depth == - static_cast(target.size()) - 1) { - leaf_index_ = SparseLeafIndexFromHasChildRank( - last_cs, last_hcr); - valid_ = true; - return true; - } - // Target has more bytes, trie key < target. - // Advance. - leaf_index_ = SparseLeafIndexFromHasChildRank( - last_cs, last_hcr); - valid_ = true; - return Advance(); - } - // Chain ends at an internal node with fanout > 1. - sparse_start = trie_->s_child_start_pos_[end_child_idx]; - sparse_end = trie_->s_child_end_pos_[end_child_idx]; - have_sparse_bounds = true; - in_dense = false; - continue; - } - // Chain mismatch. Find the divergence point. - uint16_t mismatch_pos = 0; - while (mismatch_pos < chain_len && - target_bytes[mismatch_pos] == - suffix[mismatch_pos]) { - mismatch_pos++; - } - // Push path entries up to the mismatch point. - uint64_t cur_idx = child_idx; - for (uint16_t ci = 0; ci < mismatch_pos; ci++) { - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - cur_idx = trie_->s_has_child_.Rank1(cs + 1) - 1; - } - // At the mismatch node: push the node's label and - // handle. - uint32_t mis_cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(mis_cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[mis_cs]); - - if (target_bytes[mismatch_pos] < suffix[mismatch_pos]) { - // Target < chain. Chain's path leads to the first key - // >= target. Descend to leftmost leaf from here. - if (!trie_->s_has_child_.GetBit(mis_cs)) { - leaf_index_ = SparseLeafIndex(mis_cs); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf( - false, SparseChildNodeNum(mis_cs)); - } - // target_bytes[mismatch_pos] > suffix[mismatch_pos]: - // All keys through this chain node are < target. - // Advance. - return Advance(); - } - // Target runs out before chain ends. - // Check if target's remaining bytes match the chain - // prefix. - if (target_remaining > 0) { - const uint8_t* target_bytes = - reinterpret_cast(target.data()) + - depth + 1; - int cmp = memcmp(target_bytes, suffix, target_remaining); - if (cmp < 0) { - // Target < chain prefix. Push matched portion. - uint64_t cur_idx = child_idx; - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - // Descend to leftmost leaf from this first chain - // node. - if (!trie_->s_has_child_.GetBit(cs)) { - leaf_index_ = SparseLeafIndex(cs); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, - SparseChildNodeNum(cs)); - } - if (cmp > 0) { - // Target > chain prefix at some point. We need to - // find the exact divergence point. - uint16_t mismatch_pos = 0; - while (mismatch_pos < target_remaining && - target_bytes[mismatch_pos] == - suffix[mismatch_pos]) { - mismatch_pos++; - } - // Push path entries up to the mismatch. - uint64_t cur_idx = child_idx; - for (uint16_t ci = 0; ci < mismatch_pos; ci++) { - uint32_t cs2 = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs2)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs2]); - cur_idx = trie_->s_has_child_.Rank1(cs2 + 1) - 1; - } - uint32_t mis_cs2 = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(mis_cs2)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[mis_cs2]); - return Advance(); - } - // cmp == 0: target matches chain prefix exactly. Target - // is fully consumed. Push matched nodes and check - // prefix key / descend to leftmost leaf. - uint64_t cur_idx = child_idx; - for (uint32_t ci = 0; ci < target_remaining; ci++) { - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - if (ci + 1 < target_remaining) { - cur_idx = trie_->s_has_child_.Rank1(cs + 1) - 1; - } - } - // Target consumed mid-chain. The remaining chain nodes - // form keys > target. Descend to leftmost leaf from the - // next chain node. - uint32_t last_cs = trie_->s_child_start_pos_[cur_idx]; - // This node is always internal (it's mid-chain). - return DescendToLeftmostLeaf(false, - SparseChildNodeNum(last_cs)); - } - // target_remaining == 0: target fully consumed. - // This means the target ended exactly at the parent node. - // The chain nodes are all > target. Push first chain node - // and descend to leftmost leaf. - uint32_t cs = trie_->s_child_start_pos_[child_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - if (!trie_->s_has_child_.GetBit(cs)) { - leaf_index_ = SparseLeafIndex(cs); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, SparseChildNodeNum(cs)); - } - } // if constexpr (kHasChains) - // No chain -- normal child lookup. - sparse_start = trie_->s_child_start_pos_[child_idx]; - sparse_end = trie_->s_child_end_pos_[child_idx]; - have_sparse_bounds = true; - } else { - node_num = SparseChildNodeNum(start); - have_sparse_bounds = false; - } - } else { - node_num = SparseChildNodeNum(start); - have_sparse_bounds = false; - } - in_dense = false; - continue; - } - // Leaf node. - if (depth == static_cast(target.size()) - 1) { - leaf_index_ = - SparseLeafIndexFromHasChildRank(start, has_child_rank); - valid_ = true; - return true; - } - leaf_index_ = SparseLeafIndexFromHasChildRank(start, has_child_rank); - valid_ = true; - return Advance(); - } - - // label != target_byte. Still need to push path for backtracking. - path_.push_back(LevelPos::MakeSparse(start)); - AppendKeySlot() = static_cast(label); - - if (label > target_byte) { - // Label is greater: go to leftmost leaf in subtree. - bool is_internal = trie_->s_has_child_.GetBit(start); - uint64_t has_child_rank = trie_->s_has_child_.Rank1(start + 1); - if (!is_internal) { - leaf_index_ = - SparseLeafIndexFromHasChildRank(start, has_child_rank); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, SparseChildNodeNum(start)); - } - // label < target_byte: no label >= target in this node. Backtrack. - return Advance(); - } - - // General path for nodes with fanout > 1. - uint64_t pos; - bool exact = SparseSeekLabel(start, end, target_byte, &pos); - - if (pos >= end) { - return Advance(); - } - - path_.push_back(LevelPos::MakeSparse(pos)); - AppendKeySlot() = static_cast(trie_->s_labels_data_[pos]); - - bool is_internal = trie_->s_has_child_.GetBit(pos); - uint64_t has_child_rank = trie_->s_has_child_.Rank1(pos + 1); - - if (!exact) { - if (!is_internal) { - leaf_index_ = SparseLeafIndexFromHasChildRank(pos, has_child_rank); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, SparseChildNodeNum(pos)); - } - - if (!is_internal) { - // Check if target is fully consumed. - if (depth == static_cast(target.size()) - 1) { - leaf_index_ = SparseLeafIndexFromHasChildRank(pos, has_child_rank); - valid_ = true; - return true; - } - // Target has more bytes. Trie key < target. Advance. - leaf_index_ = SparseLeafIndexFromHasChildRank(pos, has_child_rank); - valid_ = true; - return Advance(); - } - - // Descend to child: Get child bounds using Rank1 + array lookup. - // This is the key SuRF optimization - NO FindNthOneBit here! - // Reuse the already-computed has_child_rank. - { - uint64_t child_idx = has_child_rank - 1; - if (!trie_->s_child_start_pos_.empty() && - child_idx < trie_->s_child_start_pos_.size()) { - // ---- Path compression: chain skip (general path) ---- - if constexpr (kHasChains) { - if (child_idx < trie_->s_chain_bitmap_.NumBits() && - trie_->s_chain_bitmap_.GetBit(child_idx)) { - uint64_t chain_idx = - trie_->s_chain_bitmap_.Rank1(child_idx + 1) - 1; - uint16_t chain_len = trie_->s_chain_lens_[chain_idx]; - const uint8_t* suffix = trie_->s_chain_suffix_data_ + - trie_->s_chain_suffix_offsets_[chain_idx]; - uint32_t target_remaining = - static_cast(target.size()) - depth - 1; - - if (target_remaining >= chain_len) { - const uint8_t* target_bytes = - reinterpret_cast(target.data()) + depth + 1; - int cmp = memcmp(target_bytes, suffix, chain_len); - if (cmp == 0) { - // Full chain match! Push all chain nodes. - uint64_t cur_idx = child_idx; - for (uint16_t ci = 0; ci < chain_len; ci++) { - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - if (ci + 1 < chain_len) { - cur_idx = trie_->s_has_child_.Rank1(cs + 1) - 1; - } - } - depth += chain_len; - - uint32_t end_child_idx = - trie_->s_chain_end_child_idx_[chain_idx]; - if (end_child_idx == UINT32_MAX) { - uint32_t last_cs = trie_->s_child_start_pos_[cur_idx]; - uint64_t last_hcr = trie_->s_has_child_.Rank1(last_cs + 1); - if (depth == static_cast(target.size()) - 1) { - leaf_index_ = - SparseLeafIndexFromHasChildRank(last_cs, last_hcr); - valid_ = true; - return true; - } - leaf_index_ = - SparseLeafIndexFromHasChildRank(last_cs, last_hcr); - valid_ = true; - return Advance(); - } - sparse_start = trie_->s_child_start_pos_[end_child_idx]; - sparse_end = trie_->s_child_end_pos_[end_child_idx]; - have_sparse_bounds = true; - in_dense = false; - continue; - } - // Mismatch: find divergence point. - uint16_t mismatch_pos = 0; - while (mismatch_pos < chain_len && - target_bytes[mismatch_pos] == suffix[mismatch_pos]) { - mismatch_pos++; - } - uint64_t cur_idx = child_idx; - for (uint16_t ci = 0; ci < mismatch_pos; ci++) { - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - cur_idx = trie_->s_has_child_.Rank1(cs + 1) - 1; - } - uint32_t mis_cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(mis_cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[mis_cs]); - - if (target_bytes[mismatch_pos] < suffix[mismatch_pos]) { - if (!trie_->s_has_child_.GetBit(mis_cs)) { - leaf_index_ = SparseLeafIndex(mis_cs); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, - SparseChildNodeNum(mis_cs)); - } - return Advance(); - } - // Target runs out before chain ends. - if (target_remaining > 0) { - const uint8_t* target_bytes = - reinterpret_cast(target.data()) + depth + 1; - int cmp = memcmp(target_bytes, suffix, target_remaining); - if (cmp < 0) { - uint32_t cs = trie_->s_child_start_pos_[child_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - if (!trie_->s_has_child_.GetBit(cs)) { - leaf_index_ = SparseLeafIndex(cs); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, SparseChildNodeNum(cs)); - } - if (cmp > 0) { - uint16_t mp = 0; - while (mp < target_remaining && - target_bytes[mp] == suffix[mp]) { - mp++; - } - uint64_t cur_idx = child_idx; - for (uint16_t ci = 0; ci < mp; ci++) { - uint32_t cs2 = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs2)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs2]); - cur_idx = trie_->s_has_child_.Rank1(cs2 + 1) - 1; - } - uint32_t mis_cs2 = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(mis_cs2)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[mis_cs2]); - return Advance(); - } - // Prefix match: target consumed mid-chain. - uint64_t cur_idx = child_idx; - for (uint32_t ci = 0; ci < target_remaining; ci++) { - uint32_t cs = trie_->s_child_start_pos_[cur_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = - static_cast(trie_->s_labels_data_[cs]); - if (ci + 1 < target_remaining) { - cur_idx = trie_->s_has_child_.Rank1(cs + 1) - 1; - } - } - uint32_t last_cs = trie_->s_child_start_pos_[cur_idx]; - return DescendToLeftmostLeaf(false, - SparseChildNodeNum(last_cs)); - } - // target_remaining == 0: first chain node is > target. - uint32_t cs = trie_->s_child_start_pos_[child_idx]; - path_.push_back(LevelPos::MakeSparse(cs)); - AppendKeySlot() = static_cast(trie_->s_labels_data_[cs]); - if (!trie_->s_has_child_.GetBit(cs)) { - leaf_index_ = SparseLeafIndex(cs); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, SparseChildNodeNum(cs)); - } - } // if constexpr (kHasChains) - // No chain -- normal child lookup. - sparse_start = trie_->s_child_start_pos_[child_idx]; - sparse_end = trie_->s_child_end_pos_[child_idx]; - have_sparse_bounds = true; - } else { - node_num = SparseChildNodeNum(pos); - have_sparse_bounds = false; - } - } - in_dense = false; - } - } - - // Target key fully consumed. Check if current node is a prefix key. - if (in_dense) { - if (trie_->d_is_prefix_key_.NumBits() > 0 && - node_num < trie_->d_is_prefix_key_.NumBits() && - trie_->d_is_prefix_key_.GetBit(node_num)) { - is_at_prefix_key_ = true; - leaf_index_ = DensePrefixKeyLeafIndex(node_num); - valid_ = true; - return true; - } - } else { - // For prefix key check, we need node_num. Compute it if we were tracking - // positions directly. - if (have_sparse_bounds) { - // Compute node_num from sparse_start position. - // node_num = number of nodes before this position = Rank1(s_louds, pos) - // But sparse_start IS the start of the node, so: - // node_num = Rank1(s_louds_, sparse_start + 1) - 1 - node_num = trie_->s_louds_.Rank1(sparse_start + 1) - 1; - } - if (trie_->s_is_prefix_key_.NumBits() > 0 && - node_num < trie_->s_is_prefix_key_.NumBits() && - trie_->s_is_prefix_key_.GetBit(node_num)) { - is_at_prefix_key_ = true; - leaf_index_ = SparsePrefixKeyLeafIndex(node_num); - valid_ = true; - return true; - } - } - - // Descend to leftmost leaf. - return DescendToLeftmostLeaf(in_dense, node_num); -} - -bool LoudsTrieIterator::Next() { - if (!valid_) { - return false; - } - - if (is_at_prefix_key_) { - is_at_prefix_key_ = false; - // The prefix key is at a node that also has children. The next leaf - // is the leftmost child leaf. - bool in_dense; - uint64_t node_num; - if (path_.empty()) { - // Root is the prefix key. - in_dense = (trie_->cutoff_level_ > 0); - node_num = 0; - } else { - // The last path entry is the label that leads TO this prefix key node. - // The prefix key node IS the child of that label. - auto last = path_.back(); - if (last.is_dense()) { - // Compute label_rank once and reuse for child node lookup. - uint64_t lr = trie_->d_labels_.Rank1(last.pos() + 1) - 1; - uint64_t child = DenseChildNodeNumFromRank(lr); - in_dense = (child < trie_->dense_node_count_); - node_num = in_dense ? child : child - trie_->dense_node_count_; - } else { - node_num = SparseChildNodeNum(last.pos()); - in_dense = false; - } - } - - // Find leftmost child leaf (NOT checking prefix key again, since we - // just came from it). - if (in_dense) { - uint64_t base = node_num * 256; - if (base >= trie_->d_labels_.NumBits()) { - return Advance(); - } - uint64_t first = trie_->d_labels_.NextSetBit(base); - if (first >= base + 256 || first >= trie_->d_labels_.NumBits()) { - return Advance(); - } - path_.push_back(LevelPos::MakeDense(first)); - AppendKeySlot() = static_cast(first % 256); - - uint64_t label_rank = trie_->d_labels_.Rank1(first + 1) - 1; - if (!trie_->d_has_child_.GetBit(label_rank)) { - leaf_index_ = DenseLeafIndexFromRank(first, label_rank); - valid_ = true; - return true; - } - uint64_t child = DenseChildNodeNumFromRank(label_rank); - bool cd = (child < trie_->dense_node_count_); - return DescendToLeftmostLeaf( - cd, cd ? child : child - trie_->dense_node_count_); - } else { - uint64_t start = SparseNodeStartPos(node_num); - if (start >= trie_->s_labels_size_) { - return Advance(); - } - path_.push_back(LevelPos::MakeSparse(start)); - AppendKeySlot() = static_cast(trie_->s_labels_data_[start]); - - if (!trie_->s_has_child_.GetBit(start)) { - leaf_index_ = SparseLeafIndex(start); - valid_ = true; - return true; - } - return DescendToLeftmostLeaf(false, SparseChildNodeNum(start)); - } - } - - return Advance(); -} - -bool LoudsTrieIterator::Advance() { - // Backtrack up the path to find the next sibling, then descend to the - // leftmost leaf in that subtree. When a sibling is found, replace the - // top of the path stack in-place instead of pop+push (avoids 2 autovector - // size modifications + 2 key_len_ changes per level in the common case). - while (!path_.empty()) { - auto& cur = path_.back(); - - if (cur.is_dense()) { - uint64_t cur_pos = cur.pos(); - uint64_t node_num = cur_pos / 256; - uint64_t node_end = (node_num + 1) * 256; - uint64_t next = trie_->d_labels_.NextSetBit(cur_pos + 1); - - if (next < node_end && next < trie_->d_labels_.NumBits()) { - assert(key_len_ > 0); - cur = LevelPos::MakeDense(next); - key_buf_[key_len_ - 1] = static_cast(next % 256); - - uint64_t label_rank = trie_->d_labels_.Rank1(next + 1) - 1; - if (!trie_->d_has_child_.GetBit(label_rank)) { - leaf_index_ = DenseLeafIndexFromRank(next, label_rank); - valid_ = true; - return true; - } - - uint64_t child = DenseChildNodeNumFromRank(label_rank); - bool cd = (child < trie_->dense_node_count_); - return DescendToLeftmostLeaf( - cd, cd ? child : child - trie_->dense_node_count_); - } - } else { - uint64_t next_pos = cur.pos() + 1; - if (next_pos < trie_->s_labels_size_ && - !trie_->s_louds_.GetBit(next_pos)) { - assert(key_len_ > 0); - cur = LevelPos::MakeSparse(next_pos); - key_buf_[key_len_ - 1] = - static_cast(trie_->s_labels_data_[next_pos]); - - if (!trie_->s_has_child_.GetBit(next_pos)) { - leaf_index_ = SparseLeafIndex(next_pos); - valid_ = true; - return true; - } - - return DescendToLeftmostLeaf(false, SparseChildNodeNum(next_pos)); - } - } - - // No sibling found at this level -- pop and continue backtracking. - path_.pop_back(); - if (key_len_ > 0) { - key_len_--; - } - } - - valid_ = false; - return false; -} - -bool LoudsTrieIterator::DescendToRightmostLeaf(bool in_dense, - uint64_t node_num) { - // Mirror of DescendToLeftmostLeaf: go to the rightmost (largest) child - // at each level. Does NOT check prefix keys -- they are the smallest leaf - // at a node, so they come last in reverse order and are handled by - // Retreat() when backtracking past all children. - - // Sparse child end position from the precomputed lookup table. When - // available, avoids SparseNodeEndPos's NextSetBit scan. - uint64_t sparse_end = 0; - bool have_sparse_end = false; - - while (true) { - if (in_dense) { - // Find the LAST set bit in this node's 256-bit bitmap. - uint64_t node_end = (node_num + 1) * 256; - if (node_end > trie_->d_labels_.NumBits()) { - node_end = trie_->d_labels_.NumBits(); - } - uint64_t base = node_num * 256; - if (base >= trie_->d_labels_.NumBits()) { - valid_ = false; - return false; - } - uint64_t last = trie_->d_labels_.PrevSetBit(node_end); - if (last >= trie_->d_labels_.NumBits() || last < base) { - // Empty node -- shouldn't happen in a valid trie, but guard anyway. - valid_ = false; - return false; - } - - path_.push_back(LevelPos::MakeDense(last)); - AppendKeySlot() = static_cast(last % 256); - - uint64_t label_rank = trie_->d_labels_.Rank1(last + 1) - 1; - if (!trie_->d_has_child_.GetBit(label_rank)) { - leaf_index_ = DenseLeafIndexFromRank(last, label_rank); - valid_ = true; - return true; - } - - uint64_t child = DenseChildNodeNumFromRank(label_rank); - if (child < trie_->dense_node_count_) { - node_num = child; - in_dense = true; - } else { - node_num = child - trie_->dense_node_count_; - in_dense = false; - } - } else { - // Sparse: last label is at node_end - 1. - uint64_t end; - if (have_sparse_end) { - end = sparse_end; - have_sparse_end = false; - } else { - uint64_t start = SparseNodeStartPos(node_num); - if (start >= trie_->s_labels_size_) { - valid_ = false; - return false; - } - end = SparseNodeEndPos(start); - } - uint64_t last_pos = end - 1; - - path_.push_back(LevelPos::MakeSparse(last_pos)); - AppendKeySlot() = static_cast(trie_->s_labels_data_[last_pos]); - - if (!trie_->s_has_child_.GetBit(last_pos)) { - leaf_index_ = SparseLeafIndex(last_pos); - valid_ = true; - return true; - } - - // Use the precomputed child lookup table when available to avoid - // SparseNodeEndPos's NextSetBit scan on the next iteration. - uint64_t child_idx = trie_->s_has_child_.Rank1(last_pos + 1) - 1; - if (child_idx < trie_->s_child_end_pos_.size()) { - sparse_end = trie_->s_child_end_pos_[child_idx]; - have_sparse_end = true; - } - node_num = trie_->dense_child_count_ + child_idx; - in_dense = false; - } - } -} - -bool LoudsTrieIterator::SeekToLast() { - valid_ = false; - leaf_index_ = 0; - key_len_ = 0; - path_.clear(); - is_at_prefix_key_ = false; - - if (trie_->NumKeys() == 0) { - return false; - } - - bool in_dense = (trie_->cutoff_level_ > 0); - return DescendToRightmostLeaf(in_dense, /*node_num=*/0); -} - -bool LoudsTrieIterator::Retreat() { - // Backtrack up the path to find the previous sibling, then descend to the - // rightmost leaf in that subtree. When a sibling is found, replace the - // top of the path stack in-place (same optimization as Advance). When no - // previous sibling exists, check the prefix key before popping. - while (!path_.empty()) { - auto& cur = path_.back(); - - if (cur.is_dense()) { - uint64_t cur_pos = cur.pos(); - uint64_t node_num = cur_pos / 256; - uint64_t node_base = node_num * 256; - - // Find previous set bit within the same node's 256-bit bitmap. - uint64_t prev = trie_->d_labels_.PrevSetBit(cur_pos); - if (prev < trie_->d_labels_.NumBits() && prev >= node_base) { - // Found a previous sibling -- replace in-place. - assert(key_len_ > 0); - cur = LevelPos::MakeDense(prev); - key_buf_[key_len_ - 1] = static_cast(prev % 256); - - uint64_t label_rank = trie_->d_labels_.Rank1(prev + 1) - 1; - if (!trie_->d_has_child_.GetBit(label_rank)) { - leaf_index_ = DenseLeafIndexFromRank(prev, label_rank); - valid_ = true; - return true; - } - - uint64_t child = DenseChildNodeNumFromRank(label_rank); - bool cd = (child < trie_->dense_node_count_); - return DescendToRightmostLeaf( - cd, cd ? child : child - trie_->dense_node_count_); - } - - // No previous sibling -- check prefix key at this node. - // Prefix keys are the smallest leaf at a node, so in reverse order - // they come after all children. - if (trie_->d_is_prefix_key_.NumBits() > 0 && - node_num < trie_->d_is_prefix_key_.NumBits() && - trie_->d_is_prefix_key_.GetBit(node_num)) { - // Pop the current level before returning the prefix key -- the - // prefix key doesn't have its own path entry. - path_.pop_back(); - if (key_len_ > 0) { - key_len_--; - } - is_at_prefix_key_ = true; - leaf_index_ = DensePrefixKeyLeafIndex(node_num); - valid_ = true; - return true; - } - } else { - // Sparse: check if current position is the first label of its node. - uint64_t cur_sparse_pos = cur.pos(); - if (!trie_->s_louds_.GetBit(cur_sparse_pos)) { - // Previous sibling exists at cur.pos() - 1 -- replace in-place. - assert(key_len_ > 0); - assert(cur_sparse_pos > 0); - uint64_t prev_pos = cur_sparse_pos - 1; - cur = LevelPos::MakeSparse(prev_pos); - key_buf_[key_len_ - 1] = - static_cast(trie_->s_labels_data_[prev_pos]); - - if (!trie_->s_has_child_.GetBit(prev_pos)) { - leaf_index_ = SparseLeafIndex(prev_pos); - valid_ = true; - return true; - } - - return DescendToRightmostLeaf(false, SparseChildNodeNum(prev_pos)); - } - - // At first label -- no previous sibling. Check prefix key. - uint64_t sparse_node = SparseNodeNum(cur_sparse_pos); - if (trie_->s_is_prefix_key_.NumBits() > 0 && - sparse_node < trie_->s_is_prefix_key_.NumBits() && - trie_->s_is_prefix_key_.GetBit(sparse_node)) { - path_.pop_back(); - if (key_len_ > 0) { - key_len_--; - } - is_at_prefix_key_ = true; - leaf_index_ = SparsePrefixKeyLeafIndex(sparse_node); - valid_ = true; - return true; - } - } - - // No sibling and no prefix key -- pop and continue backtracking. - path_.pop_back(); - if (key_len_ > 0) { - key_len_--; - } - } - - valid_ = false; - return false; -} - -bool LoudsTrieIterator::Prev() { - if (!valid_) { - return false; - } - // Clear prefix key flag before retreating. If we're at a prefix key, the - // path is already at the parent level (prefix keys don't have their own - // path entry), so Retreat() continues backtracking from there. - is_at_prefix_key_ = false; - return Retreat(); -} - -TrieBlockHandle LoudsTrieIterator::Value() const { - assert(valid_); - return trie_->GetHandle(leaf_index_); -} - -// Explicit template instantiations for SeekImpl. -template bool LoudsTrieIterator::SeekImpl(const Slice&); -template bool LoudsTrieIterator::SeekImpl(const Slice&); - -} // namespace trie_index -} // namespace ROCKSDB_NAMESPACE diff --git a/utilities/trie_index/louds_trie.h b/utilities/trie_index/louds_trie.h deleted file mode 100644 index 5949d8e074d3..000000000000 --- a/utilities/trie_index/louds_trie.h +++ /dev/null @@ -1,713 +0,0 @@ -// Copyright (c) Meta Platforms, Inc. and affiliates. -// This source code is licensed under both the GPLv2 (found in the -// COPYING file in the root directory) and Apache 2.0 License -// (found in the LICENSE.Apache file in the root directory). -// -// ***************************************************************** -// EXPERIMENTAL - subject to change while under development -// ***************************************************************** -// -// Fast Succinct Trie (FST) implementation based on the LOUDS (Level-Order -// Unary Degree Sequence) encoding, inspired by the SuRF paper (Zhang et al., -// SIGMOD 2018). The trie uses a hybrid encoding: -// -// - LOUDS-Dense: For the upper levels of the trie (levels close to the root) -// where the fanout tends to be high. Uses 256-bit bitmaps per node (one bit -// per possible byte label), achieving excellent cache locality and O(1) -// child lookup via popcount. -// -// - LOUDS-Sparse: For the lower levels of the trie where fanout is typically -// low. Uses compact label arrays and bitvectors, achieving better space -// efficiency than the dense encoding for sparse regions. -// -// The boundary between dense and sparse levels (the "cutoff level") is chosen -// to minimize total space: dense levels use 256 bits per node regardless of -// fanout, while sparse levels use ~10 bits per edge. When the average fanout -// drops below ~25 children per node, sparse becomes more efficient. -// -// Key design decisions: -// - Immutable: Built once from sorted keys during SST file construction. -// - Flat-array layout: The entire trie is stored as a sequence of bitvectors, -// making serialization trivial and enabling zero-copy reads from disk. -// - Leaf-indexed: Each trie leaf maps to a data block handle via packed -// uint32_t offset/size arrays, indexed by the leaf's BFS ordinal. -// - Key reconstruction: The separator key is reconstructed by tracing -// the path from root to the current leaf, collecting byte labels at -// each level. Dense levels encode the label in the bit position -// (pos % 256), sparse levels store it in the label array. -// -// Leaf ordinal computation (SuRF formulas): -// - Dense leaf: rank1(d_labels, pos+1) - rank1(d_has_child, rank1(d_labels, -// pos+1)) + rank1(d_is_prefix_key, node_num+1) - 1 -// where pos = node_num * 256 + label_byte -// - Sparse leaf: ((label_pos + 1) - rank1(s_has_child, label_pos + 1)) + -// rank1(s_is_prefix_key, node_num+1) + dense_leaf_count - 1 - -#pragma once - -#include -#include -#include -#include -#include - -#include "rocksdb/slice.h" -#include "rocksdb/status.h" -#include "util/autovector.h" -#include "utilities/trie_index/bitvector.h" - -namespace ROCKSDB_NAMESPACE { -namespace trie_index { - -// ============================================================================ -// Forward declarations -// ============================================================================ -class LoudsTrie; - -// ============================================================================ -// BlockHandle: offset and size of a data block in the SST file. -// Matches IndexFactoryBuilder::BlockHandle but defined locally to avoid -// header dependencies in the core trie implementation. -// ============================================================================ -struct TrieBlockHandle { - uint64_t offset = 0; - uint64_t size = 0; -}; - -// ============================================================================ -// LoudsTrieBuilder: Constructs a LOUDS-encoded trie from sorted keys. -// -// Usage: -// LoudsTrieBuilder builder; -// for each data block in sorted order: -// builder.AddKey(separator_key, block_handle); -// builder.Finish(); -// Slice serialized = builder.GetSerializedData(); -// -// The builder collects all separator keys, then in Finish() it: -// 1. Determines the optimal cutoff level between dense and sparse encoding. -// 2. Constructs LOUDS-Dense bitvectors for levels [0, cutoff). -// 3. Constructs LOUDS-Sparse bitvectors for levels [cutoff, max_depth). -// 4. Serializes everything into a flat buffer. -// -// All keys must be added in sorted order (according to the comparator used -// by the SST file). -// ============================================================================ -class LoudsTrieBuilder { - public: - LoudsTrieBuilder(); - - // Add a separator key and its associated data block handle. - // Keys must be added in sorted (ascending) order. - // This is the basic form used when seqno encoding is not active. - void AddKey(const Slice& key, const TrieBlockHandle& handle); - - // Add a separator key with seqno side-table metadata. Used when - // has_seqno_encoding_ is true. The seqno is stored in a side-table - // alongside the trie (NOT encoded into the key). block_count is the - // number of consecutive data blocks that share this separator key - // (1 = no overflow, >1 = same-user-key run). - // - // IMPORTANT: When block_count > 1, the overflow blocks must be added - // via AddOverflowBlock() immediately after this call, before calling - // AddKey() for the next separator. - void AddKeyWithSeqno(const Slice& key, const TrieBlockHandle& handle, - uint64_t seqno, uint32_t block_count); - - // Add an overflow block for the most recently added key (same separator). - // Called block_count-1 times after AddKeyWithSeqno() with block_count > 1. - // Each overflow block has its own handle and seqno. - void AddOverflowBlock(const TrieBlockHandle& handle, uint64_t seqno); - - // Set the seqno encoding flag. Must be called before Finish(). - // When true, the serialized trie will include a seqno side-table after - // the handle arrays, enabling post-seek correction for same-user-key - // block boundaries. - void SetHasSeqnoEncoding(bool has_seqno_encoding) { - has_seqno_encoding_ = has_seqno_encoding; - } - - // Finalize the trie construction. After this call, GetSerializedData() - // returns the serialized trie. - void Finish(); - - // Get the serialized trie data. Valid only after Finish(). - Slice GetSerializedData() const { return Slice(serialized_data_); } - - private: - // Determine the optimal cutoff level between dense and sparse encoding. - // Returns the first level where sparse encoding is more space-efficient. - uint32_t ComputeCutoffLevel() const; - - // Serialize all trie data structures into serialized_data_. - void SerializeAll(); - - // ---- Input data ---- - std::vector keys_; - std::vector handles_; - - // ---- Seqno side-table data (populated by AddKeyWithSeqno/AddOverflowBlock) - // ---- - // - // Per-key (one entry per call to AddKey/AddKeyWithSeqno): - // seqnos_[i]: seqno for the i-th separator (0 = sentinel for non-boundary) - // block_counts_[i]: how many consecutive blocks share this separator - // (1 = normal, >1 = same-user-key run with overflows) - // - // Per-overflow-block (one entry per call to AddOverflowBlock): - // overflow_handles_[j]: handle for the j-th overflow block - // overflow_seqnos_[j]: seqno for the j-th overflow block - std::vector seqnos_; - std::vector block_counts_; - std::vector overflow_handles_; - std::vector overflow_seqnos_; - - // ---- Trie structure (built during Finish()) ---- - - // Cutoff level: levels [0, cutoff_level_) use dense, rest use sparse. - uint32_t cutoff_level_; - - // Max key depth (length of longest key). - uint32_t max_depth_; - - // LOUDS-Dense bitvectors (all levels concatenated into single bitvectors): - // d_labels_: 256-bit bitmaps concatenated for all nodes across all dense - // levels. Bit (node_num * 256 + label) is set if node has child with - // that label. - // d_has_child_: One bit per set bit in d_labels_. Set if the child - // is an internal node (has further children), clear if it's a leaf. - // d_is_prefix_key_: One bit per node across all dense levels. Set if - // the path from root to this node forms a valid key (prefix match). - BitvectorBuilder d_labels_; - BitvectorBuilder d_has_child_; - BitvectorBuilder d_is_prefix_key_; - - // LOUDS-Sparse arrays (all sparse levels concatenated): - // s_labels_: Byte labels of all edges, in level-order. - // s_has_child_: One bit per label. Set if the child is internal. - // s_louds_: One bit per label. Set at the first label of each node - // (marks node boundaries in the label array). - // s_is_prefix_key_: One bit per node in the sparse region. Set if the - // path to this node forms a valid key. - std::vector s_labels_; - BitvectorBuilder s_has_child_; - BitvectorBuilder s_louds_; - BitvectorBuilder s_is_prefix_key_; - - // Total number of leaves in the dense section. - uint64_t dense_leaf_count_; - - // Total number of nodes in all dense levels combined. - uint64_t dense_node_count_; - - // Number of sparse root nodes: internal children at the last dense level - // that cross into the sparse region. See LoudsTrie::dense_child_count_. - uint64_t dense_child_count_; - - // Whether separator keys include seqno encoding. Written to the serialized - // header so the reader can detect it. - bool has_seqno_encoding_; - - // ---- Serialized output ---- - std::string serialized_data_; -}; - -// ============================================================================ -// LoudsTrie: Immutable LOUDS-encoded trie for reading. -// -// Deserialized from a flat buffer (e.g., read from an SST meta-block). -// Supports Seek (find the first leaf >= target key) and Next (advance to -// the next leaf in sorted order). -// -// The trie does NOT own the underlying data when initialized from external -// memory (e.g., a block cache entry). The caller must ensure the data remains -// valid for the lifetime of this object. -// -// The trie is organized as: -// - Dense levels [0, cutoff_level_): bitvectors d_labels_, d_has_child_, -// d_is_prefix_key_ with 256-bit bitmaps per node. -// - Sparse levels [cutoff_level_, max_depth_]: arrays s_labels_, -// s_has_child_, s_louds_, s_is_prefix_key_. -// ============================================================================ -class LoudsTrie { - public: - LoudsTrie(); - - // LoudsTrie contains Bitvector members (which hold raw pointers into - // owned or external memory) and a raw pointer to sparse labels data. - // Copying would create dangling pointers or aliased external references. - // - // Move safety: when aligned_copy_ is non-empty, bitvectors and raw - // pointers reference its buffer. The C++ standard guarantees that - // std::string's noexcept move constructor transfers the heap buffer - // without reallocation (noexcept precludes allocation, and COW is - // forbidden since C++11). Trie data always exceeds the SSO threshold - // (hundreds to thousands of bytes), so aligned_copy_ is always - // heap-allocated, and move always preserves the buffer address. - ~LoudsTrie() = default; - - LoudsTrie(const LoudsTrie&) = delete; - LoudsTrie& operator=(const LoudsTrie&) = delete; - LoudsTrie(LoudsTrie&&) = default; - LoudsTrie& operator=(LoudsTrie&&) = default; - - // Initialize from serialized data. Returns Status::OK() on success, - // or Status::Corruption() if the data is malformed. - Status InitFromData(const Slice& data); - - // ---- Accessors ---- - uint64_t NumKeys() const { return num_keys_; } - uint32_t CutoffLevel() const { return cutoff_level_; } - uint32_t MaxDepth() const { return max_depth_; } - - // Whether this trie was built with a seqno side-table (enabling post-seek - // correction for same-user-key block boundaries). When true, the serialized - // data includes per-leaf seqno and block count arrays, plus overflow block - // metadata for runs of blocks sharing the same separator key. - bool HasSeqnoEncoding() const { return has_seqno_encoding_; } - - // Get the block handle for the i-th leaf (0-indexed). - TrieBlockHandle GetHandle(uint64_t leaf_index) const; - - // Whether this trie has path-compression chains. Used by the iterator - // to select a specialized Seek implementation at construction time, - // avoiding any per-level overhead when chains are absent. - bool HasChains() const { return !s_chain_lens_.empty(); } - - // Approximate heap memory used by auxiliary data structures (child position - // lookup tables). Does not include the serialized data itself (which is - // typically owned by the block cache). - size_t ApproximateAuxMemoryUsage() const { - return (s_child_start_pos_.capacity() + s_child_end_pos_.capacity()) * - sizeof(uint32_t); - } - - // Allow the iterator to access internal bitvectors directly for - // performance-critical rank/select operations during traversal. - friend class LoudsTrieIterator; - - private: - // Number of keys (leaves). - uint64_t num_keys_; - - // Cutoff level between dense and sparse. - uint32_t cutoff_level_; - - // Maximum key depth. - uint32_t max_depth_; - - // True if the trie includes a seqno side-table for post-seek correction. - // Set from the flags field in the serialized header. - bool has_seqno_encoding_; - - // Dense leaf count (leaves in levels [0, cutoff_level_)). - uint64_t dense_leaf_count_; - - // Total number of nodes across all dense levels. - uint64_t dense_node_count_; - - // Number of sparse root nodes: internal children at the last dense level - // (cutoff_level_ - 1) that cross into the sparse region. Used to offset - // sparse node numbering so that children of sparse internal labels are - // numbered after these root sparse nodes. When cutoff_level_ == 0, this - // is set to 1 (the root itself). - uint64_t dense_child_count_; - - // LOUDS-Dense bitvectors (all dense levels concatenated). - Bitvector d_labels_; - Bitvector d_has_child_; - Bitvector d_is_prefix_key_; - - // LOUDS-Sparse (all sparse levels concatenated). - const uint8_t* s_labels_data_; - uint64_t s_labels_size_; - Bitvector s_has_child_; - Bitvector s_louds_; - Bitvector s_is_prefix_key_; - - // SuRF-style child position lookup tables for Select-free traversal. - // Instead of computing FindNthOneBit(node_num) during traversal, we - // precompute child start/end positions indexed by internal label rank. This - // allows traversal using only Rank1 (O(1)) and array lookup (O(1)). - // - // For the k-th internal label (has_child[pos]=1, where k = Rank1(pos+1)-1): - // s_child_start_pos_[k] = start position of child node - // s_child_end_pos_[k] = end position (exclusive) of child node - // - // Memory overhead: 8 bytes per internal node (2 x uint32_t). - std::vector s_child_start_pos_; - std::vector s_child_end_pos_; - - // Path compression: chain metadata for fanout-1 chains in the sparse region. - // - // A "chain" is a sequence of >= 2 consecutive fanout-1 nodes (nodes with - // exactly one label that is internal) starting from the child of an internal - // label. Chains are common in tries with long shared prefixes (e.g., - // zero-padded numeric keys, URL paths). - // - // For the k-th internal label (same indexing as s_child_start_pos_): - // Storage uses a bitmap (1 bit per internal label) for O(1) chain detection, - // plus compact arrays indexed by chain ordinal (Rank1 on the bitmap). - // - // Lookup during Seek: - // 1. s_chain_bitmap_.GetBit(child_idx) -- has chain? - // 2. chain_idx = s_chain_bitmap_.Rank1(child_idx + 1) - 1 - // 3. s_chain_lens_[chain_idx], s_chain_suffix_offsets_[chain_idx], etc. - // - // Space overhead: 1 bit per internal label (bitmap) + 10 bytes per chain - // (offset + len + end_child_idx) + suffix bytes. For key sets with few - // chains (e.g., random hex), overhead is < 1 byte per internal label. - Bitvector s_chain_bitmap_; - std::vector s_chain_suffix_offsets_; - std::vector s_chain_lens_; - std::vector s_chain_end_child_idx_; - const uint8_t* s_chain_suffix_data_; - uint64_t s_chain_suffix_size_; - - // Block handles: packed uint32_t arrays for data block offsets and sizes. - // BFS leaf order does not necessarily match key-sorted order (deeper leaves - // appear later in BFS even if they precede shallower leaves - // lexicographically), so offsets are NOT monotonically non-decreasing and - // cannot use Elias-Fano encoding. Instead, we store offsets and sizes as - // packed uint32_t arrays for O(1) random access. - // - // uint32_t limits individual values to ~4 GB, which is sufficient since - // RocksDB SST files are typically 64 MB to 1 GB and never exceed 4 GB. - const uint32_t* handle_offsets_; - const uint32_t* handle_sizes_; - - // ---- Seqno side-table (deserialized when has_seqno_encoding_ is true) ---- - // - // The side-table enables post-seek correction for same-user-key block - // boundaries. It stores per-leaf seqno and block count data in BFS leaf - // order, plus overflow block handles/seqnos for runs where the same - // separator maps to multiple blocks. - // - // leaf_seqnos_[i]: seqno for the i-th leaf (BFS order). - // For non-boundary leaves: stores 0 (sentinel meaning "no seqno - // correction needed"), matching the standard index's user-key-only - // comparison mode. - // For same-user-key boundary and last-block leaves: stores the actual - // tag ((sequence_number << 8) | value_type). - // - // leaf_block_counts_[i]: how many consecutive blocks share this separator. - // 1 = no overflow (the common case). >1 = same-user-key run. - // - // overflow_offsets_/overflow_sizes_/overflow_seqnos_: packed arrays for - // the overflow blocks (total count = sum of (block_count-1) for all - // leaves). - // - // overflow_base_[i]: prefix sum of (block_count-1) for leaves [0, i), - // precomputed during InitFromData() for O(1) random access into the - // overflow arrays. overflow_base_[i] is the starting index into the - // overflow arrays for leaf i's overflow blocks. - const uint64_t* leaf_seqnos_; - const uint32_t* leaf_block_counts_; - const uint32_t* overflow_offsets_; - const uint32_t* overflow_sizes_; - const uint64_t* overflow_seqnos_; - uint32_t num_overflow_blocks_; - std::vector overflow_base_; - - public: - // ---- Seqno side-table accessors (used by TrieIndexIterator) ---- - - // Get the seqno for the i-th leaf (BFS order). Returns 0 (sentinel) for - // non-boundary leaves. - uint64_t GetLeafSeqno(uint64_t leaf_index) const { - assert(has_seqno_encoding_ && leaf_seqnos_ != nullptr); - assert(leaf_index < num_keys_); - return leaf_seqnos_[leaf_index]; - } - - // Get the block count for the i-th leaf. Returns 1 for normal leaves. - uint32_t GetLeafBlockCount(uint64_t leaf_index) const { - assert(has_seqno_encoding_ && leaf_block_counts_ != nullptr); - assert(leaf_index < num_keys_); - return leaf_block_counts_[leaf_index]; - } - - // Get the overflow base (starting index into overflow arrays) for leaf i. - uint32_t GetOverflowBase(uint64_t leaf_index) const { - assert(has_seqno_encoding_); - assert(leaf_index < overflow_base_.size()); - return overflow_base_[leaf_index]; - } - - // Get the handle for the j-th overflow block. - TrieBlockHandle GetOverflowHandle(uint32_t overflow_index) const { - assert(overflow_index < num_overflow_blocks_); - return TrieBlockHandle{overflow_offsets_[overflow_index], - overflow_sizes_[overflow_index]}; - } - - // Get the seqno for the j-th overflow block. - uint64_t GetOverflowSeqno(uint32_t overflow_index) const { - assert(overflow_index < num_overflow_blocks_); - return overflow_seqnos_[overflow_index]; - } - - private: - // Aligned copy of the serialized trie data, used when the input data from - // the block reader is not 8-byte aligned (e.g., mmap at an unaligned file - // offset). All Bitvector and raw pointer members reference this buffer - // when non-empty. std::string::data() returns memory from new[]/malloc, - // which is aligned to at least alignof(max_align_t) >= 8. - std::string aligned_copy_; -}; - -// ============================================================================ -// LoudsTrieIterator: Iterates over the leaves of a LoudsTrie. -// -// Supports bidirectional iteration: Seek/SeekToFirst/Next (forward) and -// SeekToLast/Prev (reverse). -// Reconstructs the separator key from the trie path at each position. -// -// The iterator maintains a stack of positions through the trie (one per -// level from root to current leaf). This enables: -// - Key reconstruction: collecting the label byte at each level. -// - Backtracking for Next(): pop to parent, advance to next sibling. -// - Leaf ordinal computation: using rank formulas on bitvectors. -// -// Design follows the SuRF reference implementation for correctness. -// ============================================================================ -class LoudsTrieIterator { - public: - explicit LoudsTrieIterator(const LoudsTrie* trie); - - // Position on the very first leaf (smallest key) by descending from the - // root to the leftmost leaf. More efficient than Seek(Slice()) because it - // skips SeekImpl's target-consumption loop and its redundant prefix key - // check at root (DescendToLeftmostLeaf handles prefix keys at every node). - // Returns true if positioned on a valid leaf. - bool SeekToFirst(); - - // Position on the very last leaf (largest key) by descending from the - // root to the rightmost leaf. Returns true if positioned on a valid leaf. - bool SeekToLast(); - - // Move to the previous leaf in sorted order. - // Returns true if positioned on a valid leaf, false if before the first. - bool Prev(); - - // Seek to the first leaf whose key is >= `target`. - // Returns true if positioned on a valid leaf. - // - // Dispatches to a specialized implementation selected at construction time - // based on whether the trie has path-compression chains. This eliminates - // all chain-related code from the instruction cache when chains are absent, - // following the same pattern as RocksDB's BlockIter::ParseNextKey template. - bool Seek(const Slice& target) { - if (has_chains_) { - return SeekImpl(target); - } - return SeekImpl(target); - } - - // Advance to the next leaf in sorted order. - // Returns true if positioned on a valid leaf. - bool Next(); - - // Check if the iterator is positioned on a valid leaf. - bool Valid() const { return valid_; } - - // Get the current separator key. Valid only when Valid() is true. - // The returned Slice is valid until the next Seek/Next call. - Slice Key() const { return Slice(key_buf_.get(), key_len_); } - - // Get the current leaf index (for mapping to block handles). - uint64_t LeafIndex() const { return leaf_index_; } - - // Get the block handle for the current leaf. - TrieBlockHandle Value() const; - - private: - // Position within a single trie level. The iterator maintains a stack - // of these from root to the current position. - // - // Packed into 8 bytes by encoding the is_dense flag in bit 63 of the - // position value. Since bitvector positions and label array indices are - // well under 2^63, this is safe. This halves the per-level memory from - // 16 bytes (with alignment padding) to 8 bytes, improving cache - // utilization for the path_ stack. - struct LevelPos { - // Position in the bitvector/label array at this level, with the - // is_dense flag encoded in the high bit (bit 63). - // - Dense: bit position in d_labels_ (= node_num * 256 + - // label_byte). The label byte is pos % 256. - // - Sparse: index into s_labels_ array. - uint64_t pos_and_flag; - - static constexpr uint64_t kDenseFlag = uint64_t(1) << 63; - - uint64_t pos() const { return pos_and_flag & ~kDenseFlag; } - bool is_dense() const { return (pos_and_flag & kDenseFlag) != 0; } - - static LevelPos MakeDense(uint64_t p) { return {p | kDenseFlag}; } - static LevelPos MakeSparse(uint64_t p) { return {p}; } - }; - - // --- Dense level helpers --- - - // Seek within a dense node for a target label byte. - // Sets result to the position of the label if found, or the position of - // the next label >= target_byte if not found. - // Returns true if the exact label was found, false if we landed on a - // label > target_byte (or no label exists). - bool DenseSeekLabel(uint64_t node_num, uint8_t target_byte, - uint64_t* out_pos); - - // Compute the child node number for a dense internal child with the given - // label_rank. Takes a pre-computed label_rank to avoid redundant - // Rank1(d_labels_) calls in hot paths where label_rank was already computed - // for has_child checking. - uint64_t DenseChildNodeNumFromRank(uint64_t label_rank) const; - - // Compute the leaf ordinal for a dense leaf at `pos`. - // leaf_idx = rank1(d_labels_, pos+1) - rank1(d_has_child_, - // rank1(d_labels_, pos+1)) + rank1(d_is_prefix_key_, node_num+1) - 1 - // Takes a pre-computed label_rank. - uint64_t DenseLeafIndexFromRank(uint64_t pos, uint64_t label_rank) const; - - // Same as DenseLeafIndexFromRank but also takes a pre-computed - // d_has_child_.Rank1(label_rank + 1) to avoid redundant rank call. - uint64_t DenseLeafIndexFromRankAndHasChildRank(uint64_t pos, - uint64_t label_rank, - uint64_t has_child_rank) const; - - // Compute the leaf ordinal for a dense prefix key at node `node_num`. - // A prefix key leaf comes before any child leaves of that node. - uint64_t DensePrefixKeyLeafIndex(uint64_t node_num) const; - - // --- Sparse level helpers --- - - // Seek within a sparse node starting at `node_start_pos` for label byte. - // Returns true if the exact label was found, false otherwise. Writes the - // position to `out_pos`. `node_end_pos` is one past the last label - // position of this node. - bool SparseSeekLabel(uint64_t node_start_pos, uint64_t node_end_pos, - uint8_t target_byte, uint64_t* out_pos); - - // Compute the child node number for a sparse internal child at `pos`. - // child_node_num = dense_child_count_ + rank1(s_has_child_, pos+1) - 1 - // (offset by dense_child_count_ because children of sparse internal labels - // are numbered after the root sparse nodes). - uint64_t SparseChildNodeNum(uint64_t pos) const; - - // Compute the leaf ordinal for a sparse leaf at `pos`. - uint64_t SparseLeafIndex(uint64_t pos) const; - - // Same as SparseLeafIndex but takes a pre-computed - // s_has_child_.Rank1(pos + 1) to avoid redundant rank call. - uint64_t SparseLeafIndexFromHasChildRank(uint64_t pos, - uint64_t has_child_rank) const; - - // Compute the leaf ordinal for a sparse prefix key at sparse node - // `sparse_node_num`. The sparse_node_num is 0-indexed among sparse nodes - // only (not including dense nodes). - uint64_t SparsePrefixKeyLeafIndex(uint64_t sparse_node_num) const; - - // Get the sparse node number (0-indexed among sparse nodes) from a - // position in the s_labels_ array. - uint64_t SparseNodeNum(uint64_t pos) const; - - // Get the start position (in s_labels_) for sparse node `sparse_node_num`. - uint64_t SparseNodeStartPos(uint64_t sparse_node_num) const; - - // Get the end position (one past last label) for a sparse node starting - // at `start_pos`. - uint64_t SparseNodeEndPos(uint64_t start_pos) const; - - // --- Traversal helpers --- - - // Descend from the given node to the leftmost leaf in its subtree, - // pushing entries onto path_ and building key_buf_. Sets - // leaf_index_ and valid_. Returns true if a leaf was found. - bool DescendToLeftmostLeaf(bool in_dense, uint64_t node_num); - - // Descend from the given node to the rightmost leaf in its subtree, - // pushing entries onto path_ and building key_buf_. Sets - // leaf_index_ and valid_. Returns true if a leaf was found. - // Unlike DescendToLeftmostLeaf, this does NOT check prefix keys -- - // prefix keys are the smallest leaf at a node, so in reverse order - // they are visited last (handled by Retreat when backtracking). - bool DescendToRightmostLeaf(bool in_dense, uint64_t node_num); - - // Advance to the next valid leaf by backtracking up the trie path - // and finding the next sibling label, then descending to the leftmost - // leaf in that subtree. Used by Next() and SeekImpl(). - // Returns true if a next leaf was found. - bool Advance(); - - // Retreat to the previous valid leaf by backtracking up the trie path - // and finding the previous sibling label, then descending to the - // rightmost leaf in that subtree. When no previous sibling exists, - // checks if the current node has a prefix key (which is the smallest - // leaf at a node, so it comes last in reverse order). - // Used by Prev(). - // Returns true if a previous leaf was found. - bool Retreat(); - - // Seek implementation, templated on whether path-compression chains exist. - // When kHasChains=false, the compiler eliminates ALL chain-related code, - // keeping the i-cache footprint minimal for tries without chains. - // This follows the same pattern as RocksDB's BlockIter::ParseNextKey. - template - bool SeekImpl(const Slice& target); - - // True if the trie has path-compression chains. Set once in the constructor - // and used by Seek() to dispatch to the correct specialization. - bool has_chains_; - - const LoudsTrie* trie_; - bool valid_; - uint64_t leaf_index_; - - // The reconstructed key at the current position. Each byte corresponds - // to a level in path_. For dense levels, the byte is pos % 256; for - // sparse levels, the byte is s_labels_[pos]. - // - // Key reconstruction appends one byte per trie level in the Seek/Next - // hot loop, so the append operation must be as cheap as possible -- a - // single inlined store + increment with no function call overhead. The - // buffer is heap-allocated once in the constructor to MaxDepth()+1 bytes. - // - // All append sites go through AppendKeySlot() which validates bounds in - // debug builds. In release builds it compiles to the same single store + - // increment with no overhead (the assert is elided). - std::unique_ptr key_buf_; - uint32_t key_len_; - uint32_t key_cap_; - - // Returns a reference to the next key buffer slot, advancing key_len_. - // Validates in debug builds that the buffer has space. A corrupted trie - // with depth > max_depth_ would overflow key_buf_ without this check. - char& AppendKeySlot() { - assert(key_len_ < key_cap_ && - "key_buf_ overflow: trie depth exceeds max_depth_"); - return key_buf_[key_len_++]; - } - - // Stack of positions from root to current leaf. path_[i] holds the - // position at depth i in the trie. path_.size() equals key_len_ for - // both child leaves and prefix keys (in which case the last path_ - // entry's label is not appended to - // key_buf_ since the node itself is the key). - // - // For a prefix key, we mark it by setting is_at_prefix_key_ = true and - // the path_ only goes up to the prefix key node level (no label entry - // for the prefix key itself since the key terminates at the node, not - // at a child edge). - // - // Uses autovector with 24 inline slots to avoid heap allocation for - // tries up to 24 levels deep. Most real-world key sets have depth < 24. - autovector path_; - - // True if the current leaf is a prefix key (the key terminates at a - // node that also has children). In this case, path_.size() == depth - // of the node and key_buf_[0..key_len_) holds the prefix key value. - bool is_at_prefix_key_; -}; - -} // namespace trie_index -} // namespace ROCKSDB_NAMESPACE diff --git a/utilities/trie_index/trie.cc b/utilities/trie_index/trie.cc new file mode 100644 index 000000000000..8ce598418b57 --- /dev/null +++ b/utilities/trie_index/trie.cc @@ -0,0 +1,1549 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// This source code is licensed under both the GPLv2 (found in the +// COPYING file in the root directory) and Apache 2.0 License +// (found in the LICENSE.Apache file in the root directory). + +#include "utilities/trie_index/trie.h" + +#include +#include +#include +#include + +#include "util/cast_util.h" +#include "util/coding.h" +#include "util/math.h" + +namespace ROCKSDB_NAMESPACE { +namespace trie_index { + +namespace { + +constexpr uint32_t kMagic = 0x54524945; // "TRIE" +constexpr uint32_t kVersion = 2; +constexpr uint64_t kExplicitSizes = ~uint64_t{0}; + +// Columns. +constexpr uint32_t kGroupBits = 6; +constexpr uint64_t kGroupSize = uint64_t{1} << kGroupBits; +constexpr size_t kColumnHeaderSize = 8; +constexpr size_t kGroupHeaderSize = 16; +constexpr size_t kColumnPadding = 8; +// A residual is extracted with one unaligned 8-byte load and a shift of up to +// 7 bits. +constexpr uint32_t kMaxColumnWidth = 56; +// A group's layout field holds its residual width in the low kWidthBits bits +// and the position of its residuals, in 8-byte words, above them. +constexpr uint32_t kWidthBits = 6; +constexpr uint32_t kWidthMask = (1u << kWidthBits) - 1; +constexpr uint64_t kMaxResidualWords = uint64_t{1} << (32 - kWidthBits); +// A linear fit keeps base and slope arithmetic within int64_t for values +// below this; larger values are stored with a zero slope. +constexpr uint64_t kMaxLinearValue = uint64_t{1} << 62; + +// Side tables. +constexpr size_t kSideTableEntrySize = 16; + +// Trie node header. +constexpr uint8_t kKindMask = 3; +constexpr uint8_t kKindLeaf = 0; +constexpr uint8_t kKindSparse = 1; +constexpr uint8_t kKindBitmap = 2; +constexpr uint8_t kTerminalBit = 4; +constexpr int kPointerWidthShift = 3; +constexpr uint8_t kReservedNodeBits = 0xE0; +constexpr uint8_t kLongPathMarker = 255; +constexpr size_t kBitmapBytes = 32; +// Nodes with at least this many children use a bitmap instead of a label +// list. +constexpr size_t kMinBitmapChildren = 32; + +// Bucket entry header. +constexpr uint8_t kLongEntryMarker = 0xFF; +constexpr uint32_t kShortEntryLimit = 15; + +// An iterator reserves room for keys up to this long when it is created. +// Longer keys grow the buffer when reached, so one huge separator does not +// make every iterator over the index allocate its length. +constexpr size_t kInitialKeyCapacity = 256; + +// Header field positions. +constexpr size_t kMagicPos = 0; +constexpr size_t kVersionPos = 4; +constexpr size_t kFlagsPos = 8; +constexpr size_t kBucketSizePos = 12; +constexpr size_t kNumBlocksPos = 16; +constexpr size_t kNumKeysPos = 24; +constexpr size_t kPrefixLenPos = 32; +constexpr size_t kMaxKeyLenPos = 36; +constexpr size_t kIdWidthPos = 40; +constexpr size_t kTrieRootPos = 44; +constexpr size_t kSizeGapPos = 48; +constexpr size_t kSectionsPos = 56; +// Indexes of the section offsets in the header. +constexpr size_t kTrieSection = 0; +constexpr size_t kBucketTableSection = 1; +constexpr size_t kBucketDataSection = 2; +constexpr size_t kOffsetTableSection = 3; +constexpr size_t kSizeTableSection = 4; +constexpr size_t kRunTableSection = 5; +constexpr size_t kTagTableSection = 6; +constexpr size_t kEnd = 7; +constexpr size_t kNumSectionOffsets = 8; +static_assert(kSectionsPos + kNumSectionOffsets * 8 == kTrieHeaderSize); + +Status Corrupt(const char* what) { + return Status::Corruption("Trie index", what); +} + +// Number of bytes, 1 to 4, needed to store v, which must fit 32 bits. +uint32_t BytesFor(uint64_t v) { + if (v <= 0xFF) { + return 1; + } else if (v <= 0xFFFF) { + return 2; + } else if (v <= 0xFFFFFF) { + return 3; + } + return 4; +} + +uint32_t BitsFor(uint64_t v) { + return v == 0 ? 0 : lossless_cast(FloorLog2(v) + 1); +} + +uint8_t Byte(const char* p, size_t i) { return lossless_cast(p[i]); } + +// Label of bit `bit` of word `word` of a bitmap node. +uint8_t BitmapLabel(int word, int bit) { + assert(word >= 0 && word < 4 && bit >= 0 && bit < 64); + return static_cast(64 * word + bit); +} + +uint32_t LoadLE(const char* p, uint32_t width) { + switch (width) { + case 1: + return Byte(p, 0); + case 2: + return DecodeFixed16(p); + case 3: + return DecodeFixed16(p) | (uint32_t{Byte(p, 2)} << 16); + default: + return DecodeFixed32(p); + } +} + +void PutLE(std::string* dst, uint64_t v, uint32_t width) { + for (uint32_t i = 0; i < width; i++) { + dst->push_back(static_cast((v >> (8 * i)) & 0xFF)); + } +} + +// Length of the common prefix of a and b, looking at most n bytes. +size_t CommonPrefix(const char* a, const char* b, size_t n) { + size_t i = 0; + while (i + 8 <= n) { + uint64_t diff = DecodeFixed64(a + i) ^ DecodeFixed64(b + i); + if (diff != 0) { + return i + lossless_cast(CountTrailingZeroBits(diff)) / 8; + } + i += 8; + } + while (i < n && a[i] == b[i]) { + i++; + } + return i; +} + +// Appends values as a column. +Status AppendColumn(const std::vector& values, std::string* out) { + const uint64_t count = values.size(); + const uint64_t num_groups = (count + kGroupSize - 1) >> kGroupBits; + PutFixed64(out, count); + size_t group_pos = out->size(); + const size_t residuals_pos = group_pos + num_groups * kGroupHeaderSize; + out->resize(residuals_pos + kColumnPadding); + uint64_t words = 0; + for (uint64_t g = 0; g < num_groups; g++) { + const uint64_t lo = g << kGroupBits; + const uint64_t n = std::min(count - lo, kGroupSize); + const uint64_t* v = values.data() + lo; + bool linear = n > 1 && v[n - 1] < kMaxLinearValue; + for (uint64_t j = 1; linear && j < n; j++) { + linear = v[j] >= v[j - 1]; + } + uint64_t slope = linear ? (v[n - 1] - v[0]) / (n - 1) : 0; + uint64_t base = v[0]; + if (linear && slope <= std::numeric_limits::max()) { + // v[j] - slope * j can be negative, but stays within int64_t. + int64_t min_offset = std::numeric_limits::max(); + for (uint64_t j = 0; j < n; j++) { + min_offset = + std::min(min_offset, lossless_cast(v[j]) - + lossless_cast(slope * j)); + } + base = lossless_cast(min_offset); + } else { + slope = 0; + for (uint64_t j = 0; j < n; j++) { + base = std::min(base, v[j]); + } + } + // Unsigned wraparound gives the exact non-negative residuals. + uint64_t max_residual = 0; + for (uint64_t j = 0; j < n; j++) { + max_residual = std::max(max_residual, v[j] - slope * j - base); + } + const uint32_t width = BitsFor(max_residual); + if (width > kMaxColumnWidth) { + return Status::NotSupported("Trie index: column value range too large"); + } + if (words >= kMaxResidualWords) { + return Status::NotSupported("Trie index: column too large"); + } + char* group = &(*out)[group_pos]; + EncodeFixed64(group, base); + // Both fit 32 bits, as checked above. + EncodeFixed32(group + 8, static_cast(slope)); + EncodeFixed32(group + 12, + static_cast(words << kWidthBits | width)); + group_pos += kGroupHeaderSize; + // The zero padding after the last group keeps the 8-byte stores in + // bounds. + const size_t start = residuals_pos + words * 8; + words += (n * width + 63) / 64; + out->resize(residuals_pos + words * 8 + kColumnPadding); + for (uint64_t j = 0; j < n && width > 0; j++) { + const uint64_t bit = j * width; + char* p = &(*out)[start + bit / 8]; + EncodeFixed64( + p, DecodeFixed64(p) | ((v[j] - slope * j - base) << (bit % 8))); + } + } + return Status::OK(); +} + +// Writes a path-compressed trie over sorted, distinct heads, children before +// parents, without recursion so that deep tries cannot exhaust the stack. +class TrieWriter { + public: + TrieWriter(const std::vector& heads, uint32_t id_width) + : heads_(heads), id_width_(id_width) {} + + Status Write(std::string* out, uint32_t* root) { + std::vector stack; + stack.push_back(MakeFrame(0, heads_.size(), 0)); + uint64_t last_pos = 0; + while (!stack.empty()) { + Frame& top = stack.back(); + if (top.next < top.end) { + const size_t d = top.path_end; + const uint8_t label = Byte(heads_[top.next].data(), d); + uint64_t j = top.next + 1; + while (j < top.end && Byte(heads_[j].data(), d) == label) { + j++; + } + top.labels.push_back(label); + const uint64_t child_begin = top.next; + top.next = j; + // Invalidates top. + stack.push_back(MakeFrame(child_begin, j, d + 1)); + continue; + } + Status s = Emit(top, out, &last_pos); + if (!s.ok()) { + return s; + } + stack.pop_back(); + if (!stack.empty()) { + stack.back().children.push_back(last_pos); + } + } + if (out->size() > std::numeric_limits::max()) { + return Status::NotSupported("Trie index: trie too large"); + } + *root = static_cast(last_pos); + return Status::OK(); + } + + private: + // Heads [begin, end), which share their first `depth` bytes. + struct Frame { + uint64_t begin = 0; + uint64_t end = 0; + size_t depth = 0; + size_t path_end = 0; + bool terminal = false; + uint64_t next = 0; + std::vector labels; + std::vector children; + }; + + Frame MakeFrame(uint64_t begin, uint64_t end, size_t depth) const { + const Slice& first = heads_[begin]; + const Slice& last = heads_[end - 1]; + Frame f; + f.begin = begin; + f.end = end; + f.depth = depth; + f.path_end = + depth + CommonPrefix(first.data() + depth, last.data() + depth, + std::min(first.size(), last.size()) - depth); + // Only the first head of a sorted range can end at the shared path. + f.terminal = first.size() == f.path_end; + f.next = f.terminal ? begin + 1 : begin; + return f; + } + + Status Emit(const Frame& f, std::string* out, uint64_t* pos) const { + *pos = out->size(); + uint64_t max_distance = 0; + for (uint64_t child : f.children) { + max_distance = std::max(max_distance, *pos - child); + } + if (max_distance > std::numeric_limits::max()) { + return Status::NotSupported("Trie index: trie too large"); + } + const uint32_t pointer_width = BytesFor(max_distance); + uint8_t kind = kKindLeaf; + if (f.labels.size() >= kMinBitmapChildren) { + kind = kKindBitmap; + } else if (!f.labels.empty()) { + kind = kKindSparse; + } + out->push_back( + static_cast(kind | (f.terminal ? kTerminalBit : 0) | + ((pointer_width - 1) << kPointerWidthShift))); + const size_t path_len = f.path_end - f.depth; + if (path_len < kLongPathMarker) { + out->push_back(static_cast(path_len)); + } else { + out->push_back(lossless_cast(kLongPathMarker)); + PutVarint32(out, static_cast(path_len)); + } + out->append(heads_[f.begin].data() + f.depth, path_len); + if (f.terminal) { + PutLE(out, f.begin, id_width_); + } + if (kind == kKindSparse) { + out->push_back(static_cast(f.labels.size() - 1)); + for (uint8_t label : f.labels) { + out->push_back(lossless_cast(label)); + } + } else if (kind == kKindBitmap) { + uint64_t words[4] = {0, 0, 0, 0}; + for (uint8_t label : f.labels) { + words[label / 64] |= uint64_t{1} << (label % 64); + } + for (uint64_t w : words) { + PutFixed64(out, w); + } + } + for (uint64_t child : f.children) { + PutLE(out, *pos - child, pointer_width); + } + return Status::OK(); + } + + const std::vector& heads_; + const uint32_t id_width_; +}; + +} // namespace + +// ============================================================================ +// TrieBuilder +// ============================================================================ + +TrieBuilder::TrieBuilder(uint32_t bucket_size) : bucket_size_(bucket_size) {} + +Slice TrieBuilder::LastSeparator() const { + const size_t num_keys = key_ends_.size(); + if (num_keys == 0) { + return Slice(); + } + const size_t begin = num_keys > 1 ? key_ends_[num_keys - 2] : 0; + return Slice(key_bytes_.data() + begin, key_ends_.back() - begin); +} + +void TrieBuilder::AddBlock(const Slice& separator, + const TrieBlockHandle& handle, uint64_t tag) { + const uint64_t num_keys = key_ends_.size(); + const int cmp = num_keys > 0 ? separator.compare(LastSeparator()) : 1; + if (cmp > 0) { + key_bytes_.append(separator.data(), separator.size()); + key_ends_.push_back(key_bytes_.size()); + } else { + out_of_order_ |= cmp < 0; + if (!runs_.empty() && runs_.back().key_index == num_keys - 1) { + runs_.back().extra_blocks++; + } else { + runs_.push_back({num_keys - 1, 1}); + } + } + if (tag != 0) { + tags_.push_back({handles_.size(), tag}); + } + handles_.push_back(handle); +} + +Status TrieBuilder::Finish(std::string* output) { + if (out_of_order_) { + return Status::InvalidArgument("Trie index: separators out of order"); + } + if (bucket_size_ == 0) { + return Status::InvalidArgument("Trie index: bucket size must be > 0"); + } + const uint64_t num_blocks = handles_.size(); + const uint64_t num_keys = key_ends_.size(); + const uint64_t num_buckets = (num_keys + bucket_size_ - 1) / bucket_size_; + if (num_buckets > uint64_t{std::numeric_limits::max()} + 1) { + return Status::NotSupported("Trie index: too many buckets"); + } + auto key = [&](uint64_t i) { + const size_t begin = i > 0 ? key_ends_[i - 1] : 0; + return Slice(key_bytes_.data() + begin, key_ends_[i] - begin); + }; + + size_t prefix_len = 0; + size_t max_key_len = 0; + if (num_keys > 0) { + const Slice first = key(0); + const Slice last = key(num_keys - 1); + prefix_len = CommonPrefix(first.data(), last.data(), + std::min(first.size(), last.size())); + for (uint64_t i = 0; i < num_keys; i++) { + max_key_len = std::max(max_key_len, key(i).size()); + } + } + if (max_key_len > std::numeric_limits::max()) { + return Status::NotSupported("Trie index: separator too long"); + } + + // Trie over the bucket heads, without the shared prefix. + const uint32_t id_width = BytesFor(num_buckets > 0 ? num_buckets - 1 : 0); + std::string trie; + uint32_t trie_root = 0; + if (num_keys > 0) { + std::vector heads; + heads.reserve(num_buckets); + for (uint64_t b = 0; b < num_buckets; b++) { + Slice head = key(b * bucket_size_); + head.remove_prefix(prefix_len); + heads.push_back(head); + } + Status s = TrieWriter(heads, id_width).Write(&trie, &trie_root); + if (!s.ok()) { + return s; + } + } + + // Front-coded buckets. + std::string bucket_data; + std::vector bucket_positions; + bucket_positions.reserve(num_buckets + 1); + for (uint64_t b = 0; b < num_buckets; b++) { + bucket_positions.push_back(bucket_data.size()); + const uint64_t begin = b * bucket_size_; + const uint64_t end = std::min(num_keys, begin + bucket_size_); + for (uint64_t i = begin + 1; i < end; i++) { + const Slice prev = key(i - 1); + const Slice cur = key(i); + const size_t shared = CommonPrefix(prev.data(), cur.data(), + std::min(prev.size(), cur.size())); + const size_t drop = prev.size() - shared; + const size_t add = cur.size() - shared; + if (drop < kShortEntryLimit && add < kShortEntryLimit) { + bucket_data.push_back(static_cast((drop << 4) | add)); + } else { + bucket_data.push_back(lossless_cast(kLongEntryMarker)); + PutVarint32(&bucket_data, static_cast(drop)); + PutVarint32(&bucket_data, static_cast(add)); + } + bucket_data.append(cur.data() + shared, add); + } + } + bucket_positions.push_back(bucket_data.size()); + + // Block offsets, with sizes derived when every gap between blocks is the + // same. + std::vector offsets(num_blocks + 1, 0); + bool derived = true; + uint64_t size_gap = 0; + for (uint64_t i = 0; i < num_blocks; i++) { + const TrieBlockHandle& h = handles_[i]; + if (i > 0 && h.offset < offsets[i - 1]) { + return Status::InvalidArgument("Trie index: block offsets decrease"); + } + offsets[i] = h.offset; + if (h.size > std::numeric_limits::max() - h.offset) { + derived = false; + continue; + } + const uint64_t block_end = h.offset + h.size; + if (i + 1 < num_blocks) { + const uint64_t next = handles_[i + 1].offset; + if (next < block_end) { + derived = false; + } else if (i == 0) { + size_gap = next - block_end; + } else if (next - block_end != size_gap) { + derived = false; + } + } else if (derived && + size_gap <= std::numeric_limits::max() - block_end) { + offsets[num_blocks] = block_end + size_gap; + } else { + derived = false; + } + } + if (!derived) { + // Sizes are stored, so the last offset only has to keep the column + // non-decreasing. + offsets[num_blocks] = offsets[num_blocks - 1]; + } + + std::string out; + out.resize(kTrieHeaderSize); + uint64_t sections[kNumSectionOffsets] = {}; + if (num_keys > 0) { + out.append(key(0).data(), prefix_len); + } + sections[kTrieSection] = out.size(); + out.append(trie); + sections[kBucketTableSection] = out.size(); + Status s = AppendColumn(bucket_positions, &out); + if (!s.ok()) { + return s; + } + sections[kBucketDataSection] = out.size(); + out.append(bucket_data); + sections[kOffsetTableSection] = out.size(); + s = AppendColumn(offsets, &out); + if (!s.ok()) { + return s; + } + if (!derived) { + sections[kSizeTableSection] = out.size(); + std::vector sizes(num_blocks); + for (uint64_t i = 0; i < num_blocks; i++) { + sizes[i] = handles_[i].size; + } + s = AppendColumn(sizes, &out); + if (!s.ok()) { + return s; + } + } + sections[kRunTableSection] = out.size(); + PutFixed64(&out, runs_.size()); + uint64_t extra_blocks = 0; + for (const Run& run : runs_) { + extra_blocks += run.extra_blocks; + PutFixed64(&out, run.key_index); + PutFixed64(&out, extra_blocks); + } + sections[kTagTableSection] = out.size(); + PutFixed64(&out, tags_.size()); + for (const Tag& tag : tags_) { + PutFixed64(&out, tag.block_index); + PutFixed64(&out, tag.tag); + } + sections[kEnd] = out.size(); + + char* h = out.data(); + EncodeFixed32(h + kMagicPos, kMagic); + EncodeFixed32(h + kVersionPos, kVersion); + EncodeFixed32(h + kFlagsPos, 0); + EncodeFixed32(h + kBucketSizePos, bucket_size_); + EncodeFixed64(h + kNumBlocksPos, num_blocks); + EncodeFixed64(h + kNumKeysPos, num_keys); + EncodeFixed32(h + kPrefixLenPos, static_cast(prefix_len)); + EncodeFixed32(h + kMaxKeyLenPos, static_cast(max_key_len)); + EncodeFixed32(h + kIdWidthPos, id_width); + EncodeFixed32(h + kTrieRootPos, trie_root); + EncodeFixed64(h + kSizeGapPos, derived ? size_gap : kExplicitSizes); + for (size_t i = 0; i < kNumSectionOffsets; i++) { + EncodeFixed64(h + kSectionsPos + 8 * i, sections[i]); + } + *output = std::move(out); + return Status::OK(); +} + +uint64_t TrieBuilder::EntrySize(uint64_t drop, uint64_t add) { + return add + (drop < kShortEntryLimit && add < kShortEntryLimit + ? 1 + : 1 + VarintLength(drop) + VarintLength(add)); +} + +uint64_t TrieBuilder::SizeBound(const SizeStats& stats, uint32_t bucket_size) { + assert(bucket_size > 0); + // A group's residuals never exceed the range of its values. + auto column = [](uint64_t count, uint64_t max_range) { + const uint64_t num_groups = (count + kGroupSize - 1) >> kGroupBits; + return kColumnHeaderSize + num_groups * kGroupHeaderSize + + num_groups * 8 * BitsFor(max_range) + kColumnPadding; + }; + const uint64_t n = stats.num_blocks; + const uint64_t num_buckets = (n + bucket_size - 1) / bucket_size; + // Path and label bytes never exceed the suffix bytes. A trie with b + // terminal nodes has fewer than 2b nodes, each costing at most 17 bytes + // besides its path: header, path length, bucket id, child count, the + // pointer to it, and its share of a bitmap. + const uint64_t trie = stats.suffix_bytes + 17 * 2 * num_buckets; + // Offsets come in block order, so a group of 64 offsets spans at most 64 + // steps (the last one to the sentinel), and a size is at most the largest. + // Finish() fails for residuals wider than kMaxColumnWidth bits, so clamping + // each term there keeps the product from wrapping and still bounds them. + constexpr uint64_t kMaxRange = uint64_t{1} << kMaxColumnWidth; + const uint64_t max_step = std::min(stats.max_block_step, kMaxRange) + + std::min(stats.max_block_size, kMaxRange); + uint64_t size = kTrieHeaderSize + stats.max_key_len + trie + + column(num_buckets + 1, stats.entry_bytes) + + stats.entry_bytes + column(n + 1, kGroupSize * max_step) + + 2 * 8 + + kSideTableEntrySize * (stats.duplicates + stats.tagged); + if (stats.irregular_gaps) { + size += column(n, stats.max_block_size); + } + return size; +} + +// ============================================================================ +// Trie +// ============================================================================ + +Status Trie::Column::Init(const char* begin, const char* end, + uint64_t expected_count) { + const uint64_t avail = lossless_cast(end - begin); + if (avail < kColumnHeaderSize + kColumnPadding) { + return Corrupt("truncated column"); + } + count = DecodeFixed64(begin); + if (count != expected_count) { + return Corrupt("bad column header"); + } + const uint64_t num_groups = (count + kGroupSize - 1) >> kGroupBits; + if (num_groups > + (avail - kColumnHeaderSize - kColumnPadding) / kGroupHeaderSize) { + return Corrupt("bad column size"); + } + groups = begin + kColumnHeaderSize; + residuals = groups + num_groups * kGroupHeaderSize; + // The groups' residuals must follow each other with no gaps. + uint64_t words = 0; + for (uint64_t g = 0; g < num_groups; g++) { + const uint32_t layout = DecodeFixed32(groups + g * kGroupHeaderSize + 12); + const uint32_t width = layout & kWidthMask; + if (width > kMaxColumnWidth || layout >> kWidthBits != words) { + return Corrupt("bad column group"); + } + words += + (std::min(count - (g << kGroupBits), kGroupSize) * width + 63) / 64; + } + if (lossless_cast(end - residuals) != words * 8 + kColumnPadding) { + return Corrupt("bad column size"); + } + return Status::OK(); +} + +uint64_t Trie::Column::Get(uint64_t i) const { + const char* group = groups + (i >> kGroupBits) * kGroupHeaderSize; + const uint64_t j = i & (kGroupSize - 1); + const uint32_t layout = DecodeFixed32(group + 12); + const uint32_t width = layout & kWidthMask; + // With width 0 this reads bytes that are masked away. + const uint64_t bit = uint64_t{layout >> kWidthBits} * 64 + j * width; + const uint64_t residual = (DecodeFixed64(residuals + bit / 8) >> (bit % 8)) & + ((uint64_t{1} << width) - 1); + return DecodeFixed64(group) + uint64_t{DecodeFixed32(group + 8)} * j + + residual; +} + +// Parsed trie node. +struct TrieIterator::Node { + const char* path = nullptr; + const uint8_t* labels = nullptr; + const char* bitmap = nullptr; + const char* pointers = nullptr; + uint32_t pos = 0; + uint32_t path_len = 0; + uint32_t bucket_id = 0; + uint32_t child_count = 0; + uint32_t pointer_width = 0; + uint8_t kind = 0; + bool terminal = false; + + // Parses the node at pos, which Init() has validated. + void Parse(const char* trie, uint32_t node_pos, uint32_t id_width) { + const char* p = trie + node_pos; + const uint8_t header = Byte(p++, 0); + pos = node_pos; + kind = header & kKindMask; + terminal = (header & kTerminalBit) != 0; + pointer_width = ((header >> kPointerWidthShift) & 3) + 1; + path_len = Byte(p++, 0); + if (path_len == kLongPathMarker) { + p = GetVarint32Ptr(p, p + 5, &path_len); + } + path = p; + p += path_len; + if (terminal) { + bucket_id = LoadLE(p, id_width); + p += id_width; + } + if (kind == kKindSparse) { + child_count = Byte(p, 0) + 1u; + labels = lossless_cast(p + 1); + p += 1 + child_count; + } else if (kind == kKindBitmap) { + bitmap = p; + child_count = 0; + for (int i = 0; i < 4; i++) { + child_count += BitsSetToOne(DecodeFixed64(p + 8 * i)); + } + p += kBitmapBytes; + } else { + child_count = 0; + } + pointers = p; + } + + uint32_t Child(uint32_t index) const { + return pos - LoadLE(pointers + index * pointer_width, pointer_width); + } + + // Labels of the first and last child. + uint8_t FirstLabel() const { + if (kind == kKindSparse) { + return labels[0]; + } + for (int i = 0; i < 4; i++) { + const uint64_t w = DecodeFixed64(bitmap + 8 * i); + if (w != 0) { + return BitmapLabel(i, CountTrailingZeroBits(w)); + } + } + return 0; + } + uint8_t LastLabel() const { + if (kind == kKindSparse) { + return labels[child_count - 1]; + } + for (int i = 3; i >= 0; i--) { + const uint64_t w = DecodeFixed64(bitmap + 8 * i); + if (w != 0) { + return BitmapLabel(i, FloorLog2(w)); + } + } + return 0; + } +}; + +Status Trie::Init(const Slice& data) { + Status s = Parse(data); + if (!s.ok()) { + *this = Trie(); + } + return s; +} + +Status Trie::Parse(const Slice& data) { + const char* p = data.data(); + // Checked before the header size, which other versions need not share. + if (data.size() < kVersionPos + 4 || DecodeFixed32(p + kMagicPos) != kMagic) { + return Corrupt("bad magic number"); + } + const uint32_t version = DecodeFixed32(p + kVersionPos); + if (version != kVersion) { + return Status::NotSupported( + "Trie index: format version " + std::to_string(version) + + " is not readable by this release; open with " + "index_mode=kStandardDefault and compact to rewrite the trie"); + } + if (data.size() < kTrieHeaderSize) { + return Corrupt("too short"); + } + if (DecodeFixed32(p + kFlagsPos) != 0) { + return Corrupt("unknown flags"); + } + bucket_size_ = DecodeFixed32(p + kBucketSizePos); + num_blocks_ = DecodeFixed64(p + kNumBlocksPos); + num_keys_ = DecodeFixed64(p + kNumKeysPos); + prefix_len_ = DecodeFixed32(p + kPrefixLenPos); + max_key_len_ = DecodeFixed32(p + kMaxKeyLenPos); + id_width_ = DecodeFixed32(p + kIdWidthPos); + trie_root_ = DecodeFixed32(p + kTrieRootPos); + size_gap_ = DecodeFixed64(p + kSizeGapPos); + explicit_sizes_ = size_gap_ == kExplicitSizes; + // Every block takes at least a quarter byte of the offset column, which + // also keeps the arithmetic below from overflowing. + if (bucket_size_ == 0 || num_keys_ > num_blocks_ || + (num_blocks_ == 0) != (num_keys_ == 0) || num_blocks_ / 4 > data.size() || + prefix_len_ > max_key_len_ || id_width_ < 1 || id_width_ > 4) { + return Corrupt("bad header"); + } + num_buckets_ = num_keys_ == 0 ? 0 : (num_keys_ - 1) / bucket_size_ + 1; + if (num_buckets_ > 0 && num_buckets_ - 1 >= uint64_t{1} << (8 * id_width_)) { + return Corrupt("bucket id width too small"); + } + + uint64_t sections[kNumSectionOffsets]; + for (size_t i = 0; i < kNumSectionOffsets; i++) { + sections[i] = DecodeFixed64(p + kSectionsPos + 8 * i); + } + if (sections[kTrieSection] != kTrieHeaderSize + uint64_t{prefix_len_} || + (sections[kSizeTableSection] != 0) != explicit_sizes_ || + sections[kEnd] != data.size()) { + return Corrupt("bad section offsets"); + } + uint64_t prev = sections[kTrieSection]; + for (size_t i = kTrieSection + 1; i < kNumSectionOffsets; i++) { + if (i == kSizeTableSection && !explicit_sizes_) { + continue; + } + if (sections[i] < prev) { + return Corrupt("bad section offsets"); + } + prev = sections[i]; + } + + prefix_ = p + kTrieHeaderSize; + trie_ = p + sections[kTrieSection]; + bucket_data_ = p + sections[kBucketDataSection]; + runs_ = p + sections[kRunTableSection]; + tags_ = p + sections[kTagTableSection]; + + Status s = + bucket_table_.Init(p + sections[kBucketTableSection], + p + sections[kBucketDataSection], num_buckets_ + 1); + if (!s.ok()) { + return s; + } + const uint64_t trie_size = + sections[kBucketTableSection] - sections[kTrieSection]; + const uint64_t bucket_data_size = + sections[kOffsetTableSection] - sections[kBucketDataSection]; + // A separator is the shared prefix, a bucket head spelled by the trie and + // suffixes from its bucket, so those bytes bound its length. A larger + // max_key_len is corrupt, and would loosen the length checks below and in + // the iterator. + if (max_key_len_ - prefix_len_ > trie_size + bucket_data_size) { + return Corrupt("max key length exceeds index contents"); + } + bucket_data_size_ = bucket_data_size; + const uint64_t offset_table_end = explicit_sizes_ + ? sections[kSizeTableSection] + : sections[kRunTableSection]; + s = offset_table_.Init(p + sections[kOffsetTableSection], + p + offset_table_end, num_blocks_ + 1); + if (!s.ok()) { + return s; + } + if (explicit_sizes_) { + s = size_table_.Init(p + sections[kSizeTableSection], + p + sections[kRunTableSection], num_blocks_); + if (!s.ok()) { + return s; + } + } + s = ValidateSideTables( + sections[kTagTableSection] - sections[kRunTableSection], + sections[kEnd] - sections[kTagTableSection]); + if (!s.ok()) { + return s; + } + return ValidateTrie(trie_size); +} + +Status Trie::ValidateSideTables(uint64_t run_table_size, + uint64_t tag_table_size) { + if (run_table_size < 8 || tag_table_size < 8) { + return Corrupt("truncated side table"); + } + num_runs_ = DecodeFixed64(runs_); + num_tags_ = DecodeFixed64(tags_); + if (num_runs_ > run_table_size / kSideTableEntrySize || + run_table_size != 8 + num_runs_ * kSideTableEntrySize || + num_tags_ > tag_table_size / kSideTableEntrySize || + tag_table_size != 8 + num_tags_ * kSideTableEntrySize) { + return Corrupt("bad side table size"); + } + uint64_t prev_key = 0; + uint64_t prev_extra = 0; + for (uint64_t r = 0; r < num_runs_; r++) { + const char* e = runs_ + 8 + r * kSideTableEntrySize; + const uint64_t key_index = DecodeFixed64(e); + const uint64_t extra = DecodeFixed64(e + 8); + if (key_index >= num_keys_ || (r > 0 && key_index <= prev_key) || + extra <= prev_extra) { + return Corrupt("bad run table"); + } + prev_key = key_index; + prev_extra = extra; + } + if (prev_extra != num_blocks_ - num_keys_) { + return Corrupt("run table does not match block count"); + } + uint64_t prev_block = 0; + for (uint64_t t = 0; t < num_tags_; t++) { + const char* e = tags_ + 8 + t * kSideTableEntrySize; + const uint64_t block = DecodeFixed64(e); + if (block >= num_blocks_ || (t > 0 && block <= prev_block) || + DecodeFixed64(e + 8) == 0) { + return Corrupt("bad tag table"); + } + prev_block = block; + } + return Status::OK(); +} + +uint64_t Trie::SideTableLowerBound(const char* table, uint64_t count, + uint64_t value) { + uint64_t lo = 0; + uint64_t hi = count; + while (lo < hi) { + const uint64_t mid = lo + (hi - lo) / 2; + if (DecodeFixed64(table + 8 + mid * kSideTableEntrySize) < value) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; +} + +uint64_t Trie::RunKey(uint64_t r) const { + return DecodeFixed64(runs_ + 8 + r * kSideTableEntrySize); +} + +uint64_t Trie::RunExtraThrough(uint64_t r) const { + return DecodeFixed64(runs_ + 16 + r * kSideTableEntrySize); +} + +uint64_t Trie::TagBlock(uint64_t t) const { + return DecodeFixed64(tags_ + 8 + t * kSideTableEntrySize); +} + +uint64_t Trie::TagValue(uint64_t t) const { + return DecodeFixed64(tags_ + 16 + t * kSideTableEntrySize); +} + +Status Trie::ValidateTrie(uint64_t trie_size) { + if (num_keys_ == 0) { + return trie_size == 0 && trie_root_ == 0 ? Status::OK() + : Corrupt("unexpected trie"); + } + if (trie_root_ >= trie_size) { + return Corrupt("bad trie root"); + } + const uint64_t max_head_len = max_key_len_ - prefix_len_; + // Pre-order walk. A valid trie spells the bucket heads in bucket order. + struct Visit { + uint32_t pos; + uint64_t head_len; + }; + std::vector stack; + stack.push_back({trie_root_, 0}); + uint64_t next_bucket = 0; + uint64_t visited = 0; + while (!stack.empty()) { + const Visit v = stack.back(); + stack.pop_back(); + // Every node takes at least two bytes. + if (++visited > trie_size / 2) { + return Corrupt("trie has too many nodes"); + } + const char* p = trie_ + v.pos; + const char* end = trie_ + trie_size; + const uint8_t header = Byte(p++, 0); + const uint8_t kind = header & kKindMask; + const bool terminal = (header & kTerminalBit) != 0; + const uint32_t pointer_width = ((header >> kPointerWidthShift) & 3) + 1; + if ((header & kReservedNodeBits) != 0 || kind > kKindBitmap || + (kind == kKindLeaf && !terminal) || p >= end) { + return Corrupt("bad trie node"); + } + uint32_t path_len = Byte(p++, 0); + if (path_len == kLongPathMarker) { + p = GetVarint32Ptr(p, end, &path_len); + if (p == nullptr) { + return Corrupt("bad trie node"); + } + } + if (path_len > lossless_cast(end - p) || + path_len > max_head_len - v.head_len) { + return Corrupt("bad trie node path"); + } + p += path_len; + const uint64_t head_len = v.head_len + path_len; + if (terminal) { + if (lossless_cast(end - p) < id_width_ || + LoadLE(p, id_width_) != next_bucket) { + return Corrupt("bad trie bucket id"); + } + next_bucket++; + p += id_width_; + } + uint32_t child_count = 0; + if (kind == kKindSparse) { + if (p >= end) { + return Corrupt("bad trie node"); + } + child_count = Byte(p, 0) + 1u; + if (child_count > lossless_cast(end - p - 1)) { + return Corrupt("bad trie node"); + } + for (uint32_t i = 1; i < child_count; i++) { + if (Byte(p + 1, i) <= Byte(p + 1, i - 1)) { + return Corrupt("trie labels out of order"); + } + } + p += 1 + child_count; + } else if (kind == kKindBitmap) { + if (lossless_cast(end - p) < kBitmapBytes) { + return Corrupt("bad trie node"); + } + for (int i = 0; i < 4; i++) { + child_count += BitsSetToOne(DecodeFixed64(p + 8 * i)); + } + if (child_count == 0) { + return Corrupt("empty trie bitmap"); + } + p += kBitmapBytes; + } + if (child_count > 0 && head_len >= max_head_len) { + return Corrupt("bad trie node path"); + } + if (uint64_t{child_count} * pointer_width > + lossless_cast(end - p)) { + return Corrupt("bad trie node"); + } + // Push children last first, so the first child is visited next. + for (uint32_t i = child_count; i-- > 0;) { + const uint32_t pointer = LoadLE(p + i * pointer_width, pointer_width); + if (pointer == 0 || pointer > v.pos) { + return Corrupt("bad trie pointer"); + } + stack.push_back({v.pos - pointer, head_len + 1}); + } + } + if (next_bucket != num_buckets_) { + return Corrupt("trie does not match bucket count"); + } + return Status::OK(); +} + +// ============================================================================ +// TrieIterator +// ============================================================================ + +TrieIterator::TrieIterator(const Trie* index) + : index_(index), block_(index->num_blocks_) { + key_.reserve(std::max( + index_->prefix_len_, + std::min(index_->max_key_len_, kInitialKeyCapacity))); + key_.assign(index_->prefix_, index_->prefix_len_); +} + +void TrieIterator::Invalidate() { block_ = index_->num_blocks_; } + +Status TrieIterator::EnterBlock(uint64_t block) { + const Trie& index = *index_; + const uint64_t offset = index.offset_table_.Get(block); + uint64_t size; + if (index.explicit_sizes_) { + size = index.size_table_.Get(block); + } else { + // Init() does not check the offsets, so a corrupt column could make the + // size wrap around. + const uint64_t next = index.offset_table_.Get(block + 1); + if (next < offset || next - offset < index.size_gap_) { + Invalidate(); + return Corrupt("bad offset table"); + } + size = next - offset - index.size_gap_; + } + block_ = block; + handle_ = TrieBlockHandle{offset, size}; + return Status::OK(); +} + +void TrieIterator::Descend(const char* target, size_t target_len, + Candidate* floor, Candidate* ceiling) const { + floor->spell = Spell::kNone; + ceiling->spell = Spell::kNone; + uint32_t pos = index_->trie_root_; + size_t depth = 0; + Node node; + for (;;) { + node.Parse(index_->trie_, pos, index_->id_width_); + const size_t n = std::min(node.path_len, target_len - depth); + const size_t matched = CommonPrefix(target + depth, node.path, n); + if (matched < n) { + // Every head below this node is on one side of target. + if (Byte(target, depth + matched) < Byte(node.path, matched)) { + *ceiling = Candidate{depth, pos, Spell::kSmallest}; + } else { + *floor = Candidate{depth, pos, Spell::kLargest}; + } + return; + } + if (n < node.path_len) { + // Target ends inside the path: it precedes every head below. + *ceiling = Candidate{depth, pos, Spell::kSmallest}; + return; + } + const size_t label_depth = depth + node.path_len; + if (node.terminal) { + if (label_depth == target_len) { + *ceiling = Candidate{label_depth, pos, Spell::kTerminal}; + return; + } + *floor = Candidate{label_depth, pos, Spell::kTerminal}; + } else if (label_depth == target_len) { + *ceiling = Candidate{depth, pos, Spell::kSmallest}; + return; + } + if (node.child_count == 0) { + return; + } + const uint8_t c = Byte(target, label_depth); + uint32_t rank; + bool present; + bool has_below = false; + bool has_above = false; + uint8_t below = 0; + uint8_t above = 0; + if (node.kind == kKindSparse) { + rank = 0; + while (rank < node.child_count && node.labels[rank] < c) { + rank++; + } + present = rank < node.child_count && node.labels[rank] == c; + if (rank > 0) { + has_below = true; + below = node.labels[rank - 1]; + } + const uint32_t next = present ? rank + 1 : rank; + if (next < node.child_count) { + has_above = true; + above = node.labels[next]; + } + } else { + uint64_t words[4]; + for (int i = 0; i < 4; i++) { + words[i] = DecodeFixed64(node.bitmap + 8 * i); + } + const int wi = c / 64; + const int bi = c % 64; + rank = 0; + for (int i = 0; i < wi; i++) { + rank += BitsSetToOne(words[i]); + } + const uint64_t lower = words[wi] & ((uint64_t{1} << bi) - 1); + rank += BitsSetToOne(lower); + present = ((words[wi] >> bi) & 1) != 0; + if (lower != 0) { + has_below = true; + below = BitmapLabel(wi, FloorLog2(lower)); + } else { + for (int i = wi - 1; i >= 0; i--) { + if (words[i] != 0) { + has_below = true; + below = BitmapLabel(i, FloorLog2(words[i])); + break; + } + } + } + const uint64_t higher = + bi == 63 ? 0 : (words[wi] & ~((uint64_t{2} << bi) - 1)); + if (higher != 0) { + has_above = true; + above = BitmapLabel(wi, CountTrailingZeroBits(higher)); + } else { + for (int i = wi + 1; i < 4; i++) { + if (words[i] != 0) { + has_above = true; + above = BitmapLabel(i, CountTrailingZeroBits(words[i])); + break; + } + } + } + } + if (has_below) { + *floor = Candidate{label_depth, node.Child(rank - 1), Spell::kLargest, + true, below}; + } + if (has_above) { + *ceiling = Candidate{label_depth, node.Child(present ? rank + 1 : rank), + Spell::kSmallest, true, above}; + } + if (!present) { + return; + } + pos = node.Child(rank); + depth = label_depth + 1; + } +} + +uint64_t TrieIterator::SpellHead(const Candidate& candidate, + const char* target) { + key_.resize(index_->prefix_len_); + key_.append(target, candidate.depth); + Node node; + if (candidate.spell == Spell::kTerminal) { + node.Parse(index_->trie_, candidate.node, index_->id_width_); + return node.bucket_id; + } + if (candidate.has_label) { + key_.push_back(lossless_cast(candidate.label)); + } + uint32_t pos = candidate.node; + for (;;) { + node.Parse(index_->trie_, pos, index_->id_width_); + key_.append(node.path, node.path_len); + if (candidate.spell == Spell::kSmallest) { + if (node.terminal) { + return node.bucket_id; + } + key_.push_back(lossless_cast(node.FirstLabel())); + pos = node.Child(0); + } else { + if (node.child_count == 0) { + return node.bucket_id; + } + key_.push_back(lossless_cast(node.LastLabel())); + pos = node.Child(node.child_count - 1); + } + } +} + +void TrieIterator::EnterBucket(uint64_t bucket) { + bucket_ = bucket; + key_index_ = bucket * index_->bucket_size_; + // Init() does not check bucket positions. Clamping keeps a corrupt one + // inside the bucket data, where DecodeNext() bounds-checks every separator. + const uint64_t size = index_->bucket_data_size_; + bucket_pos_ = + index_->bucket_data_ + std::min(index_->bucket_table_.Get(bucket), size); + bucket_end_ = index_->bucket_data_ + + std::min(index_->bucket_table_.Get(bucket + 1), size); +} + +bool TrieIterator::EnterBucketAt(const Candidate& candidate, const char* target, + uint64_t bucket) { + if (candidate.spell == Spell::kNone || + SpellHead(candidate, target) != bucket) { + Invalidate(); + return false; + } + EnterBucket(bucket); + return true; +} + +void TrieIterator::EnterKey() { + uint64_t extra_before = 0; + uint64_t extra = 0; + const uint64_t num_runs = index_->num_runs_; + if (num_runs > 0) { + // First run at or after the current key. + const uint64_t lo = + Trie::SideTableLowerBound(index_->runs_, num_runs, key_index_); + if (lo > 0) { + extra_before = index_->RunExtraThrough(lo - 1); + } + if (lo < num_runs && index_->RunKey(lo) == key_index_) { + extra = index_->RunExtraThrough(lo) - extra_before; + } + } + run_begin_ = key_index_ + extra_before; + run_end_ = run_begin_ + 1 + extra; +} + +Status TrieIterator::ApplyTag(uint64_t tag) { + const uint64_t num_tags = index_->num_tags_; + if (num_tags == 0) { + return EnterBlock(run_begin_); + } + // First tagged block at or after the current one. + uint64_t lo = Trie::SideTableLowerBound(index_->tags_, num_tags, run_begin_); + // A block without a tag sorts after every tag of its key. + if (lo == num_tags || index_->TagBlock(lo) != run_begin_) { + return EnterBlock(run_begin_); + } + for (uint64_t b = run_begin_; b < run_end_; b++) { + uint64_t block_tag = 0; + if (lo < num_tags && index_->TagBlock(lo) == b) { + block_tag = index_->TagValue(lo); + lo++; + } + if (block_tag == 0 || tag >= block_tag) { + return EnterBlock(b); + } + } + // Every block of this key precedes the target. + return NextKey(); +} + +Status TrieIterator::DecodeNext(size_t* shared, const char** suffix, + size_t* suffix_len) { + if (bucket_pos_ >= bucket_end_) { + return Corrupt("truncated bucket"); + } + const uint8_t header = Byte(bucket_pos_++, 0); + uint32_t drop = header >> 4; + uint32_t add = header & 0xF; + if (header == kLongEntryMarker) { + const char* p = GetVarint32Ptr(bucket_pos_, bucket_end_, &drop); + if (p != nullptr) { + p = GetVarint32Ptr(p, bucket_end_, &add); + } + if (p == nullptr) { + return Corrupt("bad bucket entry"); + } + bucket_pos_ = p; + } + const size_t prefix_len = index_->prefix_len_; + const size_t len = key_.size() - prefix_len; + if (drop > len || add > lossless_cast(bucket_end_ - bucket_pos_) || + len - drop + add > index_->max_key_len_ - prefix_len) { + return Corrupt("bad bucket entry"); + } + *shared = len - drop; + key_.resize(prefix_len + *shared); + key_.append(bucket_pos_, add); + *suffix = bucket_pos_; + *suffix_len = add; + bucket_pos_ += add; + key_index_++; + return Status::OK(); +} + +Status TrieIterator::DecodeNext() { + size_t shared; + const char* suffix; + size_t suffix_len; + return DecodeNext(&shared, &suffix, &suffix_len); +} + +Status TrieIterator::ScanBucket(const char* target, size_t target_len, + bool* found, bool* equal) { + const size_t prefix_len = index_->prefix_len_; + const uint64_t end = + std::min(index_->num_keys_, key_index_ + index_->bucket_size_); + // Length of the prefix the current separator shares with target. The + // current separator is always below target. + size_t matched = CommonPrefix(key_.data() + prefix_len, target, + std::min(key_.size() - prefix_len, target_len)); + while (key_index_ + 1 < end) { + size_t shared; + const char* suffix; + size_t suffix_len; + Status s = DecodeNext(&shared, &suffix, &suffix_len); + if (!s.ok()) { + return s; + } + if (shared > matched) { + // Same byte as the previous separator where it fell below target. + continue; + } + if (shared < matched) { + // Larger than the previous separator where it matched target. + *found = true; + *equal = false; + return Status::OK(); + } + const size_t rest = target_len - matched; + const size_t n = std::min(suffix_len, rest); + const size_t common = CommonPrefix(suffix, target + matched, n); + if (common < n) { + if (Byte(suffix, common) > Byte(target, matched + common)) { + *found = true; + *equal = false; + return Status::OK(); + } + matched += common; + } else if (rest <= suffix_len) { + // Target is a prefix of this separator, or equal to it. + *found = true; + *equal = rest == suffix_len; + return Status::OK(); + } else { + matched += suffix_len; + } + } + *found = false; + return Status::OK(); +} + +Status TrieIterator::SeekToFirst() { + if (index_->num_keys_ == 0) { + Invalidate(); + return Status::OK(); + } + const Candidate first{0, index_->trie_root_, Spell::kSmallest}; + if (SpellHead(first, nullptr) != 0) { + Invalidate(); + return Corrupt("first bucket head not found"); + } + EnterBucket(0); + EnterKey(); + return EnterBlock(run_begin_); +} + +Status TrieIterator::Seek(const Slice& key, uint64_t tag) { + if (index_->num_keys_ == 0) { + Invalidate(); + return Status::OK(); + } + const char* target = key.data(); + size_t target_len = key.size(); + const size_t prefix_len = index_->prefix_len_; + if (prefix_len > 0) { + const int cmp = + memcmp(target, index_->prefix_, std::min(prefix_len, target_len)); + if (cmp > 0) { + Invalidate(); + return Status::OK(); + } + if (cmp < 0 || target_len < prefix_len) { + return SeekToFirst(); + } + target += prefix_len; + target_len -= prefix_len; + } + Candidate floor; + Candidate ceiling; + Descend(target, target_len, &floor, &ceiling); + if (ceiling.spell == Spell::kTerminal) { + // Target is a bucket head, so the floor's bucket holds only smaller keys. + EnterBucket(SpellHead(ceiling, target)); + EnterKey(); + return ApplyTag(tag); + } + if (floor.spell == Spell::kNone) { + // Every head is greater than target, so tags do not matter. + return SeekToFirst(); + } + const uint64_t bucket = SpellHead(floor, target); + EnterBucket(bucket); + bool found; + bool equal = false; + Status s = ScanBucket(target, target_len, &found, &equal); + if (!s.ok()) { + Invalidate(); + return s; + } + if (!found) { + if (bucket + 1 == index_->num_buckets_) { + Invalidate(); + return Status::OK(); + } + // The next head is greater than target, so tags do not matter. + if (!EnterBucketAt(ceiling, target, bucket + 1)) { + return Corrupt("bucket head out of order"); + } + } + EnterKey(); + return equal ? ApplyTag(tag) : EnterBlock(run_begin_); +} + +Status TrieIterator::NextKey() { + if (key_index_ + 1 >= index_->num_keys_) { + Invalidate(); + return Status::OK(); + } + const uint64_t bucket_end = + std::min(index_->num_keys_, (bucket_ + 1) * index_->bucket_size_); + if (key_index_ + 1 < bucket_end) { + Status s = DecodeNext(); + if (!s.ok()) { + Invalidate(); + return s; + } + } else { + // The next bucket's head is the smallest head after the current key. + const size_t prefix_len = index_->prefix_len_; + scratch_.assign(key_.data() + prefix_len, key_.size() - prefix_len); + scratch_.push_back('\0'); + Candidate floor; + Candidate ceiling; + Descend(scratch_.data(), scratch_.size(), &floor, &ceiling); + if (!EnterBucketAt(ceiling, scratch_.data(), bucket_ + 1)) { + return Corrupt("bucket head out of order"); + } + } + EnterKey(); + return EnterBlock(run_begin_); +} + +Status TrieIterator::AdvanceTo(uint64_t key_index) { + while (key_index_ < key_index) { + Status s = DecodeNext(); + if (!s.ok()) { + Invalidate(); + return s; + } + } + return Status::OK(); +} + +Status TrieIterator::EnterKeyFromEnd() { + EnterKey(); + return EnterBlock(run_end_ - 1); +} + +Status TrieIterator::SeekToLast() { + if (index_->num_keys_ == 0) { + Invalidate(); + return Status::OK(); + } + const uint64_t last_bucket = index_->num_buckets_ - 1; + const Candidate last{0, index_->trie_root_, Spell::kLargest}; + if (SpellHead(last, nullptr) != last_bucket) { + Invalidate(); + return Corrupt("last bucket head not found"); + } + EnterBucket(last_bucket); + Status s = AdvanceTo(index_->num_keys_ - 1); + return s.ok() ? EnterKeyFromEnd() : s; +} + +Status TrieIterator::Prev() { + if (!Valid()) { + return Status::OK(); + } + if (block_ > run_begin_) { + return EnterBlock(block_ - 1); + } + if (key_index_ == 0) { + Invalidate(); + return Status::OK(); + } + // The previous key's bucket starts with the largest head below the current + // key. Front coding only decodes forward, so decode from that head. + const uint64_t prev_key = key_index_ - 1; + const size_t prefix_len = index_->prefix_len_; + scratch_.assign(key_.data() + prefix_len, key_.size() - prefix_len); + Candidate floor; + Candidate ceiling; + Descend(scratch_.data(), scratch_.size(), &floor, &ceiling); + if (!EnterBucketAt(floor, scratch_.data(), prev_key / index_->bucket_size_)) { + return Corrupt("bucket head out of order"); + } + Status s = AdvanceTo(prev_key); + return s.ok() ? EnterKeyFromEnd() : s; +} + +Status TrieIterator::Next() { + if (!Valid()) { + return Status::OK(); + } + if (block_ + 1 < run_end_) { + return EnterBlock(block_ + 1); + } + return NextKey(); +} + +} // namespace trie_index +} // namespace ROCKSDB_NAMESPACE diff --git a/utilities/trie_index/trie.h b/utilities/trie_index/trie.h new file mode 100644 index 000000000000..1df7b86f49a7 --- /dev/null +++ b/utilities/trie_index/trie.h @@ -0,0 +1,355 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// This source code is licensed under both the GPLv2 (found in the +// COPYING file in the root directory) and Apache 2.0 License +// (found in the LICENSE.Apache file in the root directory). +// +// ***************************************************************** +// EXPERIMENTAL - subject to change while under development +// ***************************************************************** +// +// Serialized trie index. +// +// Maps the separator of every data block in an SST to the block's handle. +// Separators are user keys (bytewise ordered) plus an optional tag (packed +// sequence number and value type), and Seek finds the first block whose +// separator is at or after a target in internal key order. +// +// Structure, following the block-tree design used by Lucene and tantivy: +// - The bytes shared by every separator are stored once. +// - Distinct separators are grouped into fixed-size buckets. Inside a bucket +// each separator is front coded against the previous one. +// - A path-compressed byte trie over the first separator of each bucket (the +// bucket head) routes a Seek to one bucket. Trie nodes are typed (leaf, +// sparse or bitmap), hold their compressed path bytes inline, and point to +// children through narrow backward-relative offsets. +// - Block offsets are stored as a linear frame-of-reference column, so +// there is no file size limit. Block sizes are derived from consecutive +// offsets when the gaps between blocks are uniform. +// - Blocks that share a separator (one user key spanning several blocks) +// and nonzero tags are stored in small sparse side tables. +// +// Serialized layout. Integers are little-endian and read with unaligned +// loads, so the index is used in place at any alignment and on any host. +// +// Header (kTrieHeaderSize bytes): +// u32 magic "TRIE" +// u32 version 2. Version 1 was an earlier layout that readers +// reject. +// u32 flags 0; readers reject any other value +// u32 bucket_size separators per bucket, at least 1 +// u64 num_blocks data blocks indexed +// u64 num_keys distinct separators, at most num_blocks +// u32 prefix_len length of the prefix shared by every separator +// u32 max_key_len length of the longest separator +// u32 id_width bytes per bucket id in trie nodes, 1 to 4 +// u32 trie_root position of the root node within the trie section +// u64 size_gap block size = next block offset - offset - size_gap, +// or all ones when sizes are stored explicitly +// u64 trie_offset, bucket_table_offset, bucket_data_offset, +// offset_table_offset, size_table_offset, run_table_offset, +// tag_table_offset, end_offset +// section boundaries from the start of the index. +// Sections are contiguous and in this order; +// size_table_offset is 0 when sizes are derived, and +// end_offset is the size of the index. +// Shared prefix bytes +// Trie section +// Bucket table: column of num_buckets + 1 positions in the bucket data +// Bucket data +// Offset table: column of num_blocks + 1 non-decreasing offsets. When +// sizes are derived, the last one is where a next block would start; +// otherwise it repeats the last block's offset +// Size table (only when sizes are explicit): column of num_blocks sizes +// Run table: u64 count, then count entries of +// {u64 key_index, u64 extra_blocks_through}. One entry per separator +// shared by more than one block, in key order. extra_blocks_through is +// the total number of extra blocks of this run and all earlier runs. +// Tag table: u64 count, then count entries of {u64 block_index, u64 tag} +// for every block with a nonzero tag, in block order. It ends the index. +// +// Column (linear frame of reference, groups of 64 values): +// u64 count, then per group {u64 base, u32 slope, u32 layout}, then the +// groups' residuals back to back, then 8 zero bytes. layout is +// words << 6 | width: the group's residuals are `width` bits each (at most +// 56), packed little-endian from 8-byte word `words` of the residuals, +// which is the total size of the earlier groups' residuals. A group of n +// values takes ceil(n * width / 64) words. +// Value i is base + slope * (i % 64) + residual(i) in 64-bit unsigned +// arithmetic, where base, slope and width belong to group i / 64. +// +// Trie node: +// u8 header: bits 0-1 kind (0 leaf, 1 sparse, 2 bitmap), bit 2 terminal, +// bits 3-4 child pointer width - 1, bits 5-7 zero. A leaf is terminal. +// Path length: u8, or 255 followed by a varint32. +// Path bytes: the compressed edge leading into this node's labels. +// Terminal nodes: bucket id (id_width bytes). The node's path spells that +// bucket's head. +// Sparse nodes: u8 child count - 1, then the child labels in ascending +// order. +// Bitmap nodes: 256-bit label bitmap as four u64 words. +// Child pointers, one per label, pointer width bytes each. A child starts +// `pointer` bytes before its parent; children are always written first. +// +// Bucket data: for every bucket, its separators after the head, each as +// u8 (drop << 4 | add) when drop < 15 and add < 15, otherwise 0xFF +// followed by varint32 drop and varint32 add; then `add` suffix bytes. +// The separator is the previous one minus its last `drop` bytes, plus the +// suffix. + +#pragma once + +#include +#include +#include +#include + +#include "rocksdb/slice.h" +#include "rocksdb/status.h" + +namespace ROCKSDB_NAMESPACE { +namespace trie_index { + +inline constexpr size_t kTrieHeaderSize = 120; + +// Offset and size of a data block in the SST file. +struct TrieBlockHandle { + uint64_t offset = 0; + uint64_t size = 0; +}; + +// Builds a trie index from the separators of an SST's data blocks, added in +// block order. +class TrieBuilder { + public: + // bucket_size is the number of distinct separators per bucket. Larger + // buckets make the index smaller and Seek slower. + explicit TrieBuilder(uint32_t bucket_size); + + // Adds the next data block. Separators must be bytewise non-decreasing, + // and so must block offsets. + // Consecutive equal separators describe one user key spanning several + // blocks. tag is 0 when the block can be found by user key alone, + // otherwise the packed sequence number and type of the block's last key. + void AddBlock(const Slice& separator, const TrieBlockHandle& handle, + uint64_t tag); + + // Blocks added so far, and the last distinct separator added, which is the + // most recent block's separator when separators are in order (empty when + // there are no blocks). The separator stays valid until the next + // AddBlock(). + uint64_t NumBlocks() const { return handles_.size(); } + Slice LastSeparator() const; + + // Serializes the index into *output. Fails if separators or offsets were + // added out of order, or a column value or the trie does not fit the + // format. + Status Finish(std::string* output); + + // Bytes a separator takes in its bucket when it drops `drop` bytes from the + // end of the previous separator and adds `add`. + static uint64_t EntrySize(uint64_t drop, uint64_t add); + + // Running totals over the blocks added so far, from which SizeBound() + // bounds the size Finish() produces. + struct SizeStats { + uint64_t num_blocks = 0; + // Bytes of each separator past the prefix it shares with the previous + // one, and the EntrySize() of every separator. + uint64_t suffix_bytes = 0; + uint64_t entry_bytes = 0; + uint64_t max_key_len = 0; + // Separators equal to the previous one, and blocks that may get a + // nonzero tag. + uint64_t duplicates = 0; + uint64_t tagged = 0; + // Largest distance between consecutive block offsets, largest block + // size, and whether the gaps between blocks differ. + uint64_t max_block_step = 0; + uint64_t max_block_size = 0; + bool irregular_gaps = false; + }; + // bucket_size is the builder's, and must be positive. + static uint64_t SizeBound(const SizeStats& stats, uint32_t bucket_size); + + private: + const uint32_t bucket_size_; + // Distinct separators, concatenated, and the end of each one. + std::string key_bytes_; + std::vector key_ends_; + // A distinct separator shared by more than one block. + struct Run { + uint64_t key_index; + // Number of blocks sharing the separator, minus one. + uint64_t extra_blocks; + }; + struct Tag { + uint64_t block_index; + uint64_t tag; + }; + std::vector runs_; + std::vector handles_; + // Every block with a nonzero tag. + std::vector tags_; + bool out_of_order_ = false; +}; + +// Read-only view of a serialized trie index. Init() validates the header, +// section and column layouts, side tables and trie. It does not visit every +// bucket position, block offset or separator; the iterator keeps reads of +// them inside the index and returns Corruption for a separator out of bounds +// or a block size that would wrap around. Other wrong bytes give a wrong key +// or handle, as the index has no checksum of its own. +class Trie { + public: + Trie() = default; + + // Parses data, which must stay valid and unchanged for the lifetime of + // this object and its iterators. On failure the index is left empty. + Status Init(const Slice& data); + + private: + friend class TrieIterator; + friend class TrieTest; + + uint64_t NumBlocks() const { return num_blocks_; } + + // One frame-of-reference column. + struct Column { + const char* groups = nullptr; + const char* residuals = nullptr; + uint64_t count = 0; + + // Parses a column of expected_count values occupying exactly + // [begin, end). + Status Init(const char* begin, const char* end, uint64_t expected_count); + uint64_t Get(uint64_t i) const; + }; + + Status ValidateSideTables(uint64_t run_table_size, uint64_t tag_table_size); + // Index of the first entry of a side table, of count entries, whose first + // field is at least value. + static uint64_t SideTableLowerBound(const char* table, uint64_t count, + uint64_t value); + uint64_t RunKey(uint64_t r) const; + uint64_t RunExtraThrough(uint64_t r) const; + uint64_t TagBlock(uint64_t t) const; + uint64_t TagValue(uint64_t t) const; + Status Parse(const Slice& data); + Status ValidateTrie(uint64_t trie_size); + + uint64_t num_blocks_ = 0; + uint64_t num_keys_ = 0; + uint64_t num_buckets_ = 0; + uint64_t size_gap_ = 0; + const char* prefix_ = nullptr; + const char* trie_ = nullptr; + const char* bucket_data_ = nullptr; + const char* runs_ = nullptr; + const char* tags_ = nullptr; + uint64_t num_runs_ = 0; + uint64_t num_tags_ = 0; + uint64_t bucket_data_size_ = 0; + Column bucket_table_; + Column offset_table_; + Column size_table_; + uint32_t bucket_size_ = 0; + uint32_t prefix_len_ = 0; + uint32_t max_key_len_ = 0; + uint32_t id_width_ = 0; + uint32_t trie_root_ = 0; + bool explicit_sizes_ = false; +}; + +// Iterates the blocks of a Trie in order. Positioning methods return +// Corruption if a separator or block handle they read is malformed, and leave +// the iterator invalid. +class TrieIterator { + public: + explicit TrieIterator(const Trie* index); + + Status SeekToFirst(); + Status SeekToLast(); + // Positions at the first block whose separator is at or after (key, tag) + // in internal key order: by key ascending, then by tag descending, where a + // block without a tag sorts after every tag of its key. + Status Seek(const Slice& key, uint64_t tag); + Status Next(); + Status Prev(); + + bool Valid() const { return block_ < index_->num_blocks_; } + // The current block's separator. Valid until the next positioning call. + Slice key() const { return Slice(key_); } + TrieBlockHandle value() const { return handle_; } + + private: + // A trie position from which a bucket head can be spelled out: the + // terminal node itself, or the smallest or largest head below a node. + enum class Spell : uint8_t { kNone, kTerminal, kSmallest, kLargest }; + struct Candidate { + // Number of target bytes (after the shared prefix) that lead to node. + size_t depth = 0; + uint32_t node = 0; + Spell spell = Spell::kNone; + // Whether the label leading from the parent into node follows the depth + // bytes, rather than being included in them. + bool has_label = false; + uint8_t label = 0; + }; + struct Node; + + void Invalidate(); + // Finds the bucket heads closest below (floor) and at or above (ceiling) + // target, which excludes the shared prefix. + void Descend(const char* target, size_t target_len, Candidate* floor, + Candidate* ceiling) const; + // Writes the head that candidate leads to into key_ and returns its bucket + // id. target holds at least candidate.depth bytes. + uint64_t SpellHead(const Candidate& candidate, const char* target); + // Positions on the head of bucket, whose key is already in key_. + void EnterBucket(uint64_t bucket); + // Spells the head candidate leads to, which must be bucket's, and positions + // on it. Returns false, leaving the iterator invalid, if the trie leads + // elsewhere, which means it is corrupt. + bool EnterBucketAt(const Candidate& candidate, const char* target, + uint64_t bucket); + // Decodes the next separator in the bucket into key_. Returns the length + // shared with the previous separator, excluding the shared prefix, and the + // suffix through *suffix and *suffix_len. + Status DecodeNext(size_t* shared, const char** suffix, size_t* suffix_len); + // The same, for callers that only need key_. + Status DecodeNext(); + // Moves from the bucket head to the first separator >= target in this + // bucket. Sets *found to false if every separator is smaller, and *equal to + // whether the separator found equals target. + Status ScanBucket(const char* target, size_t target_len, bool* found, + bool* equal); + // Sets the block range of the current key index. + void EnterKey(); + // Positions on block and reads its handle. + Status EnterBlock(uint64_t block); + // After positioning on a separator equal to the Seek target, moves past its + // blocks whose tags are greater than tag. + Status ApplyTag(uint64_t tag); + Status NextKey(); + // Decodes forward in the current bucket until the current key index is + // key_index. + Status AdvanceTo(uint64_t key_index); + // Positions on the last block of the current key. + Status EnterKeyFromEnd(); + + const Trie* index_; + std::string key_; + std::string scratch_; + const char* bucket_pos_ = nullptr; + const char* bucket_end_ = nullptr; + uint64_t bucket_ = 0; + uint64_t key_index_ = 0; + // Blocks [run_begin_, run_end_) share the current separator. + uint64_t run_begin_ = 0; + uint64_t run_end_ = 0; + uint64_t block_; + TrieBlockHandle handle_; +}; + +} // namespace trie_index +} // namespace ROCKSDB_NAMESPACE diff --git a/utilities/trie_index/trie_index_db_test.cc b/utilities/trie_index/trie_index_db_test.cc index 80d2fdbb1dfd..4df562d77f0c 100644 --- a/utilities/trie_index/trie_index_db_test.cc +++ b/utilities/trie_index/trie_index_db_test.cc @@ -3797,6 +3797,65 @@ TEST_P(TrieIndexDBTest, ReopenWithoutTrieUDI) { ASSERT_EQ(keys[1], "key_b"); } +TEST_P(TrieIndexDBTest, EmptyUserKey) { + ASSERT_OK(OpenDB(/*block_size=*/256)); + // The empty key's value fills a data block, so that block's separator is + // the empty key. + ASSERT_OK(db_->Put(WriteOptions(), "", std::string(1000, 'E'))); + std::vector expected_keys = {""}; + for (char c = 'b'; c <= 'f'; c++) { + expected_keys.emplace_back(1, c); + ASSERT_OK( + db_->Put(WriteOptions(), expected_keys.back(), std::string(300, c))); + } + ASSERT_OK(db_->Flush(FlushOptions())); + ASSERT_NO_FATAL_FAILURE(VerifyGetBothIndexes("", std::string(1000, 'E'))); + ASSERT_NO_FATAL_FAILURE(VerifyGetBothIndexes("d", std::string(300, 'd'))); + ASSERT_NO_FATAL_FAILURE(VerifyScanBothIndexes(expected_keys)); +} + +TEST_P(TrieIndexDBTest, OnlyEmptyUserKey) { + ASSERT_OK(OpenDB()); + ASSERT_OK(db_->Put(WriteOptions(), "", "v")); + ASSERT_OK(db_->Flush(FlushOptions())); + ASSERT_NO_FATAL_FAILURE(VerifyGetBothIndexes("", "v")); + ASSERT_NO_FATAL_FAILURE(VerifyScanBothIndexes({""})); +} + +TEST_P(TrieIndexDBTest, OpenRejectsNonBytewiseComparator) { + options_.comparator = ReverseBytewiseComparator(); + ASSERT_EQ(OpenDB().code(), Status::Code::kNotSupported); +} + +TEST_P(TrieIndexDBTest, CreateColumnFamilyRejectsNonBytewiseComparator) { + ASSERT_OK(OpenDB()); + ColumnFamilyOptions cf_options(options_); + cf_options.comparator = ReverseBytewiseComparator(); + ColumnFamilyHandle* handle = nullptr; + ASSERT_EQ(db_->CreateColumnFamily(cf_options, "reverse", &handle).code(), + Status::Code::kNotSupported); + ASSERT_EQ(handle, nullptr); +} + +// kStandardOnly ignores the factory, so the trie's comparator requirement must +// not stop a DB from opening and working. +TEST_F(TrieIndexDBTest, StandardOnlyIgnoresTrieComparatorRequirement) { + options_.comparator = ReverseBytewiseComparator(); + ASSERT_OK(OpenDBImpl(/*block_size=*/0, + BlockBasedTableOptions::IndexMode::kStandardOnly)); + ASSERT_OK(db_->Put(WriteOptions(), "a", "1")); + ASSERT_OK(db_->Put(WriteOptions(), "b", "2")); + ASSERT_OK(db_->Flush(FlushOptions())); + std::string value; + ASSERT_OK(db_->Get(ReadOptions(), "a", &value)); + ASSERT_EQ(value, "1"); + std::unique_ptr it(db_->NewIterator(ReadOptions())); + it->SeekToFirst(); + ASSERT_TRUE(it->Valid()); + ASSERT_EQ(it->key(), "b"); + ASSERT_OK(it->status()); +} + // Mixed SSTs: some written with UDI, some without. Both should be readable // through both index paths. TEST_P(TrieIndexDBTest, MixedSSTsWithAndWithoutUDI) { @@ -4713,6 +4772,144 @@ TEST_F(TrieIndexDBTest, CustomOnlySSTsSurviveLoweredIndexMode) { ASSERT_TRUE(iter->status().IsInvalidArgument()); } +// The trie reads its block in place. With mmap reads, a block that is not in +// the block cache points into the table file's mapping, which lives as long as +// the table reader; a cached block is a heap copy. Table readers are evicted +// and reopened here, with and without the index in the block cache. +TEST_P(TrieIndexDBTest, MmapReads) { + for (bool cache_index : {false, true}) { + SCOPED_TRACE(cache_index); + options_.allow_mmap_reads = true; + // The smallest table cache, 10 readers, for 16 SSTs. + options_.max_open_files = 20; + options_.disable_auto_compactions = true; + options_.create_if_missing = true; + BlockBasedTableOptions table_options; + table_options.user_defined_index_factory = trie_factory_; + table_options.index_mode = GetIndexMode(); + table_options.block_size = 256; + table_options.cache_index_and_filter_blocks = cache_index; + options_.table_factory.reset(NewBlockBasedTableFactory(table_options)); + last_options_ = options_; + ASSERT_OK(DB::Open(options_, dbname_, &db_)); + + std::vector keys; + for (int sst = 0; sst < 16; sst++) { + for (int i = sst; i < 800; i += 16) { + char key[16]; + snprintf(key, sizeof(key), "key_%04d", i); + ASSERT_OK(db_->Put(WriteOptions(), key, std::string(40, 'a' + sst))); + } + ASSERT_OK(db_->Flush(FlushOptions())); + } + for (int i = 0; i < 800; i++) { + char key[16]; + snprintf(key, sizeof(key), "key_%04d", i); + keys.emplace_back(key); + } + for (int round = 0; round < 3; round++) { + SCOPED_TRACE(round); + if (round == 1) { + ASSERT_OK(db_->CompactRange(CompactRangeOptions(), nullptr, nullptr)); + } else if (round == 2) { + ASSERT_OK(db_->Close()); + db_.reset(); + ASSERT_OK(DB::Open(options_, dbname_, &db_)); + } + std::string value; + for (size_t i = 0; i < keys.size(); i++) { + ASSERT_OK(db_->Get(TrieIndexReadOptions(), keys[i], &value)); + ASSERT_EQ(value, std::string(40, static_cast('a' + i % 16))); + } + VerifyScanBothIndexes(keys); + } + ASSERT_OK(db_->Close()); + db_.reset(); + ASSERT_OK(DestroyDB(dbname_, options_)); + } +} + +// A trie block of another format version, such as one written by an earlier +// release, cannot be read. kStandardDefault falls back to the standard index; +// modes that need the trie fail with NotSupported. +TEST_F(TrieIndexDBTest, OtherTrieVersionFallsBackOrFails) { + // Writes the trie's blocks with version 1 in the header. + class OldVersionFactory : public TrieIndexFactory { + public: + Status NewBuilder( + const IndexFactoryOptions& options, + std::unique_ptr& builder) const override { + std::unique_ptr trie; + Status s = TrieIndexFactory::NewBuilder(options, trie); + if (s.ok()) { + builder = std::make_unique(std::move(trie)); + } + return s; + } + + private: + class Builder : public IndexFactoryBuilder { + public: + explicit Builder(std::unique_ptr trie) + : trie_(std::move(trie)) {} + Slice AddIndexEntry(const Slice& last_key_in_current_block, + const Slice* first_key_in_next_block, + const BlockHandle& block_handle, + std::string* separator_scratch, + const IndexEntryContext& context) override { + return trie_->AddIndexEntry(last_key_in_current_block, + first_key_in_next_block, block_handle, + separator_scratch, context); + } + void OnKeyAdded(const Slice& key, ValueType type, + const Slice& value) override { + trie_->OnKeyAdded(key, type, value); + } + Status Finish(Slice* index_contents) override { + Status s = trie_->Finish(index_contents); + if (s.ok()) { + data_ = index_contents->ToString(); + EncodeFixed32(&data_[4], 1); + *index_contents = data_; + } + return s; + } + uint64_t EstimatedSize() const override { return trie_->EstimatedSize(); } + + private: + std::unique_ptr trie_; + std::string data_; + }; + }; + + using IndexMode = BlockBasedTableOptions::IndexMode; + trie_factory_ = std::make_shared(); + ASSERT_OK(OpenDBImpl(/*block_size=*/128, IndexMode::kStandardDefault)); + WriteSequentialKeys(0, 30); + ASSERT_OK(db_->Flush(FlushOptions())); + ASSERT_OK(db_->Close()); + db_.reset(); + + trie_factory_ = std::make_shared(); + options_.statistics = CreateDBStatistics(); + ASSERT_OK(OpenDBImpl(/*block_size=*/128, IndexMode::kStandardDefault)); + std::string value; + ASSERT_OK(db_->Get(ReadOptions(), "key_0015", &value)); + ASSERT_EQ(value, "val_0015"); + ASSERT_EQ(options_.statistics->getTickerCount( + SST_USER_DEFINED_INDEX_LOAD_FAIL_COUNT), + 1u); + ASSERT_OK(db_->Close()); + db_.reset(); + + for (IndexMode mode : + {IndexMode::kStandardRequired, IndexMode::kCustomDefault}) { + SCOPED_TRACE(static_cast(mode)); + ASSERT_EQ(OpenDBImpl(/*block_size=*/128, mode).code(), + Status::Code::kNotSupported); + } +} + TEST_P(TrieIndexDBTest, RollbackFromPrimaryWithoutCompactSucceeds) { // Verifies that removing UDI from primary-mode SSTs WITHOUT compacting // first still works. The standard index is always fully populated (even diff --git a/utilities/trie_index/trie_index_factory.cc b/utilities/trie_index/trie_index_factory.cc index a1e37cc5e690..637efdb097aa 100644 --- a/utilities/trie_index/trie_index_factory.cc +++ b/utilities/trie_index/trie_index_factory.cc @@ -6,18 +6,48 @@ #include "utilities/trie_index/trie_index_factory.h" #include -#include -#include +#include #include "db/dbformat.h" #include "rocksdb/comparator.h" #include "rocksdb/utilities/object_registry.h" #include "util/cast_util.h" -#include "util/coding.h" namespace ROCKSDB_NAMESPACE { namespace trie_index { +namespace { +// Separators per bucket. Larger buckets make the index smaller and Seek +// slower; 8 was the best tradeoff on the key sets measured. The value is +// stored in each index, so it can change without a format change. +constexpr uint32_t kTrieBucketSize = 8; + +// The trie orders separators byte by byte, so it requires a bytewise +// comparator. Null means the default, BytewiseComparator. +Status RequireBytewiseComparator(const Comparator* comparator) { + if (comparator != nullptr && comparator != BytewiseComparator()) { + return Status::NotSupported( + "TrieIndexFactory requires BytewiseComparator; got: ", + comparator->Name()); + } + return Status::OK(); +} + +// A block's tag is 0 when the block can be found by user key alone, and the +// real tag of its last key otherwise: +// - when it shares its separator with the previous block, the tags tell the +// blocks of that user key apart; +// - for the last block, whose separator is the last user key itself, the tag +// lets a Seek for that user key with a smaller sequence number move past it, +// as the standard index does with its full internal key. +// A separator between different user keys needs no tag: the standard index +// with index_key_is_user_key also compares user keys only there. +uint64_t BlockTag(bool shares_previous_separator, bool has_next_block, + uint64_t last_key_tag) { + return shares_previous_separator || !has_next_block ? last_key_tag : 0; +} +} // namespace + int RegisterBuiltinTrieIndexFactory(ObjectLibrary& library, const std::string& /*arg*/) { library.AddFactory( @@ -36,7 +66,9 @@ int RegisterBuiltinTrieIndexFactory(ObjectLibrary& library, // ============================================================================ TrieIndexBuilder::TrieIndexBuilder(const Comparator* comparator) - : comparator_(comparator), finished_(false) {} + : comparator_(comparator), + finished_(false), + trie_builder_(kTrieBucketSize) {} Slice TrieIndexBuilder::AddIndexEntry(const Slice& last_key_in_current_block, const Slice* first_key_in_next_block, @@ -52,36 +84,19 @@ Slice TrieIndexBuilder::AddIndexEntry(const Slice& last_key_in_current_block, // If first_key_in_next_block is nullptr, this is the last block -- use a // short successor of the last key. Slice separator; - // True when last_key and first_key_in_next_block are the same user key - // (same-user-key block boundary). Computed once and reused below for - // both the sticky flag and the per-entry seqno decision. + // True when last_key and first_key_in_next_block are the same user key: the + // separator is then that user key for both blocks, so this block needs its + // real tag. bool same_user_key = false; if (first_key_in_next_block != nullptr) { same_user_key = comparator_->Compare(last_key_in_current_block, *first_key_in_next_block) == 0; - *separator_scratch = last_key_in_current_block.ToString(); + separator_scratch->assign(last_key_in_current_block.data(), + last_key_in_current_block.size()); comparator_->FindShortestSeparator(separator_scratch, *first_key_in_next_block); separator = Slice(*separator_scratch); - - // Detect same-user-key block boundary: if the two user keys are - // identical, FindShortestSeparator returns the same key for both sides, - // making it impossible to distinguish the two blocks. Such an entry - // carries a real tag rather than the 0 sentinel. - - // Edge case: FindShortestSeparator may fail to shorten the key even when - // the user keys are different. Example: FindShortestSeparator("abc","abd") - // returns "abc" unchanged because incrementing 'c' would yield "abd" which - // is not < limit. When the resulting separator matches the previous entry's - // separator, the blocks will be grouped into the same run in Finish(). - // We must mark this as a same-user-key boundary so it gets a real seqno - // rather than kMaxSequenceNumber (which would trigger the overflow block - // assertion in Finish()). - if (!same_user_key && !buffered_entries_.empty() && - buffered_entries_.back().separator_key == *separator_scratch) { - same_user_key = true; - } } else { // Last block: use the last key itself as the separator, NOT a shortened // successor. This matches the standard ShortenedIndexBuilder behavior @@ -100,56 +115,29 @@ Slice TrieIndexBuilder::AddIndexEntry(const Slice& last_key_in_current_block, // desynchronization -- the trie index returns a valid block while the // standard index correctly reports no match. separator = last_key_in_current_block; - - // Edge case: if this last block's separator matches the previous entry's - // separator, they share the same user key (same-user-key run boundary). - if (!buffered_entries_.empty() && - comparator_->Compare(buffered_entries_.back().separator_key, - separator) == 0) { - same_user_key = true; - } - } - - // Buffer the entry for deferred trie construction in Finish(). - // We buffer rather than adding to the trie immediately because the - // all-or-nothing seqno encoding decision is made at Finish() time. - TrieBlockHandle handle; - handle.offset = block_handle.offset; - handle.size = block_handle.size; - - BufferedEntry entry; - entry.separator_key = separator.ToString(); - if (same_user_key) { - // Same-user-key boundary: store the real tag for correct block - // selection within the overflow run. - entry.tag = last_key_tag; - } else if (first_key_in_next_block == nullptr) { - // Last block: store the real tag. The standard index stores the full - // internal key (user key + seqno) as the last block's separator. The - // real tag ensures the post-seek correction correctly handles seeks - // where the target user key matches but the seqno differs. - entry.tag = last_key_tag; - } else { - // Non-boundary separator between blocks with different user keys. - // Store 0 (sentinel meaning "no seqno correction needed"). When the - // standard index has index_key_is_user_key=true, it compares user keys - // only and always stays on equal user keys. The trie matches this by - // ensuring target_tag < 0 is always false. - entry.tag = 0; } - entry.handle = handle; - has_unsupported_handle_ |= - handle.offset > UINT32_MAX || handle.size > UINT32_MAX; - const Slice previous_separator = - buffered_entries_.empty() ? Slice() - : Slice(buffered_entries_.back().separator_key); - UpdateSizeEstimate(entry.separator_key, previous_separator); - buffered_entries_.push_back(std::move(entry)); - + const uint64_t tag = + BlockTag(same_user_key || SharesPreviousSeparator(separator), + first_key_in_next_block != nullptr, last_key_tag); + const TrieBlockHandle handle{block_handle.offset, block_handle.size}; + TrackHandle(handle); + // Before AddBlock(), which may move the bytes LastSeparator() points to. + UpdateSizeEstimate(separator, trie_builder_.LastSeparator(), tag != 0); + trie_builder_.AddBlock(separator, handle, tag); return separator; } +bool TrieIndexBuilder::SharesPreviousSeparator(const Slice& separator) const { + // FindShortestSeparator() can leave a separator unshortened, e.g. "abc" + // before "abd", so it can equal the previous block's separator even when + // this block's last key and the next block's first key differ. The last + // block's separator is its last user key, which equals the previous + // separator when one user key spans both blocks. + return trie_builder_.NumBlocks() > 0 && + trie_builder_.LastSeparator() == separator; +} + void TrieIndexBuilder::OnKeyAdded(const Slice& /*key*/, ValueType /*type*/, const Slice& /*value*/) { // No-op: the trie is built from separator keys in AddIndexEntry(), not @@ -161,110 +149,81 @@ Status TrieIndexBuilder::Finish(Slice* index_contents) { return Status::InvalidArgument("TrieIndexBuilder::Finish called twice"); } finished_ = true; - if (has_unsupported_handle_) { - return Status::NotSupported( - "Trie index block offset or size exceeds the 32-bit encoding limit"); - } if (has_unprepared_entry_) { return Status::InvalidArgument( "TrieIndexBuilder::FinishAddEntry called without a matching " "PrepareAddEntry"); } - - // Every staged entry needs the seqno side-table so the post-seek - // correction can disambiguate the last block and any same-user-key run, - // so this is on whenever the trie is non-empty. It costs 8 bytes per leaf. - bool use_seqno = !buffered_entries_.empty(); - trie_builder_.SetHasSeqnoEncoding(use_seqno); - - if (use_seqno) { - // Feed de-duplicated separators to the trie with seqno side-table metadata. - // Consecutive identical separators form a "run" -- only the first - // occurrence goes into the trie (as the primary block). The remaining - // blocks in the run are stored as overflow blocks in the side-table. - // - // For non-boundary separators (different user keys), the tag is 0 - // (sentinel meaning "no seqno correction needed"), matching the standard - // index's user-key-only comparison mode. For the last block, the real - // tag is stored to match the standard index's full internal key behavior. - size_t i = 0; - while (i < buffered_entries_.size()) { - const auto& entry = buffered_entries_[i]; - - // Count how many consecutive entries share this separator key. - size_t run_start = i; - size_t run_end = i + 1; - while (run_end < buffered_entries_.size() && - buffered_entries_[run_end].separator_key == entry.separator_key) { - run_end++; - } - uint32_t block_count = static_cast(run_end - run_start); - - // Non-boundary entries have tag=0 (sentinel meaning "no seqno - // correction needed"). Same-user-key boundary and last-block entries - // have real tags. The trie builder stores these directly in the seqno - // side-table. - // - // For boundary separators (same user key), store the actual packed - // tag for correct seqno-based block selection. - uint64_t tag = entry.tag; - - // Add the primary (first) block for this separator. - trie_builder_.AddKeyWithSeqno(Slice(entry.separator_key), entry.handle, - tag, block_count); - - // Add overflow blocks (2nd, 3rd, ... in the run). - // Overflow blocks only exist within same-key runs, so their packed - // tags come from last_key_tag in AddIndexEntry. - // The tag may be 0 when bottommost compaction zeroes all sequence - // numbers -- this is valid; see AddOverflowBlock comment. - for (size_t j = run_start + 1; j < run_end; j++) { - assert(buffered_entries_[j].tag != - PackSequenceAndType(kMaxSequenceNumber, kValueTypeForSeek)); - trie_builder_.AddOverflowBlock(buffered_entries_[j].handle, - buffered_entries_[j].tag); - } - - i = run_end; - } + // The tags have the meaning TrieBuilder::AddBlock() expects: 0 for a block + // found by user key alone, otherwise the tag of the block's last key. + Status s = trie_builder_.Finish(&data_); + if (s.ok()) { + *index_contents = data_; } - - // Release buffered entries -- no longer needed after feeding to the trie. - buffered_entries_.clear(); - buffered_entries_.shrink_to_fit(); - - // Always finish the trie builder, even with 0 keys -- this produces a valid - // serialized trie that can be parsed by NewReader. Without this, an empty - // Slice would be returned, causing InitFromData to fail with "data too short - // for header". - trie_builder_.Finish(); - *index_contents = trie_builder_.GetSerializedData(); - return Status::OK(); + return s; } -// ============================================================================ -// TrieIndexIterator -// ============================================================================ - uint64_t TrieIndexBuilder::EstimatedSize() const { if (finished_) { - return trie_builder_.GetSerializedData().size(); + return data_.size(); } - // Dense/sparse labels, bitvectors, child positions, and capped chain metadata - // use at most 12 bytes per distinct trie edge. Each handle and seqno record - // needs at most 20 bytes. Reserve 1 KiB for headers and alignment. Shared - // prefixes contribute only once, just as in LoudsTrieBuilder::Finish(). - return 1024 + estimated_trie_edges_ * 12 + estimated_num_entries_ * 20; + TrieBuilder::SizeStats stats = estimate_; + stats.max_block_step = max_block_step_.LoadRelaxed(); + stats.max_block_size = max_block_size_.LoadRelaxed(); + stats.irregular_gaps = irregular_block_gaps_.LoadRelaxed(); + return TrieBuilder::SizeBound(stats, kTrieBucketSize); } size_t TrieIndexBuilder::UpdateSizeEstimate(const Slice& separator, - const Slice& previous_separator) { + const Slice& previous_separator, + bool may_have_tag) { const size_t shared_prefix = previous_separator.difference_offset(separator); - estimated_trie_edges_ += separator.size() - shared_prefix; - ++estimated_num_entries_; + const size_t drop = previous_separator.size() - shared_prefix; + const size_t add = separator.size() - shared_prefix; + // Matches the entry encoding in trie.h. The first separator and the bucket + // heads are charged too, which only overestimates. + estimate_.suffix_bytes += add; + estimate_.entry_bytes += TrieBuilder::EntrySize(drop, add); + estimate_.max_key_len = + std::max(estimate_.max_key_len, separator.size()); + if (estimate_.num_blocks > 0 && drop == 0 && add == 0) { + ++estimate_.duplicates; + may_have_tag = true; + } + if (may_have_tag) { + ++estimate_.tagged; + } + ++estimate_.num_blocks; return shared_prefix; } +void TrieIndexBuilder::TrackHandle(const TrieBlockHandle& handle) { + max_block_size_.StoreRelaxed( + std::max(max_block_size_.LoadRelaxed(), handle.size)); + // Mirrors the derived-size rule in TrieBuilder::Finish(): no block may end + // past 2^64 - 1, the gaps between blocks must be equal, and the sentinel + // after the last block, its end plus the gap, must fit too. Wrapping + // arithmetic only matters once sizes are already explicit. + const uint64_t max = std::numeric_limits::max(); + bool irregular = handle.size > max - handle.offset; + if (trie_builder_.NumBlocks() > 0) { + const uint64_t step = handle.offset - previous_handle_.offset; + max_block_step_.StoreRelaxed(std::max(max_block_step_.LoadRelaxed(), step)); + const uint64_t previous_end = + previous_handle_.offset + previous_handle_.size; + const uint64_t gap = handle.offset - previous_end; + irregular = irregular || handle.offset < previous_end || + (trie_builder_.NumBlocks() > 1 && gap != block_gap_); + block_gap_ = gap; + } + // Any block may be the last. If this one is not, a sentinel that does not + // fit means the next block's gap differs, so flagging it now is exact. + if (irregular || block_gap_ > max - handle.offset - handle.size) { + irregular_block_gaps_.StoreRelaxed(true); + } + previous_handle_ = handle; +} + // --------------------------------------------------------------------------- // Parallel compression protocol // --------------------------------------------------------------------------- @@ -292,17 +251,22 @@ void TrieIndexBuilder::PrepareAddEntry(const Slice& last_key_in_current_block, comparator_->Compare(last_key_in_current_block, *first_key_in_next_block) == 0; - p->separator_key = last_key_in_current_block.ToString(); + p->separator_key.assign(last_key_in_current_block.data(), + last_key_in_current_block.size()); comparator_->FindShortestSeparator(&p->separator_key, *first_key_in_next_block); } else { // Last block: store the last key itself (no shortening). - p->separator_key = last_key_in_current_block.ToString(); + p->separator_key.assign(last_key_in_current_block.data(), + last_key_in_current_block.size()); } p->valid = true; + // A separator equal to the previous one gets a tag in FinishAddEntry, which + // UpdateSizeEstimate detects. const size_t shared_prefix = - UpdateSizeEstimate(p->separator_key, previous_prepared_separator_); + UpdateSizeEstimate(p->separator_key, previous_prepared_separator_, + p->same_user_key_initial || !p->has_next_block); // The writer can move the prepared key as soon as this callback returns. // Keep one emit-owned copy, reusing the prefix bytes already stored here. previous_prepared_separator_.replace(shared_prefix, std::string::npos, @@ -323,66 +287,29 @@ void TrieIndexBuilder::FinishAddEntry(const BlockHandle& block_handle, return; } - // Recheck the buffer-state-dependent branch of the same-user-key - // detection. Now that we run on the BG writer thread in commit order, - // buffered_entries_.back() is the immediately preceding block's entry. - // - // The branch structure mirrors the serial AddIndexEntry exactly, including - // which comparison primitive each branch uses, so the two paths cannot - // diverge on a comparator where byte equality and Compare() == 0 disagree. - bool same_user_key = p->same_user_key_initial; - if (p->has_next_block) { - // Non-last block: a separator FindShortestSeparator could not shorten can - // still collide with the previous entry's separator. - if (!same_user_key && !buffered_entries_.empty() && - buffered_entries_.back().separator_key == p->separator_key) { - same_user_key = true; - } - } else if (!buffered_entries_.empty() && - comparator_->Compare(buffered_entries_.back().separator_key, - p->separator_key) == 0) { - // Last block: its separator matching the previous entry's means they - // share a user key. - same_user_key = true; - } - - BufferedEntry be; - be.separator_key = std::move(p->separator_key); - if (same_user_key) { - // Same-user-key run: real tag distinguishes blocks within the run. - be.tag = p->last_key_tag; - } else if (!p->has_next_block) { - // Last block: real tag covers any post-seek correction by seqno. - be.tag = p->last_key_tag; - } else { - // Distinct-user-key separator: 0 sentinel -- no seqno correction - // needed because the separator already disambiguates blocks. - be.tag = 0; - } - be.handle.offset = block_handle.offset; - be.handle.size = block_handle.size; - has_unsupported_handle_ |= - be.handle.offset > UINT32_MAX || be.handle.size > UINT32_MAX; - - buffered_entries_.push_back(std::move(be)); + // The half of the tag decision that depends on the previous block runs here: + // on the BG writer thread, in commit order, trie_builder_ holds every block + // before this one. It is the same code as the serial AddIndexEntry(). + const Slice separator(p->separator_key); + const uint64_t tag = + BlockTag(p->same_user_key_initial || SharesPreviousSeparator(separator), + p->has_next_block, p->last_key_tag); + const TrieBlockHandle handle{block_handle.offset, block_handle.size}; + TrackHandle(handle); + trie_builder_.AddBlock(separator, handle, tag); // Reset valid so the same prepared slot can be reused (the table // builder rotates BlockReps in a ring buffer). p->valid = false; } -TrieIndexIterator::TrieIndexIterator(const LoudsTrie* trie, - const Comparator* comparator, - bool has_seqno_encoding) - : comparator_(comparator), - iter_(trie), - trie_(trie), - current_scan_idx_(0), - prepared_(false), - has_seqno_encoding_(has_seqno_encoding), - overflow_run_index_(0), - overflow_run_size_(1), - overflow_base_idx_(0) {} +// ============================================================================ +// TrieIndexIterator +// ============================================================================ + +TrieIndexIterator::TrieIndexIterator(const Trie* trie, + const Comparator* comparator) + : comparator_(comparator), iter_(trie) {} void TrieIndexIterator::Prepare(const ScanOptions scan_opts[], size_t num_opts) { @@ -395,71 +322,24 @@ void TrieIndexIterator::Prepare(const ScanOptions scan_opts[], prepared_ = true; } -Status TrieIndexIterator::SeekToFirstAndGetResult(IterateResult* result) { - ResetOverflowState(); - - if (!iter_.SeekToFirst()) { - result->bound_check_result = IterBoundCheck::kUnknown; +Status TrieIndexIterator::SetResult(const Status& s, IterBoundCheck bounds, + IterateResult* result) const { + if (s.ok() && iter_.Valid()) { + result->key = iter_.key(); + result->bound_check_result = bounds; + } else { + // Running off the end of this SST says nothing about the upper bound: + // the next SST on the level may still hold in-bound keys, and + // kOutOfBound would make LevelIterator stop scanning the level. result->key = Slice(); - return Status::OK(); - } - - CopyTrieKeyToResult(result); - SetupOverflowForCurrentLeaf(/*position_at_last=*/false); - result->bound_check_result = IterBoundCheck::kInbound; - return Status::OK(); -} - -Status TrieIndexIterator::SeekToLastAndGetResult(IterateResult* result) { - ResetOverflowState(); - - if (!iter_.SeekToLast()) { result->bound_check_result = IterBoundCheck::kUnknown; - result->key = Slice(); - return Status::OK(); } - - CopyTrieKeyToResult(result); - SetupOverflowForCurrentLeaf(/*position_at_last=*/true); - result->bound_check_result = IterBoundCheck::kInbound; - return Status::OK(); -} - -Status TrieIndexIterator::PrevAndGetResult(IterateResult* result) { - // Upper-bound checks are based on forward movement. While moving backward, - // DBIter's lower-bound handling decides when to stop. - // Overflow fast path: key doesn't change within the same run. - if (overflow_run_index_ > 0) { - overflow_run_index_--; - result->key = Slice(current_key_scratch_); - result->bound_check_result = IterBoundCheck::kInbound; - return Status::OK(); - } - - // Non-overflow: current_key_scratch_ is about to be overwritten, so swap - // into prev_key_scratch_ in O(1) instead of copying. - std::swap(prev_key_scratch_, current_key_scratch_); - - // Move to the previous trie leaf. - ResetOverflowState(); - - if (!iter_.Prev()) { - result->bound_check_result = IterBoundCheck::kUnknown; - result->key = Slice(); - return Status::OK(); - } - - CopyTrieKeyToResult(result); - SetupOverflowForCurrentLeaf(/*position_at_last=*/true); - result->bound_check_result = IterBoundCheck::kInbound; - return Status::OK(); + return s; } Status TrieIndexIterator::SeekAndGetResult(const Slice& target, IterateResult* result, const SeekContext& context) { - uint64_t target_packed = context.target_tag; - // Advance current_scan_idx_ past any scans whose limit <= target. // This handles the multi-scan case where the caller seeks into a later // scan range after the previous scan returned kOutOfBound. @@ -474,160 +354,38 @@ Status TrieIndexIterator::SeekAndGetResult(const Slice& target, } } } - - ResetOverflowState(); - - // Always seek with user key only -- the trie stores user-key separators. - // When seqno encoding is active, post-seek correction handles the seqno. - if (!iter_.Seek(target)) { - // No leaf has a key >= target: the target is past all blocks in this SST. - // Return kUnknown (not kOutOfBound) because exhausting this SST's trie - // says nothing about the upper bound -- the next SST on the level may - // still contain in-bound keys. kOutOfBound would cause LevelIterator to - // stop scanning the level prematurely. - result->bound_check_result = IterBoundCheck::kUnknown; - result->key = Slice(); - return Status::OK(); - } - - // Set the result key (always a user key, no suffix stripping needed). - // Reuse current_key_scratch_ capacity -- avoids heap allocation after warmup. - { - Slice trie_key = iter_.Key(); - current_key_scratch_.assign(trie_key.data(), trie_key.size()); - result->key = Slice(current_key_scratch_); - } - - // ---- Post-seek correction for seqno side-table ---- - // - // When has_seqno_encoding_ is true, the leaf we landed on has a seqno - // side-table entry. We use it to determine if this is the right block - // for the given (target, target_packed). - // - // For same-user-key boundaries: the tag determines which block - // within a run of same-key blocks is correct. If target_packed < leaf_packed, - // advance through overflow blocks. - // - // For non-boundary separators: leaf_seqno is 0. The comparison - // target_packed < 0 is always false, so no advancement occurs. This matches - // the standard index's index_key_is_user_key=true mode where equal user - // keys always match without seqno comparison. - // - // For the last block: leaf_seqno stores the real tag of the last key. - // This matches the standard index which stores the full internal key - // as the last block's separator. - if (has_seqno_encoding_ && iter_.Valid()) { - uint64_t leaf_idx = iter_.LeafIndex(); - uint64_t leaf_seqno = trie_->GetLeafSeqno(leaf_idx); - - if (target_packed < leaf_seqno && - comparator_->Compare(Slice(current_key_scratch_), target) == 0) { - // Target's internal key is AFTER the separator (lower tag = - // later in internal key order for same user key). Advance through - // overflow blocks. - uint32_t block_count = trie_->GetLeafBlockCount(leaf_idx); - uint32_t base = trie_->GetOverflowBase(leaf_idx); - - bool found = false; - for (uint32_t oi = 0; oi < block_count - 1; oi++) { - uint64_t ov_seqno = trie_->GetOverflowSeqno(base + oi); - if (ov_seqno == 0 || target_packed >= ov_seqno) { - // This overflow block is the right one. - overflow_run_index_ = oi + 1; // 1-based (0 = primary) - overflow_run_size_ = block_count; - overflow_base_idx_ = base; - found = true; - break; - } - } - - if (!found) { - // target_packed is below all tags in this run. Advance to - // the next trie leaf (the block after the run). - if (!iter_.Next()) { - // Exhausted all blocks: target is past the end of this SST. - // Return kUnknown -- see comment in Seek path above. - result->bound_check_result = IterBoundCheck::kUnknown; - result->key = Slice(); - return Status::OK(); - } - // Update key and overflow state for the new leaf. - { - Slice trie_key = iter_.Key(); - current_key_scratch_.assign(trie_key.data(), trie_key.size()); - result->key = Slice(current_key_scratch_); - } - overflow_run_index_ = 0; - overflow_run_size_ = 1; - overflow_base_idx_ = 0; - // Check if the new leaf also has overflow (unlikely but possible - // with adjacent same-key runs for different user keys). - // iter_.Valid() is guaranteed here -- Next() returned true above. - if (has_seqno_encoding_) { - uint64_t new_leaf = iter_.LeafIndex(); - overflow_run_size_ = trie_->GetLeafBlockCount(new_leaf); - overflow_base_idx_ = trie_->GetOverflowBase(new_leaf); - } - } - } else { - // Right block (common path). Set overflow state in case this leaf - // has a run (for subsequent Next() calls). - uint32_t block_count = trie_->GetLeafBlockCount(leaf_idx); - overflow_run_index_ = 0; - overflow_run_size_ = block_count; - overflow_base_idx_ = trie_->GetOverflowBase(leaf_idx); - } - } - - result->bound_check_result = CheckBounds(target); - return Status::OK(); + Status s = iter_.Seek(target, context.target_tag); + return SetResult(s, CheckBounds(target), result); } Status TrieIndexIterator::NextAndGetResult(IterateResult* result) { - // Overflow fast path: key doesn't change within the same run, so - // current_key_scratch_ can be passed directly to CheckBounds (no copy). - if (overflow_run_index_ + 1 < overflow_run_size_) { - overflow_run_index_++; - result->key = Slice(current_key_scratch_); - result->bound_check_result = CheckBounds(Slice(current_key_scratch_)); - return Status::OK(); - } - - // Non-overflow: current_key_scratch_ is about to be overwritten, so swap - // into prev_key_scratch_ in O(1) instead of copying. - std::swap(prev_key_scratch_, current_key_scratch_); + // The bound check uses the separator before the step, so run it first. + const IterBoundCheck bounds = CheckBounds(iter_.key()); + return SetResult(iter_.Next(), bounds, result); +} - // Advance to the next trie leaf. - ResetOverflowState(); +Status TrieIndexIterator::SeekToFirstAndGetResult(IterateResult* result) { + return SetResult(iter_.SeekToFirst(), IterBoundCheck::kInbound, result); +} - if (!iter_.Next()) { - result->bound_check_result = IterBoundCheck::kUnknown; - result->key = Slice(); - return Status::OK(); - } +Status TrieIndexIterator::SeekToLastAndGetResult(IterateResult* result) { + return SetResult(iter_.SeekToLast(), IterBoundCheck::kInbound, result); +} - CopyTrieKeyToResult(result); - SetupOverflowForCurrentLeaf(/*position_at_last=*/false); - result->bound_check_result = CheckBounds(Slice(prev_key_scratch_)); - return Status::OK(); +Status TrieIndexIterator::PrevAndGetResult(IterateResult* result) { + // Upper-bound checks are based on forward movement. While moving backward, + // DBIter's lower-bound handling decides when to stop. + return SetResult(iter_.Prev(), IterBoundCheck::kInbound, result); } IndexFactoryBuilder::BlockHandle TrieIndexIterator::value() { - if (overflow_run_index_ == 0) { - // Primary block -- use the trie leaf's handle. - auto handle = iter_.Value(); - return IndexFactoryBuilder::BlockHandle{handle.offset, handle.size}; - } - // Overflow block -- use the side-table handle. - // overflow_run_index_ is 1-based, overflow array is 0-based. - uint32_t overflow_idx = overflow_base_idx_ + overflow_run_index_ - 1; - auto handle = trie_->GetOverflowHandle(overflow_idx); + const TrieBlockHandle handle = iter_.value(); return IndexFactoryBuilder::BlockHandle{handle.offset, handle.size}; } IterBoundCheck TrieIndexIterator::CheckBounds( const Slice& reference_key) const { - if (!prepared_ || scan_opts_.empty()) { + if (!HasBounds()) { // No bounds to check -- always in-bound. return IterBoundCheck::kInbound; } @@ -668,42 +426,39 @@ TrieIndexReader::TrieIndexReader(const Comparator* comparator) : comparator_(comparator) {} Status TrieIndexReader::InitFromSlice(const Slice& data) { - return trie_.InitFromData(data); + return trie_.Init(data); } std::unique_ptr TrieIndexReader::NewIterator( const ReadOptions& /*read_options*/) { - return std::make_unique(&trie_, comparator_, - trie_.HasSeqnoEncoding()); + return std::make_unique(&trie_, comparator_); } -size_t TrieIndexReader::ApproximateMemoryUsage() const { - // The serialized data is owned by the table reader or block cache. Only - // report the child-position lookup tables allocated by InitFromData(). - return trie_.ApproximateAuxMemoryUsage(); -} +size_t TrieIndexReader::ApproximateMemoryUsage() const { return sizeof(*this); } // ============================================================================ // TrieIndexFactory // ============================================================================ +Status TrieIndexFactory::ValidateOptions( + const DBOptions& db_opts, const ColumnFamilyOptions& cf_opts) const { + // NewBuilder() and NewReader() reject other comparators too; failing here + // stops DB::Open instead of every flush and compaction. + Status s = RequireBytewiseComparator(cf_opts.comparator); + if (!s.ok()) { + return s; + } + return IndexFactory::ValidateOptions(db_opts, cf_opts); +} + Status TrieIndexFactory::NewBuilder( const IndexFactoryOptions& option, std::unique_ptr& builder) const { - // The trie traverses keys byte-by-byte in lexicographic order, so it - // requires a bytewise comparator. Non-bytewise comparators (e.g., - // ReverseBytewiseComparator or custom comparators) would produce separator - // keys in a different order than the trie's byte-level traversal, causing - // incorrect Seek results. - if (option.comparator != nullptr && - option.comparator != BytewiseComparator()) { - return Status::NotSupported( - "TrieIndexFactory requires BytewiseComparator; got: ", - option.comparator->Name()); + Status s = RequireBytewiseComparator(option.comparator); + if (!s.ok()) { + return s; } - // Default to BytewiseComparator when null. The trie requires a bytewise - // comparator for separator key ordering; null would cause a dereference - // crash in AddIndexEntry when comparing keys. + // AddIndexEntry() compares keys, so a null comparator becomes the default. const Comparator* cmp = option.comparator ? option.comparator : BytewiseComparator(); builder = std::make_unique(cmp); @@ -713,14 +468,14 @@ Status TrieIndexFactory::NewBuilder( Status TrieIndexFactory::NewReader( const IndexFactoryOptions& option, Slice& index_block, std::unique_ptr& reader) const { + Status s = RequireBytewiseComparator(option.comparator); + if (!s.ok()) { + return s; + } const Comparator* cmp = option.comparator ? option.comparator : BytewiseComparator(); - if (cmp != BytewiseComparator()) { - return Status::NotSupported( - "TrieIndexFactory requires BytewiseComparator; got: ", cmp->Name()); - } auto trie_reader = std::make_unique(cmp); - Status s = trie_reader->InitFromSlice(index_block); + s = trie_reader->InitFromSlice(index_block); if (!s.ok()) { return s; } diff --git a/utilities/trie_index/trie_index_factory.h b/utilities/trie_index/trie_index_factory.h index e60210681786..412c94e59f7f 100644 --- a/utilities/trie_index/trie_index_factory.h +++ b/utilities/trie_index/trie_index_factory.h @@ -10,10 +10,8 @@ // Trie-based User Defined Index (UDI) for RocksDB's block-based tables. // // This provides a TrieIndexFactory that implements the IndexFactory -// interface, building a Fast Succinct Trie (FST) index from the separator keys -// generated during SST file construction. Based on the SuRF paper results, the -// trie is expected to achieve significant space reduction compared to the -// standard index while providing comparable Seek() performance. +// interface, building a trie index (see trie.h) from the separator keys +// generated during SST file construction. // // Usage: // auto trie_factory = std::make_shared(); @@ -29,13 +27,6 @@ // auto iter = db->NewIterator(ro); // // In kCustomDefault/kCustomOnly mode, all reads use the trie automatically. -// -// The trie format stores data-block offsets and sizes as 32-bit integers. -// If either exceeds UINT32_MAX, SST construction fails with NotSupported in -// every mode that builds the trie, including kStandardDefault. File-size -// targets do not guarantee this limit: L0 outputs and SstFileWriter -// outputs are not split at the target size. Bound those outputs explicitly -// or use kStandardOnly to stop building the trie. #pragma once @@ -47,7 +38,8 @@ #include "rocksdb/comparator.h" #include "rocksdb/index_factory.h" #include "rocksdb/types.h" -#include "utilities/trie_index/louds_trie.h" +#include "util/atomic.h" +#include "utilities/trie_index/trie.h" namespace ROCKSDB_NAMESPACE { class ObjectLibrary; @@ -59,28 +51,28 @@ int RegisterBuiltinTrieIndexFactory(ObjectLibrary& library, const std::string& arg); // ============================================================================ -// TrieIndexBuilder: Implements IndexFactoryBuilder using LoudsTrieBuilder. +// TrieIndexBuilder: Implements IndexFactoryBuilder using TrieBuilder. // // During SST file construction, RocksDB calls: // 1. OnKeyAdded() for each key-value pair. // 2. AddIndexEntry() at each data block boundary. // 3. Finish() to serialize the index. // -// The trie builder collects the separator keys from AddIndexEntry() and -// builds a LOUDS-encoded trie during Finish(). +// The separator of each data block goes straight into a TrieBuilder, which +// serializes the index in Finish(). // ============================================================================ class TrieIndexBuilder final : public IndexFactoryBuilder { public: explicit TrieIndexBuilder(const Comparator* comparator); ~TrieIndexBuilder() override = default; - // Called at each data block boundary. We compute a short separator - // between last_key_in_current_block and first_key_in_next_block, buffer - // it for deferred trie construction in Finish(), and return it. + // Called at each data block boundary. Computes a short separator between + // last_key_in_current_block and first_key_in_next_block, adds it to the + // trie with the block's handle, and returns it. // - // The sequence numbers from context are used to detect same-user-key - // block boundaries, which decide whether an entry carries a real tag or - // the 0 sentinel. Every non-empty trie emits the seqno side-table. + // A block gets the real tag of its last key when that user key continues + // into the next block, when its separator equals the previous block's, or + // when it is the last block; other blocks get the 0 tag. Slice AddIndexEntry(const Slice& last_key_in_current_block, const Slice* first_key_in_next_block, const BlockHandle& block_handle, @@ -92,24 +84,23 @@ class TrieIndexBuilder final : public IndexFactoryBuilder { void OnKeyAdded(const Slice& key, ValueType type, const Slice& value) override; - // Finalize the trie and return the serialized index data. Returns - // NotSupported if any data-block offset or size exceeds UINT32_MAX; - // the serialized trie encodes these fields as 32-bit integers. + // Finalize the trie and return the serialized index data. Status Finish(Slice* index_contents) override; - // Returns an estimate of the current serialized index size. The estimate is - // updated only by emit-thread callbacks, so concurrent FinishAddEntry calls - // do not require synchronization. + // Returns an estimate of the current serialized index size. Separator + // statistics are updated only by emit-thread callbacks and block statistics + // are relaxed atomics, so concurrent FinishAddEntry calls do not require + // synchronization. uint64_t EstimatedSize() const override; // --- Parallel compression protocol --- // // The trie builder participates in the parallel pipeline by splitting // AddIndexEntry into Prepare (emit thread, computes the separator) and - // Finish (write thread, attaches the now-known block handle and appends - // to buffered_entries_). Because FinishAddEntry runs serially on the BG - // worker thread (the table builder commits blocks in order), the entry - // buffer mutation needs no locking. + // Finish (write thread, attaches the now-known block handle and adds the + // block to trie_builder_). Because FinishAddEntry runs serially on the BG + // worker thread (the table builder commits blocks in order), trie_builder_ + // needs no locking. bool SupportsParallelAddEntry() const override { return true; } std::unique_ptr CreatePreparedAddEntry() override; void PrepareAddEntry(const Slice& last_key_in_current_block, @@ -122,35 +113,47 @@ class TrieIndexBuilder final : public IndexFactoryBuilder { private: const Comparator* comparator_; - LoudsTrieBuilder trie_builder_; + // Serialized index. + std::string data_; bool finished_; - bool has_unsupported_handle_ = false; // Set when FinishAddEntry receives an entry with nothing prepared in it, so // Finish() fails instead of writing an index that omits that data block. bool has_unprepared_entry_ = false; + // Whether the gaps between blocks differ, so that sizes are stored. One of + // the block statistics described below; kept here to avoid padding. + RelaxedAtomic irregular_block_gaps_{false}; - // Buffered separator entries: (separator_key, tag, handle). - // The separator_key is the user-key-only separator computed by - // FindShortestSeparator. The tag field stores: - // - For same-user-key boundaries: the real tag of last_key - // - For the last block: the real tag of last_key - // - For intermediate non-boundary entries: 0 (sentinel) - struct BufferedEntry { - std::string separator_key; - uint64_t tag{}; - TrieBlockHandle handle; - }; - std::vector buffered_entries_; + // Blocks added so far. Owned by the thread that adds blocks (the BG writer + // during parallel builds). + TrieBuilder trie_builder_; // EstimatedSize() is called on the emit thread. Keep its state disjoint from - // buffered_entries_, which the BG writer mutates during parallel builds. - // Adjacent separators charge their common prefix only once. - uint64_t estimated_trie_edges_ = 0; - uint64_t estimated_num_entries_ = 0; + // trie_builder_, which the BG writer mutates during parallel builds. + // Separator statistics of the entries prepared so far. Its block statistics + // are unused; the atomics below hold them. + TrieBuilder::SizeStats estimate_; std::string previous_prepared_separator_; - // Returns the shared prefix length for reusing emit-owned separator bytes. + // Adds a separator to estimate_ and returns its shared prefix length with + // the previous one, for reusing emit-owned separator bytes. may_have_tag is + // whether the entry may get a nonzero tag. size_t UpdateSizeEstimate(const Slice& separator, - const Slice& previous_separator); + const Slice& previous_separator, bool may_have_tag); + + // Statistics of the block handles added so far, with irregular_block_gaps_ + // above. Written by the thread that adds handles (the BG writer during + // parallel builds) and read by EstimatedSize() on the emit thread. They + // bound the widths of the offset and size columns. Blocks still being + // compressed in a parallel build are charged like the largest block so far, + // so the estimate can fall short until their handles arrive. + RelaxedAtomic max_block_step_{0}; + RelaxedAtomic max_block_size_{0}; + // Owned by the thread that adds handles. + TrieBlockHandle previous_handle_; + uint64_t block_gap_ = 0; + // Records the next block's handle in the statistics above. + void TrackHandle(const TrieBlockHandle& handle); + // Whether separator equals the separator of the block added last. + bool SharesPreviousSeparator(const Slice& separator) const; // Staged data for the parallel AddIndexEntry protocol. Populated by // PrepareAddEntry on the emit thread, consumed by FinishAddEntry on @@ -158,16 +161,16 @@ class TrieIndexBuilder final : public IndexFactoryBuilder { // runs and is filled in by Finish. struct PreparedTrieEntry : public IndexFactoryBuilder::PreparedAddEntry { std::string separator_key; - // Real tag of last_key_in_current_block. Used as entry.tag for last - // blocks and for same-user-key runs. + // Real tag of last_key_in_current_block, for last blocks and for blocks + // sharing a separator with the previous one. uint64_t last_key_tag = 0; // True if first_key_in_next_block was non-null at Prepare time. bool has_next_block = false; // True if last_key and first_key_in_next_block share the same user // key (computed from inputs alone). Finish may upgrade this to true // if FindShortestSeparator failed to shorten and the result matches - // the previous buffered entry's separator (a check that depends on - // buffer state and so must run during Finish). + // the previous block's separator, which only the thread adding blocks + // can check. bool same_user_key_initial = false; // Set by PrepareAddEntry and cleared by FinishAddEntry. Finishing an // entry while this is false fails the build. @@ -176,160 +179,77 @@ class TrieIndexBuilder final : public IndexFactoryBuilder { }; // ============================================================================ -// TrieIndexIterator: Implements IndexFactoryIterator using -// LoudsTrieIterator. +// TrieIndexIterator: Implements IndexFactoryIterator using TrieIterator. // -// Wraps LoudsTrieIterator and adapts it to the UDI iterator interface, -// handling bounds checking against ScanOptions. +// Wraps TrieIterator and adapts it to the UDI iterator interface, handling +// bounds checking against ScanOptions. // ============================================================================ class TrieIndexIterator final : public IndexFactoryIterator { public: - // @param has_seqno_encoding: true if the trie was built with a seqno - // side-table (enabling post-seek correction for same-user-key boundaries). - // This flag is read from the trie's serialized header. - TrieIndexIterator(const LoudsTrie* trie, const Comparator* comparator, - bool has_seqno_encoding); - ~TrieIndexIterator() override = default; + TrieIndexIterator(const Trie* trie, const Comparator* comparator); // Prepare for a batch of scans. Stores scan bounds for later use. void Prepare(const ScanOptions scan_opts[], size_t num_opts) override; - - // Position at the very first index entry. Descends directly to the - // leftmost leaf without a full seek traversal. - Status SeekToFirstAndGetResult(IterateResult* result) override; - - // Position at the very last index entry. Descends directly to the - // rightmost leaf without a full seek traversal. - Status SeekToLastAndGetResult(IterateResult* result) override; - - // Move to the previous index entry. When in an overflow run, decrements - // within the run before moving to the previous trie leaf. - Status PrevAndGetResult(IterateResult* result) override; - - // Seek to the first index entry >= target. When has_seqno_encoding_ is - // true, the trie is searched with user_key only, then post-seek correction - // uses target_tag from context to advance through overflow blocks. Status SeekAndGetResult(const Slice& target, IterateResult* result, const SeekContext& context) override; - - // Advance to the next index entry. When in an overflow run, advances - // within the run before moving to the next trie leaf. Status NextAndGetResult(IterateResult* result) override; - - // Return the BlockHandle of the current block. When positioned on an - // overflow block, returns the overflow block's handle instead of the - // trie leaf's handle. + Status SeekToFirstAndGetResult(IterateResult* result) override; + Status SeekToLastAndGetResult(IterateResult* result) override; + Status PrevAndGetResult(IterateResult* result) override; IndexFactoryBuilder::BlockHandle value() override; private: - // Check if the current block is within the active scan bounds. - // reference_key is the key to compare against the limit: for Seek this - // is the seek target, for Next this is the previous separator key. - // The trie stores separator keys (upper bounds on block contents), not - // first-in-block keys, so we cannot compare the current separator against - // the limit directly -- see the UDI API contract in user_defined_index.h. + // Whether CheckBounds() can return anything but kInbound. + bool HasBounds() const { return prepared_ && !scan_opts_.empty(); } + + // Checks whether the current block may hold keys in the current scan. + // reference_key is the key to compare against the limit: for Seek this is + // the seek target, for Next this is the previous separator key. Separators + // are upper bounds on block contents, not first-in-block keys, so the + // current separator cannot be compared against the limit directly -- see + // the UDI API contract in user_defined_index.h. IterBoundCheck CheckBounds(const Slice& reference_key) const; + // Fills *result from the iterator position after a call that returned s. + // bounds is the bound check result to report when positioned. + Status SetResult(const Status& s, IterBoundCheck bounds, + IterateResult* result) const; + const Comparator* comparator_; - LoudsTrieIterator iter_; - // Pointer to the trie for side-table access during overflow iteration. - const LoudsTrie* trie_; - // Scratch space for the current separator key (reconstructed from trie). - std::string current_key_scratch_; - // Previous separator key, used as reference for Next() bounds checking. - // Always valid when NextAndGetResult is called, since Next requires a - // preceding Seek or Next which sets current_key_scratch_ (copied here). - std::string prev_key_scratch_; - - // Active scan options (from Prepare()). + TrieIterator iter_; std::vector scan_opts_; - size_t current_scan_idx_; - bool prepared_; - - // True if the trie was built with a seqno side-table. When true: - // - SeekAndGetResult seeks with user key only, then does post-seek - // correction using seqno comparison - // - value() may return overflow block handles - // - NextAndGetResult advances within overflow runs before trie leaves - bool has_seqno_encoding_; - - // Reset overflow state to the default single-block position. - void ResetOverflowState() { - overflow_run_index_ = 0; - overflow_run_size_ = 1; - overflow_base_idx_ = 0; - } - - // Set up overflow state for the current trie leaf. When position_at_last - // is true (reverse iteration), positions at the last block in the run. - void SetupOverflowForCurrentLeaf(bool position_at_last) { - if (has_seqno_encoding_ && iter_.Valid()) { - uint64_t leaf_idx = iter_.LeafIndex(); - uint32_t block_count = trie_->GetLeafBlockCount(leaf_idx); - overflow_run_size_ = block_count; - overflow_base_idx_ = trie_->GetOverflowBase(leaf_idx); - if (position_at_last) { - overflow_run_index_ = block_count - 1; - } - } - } - - // Copy the current trie key into current_key_scratch_ and set the result. - void CopyTrieKeyToResult(IterateResult* result) { - Slice trie_key = iter_.Key(); - current_key_scratch_.assign(trie_key.data(), trie_key.size()); - result->key = Slice(current_key_scratch_); - } - - // ---- Overflow run state ---- - // - // When the trie iterator lands on a leaf that has block_count > 1 (a - // same-user-key run), these track the current position within the run: - // - // overflow_run_index_: 0 = primary block (trie leaf), 1+ = overflow block. - // Used by value() to select the right handle. - // overflow_run_size_: total blocks for current leaf (from leaf_block_counts). - // 1 = no overflow (the common case). - // overflow_base_idx_: starting index into overflow arrays for this leaf - // (from overflow_base_ prefix sum). - uint32_t overflow_run_index_; - uint32_t overflow_run_size_; - uint32_t overflow_base_idx_; + size_t current_scan_idx_ = 0; + bool prepared_ = false; }; // ============================================================================ // TrieIndexReader: Implements IndexFactoryReader. -// -// Owns (or references) the deserialized LoudsTrie and creates iterators -// for read operations. // ============================================================================ class TrieIndexReader : public IndexFactoryReader { public: explicit TrieIndexReader(const Comparator* comparator); - ~TrieIndexReader() override = default; // Initialize from serialized index data. The data must remain valid for // the lifetime of this reader (it's typically a block cache entry). Status InitFromSlice(const Slice& data); - // Create a new iterator for scanning. std::unique_ptr NewIterator( const ReadOptions& read_options) override; - // Approximate memory usage of the deserialized trie. + // The index is used in place, so only this object is charged. size_t ApproximateMemoryUsage() const override; private: const Comparator* comparator_; - LoudsTrie trie_; + Trie trie_; }; // ============================================================================ // TrieIndexFactory: Implements IndexFactory. // // Factory for creating TrieIndexBuilder (during SST file writes) and -// TrieIndexReader (during SST file reads). Registered as a Customizable -// with name "trie_index". +// TrieIndexReader (during SST file reads). Registered as a Customizable with +// name "trie_index". // ============================================================================ class TrieIndexFactory : public IndexFactory { public: @@ -347,6 +267,10 @@ class TrieIndexFactory : public IndexFactory { Status NewReader(const IndexFactoryOptions& option, Slice& index_block, std::unique_ptr& reader) const override; + + // Rejects a comparator other than BytewiseComparator. + Status ValidateOptions(const DBOptions& db_opts, + const ColumnFamilyOptions& cf_opts) const override; }; } // namespace trie_index diff --git a/utilities/trie_index/trie_index_test.cc b/utilities/trie_index/trie_index_test.cc index 06e608f66b12..cbd0662b6b3d 100644 --- a/utilities/trie_index/trie_index_test.cc +++ b/utilities/trie_index/trie_index_test.cc @@ -13,6 +13,7 @@ #include #include #include +#include #include #include #include @@ -23,6 +24,7 @@ #include "options/options_helper.h" #include "port/port.h" #include "rocksdb/comparator.h" +#include "rocksdb/convenience.h" #include "rocksdb/options.h" #include "rocksdb/slice.h" #include "rocksdb/sst_file_reader.h" @@ -38,8 +40,7 @@ #include "test_util/testutil.h" #include "util/random.h" #include "utilities/merge_operators.h" -#include "utilities/trie_index/bitvector.h" -#include "utilities/trie_index/louds_trie.h" +#include "utilities/trie_index/trie.h" #include "utilities/trie_index/trie_index_factory.h" namespace ROCKSDB_NAMESPACE { @@ -59,2299 +60,34 @@ IndexFactoryBuilder::IndexEntryContext EntryCtx(SequenceNumber last_seq, first_seq ? ((static_cast(first_seq) << 8) | 1) : 0}; } -// Bitvector tests -// ============================================================================ - -class BitvectorTest : public testing::Test {}; - -TEST_F(BitvectorTest, EmptyBitvector) { - BitvectorBuilder builder; - ASSERT_EQ(builder.NumBits(), 0u); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 0u); - ASSERT_EQ(bv.NumOnes(), 0u); - ASSERT_EQ(bv.NumZeros(), 0u); -} - -TEST_F(BitvectorTest, SingleBit) { - // Single 1-bit. - { - BitvectorBuilder builder; - builder.Append(true); - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 1u); - ASSERT_EQ(bv.NumOnes(), 1u); - ASSERT_TRUE(bv.GetBit(0)); - ASSERT_EQ(bv.Rank1(0), 0u); - ASSERT_EQ(bv.Rank1(1), 1u); - ASSERT_EQ(bv.FindNthOneBit(0), 0u); - ASSERT_EQ(bv.FindNthZeroBit(0), 1u); // No 0-bits; returns num_bits_. - } - // Single 0-bit. - { - BitvectorBuilder builder; - builder.Append(false); - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 1u); - ASSERT_EQ(bv.NumOnes(), 0u); - ASSERT_FALSE(bv.GetBit(0)); - ASSERT_EQ(bv.Rank1(0), 0u); - ASSERT_EQ(bv.Rank1(1), 0u); - ASSERT_EQ(bv.Rank0(1), 1u); - ASSERT_EQ(bv.FindNthZeroBit(0), 0u); - ASSERT_EQ(bv.FindNthOneBit(0), 1u); // No 1-bits; returns num_bits_. - } -} - -TEST_F(BitvectorTest, RankAndSelect) { - // Build a bitvector: 1010 1010 1010 ... (alternating bits). - const uint64_t n = 1000; - BitvectorBuilder builder; - for (uint64_t i = 0; i < n; i++) { - builder.Append(i % 2 == 0); // Even positions are 1. - } - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), n); - ASSERT_EQ(bv.NumOnes(), n / 2); - - // Verify rank at various positions. - for (uint64_t pos = 0; pos <= n; pos++) { - uint64_t expected_ones = (pos + 1) / 2; // Count of even numbers in [0,pos) - ASSERT_EQ(bv.Rank1(pos), expected_ones) << "Rank1 failed at pos=" << pos; - ASSERT_EQ(bv.Rank0(pos), pos - expected_ones) - << "Rank0 failed at pos=" << pos; - } - - // Verify select. - for (uint64_t i = 0; i < n / 2; i++) { - ASSERT_EQ(bv.FindNthOneBit(i), i * 2) << "FindNthOneBit failed at i=" << i; - ASSERT_EQ(bv.FindNthZeroBit(i), i * 2 + 1) - << "FindNthZeroBit failed at i=" << i; - } - - // Out of range selects. - ASSERT_EQ(bv.FindNthOneBit(n / 2), n); - ASSERT_EQ(bv.FindNthZeroBit(n / 2), n); -} - -TEST_F(BitvectorTest, AllOnes) { - const uint64_t n = 512; // Exactly two rank sample intervals (256 each). - BitvectorBuilder builder; - builder.AppendMultiple(true, n); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), n); - ASSERT_EQ(bv.NumOnes(), n); - - for (uint64_t pos = 0; pos <= n; pos++) { - ASSERT_EQ(bv.Rank1(pos), pos); - ASSERT_EQ(bv.Rank0(pos), 0u); - } - for (uint64_t i = 0; i < n; i++) { - ASSERT_EQ(bv.FindNthOneBit(i), i); - } - ASSERT_EQ(bv.FindNthZeroBit(0), n); // No 0-bits. -} - -TEST_F(BitvectorTest, AllZeros) { - const uint64_t n = 512; - BitvectorBuilder builder; - builder.AppendMultiple(false, n); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), n); - ASSERT_EQ(bv.NumOnes(), 0u); - - for (uint64_t pos = 0; pos <= n; pos++) { - ASSERT_EQ(bv.Rank1(pos), 0u); - ASSERT_EQ(bv.Rank0(pos), pos); - } - for (uint64_t i = 0; i < n; i++) { - ASSERT_EQ(bv.FindNthZeroBit(i), i); - } - ASSERT_EQ(bv.FindNthOneBit(0), n); // No 1-bits. -} - -TEST_F(BitvectorTest, LargeBitvector) { - // Test with a bitvector larger than one rank sample (kBitsPerRankSample). - const uint64_t n = 2000; - BitvectorBuilder builder; - uint64_t expected_ones = 0; - for (uint64_t i = 0; i < n; i++) { - bool bit = (i % 3 == 0); - builder.Append(bit); - if (bit) { - expected_ones++; - } - } - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), n); - ASSERT_EQ(bv.NumOnes(), expected_ones); - - // Verify rank at boundaries. - ASSERT_EQ(bv.Rank1(0), 0u); - ASSERT_EQ(bv.Rank1(n), expected_ones); - - // Verify select for a subset of positions. - uint64_t ones_so_far = 0; - for (uint64_t i = 0; i < n; i++) { - if (bv.GetBit(i)) { - ASSERT_EQ(bv.FindNthOneBit(ones_so_far), i) - << "FindNthOneBit failed at ones_so_far=" << ones_so_far; - ones_so_far++; - } - } - ASSERT_EQ(ones_so_far, expected_ones); -} - -TEST_F(BitvectorTest, SerializeAndDeserialize) { - const uint64_t n = 1500; - BitvectorBuilder builder; - for (uint64_t i = 0; i < n; i++) { - builder.Append(i % 5 == 0); - } - - Bitvector original; - original.BuildFrom(builder); - - // Serialize. - std::string serialized; - original.EncodeTo(&serialized); - ASSERT_GT(serialized.size(), 0u); - - // Deserialize. - Bitvector deserialized; - size_t consumed = 0; - ASSERT_OK(deserialized.InitFromData(serialized.data(), serialized.size(), - &consumed)); - ASSERT_GT(consumed, 0u); - ASSERT_EQ(consumed, serialized.size()); - - // Verify match. - ASSERT_EQ(deserialized.NumBits(), original.NumBits()); - ASSERT_EQ(deserialized.NumOnes(), original.NumOnes()); - - for (uint64_t pos = 0; pos <= n; pos++) { - ASSERT_EQ(deserialized.Rank1(pos), original.Rank1(pos)) - << "Rank1 mismatch at pos=" << pos; - } - - for (uint64_t i = 0; i < original.NumOnes(); i++) { - ASSERT_EQ(deserialized.FindNthOneBit(i), original.FindNthOneBit(i)) - << "FindNthOneBit mismatch at i=" << i; - } -} - -TEST_F(BitvectorTest, SerializeRoundTripNonAligned) { - // Test with a size that's not a multiple of 64. - const uint64_t n = 100; - BitvectorBuilder builder; - for (uint64_t i = 0; i < n; i++) { - builder.Append(i % 7 == 0); - } - - Bitvector original; - original.BuildFrom(builder); - - std::string serialized; - original.EncodeTo(&serialized); - - Bitvector deserialized; - size_t consumed = 0; - ASSERT_OK(deserialized.InitFromData(serialized.data(), serialized.size(), - &consumed)); - ASSERT_EQ(consumed, serialized.size()); - ASSERT_EQ(deserialized.NumBits(), n); - ASSERT_EQ(deserialized.NumOnes(), original.NumOnes()); - - // Verify every bit matches. - for (uint64_t i = 0; i < n; i++) { - ASSERT_EQ(deserialized.GetBit(i), original.GetBit(i)) - << "Bit mismatch at position " << i; - } -} - -TEST_F(BitvectorTest, AppendWord) { - // Test the AppendWord API used by the dense bitmap builder. - BitvectorBuilder builder; - - // Append 4 words to simulate a 256-bit dense node bitmap. - // Word 0: bits 0 and 3 set (labels 0x00 and 0x03). - builder.AppendWord(0x9, 64); // 0b1001 - builder.AppendWord(0x0, 64); // all zeros - builder.AppendWord(0x0, 64); // all zeros - builder.AppendWord(0x80000000'00000000ULL, 64); // bit 63 set (label 0xFF) - - ASSERT_EQ(builder.NumBits(), 256u); - ASSERT_TRUE(builder.GetBit(0)); // label 0x00 - ASSERT_TRUE(builder.GetBit(3)); // label 0x03 - ASSERT_FALSE(builder.GetBit(1)); - ASSERT_FALSE(builder.GetBit(64)); - ASSERT_TRUE(builder.GetBit(255)); // label 0xFF - - // Build and verify rank/select work correctly. - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 256u); - ASSERT_EQ(bv.NumOnes(), 3u); // bits 0, 3, 255 - ASSERT_EQ(bv.FindNthOneBit(0), 0u); - ASSERT_EQ(bv.FindNthOneBit(1), 3u); - ASSERT_EQ(bv.FindNthOneBit(2), 255u); -} - -TEST_F(BitvectorTest, AppendWordAndAppendInterleaved) { - // Verify AppendWord works when mixed with regular Append calls, - // as long as AppendWord is called on word-aligned boundaries. - BitvectorBuilder builder; - - // Append 64 bits via AppendWord. - builder.AppendWord(0xAAAAAAAA'AAAAAAAAULL, 64); - ASSERT_EQ(builder.NumBits(), 64u); - - // Append 64 more bits one at a time. - for (int i = 0; i < 64; i++) { - builder.Append(i % 2 == 0); - } - ASSERT_EQ(builder.NumBits(), 128u); - - // Another word-aligned AppendWord. - builder.AppendWord(0xFFFFFFFF'FFFFFFFFULL, 64); - ASSERT_EQ(builder.NumBits(), 192u); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 192u); - - // First word: alternating starting at bit 1 (0xAAAA... = 1010...) - ASSERT_FALSE(bv.GetBit(0)); - ASSERT_TRUE(bv.GetBit(1)); - // Second word: alternating starting at bit 0 - ASSERT_TRUE(bv.GetBit(64)); - ASSERT_FALSE(bv.GetBit(65)); - // Third word: all ones - ASSERT_TRUE(bv.GetBit(128)); - ASSERT_TRUE(bv.GetBit(191)); -} - -TEST_F(BitvectorTest, NextSetBit) { - // Build: 64 zeros, then 1, then 63 zeros, then 1. - BitvectorBuilder builder; - builder.AppendMultiple(false, 64); - builder.Append(true); - builder.AppendMultiple(false, 63); - builder.Append(true); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 129u); - - // NextSetBit from various positions. - ASSERT_EQ(bv.NextSetBit(0), 64u); // First set bit. - ASSERT_EQ(bv.NextSetBit(64), 64u); // Exact position. - ASSERT_EQ(bv.NextSetBit(65), 128u); // Next one. - ASSERT_EQ(bv.NextSetBit(128), 128u); // Exact. - ASSERT_EQ(bv.NextSetBit(129), 129u); // Past end -> num_bits_. -} - -TEST_F(BitvectorTest, NextSetBitAllZeros) { - BitvectorBuilder builder; - builder.AppendMultiple(false, 200); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NextSetBit(0), 200u); - ASSERT_EQ(bv.NextSetBit(100), 200u); - ASSERT_EQ(bv.NextSetBit(199), 200u); -} - -TEST_F(BitvectorTest, DistanceToNextSetBit) { - // Build: 1, 0, 0, 1, 0, 1 - BitvectorBuilder builder; - builder.Append(true); - builder.Append(false); - builder.Append(false); - builder.Append(true); - builder.Append(false); - builder.Append(true); - - Bitvector bv; - bv.BuildFrom(builder); - - ASSERT_EQ(bv.DistanceToNextSetBit(0), 3u); // 0 -> 3 - ASSERT_EQ(bv.DistanceToNextSetBit(3), 2u); // 3 -> 5 - // Last set bit: distance to next = num_bits_ - pos. - ASSERT_EQ(bv.DistanceToNextSetBit(5), 1u); // 5 -> 6 (num_bits_) -} - -TEST_F(BitvectorTest, MultipleRankSampleBoundaries) { - // Test rank/select across multiple rank sample boundaries. - const uint64_t n = 2048; // 8 rank samples (256 bits each). - BitvectorBuilder builder; - for (uint64_t i = 0; i < n; i++) { - // Set bit at every 100th position. - builder.Append(i % 100 == 0); - } - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumOnes(), 21u); // 0, 100, 200, ..., 2000 - - // Verify rank at sample boundaries. - ASSERT_EQ(bv.Rank1(512), 6u); // 0,100,200,300,400,500 - ASSERT_EQ(bv.Rank1(1024), 11u); // +600,700,800,900,1000 - ASSERT_EQ(bv.Rank1(1536), 16u); // +1100,1200,1300,1400,1500 - - // Verify select. - for (uint64_t i = 0; i < bv.NumOnes(); i++) { - ASSERT_EQ(bv.FindNthOneBit(i), i * 100); - } -} - -TEST_F(BitvectorTest, InitFromDataCorruption) { - // Test that InitFromData correctly rejects truncated data. - Bitvector bv; - size_t consumed; - - // Too short for header. - ASSERT_TRUE(bv.InitFromData("short", 5, &consumed).IsCorruption()); - - // Header says more words than available. - std::string bad; - uint64_t num_bits = 1024; - uint64_t num_ones = 0; - bad.append(reinterpret_cast(&num_bits), 8); - bad.append(reinterpret_cast(&num_ones), 8); - // No word data follows. - ASSERT_TRUE( - bv.InitFromData(bad.data(), bad.size(), &consumed).IsCorruption()); -} - -TEST_F(BitvectorTest, FindNthSetBitInWordExhaustive) { - // Verify FindNthSetBitInWord for every position in a word with known - // popcount. - uint64_t word = 0xDEADBEEFCAFEBABEULL; - uint64_t pc = Popcount(word); - - for (uint64_t i = 0; i < pc; i++) { - uint64_t pos = FindNthSetBitInWord(word, i); - // The pos-th bit should be set. - ASSERT_TRUE((word >> pos) & 1) - << "FindNthSetBitInWord returned non-set bit at i=" << i; - // Exactly i bits should be set before pos. - uint64_t mask = - (pos == 0) ? 0 : (pos == 64 ? ~uint64_t(0) : (uint64_t(1) << pos) - 1); - ASSERT_EQ(Popcount(word & mask), i) - << "Wrong rank at FindNthSetBitInWord result for i=" << i; - } -} - -TEST_F(BitvectorTest, MoveConstructor) { - // Build a non-trivial bitvector and verify move constructor preserves - // correctness, especially the RecomputePointers() logic. - const uint64_t n = 1000; - BitvectorBuilder builder; - for (uint64_t i = 0; i < n; i++) { - builder.Append(i % 3 == 0); - } - - Bitvector original; - original.BuildFrom(builder); - uint64_t orig_ones = original.NumOnes(); - uint64_t orig_rank_500 = original.Rank1(500); - uint64_t orig_select_10 = original.FindNthOneBit(10); - - // Move-construct a new bitvector. - Bitvector moved(std::move(original)); - - // Moved-to should preserve all data. - ASSERT_EQ(moved.NumBits(), n); - ASSERT_EQ(moved.NumOnes(), orig_ones); - ASSERT_EQ(moved.Rank1(500), orig_rank_500); - ASSERT_EQ(moved.FindNthOneBit(10), orig_select_10); - - // Verify full rank/select consistency. - for (uint64_t i = 0; i < moved.NumOnes(); i++) { - uint64_t pos = moved.FindNthOneBit(i); - ASSERT_TRUE(moved.GetBit(pos)) - << "Bit not set at FindNthOneBit(" << i << ")"; - } - - // Moved-from should be empty/zeroed. - // NOLINTBEGIN(bugprone-use-after-move) - ASSERT_EQ(original.NumBits(), 0u); - ASSERT_EQ(original.NumOnes(), 0u); - // NOLINTEND(bugprone-use-after-move) -} - -TEST_F(BitvectorTest, MoveAssignment) { - // Build two bitvectors, move-assign the first to the second. - const uint64_t n = 500; - BitvectorBuilder builder1; - for (uint64_t i = 0; i < n; i++) { - builder1.Append(i % 5 == 0); - } - Bitvector bv1; - bv1.BuildFrom(builder1); - - BitvectorBuilder builder2; - for (uint64_t i = 0; i < 200; i++) { - builder2.Append(true); - } - Bitvector bv2; - bv2.BuildFrom(builder2); - - uint64_t bv1_ones = bv1.NumOnes(); - uint64_t bv1_rank_250 = bv1.Rank1(250); - - // Move-assign bv1 to bv2. - bv2 = std::move(bv1); - - // bv2 should now have bv1's data. - ASSERT_EQ(bv2.NumBits(), n); - ASSERT_EQ(bv2.NumOnes(), bv1_ones); - ASSERT_EQ(bv2.Rank1(250), bv1_rank_250); - - // Moved-from should be empty. - ASSERT_EQ(bv1.NumBits(), 0u); // NOLINT(bugprone-use-after-move) -} - -TEST_F(BitvectorTest, SerializedSizeMatchesEncoded) { - // Verify SerializedSize() matches actual encoded output length. - const uint64_t n = 1000; - BitvectorBuilder builder; - for (uint64_t i = 0; i < n; i++) { - builder.Append(i % 7 == 0); - } - Bitvector bv; - bv.BuildFrom(builder); - - std::string encoded; - bv.EncodeTo(&encoded); - ASSERT_EQ(bv.SerializedSize(), encoded.size()); -} - -TEST_F(BitvectorTest, DistanceToNextSetBitCrossWordBoundary) { - // Test DistanceToNextSetBit where the next set bit is in a different - // 64-bit word. Build: bit 0 set, then 127 zeros, then bit 128 set. - BitvectorBuilder builder; - builder.Append(true); // pos 0 - builder.AppendMultiple(false, 127); - builder.Append(true); // pos 128 - builder.Append(false); - - Bitvector bv; - bv.BuildFrom(builder); - - // Distance from pos 0 to next set bit (pos 128) = 128. - ASSERT_EQ(bv.DistanceToNextSetBit(0), 128u); - // Distance from pos 128 is to end (no more set bits after). - // DistanceToNextSetBit returns distance to next set bit after pos, - // where pos itself must be a set bit. -} - -TEST_F(BitvectorTest, NumZeros) { - // Verify NumZeros for various patterns. - const uint64_t n = 300; - BitvectorBuilder builder; - uint64_t expected_ones = 0; - for (uint64_t i = 0; i < n; i++) { - bool bit = (i % 4 == 0); - builder.Append(bit); - if (bit) { - expected_ones++; - } - } - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), n); - ASSERT_EQ(bv.NumOnes(), expected_ones); - ASSERT_EQ(bv.NumZeros(), n - expected_ones); -} - -TEST_F(BitvectorTest, InitFromDataNumOnesExceedsNumBits) { - // Craft a header where num_ones > num_bits -- should return Corruption. - std::string buf(2 * sizeof(uint64_t), '\0'); - uint64_t num_bits = 64; - uint64_t num_ones = 65; // Invalid: more ones than bits. - memcpy(&buf[0], &num_bits, sizeof(uint64_t)); - memcpy(&buf[sizeof(uint64_t)], &num_ones, sizeof(uint64_t)); - - Bitvector bv; - size_t consumed = 0; - Status s = bv.InitFromData(buf.data(), buf.size(), &consumed); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(BitvectorTest, InitFromDataNumBitsExceedsUint32) { - // Craft a header where num_bits > UINT32_MAX -- should return Corruption. - std::string buf(2 * sizeof(uint64_t), '\0'); - uint64_t num_bits = static_cast(UINT32_MAX) + 1; - uint64_t num_ones = 0; - memcpy(&buf[0], &num_bits, sizeof(uint64_t)); - memcpy(&buf[sizeof(uint64_t)], &num_ones, sizeof(uint64_t)); - - Bitvector bv; - size_t consumed = 0; - Status s = bv.InitFromData(buf.data(), buf.size(), &consumed); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(BitvectorTest, InitFromDataWordsNotAligned) { - // Build a valid bitvector, serialize it, then read from a buffer where the - // words data starts at an odd offset (not 8-byte aligned). - BitvectorBuilder builder; - for (int i = 0; i < 128; i++) { - builder.Append(i % 3 == 0); - } - Bitvector bv_orig; - bv_orig.BuildFrom(builder); - - std::string serialized; - bv_orig.EncodeTo(&serialized); - - // Allocate buffer with 1-byte extra and offset by 1 to break alignment. - std::unique_ptr unaligned_buf(new char[serialized.size() + 1]); - memcpy(unaligned_buf.get() + 1, serialized.data(), serialized.size()); - - Bitvector bv; - size_t consumed = 0; - Status s = - bv.InitFromData(unaligned_buf.get() + 1, serialized.size(), &consumed); - // The header is 16 bytes; words start right after. If the base pointer is - // offset by 1, then the words pointer is at (base+1+16) which is odd. - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(BitvectorTest, InitFromDataTruncatedRankLUT) { - // Build a valid bitvector, serialize, then truncate after the words section - // so the rank LUT cannot be read. - BitvectorBuilder builder; - for (int i = 0; i < 512; i++) { - builder.Append(i % 5 == 0); - } - Bitvector bv_orig; - bv_orig.BuildFrom(builder); - - std::string serialized; - bv_orig.EncodeTo(&serialized); - - // Header = 16 bytes, words = ceil(512/64)*8 = 64 bytes. Truncate after that. - size_t truncated_size = 2 * sizeof(uint64_t) + ((512 + 63) / 64) * 8; - ASSERT_LT(truncated_size, serialized.size()); - - Bitvector bv; - size_t consumed = 0; - Status s = bv.InitFromData(serialized.data(), truncated_size, &consumed); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(BitvectorTest, InitFromDataTruncatedSelect1Hints) { - // Build a bitvector with enough ones to have select1 hints, then truncate - // after the rank LUT so select1 hints cannot be read. - BitvectorBuilder builder; - // 600 bits, all ones -> lots of select1 hints. - for (int i = 0; i < 600; i++) { - builder.Append(true); - } - Bitvector bv_orig; - bv_orig.BuildFrom(builder); - - std::string serialized; - bv_orig.EncodeTo(&serialized); - - // We need to find a truncation point after the rank LUT but before select1 - // hints finish. Use a binary approach: try reading at decreasing sizes. - // A valid read needs the full buffer. Remove the last 8 bytes (at minimum). - size_t truncated_size = serialized.size() - 8; - - Bitvector bv; - size_t consumed = 0; - Status s = bv.InitFromData(serialized.data(), truncated_size, &consumed); - // Should fail for either select1 or select0 hints truncation. - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(BitvectorTest, InitFromDataTruncatedSelect0Hints) { - // Build a bitvector with enough zeros to have select0 hints, then truncate - // so select0 hints cannot be fully read. - BitvectorBuilder builder; - // 600 bits, all zeros -> lots of select0 hints, zero select1 hints. - for (int i = 0; i < 600; i++) { - builder.Append(false); - } - Bitvector bv_orig; - bv_orig.BuildFrom(builder); - - std::string serialized; - bv_orig.EncodeTo(&serialized); - - // Remove the last 8 bytes. - size_t truncated_size = serialized.size() - 8; - - Bitvector bv; - size_t consumed = 0; - Status s = bv.InitFromData(serialized.data(), truncated_size, &consumed); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(BitvectorTest, BuilderReserve) { - // Verify Reserve() doesn't change behavior, just avoids reallocation. - BitvectorBuilder builder; - builder.Reserve(1000); - for (int i = 0; i < 1000; i++) { - builder.Append(i % 2 == 0); - } - ASSERT_EQ(builder.NumBits(), 1000u); - - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 1000u); - ASSERT_EQ(bv.NumOnes(), 500u); - // Verify correctness of rank/select after Reserve. - ASSERT_EQ(bv.FindNthOneBit(0), 0u); - ASSERT_EQ(bv.FindNthOneBit(1), 2u); - ASSERT_EQ(bv.FindNthZeroBit(0), 1u); -} - -TEST_F(BitvectorTest, AppendMultipleZeroCount) { - // AppendMultiple with count=0 should be a no-op. - BitvectorBuilder builder; - builder.Append(true); - builder.AppendMultiple(true, 0); - builder.AppendMultiple(false, 0); - ASSERT_EQ(builder.NumBits(), 1u); - ASSERT_TRUE(builder.GetBit(0)); -} - -TEST_F(BitvectorTest, AppendMultipleTruePartialWord) { - // AppendMultiple(true, count) starting at a non-word-aligned position. - // This exercises the partial-word fill path at the beginning. - BitvectorBuilder builder; - // Start with 7 false bits to offset by 7 within the first word. - builder.AppendMultiple(false, 7); - // Now append 200 true bits starting at bit position 7. - builder.AppendMultiple(true, 200); - // Then 5 more false bits. - builder.AppendMultiple(false, 5); - - ASSERT_EQ(builder.NumBits(), 212u); - // Verify the pattern: 7 zeros, 200 ones, 5 zeros. - for (uint64_t i = 0; i < 7; i++) { - ASSERT_FALSE(builder.GetBit(i)) << "Expected 0 at " << i; - } - for (uint64_t i = 7; i < 207; i++) { - ASSERT_TRUE(builder.GetBit(i)) << "Expected 1 at " << i; - } - for (uint64_t i = 207; i < 212; i++) { - ASSERT_FALSE(builder.GetBit(i)) << "Expected 0 at " << i; - } - - // Build the bitvector and verify rank/select correctness. - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumOnes(), 200u); - ASSERT_EQ(bv.FindNthOneBit(0), 7u); - ASSERT_EQ(bv.FindNthOneBit(199), 206u); - ASSERT_EQ(bv.FindNthZeroBit(0), 0u); - ASSERT_EQ(bv.FindNthZeroBit(7), 207u); -} - -TEST_F(BitvectorTest, FindNthZeroBitNonAligned) { - // Build a bitvector whose length is not a multiple of 64, then verify - // FindNthZeroBit handles the last partial word correctly. - BitvectorBuilder builder; - // 100 bits: alternating 1, 0. - for (int i = 0; i < 100; i++) { - builder.Append(i % 2 == 0); - } - // 100 bits, 50 ones (at even positions), 50 zeros (at odd positions). - Bitvector bv; - bv.BuildFrom(builder); - ASSERT_EQ(bv.NumBits(), 100u); - ASSERT_EQ(bv.NumOnes(), 50u); - - // The zeros are at positions 1, 3, 5, ..., 99. - for (uint64_t i = 0; i < 50; i++) { - ASSERT_EQ(bv.FindNthZeroBit(i), 2 * i + 1) << "Mismatch at i=" << i; - } - // Asking for the 50th zero-bit (0-indexed) should return num_bits_ (100). - ASSERT_EQ(bv.FindNthZeroBit(50), 100u); -} - -// ============================================================================ -// EliasFano tests -// ============================================================================ - -class EliasFanoTest : public testing::Test {}; - -TEST_F(EliasFanoTest, EmptySequence) { - EliasFano ef; - std::vector values; - ef.BuildFrom(values.data(), 0, 0); - ASSERT_EQ(ef.Count(), 0u); - ASSERT_EQ(ef.Universe(), 0u); -} - -TEST_F(EliasFanoTest, SingleElement) { - EliasFano ef; - uint64_t val = 42; - ef.BuildFrom(&val, 1, 100); - ASSERT_EQ(ef.Count(), 1u); - ASSERT_EQ(ef.Universe(), 100u); - ASSERT_EQ(ef.Access(0), 42u); -} - -TEST_F(EliasFanoTest, BuildAccessVariants) { - // Table-driven test covering different value distributions. - struct Case { - const char* name; - std::vector values; - uint64_t universe; - }; - - // Monotonic sequence (uniform spacing). - std::vector monotonic = {0, 10, 20, 30, 40, 50, - 60, 70, 80, 90, 100}; - // Constant sequence (all same value, universe = count). - std::vector constant(50, 7); - // Large universe simulating 4GB SST block offsets. - const uint64_t large_universe = uint64_t(4) * 1024 * 1024 * 1024; - std::vector large(1000); - for (size_t i = 0; i < large.size(); i++) { - large[i] = i * (large_universe / large.size()); - } - // Dense consecutive: 0, 1, 2, ..., 99. - std::vector consecutive(100); - for (size_t i = 0; i < 100; i++) { - consecutive[i] = i; - } - // Word-boundary crossing: low_bits ~17, small offsets per element. - std::vector word_boundary(100); - for (size_t i = 0; i < 100; i++) { - word_boundary[i] = i * 131072 + (i % 7); - } - - std::vector cases = { - {"MonotonicSequence", monotonic, 101}, - {"ConstantSequence", constant, 8}, - {"LargeUniverse", large, large_universe}, - {"ConsecutiveValues", consecutive, 100}, - {"WordBoundaryCrossing", word_boundary, 100 * 131072}, - }; - - for (const auto& tc : cases) { - SCOPED_TRACE(tc.name); - EliasFano ef; - ef.BuildFrom(tc.values.data(), tc.values.size(), tc.universe); - ASSERT_EQ(ef.Count(), tc.values.size()); - ASSERT_EQ(ef.Universe(), tc.universe); - for (size_t i = 0; i < tc.values.size(); i++) { - ASSERT_EQ(ef.Access(i), tc.values[i]) << "Mismatch at i=" << i; - } - } -} - -TEST_F(EliasFanoTest, SerializeDeserializeRoundTrip) { - // Build, serialize, deserialize, verify all values match. - std::vector values = {0, 5, 100, 1000, 5000, 10000, 50000}; - EliasFano original; - original.BuildFrom(values.data(), values.size(), 100000); - - std::string serialized; - original.EncodeTo(&serialized); - ASSERT_GT(serialized.size(), 0u); - ASSERT_EQ(original.SerializedSize(), serialized.size()); - - EliasFano deserialized; - size_t consumed = 0; - ASSERT_OK(deserialized.InitFromData(serialized.data(), serialized.size(), - &consumed)); - ASSERT_EQ(consumed, serialized.size()); - ASSERT_EQ(deserialized.Count(), original.Count()); - ASSERT_EQ(deserialized.Universe(), original.Universe()); - - for (size_t i = 0; i < values.size(); i++) { - ASSERT_EQ(deserialized.Access(i), values[i]) << "Mismatch at i=" << i; - } -} - -TEST_F(EliasFanoTest, SerializeDeserializeEmpty) { - EliasFano original; - std::vector empty; - original.BuildFrom(empty.data(), 0, 0); - - std::string serialized; - original.EncodeTo(&serialized); - - EliasFano deserialized; - size_t consumed = 0; - ASSERT_OK(deserialized.InitFromData(serialized.data(), serialized.size(), - &consumed)); - ASSERT_EQ(deserialized.Count(), 0u); -} - -TEST_F(EliasFanoTest, InitFromDataCorruptionTruncatedHeader) { - std::string bad_data(10, '\0'); // Too short for 3 uint64_t header. - EliasFano ef; - size_t consumed = 0; - ASSERT_TRUE(ef.InitFromData(bad_data.data(), bad_data.size(), &consumed) - .IsCorruption()); -} - -TEST_F(EliasFanoTest, InitFromDataCorruptionBadLowBits) { - // Build a valid EliasFano, then corrupt the low_bits field to > 63. - std::vector values = {0, 10, 20}; - EliasFano ef; - ef.BuildFrom(values.data(), values.size(), 100); - - std::string serialized; - ef.EncodeTo(&serialized); - - // low_bits is the third uint64_t (bytes 16..23). Set to 64. - uint64_t bad_low_bits = 64; - memcpy(&serialized[16], &bad_low_bits, sizeof(uint64_t)); - - EliasFano corrupt; - size_t consumed = 0; - ASSERT_TRUE( - corrupt.InitFromData(serialized.data(), serialized.size(), &consumed) - .IsCorruption()); -} - -TEST_F(EliasFanoTest, InitFromDataCorruptionTruncatedLowWords) { - // Build valid, then truncate so low words are incomplete. - std::vector values = {0, 100, 200, 300, 400}; - EliasFano ef; - ef.BuildFrom(values.data(), values.size(), 1000); - - std::string serialized; - ef.EncodeTo(&serialized); - - // Truncate a few bytes off the end. - std::string truncated = serialized.substr(0, serialized.size() - 4); - EliasFano corrupt; - size_t consumed = 0; - ASSERT_TRUE( - corrupt.InitFromData(truncated.data(), truncated.size(), &consumed) - .IsCorruption()); -} - -TEST_F(EliasFanoTest, LargeCount) { - // Test with 10K elements to exercise multi-word packing. - const size_t count = 10000; - const uint64_t universe = count * 1000; - std::vector values(count); - for (size_t i = 0; i < count; i++) { - values[i] = i * 1000; - } - EliasFano ef; - ef.BuildFrom(values.data(), values.size(), universe); - ASSERT_EQ(ef.Count(), count); - - // Spot-check values. - ASSERT_EQ(ef.Access(0), 0u); - ASSERT_EQ(ef.Access(count / 2), (count / 2) * 1000); - ASSERT_EQ(ef.Access(count - 1), (count - 1) * 1000); - - // Full verification. - for (size_t i = 0; i < count; i++) { - ASSERT_EQ(ef.Access(i), values[i]) << "Mismatch at i=" << i; - } -} - -TEST_F(EliasFanoTest, InitFromDataTruncatedHighBitvector) { - // Build a valid EliasFano, serialize it, then truncate so the high bitvector - // cannot be fully read. - std::vector values = {10, 20, 30, 40, 50}; - EliasFano ef; - ef.BuildFrom(values.data(), values.size(), 100); - - std::string serialized; - ef.EncodeTo(&serialized); - - // Header is 3 * 8 = 24 bytes. Truncate just a few bytes past the header - // so the high bitvector InitFromData fails. - size_t truncated_size = 3 * sizeof(uint64_t) + 4; - ASSERT_LT(truncated_size, serialized.size()); - - EliasFano ef2; - size_t consumed = 0; - Status s = ef2.InitFromData(serialized.data(), truncated_size, &consumed); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); -} - -TEST_F(EliasFanoTest, MoveConstructor) { - // Verify that move-constructing an EliasFano preserves correctness. - // Use the InitFromData path (production path) since the default move ctor - // doesn't re-seat low_words_ into owned_low_data_. When initialized from - // serialized data, low_words_ points into external memory (not - // owned_low_data_), so the default move is safe. - std::vector values = {0, 100, 200, 300, 400}; - EliasFano ef_tmp; - ef_tmp.BuildFrom(values.data(), values.size(), 500); - - std::string serialized; - ef_tmp.EncodeTo(&serialized); - - EliasFano ef1; - size_t consumed = 0; - ASSERT_OK(ef1.InitFromData(serialized.data(), serialized.size(), &consumed)); - - // Move-construct ef2 from ef1. - EliasFano ef2(std::move(ef1)); - - ASSERT_EQ(ef2.Count(), 5u); - ASSERT_EQ(ef2.Universe(), 500u); - for (size_t i = 0; i < values.size(); i++) { - ASSERT_EQ(ef2.Access(i), values[i]) << "Mismatch at i=" << i; - } -} - -TEST_F(EliasFanoTest, MoveAssignment) { - // Verify that move-assigning an EliasFano preserves correctness. - // Use the InitFromData path (see MoveConstructor comment for rationale). - std::vector values = {5, 15, 25, 35, 45}; - EliasFano ef_tmp; - ef_tmp.BuildFrom(values.data(), values.size(), 50); - - std::string serialized; - ef_tmp.EncodeTo(&serialized); - - EliasFano ef1; - size_t consumed = 0; - ASSERT_OK(ef1.InitFromData(serialized.data(), serialized.size(), &consumed)); - - // Move-assign ef1 into a fresh ef2. - EliasFano ef2; - ef2 = std::move(ef1); - - ASSERT_EQ(ef2.Count(), 5u); - ASSERT_EQ(ef2.Universe(), 50u); - for (size_t i = 0; i < values.size(); i++) { - ASSERT_EQ(ef2.Access(i), values[i]) << "Mismatch at i=" << i; - } -} - -// ============================================================================ -// LoudsTrie tests -// ============================================================================ - -class LoudsTrieTest : public testing::Test { - protected: - // Owns builder (and its serialized data), trie, and iterator. The handle - // for key[i] is {i * 100, 50} by convention. - // - // Custom move constructor is required because `iter` holds a raw pointer to - // `trie`. A default move would transfer the unique_ptr but leave it - // pointing at the moved-from trie (dangling). We re-create the iterator - // against the new trie location instead. - struct BuiltTrie { - LoudsTrieBuilder builder; - LoudsTrie trie; - std::unique_ptr iter; - - BuiltTrie() = default; - ~BuiltTrie() = default; - BuiltTrie(BuiltTrie&& other) noexcept - : builder(std::move(other.builder)), trie(std::move(other.trie)) { - if (other.iter) { - iter = std::make_unique(&trie); - } - } - BuiltTrie& operator=(BuiltTrie&&) = delete; - BuiltTrie(const BuiltTrie&) = delete; - BuiltTrie& operator=(const BuiltTrie&) = delete; - }; - - // Build a trie from pre-sorted keys using the standard {i*100, 50} handles. - // Returns a BuiltTrie whose iterator is ready for Seek/Next. - static BuiltTrie BuildTrieFromKeys(const std::vector& keys) { - BuiltTrie bt; - for (size_t i = 0; i < keys.size(); i++) { - bt.builder.AddKey(Slice(keys[i]), TrieBlockHandle{i * 100, 50}); - } - bt.builder.Finish(); - Status s = bt.trie.InitFromData(bt.builder.GetSerializedData()); - EXPECT_OK(s); - bt.iter = std::make_unique(&bt.trie); - return bt; - } - - // Seek to the first key and walk the full trie with Next, verifying every - // key string and handle offset. - static void VerifyFullScan(const BuiltTrie& bt, - const std::vector& keys) { - auto* iter = bt.iter.get(); - ASSERT_TRUE(iter->Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter->Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter->Key().ToString(), keys[i]) << "Key mismatch at i=" << i; - ASSERT_EQ(iter->Value().offset, i * 100) << "Handle mismatch at i=" << i; - if (i < keys.size() - 1) { - ASSERT_TRUE(iter->Next()) << "Next failed at i=" << i; - } - } - ASSERT_FALSE(iter->Next()); - ASSERT_FALSE(iter->Valid()); - } -}; - -// Helper: build a trie from sorted keys and verify that Seek and Next -// iterate in exactly the expected order, checking both handles and -// reconstructed keys. -static void VerifyTrieIteration(const std::vector& keys) { - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - TrieBlockHandle h{i * 100, 50}; - builder.AddKey(Slice(keys[i]), h); - } - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), keys.size()); - - LoudsTrieIterator iter(&trie); - - // Test 1: Seek to every key and verify exact match + correct handle. - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Seek(Slice(keys[i]))) - << "Seek failed for key[" << i << "]=\"" << keys[i] << "\""; - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), keys[i]) << "Key mismatch at i=" << i; - ASSERT_EQ(iter.Value().offset, i * 100) << "Handle mismatch at i=" << i; - } - - // Test 2: Seek to first key, then Next through all keys. - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter.Key().ToString(), keys[i]) - << "Key mismatch during iteration at i=" << i; - ASSERT_EQ(iter.Value().offset, i * 100) - << "Handle mismatch during iteration at i=" << i; - if (i < keys.size() - 1) { - ASSERT_TRUE(iter.Next()) << "Next failed at i=" << i; - } - } - ASSERT_FALSE(iter.Next()); - ASSERT_FALSE(iter.Valid()); - - // Test 3: Seek past all keys. - std::string past_all = keys.back() + "zzz"; - ASSERT_FALSE(iter.Seek(Slice(past_all))); - ASSERT_FALSE(iter.Valid()); - - // Test 4: Seek before all keys. - std::string before_all(1, keys[0][0] - 1); - if (keys[0][0] > 0) { - ASSERT_TRUE(iter.Seek(Slice(before_all))); - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), keys[0]); - } -} - -// Table-driven test consolidating individual VerifyTrieIteration tests. -// Each case generates a key set covering a distinct trie topology: prefix -// chains, binary content, dense/sparse boundaries, varying lengths, etc. -TEST_F(LoudsTrieTest, TrieTopologyVariants) { - auto run = [](const char* name, std::vector keys) { - SCOPED_TRACE(name); - std::sort(keys.begin(), keys.end()); - keys.erase(std::unique(keys.begin(), keys.end()), keys.end()); - ASSERT_NO_FATAL_FAILURE(VerifyTrieIteration(keys)); - }; - - // Minimal multi-key trie. - ASSERT_NO_FATAL_FAILURE(run("TwoKeys", {"aa", "ab"})); - - // All keys share a long prefix, differ only in last byte. - { - std::string prefix(50, 'x'); - std::vector keys; - keys.reserve(26); - for (char c = 'a'; c <= 'z'; c++) { - keys.push_back(prefix + c); - } - ASSERT_NO_FATAL_FAILURE( - run("IdenticalPrefixDifferentLastByte", std::move(keys))); - } - - // Keys containing non-ASCII bytes (0x00, 0xFF, etc). - ASSERT_NO_FATAL_FAILURE(run( - "BinaryKeys", {std::string("\x00\x00", 2), std::string("\x00\x01", 2), - std::string("\x00\xFF", 2), std::string("\x01\x00", 2), - std::string("\xFF\x00", 2), std::string("\xFF\xFF", 2)})); - - // One key for every possible single byte value (0x00-0xFF). - { - std::vector keys; - keys.reserve(256); - for (int b = 0; b < 256; b++) { - keys.emplace_back(1, static_cast(b)); - } - ASSERT_NO_FATAL_FAILURE(run("SingleByteAllValues", std::move(keys))); - } - - // Single path through the trie (linear chain), depth = 100. - ASSERT_NO_FATAL_FAILURE(run("DeepChain", {std::string(100, 'a')})); - - // Mix of prefix keys and non-prefix keys in the same subtree. - ASSERT_NO_FATAL_FAILURE( - run("AlternatingPrefixAndNonPrefix", {"a", "ab", "abc", "b", "c"})); - - // Prefix key at the dense/sparse boundary. - { - std::vector keys; - keys.reserve(28); - for (char c = 'a'; c <= 'z'; c++) { - keys.emplace_back(1, c); - } - keys.emplace_back("aa"); - keys.emplace_back("aab"); - ASSERT_NO_FATAL_FAILURE(run("PrefixKeyAtDenseSparseEdge", std::move(keys))); - } - - // Key reconstruction variants. - ASSERT_NO_FATAL_FAILURE( - run("KeyReconstructionBasic", {"abc", "abd", "abe", "xyz"})); - ASSERT_NO_FATAL_FAILURE( - run("KeyReconstructionSingleByte", {"a", "b", "c", "d", "e"})); - - { - std::vector keys; - std::string prefix = "common_prefix_that_is_quite_long_"; - keys.reserve(26); - for (int i = 0; i < 26; i++) { - keys.push_back(prefix + static_cast('a' + i)); - } - ASSERT_NO_FATAL_FAILURE( - run("KeyReconstructionLongSharedPrefix", std::move(keys))); - } - - ASSERT_NO_FATAL_FAILURE( - run("KeyReconstructionVaryingLengths", - {"a", "ab", "abc", "abcd", "abcde", "b", "bc", "bcd"})); - ASSERT_NO_FATAL_FAILURE(run("TwoByteKeys", {"aa", "ab", "ba", "bb", "ca"})); - - { - std::vector keys; - keys.reserve(26); - for (int c = 'a'; c <= 'z'; c++) { - keys.push_back(std::string(1, static_cast(c)) + "suffix"); - } - ASSERT_NO_FATAL_FAILURE(run("HighFanoutRoot", std::move(keys))); - } - - ASSERT_NO_FATAL_FAILURE( - run("DeepTrie", - {"abcdefghijklmnop", "abcdefghijklmnoq", "abcdefghijklmnor"})); - - // Prefix key variants. - ASSERT_NO_FATAL_FAILURE(run("PrefixKeysSimple", {"a", "ab", "b"})); - ASSERT_NO_FATAL_FAILURE(run("PrefixKeysChain", {"a", "ab", "abc", "abcd"})); - ASSERT_NO_FATAL_FAILURE( - run("PrefixKeysMultipleBranches", {"a", "aa", "ab", "ac"})); - ASSERT_NO_FATAL_FAILURE( - run("PrefixKeysAtDifferentLevels", {"a", "aa", "aab", "b", "ba"})); - - { - std::vector keys; - keys.reserve(30); - for (int i = 0; i < 30; i++) { - keys.push_back("p_" + std::to_string(i)); - } - ASSERT_NO_FATAL_FAILURE(run("PrefixKeysDiverging", std::move(keys))); - } - - ASSERT_NO_FATAL_FAILURE(run("PrefixKeyRootOnly", {"a", "ab", "ac", "b"})); - ASSERT_NO_FATAL_FAILURE( - run("MultiplePrefixKeySameNode", {"x", "xa", "xab", "xb"})); - - { - std::vector keys; - keys.reserve(20); - for (int len = 1; len <= 20; len++) { - keys.emplace_back(len, 'a'); - } - ASSERT_NO_FATAL_FAILURE(run("ManyPrefixKeys", std::move(keys))); - } - - // Dense/sparse boundary variants. - ASSERT_NO_FATAL_FAILURE(run("AllSparse", {"a", "b"})); - - { - std::vector keys; - keys.reserve(26); - for (char c = 'a'; c <= 'z'; c++) { - keys.emplace_back(1, c); - } - ASSERT_NO_FATAL_FAILURE(run("AllDense", std::move(keys))); - } - - { - std::vector keys; - for (char c = 'a'; c <= 'z'; c++) { - keys.push_back(std::string(1, c) + "1"); - keys.push_back(std::string(1, c) + "2"); - } - ASSERT_NO_FATAL_FAILURE(run("DenseSparseTransition", std::move(keys))); - } - - // Larger key sets (former deterministic stress tests). - { - std::vector keys; - char buf[32]; - for (int i = 0; i < 1000; i++) { - snprintf(buf, sizeof(buf), "key_%06d", i); - keys.emplace_back(buf); - } - ASSERT_NO_FATAL_FAILURE(run("ThousandFormattedKeys", std::move(keys))); - } - - { - std::vector keys; - const char* prefixes[] = {"alpha", "beta", "gamma", "delta", "epsilon"}; - char buf[64]; - for (const char* p : prefixes) { - for (int i = 0; i < 50; i++) { - snprintf(buf, sizeof(buf), "%s_%03d", p, i); - keys.emplace_back(buf); - } - } - ASSERT_NO_FATAL_FAILURE(run("DiversePrefixPatterns", std::move(keys))); - } - - { - std::vector keys; - keys.reserve(176); - char buf[64]; - for (char c = 'a'; c <= 'z'; c++) { - keys.emplace_back(1, c); - } - for (int i = 0; i < 100; i++) { - snprintf(buf, sizeof(buf), "medium_%04d", i); - keys.emplace_back(buf); - } - for (int i = 0; i < 50; i++) { - snprintf(buf, sizeof(buf), "this_is_a_very_long_key_%06d", i); - keys.emplace_back(buf); - } - ASSERT_NO_FATAL_FAILURE(run("MixedLengths", std::move(keys))); - } - - { - std::vector keys; - keys.reserve(200); - for (int i = 0; i < 200; i++) { - keys.push_back("item_" + std::to_string(i)); - } - ASSERT_NO_FATAL_FAILURE(run("PrefixKeyStressPatterns", std::move(keys))); - } -} - -// Randomized stress test: generates random key sets with varying sizes, -// lengths, and character distributions, verifying trie correctness for each. -// Uses a time-based seed for broad coverage across runs; seed is logged -// via SCOPED_TRACE for reproducibility. -TEST_F(LoudsTrieTest, RandomizedStress) { - uint32_t seed = static_cast( +// Seed for randomized tests; log it with SCOPED_TRACE to reproduce a failure. +uint32_t TimeSeed() { + return static_cast( std::chrono::system_clock::now().time_since_epoch().count()); - SCOPED_TRACE("seed=" + std::to_string(seed)); - Random rnd(seed); - - for (int trial = 0; trial < 50; trial++) { - SCOPED_TRACE("trial=" + std::to_string(trial)); - int num_keys = 1 + rnd.Uniform(500); - int max_key_len = 1 + rnd.Uniform(30); - std::vector keys; - keys.reserve(num_keys); - for (int i = 0; i < num_keys; i++) { - keys.push_back(rnd.RandomString(1 + rnd.Uniform(max_key_len))); - } - std::sort(keys.begin(), keys.end()); - keys.erase(std::unique(keys.begin(), keys.end()), keys.end()); - if (keys.empty()) { - continue; - } - ASSERT_NO_FATAL_FAILURE(VerifyTrieIteration(keys)); - } -} - -// --- Basic builder and trie construction --- - -TEST_F(LoudsTrieTest, EmptyTrie) { - LoudsTrieBuilder builder; - builder.Finish(); - Slice data = builder.GetSerializedData(); - ASSERT_GT(data.size(), 0u); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), 0u); -} - -TEST_F(LoudsTrieTest, SingleKey) { - LoudsTrieBuilder builder; - TrieBlockHandle h{100, 200}; - builder.AddKey(Slice("hello"), h); - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), 1u); - - // Verify handle. - auto handle = trie.GetHandle(0); - ASSERT_EQ(handle.offset, 100u); - ASSERT_EQ(handle.size, 200u); -} - -TEST_F(LoudsTrieTest, MultipleKeys) { - LoudsTrieBuilder builder; - std::vector keys = {"apple", "application", "banana", "band", - "cat"}; - for (size_t i = 0; i < keys.size(); i++) { - TrieBlockHandle h{i * 100, 50}; - builder.AddKey(Slice(keys[i]), h); - } - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), keys.size()); - - // Verify all handles via the iterator (which computes correct leaf indices). - LoudsTrieIterator iter(&trie); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Seek(Slice(keys[i]))); - ASSERT_TRUE(iter.Valid()); - auto handle = iter.Value(); - ASSERT_EQ(handle.offset, i * 100) << "Handle mismatch for key " << keys[i]; - ASSERT_EQ(handle.size, 50u); - } -} - -TEST_F(LoudsTrieTest, SharedPrefixKeys) { - // Keys with long shared prefixes -- this exercises the trie's prefix - // compression. - LoudsTrieBuilder builder; - std::vector keys; - for (int i = 0; i < 100; i++) { - std::string key = "common_prefix_" + std::to_string(i); - keys.push_back(key); - } - // Keys are already sorted lexicographically for single-digit vs double-digit - // numbers, but we need proper sorting. - std::sort(keys.begin(), keys.end()); - - for (size_t i = 0; i < keys.size(); i++) { - TrieBlockHandle h{i * 100, 50}; - builder.AddKey(Slice(keys[i]), h); - } - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), keys.size()); -} - -// --- Block handle encoding --- - -TEST_F(LoudsTrieTest, HandleRoundTrip) { - // Verify that block handles round-trip correctly through serialization. - // Handles are stored as fixed-width uint64 arrays for zero-copy loading. - LoudsTrieBuilder builder; - - // Simulate realistic SST block handles: sequential offsets with ~4KB blocks. - const uint64_t kBlockSize = 4096; - const int kNumBlocks = 100; - std::vector keys; - std::vector expected_handles; - - for (int i = 0; i < kNumBlocks; i++) { - char buf[32]; - snprintf(buf, sizeof(buf), "key_%04d", i); - keys.emplace_back(buf); - TrieBlockHandle h{static_cast(i) * kBlockSize, kBlockSize}; - expected_handles.emplace_back(h); - builder.AddKey(Slice(keys.back()), h); - } - builder.Finish(); - Slice data = builder.GetSerializedData(); - - // Handles are stored as two fixed-width uint64 arrays (offsets + sizes), - // totaling 16 bytes per handle (100 * 16 = 1600 bytes for handles). - // This trades space for O(1) zero-copy loading. - size_t fixed_handle_bytes = kNumBlocks * 16; - // Verify that handles deserialize correctly. - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), static_cast(kNumBlocks)); - - // Verify all handles round-trip correctly. - LoudsTrieIterator iter(&trie); - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (int i = 0; i < kNumBlocks; i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - auto h = iter.Value(); - ASSERT_EQ(h.offset, expected_handles[i].offset) - << "Offset mismatch at i=" << i; - ASSERT_EQ(h.size, expected_handles[i].size) << "Size mismatch at i=" << i; - if (i < kNumBlocks - 1) { - ASSERT_TRUE(iter.Next()); - } - } - - // Log for informational purposes. - fprintf(stderr, - "HandleRoundTrip: total serialized=%zu, " - "handle arrays=%zu bytes\n", - data.size(), fixed_handle_bytes); - (void)fixed_handle_bytes; -} - -TEST_F(LoudsTrieTest, PrefixKeyHandleReorderVerification) { - // Explicitly verify that BFS handle reordering is correct with prefix keys. - // Keys: "a"=0, "ab"=1, "b"=2 - // BFS leaf order: "b" first (leaf at root), then "a" (prefix at level 1), - // then "ab" (leaf at level 2). - // So handles should be: [2, 0, 1] in BFS order. - LoudsTrieBuilder builder; - builder.AddKey(Slice("a"), {100, 10}); - builder.AddKey(Slice("ab"), {200, 20}); - builder.AddKey(Slice("b"), {300, 30}); - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - - LoudsTrieIterator iter(&trie); - - // Seek to "a" -- should find "a" with handle {100, 10}. - ASSERT_TRUE(iter.Seek(Slice("a"))); - ASSERT_EQ(iter.Key().ToString(), "a"); - ASSERT_EQ(iter.Value().offset, 100u); - ASSERT_EQ(iter.Value().size, 10u); - - // Next should give "ab" with handle {200, 20}. - ASSERT_TRUE(iter.Next()); - ASSERT_EQ(iter.Key().ToString(), "ab"); - ASSERT_EQ(iter.Value().offset, 200u); - ASSERT_EQ(iter.Value().size, 20u); - - // Next should give "b" with handle {300, 30}. - ASSERT_TRUE(iter.Next()); - ASSERT_EQ(iter.Key().ToString(), "b"); - ASSERT_EQ(iter.Value().offset, 300u); - ASSERT_EQ(iter.Value().size, 30u); - - ASSERT_FALSE(iter.Next()); -} - -TEST_F(LoudsTrieTest, PrefixKeySeekBetween) { - // Seek to keys between a prefix key and its child. - // "a" < target "aa" < "ab" - LoudsTrieBuilder builder; - builder.AddKey(Slice("a"), {100, 10}); - builder.AddKey(Slice("ab"), {200, 20}); - builder.AddKey(Slice("b"), {300, 30}); - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - LoudsTrieIterator iter(&trie); - - // "aa" is between "a" (prefix key) and "ab" -> should land on "ab". - ASSERT_TRUE(iter.Seek(Slice("aa"))); - ASSERT_EQ(iter.Key().ToString(), "ab"); - - // "a" should hit the prefix key exactly. - ASSERT_TRUE(iter.Seek(Slice("a"))); - ASSERT_EQ(iter.Key().ToString(), "a"); -} - -// --- Iterator edge cases --- - -TEST_F(LoudsTrieTest, IteratorSeekExact) { - std::vector keys = {"abc", "abd", "abe", "xyz"}; - auto bt = BuildTrieFromKeys(keys); - - // Seek to exact keys. - ASSERT_TRUE(bt.iter->Seek(Slice("abc"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Value().offset, 0u); - - ASSERT_TRUE(bt.iter->Seek(Slice("xyz"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Value().offset, 300u); -} - -TEST_F(LoudsTrieTest, IteratorSeekBetweenKeys) { - auto bt = BuildTrieFromKeys({"aaa", "bbb", "ccc", "ddd"}); - - // Seek to a key between "aaa" and "bbb" -- should land on "bbb". - ASSERT_TRUE(bt.iter->Seek(Slice("ab"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Value().offset, 100u); // "bbb"'s handle. -} - -TEST_F(LoudsTrieTest, IteratorNext) { - std::vector keys = {"a", "b", "c", "d", "e"}; - auto bt = BuildTrieFromKeys(keys); - ASSERT_NO_FATAL_FAILURE(VerifyFullScan(bt, keys)); -} - -TEST_F(LoudsTrieTest, IteratorSeekPastEnd) { - auto bt = BuildTrieFromKeys({"aaa", "bbb"}); - // Seek past all keys. - ASSERT_FALSE(bt.iter->Seek(Slice("zzz"))); - ASSERT_FALSE(bt.iter->Valid()); -} - -TEST_F(LoudsTrieTest, IteratorSeekBeforeAll) { - auto bt = BuildTrieFromKeys({"bbb", "ccc", "ddd"}); - // Seek before all keys. - ASSERT_TRUE(bt.iter->Seek(Slice("aaa"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Value().offset, 0u); // "bbb"'s handle. -} - -TEST_F(LoudsTrieTest, SeekBetweenAllPairs) { - // For a small set of keys, try seeking to every possible "between" value. - std::vector keys = {"aaa", "aab", "aac", "bbb", "ccc"}; - auto bt = BuildTrieFromKeys(keys); - - // Seek to values between consecutive keys. - struct TestCase { - std::string target; - size_t expected_idx; // index into keys, or keys.size() for past-end. - }; - std::vector cases = { - {"a", 0}, // Before first key -> first key. - {"aaa", 0}, // Exact match. - {"aaaa", 1}, // Between aaa and aab -> aab. - {"aab", 1}, // Exact match. - {"aaba", 2}, // Between aab and aac -> aac. - {"aac", 2}, // Exact match. - {"aad", 3}, // Between aac and bbb -> bbb. - {"b", 3}, // Before bbb -> bbb. - {"bbb", 3}, // Exact match. - {"bbba", 4}, // Between bbb and ccc -> ccc. - {"ccc", 4}, // Exact match. - {"ccca", 5}, // Past end. - {"zzz", 5}, // Past end. - }; - - for (const auto& tc : cases) { - bool ok = bt.iter->Seek(Slice(tc.target)); - if (tc.expected_idx >= keys.size()) { - ASSERT_FALSE(ok) << "Expected past-end for target=\"" << tc.target - << "\""; - } else { - ASSERT_TRUE(ok) << "Seek failed for target=\"" << tc.target << "\""; - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Key().ToString(), keys[tc.expected_idx]) - << "Wrong key for target=\"" << tc.target << "\"" << " expected=\"" - << keys[tc.expected_idx] << "\"" << " got=\"" - << bt.iter->Key().ToString() << "\""; - } - } -} - -TEST_F(LoudsTrieTest, EmptyTargetSeek) { - // Seeking with an empty target should return the first key. - auto bt = BuildTrieFromKeys({"aaa", "bbb", "ccc"}); - - ASSERT_TRUE(bt.iter->Seek(Slice(""))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Key().ToString(), "aaa"); -} - -TEST_F(LoudsTrieTest, IteratorReSeekAfterInvalidation) { - // After iterator becomes invalid (past end), re-seeking should work. - auto bt = BuildTrieFromKeys({"aaa", "bbb", "ccc"}); - - // Go past end. - ASSERT_FALSE(bt.iter->Seek(Slice("zzz"))); - ASSERT_FALSE(bt.iter->Valid()); - - // Re-seek to valid key. - ASSERT_TRUE(bt.iter->Seek(Slice("aaa"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Key().ToString(), "aaa"); -} - -TEST_F(LoudsTrieTest, IteratorNextOnInvalid) { - // Calling Next() on an invalid iterator should return false. - auto bt = BuildTrieFromKeys({"aaa"}); - ASSERT_FALSE(bt.iter->Valid()); - ASSERT_FALSE(bt.iter->Next()); -} - -TEST_F(LoudsTrieTest, IteratorSeekEmptyStringVariousKeys) { - // Seek("") should always return the first key. - auto bt = BuildTrieFromKeys({"x", "xy", "xyz"}); - - ASSERT_TRUE(bt.iter->Seek(Slice(""))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Key().ToString(), "x"); -} - -TEST_F(LoudsTrieTest, IteratorMultipleSeeksDescending) { - // Seek to keys in reverse order (each seek resets state). - std::vector keys = {"aaa", "bbb", "ccc", "ddd"}; - auto bt = BuildTrieFromKeys(keys); - - for (int i = static_cast(keys.size()) - 1; i >= 0; i--) { - ASSERT_TRUE(bt.iter->Seek(Slice(keys[i]))); - ASSERT_EQ(bt.iter->Key().ToString(), keys[i]); - ASSERT_EQ(bt.iter->Value().offset, static_cast(i) * 100); - } -} - -TEST_F(LoudsTrieTest, IteratorSeekToLongerThanAnyKey) { - // Target key is longer than any key in the trie. - // The trie has "ab" and "cd". Seeking "abx" should find "cd" (next after - // "ab"). - auto bt = BuildTrieFromKeys({"ab", "cd"}); - - ASSERT_TRUE(bt.iter->Seek(Slice("abx"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Key().ToString(), "cd"); -} - -TEST_F(LoudsTrieTest, IteratorSeekTargetIsPrefixOfKey) { - // Target "ab" is a prefix of trie key "abc". Should find "abc". - auto bt = BuildTrieFromKeys({"abc", "abd"}); - - ASSERT_TRUE(bt.iter->Seek(Slice("ab"))); - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->Key().ToString(), "abc"); -} - -TEST_F(LoudsTrieTest, IteratorFullScanThenReSeek) { - // Full scan with Next() until end, then re-seek to middle. - auto bt = BuildTrieFromKeys({"aa", "bb", "cc", "dd", "ee"}); - - // Scan all. - ASSERT_TRUE(bt.iter->Seek(Slice("aa"))); - int count = 0; - while (bt.iter->Valid()) { - count++; - bt.iter->Next(); - } - ASSERT_EQ(count, 5); - ASSERT_FALSE(bt.iter->Valid()); - - // Re-seek to middle. - ASSERT_TRUE(bt.iter->Seek(Slice("cc"))); - ASSERT_EQ(bt.iter->Key().ToString(), "cc"); - ASSERT_EQ(bt.iter->Value().offset, 200u); -} - -// --- Serialization / deserialization --- - -TEST_F(LoudsTrieTest, InitFromDataCorruption) { - // Truncated data should return corruption status. - LoudsTrie trie; - ASSERT_TRUE(trie.InitFromData(Slice("short")).IsCorruption()); - - // Bad magic number. - std::string bad(56, '\0'); - bad[0] = 'X'; // Wrong magic. - ASSERT_TRUE(trie.InitFromData(Slice(bad)).IsCorruption()); -} - -TEST_F(LoudsTrieTest, InitFromDataMaxDepthCorruption) { - // A corrupted SST could set max_depth_ = UINT32_MAX, which would cause - // an integer overflow in LoudsTrieIterator (key_cap_ = MaxDepth() + 1 - // wraps to 0). Verify that InitFromData rejects this. - // - // Build the header manually: magic(4) + version(4) + num_keys(8) + - // cutoff_level(4) + max_depth(4) + dense_leaf_count(8) + - // dense_node_count(8) + dense_child_count(8) + flags(4) + - // reserved(4) = 56 bytes. - std::string header; - uint32_t magic = 0x54524945; // "TRIE" - uint32_t version = 1; - uint64_t num_keys = 0; - uint32_t cutoff_level = 0; - uint32_t max_depth = UINT32_MAX; // Corrupt value - uint64_t zeros = 0; - uint32_t zero32 = 0; - - header.append(reinterpret_cast(&magic), 4); - header.append(reinterpret_cast(&version), 4); - header.append(reinterpret_cast(&num_keys), 8); - header.append(reinterpret_cast(&cutoff_level), 4); - header.append(reinterpret_cast(&max_depth), 4); - header.append(reinterpret_cast(&zeros), 8); // dense_leaf_count - header.append(reinterpret_cast(&zeros), 8); // dense_node_count - header.append(reinterpret_cast(&zeros), 8); // dense_child_count - header.append(reinterpret_cast(&zero32), 4); // flags - header.append(reinterpret_cast(&zero32), 4); // reserved - // Pad with enough data to pass subsequent parsing. - header.append(256, '\0'); - - LoudsTrie trie2; - ASSERT_TRUE(trie2.InitFromData(Slice(header)).IsCorruption()); - - // Also test a value just above the limit (65537). - max_depth = 65537; - std::string header2; - header2.append(reinterpret_cast(&magic), 4); - header2.append(reinterpret_cast(&version), 4); - header2.append(reinterpret_cast(&num_keys), 8); - header2.append(reinterpret_cast(&cutoff_level), 4); - header2.append(reinterpret_cast(&max_depth), 4); - header2.append(reinterpret_cast(&zeros), 8); - header2.append(reinterpret_cast(&zeros), 8); - header2.append(reinterpret_cast(&zeros), 8); - header2.append(reinterpret_cast(&zero32), 4); // flags - header2.append(reinterpret_cast(&zero32), 4); // reserved - header2.append(256, '\0'); - - LoudsTrie trie3; - ASSERT_TRUE(trie3.InitFromData(Slice(header2)).IsCorruption()); -} - -TEST_F(LoudsTrieTest, MoveConstructor) { - // Verify that move constructor works correctly for LoudsTrie, which - // contains Bitvector members with RecomputePointers() logic. - std::vector keys = {"apple", "banana", "cherry", "date"}; - auto bt = BuildTrieFromKeys(keys); - ASSERT_EQ(bt.trie.NumKeys(), keys.size()); - - // Move-construct a new trie. - LoudsTrie moved(std::move(bt.trie)); - ASSERT_EQ(moved.NumKeys(), keys.size()); - - // Verify iteration works on the moved trie. - LoudsTrieIterator iter(&moved); - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - ASSERT_EQ(iter.Value().offset, i * 100); - if (i < keys.size() - 1) { - ASSERT_TRUE(iter.Next()); - } - } - ASSERT_FALSE(iter.Next()); -} - -TEST_F(LoudsTrieTest, SerializeDeserializeRoundTrip) { - // Build a non-trivial trie and verify that serialization round-trips. - std::vector keys = {"alpha", "beta", "gamma", "delta"}; - std::sort(keys.begin(), keys.end()); - auto bt = BuildTrieFromKeys(keys); - ASSERT_EQ(bt.trie.NumKeys(), keys.size()); - ASSERT_NO_FATAL_FAILURE(VerifyFullScan(bt, keys)); -} - -TEST_F(LoudsTrieTest, SerializeDeserializeRoundTripMisalignedData) { - // Verify that LoudsTrie::InitFromData handles non-8-byte-aligned data. - // This can happen when the SST block is read from mmap at an unaligned - // file offset. - std::vector keys = {"alpha", "beta", "gamma", "delta"}; - std::sort(keys.begin(), keys.end()); - auto bt = BuildTrieFromKeys(keys); - Slice data = bt.builder.GetSerializedData(); - - // Create a buffer with 1-byte offset to guarantee misalignment. - std::string padded; - padded.reserve(1 + data.size()); - padded.push_back('\0'); // 1-byte padding - padded.append(data.data(), data.size()); - const char* misaligned_ptr = padded.data() + 1; - ASSERT_NE(reinterpret_cast(misaligned_ptr) % 8, 0u); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(Slice(misaligned_ptr, data.size()))); - ASSERT_EQ(trie.NumKeys(), keys.size()); - - // Verify all keys and handles via iteration. - LoudsTrieIterator iter(&trie); - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), keys[i]); - ASSERT_EQ(iter.Value().offset, i * 100); - if (i < keys.size() - 1) { - iter.Next(); - } - } -} - -TEST_F(LoudsTrieTest, StressTest10KKeys) { - // 10K keys with uniform distribution. - std::vector keys; - for (int i = 0; i < 10000; i++) { - char buf[32]; - snprintf(buf, sizeof(buf), "key_%08d", i); - keys.emplace_back(buf); - } - std::sort(keys.begin(), keys.end()); - auto bt = BuildTrieFromKeys(keys); - ASSERT_EQ(bt.trie.NumKeys(), keys.size()); - - // Verify every 100th key via Seek. - for (size_t i = 0; i < keys.size(); i += 100) { - ASSERT_TRUE(bt.iter->Seek(Slice(keys[i]))) - << "Seek failed for key[" << i << "]"; - ASSERT_EQ(bt.iter->Key().ToString(), keys[i]); - ASSERT_EQ(bt.iter->Value().offset, i * 100); - } - - // Full forward scan. - ASSERT_NO_FATAL_FAILURE(VerifyFullScan(bt, keys)); -} - -TEST_F(LoudsTrieTest, MoveAssignment) { - // Test move assignment operator (move constructor was already tested). - std::vector keys = {"alpha", "beta", "gamma", "delta"}; - std::sort(keys.begin(), keys.end()); - auto bt = BuildTrieFromKeys(keys); - - // Copy serialized data so it outlives the builder. - std::string data_copy(bt.builder.GetSerializedData().data(), - bt.builder.GetSerializedData().size()); - LoudsTrie trie1; - ASSERT_OK(trie1.InitFromData(Slice(data_copy))); - ASSERT_EQ(trie1.NumKeys(), keys.size()); - - // Move-assign trie1 into a fresh trie. - LoudsTrie trie2; - trie2 = std::move(trie1); - ASSERT_EQ(trie2.NumKeys(), keys.size()); - - // Verify the moved-to trie works correctly. - LoudsTrieIterator iter(&trie2); - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - ASSERT_EQ(iter.Value().offset, i * 100); - if (i < keys.size() - 1) { - ASSERT_TRUE(iter.Next()); - } - } - ASSERT_FALSE(iter.Next()); -} - -TEST_F(LoudsTrieTest, CutoffLevelAndMaxDepthAndHasChains) { - // Verify CutoffLevel(), MaxDepth(), and HasChains() return sensible values. - std::vector keys = {"a", "ab", "abc", "abcd", "b", "bc"}; - std::sort(keys.begin(), keys.end()); - auto bt = BuildTrieFromKeys(keys); - auto& trie = bt.trie; - - // MaxDepth should equal the longest key length. - uint32_t max_key_len = 0; - for (const auto& k : keys) { - max_key_len = std::max(max_key_len, static_cast(k.size())); - } - ASSERT_EQ(trie.MaxDepth(), max_key_len); - - // CutoffLevel should be within [0, max_depth]. - ASSERT_LE(trie.CutoffLevel(), trie.MaxDepth()); -} - -TEST_F(LoudsTrieTest, LeafIndex) { - // Verify LeafIndex() returns correct indices during sequential iteration. - std::vector keys = {"cat", "cow", "dog", "fox"}; - auto bt = BuildTrieFromKeys(keys); - - ASSERT_TRUE(bt.iter->Seek(Slice(keys[0]))); - - // LeafIndex should be sequential for sequential keys. - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(bt.iter->Valid()); - ASSERT_EQ(bt.iter->LeafIndex(), i) - << "LeafIndex mismatch at key=" << keys[i]; - ASSERT_EQ(bt.iter->Value().offset, i * 100); - if (i < keys.size() - 1) { - ASSERT_TRUE(bt.iter->Next()); - } - } -} - -TEST_F(LoudsTrieTest, EmptyTrieIterator) { - // Verify iterator behavior on an empty trie. - LoudsTrieBuilder builder; - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - ASSERT_EQ(trie.NumKeys(), 0u); - - LoudsTrieIterator iter(&trie); - ASSERT_FALSE(iter.Seek(Slice("anything"))); - ASSERT_FALSE(iter.Valid()); - ASSERT_FALSE(iter.Next()); - ASSERT_FALSE(iter.Valid()); -} - -TEST_F(LoudsTrieTest, ApproximateAuxMemoryUsage) { - // Verify that ApproximateAuxMemoryUsage() returns a reasonable value - // for a trie with sparse internal nodes. - std::vector keys; - for (int i = 0; i < 200; i++) { - char buf[32]; - snprintf(buf, sizeof(buf), "key_%04d", i); - keys.emplace_back(buf); - } - auto bt = BuildTrieFromKeys(keys); - - // Aux memory should be > 0 for a trie with sparse internal nodes - // (child position lookup tables are allocated). - size_t aux_mem = bt.trie.ApproximateAuxMemoryUsage(); - ASSERT_GT(aux_mem, 0u); -} - -TEST_F(LoudsTrieTest, InitFromDataUnsupportedVersion) { - // Build a valid trie, then patch the format version to a bad value. - std::vector keys = {"a", "b", "c"}; - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - builder.AddKey(Slice(keys[i]), {i * 100, 50}); - } - builder.Finish(); - std::string data(builder.GetSerializedData().data(), - builder.GetSerializedData().size()); - - // Version is at offset 4 (after the 4-byte magic). - uint32_t bad_version = 99; - memcpy(&data[4], &bad_version, sizeof(uint32_t)); - - LoudsTrie trie; - Status s = trie.InitFromData(Slice(data)); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); - ASSERT_NE(s.ToString().find("unsupported format version"), std::string::npos); -} - -TEST_F(LoudsTrieTest, InitFromDataProgressiveTruncation) { - // Build a valid trie (no chains), serialize it, then try InitFromData with - // every length from 0 to data.size()-1 and verify all return Corruption. - // This covers ALL truncation error paths in InitFromData in a single test. - std::vector keys; - for (int i = 0; i < 20; i++) { - char buf[16]; - snprintf(buf, sizeof(buf), "k%02d", i); - keys.emplace_back(buf); - } - - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - builder.AddKey(Slice(keys[i]), {i * 100, 50}); - } - builder.Finish(); - std::string data(builder.GetSerializedData().data(), - builder.GetSerializedData().size()); - - // Verify the full data works. - { - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(Slice(data))); - } - - // Every truncated length should fail with Corruption. - for (size_t len = 0; len < data.size(); len++) { - LoudsTrie trie; - Status s = trie.InitFromData(Slice(data.data(), len)); - ASSERT_TRUE(s.IsCorruption()) << "Expected Corruption at len=" << len << "/" - << data.size() << ", got: " << s.ToString(); - } -} - -TEST_F(LoudsTrieTest, InitFromDataProgressiveTruncationWithChains) { - // Same as above, but with keys that produce path compression chains. - // Keys share a 10-char prefix "AAAAAAAAAA" + varying suffix, ensuring - // the chain suffix is >= kMinChainLength (8). - std::vector keys; - std::string prefix(10, 'A'); - for (int i = 0; i < 20; i++) { - char suffix[8]; - snprintf(suffix, sizeof(suffix), "%02d", i); - keys.push_back(prefix + suffix); - } - std::sort(keys.begin(), keys.end()); - - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - builder.AddKey(Slice(keys[i]), {i * 100, 50}); - } - builder.Finish(); - std::string data(builder.GetSerializedData().data(), - builder.GetSerializedData().size()); - - // Verify the full data works and has chains. - { - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(Slice(data))); - ASSERT_TRUE(trie.HasChains()); - } - - // Every truncated length should fail with Corruption. - for (size_t len = 0; len < data.size(); len++) { - LoudsTrie trie; - Status s = trie.InitFromData(Slice(data.data(), len)); - ASSERT_TRUE(s.IsCorruption()) << "Expected Corruption at len=" << len << "/" - << data.size() << ", got: " << s.ToString(); - } } -TEST_F(LoudsTrieTest, SparseBinarySearchPath) { - // Create a trie where the root node has >16 children in the sparse level, - // which forces SparseSeekLabel to use std::lower_bound instead of linear - // scan (kLinearScanThreshold = 16). - // Strategy: 20 two-byte keys with distinct first bytes. With cutoff=0 - // (all sparse), the root has 20 labels. - std::vector keys; - // Use bytes 'A' through 'T' (20 distinct first bytes). - for (int i = 0; i < 20; i++) { - std::string key(1, static_cast('A' + i)); - key += 'x'; // Second byte to ensure sparse leaf. - keys.push_back(key); - } - std::sort(keys.begin(), keys.end()); - - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - builder.AddKey(Slice(keys[i]), {i * 100, 50}); - } - builder.Finish(); - Slice data = builder.GetSerializedData(); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(data)); - - LoudsTrieIterator iter(&trie); - - // Seek to each key -- exercises binary search in the 20-label root node. - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Seek(Slice(keys[i]))) - << "Seek failed for key[" << i << "]"; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - ASSERT_EQ(iter.Value().offset, i * 100); - } - - // Seek to a key between existing labels (binary search inexact match). - // 'A' + 0.5 doesn't exist -- seek "AAx" which is between "Ax" and "Bx". - ASSERT_TRUE(iter.Seek(Slice("AAx"))); - // Should land on "Ax" or "Bx" -- the first key >= "AAx". - ASSERT_TRUE(iter.Valid()); - std::string result_key = iter.Key().ToString(); - ASSERT_GE(result_key, "AAx"); - - // Seek past all keys. - ASSERT_FALSE(iter.Seek(Slice("Zx"))); - ASSERT_FALSE(iter.Valid()); - - // Full forward scan to verify trie integrity. - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - if (i < keys.size() - 1) { - ASSERT_TRUE(iter.Next()); - } +// A copy of data with one to three random bytes changed. +std::string FlipRandomBytes(std::string data, Random* rnd) { + const int flips = 1 + rnd->Uniform(3); + for (int f = 0; f < flips; f++) { + data[rnd->Uniform(static_cast(data.size()))] ^= + static_cast(1 + rnd->Uniform(255)); } - ASSERT_FALSE(iter.Next()); + return data; } -TEST_F(LoudsTrieTest, ChainSeekVariousTargets) { - // Build a trie with keys that create path compression chains, then seek to - // various targets that exercise different chain-matching code paths: - // 1. Full chain match (chain ends at internal node with fanout > 1) - // 2. Chain mismatch: target < chain suffix - // 3. Chain mismatch: target > chain suffix - // 4. Target runs out before chain (cmp < 0) - // 5. Target runs out before chain (cmp > 0) - // 6. Target runs out before chain (cmp == 0, mid-chain consumption) - // 7. target_remaining == 0 (target consumed at parent) - // - // Keys: all share a long prefix that triggers chain compression. - // The prefix "AAAAAAAAA" (9 chars) + varying 2-char suffixes. - // The chain suffix must be >= kMinChainLength (8 bytes). - std::vector keys; - std::string prefix(9, 'A'); - // Add keys with suffixes that force a fanout > 1 after the chain. - keys.push_back(prefix + "Ma"); - keys.push_back(prefix + "Mb"); - keys.push_back(prefix + "Mc"); - keys.push_back(prefix + "Na"); - keys.push_back(prefix + "Nb"); - std::sort(keys.begin(), keys.end()); - - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - builder.AddKey(Slice(keys[i]), {i * 100, 50}); - } - builder.Finish(); - std::string data(builder.GetSerializedData().data(), - builder.GetSerializedData().size()); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(Slice(data))); - - LoudsTrieIterator iter(&trie); - - // 1. Exact match on each key (full chain match, internal node after chain). - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Seek(Slice(keys[i]))) - << "Seek failed for key[" << i << "]=\"" << keys[i] << "\""; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - ASSERT_EQ(iter.Value().offset, i * 100); - } - - // 2. Chain mismatch: target < chain suffix. - // Seek to prefix + "L" -- 'L' < 'M', so chain mismatch with target < suffix. - // Should land on the first key (prefix + "Ma"). - { - std::string target = prefix + "La"; - ASSERT_TRUE(iter.Seek(Slice(target))); - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), keys[0]); // "AAAAAAAAA" + "Ma" - } - - // 3. Chain mismatch: target > chain suffix. - // Seek to prefix + "Md" -- 'd' > 'c' after chain, should advance past "Mc". - { - std::string target = prefix + "Md"; - ASSERT_TRUE(iter.Seek(Slice(target))); - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), prefix + "Na"); - } - - // 4. Target runs out before chain (shorter target < chain prefix). - // Seek to prefix + "L" (shorter than chain). 'L' < 'M' -> target < chain. - { - std::string target = prefix + "L"; - ASSERT_TRUE(iter.Seek(Slice(target))); - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), keys[0]); - } - - // 5. Target runs out before chain (shorter target > chain prefix character). - // Seek to prefix + "O" -- 'O' > 'N' (the highest chain prefix). Should - // advance past all keys. - { - std::string target = prefix + "O"; - ASSERT_FALSE(iter.Seek(Slice(target))); - ASSERT_FALSE(iter.Valid()); - } - - // 6. Target consumed exactly at the parent of the chain (target_remaining - // == 0). Seek to just the prefix "AAAAAAAAA" -- shorter than any key. - // All keys are > prefix, so should land on the first key. - { - ASSERT_TRUE(iter.Seek(Slice(prefix))); - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), keys[0]); - } - - // 7. Full forward scan to verify trie integrity after all seeks. - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - if (i < keys.size() - 1) { - ASSERT_TRUE(iter.Next()); - } +// Reading a malformed index may fail only with Corruption. +bool OkOrCorruption(const Status& s) { + if (!s.ok()) { + EXPECT_EQ(s.code(), Status::Code::kCorruption) << s.ToString(); } - ASSERT_FALSE(iter.Next()); + return s.ok(); } -TEST_F(LoudsTrieTest, ChainSeekEndAtLeaf) { - // Test path compression chains that end at a leaf node (end_child_idx == - // UINT32_MAX). This requires a key whose suffix after the chain endpoint - // has no children -- i.e., the chain is the entire remaining path. - // - // Strategy: Create keys with a long shared prefix and a single-character - // difference at the very end, so the chain covers most of the key and - // terminates at a leaf. - // - // "AAAAAAAAA" (9 chars prefix) + "B" = 10-char key. The "AAAAAAAA" portion - // (8 bytes) forms the chain suffix (>= kMinChainLength). The 'B' is the - // label at the chain-ending leaf. - // - // We also add keys with different first characters to ensure the trie has - // enough structure for the budget formula. - std::vector keys; - // Main chain key. - keys.push_back(std::string(9, 'A') + "B"); - // Additional keys with a different prefix to provide sparse-level context. - for (int i = 0; i < 5; i++) { - std::string key(1, static_cast('C' + i)); - key += "x"; - keys.push_back(key); - } - std::sort(keys.begin(), keys.end()); - - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - builder.AddKey(Slice(keys[i]), {i * 100, 50}); - } - builder.Finish(); - std::string data(builder.GetSerializedData().data(), - builder.GetSerializedData().size()); - - LoudsTrie trie; - ASSERT_OK(trie.InitFromData(Slice(data))); - - LoudsTrieIterator iter(&trie); - - // Exact seek to the chain-ending-at-leaf key. - std::string chain_key = std::string(9, 'A') + "B"; - ASSERT_TRUE(iter.Seek(Slice(chain_key))); - ASSERT_TRUE(iter.Valid()); - ASSERT_EQ(iter.Key().ToString(), chain_key); - - // Seek to a target longer than the chain key -- target has more bytes but - // trie key ends at leaf. Should Advance() past it. - std::string longer_target = chain_key + "Z"; - ASSERT_TRUE(iter.Seek(Slice(longer_target))); - ASSERT_TRUE(iter.Valid()); - // Should land on the next key after the chain key. - ASSERT_GT(iter.Key().ToString(), chain_key); - - // Seek to a target that is the chain key prefix (target consumed mid-chain). - std::string mid_chain = std::string(5, 'A'); - ASSERT_TRUE(iter.Seek(Slice(mid_chain))); - ASSERT_TRUE(iter.Valid()); - // First key >= "AAAAA" is "AAAAAAAAAB". - ASSERT_EQ(iter.Key().ToString(), chain_key); - - // Full forward scan to verify integrity. - ASSERT_TRUE(iter.Seek(Slice(keys[0]))); - for (size_t i = 0; i < keys.size(); i++) { - ASSERT_TRUE(iter.Valid()) << "Invalid at i=" << i; - ASSERT_EQ(iter.Key().ToString(), keys[i]); - if (i < keys.size() - 1) { - ASSERT_TRUE(iter.Next()); - } - } - ASSERT_FALSE(iter.Next()); +// The same for opening an index, which also reports NotSupported when the +// version field changed. +bool OpenedOrRejected(const Status& s) { + return s.IsNotSupported() ? false : OkOrCorruption(s); } // ============================================================================ @@ -2442,8 +178,21 @@ class TrieIndexFactoryTest : public testing::Test { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kUnknown); } - void AssertSerializedSizeBound(const std::vector& keys, - bool parallel, uint64_t max_size_ratio) { + // Blocks of 100 bytes laid out back to back. + static std::vector UniformHandles( + size_t n) { + std::vector handles; + handles.reserve(n); + for (size_t i = 0; i < n; ++i) { + handles.push_back({i * 100, 100}); + } + return handles; + } + + void AssertSerializedSizeBound( + const std::vector& keys, + const std::vector& handles, + bool parallel, uint64_t max_size_ratio) { std::unique_ptr builder; ASSERT_OK(factory_->NewBuilder(IndexFactoryOptions(), builder)); if (parallel && !builder->SupportsParallelAddEntry()) { @@ -2453,7 +202,7 @@ class TrieIndexFactoryTest : public testing::Test { const Slice key(keys[i]); const Slice next = i + 1 < keys.size() ? Slice(keys[i + 1]) : Slice(); const Slice* next_key = i + 1 < keys.size() ? &next : nullptr; - const IndexFactoryBuilder::BlockHandle handle{i * 100, 100}; + const IndexFactoryBuilder::BlockHandle& handle = handles[i]; const IndexFactoryBuilder::IndexEntryContext context = EntryCtx(keys.size() - i, keys.size() - i - 1); std::string scratch; @@ -2516,34 +265,103 @@ TEST_F(TrieIndexFactoryTest, EstimatedSizeBoundsSerializedTrie) { std::sort(key_set.keys.begin(), key_set.keys.end()); for (bool parallel : {false, true}) { SCOPED_TRACE(parallel); - AssertSerializedSizeBound(key_set.keys, parallel, key_set.max_size_ratio); + AssertSerializedSizeBound(key_set.keys, + UniformHandles(key_set.keys.size()), parallel, + key_set.max_size_ratio); + } + } + + // Blocks of varying size and spacing, one of them huge, as with block + // alignment padding and a very large value. + std::vector irregular; + uint64_t offset = 0; + for (size_t i = 0; i < cases[9].keys.size(); ++i) { + const uint64_t size = + i == 1000 ? uint64_t{1} << 30 : 1 + random.Uniform(8000); + irregular.push_back({offset, size}); + offset += size + random.Uniform(4096); + } + for (bool parallel : {false, true}) { + SCOPED_TRACE(parallel); + AssertSerializedSizeBound(cases[9].keys, irregular, parallel, 0); + } +} + +// No real SST has handles like these, but Finish() accepts them, so the +// estimate must still bound its output. +TEST_F(TrieIndexFactoryTest, EstimatedSizeBoundsExtremeHandles) { + constexpr uint64_t kMax = std::numeric_limits::max(); + std::vector keys; + for (int i = 0; i < 63; ++i) { + keys.push_back("k" + std::to_string(100 + i)); + } + // Blocks alternating between a size of 1 and `large`, the last one large, + // with equal gaps, shifted so that the last block ends at `end` modulo + // 2^64. + auto evenly_spaced = [&](uint64_t large, uint64_t gap, uint64_t end) { + std::vector handles; + uint64_t offset = 0; + for (size_t i = 0; i < keys.size(); ++i) { + const uint64_t size = (keys.size() - 1 - i) % 2 == 0 ? large : 1; + handles.push_back({offset, size}); + offset += size + gap; + } + const uint64_t shift = end - (handles.back().offset + handles.back().size); + for (IndexFactoryBuilder::BlockHandle& handle : handles) { + handle.offset += shift; + } + return handles; + }; + // Overlapping blocks with 56-bit offset and size residuals, whose largest + // step plus largest size is 2^58. + constexpr uint64_t kStep = uint64_t{1} << 55; + constexpr uint64_t kLarge = (uint64_t{1} << 58) - kStep; + std::vector wide_residuals; + for (size_t i = 0; i < keys.size(); ++i) { + wide_residuals.push_back( + {i == 0 ? 0 : kStep + i - 1, i % 2 == 0 ? kLarge : kLarge - kStep}); + } + const struct { + const char* name; + std::vector handles; + } cases[] = { + {"wide_residuals", wide_residuals}, + {"last_block_ends_past_max", evenly_spaced(uint64_t{1} << 50, 5, 1)}, + {"sentinel_past_max", evenly_spaced(uint64_t{1} << 50, 100, kMax - 50)}, + }; + for (const auto& c : cases) { + SCOPED_TRACE(c.name); + for (bool parallel : {false, true}) { + SCOPED_TRACE(parallel); + AssertSerializedSizeBound(keys, c.handles, parallel, 0); } } } -TEST_F(TrieIndexFactoryTest, BlockHandleEncodingLimit) { +TEST_F(TrieIndexFactoryTest, BlockHandlesPast32Bits) { constexpr uint64_t kLimit = UINT32_MAX; for (bool parallel : {false, true}) { for (bool large_offset : {false, true}) { - for (size_t invalid_entry : {0U, 1U, 2U}) { - SCOPED_TRACE(testing::Message() << parallel << " " << large_offset - << " " << invalid_entry); + for (size_t large_entry : {0U, 1U, 2U}) { + SCOPED_TRACE(testing::Message() + << parallel << " " << large_offset << " " << large_entry); std::unique_ptr builder; ASSERT_OK(factory_->NewBuilder(IndexFactoryOptions(), builder)); if (parallel && !builder->SupportsParallelAddEntry()) { continue; } - // Identical separators exercise both primary and overflow handles. + // Both blocks share one separator, so they form a run. const Slice key("same"); + std::vector handles; for (size_t i = 0; i < 2; ++i) { IndexFactoryBuilder::BlockHandle handle{kLimit, kLimit}; - if (i == invalid_entry) { - if (large_offset) { - ++handle.offset; - } else { - ++handle.size; - } + // Offsets stay in block order. + if (large_offset && i >= large_entry) { + ++handle.offset; + } else if (!large_offset && i == large_entry) { + ++handle.size; } + handles.push_back(handle); std::string scratch; const Slice* next = i == 0 ? &key : nullptr; const auto context = EntryCtx(200 - i * 100, 100); @@ -2556,24 +374,17 @@ TEST_F(TrieIndexFactoryTest, BlockHandleEncodingLimit) { } } Slice contents; - Status s = builder->Finish(&contents); - if (invalid_entry < 2) { - EXPECT_TRUE(s.IsNotSupported()) << s.ToString(); - EXPECT_TRUE(contents.empty()); - } else { - ASSERT_OK(s); - std::unique_ptr reader; - ASSERT_OK( - factory_->NewReader(IndexFactoryOptions(), contents, reader)); - auto iter = reader->NewIterator(ReadOptions()); - IterateResult result; - ASSERT_OK(iter->SeekToFirstAndGetResult(&result)); - EXPECT_EQ(iter->value().offset, kLimit); - EXPECT_EQ(iter->value().size, kLimit); - ASSERT_OK(iter->NextAndGetResult(&result)); - EXPECT_EQ(iter->value().offset, kLimit); - EXPECT_EQ(iter->value().size, kLimit); - } + ASSERT_OK(builder->Finish(&contents)); + std::unique_ptr reader; + ASSERT_OK(factory_->NewReader(IndexFactoryOptions(), contents, reader)); + auto iter = reader->NewIterator(ReadOptions()); + IterateResult result; + ASSERT_OK(iter->SeekToFirstAndGetResult(&result)); + EXPECT_EQ(iter->value().offset, handles[0].offset); + EXPECT_EQ(iter->value().size, handles[0].size); + ASSERT_OK(iter->NextAndGetResult(&result)); + EXPECT_EQ(iter->value().offset, handles[1].offset); + EXPECT_EQ(iter->value().size, handles[1].size); } } } @@ -2627,10 +438,57 @@ TEST_F(TrieIndexFactoryTest, BasicBuildAndRead) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); } +TEST_F(TrieIndexFactoryTest, EmptySeparator) { + // The first block holds only the empty key. + auto ctx = BuildTrieAndGetIterator({{"", "b", 0, 95, 10, 9}, + {"b", "c", 100, 95, 9, 8}, + {"c", "", 200, 95, 8, 0}}); + ASSERT_NO_FATAL_FAILURE( + AssertFullForwardScan(ctx.iter.get(), Slice(""), {0, 100, 200})); + ASSERT_NO_FATAL_FAILURE( + AssertSeekOffset(ctx.iter.get(), "", kMaxSequenceNumber, 0)); + ASSERT_NO_FATAL_FAILURE( + AssertSeekOffset(ctx.iter.get(), "a", kMaxSequenceNumber, 100)); + ASSERT_NO_FATAL_FAILURE( + AssertSeekOffset(ctx.iter.get(), "b", kMaxSequenceNumber, 100)); + IterateResult result; + ASSERT_OK(ctx.iter->SeekToLastAndGetResult(&result)); + ASSERT_EQ(ctx.iter->value().offset, 200u); + ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); + ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); + ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); + ASSERT_EQ(result.key, ""); + ASSERT_EQ(ctx.iter->value().offset, 0u); +} + TEST_F(TrieIndexFactoryTest, FactoryName) { ASSERT_STREQ(factory_->Name(), "trie_index"); } +// An OPTIONS file names the factory, so it must be reachable by that name and +// survive a round trip through the table options string. +TEST_F(TrieIndexFactoryTest, CreateFromName) { + ConfigOptions config_options; + std::shared_ptr factory; + ASSERT_OK( + IndexFactory::CreateFromString(config_options, "trie_index", &factory)); + ASSERT_NE(factory, nullptr); + ASSERT_STREQ(factory->Name(), TrieIndexFactory::kClassName()); + + BlockBasedTableOptions table_options; + table_options.user_defined_index_factory = factory; + std::unique_ptr table_factory( + NewBlockBasedTableFactory(table_options)); + std::string serialized; + ASSERT_OK(table_factory->GetOptionString(config_options, &serialized)); + BlockBasedTableOptions parsed; + ASSERT_OK(GetBlockBasedTableOptionsFromString( + config_options, BlockBasedTableOptions(), serialized, &parsed)); + ASSERT_NE(parsed.user_defined_index_factory, nullptr); + ASSERT_STREQ(parsed.user_defined_index_factory->Name(), + TrieIndexFactory::kClassName()); +} + TEST_F(TrieIndexFactoryTest, EmptyIndex) { IndexFactoryOptions option; option.comparator = BytewiseComparator(); @@ -2657,7 +515,7 @@ TEST_F(TrieIndexFactoryTest, DoubleFinish) { // Second Finish should fail. Slice index_contents2; Status s = builder->Finish(&index_contents2); - ASSERT_TRUE(s.IsInvalidArgument()); + ASSERT_EQ(s.code(), Status::Code::kInvalidArgument) << s.ToString(); } TEST_F(TrieIndexFactoryTest, IteratorBoundsChecking) { @@ -2947,14 +805,14 @@ TEST_F(TrieIndexFactoryTest, RejectsNonBytewiseComparator) { option.comparator = ReverseBytewiseComparator(); std::unique_ptr builder; Status s = factory_->NewBuilder(option, builder); - ASSERT_TRUE(s.IsNotSupported()) << s.ToString(); + ASSERT_EQ(s.code(), Status::Code::kNotSupported) << s.ToString(); ASSERT_EQ(builder, nullptr); // NewReader should also reject non-bytewise comparators. Slice dummy_data; std::unique_ptr reader; s = factory_->NewReader(option, dummy_data, reader); - ASSERT_TRUE(s.IsNotSupported()) << s.ToString(); + ASSERT_EQ(s.code(), Status::Code::kNotSupported) << s.ToString(); ASSERT_EQ(reader, nullptr); // BytewiseComparator should be accepted. @@ -2972,8 +830,7 @@ TEST_F(TrieIndexFactoryTest, ApproximateMemoryUsageReportsOnlyAuxData) { std::unique_ptr udi_builder; ASSERT_OK(factory_->NewBuilder(option, udi_builder)); - // Build a non-trivial index with enough keys to produce sparse internal - // nodes (which allocate child position lookup tables). + // Build a non-trivial index from 100 blocks. for (int i = 0; i < 100; i++) { char last_buf[32]; char next_buf[32]; @@ -3026,7 +883,7 @@ TEST_F(TrieIndexFactoryTest, EmptyTrieIterator) { IterateResult result; ASSERT_OK(iter->SeekAndGetResult(Slice("anything"), &result, SeekCtx(kMaxSequenceNumber))); - // kUnknown: no leaf has a key >= target, so the target is past all blocks + // kUnknown: no separator is >= target, so the target is past all blocks // in this SST. We return kUnknown (not kOutOfBound) because exhausting // this SST says nothing about the upper bound -- the next SST on the level // may still have in-bound keys. @@ -3116,7 +973,102 @@ TEST_F(TrieIndexFactoryTest, NewReaderWithCorruptedData) { Slice bad_slice(bad_data); std::unique_ptr reader; Status s = factory_->NewReader(option, bad_slice, reader); - ASSERT_TRUE(s.IsCorruption()) << s.ToString(); + ASSERT_EQ(s.code(), Status::Code::kCorruption) << s.ToString(); +} + +// An index block that passes its checksum can still be malformed, e.g. when +// written by a buggy writer. The reader must reject it or read it without +// undefined behavior; run under ASAN/UBSAN to catch out-of-bounds reads. +TEST_F(TrieIndexFactoryTest, MalformedIndexIsRejectedOrReadable) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + + // Long shared prefixes give compressed trie paths, and a user key spanning + // three blocks gives a run. + std::vector keys = {""}; + for (int i = 0; i < 30; i++) { + char suffix[8]; + snprintf(suffix, sizeof(suffix), "%04d", i * 3); + keys.push_back("tenant/0001/objects/" + std::string(suffix) + + std::string(rnd.Uniform(12), 'x')); + } + std::vector blocks; + for (size_t i = 0; i < keys.size(); i++) { + const uint64_t offset = blocks.size() * 100; + if (i == keys.size() / 2) { + blocks.push_back({keys[i], keys[i], offset, 95, 30, 29}); + blocks.push_back({keys[i], keys[i], offset + 100, 95, 20, 19}); + blocks.push_back({keys[i], keys[i + 1], offset + 200, 95, 10, 9}); + } else { + blocks.push_back( + {keys[i], i + 1 < keys.size() ? keys[i + 1] : "", offset, 95, 10, 9}); + } + } + TrieTestContext ctx = BuildTrieAndGetIterator(blocks); + const std::string data = ctx.index_contents.ToString(); + + auto exercise = [&](const std::string& mutated) { + Slice contents(mutated); + std::unique_ptr reader; + if (!OpenedOrRejected( + factory_->NewReader(IndexFactoryOptions(), contents, reader))) { + return; + } + std::unique_ptr iter = + reader->NewIterator(ReadOptions()); + const size_t limit = 2 * blocks.size() + 10; + IterateResult result; + auto inbound = [&](const Status& s) { + return OkOrCorruption(s) && + result.bound_check_result == IterBoundCheck::kInbound; + }; + bool valid = inbound(iter->SeekToFirstAndGetResult(&result)); + for (size_t i = 0; valid && i < limit; i++) { + iter->value(); + valid = inbound(iter->NextAndGetResult(&result)); + } + valid = inbound(iter->SeekToLastAndGetResult(&result)); + for (size_t i = 0; valid && i < limit; i++) { + valid = inbound(iter->PrevAndGetResult(&result)); + } + for (const std::string& key : keys) { + for (SequenceNumber seq : {kMaxSequenceNumber, SequenceNumber{15}}) { + if (inbound(iter->SeekAndGetResult(key, &result, SeekCtx(seq)))) { + iter->value(); + inbound(iter->NextAndGetResult(&result)); + } + } + } + }; + + // Counts and offsets are 32- or 64-bit fields: plant values that overflow + // size arithmetic, and pairs of 32-bit counts whose sum wraps to zero. + for (size_t pos = 0; pos + 4 <= data.size(); pos += 4) { + for (uint32_t v : {0x7fffffffu, 0x80000001u, 0xffffffffu}) { + std::string mutated = data; + EncodeFixed32(&mutated[pos], v); + exercise(mutated); + } + if (pos + 8 <= data.size()) { + std::string mutated = data; + EncodeFixed32(&mutated[pos], 0x80000001u); + EncodeFixed32(&mutated[pos + 4], 0x80000001u); + exercise(mutated); + for (uint64_t v : {uint64_t{1} << 32, uint64_t{0x8000000000000002}, + uint64_t{0xffffffffffffffff}}) { + mutated = data; + EncodeFixed64(&mutated[pos], v); + exercise(mutated); + } + } + } + for (size_t len = 0; len < data.size(); len++) { + exercise(data.substr(0, len)); + } + for (int i = 0; i < 2000; i++) { + exercise(FlipRandomBytes(data, &rnd)); + } } TEST_F(TrieIndexFactoryTest, OnKeyAddedNoOp) { @@ -3228,18 +1180,18 @@ TEST_F(TrieIndexFactoryTest, SeekSucceedsButTargetPastLimit) { } // ============================================================================ -// Sequence number encoding tests +// Same-user-key boundary tests // // These tests verify the same-user-key boundary handling: when the same user -// key spans a data block boundary with different sequence numbers, the trie -// builder detects this and encodes ALL separators with sequence number suffixes -// so the trie can distinguish the two blocks. +// key spans a data block boundary with different sequence numbers, the block +// before the boundary keeps the user key as its separator and gets its last +// key's tag, so Seek can choose between the blocks by sequence number. // ============================================================================ -TEST_F(TrieIndexFactoryTest, SameUserKeyBoundaryTriggersSeqnoEncoding) { +TEST_F(TrieIndexFactoryTest, SameUserKeyBoundaryGetsTag) { // When the same user key spans two adjacent data blocks (e.g., // "foo"|seq=100 ends Block 0, "foo"|seq=50 starts Block 1), the - // builder must detect this and switch to seqno-encoded separators. + // builder must detect this and give Block 0 its last key's tag. // Without this, FindShortestSeparator("foo", "foo") = "foo" and the // trie cannot distinguish the two blocks, causing incorrect Seek. auto ctx = BuildTrieAndGetIterator({ @@ -3291,7 +1243,7 @@ TEST_F(TrieIndexFactoryTest, SameUserKeyBoundaryTriggersSeqnoEncoding) { } TEST_F(TrieIndexFactoryTest, DistinctUserKeysSeekWithMaxSeqno) { - // Seqno encoding is always active. Verify that seeking with + // No two blocks share a user key. Verify that seeking with // kMaxSequenceNumber correctly finds each block. auto ctx = BuildTrieAndGetIterator({ {"apple", "cherry", 0, 1000, 100, 50}, @@ -3326,11 +1278,9 @@ TEST_F(TrieIndexFactoryTest, MultipleSameUserKeyBoundaries) { ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("key"), 300, 0)); - // Seek "key"|seq=250 -> Block 1. In internal key order, higher seqno is - // "smaller", so "key"|300 < "key"|250. Block 0's separator "key"+enc(300) - // is less than the target -> skip. Block 1's separator "key"+enc(200) - // is >= target because enc(200) > enc(250) (lower seqno -> larger encoded - // bytes). So the first separator >= target is Block 1. + // Seek "key"|seq=250 -> Block 1. Blocks 0 and 1 form the "key" run with + // tags 300 and 200 (Block 2 gets separator "l"), and Block 1 is the first + // block of the run whose tag is at most 250. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("key"), 250, 1000)); @@ -3338,10 +1288,8 @@ TEST_F(TrieIndexFactoryTest, MultipleSameUserKeyBoundaries) { ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("key"), 200, 1000)); - // Seek "key"|seq=150 -> Block 2. "key"+enc(300) < target (skip Block 0), - // "key"+enc(200) < target (skip Block 1, because enc(200) < enc(150): - // 200 > 150 -> higher seqno -> smaller encoding). The next separator is - // Block 2's separator which is >= target -> Block 2. + // Seek "key"|seq=150 -> Block 2. Both blocks of the run have larger tags, + // so Seek moves to the next separator, "l" (Block 2). ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("key"), 150, 2000)); @@ -3365,11 +1313,9 @@ TEST_F(TrieIndexFactoryTest, MultipleSameUserKeyBoundaries) { TEST_F(TrieIndexFactoryTest, SameUserKeyWithZeroSeqnos) { // During bottommost compaction, RocksDB zeroes sequence numbers for keys // that no longer need version discrimination. When the same user key spans - // multiple data blocks and all versions have seqno=0, the overflow entries - // in the same-key run also have seqno=0. This is valid -- the reader's - // post-seek correction handles it correctly: - // - Primary leaf seqno=0, and target_packed >= 0 so no advancement - // - Next() iterates overflow blocks by index, not seqno + // multiple data blocks and all versions have seqno=0, the blocks of its run + // have no tags. This is valid: Seek lands on the first block of the run + // for any target seqno, and Next() walks the run's blocks in order. auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 1000, 0, 0}, {"key", "key", 1000, 1000, 0, 0}, @@ -3377,12 +1323,11 @@ TEST_F(TrieIndexFactoryTest, SameUserKeyWithZeroSeqnos) { {"zzz", "", 3000, 1000, 0, 0}, }); - // Seek "key"|kMaxSequenceNumber -> Block 0 (primary, seqno=0 guard prevents - // advancement through overflow blocks). + // Seek "key"|kMaxSequenceNumber -> Block 0, the first block of the run. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("key"), kMaxSequenceNumber, 0)); - // Seek "key"|seq=0 -> Block 0 (primary, target_packed=0 >= leaf_packed=0). + // Seek "key"|seq=0 -> Block 0, which has no tag. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(ctx.iter.get(), Slice("key"), 0, 0)); // Full forward scan: Block 0 -> 1 -> 2 -> 3 @@ -3411,11 +1356,11 @@ TEST_F(TrieIndexFactoryTest, SameUserKeyLastBlockZeroSeqno) { {"key", "", 1000, 1000, 0, 0}, // last block }); - // Seek -> Block 0 (primary). + // Seek -> Block 0, the first block of the run. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("key"), kMaxSequenceNumber, 0)); - // Next -> Block 1 (overflow). + // Next -> Block 1, the second block of the run. IterateResult result; ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); @@ -3426,10 +1371,10 @@ TEST_F(TrieIndexFactoryTest, SameUserKeyLastBlockZeroSeqno) { ASSERT_NE(result.bound_check_result, IterBoundCheck::kInbound); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingRoundTripSerialization) { - // Verify that the seqno encoding flag survives serialization/deserialization. - // Build a trie with a same-user-key boundary (triggers seqno encoding), - // serialize it, deserialize it, and verify the flag is set and seeks work. +TEST_F(TrieIndexFactoryTest, TagsRoundTripSerialization) { + // Verify that block tags survive serialization/deserialization. Build a + // trie with a same-user-key boundary (a run of two tagged blocks), + // serialize it, deserialize it, and verify that seeks pick by seqno. auto ctx = BuildTrieAndGetIterator({ {"x", "x", 0, 500, 10, 5}, {"x", "", 500, 500, 5, 0}, @@ -3442,10 +1387,9 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingRoundTripSerialization) { ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(ctx.iter.get(), Slice("x"), 5, 500)); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingWithBoundsChecking) { - // Verify that bounds checking still works correctly when seqno encoding - // is active. The bounds comparison should use user keys only (not the - // seqno-encoded internal representation). +TEST_F(TrieIndexFactoryTest, TagsWithBoundsChecking) { + // Verify that bounds checking works correctly on blocks with tags. The + // bounds comparison uses user keys only, never tags. auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 1000, 100, 50}, {"key", "zzz", 1000, 1000, 50, 1}, @@ -3470,37 +1414,35 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingWithBoundsChecking) { } // ============================================================================ -// Side-table comprehensive tests +// Run tests // -// These tests exercise the side-table seqno encoding for same-user-key -// boundaries. Key insight: FindShortestSeparator shortens separator keys -// when the two user keys differ. For example, FindShortestSeparator("aaa", -// "mmm") = "b". Only same-user-key boundaries produce identical separators. -// The trie de-duplicates these into overflow runs with a seqno side-table. +// These tests exercise runs: consecutive blocks that share a separator +// because one user key spans them. Key insight: FindShortestSeparator +// shortens separator keys when the two user keys differ. For example, +// FindShortestSeparator("aaa", "mmm") = "b". Only same-user-key boundaries +// (or keys that cannot be shortened) produce identical separators. // // Coverage: -// - Large overflow runs (10+ blocks with the same user key) -// - Mixed runs (multiple different user keys each with multi-block runs) +// - Long runs (10+ blocks with the same user key) +// - Mixed runs (multiple different user keys each spanning several blocks) // - Adjacent runs (two different user keys, each with their own run) // - Seek with kMaxSequenceNumber on same-user-key blocks // - Seek with seq=0 on same-user-key blocks -// - result.key is the trie separator (no encoded suffix) -// - Seeking non-existent keys when seqno encoding is active -// - Seeking past all keys / Next() past last block with seqno encoding -// - kUnknown on exhaustion with seqno encoding and overflow -// - Next() transitions: overflow->next leaf, overflow->next overflow run -// - Bounds checking during overflow run iteration -// - Size comparison: seqno-free trie same size as without the feature -// - Re-seeking after being positioned in an overflow run +// - result.key is the separator (no tag appended) +// - Seeking keys that are not separators +// - Seeking past all keys / Next() past the last block of a run +// - kUnknown on exhaustion after a run +// - Next() transitions: run->next separator, run->next run +// - Bounds checking while walking a run +// - Re-seeking after a seek that moved past a run // ============================================================================ -TEST_F(TrieIndexFactoryTest, LargeOverflowRun) { +TEST_F(TrieIndexFactoryTest, LargeRun) { // 12 blocks all with user key "key", seqnos descending from 1200 to 100. // The last "key" block transitions to "zzz", so - // FindShortestSeparator("key", "zzz") = "l". The trie has: - // "key" (11-block run, seqnos 1200..200) -> "l" (block 11) -> "zzz" (block - // 12) - // The overflow run for "key" has 10 overflow entries (blocks 1-10). + // FindShortestSeparator("key", "zzz") = "l". The separators are: + // "key" (run of blocks 0-10, tags 1200..200) -> "l" (block 11) -> + // "zzz" (block 12) IndexFactoryOptions option; option.comparator = BytewiseComparator(); @@ -3553,22 +1495,18 @@ TEST_F(TrieIndexFactoryTest, LargeOverflowRun) { SeekCtx(kMaxSequenceNumber))); ASSERT_EQ(iter->value().offset, 0u); - // Seek "key"|1100 -> leaf_seqno=1200, 1100<1200 -> advance. - // overflow[0]=1100, 1100>=1100 -> Block 1. + // Seek "key"|1100 -> Block 1, the first block of the run whose tag is at + // most 1100. ASSERT_OK(iter->SeekAndGetResult(Slice("key"), &result, SeekCtx(1100))); ASSERT_EQ(iter->value().offset, 1000u); - // Seek "key"|550 -> advance through overflow: - // overflow seqnos: 1100,1000,900,800,700,600,500,400,300,200 - // 550 < 1100..600, 550 >= 500? No, 550 is not < 600: - // Actually: 550<1100 skip, 550<1000 skip, 550<900 skip, 550<800 skip, - // 550<700 skip, 550<600 skip, 550>=500 -> found at overflow index 6. - // overflow_run_index_ = 7 (1-based) -> Block 7, offset 7000. + // Seek "key"|550 -> Block 7 (tag 500), the first block of the run whose + // tag is at most 550. ASSERT_OK(iter->SeekAndGetResult(Slice("key"), &result, SeekCtx(550))); ASSERT_EQ(iter->value().offset, 7000u); - // Seek "key"|50 -> below all "key" overflow seqnos (minimum is 200). - // All 10 overflow seqnos exhausted -> advance to next leaf "l" (block 11). + // Seek "key"|50 -> every block of the run has a larger tag (the smallest + // is 200), so Seek moves to the next separator, "l" (block 11). ASSERT_OK(iter->SeekAndGetResult(Slice("key"), &result, SeekCtx(50))); ASSERT_EQ(iter->value().offset, 11000u); @@ -3591,45 +1529,27 @@ TEST_F(TrieIndexFactoryTest, LargeOverflowRun) { } TEST_F(TrieIndexFactoryTest, MixedSameKeyRuns) { - // Two different user keys each with multi-block runs, plus distinct blocks. - // - // The trie structure (after FindShortestSeparator shortening): - // "aaa" (2-block run, seqnos 300, 200) -> "b" (block 2) -> - // "mmm" (1-block, seqno 60) -> "n" (block 4) -> "zzz" (block 5) - // - // Block 0: "aaa"|300 -> next "aaa"|200 (same-key -> separator "aaa") - // Block 1: "aaa"|200 -> next "mmm" (diff-key -> separator "b") - // Block 2: "mmm"|60 -> next "mmm"|30 (same-key -> separator "mmm") - // Block 3: "mmm"|30 -> next "zzz" (diff-key -> separator "n") - // Block 4: "zzz"|10 -> null (last -> separator "zzz") - // - // "aaa" run: blocks 0-1 (2 entries, run of 2 in trie) - // block 2 gets separator "b" (separate trie leaf) - // "mmm" run: block 3 only (run of 1 in trie) - // block 4 gets separator "n" (separate trie leaf) - // block 5 gets separator "zzz" (separate trie leaf, no shortening) + // Two different user keys that each span several blocks, plus distinct + // blocks. // - // Wait -- let me recount. We have 6 AddIndexEntry calls: - // Entry 0: "aaa" 300 next="aaa" -> sep="aaa" - // Entry 1: "aaa" 200 next="mmm" -> sep="b" - // Entry 2: "mmm" 60 next="mmm" -> sep="mmm" - // Entry 3: "mmm" 30 next="zzz" -> sep="n" - // Entry 4: "zzz" 10 next=null -> sep="zzz" + // Blocks (separator after FindShortestSeparator): + // Block 0: "aaa"|300 -> next "aaa"|200 (same-key -> separator "aaa") + // Block 1: "aaa"|200 -> next "aaa"|100 (same-key -> separator "aaa") + // Block 2: "aaa"|100 -> next "mmm"|60 (diff-key -> separator "b") + // Block 3: "mmm"|60 -> next "mmm"|30 (same-key -> separator "mmm") + // Block 4: "mmm"|30 -> next "zzz"|10 (diff-key -> separator "n") + // Block 5: "zzz"|10 -> null (last -> separator "zzz") // - // De-duplicated separators: "aaa" appears once (run=1), "b", "mmm" appears - // once (run=1), "n", "zzz". So actually there are no overflow runs here! - // The "aaa" run is only block 0, and block 1 gets separator "b". - // The "mmm" run is only block 2, and block 3 gets separator "n". - // - // For a true multi-block overflow run for "aaa", we need 3+ consecutive - // blocks ALL with "aaa" -> "aaa" boundaries. + // The last block of each user key gets a shortened separator, so the + // "aaa" run is blocks 0-1 (tags 300, 200), and "mmm" has only block 3 + // (tag 60). IndexFactoryOptions option; option.comparator = BytewiseComparator(); std::unique_ptr builder; ASSERT_OK(factory_->NewBuilder(option, builder)); - // "aaa" run: 3 blocks (all same-key boundaries). + // User key "aaa": 3 blocks (the first two end at same-key boundaries). { IndexFactoryBuilder::BlockHandle h{0, 1000}; std::string scratch; @@ -3651,7 +1571,7 @@ TEST_F(TrieIndexFactoryTest, MixedSameKeyRuns) { Slice next("mmm"); builder->AddIndexEntry(Slice("aaa"), &next, h, &scratch, EntryCtx(100, 60)); } - // "mmm" run: 2 blocks (one same-key boundary). + // User key "mmm": 2 blocks (one same-key boundary). { IndexFactoryBuilder::BlockHandle h{3000, 1000}; std::string scratch; @@ -3682,33 +1602,30 @@ TEST_F(TrieIndexFactoryTest, MixedSameKeyRuns) { auto iter = reader->NewIterator(ro); IterateResult result; - // Trie structure: "aaa"(run=2) -> "b" -> "mmm"(run=1) -> "n" -> "zzz" + // Separators: "aaa" (run of 2) -> "b" -> "mmm" -> "n" -> "zzz" // // Seek "aaa"|kMax -> Block 0. ASSERT_OK(iter->SeekAndGetResult(Slice("aaa"), &result, SeekCtx(kMaxSequenceNumber))); ASSERT_EQ(iter->value().offset, 0u); - // Seek "aaa"|250 -> leaf_seqno=300, 250<300 -> advance. - // overflow[0]=200, 250>=200 -> Block 1. + // Seek "aaa"|250 -> Block 1, the first block of the run whose tag is at + // most 250. ASSERT_OK(iter->SeekAndGetResult(Slice("aaa"), &result, SeekCtx(250))); ASSERT_EQ(iter->value().offset, 1000u); - // Seek "aaa"|50 -> leaf_seqno=300, 50<300 -> advance. - // overflow[0]=200, 50<200 -> skip. All exhausted -> next leaf "b" (block 2). + // Seek "aaa"|50 -> both blocks of the run have larger tags, so Seek moves + // to the next separator, "b" (block 2). ASSERT_OK(iter->SeekAndGetResult(Slice("aaa"), &result, SeekCtx(50))); ASSERT_EQ(iter->value().offset, 2000u); - // Seek "mmm"|kMax -> Block 3 (the "mmm" trie leaf). + // Seek "mmm"|kMax -> Block 3 (separator "mmm"). ASSERT_OK(iter->SeekAndGetResult(Slice("mmm"), &result, SeekCtx(kMaxSequenceNumber))); ASSERT_EQ(iter->value().offset, 3000u); - // Seek "mmm"|45 -> leaf_seqno=60, 45<60 -> advance. - // "mmm" has run_size=1 (only 1 block in trie), no overflow -> next leaf "n". - // Wait -- "mmm" appears only once in the separator list (entry 2), because - // entry 3 has separator "n". So "mmm" trie leaf has block_count=1. - // 45 < 60, no overflow -> advance to "n" (block 4). + // Seek "mmm"|45 -> Block 3's tag 60 is larger and "mmm" has no other + // block, so Seek moves to the next separator, "n" (block 4). ASSERT_OK(iter->SeekAndGetResult(Slice("mmm"), &result, SeekCtx(45))); ASSERT_EQ(iter->value().offset, 4000u); @@ -3729,11 +1646,10 @@ TEST_F(TrieIndexFactoryTest, MixedSameKeyRuns) { } TEST_F(TrieIndexFactoryTest, AdjacentSameKeyRuns) { - // Two adjacent same-user-key runs. The trick is that the last block of - // the first run has a different-key boundary, creating a non-run trie leaf - // between the two runs. To get true adjacency we make ALL "aaa" blocks - // transition to "aaa", and ALL "bbb" blocks transition to "bbb", with - // the boundary between runs being "aaa"->"bbb". + // Two user keys that each span two blocks, one after the other. The last + // "aaa" block has a different-key boundary, so it gets the shortened + // separator "b", a single-block separator between the "aaa" and "bbb" + // blocks. // // Entries (separator after FindShortestSeparator): // Entry 0: "aaa" 200 next="aaa" -> sep="aaa" (same-key) @@ -3741,7 +1657,8 @@ TEST_F(TrieIndexFactoryTest, AdjacentSameKeyRuns) { // Entry 2: "bbb" 80 next="bbb" -> sep="bbb" (same-key) // Entry 3: "bbb" 40 next=null -> sep="bbb" (last block, no shortening) // - // Trie: "aaa"(run=1, seq=200) -> "b"(1 block) -> "bbb"(run=2, seqnos 80,40) + // Separators: "aaa" (1 block, tag 200) -> "b" (1 block) -> "bbb" (run of 2, + // tags 80, 40) // Note: the last block joins the "bbb" run since its separator also is "bbb". auto ctx = BuildTrieAndGetIterator({ {"aaa", "aaa", 0, 1000, 200, 100}, @@ -3753,7 +1670,7 @@ TEST_F(TrieIndexFactoryTest, AdjacentSameKeyRuns) { IterateResult result; // Full scan: "aaa" (block 0) -> "b" (block 1) -> "bbb" (block 2) -> "bbb" - // overflow (block 3). + // (block 3, the second block of the run). ASSERT_OK(ctx.iter->SeekAndGetResult(Slice("aaa"), &result, SeekCtx(kMaxSequenceNumber))); ASSERT_EQ(ctx.iter->value().offset, 0u); @@ -3776,14 +1693,13 @@ TEST_F(TrieIndexFactoryTest, AdjacentSameKeyRuns) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kUnknown); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingResultKeyIsUserKey) { - // Verify that result.key from SeekAndGetResult is the trie separator key - // (no encoded seqno suffix). For same-key boundary leaves, the separator - // IS the original user key. For different-key boundaries, the separator - // is shortened. +TEST_F(TrieIndexFactoryTest, TaggedResultKeyIsUserKey) { + // Verify that result.key from SeekAndGetResult is the separator, with no + // tag appended. At a same-key boundary, the separator IS the original user + // key. For different-key boundaries, the separator is shortened. // // Entries: - // "foo" 100 next="foo" -> sep="foo" (same-key, run of 2) + // "foo" 100 next="foo" -> sep="foo" (same-key) // "foo" 50 next="zoo" -> sep="g" (FindShortestSeparator("foo","zoo")) // "zoo" 1 next=null -> sep="zoo" (last block, no successor shortening) auto ctx = BuildTrieAndGetIterator({ @@ -3794,12 +1710,12 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingResultKeyIsUserKey) { IterateResult result; - // Seek to "foo"|100 -> "foo" leaf. result.key = "foo" (3 bytes, no suffix). + // Seek to "foo"|100 -> block 0. result.key = "foo" (3 bytes, no tag). ASSERT_OK(ctx.iter->SeekAndGetResult(Slice("foo"), &result, SeekCtx(100))); ASSERT_EQ(result.key.ToString(), "foo"); ASSERT_EQ(result.key.size(), 3u); - // Next -> "g" leaf (FindShortestSeparator("foo", "zoo") = "g"). + // Next -> block 1 (FindShortestSeparator("foo", "zoo") = "g"). // result.key = "g" (1 byte). ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.key.size(), 1u); @@ -3808,10 +1724,10 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingResultKeyIsUserKey) { ASSERT_LT(result.key.ToString(), "zoo"); } -TEST_F(TrieIndexFactoryTest, SeekNonExistentKeyWithSeqnoEncoding) { - // Seeking for keys not in the trie when seqno encoding is active. +TEST_F(TrieIndexFactoryTest, SeekNonExistentKeyWithTags) { + // Seeking keys that are not separators. // - // Trie: "mmm"(run=1, seq=100) -> "n" -> "zzz" + // Separators: "mmm" (1 block, tag 100) -> "n" -> "zzz" auto ctx = BuildTrieAndGetIterator({ {"mmm", "mmm", 0, 1000, 100, 50}, {"mmm", "zzz", 1000, 1000, 50, 1}, @@ -3819,10 +1735,10 @@ TEST_F(TrieIndexFactoryTest, SeekNonExistentKeyWithSeqnoEncoding) { }); auto* iter = ctx.iter.get(); - // Seek "aaa" (before all keys) -> lands on "mmm" leaf, Block 0. + // Seek "aaa" (before all keys) -> lands on "mmm", Block 0. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(iter, Slice("aaa"), kMaxSequenceNumber, 0u)); - // Seek "nnn" (between "n" and "zzz") -> lands on "zzz" leaf, Block 2. + // Seek "nnn" (between "n" and "zzz") -> lands on "zzz", Block 2. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(iter, Slice("nnn"), kMaxSequenceNumber, 2000u)); @@ -3833,11 +1749,11 @@ TEST_F(TrieIndexFactoryTest, SeekNonExistentKeyWithSeqnoEncoding) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kUnknown); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingPastEndAndNextPastEnd) { - // Verify seeking past all keys and Next() past the last block with seqno - // encoding active. +TEST_F(TrieIndexFactoryTest, TagsPastEndAndNextPastEnd) { + // Verify seeking past all keys and Next() past the last block, which ends + // a run. // - // Trie: "key"(run=2, seqnos 10,5) + // Separators: "key" (run of 2, tags 10, 5) auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 500, 10, 5}, {"key", "", 500, 500, 5, 0}, @@ -3850,8 +1766,8 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingPastEndAndNextPastEnd) { SeekCtx(kMaxSequenceNumber))); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kUnknown); - // Seek "key"|5 -> seqno=10 on primary, 5<10 -> overflow seqno=5, 5>=5 -> - // overflow block 1. + // Seek "key"|5 -> block 1, the first block of the run whose tag is at most + // 5. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice("key"), 5, 500u)); // Next from last block in run -> past end. @@ -3862,9 +1778,9 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingPastEndAndNextPastEnd) { ASSERT_NO_FATAL_FAILURE(AssertFullForwardScan(iter, Slice("key"), {0, 500})); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingOutOfBoundWithOverflow) { - // Verify that bounds checking works correctly with seqno encoding and - // overflow blocks. +TEST_F(TrieIndexFactoryTest, OutOfBoundWithRun) { + // Verify that bounds checking works correctly while walking a run and the + // blocks after it. // // Entries: // "key" 300 next="key" -> sep="key" (same-key) @@ -3872,7 +1788,7 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingOutOfBoundWithOverflow) { // "key" 100 next="zzz" -> sep="l" // "zzz" 1 next=null -> sep="zzz" (last block, no shortening) // - // Trie: "key"(run=2, seqnos 300,200) -> "l" -> "zzz" + // Separators: "key" (run of 2, tags 300, 200) -> "l" -> "zzz" auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 1000, 300, 200}, {"key", "key", 1000, 1000, 200, 100}, @@ -3891,18 +1807,18 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingOutOfBoundWithOverflow) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 0u); - // Next -> Block 1 (overflow within "key" run), prev "key" < "zzz" -> + // Next -> Block 1 (second block of the "key" run), prev "key" < "zzz" -> // kInbound. ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 1000u); - // Next -> "l" leaf (block 2), prev "key" < "zzz" -> kInbound. + // Next -> "l" (block 2), prev "key" < "zzz" -> kInbound. ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 2000u); - // Next -> "zzz" leaf (block 3). prev key is "l", "l" < "zzz" -> kInbound. + // Next -> "zzz" (block 3). prev key is "l", "l" < "zzz" -> kInbound. ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 3000u); @@ -3917,82 +1833,23 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingOutOfBoundWithOverflow) { ASSERT_OK(ctx.iter->SeekAndGetResult(Slice("key"), &result, SeekCtx(300))); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); - // Next -> Block 1 (overflow), prev "key" < "l" -> kInbound. + // Next -> Block 1 (second block of the run), prev "key" < "l" -> kInbound. ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); - // Next -> "l" leaf (block 2). prev "key" < "l" -> kInbound (conservative). + // Next -> "l" (block 2). prev "key" < "l" -> kInbound (conservative). ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); - // Next -> "zzz" leaf (block 3). prev "l" >= "l" -> kOutOfBound. + // Next -> "zzz" (block 3). prev "l" >= "l" -> kOutOfBound. ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kOutOfBound); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingConsistentSize) { - // Verify that tries built with different seqno contexts produce the - // same serialized size (seqno encoding is always on). - IndexFactoryOptions option; - option.comparator = BytewiseComparator(); - - // Build trie with distinct keys, seqno=0. - std::unique_ptr builder_no_seq; - ASSERT_OK(factory_->NewBuilder(option, builder_no_seq)); - for (int i = 0; i < 50; i++) { - char buf[16]; - snprintf(buf, sizeof(buf), "key_%04d", i); - char next_buf[16]; - snprintf(next_buf, sizeof(next_buf), "key_%04d", i + 1); - - IndexFactoryBuilder::BlockHandle handle{static_cast(i) * 1000, - 500}; - std::string scratch; - if (i < 49) { - Slice next(next_buf); - builder_no_seq->AddIndexEntry(Slice(buf), &next, handle, &scratch, - {0, 0}); - } else { - builder_no_seq->AddIndexEntry(Slice(buf), nullptr, handle, &scratch, - {0, 0}); - } - } - Slice contents_no_seq; - ASSERT_OK(builder_no_seq->Finish(&contents_no_seq)); - size_t size_no_seq = contents_no_seq.size(); - - // Same keys but with nonzero seqnos (still distinct -> no seqno encoding). - std::unique_ptr builder_with_seq; - ASSERT_OK(factory_->NewBuilder(option, builder_with_seq)); - for (int i = 0; i < 50; i++) { - char buf[16]; - snprintf(buf, sizeof(buf), "key_%04d", i); - char next_buf[16]; - snprintf(next_buf, sizeof(next_buf), "key_%04d", i + 1); - - IndexFactoryBuilder::BlockHandle handle{static_cast(i) * 1000, - 500}; - std::string scratch; - if (i < 49) { - Slice next(next_buf); - builder_with_seq->AddIndexEntry(Slice(buf), &next, handle, &scratch, - {100, 50}); - } else { - builder_with_seq->AddIndexEntry(Slice(buf), nullptr, handle, &scratch, - {100, 0}); - } - } - Slice contents_with_seq; - ASSERT_OK(builder_with_seq->Finish(&contents_with_seq)); - size_t size_with_seq = contents_with_seq.size(); - - ASSERT_EQ(size_no_seq, size_with_seq); -} - TEST_F(TrieIndexFactoryTest, SeekWithMaxSeqOnSameKeyBlocks) { - // kMaxSequenceNumber should always land on Block 0 of a same-key run. + // kMaxSequenceNumber should always land on the first block of a run. // - // Trie: "key"(run=3, seqnos 400,300,200) -> "l" + // Separators: "key" (run of 4, tags 400, 300, 200, 100) auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 1000, 400, 300}, {"key", "key", 1000, 1000, 300, 200}, @@ -4004,10 +1861,10 @@ TEST_F(TrieIndexFactoryTest, SeekWithMaxSeqOnSameKeyBlocks) { } TEST_F(TrieIndexFactoryTest, SeekWithZeroSeqOnSameKeyBlocks) { - // seq=0 is below all overflow seqnos -> advance past the entire run. + // seq=0 is below every tag of the run, so Seek moves past the whole run. // - // Trie: "key"(run=2, seqnos 300,200) -> "l" -> "zzz" - // seq=0 < all overflow seqnos -> advance past run -> lands on "l" (block 2). + // Separators: "key" (run of 2, tags 300, 200) -> "l" -> "zzz" + // Seek "key"|0 lands on "l" (block 2). auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 1000, 300, 200}, {"key", "key", 1000, 1000, 200, 100}, @@ -4018,19 +1875,19 @@ TEST_F(TrieIndexFactoryTest, SeekWithZeroSeqOnSameKeyBlocks) { AssertSeekOffset(ctx.iter.get(), Slice("key"), 0, 2000u)); } -TEST_F(TrieIndexFactoryTest, ZeroSeqMustNotSkipLeafForSmallerUserKey) { +TEST_F(TrieIndexFactoryTest, ZeroSeqMustNotSkipSmallerUserKey) { // Validates correctness of DBIter's forward-scan reseek path. // // Block 0 ends with user key "m" and block 1 starts with the same user key, - // so the trie stores separator "m" with non-zero seqno metadata on block 0. + // so block 0 gets separator "m" and a nonzero tag. // However, a seek target for a *smaller* user key "l"|0 must still land on // block 0, because block 0 can contain keys in ("l", "m"]. // - // Verifies that seqno-based post-seek correction is NOT applied when the - // target user key is strictly less than the separator user key. The seqno - // comparison is only meaningful when user keys are equal. If the target - // user key "l" < separator "m", the seek must stay on block 0 regardless - // of seqno, because block 0 can contain keys up to "m". + // Verifies that tags are NOT compared when the target user key is strictly + // less than the separator. The tag comparison is only meaningful when user + // keys are equal. If the target user key "l" < separator "m", the seek must + // stay on block 0 regardless of seqno, because block 0 can contain keys up + // to "m". auto ctx = BuildTrieAndGetIterator({ {"m", "m", 0, 1000, 300, 200}, {"y", "zzz", 1000, 1000, 100, 1}, @@ -4040,12 +1897,12 @@ TEST_F(TrieIndexFactoryTest, ZeroSeqMustNotSkipLeafForSmallerUserKey) { ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(ctx.iter.get(), Slice("l"), 0, 0u)); } -TEST_F(TrieIndexFactoryTest, NextTransitionOverflowToOverflow) { - // Test Next() transitioning from one overflow run to the next trie leaf - // that also has an overflow run. +TEST_F(TrieIndexFactoryTest, NextFromRunToRun) { + // Test Next() walking from one run, through a single-block separator, into + // another run. // - // To get two true adjacent overflow runs, both user keys need multiple - // same-key boundaries. We need a 3rd user key to separate them in the trie. + // Both user keys need several same-key boundaries to form runs. The last + // "aaa" block gets the shortened separator "b", which sits between them. // // Entries: // "aaa" 200 next="aaa" -> sep="aaa" (same-key) @@ -4055,7 +1912,8 @@ TEST_F(TrieIndexFactoryTest, NextTransitionOverflowToOverflow) { // "bbb" 60 next="bbb" -> sep="bbb" (same-key) // "bbb" 30 next=null -> sep="bbb" (last block, no successor shortening) // - // Trie: "aaa"(run=2, seqnos 200,100) -> "b" -> "bbb"(run=3, seqnos 90,60,30) + // Separators: "aaa" (run of 2, tags 200, 100) -> "b" -> + // "bbb" (run of 3, tags 90, 60, 30) // Note: the last block's separator is "bbb" (not "c"), matching the standard // index builder's behavior with kShortenSeparators (the default). This means // the last block joins the "bbb" run, making it a 3-block run. @@ -4118,23 +1976,23 @@ TEST_F(TrieIndexFactoryTest, NextTransitionOverflowToOverflow) { ASSERT_EQ(iter->value().offset, 0u); ASSERT_EQ(result.key.ToString(), "aaa"); - ASSERT_OK(iter->NextAndGetResult(&result)); // "aaa" overflow (block 1) + ASSERT_OK(iter->NextAndGetResult(&result)); // "aaa" (block 1) ASSERT_EQ(iter->value().offset, 1000u); ASSERT_EQ(result.key.ToString(), "aaa"); - ASSERT_OK(iter->NextAndGetResult(&result)); // -> "b" leaf (block 2) + ASSERT_OK(iter->NextAndGetResult(&result)); // -> "b" (block 2) ASSERT_EQ(iter->value().offset, 2000u); ASSERT_EQ(result.key.ToString(), "b"); - ASSERT_OK(iter->NextAndGetResult(&result)); // -> "bbb" leaf (block 3) + ASSERT_OK(iter->NextAndGetResult(&result)); // -> "bbb" (block 3) ASSERT_EQ(iter->value().offset, 3000u); ASSERT_EQ(result.key.ToString(), "bbb"); - ASSERT_OK(iter->NextAndGetResult(&result)); // "bbb" overflow (block 4) + ASSERT_OK(iter->NextAndGetResult(&result)); // "bbb" (block 4) ASSERT_EQ(iter->value().offset, 4000u); ASSERT_EQ(result.key.ToString(), "bbb"); - ASSERT_OK(iter->NextAndGetResult(&result)); // "bbb" overflow (block 5) + ASSERT_OK(iter->NextAndGetResult(&result)); // "bbb" (block 5) ASSERT_EQ(iter->value().offset, 5000u); ASSERT_EQ(result.key.ToString(), "bbb"); @@ -4146,20 +2004,19 @@ TEST_F(TrieIndexFactoryTest, NextTransitionOverflowToOverflow) { TEST_F(TrieIndexFactoryTest, SingleBlockWithSeqnoActive) { // Both blocks have the same user key "x". Without successor shortening // for the last block, both separators are "x", forming a 2-block run: - // "x"(run=2, seqnos 10,5) + // "x" (run of 2, tags 10, 5) auto ctx = BuildTrieAndGetIterator({ {"x", "x", 0, 500, 10, 5}, {"x", "", 500, 500, 5, 0}, }); auto* iter = ctx.iter.get(); - // Seek "x"|10 -> leaf_seqno=10, 10>=10 -> Block 0. + // Seek "x"|10 -> Block 0, whose tag is 10. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice("x"), 10, 0u)); - // Seek "x"|7 -> 7<10 -> advance through overflow. overflow seqno=5, 7>=5 -> - // overflow block 1. + // Seek "x"|7 -> Block 1, the first block of the run whose tag is at most 7. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice("x"), 7, 500u)); - // Seek "x"|3 -> 3<10 -> advance. overflow seqno=5, 3<5 -> not found. Advance - // past run. No more leaves -> invalid. This matches the standard index + // Seek "x"|3 -> both tags are larger and there is no later separator, so + // the iterator becomes invalid. This matches the standard index // behavior: seeking for internal key "x"|3 is past the standard index's // last separator "x"|5. { @@ -4175,16 +2032,17 @@ TEST_F(TrieIndexFactoryTest, TagDistinguishesSameSeqDifferentType) { // when two blocks share the same user key and same seqno but differ in // value type, the trie could return the wrong block. // - // Scenario: separator has tag (50 << 8) | 0 = 12800 + // Scenario: block 0 has tag (50 << 8) | 0 = 12800 // (kTypeDeletion) Target has tag (50 << 8) | 1 = 12801 - // (kTypeValue) In internal key order, higher tag = smaller key, so - // target (12801) < separator (12800). The trie should stay on Block 0. + // (kTypeValue) In internal key order, higher tag = smaller key, so the + // target (12801) sorts before block 0's last key (12800). The trie should + // stay on Block 0. // // With seqno-only comparison: 50 < 50 is false -> stay. Same result. // But the reverse case matters: - // Separator tag (50 << 8) | 1 = 12801 (kTypeValue) + // Block 0 tag (50 << 8) | 1 = 12801 (kTypeValue) // Target tag (50 << 8) | 0 = 12800 (kTypeDeletion) - // target (12800) < separator (12801) -> advance to Block 1. + // target (12800) < block 0's tag (12801) -> advance to Block 1. // With seqno-only: 50 < 50 is false -> stay (WRONG -- should advance). IndexFactoryOptions option; @@ -4211,13 +2069,13 @@ TEST_F(TrieIndexFactoryTest, TagDistinguishesSameSeqDifferentType) { IterateResult result; // Target: ("x", seqno=50, kTypeValue=1), packed=12801. - // Separator packed=12801. 12801 >= 12801 -> stay on Block 0. Correct. + // Block 0 tag=12801. 12801 >= 12801 -> stay on Block 0. Correct. ASSERT_OK(iter->SeekAndGetResult(Slice("x"), &result, {(50ULL << 8) | 1})); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(iter->value().offset, 0u); // Target: ("x", seqno=50, kTypeDeletion=0), packed=12800. - // Separator packed=12801. 12800 < 12801 -> advance to Block 1. Correct. + // Block 0 tag=12801. 12800 < 12801 -> advance to Block 1. Correct. // With seqno-only comparison (without tag): 50 < 50 -> false -> // stay on Block 0 (WRONG). ASSERT_OK(iter->SeekAndGetResult(Slice("x"), &result, {(50ULL << 8) | 0})); @@ -4225,11 +2083,11 @@ TEST_F(TrieIndexFactoryTest, TagDistinguishesSameSeqDifferentType) { ASSERT_EQ(iter->value().offset, 500u); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingReSeekAfterOverflow) { - // Verify that re-seeking after being positioned in an overflow run - // correctly resets the overflow state. +TEST_F(TrieIndexFactoryTest, ReSeekAfterRun) { + // Verify that re-seeking after a seek that moved past a run lands on the + // right block each time. // - // Trie: "key"(run=2, seqnos 300,200) -> "l" -> "zzz" + // Separators: "key" (run of 2, tags 300, 200) -> "l" -> "zzz" auto ctx = BuildTrieAndGetIterator({ {"key", "key", 0, 1000, 300, 200}, {"key", "key", 1000, 1000, 200, 100}, @@ -4238,11 +2096,11 @@ TEST_F(TrieIndexFactoryTest, SeqnoEncodingReSeekAfterOverflow) { }); auto* iter = ctx.iter.get(); - // Position in overflow: 150 < all run seqnos -> advance to "l" (block 2). + // Both tags of the run are larger than 150 -> Seek moves to "l" (block 2). ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice("key"), 150, 2000u)); // Re-seek to "key"|300 -> should reset to Block 0. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice("key"), 300, 0u)); - // Re-seek to "zzz" -> "zzz" leaf (block 3), overflow state should be clean. + // Re-seek to "zzz" -> block 3. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(iter, Slice("zzz"), kMaxSequenceNumber, 3000u)); @@ -4266,7 +2124,8 @@ TEST_F(TrieIndexFactoryTest, AllFfLastKeyWithSameKeyBoundary) { // Seek "\xff\xff"|200 -> Block 0. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice(ff), 200, 0u)); - // Seek "\xff\xff"|100 -> advance to overflow Block 1. + // Seek "\xff\xff"|100 -> Block 1, the first block of the run whose tag is + // at most 100. ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(iter, Slice(ff), 100, 500u)); // Seek "\xff\xff"|50 -> past entire run -> exhausted -> kUnknown. @@ -4280,7 +2139,7 @@ TEST_F(TrieIndexFactoryTest, AllFfLastKeyWithSameKeyBoundary) { TEST_F(TrieIndexFactoryTest, AllFfLastKeyWithoutPrecedingSameKey) { // Complementary test: all-0xFF last key but NO same-user-key boundary. - // No seqno encoding is triggered. + // Each separator ("b" and "\xff\xff") has a single block. std::string ff("\xff\xff", 2); auto ctx = BuildTrieAndGetIterator({ {"aaa", ff, 0, 500, 100, 50}, @@ -4290,16 +2149,14 @@ TEST_F(TrieIndexFactoryTest, AllFfLastKeyWithoutPrecedingSameKey) { AssertFullForwardScan(ctx.iter.get(), Slice("aaa"), {0, 500})); } -TEST_F(TrieIndexFactoryTest, SeqnoEncodingRandomized) { +TEST_F(TrieIndexFactoryTest, TagsRandomized) { // Randomized test: generate random block layouts with a mix of same-user-key // and distinct-key boundaries, build a trie, then verify: // 1. Full forward scan visits all blocks in offset order // 2. Seeking each block's key with its exact seqno returns the right offset // 3. Seeking with kMaxSequenceNumber returns the first block for that key // - // Uses Random from util/random.h per project guidelines. - uint32_t seed = static_cast( - std::chrono::steady_clock::now().time_since_epoch().count()); + const uint32_t seed = TimeSeed(); SCOPED_TRACE("seed=" + std::to_string(seed)); Random rnd(seed); @@ -4713,233 +2570,6 @@ TEST_F(TrieIndexSSTTest, MixedKeyTypesWithCompressionDict) { ASSERT_OK(s); } -// ============================================================================ -// Micro-benchmarks: trie Seek vs RocksDB IndexBlockIter. -// -// Compares the trie data structure against the actual RocksDB binary search -// index (IndexBlockIter) which is used in production. The IndexBlockIter -// operates on a prefix-compressed block with InternalKeys, varint decoding, -// and the InternalKeyComparator abstraction -- making it a realistic baseline. -// ============================================================================ - -class TrieSeekBenchmark : public testing::Test { - protected: - // Generate sorted fixed-width keys that mimic realistic RocksDB separator - // keys. Uses snprintf to create keys like "key00000000NNNNNNNN" where N is - // the zero-padded index. This produces keys with moderate prefix sharing - // (the "key" prefix + some leading zeros) while ensuring each key differs - // at multiple byte positions, matching the distribution that SST separator - // keys would have in practice. - static std::vector GenerateKeys(size_t count, - size_t key_size = 16) { - std::vector keys; - keys.reserve(count); - char buf[64]; - for (size_t i = 0; i < count; i++) { - snprintf(buf, sizeof(buf), "%0*zu", static_cast(key_size), i); - keys.emplace_back(buf, key_size); - } - return keys; - } - - // Build a trie from sorted separator keys. Returns {trie, serialized_data}. - // The serialized data string must outlive the trie since the trie stores - // zero-copy pointers into it. - static std::pair BuildTrie( - const std::vector& keys) { - LoudsTrieBuilder builder; - for (size_t i = 0; i < keys.size(); i++) { - TrieBlockHandle handle{i * 4096, 4096}; - builder.AddKey(Slice(keys[i]), handle); - } - builder.Finish(); - - std::string data(builder.GetSerializedData().data(), - builder.GetSerializedData().size()); - LoudsTrie trie; - Status s = trie.InitFromData(Slice(data)); - assert(s.ok()); - (void)s; - return {std::move(trie), std::move(data)}; - } - - // Generate random lookup keys (existing keys from the set). - static std::vector GenerateRandomIndices(size_t count, - size_t num_keys, - uint64_t seed = 42) { - std::mt19937_64 rng(seed); - std::uniform_int_distribution dist(0, num_keys - 1); - std::vector indices(count); - for (auto& idx : indices) { - idx = dist(rng); - } - return indices; - } -}; - -// Benchmark the trie against the actual RocksDB IndexBlockIter, which is -// the real binary search index used in production. Both sides replicate -// their full production code paths: -// -// Trie (via UDI wrapper): -// ParseInternalKey -> Seek(user_key) -> Key().ToString() (string copy) -> -// CheckBounds (no-op without Prepare) -> Value() -> BlockHandle -// -// IndexBlockIter: -// Seek(internal_key) -> value() -> IndexValue with BlockHandle -// -// The IndexBlockIter operates on a prefix-compressed block with InternalKeys, -// varint decoding, and the InternalKeyComparator -- matching production -// exactly. -TEST_F(TrieSeekBenchmark, TrieVsRealIndexBlockIter) { - static constexpr size_t kNumLookups = 200000; - static constexpr size_t kKeySize = 16; - static constexpr size_t kKeyCounts[] = {100, 500, 1000, 5000, 10000, 32000}; - - fprintf(stderr, - "\n====== Trie vs Real IndexBlockIter: Scaling (%zu-byte keys, %zu " - "lookups) ======\n", - kKeySize, kNumLookups); - fprintf(stderr, " %8s %12s %12s %8s %s\n", "Keys", "Trie ns/op", - "IBI ns/op", "vs IBI", "Winner"); - fprintf(stderr, " %8s %12s %12s %8s %s\n", "--------", "------------", - "------------", "--------", "----------"); - - for (size_t num_keys : kKeyCounts) { - auto keys = GenerateKeys(num_keys, kKeySize); - - // ---- Build trie ---- - auto [trie, trie_data] = BuildTrie(keys); - LoudsTrieIterator trie_iter(&trie); - - // ---- Build real RocksDB index block ---- - // Use restart_interval=1 (the default for index blocks). - const int kRestartInterval = 1; - BlockBuilder index_builder( - kRestartInterval, - /*use_delta_encoding=*/true, - /*use_value_delta_encoding=*/false, - BlockBasedTableOptions::kDataBlockBinarySearch, - /*data_block_hash_table_util_ratio=*/0.75, /*ts_sz=*/0, - /*persist_user_defined_timestamps=*/true, /*is_user_key=*/false, - /*use_separated_kv_storage=*/false, /*statistics=*/nullptr, - /*uniform_cv_threshold=*/-1.0, /*use_common_prefix=*/false); - - // Convert user keys to InternalKeys and add to the index block. - std::vector internal_keys; - internal_keys.reserve(num_keys); - for (size_t i = 0; i < num_keys; i++) { - std::string ikey = keys[i]; - AppendInternalKeyFooter(&ikey, /*s=*/0, kTypeValue); - internal_keys.push_back(ikey); - - BlockHandle handle(i * 4096, 4096); - IndexValue entry(handle, Slice()); - std::string encoded_value; - entry.EncodeTo(&encoded_value, /*have_first_key=*/false, - /*previous_handle=*/nullptr); - index_builder.Add(Slice(internal_keys.back()), Slice(encoded_value)); - } - - // Finish the block and create a Block reader. - Slice raw_block = index_builder.Finish(); - std::string block_data(raw_block.data(), raw_block.size()); - BlockContents contents; - contents.data = Slice(block_data); - Block index_block(std::move(contents)); - - // ---- Generate random lookup targets ---- - auto indices = GenerateRandomIndices(kNumLookups, num_keys); - - // Pre-build InternalKey seek targets (user_key + kMaxSequenceNumber). - // Both sides use the same InternalKey targets -- the trie path parses - // them to extract the user key (matching production), while - // IndexBlockIter uses them directly. - std::vector seek_ikeys; - seek_ikeys.reserve(kNumLookups); - for (auto idx : indices) { - std::string sk = keys[idx]; - AppendInternalKeyFooter(&sk, kMaxSequenceNumber, kValueTypeForSeek); - seek_ikeys.push_back(std::move(sk)); - } - std::vector seek_ikey_slices; - seek_ikey_slices.reserve(kNumLookups); - for (const auto& sk : seek_ikeys) { - seek_ikey_slices.emplace_back(sk); - } - - // ---- Benchmark: Trie Seek (full production path) ---- - // Replicates IndexFactoryIteratorWrapper::Seek() -> - // TrieIndexIterator::SeekAndGetResult(): - // 1. ParseInternalKey to extract user_key - // 2. trie_iter.Seek(user_key) - // 3. trie_iter.Key().ToString() -- string copy (result must outlive - // call) - // 4. trie_iter.Value() -- look up BlockHandle from packed arrays - volatile uint64_t trie_checksum = 0; - std::string key_scratch; // Reused scratch, like current_key_scratch_ - auto t0 = std::chrono::high_resolution_clock::now(); - for (size_t i = 0; i < kNumLookups; i++) { - // Step 1: Parse InternalKey -> user key (same as wrapper) - ParsedInternalKey pkey; - ASSERT_OK(ParseInternalKey(seek_ikey_slices[i], &pkey, - /*log_err_key=*/false)); - - // Step 2: Seek on the trie with user key - trie_iter.Seek(pkey.user_key); - - if (trie_iter.Valid()) { - // Step 3: Materialize the key (string copy, matches production) - key_scratch = trie_iter.Key().ToString(); - - // Step 4: Get BlockHandle (two uint32 array reads) - TrieBlockHandle handle = trie_iter.Value(); - trie_checksum = trie_checksum + handle.offset; - } - } - auto t1 = std::chrono::high_resolution_clock::now(); - double trie_ns = - static_cast( - std::chrono::duration_cast(t1 - t0) - .count()) / - static_cast(kNumLookups); - - // ---- Benchmark: Real IndexBlockIter Seek ---- - // This is exactly what production does: Seek(InternalKey) -> value(). - std::unique_ptr ibi_iter(index_block.NewIndexIterator( - BytewiseComparator(), kDisableGlobalSequenceNumber, - /*iter=*/nullptr, /*stats=*/nullptr, - /*total_order_seek=*/true, /*have_first_key=*/false, - /*key_includes_seq=*/true, /*value_is_full=*/true, - /*block_contents_pinned=*/false, - /*user_defined_timestamps_persisted=*/true, /*prefix_index=*/nullptr, - BlockBasedTableOptions::kBinary, /*value_delta_escape=*/false)); - - volatile uint64_t ibi_checksum = 0; - auto t2 = std::chrono::high_resolution_clock::now(); - for (size_t i = 0; i < kNumLookups; i++) { - ibi_iter->Seek(seek_ikey_slices[i]); - if (ibi_iter->Valid()) { - IndexValue val = ibi_iter->value(); - ibi_checksum = ibi_checksum + val.handle.offset(); - } - } - auto t3 = std::chrono::high_resolution_clock::now(); - double ibi_ns = - static_cast( - std::chrono::duration_cast(t3 - t2) - .count()) / - static_cast(kNumLookups); - - double vs_ibi = (ibi_ns - trie_ns) / ibi_ns * 100.0; - const char* winner = trie_ns <= ibi_ns ? "TRIE" : "IBI"; - fprintf(stderr, " %8zu %10.1f %10.1f %+6.1f%% %s\n", num_keys, - trie_ns, ibi_ns, vs_ibi, winner); - } - - fprintf(stderr, "\n"); -} - // ============================================================================ // Mixed key type tests -- verifies UDI works with Delete, Merge, etc. // ============================================================================ @@ -5275,28 +2905,12 @@ TEST_F(TrieIndexFactoryTest, WrapperNextAndGetResultReturnsInternalKey) { EXPECT_FALSE(valid); } -// Verifies that overflow blocks are BFS-reordered alongside primary handles. -// If overflow blocks were stored in key-sorted order instead of BFS order, -// the overflow_base_ prefix sum would map overflow blocks to the wrong -// leaves when separator keys have different lengths. -// -// Key design: -// Trie entries: "ab"(bc=2), "ac"(bc=1), "b"(bc=2), "c"(bc=1), "e"(bc=1) -// -// Trie structure: -// Root: ['a'(internal), 'b'(leaf), 'c'(leaf), 'e'(leaf)] -// Level 1 under 'a': ['b'(leaf="ab"), 'c'(leaf="ac")] -// -// BFS leaf order: "b"(0), "c"(1), "e"(2), "ab"(3), "ac"(4) -// Key-sorted order: "ab"(0), "ac"(1), "b"(2), "c"(3), "e"(4) -// -// Both "ab" and "b" have overflow runs. Key-sorted overflow is -// ["ab"-overflow, "b"-overflow]. BFS-reordered overflow must be -// ["b"-overflow, "ab"-overflow]. -// -// Without fix: Seek("b") at low seqno gets "ab"'s overflow data (wrong -// offset and seqno), while Seek("ab") at low seqno gets "b"'s overflow. -TEST_F(TrieIndexFactoryTest, OverflowBfsReordering) { +// Runs on separators of different lengths ("ab" and "b") map every block to +// its own handle in key order, and Seek picks the right block of each run by +// seqno. In each run, the second block's last key is a prefix of the next +// block's first key ("ab" of "abc", "b" of "ba"), so FindShortestSeparator +// leaves it unchanged and the block joins the run. +TEST_F(TrieIndexFactoryTest, RunsOnPrefixKeysKeepBlockOrder) { IndexFactoryOptions option; option.comparator = BytewiseComparator(); @@ -5327,7 +2941,7 @@ TEST_F(TrieIndexFactoryTest, OverflowBfsReordering) { // Block 2: last="abc", next="b" // FindShortestSeparator("abc","b"): diff_index=0, 'a' vs 'b', // limit.size()-1=0 so fallback: increment start[1] 'b'->'c' -> "ac" - // -> sep="ac", different key, seqno=kMax->0 + // -> sep="ac", different key, tag 0 { IndexFactoryBuilder::BlockHandle h{200, 100}; Slice next("b"); @@ -5356,7 +2970,7 @@ TEST_F(TrieIndexFactoryTest, OverflowBfsReordering) { // Block 5: last="ba", next="d" // FindShortestSeparator("ba","d"): diff_index=0, 'b' vs 'd', // start_byte+1='c' < 'd' -> increment and truncate -> "c" - // -> sep="c", different key, seqno=kMax->0 + // -> sep="c", different key, tag 0 { IndexFactoryBuilder::BlockHandle h{500, 100}; Slice next("d"); @@ -5365,25 +2979,20 @@ TEST_F(TrieIndexFactoryTest, OverflowBfsReordering) { ASSERT_EQ(scratch, "c") << "Block 5 separator"; } // Block 6: last="d", next=null (last block, no successor shortening) - // -> sep="d", different from prev "c", seqno=kMax->0 + // -> sep="d", different from prev "c", last block keeps its tag, seqno=100 { IndexFactoryBuilder::BlockHandle h{600, 100}; sep = builder->AddIndexEntry(Slice("d"), nullptr, h, &scratch, EntryCtx(100, 100)); } - // After Finish(), trie entries (key-sorted): - // ki=0: "ab" bc=2, primary={offset=0,seqno=500}, - // overflow=[{offset=100,seqno=400}] - // ki=1: "ac" bc=1, {offset=200,seqno=0} - // ki=2: "b" bc=2, primary={offset=300,seqno=300}, - // overflow=[{offset=400,seqno=200}] - // ki=3: "c" bc=1, {offset=500,seqno=0} - // ki=4: "d" bc=1, {offset=600,seqno=0} - // - // BFS leaf order: "b"(0), "c"(1), "d"(2), "ab"(3), "ac"(4) - // BFS-reordered overflow: [{offset=400,seqno=200}, {offset=100,seqno=400}] - // overflow_base_: leaf0("b")=0, leaf3("ab")=1 + // After Finish(), separators in key order with their blocks + // {offset, tag seqno}: + // "ab": {0, 500}, {100, 400} + // "ac": {200, 0} + // "b": {300, 300}, {400, 200} + // "c": {500, 0} + // "d": {600, 100} Slice index_contents; ASSERT_OK(builder->Finish(&index_contents)); @@ -5396,79 +3005,73 @@ TEST_F(TrieIndexFactoryTest, OverflowBfsReordering) { IterateResult result; // --- Full forward scan: verify all block offsets in key order --- - // Expected order: ab(0), ab-overflow(100), ac(200), b(300), - // b-overflow(400), c(500), e(600) + // Expected offsets: ab 0, ab 100, ac 200, b 300, b 400, c 500, d 600 ASSERT_OK(iter->SeekToFirstAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "ab"); - EXPECT_EQ(iter->value().offset, 0u) << "SeekToFirst: ab primary"; + EXPECT_EQ(iter->value().offset, 0u) << "SeekToFirst: ab first block"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "ab"); - EXPECT_EQ(iter->value().offset, 100u) << "Next: ab overflow"; + EXPECT_EQ(iter->value().offset, 100u) << "Next: ab second block"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "ac"); - EXPECT_EQ(iter->value().offset, 200u) << "Next: ac primary"; + EXPECT_EQ(iter->value().offset, 200u) << "Next: ac"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "b"); - EXPECT_EQ(iter->value().offset, 300u) << "Next: b primary"; + EXPECT_EQ(iter->value().offset, 300u) << "Next: b first block"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "b"); - EXPECT_EQ(iter->value().offset, 400u) << "Next: b overflow"; + EXPECT_EQ(iter->value().offset, 400u) << "Next: b second block"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "c"); - EXPECT_EQ(iter->value().offset, 500u) << "Next: c primary"; + EXPECT_EQ(iter->value().offset, 500u) << "Next: c"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.key.ToString(), "d"); - EXPECT_EQ(iter->value().offset, 600u) << "Next: d primary"; + EXPECT_EQ(iter->value().offset, 600u) << "Next: d"; ASSERT_OK(iter->NextAndGetResult(&result)); EXPECT_EQ(result.bound_check_result, IterBoundCheck::kUnknown) << "past end"; - // --- Seek-based tests: verify overflow data is correctly associated --- + // --- Seek-based tests: verify each run's tags select the right block --- - // "ab" primary: Seek("ab", kMax) -> offset=0 + // "ab" first block: Seek("ab", kMax) -> offset=0 ASSERT_OK(iter->SeekAndGetResult(Slice("ab"), &result, SeekCtx(kMaxSequenceNumber))); EXPECT_EQ(result.bound_check_result, IterBoundCheck::kInbound); - EXPECT_EQ(iter->value().offset, 0u) << "Seek(ab,kMax): ab primary"; + EXPECT_EQ(iter->value().offset, 0u) << "Seek(ab,kMax): ab first block"; - // "ab" overflow: Seek("ab", 400) -> offset=100 - // Primary seqno=500, 400<500 -> advance to overflow. Overflow seqno=400, - // 400>=400 -> match. + // "ab" second block: Seek("ab", 400) -> offset=100, the first block of the + // run whose tag is at most 400. ASSERT_OK(iter->SeekAndGetResult(Slice("ab"), &result, SeekCtx(400))); EXPECT_EQ(result.bound_check_result, IterBoundCheck::kInbound); - EXPECT_EQ(iter->value().offset, 100u) << "Seek(ab,400): ab overflow"; + EXPECT_EQ(iter->value().offset, 100u) << "Seek(ab,400): ab second block"; // "ab" advance past run: Seek("ab", 50) -> advances to "ac" at offset=200 - // Primary seqno=500, 50<500 -> overflow seqno=400, 50<400 -> exhaust run -> - // advance to next trie leaf "ac". + // Both tags of the run are larger than 50, so Seek moves to the next + // separator. ASSERT_OK(iter->SeekAndGetResult(Slice("ab"), &result, SeekCtx(50))); EXPECT_EQ(result.key.ToString(), "ac") << "Seek(ab,50): should advance past ab run"; EXPECT_EQ(iter->value().offset, 200u) << "Seek(ab,50): expected ac (200)"; - // "b" primary: Seek("b", kMax) -> offset=300 + // "b" first block: Seek("b", kMax) -> offset=300 ASSERT_OK( iter->SeekAndGetResult(Slice("b"), &result, SeekCtx(kMaxSequenceNumber))); EXPECT_EQ(result.bound_check_result, IterBoundCheck::kInbound); - EXPECT_EQ(iter->value().offset, 300u) << "Seek(b,kMax): b primary"; - - // "b" overflow: Seek("b", 200) -> offset=400 - // Primary seqno=300, 200<300 -> advance to overflow. Overflow seqno=200, - // 200>=200 -> match. - // If overflow blocks were NOT BFS-reordered, overflow[0] would contain - // ab's data {offset=100,seqno=400}, and 200<400 would fail to match, - // incorrectly advancing to "c". + EXPECT_EQ(iter->value().offset, 300u) << "Seek(b,kMax): b first block"; + + // "b" second block: Seek("b", 200) -> offset=400, the first block of the + // run whose tag is at most 200. Seek must not move on to "c". ASSERT_OK(iter->SeekAndGetResult(Slice("b"), &result, SeekCtx(200))); EXPECT_EQ(result.bound_check_result, IterBoundCheck::kInbound); EXPECT_EQ(result.key.ToString(), "b") << "Seek(b,200): must stay on b, not advance"; - EXPECT_EQ(iter->value().offset, 400u) << "Seek(b,200): b overflow"; + EXPECT_EQ(iter->value().offset, 400u) << "Seek(b,200): b second block"; // "b" advance past run: Seek("b", 50) -> advances to "c" at offset=500 ASSERT_OK(iter->SeekAndGetResult(Slice("b"), &result, SeekCtx(50))); @@ -5478,15 +3081,9 @@ TEST_F(TrieIndexFactoryTest, OverflowBfsReordering) { } TEST_F(TrieIndexFactoryTest, SeekToLastPrefixKeyOnlyTrie) { - // DescendToRightmostLeaf skips prefix keys (they are the smallest leaf at - // a node). This test verifies that SeekToLast correctly lands on the - // rightmost regular leaf ("ab"), and Prev correctly visits the prefix key - // ("a") during backtracking. - // - // This trie has separators "a" and "ab". Separator "a" is a prefix of "ab", - // creating a prefix key at the "a" node. When we SeekToLast, we should - // land on "ab" (the rightmost leaf), and Prev from there should land on - // "a" (the prefix key). + // SeekToLast and Prev when one separator is a prefix of the other. The + // separators are "a" and "ab", so SeekToLast must land on "ab", the + // largest separator, and Prev from there must land on "a". auto ctx = BuildTrieAndGetIterator({ {"a", "ab", 0, 500, 0, 0}, {"ab", "", 500, 500, 0, 0}, @@ -5499,7 +3096,7 @@ TEST_F(TrieIndexFactoryTest, SeekToLastPrefixKeyOnlyTrie) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(result.key.ToString(), "ab"); - // Prev should land on "a" (the prefix key). + // Prev should land on "a" (a prefix of "ab"). ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(result.key.ToString(), "a"); @@ -5510,9 +3107,9 @@ TEST_F(TrieIndexFactoryTest, SeekToLastPrefixKeyOnlyTrie) { } TEST_F(TrieIndexFactoryTest, SeekToLastAndPrevWithPrefixKeys) { - // Verifies SeekToLast and Prev correctly handle prefix keys in a trie - // with multiple levels. Separators "a", "ab", "abc", "b" create prefix - // keys at the "a" and "ab" nodes. + // Verifies SeekToLast and Prev when separators are prefixes of later ones. + // The last keys "a" and "ab" are prefixes of the next block's first key, + // so they stay unshortened: the separators are "a", "ab", "ac" and "b". auto ctx = BuildTrieAndGetIterator({ {"a", "ab", 0, 500, 0, 0}, {"ab", "abc", 500, 500, 0, 0}, @@ -5522,7 +3119,7 @@ TEST_F(TrieIndexFactoryTest, SeekToLastAndPrevWithPrefixKeys) { IterateResult result; - // SeekToLast should land on "b" (rightmost leaf). + // SeekToLast should land on "b" (the largest separator). ASSERT_OK(ctx.iter->SeekToLastAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(result.key.ToString(), "b"); @@ -5566,18 +3163,20 @@ TEST_F(TrieIndexFactoryTest, LastBlockSeekWithRealSeqno) { ASSERT_NO_FATAL_FAILURE(AssertSeekOffset(ctx.iter.get(), Slice("cherry"), kMaxSequenceNumber, 1000)); - // Seek "cherry" with a real seqno (200) -- the last block's separator has + // Seek "cherry" with a real seqno (200) -- the last block's tag has // seqno=50. In internal key order, 200 > 50 means target is SMALLER (higher - // seqno = smaller key). So target <= separator -> should stay on this block. + // seqno = smaller key). So target <= the block's last key -> should stay on + // this block. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("cherry"), 200, 1000)); - // Seek "cherry" with seqno=50 (equal to separator) -- should stay. + // Seek "cherry" with seqno=50 (equal to the block's tag) -- should stay. ASSERT_NO_FATAL_FAILURE( AssertSeekOffset(ctx.iter.get(), Slice("cherry"), 50, 1000)); - // Seek "cherry" with seqno=10 (less than separator seqno=50) -- target is - // LARGER in internal key order. target > separator -> should advance past. + // Seek "cherry" with seqno=10 (less than the tag's seqno=50) -- target is + // LARGER in internal key order. target > the block's last key -> should + // advance past. // But this is the last block -- no next block -> past end. IterateResult result; ASSERT_OK(ctx.iter->SeekAndGetResult(Slice("cherry"), &result, SeekCtx(10))); @@ -5620,7 +3219,6 @@ TEST_F(TrieIndexFactoryTest, NonBoundarySeparatorSeekWhenShorteningFails) { // on separator "acc" with tag=0 (non-boundary sentinel). The post-seek // correction must NOT advance past it regardless of the target seqno. // - // Seqno encoding is always active for a non-empty trie. // The key arrangement: // Block 0: separator = FindShortestSeparator("aaa","acc") = "ab" // Block 1: separator = FindShortestSeparator("acc","acd") = "acc" @@ -5654,11 +3252,13 @@ TEST_F(TrieIndexFactoryTest, NonBoundarySeparatorSeekWhenShorteningFails) { AssertFullForwardScan(ctx.iter.get(), Slice("aaa"), {0, 1000, 2000}); } -TEST_F(TrieIndexFactoryTest, PrevWithinOverflowRun) { - // Exercises PrevAndGetResult when positioned mid-overflow: the fast path - // that decrements overflow_run_index_ without calling iter_.Prev(). +TEST_F(TrieIndexFactoryTest, PrevWithinRun) { + // Exercises PrevAndGetResult into and within a run: Prev from the block + // after a run lands on the run's last block, and Prev within the run walks + // back to its first block. auto ctx = BuildTrieAndGetIterator({ - // Block 0-2: "key" spans 3 blocks (overflow run). + // Blocks 0-2: user key "key" spans 3 blocks. Blocks 0-1 share + // separator "key" (a run); block 2 gets separator "l". {"key", "key", 0, 1000, 300, 200}, {"key", "key", 1000, 1000, 200, 100}, {"key", "zzz", 2000, 1000, 100, 50}, @@ -5668,12 +3268,12 @@ TEST_F(TrieIndexFactoryTest, PrevWithinOverflowRun) { IterateResult result; - // Seek to "key"|100 -- should land on the last block in the "key" overflow - // run (offset 2000) since seqno=100 matches that block's tag. + // Seek to "key"|100 -- both blocks of the "key" run have larger tags (300, + // 200), so Seek moves to the next separator, "l" (block 2, offset 2000). ASSERT_OK(ctx.iter->SeekAndGetResult(Slice("key"), &result, SeekCtx(100))); ASSERT_EQ(ctx.iter->value().offset, 2000u); - // Prev should go back to the middle of the overflow run (offset 1000). + // Prev should go back to the last block of the "key" run (offset 1000). ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 1000u); @@ -5688,14 +3288,13 @@ TEST_F(TrieIndexFactoryTest, PrevWithinOverflowRun) { ASSERT_NE(result.bound_check_result, IterBoundCheck::kInbound); } -TEST_F(TrieIndexFactoryTest, SeekToLastWithOverflowRun) { - // SeekToLast on a trie where the last leaf has block_count > 1 should - // position at the last overflow block. +TEST_F(TrieIndexFactoryTest, SeekToLastWithRun) { + // SeekToLast when the last separator has a run of several blocks should + // position at the last block of the run. auto ctx = BuildTrieAndGetIterator({ // Block 0: "aaa" standalone. {"aaa", "zzz", 0, 1000, 100, 50}, - // Block 1-3: "zzz" spans 3 blocks (overflow run) -- this is the last - // leaf. + // Block 1-3: "zzz" spans 3 blocks, a run on the last separator. {"zzz", "zzz", 1000, 1000, 50, 30}, {"zzz", "zzz", 2000, 1000, 30, 10}, {"zzz", "", 3000, 1000, 10, 0}, @@ -5703,23 +3302,23 @@ TEST_F(TrieIndexFactoryTest, SeekToLastWithOverflowRun) { IterateResult result; - // SeekToLast should land on the last overflow block (offset 3000). + // SeekToLast should land on the last block of the run (offset 3000). ASSERT_OK(ctx.iter->SeekToLastAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(result.key.ToString(), "zzz"); ASSERT_EQ(ctx.iter->value().offset, 3000u); - // Prev within overflow: offset 2000. + // Prev within the run: offset 2000. ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 2000u); - // Prev within overflow: offset 1000. + // Prev within the run: offset 1000. ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 1000u); - // Prev to previous trie leaf: "aaa" at offset 0. + // Prev to the previous separator: block 0 at offset 0. ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 0u); @@ -5729,9 +3328,10 @@ TEST_F(TrieIndexFactoryTest, SeekToLastWithOverflowRun) { ASSERT_NE(result.bound_check_result, IterBoundCheck::kInbound); } -TEST_F(TrieIndexFactoryTest, PrevLandsOnLeafWithOverflow) { - // Prev from a standalone leaf should land on a previous leaf that has an - // overflow run, positioning at the last block in that run. +TEST_F(TrieIndexFactoryTest, PrevAcrossSameUserKeyBlocks) { + // Prev from a standalone block steps back one block at a time, including + // across the two blocks of user key "aaa". Only block 0 keeps separator + // "aaa"; block 1, the last "aaa" block, gets the shortened separator "b". auto ctx = BuildTrieAndGetIterator({ // Block 0-1: "aaa" spans 2 blocks. {"aaa", "aaa", 0, 1000, 200, 100}, @@ -5753,13 +3353,12 @@ TEST_F(TrieIndexFactoryTest, PrevLandsOnLeafWithOverflow) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 2000u); - // Prev to "aaa" -- should land on the LAST block in the overflow run - // (offset 1000), not the primary block (offset 0). + // Prev to block 1 (offset 1000), the last "aaa" block, not block 0. ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 1000u); - // Prev within "aaa" overflow: offset 0. + // Prev to block 0: offset 0. ASSERT_OK(ctx.iter->PrevAndGetResult(&result)); ASSERT_EQ(result.bound_check_result, IterBoundCheck::kInbound); ASSERT_EQ(ctx.iter->value().offset, 0u); @@ -5805,10 +3404,10 @@ TEST_F(TrieIndexFactoryTest, SeekToLastOnEmptyTrie) { ASSERT_EQ(result.bound_check_result, IterBoundCheck::kUnknown); } -TEST_F(TrieIndexFactoryTest, OverflowExhaustionThenForwardScanThroughOverflow) { - // When SeekAndGetResult exhausts an overflow run and advances to the next - // trie leaf, verify that subsequent Next() calls correctly traverse through - // any later overflow runs. +TEST_F(TrieIndexFactoryTest, RunExhaustionThenScanThroughRun) { + // When SeekAndGetResult moves past a run because every tag in it is + // larger than the target, verify that subsequent Next() calls correctly + // walk through a later run. // // Layout: // Blocks 0-1: "key" same-key run (2 blocks) @@ -5824,19 +3423,19 @@ TEST_F(TrieIndexFactoryTest, OverflowExhaustionThenForwardScanThroughOverflow) { IterateResult result; - // Seek "key"|1 -- seqno=1 is below all overflow tags in the "key" run - // (300, 200), so it exhausts the run and advances to the next leaf. - // FindShortestSeparator("key", "zzz") = "l", so the next leaf is "l". + // Seek "key"|1 -- seqno=1 is below both tags of the "key" run (300, 200), + // so Seek moves to the next separator. + // FindShortestSeparator("key", "zzz") = "l", so that separator is "l". ASSERT_OK(ctx.iter->SeekAndGetResult(Slice("key"), &result, SeekCtx(1))); ASSERT_EQ(result.key.ToString(), "l"); ASSERT_EQ(ctx.iter->value().offset, 2000u); - // Next advances to "zzz" primary (offset 3000). + // Next advances to the first "zzz" block (offset 3000). ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(result.key.ToString(), "zzz"); ASSERT_EQ(ctx.iter->value().offset, 3000u); - // Next advances through "zzz" overflow (offset 4000). + // Next advances to the second "zzz" block (offset 4000). ASSERT_OK(ctx.iter->NextAndGetResult(&result)); ASSERT_EQ(ctx.iter->value().offset, 4000u); @@ -5922,9 +3521,8 @@ TEST_F(TrieIndexFactoryTest, ParallelMatchesSerialOutput) { {"bbb", "ccc", 210, 100, 60, 50}, {"ccc", "ddd", 315, 100, 55, 45}, // Four blocks split inside user key "kkk". The first three all get - // separator "kkk", so Finish() has to collapse them into one - // multi-block run with a seqno side-table entry per block. A run - // longer than two is the only way that path gets exercised. + // separator "kkk" and a tag each, forming a run of three blocks, + // which both paths must build identically. The fourth gets "l". {"kkk", "kkk", 420, 100, 44, 40}, {"kkk", "kkk", 525, 100, 38, 34}, {"kkk", "kkk", 630, 100, 32, 28}, @@ -6138,8 +3736,8 @@ TEST_F(TrieIndexFactoryTest, ParallelPipelineOverlapsPrepareAndFinish) { } TEST_F(TrieIndexFactoryTest, ParallelEstimatedSizeStaysMonotonic) { - // EstimatedSize() uses state updated by emit-thread PrepareAddEntry calls, - // disjoint from state committed by FinishAddEntry on the BG worker. + // EstimatedSize() runs on the emit thread while FinishAddEntry commits + // entries on the BG worker, and never decreases. IndexFactoryOptions option; option.comparator = BytewiseComparator(); std::unique_ptr builder; @@ -6166,9 +3764,10 @@ TEST_F(TrieIndexFactoryTest, ParallelEstimatedSizeStaysMonotonic) { std::string scratch; builder->FinishAddEntry(h, p.get(), &scratch, /*skip_delta_encoding=*/false); - EXPECT_EQ(builder->EstimatedSize(), prepared_size) - << "BG-thread FinishAddEntry must not mutate estimation state"; - last_size = prepared_size; + const uint64_t finished_size = builder->EstimatedSize(); + EXPECT_GE(finished_size, prepared_size) + << "EstimatedSize should not shrink after FinishAddEntry"; + last_size = finished_size; } } @@ -6216,6 +3815,795 @@ TEST_F(TrieIndexFactoryTest, ParallelFinishWithoutPrepareFailsBuild) { } } +// ============================================================================ +// Trie +// ============================================================================ + +class TrieTest : public testing::Test { + protected: + struct Block { + std::string separator; + TrieBlockHandle handle; + uint64_t tag; + }; + + // Tag of a Seek target that lands on the first block of its key. + const uint64_t kSeekTag = + PackSequenceAndType(kMaxSequenceNumber, kValueTypeForSeek); + + static uint64_t NumBlocks(const Trie& index) { + return index.Trie::NumBlocks(); + } + + // Serializes blocks with the given bucket size, failing the test on error. + static std::string Build(const std::vector& blocks, + uint32_t bucket_size) { + TrieBuilder builder(bucket_size); + for (const Block& b : blocks) { + builder.AddBlock(b.separator, b.handle, b.tag); + } + std::string data; + Status s = builder.Finish(&data); + EXPECT_OK(s); + return data; + } + + // One block per separator, laid out back to back with random sizes and a + // 5-byte gap between blocks. + static std::vector ContiguousBlocks( + const std::vector& separators, Random* rnd) { + std::vector blocks; + uint64_t offset = rnd->Uniform(1000); + for (const std::string& sep : separators) { + const uint64_t size = 100 + rnd->Uniform(4000); + blocks.push_back({sep, {offset, size}, 0}); + offset += size + 5; + } + return blocks; + } + + // Blocks shaped like the factory's output: some separators span several + // blocks with descending tags (the last block of such a run may have none), + // and gaps are either uniform (derived sizes) or not (explicit sizes), + // sometimes starting past 4 GiB. + static std::vector RandomBlocks( + const std::vector& separators, Random* rnd) { + std::vector blocks; + const bool uniform_gaps = rnd->OneIn(2); + uint64_t offset = rnd->OneIn(4) ? uint64_t{1} << (32 + rnd->Uniform(20)) + : rnd->Uniform(1000); + for (const std::string& sep : separators) { + const uint32_t run = + rnd->OneIn(3) ? 2 + rnd->Uniform(rnd->OneIn(5) ? 40 : 5) : 1; + uint64_t tag = + PackSequenceAndType(1000 + rnd->Uniform(1 << 30), kTypeValue); + for (uint32_t j = 0; j < run; j++) { + const bool tagged = j + 1 < run || rnd->OneIn(3); + const uint64_t size = 1 + rnd->Uniform(5000); + blocks.push_back({sep, {offset, size}, tagged ? tag : 0}); + offset += size + (uniform_gaps ? 5 : rnd->Uniform(20)); + tag -= 1 + rnd->Uniform(256); + } + } + return blocks; + } + + // Sorted distinct non-empty keys from one of several shapes: a tiny + // alphabet (deep shared prefixes, keys that prefix other keys), full-range + // bytes (wide fanout), and either of those behind a long common prefix. + static std::vector RandomKeys(Random* rnd, size_t n) { + const uint32_t shape = rnd->Uniform(4); + const std::string common = shape >= 2 ? rnd->RandomString(30) : ""; + std::set keys; + while (keys.size() < n) { + std::string k = common; + const size_t len = 1 + rnd->Uniform(shape % 2 == 0 ? 12 : 3); + for (size_t i = 0; i < len; i++) { + k.push_back(shape % 2 == 0 + ? static_cast("ab\0\xff"[rnd->Uniform(4)]) + : static_cast(rnd->Uniform(256))); + } + keys.insert(k); + } + return {keys.begin(), keys.end()}; + } + + // Seek targets around each separator plus random strings. + static std::vector SeekTargets( + const std::vector& separators, Random* rnd) { + std::vector targets = {""}; + for (const std::string& sep : separators) { + targets.push_back(sep); + targets.push_back(sep + std::string(1, '\0')); + targets.push_back(sep + "a"); + targets.push_back( + sep.substr(0, rnd->Uniform(static_cast(sep.size()) + 1))); + if (!sep.empty()) { + std::string bumped = sep; + bumped.back() = + static_cast(bumped.back() + (rnd->OneIn(2) ? 1 : -1)); + targets.push_back(bumped); + } + } + for (int i = 0; i < 100; i++) { + targets.push_back(rnd->RandomBinaryString(rnd->Uniform(40))); + } + return targets; + } + + // Index of the first block at or after (target, tag) in internal key + // order, or blocks.size(). + static size_t ExpectedSeek(const std::vector& blocks, + const std::string& target, uint64_t tag) { + for (size_t i = 0; i < blocks.size(); i++) { + const int cmp = Slice(blocks[i].separator).compare(target); + if (cmp > 0 || + (cmp == 0 && (blocks[i].tag == 0 || tag >= blocks[i].tag))) { + return i; + } + } + return blocks.size(); + } + + static void ExpectAt(const TrieIterator& it, const std::vector& blocks, + size_t i) { + if (i == blocks.size()) { + ASSERT_FALSE(it.Valid()); + return; + } + ASSERT_TRUE(it.Valid()); + ASSERT_EQ(it.key(), blocks[i].separator); + ASSERT_EQ(it.value().offset, blocks[i].handle.offset); + ASSERT_EQ(it.value().size, blocks[i].handle.size); + } + + // Seek tags for user-key-only lookups and for an old snapshot. + static std::vector BasicSeekTags() { + return {PackSequenceAndType(kMaxSequenceNumber, kValueTypeForSeek), + PackSequenceAndType(1, kValueTypeForSeek)}; + } + + // Checks a forward scan and Seek to every (target, tag) against the + // blocks. + static void Verify(const std::vector& blocks, uint32_t bucket_size, + const std::vector& targets, + const std::vector& tags) { + SCOPED_TRACE("bucket_size=" + std::to_string(bucket_size)); + std::string data = Build(blocks, bucket_size); + Trie index; + ASSERT_OK(index.Init(data)); + ASSERT_EQ(NumBlocks(index), blocks.size()); + TrieIterator it(&index); + ASSERT_OK(it.SeekToFirst()); + for (size_t i = 0; i <= blocks.size(); i++) { + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, i)); + ASSERT_OK(it.Next()); + } + ASSERT_OK(it.SeekToLast()); + for (size_t i = blocks.size(); i-- > 0;) { + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, i)); + ASSERT_OK(it.Prev()); + } + ASSERT_FALSE(it.Valid()); + for (const std::string& target : targets) { + for (uint64_t tag : tags) { + SCOPED_TRACE(Slice(target).ToString(true) + " tag " + + std::to_string(tag)); + ASSERT_OK(it.Seek(target, tag)); + const size_t expected = ExpectedSeek(blocks, target, tag); + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, expected)); + if (expected == blocks.size()) { + continue; + } + // Step back and forth from the Seek position. + ASSERT_OK(it.Prev()); + if (expected == 0) { + ASSERT_FALSE(it.Valid()); + continue; + } + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, expected - 1)); + ASSERT_OK(it.Next()); + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, expected)); + } + } + } + + // Seeks to random (target, tag) pairs, then walks Next and Prev at random + // across runs and bucket boundaries, checking every position. + static void VerifyWalks(const std::vector& blocks, + uint32_t bucket_size, + const std::vector& targets, + const std::vector& tags, Random* rnd) { + SCOPED_TRACE("bucket_size=" + std::to_string(bucket_size)); + std::string data = Build(blocks, bucket_size); + Trie index; + ASSERT_OK(index.Init(data)); + TrieIterator it(&index); + for (int i = 0; i < 200; i++) { + const std::string& target = + targets[rnd->Uniform(static_cast(targets.size()))]; + const uint64_t tag = tags[rnd->Uniform(static_cast(tags.size()))]; + SCOPED_TRACE(Slice(target).ToString(true) + " tag " + + std::to_string(tag)); + ASSERT_OK(it.Seek(target, tag)); + size_t pos = ExpectedSeek(blocks, target, tag); + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, pos)); + for (int step = 0; step < 30 && pos < blocks.size(); step++) { + if (rnd->OneIn(2)) { + ASSERT_OK(it.Next()); + pos++; + } else { + ASSERT_OK(it.Prev()); + pos = pos == 0 ? blocks.size() : pos - 1; + } + ASSERT_NO_FATAL_FAILURE(ExpectAt(it, blocks, pos)); + } + } + } +}; + +TEST_F(TrieTest, EmptySeparator) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + std::vector> cases = { + {""}, + {"", std::string(1, '\0')}, + {"", "a"}, + {"", "a", "ab", "b"}, + }; + for (int i = 0; i < 10; i++) { + std::vector keys = RandomKeys(&rnd, 1 + rnd.Uniform(100)); + keys.insert(keys.begin(), ""); + cases.push_back(keys); + } + for (const auto& keys : cases) { + std::vector blocks = ContiguousBlocks(keys, &rnd); + std::vector targets = SeekTargets(keys, &rnd); + for (uint32_t bucket_size : {1u, 2u, 8u}) { + ASSERT_NO_FATAL_FAILURE( + Verify(blocks, bucket_size, targets, BasicSeekTags())); + } + } +} + +// One user key spanning several blocks repeats its separator; each block of +// the run carries the tag of its last key, newest first. +TEST_F(TrieTest, SameKeyRunsAndTags) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + for (int round = 0; round < 30; round++) { + std::vector keys = RandomKeys(&rnd, 1 + rnd.Uniform(300)); + std::vector separators; + std::vector block_tags; + for (const std::string& key : keys) { + const uint32_t run = rnd.OneIn(3) ? 2 + rnd.Uniform(4) : 1; + SequenceNumber seq = 1000 + rnd.Uniform(1000); + for (uint32_t j = 0; j < run; j++) { + separators.push_back(key); + uint64_t tag = 0; + if (run > 1 ? (j + 1 < run || rnd.OneIn(2)) : rnd.OneIn(8)) { + tag = PackSequenceAndType(seq, kTypeValue); + } + block_tags.push_back(tag); + seq -= 1 + rnd.Uniform(100); + } + } + std::vector blocks = ContiguousBlocks(separators, &rnd); + std::vector tags = BasicSeekTags(); + tags.push_back(0); + for (size_t i = 0; i < blocks.size(); i++) { + blocks[i].tag = block_tags[i]; + if (block_tags[i] != 0) { + tags.push_back(block_tags[i]); + tags.push_back(block_tags[i] - 1); + tags.push_back(block_tags[i] + 1); + } + } + // Keep the cross product of targets and tags small. + std::vector targets = keys; + targets.emplace_back(); + std::sort(tags.begin(), tags.end()); + tags.erase(std::unique(tags.begin(), tags.end()), tags.end()); + if (tags.size() > 60) { + tags.resize(60); + } + for (uint32_t bucket_size : {1u, 3u, 8u}) { + ASSERT_NO_FATAL_FAILURE(Verify(blocks, bucket_size, targets, tags)); + } + } +} + +TEST_F(TrieTest, LargeHandlesAndIrregularGaps) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + std::vector keys = RandomKeys(&rnd, 500); + std::vector targets = SeekTargets(keys, &rnd); + + // Offsets past 4 GiB with a uniform gap between blocks. + std::vector high = ContiguousBlocks(keys, &rnd); + for (Block& b : high) { + b.handle.offset += uint64_t{5} << 40; + } + // Gaps of varying size, as with block alignment padding. + std::vector padded = ContiguousBlocks(keys, &rnd); + uint64_t offset = 0; + for (Block& b : padded) { + b.handle.offset = offset; + offset += b.handle.size + rnd.Uniform(4096); + } + // Blocks larger than 4 GiB. + std::vector huge = ContiguousBlocks(keys, &rnd); + offset = 0; + for (Block& b : huge) { + b.handle.offset = offset; + b.handle.size = (uint64_t{1} << 33) + rnd.Uniform(1 << 20); + offset += b.handle.size + 5; + } + for (const auto* blocks : {&high, &padded, &huge}) { + ASSERT_NO_FATAL_FAILURE(Verify(*blocks, 8, targets, BasicSeekTags())); + } +} + +// A handle whose end does not fit in 64 bits is stored as is, with explicit +// sizes, whether or not it is the last block. (Sizes still have to fit one +// column, so they are all near the maximum.) +TEST_F(TrieTest, HandleEndPastUint64) { + const uint64_t max = std::numeric_limits::max(); + const std::vector> cases = { + {{"a", {5, max}, 0}}, + {{"a", {0, max}, 0}, {"b", {1, max}, 0}}, + {{"a", {0, max}, 0}, {"b", {1, max - 1}, 0}}, + }; + for (const std::vector& blocks : cases) { + SCOPED_TRACE("num_blocks=" + std::to_string(blocks.size())); + ASSERT_NO_FATAL_FAILURE( + Verify(blocks, 8, {"", "a", "b", "c"}, BasicSeekTags())); + } +} + +// Column residuals are limited to 56 bits, far beyond any real SST. +TEST_F(TrieTest, ColumnRangeLimit) { + TrieBuilder builder(8); + builder.AddBlock("a", {0, 1}, 0); + builder.AddBlock("b", {1, 1}, 0); + builder.AddBlock("c", {uint64_t{1} << 60, 1}, 0); + std::string data; + ASSERT_EQ(builder.Finish(&data).code(), Status::Code::kNotSupported); +} + +TEST_F(TrieTest, NoBlocks) { + std::string data = Build({}, 8); + Trie index; + ASSERT_OK(index.Init(data)); + ASSERT_EQ(NumBlocks(index), 0u); + TrieIterator it(&index); + ASSERT_OK(it.SeekToFirst()); + ASSERT_FALSE(it.Valid()); + ASSERT_OK(it.SeekToLast()); + ASSERT_FALSE(it.Valid()); + ASSERT_OK(it.Seek("a", kSeekTag)); + ASSERT_FALSE(it.Valid()); +} + +// Version 1 was an earlier layout, so its indexes are rejected too. The +// version is checked before the header size, as other versions may have a +// shorter header. +TEST_F(TrieTest, RejectsOtherVersions) { + std::string data = Build({{"a", {0, 10}, 0}}, 8); + for (uint32_t version : {0u, 1u, 3u}) { + SCOPED_TRACE(version); + EncodeFixed32(&data[4], version); + for (size_t len : {data.size(), size_t{8}}) { + Trie index; + ASSERT_EQ(index.Init(Slice(data.data(), len)).code(), + Status::Code::kNotSupported); + } + } +} + +// Init() rejects a max_key_len that the separator bytes stored in the index +// cannot reach. +TEST_F(TrieTest, RejectsMaxKeyLenBeyondContent) { + const std::string data = + Build({{"apple", {0, 100}, 0}, {"apricot", {105, 100}, 0}}, 8); + Trie index; + ASSERT_OK(index.Init(data)); + for (uint32_t max_key_len : {uint32_t{1} << 20, uint32_t{0xFFFFFFFF}}) { + SCOPED_TRACE(max_key_len); + std::string mutated = data; + // Header field max_key_len. + EncodeFixed32(&mutated[36], max_key_len); + ASSERT_EQ(index.Init(mutated).code(), Status::Code::kCorruption); + } +} + +// The header's end offset is the index size, so bytes after it mean the +// header or the block boundaries are wrong. +TEST_F(TrieTest, RejectsBytesAfterEnd) { + const std::string data = + Build({{"apple", {0, 100}, 0}, {"apricot", {105, 100}, 0}}, 8); + Trie index; + ASSERT_OK(index.Init(data)); + for (size_t extra : {1, 8, 16, 17}) { + SCOPED_TRACE(extra); + const std::string longer = data + std::string(extra, '\0'); + ASSERT_EQ(index.Init(longer).code(), Status::Code::kCorruption); + } +} + +// Init() reads the header before it validates the sections, so a failure +// must not leave an iterator reading a half parsed index, nor a stale one +// from an earlier Init(). +TEST_F(TrieTest, FailedInitLeavesIndexEmpty) { + const std::string data = + Build({{"apple", {0, 100}, 0}, {"apricot", {105, 100}, 0}}, 8); + // The end offset no longer matches the size. + const std::string longer = data + '\0'; + for (bool reinit : {false, true}) { + SCOPED_TRACE(reinit); + Trie index; + if (reinit) { + ASSERT_OK(index.Init(data)); + } + ASSERT_EQ(index.Init(longer).code(), Status::Code::kCorruption); + TrieIterator it(&index); + ASSERT_OK(it.SeekToFirst()); + ASSERT_FALSE(it.Valid()); + ASSERT_OK(it.SeekToLast()); + ASSERT_FALSE(it.Valid()); + ASSERT_OK(it.Seek("apple", kSeekTag)); + ASSERT_FALSE(it.Valid()); + } +} + +TEST_F(TrieTest, SeparatorsOutOfOrder) { + TrieBuilder builder(8); + builder.AddBlock("b", {0, 10}, 0); + builder.AddBlock("a", {15, 10}, 0); + std::string data; + ASSERT_EQ(builder.Finish(&data).code(), Status::Code::kInvalidArgument); +} + +// A block much larger than the others widens only the residuals of its own +// group of 64 offsets. +TEST_F(TrieTest, OutlierBlockWidensOnlyItsGroup) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + std::vector keys; + keys.reserve(2048); + for (int i = 0; i < 2048; i++) { + keys.push_back("key" + std::to_string(10000 + i)); + } + std::vector blocks = ContiguousBlocks(keys, &rnd); + const size_t uniform_size = Build(blocks, 8).size(); + blocks[1000].handle.size += 1 << 20; + for (size_t i = 1001; i < blocks.size(); i++) { + blocks[i].handle.offset += 1 << 20; + } + ASSERT_LE(Build(blocks, 8).size(), uniform_size + 64 * 56 / 8); +} + +// Decreasing offsets would wrap a derived block size around. The writer +// refuses them. Init() does not read every offset, so the iterator reports +// Corruption when it reaches such a block. +TEST_F(TrieTest, DecreasingOffsets) { + TrieBuilder builder(8); + builder.AddBlock("a", {100, 10}, 0); + builder.AddBlock("b", {50, 10}, 0); + std::string data; + ASSERT_EQ(builder.Finish(&data).code(), Status::Code::kInvalidArgument); + + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + std::vector keys; + keys.reserve(128); + for (int i = 0; i < 128; i++) { + keys.push_back("key" + std::to_string(1000 + i)); + } + data = Build(ContiguousBlocks(keys, &rnd), 8); + // Zero the base of the second group of offsets, the header's fourth + // section offset locating the offset table. Block 64 now starts before + // block 63. + const uint64_t offset_table = DecodeFixed64(&data[56 + 8 * 3]); + EncodeFixed64(&data[offset_table + 8 + 16], 0); + Trie index; + ASSERT_OK(index.Init(data)); + TrieIterator it(&index); + ASSERT_OK(it.SeekToFirst()); + for (int i = 0; i < 62; i++) { + ASSERT_OK(it.Next()); + } + ASSERT_EQ(it.key(), keys[62]); + ASSERT_EQ(it.Next().code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); + ASSERT_EQ(it.Seek(keys[63], kSeekTag).code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); + ASSERT_OK(it.Seek(keys[64], kSeekTag)); + ASSERT_EQ(it.key(), keys[64]); + ASSERT_EQ(it.Prev().code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); +} + +// The serialized bytes of a small index, written out by hand from the format +// description in trie.h. Changing them breaks existing SST files. +TEST_F(TrieTest, GoldenBytes) { + std::string data = Build({{"ab", {0, 10}, 0}, {"ac", {15, 20}, 257}}, 8); + std::string expected; + ASSERT_TRUE(Slice( + // magic, version, flags, bucket_size + "45495254" + "02000000" + "00000000" + "08000000" + // num_blocks, num_keys + "0200000000000000" + "0200000000000000" + // prefix_len, max_key_len, id_width, trie_root + "01000000" + "02000000" + "01000000" + "00000000" + // size_gap + "0500000000000000" + // trie, bucket table, bucket data, offset table, size table + // (none), run table, tag table, end + "7900000000000000" + "7D00000000000000" + "9D00000000000000" + "9F00000000000000" + "0000000000000000" + "C700000000000000" + "CF00000000000000" + "E700000000000000" + // shared prefix "a" + "61" + // trie: terminal leaf, 1-byte pointers, path "b", bucket 0 + "04" + "01" + "62" + "00" + // bucket table: 2 values as 0 + 2 * i, residuals of width 0 + // from word 0, then padding + "0200000000000000" + "0000000000000000" + "02000000" + "00000000" + "0000000000000000" + // bucket data: "ac" after "ab" drops 1 byte and adds "c" + "11" + "63" + // offset table: offsets 0, 15 and 40 (35 plus the 5-byte gap) + // as -5 + 20 * i + residuals 5, 0, 5 of width 3 from word 0, + // then padding + "0300000000000000" + "FBFFFFFFFFFFFFFF" + "14000000" + "03000000" + "4501000000000000" + "0000000000000000" + // run table: empty + "0000000000000000" + // tag table: block 1 has tag 257 + "0100000000000000" + "0100000000000000" + "0101000000000000") + .DecodeHex(&expected)); + ASSERT_EQ(Slice(data).ToString(true), Slice(expected).ToString(true)); +} + +// Init() does not check the separators inside buckets. The iterator +// bounds-checks each one as it decodes it and reports one that is out of +// bounds as Corruption, leaving itself invalid. +TEST_F(TrieTest, MalformedBucketEntryFailsOnRead) { + std::vector blocks; + blocks.reserve(4); + for (uint64_t i = 0; i < 4; i++) { + blocks.push_back({"k" + std::to_string(i), {i * 105, 100}, 0}); + } + std::string data = Build(blocks, 8); + // The first entry after bucket 0's head, located by the header's third + // section offset. It now drops 15 bytes from a 1-byte suffix. + const uint64_t bucket_data = DecodeFixed64(&data[56 + 8 * 2]); + data[bucket_data] = '\xF0'; + Trie index; + ASSERT_OK(index.Init(data)); + TrieIterator it(&index); + ASSERT_OK(it.SeekToFirst()); + ASSERT_EQ(it.key(), "k0"); + ASSERT_EQ(it.Next().code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); + ASSERT_EQ(it.Seek("k2", kSeekTag).code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); + ASSERT_EQ(it.SeekToLast().code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); +} + +// Init() does not check bucket positions either. A position past the bucket +// data is clamped to its end, so the bucket reads as truncated. +TEST_F(TrieTest, BadBucketPositionFailsOnRead) { + std::vector blocks; + blocks.reserve(4); + for (uint64_t i = 0; i < 4; i++) { + blocks.push_back({"k" + std::to_string(i), {i * 105, 100}, 0}); + } + std::string data = Build(blocks, 8); + // The base of the only group of the bucket table, located by the header's + // second section offset. + const uint64_t bucket_table = DecodeFixed64(&data[56 + 8]); + EncodeFixed64(&data[bucket_table + 8], uint64_t{1} << 40); + Trie index; + ASSERT_OK(index.Init(data)); + TrieIterator it(&index); + // The bucket head comes from the trie, so only the entries after it fail. + ASSERT_OK(it.SeekToFirst()); + ASSERT_EQ(it.key(), "k0"); + ASSERT_EQ(it.Next().code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); + ASSERT_EQ(it.Seek("k2", kSeekTag).code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); +} + +// Prev re-decodes the previous key's bucket from its head, so it meets a +// malformed entry the iterator has not decoded yet. +TEST_F(TrieTest, MalformedBucketEntryFailsOnPrev) { + std::vector blocks; + blocks.reserve(4); + for (uint64_t i = 0; i < 4; i++) { + blocks.push_back({"k" + std::to_string(i), {i * 105, 100}, 0}); + } + // Buckets {k0, k1} and {k2, k3}. Break k1, the only entry of bucket 0. + std::string data = Build(blocks, 2); + const uint64_t bucket_data = DecodeFixed64(&data[56 + 8 * 2]); + data[bucket_data] = '\xF0'; + Trie index; + ASSERT_OK(index.Init(data)); + TrieIterator it(&index); + // k2 is a bucket head, so the Seek decodes no entry. + ASSERT_OK(it.Seek("k2", kSeekTag)); + ASSERT_EQ(it.key(), "k2"); + ASSERT_EQ(it.Prev().code(), Status::Code::kCorruption); + ASSERT_FALSE(it.Valid()); +} + +// Malformed indexes must be rejected by Init() or iterate without crashing. +// Run under ASAN to catch out-of-bounds reads. +TEST_F(TrieTest, MalformedData) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + + auto exercise = [&](const std::string& data) { + Trie index; + if (!OpenedOrRejected(index.Init(data))) { + return; + } + TrieIterator it(&index); + const uint64_t limit = 2 * NumBlocks(index) + 10; + bool ok = OkOrCorruption(it.SeekToFirst()); + for (uint64_t i = 0; ok && it.Valid() && i < limit; i++) { + // Sizes are derived here, and never wrap around. + const TrieBlockHandle handle = it.value(); + ASSERT_GE(handle.offset + handle.size, handle.offset); + ok = OkOrCorruption(it.Next()); + } + ok = OkOrCorruption(it.SeekToLast()); + for (uint64_t i = 0; ok && it.Valid() && i < limit; i++) { + ok = OkOrCorruption(it.Prev()); + } + for (int i = 0; i < 20; i++) { + ok = OkOrCorruption( + it.Seek(rnd.RandomBinaryString(rnd.Uniform(12)), kSeekTag)); + if (ok && it.Valid()) { + OkOrCorruption(it.Prev()); + } + } + }; + + for (int round = 0; round < 20; round++) { + std::vector separators = RandomKeys(&rnd, 1 + rnd.Uniform(60)); + if (rnd.OneIn(2)) { + separators.insert( + separators.begin() + rnd.Uniform(static_cast(separators.size())), + separators[0]); + std::sort(separators.begin(), separators.end()); + } + std::vector blocks = ContiguousBlocks(separators, &rnd); + blocks.back().tag = 1 << 8; + const std::string data = Build(blocks, 1 + rnd.Uniform(8)); + + for (size_t len = 0; len < data.size(); len++) { + Trie index; + ASSERT_EQ(index.Init(Slice(data.data(), len)).code(), + Status::Code::kCorruption); + } + for (size_t pos = 0; pos < kTrieHeaderSize; pos++) { + for (char v : {'\0', '\xff', static_cast(data[pos] ^ 1)}) { + std::string mutated = data; + mutated[pos] = v; + exercise(mutated); + } + } + for (int i = 0; i < 300; i++) { + exercise(FlipRandomBytes(data, &rnd)); + } + } +} + +TEST_F(TrieTest, RandomizedSeekAndScan) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + for (int round = 0; round < 40; round++) { + const size_t n = 1 + rnd.Uniform(round < 20 ? 40 : 1500); + std::vector keys = RandomKeys(&rnd, n); + std::vector blocks = ContiguousBlocks(keys, &rnd); + std::vector targets = SeekTargets(keys, &rnd); + for (uint32_t bucket_size : {1u, 2u, 3u, 8u, 64u}) { + ASSERT_NO_FATAL_FAILURE( + Verify(blocks, bucket_size, targets, BasicSeekTags())); + } + } +} + +// Like RandomizedSeekAndScan, with the shapes it leaves out: separators that +// span several blocks, tags, explicit sizes, offsets past 4 GiB, the empty +// key and a separator equal to the shared prefix. Also walks Next and Prev. +TEST_F(TrieTest, RandomizedRunsTagsAndSizes) { + const uint32_t seed = TimeSeed(); + SCOPED_TRACE("seed=" + std::to_string(seed)); + Random rnd(seed); + for (int round = 0; round < 30; round++) { + std::vector keys = + RandomKeys(&rnd, 1 + rnd.Uniform(round < 15 ? 30 : 200)); + const size_t shared = Slice(keys.front()).difference_offset(keys.back()); + if (rnd.OneIn(3) && shared < keys.front().size()) { + keys.insert(keys.begin(), keys.front().substr(0, shared)); + } else if (rnd.OneIn(2)) { + keys.insert(keys.begin(), ""); + } + std::vector blocks = RandomBlocks(keys, &rnd); + std::vector targets = SeekTargets(keys, &rnd); + std::vector tags = BasicSeekTags(); + for (int i = 0; i < 2; i++) { + const uint64_t tag = + blocks[rnd.Uniform(static_cast(blocks.size()))].tag; + tags.insert(tags.end(), {tag, tag + 1, tag > 0 ? tag - 1 : 0}); + } + for (uint32_t bucket_size : {1u, 2u, 3u, 8u, 64u}) { + ASSERT_NO_FATAL_FAILURE(Verify(blocks, bucket_size, targets, tags)); + ASSERT_NO_FATAL_FAILURE( + VerifyWalks(blocks, bucket_size, targets, tags, &rnd)); + } + } +} + +TEST_F(TrieTest, SingleBlock) { + std::string data = Build({{"apple", {0, 4000}, 0}}, 8); + Trie index; + ASSERT_OK(index.Init(data)); + ASSERT_EQ(NumBlocks(index), 1u); + TrieIterator it(&index); + + for (const char* target : {"", "aardvark", "apple"}) { + SCOPED_TRACE(target); + ASSERT_OK(it.Seek(target, kSeekTag)); + ASSERT_TRUE(it.Valid()); + ASSERT_EQ(it.key(), "apple"); + ASSERT_EQ(it.value().offset, 0u); + ASSERT_EQ(it.value().size, 4000u); + } + for (const char* target : {"apples", "b"}) { + SCOPED_TRACE(target); + ASSERT_OK(it.Seek(target, kSeekTag)); + ASSERT_FALSE(it.Valid()); + } +} + } // namespace trie_index } // namespace ROCKSDB_NAMESPACE