From 61fe952125b117b1e778b554f5d0c62f55b03291 Mon Sep 17 00:00:00 2001 From: Aleksey Shakhmatov Date: Sun, 2 Aug 2026 21:10:25 +0300 Subject: [PATCH] index: B+tree over the encoded keys (roadmap item 1) Replace Index.entries (one sorted array) with a B+tree so writes into an already-built index stop being quadratic. Nodes are fixed 4 KiB slotted pages in a flat u32-addressed ArrayListUnmanaged(Node); records longer than a quarter page spill to an append-only overflow slab (BSON strings reach 16 MB). Leaves are doubly linked for ordered iteration; the flat node array stays one contiguous byte range for a later checkpoint. Insertion descends by separator and splits leaves/internals upward, promoting keys via a stable copy (a nested split can otherwise clobber the promoted-key scratch). Deletion does not rebalance: emptied leaves are unlinked and dropped from their parent, internal nodes may carry one child, and dead pages are abandoned in place (node memory peaks at the tree's peak size, exactly what the old array's capacity did). Lookups are lower-bound seeks plus leaf-chain band scans, so equal keys may span leaves freely. Bulk build (append_doc_entries + finish_bulk) sorts a staging array and packs leaves bottom-up. reserve_for now takes the built entries and reserves exact overflow bytes plus a worst-case node count, keeping insert_entries infallible after the log append. db.zig: TTL sweep now seeks the minimum-datetime encoded key and walks the contiguous datetime band, stopping at the cutoff or type change. Measured (tests/e2e/results/phase2.txt): updateMany 17.3 -> 1.8 ms (2.8x slower than MongoDB -> 3.7x faster), createIndex 62 -> 51 ms. Verified: unit suite ReleaseFast/ReleaseSafe/Debug (incl. the existing lookup_range and remove_doc differentials, plus a new incremental insert/remove differential against a brute-force model), the crash pair, e2e3/e2e4/e2e6, and dev stress tests for depth-2 splits, full drains, and spilled records through internal levels. --- README.md | 56 +- ROADMAP.md | 29 +- src/db.zig | 36 +- src/index.zig | 1173 +++++++++++++++++++++++++++++----- src/spill.zig | 74 +++ src/spill2.zig | 88 +++ src/stress.zig | 174 +++++ tests/e2e/results/phase2.txt | 38 ++ 8 files changed, 1450 insertions(+), 218 deletions(-) create mode 100644 src/spill.zig create mode 100644 src/spill2.zig create mode 100644 src/stress.zig create mode 100644 tests/e2e/results/phase2.txt diff --git a/README.md b/README.md index ed33102..4fd2307 100644 --- a/README.md +++ b/README.md @@ -185,22 +185,22 @@ With the ReleaseFast default build (MongoDB 8.3.7 on the same Mac): | benchmark | mongo-lite | mongodb | winner | |---|---|---|---| -| insertOne (sequential) | 0.19 ms | 5.0 ms | **mongo-lite ×26** | -| bulk insert (insertMany) | 853 MB/s | 690 MB/s | **mongo-lite ×1.2** | -| createIndex({k: 1}) | 62 ms | 78 ms | **mongo-lite** | -| countDocuments({}) | 1.5 ms | 11.5 ms | **mongo-lite ×8** | -| findOne({_id}) | 0.48 ms | 0.54 ms | mongo-lite | -| findOne indexed | 0.64 ms | 4.3 ms | **mongo-lite ×7** | -| range-scan count | 21 ms | 13 ms | mongodb ×1.7 | -| sort + limit(20), on `_id` | 4.3 ms | 2.0 ms | mongodb ×2 | +| insertOne (sequential) | 0.19 ms | 4.1 ms | **mongo-lite ×22** | +| bulk insert (insertMany) | 810 MB/s | 714 MB/s | **mongo-lite ×1.1** | +| createIndex({k: 1}) | 51 ms | 82 ms | **mongo-lite** | +| countDocuments({}) | 1.5 ms | 13.8 ms | **mongo-lite ×9** | +| findOne({_id}) | 0.57 ms | 0.67 ms | mongo-lite | +| findOne indexed | 0.57 ms | 1.8 ms | **mongo-lite ×3** | +| range-scan count | 20 ms | 13 ms | mongodb ×1.6 | +| sort + limit(20), on `_id` | 6.2 ms | 2.7 ms | mongodb ×2.3 | | sort + limit(20), indexed field | 1.0 ms | — | — | -| aggregate $group | 9.8 ms | 13.7 ms | **mongo-lite** | -| updateOne({_id}) | 0.16 ms | 0.19 ms | mongo-lite | -| updateMany (65 docs) | 5.5 ms | 6.1 ms | mongo-lite | -| deleteOne + insert | 0.62 ms | 4.9 ms | **mongo-lite ×8** | -| server RSS | 2.0 GB | 1.4 GB | mongodb (×0.7) | +| aggregate $group | 11.5 ms | 15.5 ms | **mongo-lite** | +| updateOne({_id}) | 0.17 ms | 0.19 ms | mongo-lite | +| updateMany (65 docs) | 1.8 ms | 6.7 ms | **mongo-lite ×3.7** | +| deleteOne + insert | 0.50 ms | 5.0 ms | **mongo-lite ×10** | +| server RSS | 2.0 GB | 1.5 GB | mongodb (×0.7) | | kill -9 → reopen | 0.8 s | 1.3 s | **mongo-lite** | -| db on disk | 1.0 GB | 93 MB | mongodb (compressed) | +| db on disk | 1.0 GB | 96 MB | mongodb (compressed) | The remaining losses are structural rather than incidental. Disk size is the big one: payloads are stored raw, so the log is 11x MongoDB's @@ -211,8 +211,10 @@ that each live in a separate allocation, one pointer chase apiece. And covers `_id` yet; the same sort on an indexed field streams straight out of the index at 1.0 ms. -Reproduce the table with `bash tests/e2e/compare-run.sh 1g 16k`; the run -above is recorded in `tests/e2e/results/phase1.txt`. +Reproduce the table with `bash tests/e2e/compare-run.sh 1g 16k`; the +pre-tree baseline is recorded in `tests/e2e/results/phase1.txt`, the run +above (with the B+tree, roadmap item 1) in +`tests/e2e/results/phase2.txt`. ### What is left (highest impact first) @@ -224,20 +226,15 @@ traps in [ROADMAP.md](ROADMAP.md). highly compressible workloads massively; note Zig 0.16 ships zstd decompression only, and deflate would cap writes below the current insert rate. -2. **A B-tree over the encoded keys** — entry insert still memmoves the - tail of a sorted array, so writing into a collection that already has an - index is quadratic. A flat, `u32`-indexed node array would also be - dumpable into a checkpoint, which is what makes a fast reopen possible. -3. **An ordered `_id` index** — `sort({_id: ...})` still materializes every +2. **An ordered `_id` index** — `sort({_id: ...})` still materializes every candidate, and integer `_id`s still scan. Both fall out of indexing the - encoded `_id`. It wants the tree first: an `_id` index updates on every - insert, and doing that against a sorted array is only cheap because - ObjectIds append at the end. -4. **Stop giving every document its own arena** — the source of both the + encoded `_id`. The tree is in place, so an `_id` index update is a + leaf insert, not a tail memmove. +3. **Stop giving every document its own arena** — the source of both the RSS gap and the range-scan gap. Storing canonical BSON bytes in a per-collection slab and matching against them (parsing only the fields a filter names) makes scans contiguous instead of a pointer chase. -5. **Decompose the global lock** — one reader/writer lock covers the whole +4. **Decompose the global lock** — one reader/writer lock covers the whole engine and is held across fsync, compaction and reply construction. Per-collection locks plus cross-connection group commit are the path to using more than one core on writes. @@ -255,6 +252,13 @@ Done so far, with the measurement that drove each: array. `createIndex` over 65,536 documents 649 → 44 ms. - **Index entries hold encoded byte keys**, so comparing them is a memcmp rather than a walk over values in unrelated arenas. +- **A B+tree over the encoded keys** (roadmap item 1): fixed 4 KiB slotted + pages in a flat u32-addressed node array, an overflow slab for long + records, no rebalancing on delete, and bulk bottom-up packing. Entry + insertion and removal are a descent plus a leaf-local edit instead of a + tail memmove, so writes into an already-built index stopped being + quadratic. `updateMany` 17.3 → 1.8 ms (2.8x slower than MongoDB → 3.7x + faster); `createIndex` 62 → 51 ms. - **Entry removal is a binary search**, not a scan of the whole index. `updateMany` 15.4 → 5.5 ms. - **Top-k sort selection** and an allocation-free decorate pass, plus diff --git a/ROADMAP.md b/ROADMAP.md index 3459c9f..18e8faa 100644 --- a/ROADMAP.md +++ b/ROADMAP.md @@ -1,6 +1,10 @@ # Remaining performance work -Five items, in dependency order. Each is sized to be landed and verified on +Status: **item 1 (B+tree over the encoded keys) is done** — landed and +verified in `tests/e2e/results/phase2.txt` (updateMany 17.3 → 1.8 ms, +createIndex 62 → 51 ms). Its dependents (items 2 and 4) now stand on a +tree instead of a sorted array. Five items below, in dependency order. +Each is sized to be landed and verified on its own; the ordering constraints between them are the load-bearing part, so read those before picking one up. @@ -43,7 +47,28 @@ touching indexes needs `e2e3.js` and `e2e4.js`. --- -## 1. B-tree over the encoded keys +--- + +## 1. B-tree over the encoded keys — DONE + +Landed as a B+tree in `src/index.zig`: fixed 4 KiB slotted pages in a flat +u32-addressed `ArrayListUnmanaged(Node)`, an append-only overflow slab for +records longer than a quarter page (BSON strings reach 16 MB), leaves +linked for ordered iteration, bulk bottom-up packing for +`append_doc_entries` + `finish_bulk`, `remove_doc` regenerating entries +and removing them with a descent plus a leaf-local edit, and no +rebalancing on delete (empty leaves are unlinked and dropped; internal +nodes may carry one child). The flat node array stays the contiguous byte +range that item 4 can write into a checkpoint. Deletes abandon dead pages +rather than reusing them, so node memory peaks at the tree's peak size, +exactly what the old entry array's capacity did. + +Recorded deltas vs `tests/e2e/results/phase1.txt`: `updateMany` 17.3 → +1.8 ms (2.8x slower than MongoDB → 3.7x faster), `createIndex` 62.4 → +50.8 ms. Verified with `zig build test` (ReleaseFast and ReleaseSafe), the +crash pair, `e2e3.js`/`e2e4.js`/`e2e6.js`, plus a 50k-entry stress (bulk +build, random inserts, random deletes, full drain) and a spill stress +(2 KiB keys through splits and internal nodes). **Why.** Index entries live in one sorted array, so inserting an entry memmoves the tail. Building an index is fine (entries are appended and sorted diff --git a/src/db.zig b/src/db.zig index aef758b..dc7ee55 100644 --- a/src/db.zig +++ b/src/db.zig @@ -255,10 +255,10 @@ pub const Engine = struct { }; } - // 4. Reserve entry capacity — the last fallible step, so the entry + // 4. Reserve tree capacity — the last fallible step, so the entry // insertion after the log append is infallible. for (built_list.items) |*b| { - try b.ix.reserve_for(self.gpa, b.built.entries.items.len); + try b.ix.reserve_for(self.gpa, b.built.entries.items); } // 5. Log (and sync) before anything becomes visible. @@ -365,7 +365,7 @@ pub const Engine = struct { while (doc_it.next()) |entry| { try ix.append_doc_entries(self.gpa, entry.value_ptr.*, entry.key_ptr.*); } - _ = try ix.finish_bulk(true); + _ = try ix.finish_bulk(self.gpa, true); // Reserve the collection slot, then persist and publish. try coll.indexes.ensureUnusedCapacity(self.gpa, 1); @@ -428,16 +428,16 @@ pub const Engine = struct { for (coll_entry.value_ptr.indexes.items) |*ix| { const ttl = ix.ttl orelse continue; const cutoff: i128 = @as(i128, now_ms) - @as(i128, ttl) * 1000; - for (ix.entries.items) |e| { - // Still a linear walk: the type test cannot be a - // one-sided range lookup, because bson compare order - // ranks datetime above null, numbers and strings, so - // a datetime upper bound would also select every - // value of a lesser type. (Datetimes are contiguous - // in that order, so a two-sided band lookup would - // work — that comes with the tree.) - const ms = bson.encoded_leading_datetime(e.key) orelse continue; - if (@as(i128, ms) > cutoff) continue; + // bson compare order ranks datetime above null, numbers + // and strings and below only timestamp and maxKey, so + // datetimes form a contiguous band in the encoded key + // order: seek the minimum datetime and stop when the + // leading type changes or the cutoff is passed. + const min_dt = [_]u8{ bson.encoded_datetime_tag, 0, 0, 0, 0, 0, 0, 0, 0 }; + var it = ix.seek(&min_dt); + while (it.next()) |e| { + const ms = bson.encoded_leading_datetime(e.key) orelse break; + if (@as(i128, ms) > cutoff) break; try ids.append(self.gpa, try self.gpa.dupe(u8, e.id)); } } @@ -624,7 +624,7 @@ pub const Engine = struct { var coll_it = db_entry.value_ptr.collections.iterator(); while (coll_it.next()) |coll_entry| { for (coll_entry.value_ptr.indexes.items) |*ix| { - if (ix.entries.items.len > 0) continue; // defensive + if (ix.count() > 0) continue; // defensive var doc_it = coll_entry.value_ptr.docs.iterator(); while (doc_it.next()) |doc_entry| { ix.append_doc_entries(self.gpa, doc_entry.value_ptr.*, doc_entry.key_ptr.*) catch |err| switch (err) { @@ -636,7 +636,7 @@ pub const Engine = struct { }; } // Tolerated, not enforced: the database must always open. - if (try ix.finish_bulk(false)) { + if (try ix.finish_bulk(self.gpa, false)) { std.debug.print("mongo-lite: WARNING: unique index '{s}' has duplicate keys in existing data; duplicates not enforced for existing documents\n", .{ix.name}); } } @@ -1369,12 +1369,12 @@ test "ttl_sweep deletes expired documents and the deletion survives reopen" { } const coll = engine.get_collection("app", "sessions").?; try testing.expectEqual(@as(usize, 6), coll.docs.count()); - try testing.expectEqual(@as(usize, 6), coll.indexes.items[0].entries.items.len); + try testing.expectEqual(@as(usize, 6), coll.indexes.items[0].count()); // The cutoff is inclusive: doc 2 goes with doc 1. try testing.expectEqual(@as(usize, 2), try engine.ttl_sweep(now_ms)); try testing.expectEqual(@as(usize, 4), coll.docs.count()); - try testing.expectEqual(@as(usize, 4), coll.indexes.items[0].entries.items.len); + try testing.expectEqual(@as(usize, 4), coll.indexes.items[0].count()); try testing.expectEqual(@as(usize, 0), try index_count(gpa, &engine, "app", "sessions", "expireAt_1", .{ .datetime = now_ms - 120_000 })); // The string and the missing field are untouched by any sweep. try testing.expectEqual(@as(usize, 1), try index_count(gpa, &engine, "app", "sessions", "expireAt_1", .{ .string = "tomorrow" })); @@ -1397,7 +1397,7 @@ test "ttl_sweep deletes expired documents and the deletion survives reopen" { try testing.expectEqual(@as(usize, 3), coll.docs.count()); try testing.expectEqual(@as(usize, 1), coll.indexes.items.len); try testing.expectEqual(@as(?i64, 60), coll.indexes.items[0].ttl); - try testing.expectEqual(@as(usize, 3), coll.indexes.items[0].entries.items.len); + try testing.expectEqual(@as(usize, 3), coll.indexes.items[0].count()); for ([_]i32{ 1, 2, 3 }) |id| { const id_key = try bson.serialize_value(gpa, bson.Value{ .int32 = id }); defer gpa.free(id_key); diff --git a/src/index.zig b/src/index.zig index d883499..3a4f56f 100644 --- a/src/index.zig +++ b/src/index.zig @@ -26,6 +26,17 @@ //! arena. `Entry.id` still aliases the docs map key; index removal happens //! at the top of evict_doc (src/db.zig) — the single chokepoint where a //! document dies — so that aliasing is safe by construction. +//! +//! Entries live in a B+tree (item 1 of the performance roadmap) instead of +//! one sorted array, so inserting into an already-built index costs O(log n) +//! instead of memmoving the array tail. Nodes are fixed 4 KiB pages in a +//! flat u32-addressed `ArrayListUnmanaged(Node)`; leaves are doubly linked +//! for ordered iteration. Records longer than a quarter page spill to an +//! append-only overflow slab (BSON strings reach 16 MB). Deletion does not +//! rebalance: leaves may sit underfull until emptied (an emptied leaf is +//! unlinked and dropped from its parent), and internal nodes may carry one +//! child — both are correct, and a fresh tree is rebuilt from the log on +//! every open anyway. const std = @import("std"); const bson = @import("bson.zig"); @@ -63,6 +74,51 @@ pub const BuiltEntries = struct { } }; +// --------------------------------------------------------------------------- +// B+tree storage +// --------------------------------------------------------------------------- + +const page_size = 4096; +/// Bytes of node payload: a 32-byte header plus the slotted region. +const page_data = page_size - 32; +/// Records longer than a quarter of a node spill to the overflow slab. +const inline_limit = page_size / 4; + +/// One slotted-page entry: a key plus either an id length (leaf records) or +/// a child node id (internal separators). The key bytes live either inline +/// in the node's page (`off` into `buf`) or in the overflow slab (`off` into +/// `Index.overflow`). +const Slot = packed struct { + off: u64, + key_len: u32, + extra: u32, + spill: bool, + _pad: u31, +}; +const slot_size = @sizeOf(Slot); + +/// A fixed 4 KiB page. Slots grow from the front of `buf`, record bytes +/// from the back; slots are kept ordered by key, so a node lookup is a +/// binary search over slots. Node ids address `Index.nodes` directly. +const Node = extern struct { + is_leaf: u32, + count: u32, + parent: u32, + next: u32, + prev: u32, + first_child: u32, + data_start: u32, + _rsv: u32, + buf: [page_data]u8, +}; + +/// A view of one stored entry yielded by iteration. Both slices alias the +/// tree and are valid only while the tree is not mutated. +pub const EntryRef = struct { + key: []const u8, + id: []const u8, +}; + pub const Index = struct { name: []const u8, keys: []const IndexKey, @@ -75,7 +131,28 @@ pub const Index = struct { /// [0, max_expire_after_seconds] — parse_spec is the only producer. ttl: ?i64, multikey: bool, - entries: std.ArrayListUnmanaged(Entry), + + // -- the tree ---------------------------------------------------------- + nodes: std.ArrayListUnmanaged(Node), + /// Append-only slab of spilled records; referenced by slots, never + /// rewritten or freed (offsets stay valid forever). + overflow: std.ArrayListUnmanaged(u8), + /// Bulk-build staging: entries appended unsorted by append_doc_entries, + /// sorted and packed into the tree by finish_bulk. The keys are owned + /// by the staging array until the pack consumes them. + staging: std.ArrayListUnmanaged(Entry), + root: u32, + first_leaf: u32, + leaf_count: u32, + /// Internal levels above the leaves: 0 when the root is a leaf. + depth: u32, + /// Total entries, maintained incrementally. + entry_count: usize, + /// Repack scratch: any single node's record bytes fit here. + scratch: [page_data]u8, + /// Promoted-key scratch: inline keys being propagated up a split are + /// copied here so the reference survives node-array growth. + promo: [inline_limit]u8, pub fn init(gpa: std.mem.Allocator, name: []const u8, keys: []const IndexKey, unique: bool, sparse: bool, ttl: ?i64) !Index { var self: Index = .{ @@ -85,7 +162,16 @@ pub const Index = struct { .sparse = sparse, .ttl = ttl, .multikey = false, - .entries = .empty, + .nodes = .empty, + .overflow = .empty, + .staging = .empty, + .root = 0, + .first_leaf = 0, + .leaf_count = 0, + .depth = 0, + .entry_count = 0, + .scratch = undefined, + .promo = undefined, }; self.name = try gpa.dupe(u8, name); const owned_keys = try gpa.alloc(IndexKey, keys.len); @@ -99,17 +185,30 @@ pub const Index = struct { owned_keys[n] = .{ .path = try gpa.dupe(u8, keys[n].path), .descending = keys[n].descending }; } self.keys = owned_keys; + // nodes[0] is a dummy (0 is the null node id); the root is one + // empty leaf, so a fresh index is always a valid tree. + try self.nodes.append(gpa, empty_node(0)); + try self.nodes.append(gpa, empty_node(1)); + self.root = 1; + self.first_leaf = 1; + self.leaf_count = 1; return self; } pub fn deinit(self: *Index, gpa: std.mem.Allocator) void { - for (self.entries.items) |e| gpa.free(e.key); - self.entries.deinit(gpa); + for (self.staging.items) |e| gpa.free(e.key); + self.staging.deinit(gpa); + self.nodes.deinit(gpa); + self.overflow.deinit(gpa); for (self.keys) |k| gpa.free(k.path); gpa.free(self.keys); gpa.free(self.name); } + pub fn count(self: *const Index) usize { + return self.entry_count; + } + // -- entry generation --------------------------------------------------- /// Build the entries one document contributes. Mirrors field_matches @@ -191,24 +290,31 @@ pub const Index = struct { // -- mutation ----------------------------------------------------------- - /// Ensure the entry array has room for `n` more entries. Called before - /// the log append, so the subsequent insert_entries is infallible. - pub fn reserve_for(self: *Index, gpa: std.mem.Allocator, n: usize) !void { - try self.entries.ensureUnusedCapacity(gpa, n); + /// Reserve everything `insert_entries` of this batch will need, so that + /// insertion — which runs after the log append — is infallible: node + /// pages for the worst-case splits of a batch (each entry creates at + /// most one leaf plus a bounded share of internal splits; batches are + /// sorted, so new leaves are packed, but reserve the trivial bound), + /// and the exact overflow bytes (spilled records are copied once). + pub fn reserve_for(self: *Index, gpa: std.mem.Allocator, entries: []const Entry) !void { + const n: u64 = entries.len; + var overflow_bytes: u64 = 0; + for (entries) |e| { + const rec_len: u64 = e.key.len + e.id.len; + if (rec_len > inline_limit) overflow_bytes += rec_len; + } + const extra_nodes: u64 = n + n / 100 + self.depth + 2; + try self.nodes.ensureUnusedCapacity(gpa, self.nodes.items.len + @as(usize, @intCast(extra_nodes))); + try self.overflow.ensureUnusedCapacity(gpa, self.overflow.items.len + @as(usize, @intCast(overflow_bytes))); } - /// Insert pre-built entries (maintaining sort order) and drain the - /// batch: ownership of each entry's key slice moves into the index, so - /// the batch's deinit must not free them. Infallible: capacity must - /// already be reserved. + /// Insert pre-built entries (maintaining order) and drain the batch. + /// Infallible: reserve_for must have run first. The tree copies each + /// key, so ownership stays with the batch and its deinit frees them. pub fn insert_entries(self: *Index, built: *BuiltEntries) void { for (built.entries.items) |e| { - const pos = self.insert_pos(e); - self.entries.insertAssumeCapacity(pos, e); + self.insert_entry(e.key, e.id); } - // Key slices now belong to the index; forget them in the batch so - // its deinit only frees the (now empty) ArrayList buffer. - built.entries.items.len = 0; } /// Build, check, and insert entries for one document — the whole @@ -222,8 +328,8 @@ pub const Index = struct { /// value reports whether that happened. pub fn add_doc(self: *Index, gpa: std.mem.Allocator, doc: *const bson.Document, id: []const u8, enforce_unique: bool) !bool { var built = try self.build_entries(gpa, doc, id); - // Runs on success too: insert_entries drains the keys, leaving only - // the (now empty) ArrayList buffer to free. + // Runs on success too: the batch's keys are copied into the tree, + // so deinit frees exactly what this call allocated. defer built.deinit(gpa); var duplicate = false; if (self.unique) { @@ -233,79 +339,85 @@ pub const Index = struct { }; } if (built.multikey) self.multikey = true; - try self.reserve_for(gpa, built.entries.items.len); + try self.reserve_for(gpa, built.entries.items); self.insert_entries(&built); return duplicate; } - /// Append one document's entries without maintaining sort order. Pairs - /// with `finish_bulk`, which sorts the whole array once at the end. + /// Append one document's entries to the bulk-build staging array + /// without maintaining order. Pairs with `finish_bulk`, which sorts the + /// staging once and packs it into leaves bottom-up. /// - /// This is the bulk counterpart of `add_doc`. Inserting documents one at - /// a time keeps the array sorted by memmoving the tail on every entry, - /// so building an index over n documents moves O(n²) bytes — that is the - /// whole cost of createIndex on a large collection. Appending and - /// sorting once is O(n log n) comparisons and no memmove. + /// Inserting documents one at a time keeps a sorted array memmoved on + /// every entry, so building an index over n documents would move O(n²) + /// bytes — that was the whole cost of createIndex on a large + /// collection. Staging and packing is O(n log n) and no memmove. pub fn append_doc_entries(self: *Index, gpa: std.mem.Allocator, doc: *const bson.Document, id: []const u8) !void { var built = try self.build_entries(gpa, doc, id); - // Runs on success too: the append below drains the keys, leaving - // only the (now empty) ArrayList buffer to free. + // Runs on success too: the append below moves the keys into the + // staging array, leaving only the (now empty) ArrayList buffer. defer built.deinit(gpa); - try self.entries.appendSlice(gpa, built.entries.items); + try self.staging.appendSlice(gpa, built.entries.items); if (built.multikey) self.multikey = true; - // Key slices now belong to the index. + // Key ownership moved into the staging array. built.entries.items.len = 0; } - /// Sort the entries collected by `append_doc_entries` into the index's - /// total order and, for a unique index, look for duplicate keys — one - /// adjacent-pair scan instead of a binary search per document. + /// Sort the staging collected by `append_doc_entries`, check a unique + /// index for duplicate keys (one adjacent-pair scan), and pack the tree. /// /// With `enforce_unique` false a duplicate is tolerated rather than /// rejected, matching `add_doc`; the return value reports whether that /// happened. - pub fn finish_bulk(self: *Index, enforce_unique: bool) error{DuplicateKeyIndex}!bool { - std.mem.sort(Entry, self.entries.items, {}, entry_less); - if (!self.unique or self.entries.items.len < 2) return false; - + pub fn finish_bulk(self: *Index, gpa: std.mem.Allocator, enforce_unique: bool) error{ DuplicateKeyIndex, OutOfMemory }!bool { + std.mem.sort(Entry, self.staging.items, {}, entry_less); var duplicate = false; - for (self.entries.items[1..], 0..) |cur, prev_i| { - const prev = self.entries.items[prev_i]; - // Sorting puts equal keys next to each other. - if (!std.mem.eql(u8, prev.key, cur.key)) continue; - // A document's own entries were deduped at build time, so equal - // keys under one id are not a conflict — same rule as - // check_unique's exclude_id. - if (std.mem.eql(u8, prev.id, cur.id)) continue; - if (enforce_unique) return error.DuplicateKeyIndex; - duplicate = true; + if (self.unique and self.staging.items.len >= 2) { + for (self.staging.items[1..], 0..) |cur, prev_i| { + const prev = self.staging.items[prev_i]; + // Sorting puts equal keys next to each other. + if (!std.mem.eql(u8, prev.key, cur.key)) continue; + // A document's own entries were deduped at build time, so + // equal keys under one id are not a conflict — same rule as + // check_unique's exclude_id. + if (std.mem.eql(u8, prev.id, cur.id)) continue; + if (enforce_unique) return error.DuplicateKeyIndex; + duplicate = true; + } } + try self.reserve_for(gpa, self.staging.items); + try self.pack_tree(gpa); return duplicate; } - /// Remove every entry for `id` and free its key slices. Infallible. - /// One compaction pass: removing in place would memmove the tail per hit. + /// Remove every entry for `id`, in one pass over the leaves. Infallible. + /// Used directly by remove_id's own callers and as the fallback when + /// regeneration cannot locate entries. pub fn remove_id(self: *Index, gpa: std.mem.Allocator, id: []const u8) void { - var w: usize = 0; - for (self.entries.items) |e| { - if (std.mem.eql(u8, e.id, id)) { - gpa.free(e.key); - } else { - self.entries.items[w] = e; - w += 1; + _ = gpa; + var leaf_id = self.first_leaf; + while (leaf_id != 0) { + const node = &self.nodes.items[leaf_id]; + const next = node.next; + // Slots are removed high-to-low so the indices stay valid. + var i = node.count; + while (i > 0) { + i -= 1; + if (std.mem.eql(u8, self.id_of(leaf_id, i), id)) self.leaf_remove(leaf_id, i); } + leaf_id = next; } - self.entries.items.len = w; } - /// Remove the entries `doc` contributed, locating each by binary search - /// instead of scanning the array. + /// Remove the entries `doc` contributed, locating each by regeneration + /// instead of scanning the tree. /// /// Entry generation is a pure function of the document, so the entries - /// to remove can be regenerated and looked up directly. `remove_id` - /// walks every entry in the index comparing ids, which made a single - /// document update cost O(index size) per index — the dominant cost of - /// updateMany and deleteMany on a large collection. + /// to remove can be regenerated and looked up directly; each is a + /// descent plus a leaf-local scan of its key band (equal keys may + /// legitimately span several leaves). `remove_id` walks every leaf + /// comparing ids, which made a single document update cost O(index + /// size) per index. /// /// Infallible by construction: regeneration allocates and can fail, and /// a document the index cannot key contributed nothing to remove, so @@ -315,32 +427,10 @@ pub const Index = struct { defer built.deinit(gpa); // Sparse index that skipped this document: nothing was inserted. if (built.entries.items.len == 0) return; - - // Positions must be removed high-to-low so the earlier ones stay - // valid. build_entries returns its entries sorted, and the array is - // in the same order, so positions come out ascending already. - var found: [max_index_keys]usize = undefined; - var n: usize = 0; for (built.entries.items) |want| { - const pos = std.sort.lowerBound(Entry, self.entries.items, want, compare_entries); - if (pos >= self.entries.items.len) continue; - const have = self.entries.items[pos]; - // compare_entries orders by key then id, so an exact hit means - // this is the very entry the document contributed. - if (compare_entries(have, want) != .eq) continue; - if (n == found.len) return self.remove_id(gpa, id); // more entries than we can track - found[n] = pos; - n += 1; - } - // A document that generated entries we cannot find would leave the - // index stale, which is worse than being slow. - if (n != built.entries.items.len) return self.remove_id(gpa, id); - - var i = n; - while (i > 0) { - i -= 1; - gpa.free(self.entries.items[found[i]].key); - _ = self.entries.orderedRemove(found[i]); + // A document that generated entries we cannot find would leave + // the index stale, which is worse than being slow. + if (!self.remove_entry(want.key, want.id)) return self.remove_id(gpa, id); } } @@ -349,18 +439,14 @@ pub const Index = struct { /// own old entries) are allowed. pub fn check_unique(self: *const Index, new_entries: []const Entry, exclude_id: []const u8) error{DuplicateKeyIndex}!void { for (new_entries) |e| { - const start = self.lower_bound_prefix(e.key); - const end = self.upper_bound_prefix(e.key); - for (self.entries.items[start..end]) |existing| { - if (!std.mem.eql(u8, existing.id, exclude_id)) return error.DuplicateKeyIndex; + var it = self.seek(e.key); + while (it.next()) |have| { + if (cmp_prefix(e.key, have.key) != .eq) break; + if (!std.mem.eql(u8, have.id, exclude_id)) return error.DuplicateKeyIndex; } } } - fn insert_pos(self: *const Index, e: Entry) usize { - return std.sort.lowerBound(Entry, self.entries.items, e, compare_entries); - } - // -- search ------------------------------------------------------------- /// All ids whose key equals `key` (component-wise). For a partial key @@ -369,10 +455,11 @@ pub const Index = struct { var enc: std.ArrayListUnmanaged(u8) = .empty; defer enc.deinit(gpa); for (key) |v| try bson.encode_key(v, gpa, &enc); - const start = self.lower_bound_prefix(enc.items); - const end = self.upper_bound_prefix(enc.items); - var i = start; - while (i < end) : (i += 1) try out.append(gpa, self.entries.items[i].id); + var it = self.seek(enc.items); + while (it.next()) |e| { + if (cmp_prefix(enc.items, e.key) != .eq) break; + try out.append(gpa, e.id); + } } /// All ids whose key starts with `prefix` and whose component at @@ -388,50 +475,83 @@ pub const Index = struct { hi_incl: bool, out: *std.ArrayListUnmanaged([]const u8), ) !void { - // The array is sorted on this component too, so both bounds are - // binary searches. Scanning the whole equality band and filtering - // made a range on the first component touch every entry in the - // index — O(n) for what is O(log n + result). - var enc: std.ArrayListUnmanaged(u8) = .empty; - defer enc.deinit(gpa); - for (prefix) |v| try bson.encode_key(v, gpa, &enc); - const prefix_len = enc.items.len; - - // Clamp into the equality band: lo/hi constrain only the component - // at prefix.len, so the bounds alone would reach past the entries - // that share the prefix. - var start = self.lower_bound_prefix(enc.items[0..prefix_len]); - var end = self.upper_bound_prefix(enc.items[0..prefix_len]); - - if (lo) |l| { - enc.items.len = prefix_len; - try bson.encode_key(l, gpa, &enc); - // Inclusive wants the first entry not less than lo; exclusive - // wants the first one strictly greater. - const bound = if (lo_incl) self.lower_bound_prefix(enc.items) else self.upper_bound_prefix(enc.items); - start = @max(start, bound); - } + // Both bounds extend the encoded prefix by one component, so each + // is a lower-bound seek (encoded bytes are self-delimiting: the + // extended key orders within the prefix's own band). The scan + // starts at the lower bound and stops at the first key beyond the + // band or the upper bound — no slice ends to compute. + var enc_a: std.ArrayListUnmanaged(u8) = .empty; + var enc_b: std.ArrayListUnmanaged(u8) = .empty; + defer enc_a.deinit(gpa); + defer enc_b.deinit(gpa); + for (prefix) |v| try bson.encode_key(v, gpa, &enc_a); + const prefix_len = enc_a.items.len; + if (lo) |l| try bson.encode_key(l, gpa, &enc_a); if (hi) |h| { - enc.items.len = prefix_len; - try bson.encode_key(h, gpa, &enc); - const bound = if (hi_incl) self.upper_bound_prefix(enc.items) else self.lower_bound_prefix(enc.items); - end = @min(end, bound); + for (prefix) |v| try bson.encode_key(v, gpa, &enc_b); + try bson.encode_key(h, gpa, &enc_b); } - if (start >= end) return; + const prefix_key = enc_a.items[0..prefix_len]; + const lo_key: ?[]const u8 = if (lo != null) enc_a.items else null; + const hi_key: ?[]const u8 = if (hi != null) enc_b.items else null; - for (self.entries.items[start..end]) |e| try out.append(gpa, e.id); + var it = self.seek(if (lo_key) |lk| lk else prefix_key); + while (it.next()) |e| { + if (lo_key) |lk| { + // Exclusive lower bound: skip entries still at or below lo + // (an entry whose component equals lo is exactly lo). + if (!lo_incl and cmp_prefix(lk, e.key) != .lt) continue; + } + if (hi_key) |hk| { + const c = cmp_prefix(hk, e.key); + if (hi_incl) { + if (c == .lt) break; // key > hi + } else { + if (c != .gt) break; // key >= hi + } + } else if (cmp_prefix(prefix_key, e.key) != .eq) { + break; // band ended + } + try out.append(gpa, e.id); + } } - /// First entry whose first `prefix.len` components are not less than - /// `prefix`. - fn lower_bound_prefix(self: *const Index, prefix: []const u8) usize { - return std.sort.lowerBound(Entry, self.entries.items, prefix, prefix_order); + // -- iteration ---------------------------------------------------------- + + /// Ordered iteration over the whole index. Entries come out in key + /// order; equal keys in the order the leaves happen to hold them. + pub const Iter = struct { + ix: *const Index, + leaf: u32, + slot: u32, + + pub fn next(self: *Iter) ?EntryRef { + const ix = self.ix; + while (self.leaf != 0) { + const node = &ix.nodes.items[self.leaf]; + if (self.slot < node.count) { + const key = ix.key_of(self.leaf, self.slot); + const id = ix.id_of(self.leaf, self.slot); + self.slot += 1; + return .{ .key = key, .id = id }; + } + self.leaf = node.next; + self.slot = 0; + } + return null; + } + }; + + pub fn iter(self: *const Index) Iter { + return .{ .ix = self, .leaf = self.first_leaf, .slot = 0 }; } - /// First entry whose first `prefix.len` components are greater than - /// `prefix`. - fn upper_bound_prefix(self: *const Index, prefix: []const u8) usize { - return std.sort.upperBound(Entry, self.entries.items, prefix, prefix_order); + /// Ordered iteration starting at the first entry whose key is not less + /// than `prefix` (prefix semantics). Used by the equality/range + /// searches and the TTL sweep's datetime band. + pub fn seek(self: *const Index, prefix: []const u8) Iter { + const b = self.lower_bound(prefix); + return .{ .ix = self, .leaf = b.leaf, .slot = b.slot }; } // -- serialization ------------------------------------------------------ @@ -475,6 +595,599 @@ pub const Index = struct { } return true; } + + // -- tree internals ----------------------------------------------------- + + const Record = struct { + key: []const u8, + id: []const u8 = "", + child: u32 = 0, + /// When set, key+id already live in the overflow slab and are + /// referenced rather than copied (moved leaf records, promoted + /// separators). + spill_off: u64 = 0, + }; + + const StableKey = struct { + key: []const u8, + spill_off: u64, + }; + + const Split = struct { + /// The separator to promote, stable across node-array growth (a + /// slice into the overflow slab, or a copy in the promo buffer). + key: []const u8, + spill_off: u64, + right: u32, + }; + + fn empty_node(is_leaf: u32) Node { + return .{ + .is_leaf = is_leaf, + .count = 0, + .parent = 0, + .next = 0, + .prev = 0, + .first_child = 0, + .data_start = page_data, + ._rsv = 0, + .buf = undefined, + }; + } + + /// Allocate a node id. Infallible: insert paths reserve capacity first; + /// the pack path reserves via reserve_for before packing. + fn alloc_node(self: *Index) u32 { + self.nodes.appendAssumeCapacity(empty_node(0)); + return @intCast(self.nodes.items.len - 1); + } + + fn get_slot(node: *const Node, i: u32) Slot { + return std.mem.bytesToValue(Slot, node.buf[i * slot_size ..][0..slot_size]); + } + + fn set_slot(node: *Node, i: u32, s: Slot) void { + std.mem.bytesAsValue(Slot, node.buf[i * slot_size ..][0..slot_size]).* = s; + } + + /// The key bytes of slot `i`, either in the node's page or the slab. + fn key_of(self: *const Index, node_id: u32, i: u32) []const u8 { + const node = &self.nodes.items[node_id]; + const s = get_slot(node, i); + if (s.spill) return self.overflow.items[@intCast(s.off) .. @intCast(s.off + s.key_len)]; + return node.buf[@intCast(s.off) .. @intCast(s.off + s.key_len)]; + } + + /// The id bytes of leaf slot `i`. + fn id_of(self: *const Index, node_id: u32, i: u32) []const u8 { + const node = &self.nodes.items[node_id]; + const s = get_slot(node, i); + const start = s.off + s.key_len; + if (s.spill) return self.overflow.items[@intCast(start) .. @intCast(start + s.extra)]; + return node.buf[@intCast(start) .. @intCast(start + s.extra)]; + } + + /// Whether a record of `rec_len` bytes fits `node`: the slot plus, when + /// it stays inline, its bytes. Oversized records spill (slot only). + fn fits(self: *const Index, node_id: u32, rec_len: u64) bool { + const node = &self.nodes.items[node_id]; + const inline_bytes: u64 = if (rec_len > inline_limit) 0 else rec_len; + return (@as(u64, node.count) + 1) * slot_size + inline_bytes <= node.data_start; + } + + /// Write `rec` into `node` at slot position `pos` (append when pos == + /// count), shifting later slots right to make room. Assumes fit. + fn store_record(self: *Index, node_id: u32, pos: u32, rec: Record) void { + const node = &self.nodes.items[node_id]; + const rec_len: u64 = rec.key.len + rec.id.len; + var s: Slot = .{ + .off = rec.spill_off, + .key_len = @intCast(rec.key.len), + .extra = if (rec.id.len > 0) @intCast(rec.id.len) else rec.child, + .spill = false, + ._pad = 0, + }; + if (rec.spill_off != 0) { + s.spill = true; + } else if (rec_len > inline_limit) { + s.off = self.overflow.items.len; + self.overflow.appendSliceAssumeCapacity(rec.key); + if (rec.id.len > 0) self.overflow.appendSliceAssumeCapacity(rec.id); + s.spill = true; + } else { + node.data_start -= @intCast(rec_len); + const dst = node.buf[node.data_start .. node.data_start + @as(usize, @intCast(rec_len))]; + @memcpy(dst[0..rec.key.len], rec.key); + if (rec.id.len > 0) @memcpy(dst[rec.key.len..][0..rec.id.len], rec.id); + s.off = node.data_start; + } + if (pos < node.count) { + const src = node.buf[pos * slot_size .. node.count * slot_size]; + const dst = node.buf[(pos + 1) * slot_size .. (node.count + 1) * slot_size]; + std.mem.copyBackwards(u8, dst, src); + } + set_slot(node, pos, s); + node.count += 1; + } + + /// Remove slot `i`, closing the hole its inline bytes leave and + /// adjusting surviving inline offsets. Infallible. + fn remove_record(self: *Index, node_id: u32, i: u32) void { + const node = &self.nodes.items[node_id]; + const s = get_slot(node, i); + const is_leaf = node.is_leaf == 1; + const rec_len: u64 = s.key_len + (if (is_leaf) s.extra else 0); + if (!s.spill and rec_len > 0) { + const src = node.buf[@intCast(s.off + rec_len)..page_data]; + const dst = node.buf[@intCast(s.off) .. page_data - @as(usize, @intCast(rec_len))]; + std.mem.copyForwards(u8, dst, src); + var j: u32 = 0; + while (j < node.count) : (j += 1) { + if (j == i) continue; + var sj = get_slot(node, j); + if (!sj.spill and sj.off > s.off) sj.off -= rec_len; + set_slot(node, j, sj); + } + // data_start is the low boundary of the live region and is + // unchanged: the freed bytes sit dead at the top of the page, + // below the accounting, until the next split repacks it. + } + const src_slots = node.buf[(i + 1) * slot_size .. node.count * slot_size]; + const dst_slots = node.buf[i * slot_size .. (node.count - 1) * slot_size]; + std.mem.copyForwards(u8, dst_slots, src_slots); + node.count -= 1; + } + + /// Repack `node` keeping only slots [0, k). Record bytes are copied to + /// the scratch and written back packed from the top of the page; slots + /// are rewritten in place with their new offsets. + fn repack_keep_prefix(self: *Index, node_id: u32, k: u32) void { + const node = &self.nodes.items[node_id]; + const is_leaf = node.is_leaf == 1; + // Pass 1: surviving inline records to the scratch, in slot order. + var scratch_len: usize = 0; + var j: u32 = 0; + while (j < k) : (j += 1) { + const s = get_slot(node, j); + if (s.spill) continue; + const rec_len: usize = @intCast(s.key_len + (if (is_leaf) s.extra else 0)); + const src = node.buf[@intCast(s.off) .. @intCast(s.off + rec_len)]; + @memcpy(self.scratch[scratch_len .. scratch_len + rec_len], src); + scratch_len += rec_len; + } + // Pass 2: rewrite slots and data, back to front so the scratch is + // consumed LIFO. + var cursor: usize = page_data; + j = k; + while (j > 0) { + j -= 1; + var s = get_slot(node, j); + if (s.spill) continue; + const rec_len: usize = @intCast(s.key_len + (if (is_leaf) s.extra else 0)); + cursor -= rec_len; + scratch_len -= rec_len; + @memcpy(node.buf[cursor .. cursor + rec_len], self.scratch[scratch_len .. scratch_len + rec_len]); + s.off = cursor; + set_slot(node, j, s); + } + node.data_start = @intCast(cursor); + node.count = k; + } + + /// The key at (node, i) in a form stable across node-array growth. + fn stable_key(self: *Index, node_id: u32, i: u32) StableKey { + const node = &self.nodes.items[node_id]; + const s = get_slot(node, i); + if (s.spill) return .{ .key = self.overflow.items[@intCast(s.off) .. @intCast(s.off + s.key_len)], .spill_off = s.off }; + const k = self.key_of(node_id, i); + @memcpy(self.promo[0..k.len], k); + return .{ .key = self.promo[0..k.len], .spill_off = 0 }; + } + + /// Entry position in a leaf, by (key, id). + fn leaf_pos(self: *const Index, leaf_id: u32, key: []const u8, id: []const u8) u32 { + const node = &self.nodes.items[leaf_id]; + var lo: u32 = 0; + var hi: u32 = node.count; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + const o = std.mem.order(u8, self.key_of(leaf_id, mid), key); + const less = o == .lt or (o == .eq and std.mem.order(u8, self.id_of(leaf_id, mid), id) == .lt); + if (less) lo = mid + 1 else hi = mid; + } + return lo; + } + + /// Separator position in an internal node: after any equal keys, so the + /// "last separator <= key" descent lands on the newest right child. + fn separator_pos(self: *const Index, node_id: u32, key: []const u8) u32 { + const node = &self.nodes.items[node_id]; + var lo: u32 = 0; + var hi: u32 = node.count; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (std.mem.order(u8, self.key_of(node_id, mid), key) != .gt) lo = mid + 1 else hi = mid; + } + return lo; + } + + /// The child holding the range `key` sorts into: right of the last + /// separator that is <= key (equal keys live right of equal separators). + fn descend_insert(self: *const Index, node_id: u32, key: []const u8) u32 { + const node = &self.nodes.items[node_id]; + var lo: u32 = 0; + var hi: u32 = node.count; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (std.mem.order(u8, self.key_of(node_id, mid), key) != .gt) lo = mid + 1 else hi = mid; + } + if (lo == 0) return node.first_child; + if (lo == node.count) { + if (node.count == 0) return node.first_child; + return get_slot(node, node.count - 1).extra; + } + return get_slot(node, lo - 1).extra; + } + + /// The child containing the lower bound of `prefix`: the child left of + /// the first separator not less than the prefix (prefix semantics), or + /// the rightmost child when every separator is less. + fn descend_lower(self: *const Index, node_id: u32, prefix: []const u8) u32 { + const node = &self.nodes.items[node_id]; + var lo: u32 = 0; + var hi: u32 = node.count; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (cmp_prefix(prefix, self.key_of(node_id, mid)) == .gt) lo = mid + 1 else hi = mid; + } + if (lo == node.count) { + // A one-child internal node (deletion does not rebalance) has + // no separators; its only child is first_child. + if (node.count == 0) return node.first_child; + return get_slot(node, node.count - 1).extra; + } + if (lo == 0) return node.first_child; + return get_slot(node, lo - 1).extra; + } + + /// Position in a leaf of the first slot whose key is not less than + /// `prefix` (prefix semantics). + fn leaf_lower(self: *const Index, leaf_id: u32, prefix: []const u8) u32 { + const node = &self.nodes.items[leaf_id]; + var lo: u32 = 0; + var hi: u32 = node.count; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (cmp_prefix(prefix, self.key_of(leaf_id, mid)) == .gt) lo = mid + 1 else hi = mid; + } + return lo; + } + + /// The (leaf, slot) of the first entry whose key is not less than + /// `prefix`. + fn lower_bound(self: *const Index, prefix: []const u8) struct { leaf: u32, slot: u32 } { + var node_id = self.root; + while (self.nodes.items[node_id].is_leaf == 0) node_id = self.descend_lower(node_id, prefix); + return .{ .leaf = node_id, .slot = self.leaf_lower(node_id, prefix) }; + } + + /// Insert one entry, splitting upward as needed. Infallible. + fn insert_entry(self: *Index, key: []const u8, id: []const u8) void { + if (self.insert_rec(self.root, key, id)) |up| { + // The root split: a new root with the two halves as children. + const new_root = self.alloc_node(); + self.nodes.items[new_root].first_child = self.root; + self.nodes.items[self.root].parent = new_root; + self.store_record(new_root, 0, .{ .key = up.key, .child = up.right, .spill_off = up.spill_off }); + self.nodes.items[up.right].parent = new_root; + self.root = new_root; + self.depth += 1; + } + } + + /// Descend and insert; return the split to promote at the level above, + /// or null when the subtree absorbed the record. + fn insert_rec(self: *Index, node_id: u32, key: []const u8, id: []const u8) ?Split { + const node = &self.nodes.items[node_id]; + if (node.is_leaf == 1) { + if (self.fits(node_id, key.len + id.len)) { + self.store_record(node_id, self.leaf_pos(node_id, key, id), .{ .key = key, .id = id }); + self.entry_count += 1; + return null; + } + return self.split_leaf(node_id, key, id); + } + const child = self.descend_insert(node_id, key); + const res = self.insert_rec(child, key, id) orelse return null; + return self.insert_separator(node_id, res); + } + + /// Split a full leaf: the right half (including any records equal to + /// the boundary key — lookups scan whole key bands, so equal keys may + /// live on both sides) moves to a new leaf, the boundary key is + /// promoted, and the new record is stored in whichever half holds its + /// position. + fn split_leaf(self: *Index, leaf_id: u32, key: []const u8, id: []const u8) ?Split { + const node = &self.nodes.items[leaf_id]; + const mid = @max(1, node.count / 2); + const right_id = self.alloc_node(); + { + const right = &self.nodes.items[right_id]; + right.is_leaf = 1; + right.next = node.next; + right.prev = leaf_id; + right.parent = node.parent; + } + // Move the right half; spilled records keep their slab reference. + var i: u32 = mid; + while (i < node.count) : (i += 1) { + const s = get_slot(&self.nodes.items[leaf_id], i); + self.store_record(right_id, @intCast(self.nodes.items[right_id].count), .{ + .key = self.key_of(leaf_id, i), + .id = self.id_of(leaf_id, i), + .spill_off = if (s.spill) s.off else 0, + }); + } + // Repack the left half in place. + self.repack_keep_prefix(leaf_id, mid); + // Link the chain. + const node2 = &self.nodes.items[leaf_id]; + if (node2.next != 0) self.nodes.items[node2.next].prev = right_id; + node2.next = right_id; + self.leaf_count += 1; + // Promote the right leaf's first key. + const sep = self.stable_key(right_id, 0); + // Store the new record in the correct half. + if (std.mem.order(u8, key, sep.key) == .lt) { + self.store_record(leaf_id, self.leaf_pos(leaf_id, key, id), .{ .key = key, .id = id }); + } else { + self.store_record(right_id, self.leaf_pos(right_id, key, id), .{ .key = key, .id = id }); + } + self.entry_count += 1; + return .{ .key = sep.key, .spill_off = sep.spill_off, .right = right_id }; + } + + /// Insert a promoted separator into an internal node, splitting it when + /// full. Returns the next promotion, or null. + fn insert_separator(self: *Index, node_id: u32, split: Split) ?Split { + // The incoming key may live in the promo buffer, which a nested + // split_internal (below) would overwrite with its own promoted key; + // spilled keys already live in the immutable slab. Copy inline keys + // to a stack buffer so they survive. + var local: [inline_limit]u8 = undefined; + const key: []const u8 = if (split.spill_off != 0) split.key else blk: { + @memcpy(local[0..split.key.len], split.key); + break :blk local[0..split.key.len]; + }; + if (self.fits(node_id, key.len)) { + self.store_record(node_id, self.separator_pos(node_id, key), .{ + .key = key, + .child = split.right, + .spill_off = split.spill_off, + }); + self.nodes.items[split.right].parent = node_id; + return null; + } + const up = self.split_internal(node_id); + // The node split; the separator goes into whichever half holds its + // position. Both halves have room (they are half-full). + if (std.mem.order(u8, key, up.key) == .lt) { + self.store_record(node_id, self.separator_pos(node_id, key), .{ + .key = key, + .child = split.right, + .spill_off = split.spill_off, + }); + self.nodes.items[split.right].parent = node_id; + } else { + self.store_record(up.right, self.separator_pos(up.right, key), .{ + .key = key, + .child = split.right, + .spill_off = split.spill_off, + }); + self.nodes.items[split.right].parent = up.right; + } + return up; + } + + /// Split a full internal node: the middle separator is promoted, the + /// first half stays, the rest moves to a new right node. + fn split_internal(self: *Index, node_id: u32) Split { + const node = &self.nodes.items[node_id]; + const s = node.count / 2; + // Copy the promoted key before the repack rewrites the page. + const mid_key = self.stable_key(node_id, s); + const right_id = self.alloc_node(); + { + const right = &self.nodes.items[right_id]; + right.is_leaf = 0; + right.parent = node.parent; + right.first_child = get_slot(node, s).extra; + // The moved subtrees now live under the right node. + self.nodes.items[right.first_child].parent = right_id; + } + var i: u32 = s + 1; + while (i < node.count) : (i += 1) { + const slot_i = get_slot(&self.nodes.items[node_id], i); + self.nodes.items[slot_i.extra].parent = right_id; + self.store_record(right_id, @intCast(self.nodes.items[right_id].count), .{ + .key = self.key_of(node_id, i), + .child = slot_i.extra, + .spill_off = if (slot_i.spill) slot_i.off else 0, + }); + } + self.repack_keep_prefix(node_id, s); + return .{ .key = mid_key.key, .spill_off = mid_key.spill_off, .right = right_id }; + } + + /// Remove the exact entry (key, id). Equal keys may span several + /// leaves, so the entry is found by scanning its key band from the + /// lower bound rather than by a straight descent. + fn remove_entry(self: *Index, key: []const u8, id: []const u8) bool { + var it = self.seek(key); + while (it.next()) |e| { + const c = std.mem.order(u8, e.key, key); + if (c == .gt) return false; // band ended, not found + if (c != .eq) continue; + if (std.mem.eql(u8, e.id, id)) { + self.leaf_remove(it.leaf, it.slot - 1); + return true; + } + } + return false; + } + + /// Remove leaf slot and, when the leaf empties, unlink it and drop it + /// from its parent, cascading up while internal nodes empty. + fn leaf_remove(self: *Index, leaf_id: u32, slot_idx: u32) void { + self.remove_record(leaf_id, slot_idx); + self.entry_count -= 1; + const node = &self.nodes.items[leaf_id]; + if (node.count > 0) return; + // Empty leaf: unlink and drop from the parent (unless it is the + // root, which stays as the empty root leaf). + if (leaf_id == self.root) return; + if (node.prev != 0) self.nodes.items[node.prev].next = node.next; + if (node.next != 0) self.nodes.items[node.next].prev = node.prev; + if (leaf_id == self.first_leaf) self.first_leaf = node.next; + self.leaf_count -= 1; + var child = leaf_id; + var parent = node.parent; + while (parent != 0) { + self.drop_child(parent, child); + const pnode = &self.nodes.items[parent]; + if (pnode.count == 0 and pnode.first_child == 0) { + if (parent == self.root) { + self.replace_root_with_leaf(); + return; + } + child = parent; + parent = pnode.parent; + } else { + return; + } + } + } + + /// Remove `child` from `parent`'s child list. Infallible; `child`'s + /// node is abandoned in place (its page is simply never referenced + /// again — node ids are append-only, so this is a leak of at most the + /// peak tree size, exactly what the old entry array's capacity was). + fn drop_child(self: *Index, parent_id: u32, child_id: u32) void { + const pnode = &self.nodes.items[parent_id]; + if (pnode.first_child == child_id) { + if (pnode.count > 0) { + const s0 = get_slot(pnode, 0); + pnode.first_child = s0.extra; + self.remove_record(parent_id, 0); + } else { + pnode.first_child = 0; + } + return; + } + var j: u32 = 0; + while (j < pnode.count) : (j += 1) { + if (get_slot(pnode, j).extra == child_id) { + self.remove_record(parent_id, j); + return; + } + } + } + + /// The internal root emptied: swap in a fresh empty leaf as the root. + fn replace_root_with_leaf(self: *Index) void { + self.root = self.alloc_node(); + self.nodes.items[self.root].is_leaf = 1; + self.nodes.items[self.root].parent = 0; + self.first_leaf = self.root; + self.leaf_count = 1; + self.depth = 0; + } + + /// Pack the sorted staging array into a fresh tree: leaves filled in + /// key order, interior levels built bottom-up. Frees the staging keys + /// and clears the staging array. + fn pack_tree(self: *Index, gpa: std.mem.Allocator) !void { + defer { + for (self.staging.items) |e| gpa.free(e.key); + self.staging.clearRetainingCapacity(); + } + if (self.staging.items.len == 0) return; + + const Level = struct { + id: u32, + first_key: []const u8, // aliases a leaf's page or the slab + }; + var level: std.ArrayListUnmanaged(Level) = .empty; + defer level.deinit(gpa); + + // Leaves. + var prev: u32 = 0; + var lit = self.staging.items; + while (lit.len > 0) { + const leaf = self.alloc_node(); + self.nodes.items[leaf].is_leaf = 1; + const first_key = lit[0].key; + // Fill until the next record would not fit. + var used: usize = 0; + var slots: usize = 0; + var n: usize = 0; + while (n < lit.len) : (n += 1) { + const rec_len: u64 = lit[n].key.len + lit[n].id.len; + const inline_bytes: usize = if (rec_len > inline_limit) 0 else @intCast(rec_len); + if ((slots + 1) * slot_size + used + inline_bytes > page_data) break; + self.store_record(leaf, @intCast(slots), .{ .key = lit[n].key, .id = lit[n].id }); + slots += 1; + used += inline_bytes; + } + self.nodes.items[leaf].prev = prev; + if (prev != 0) self.nodes.items[prev].next = leaf; + prev = leaf; + self.leaf_count += 1; + try level.append(gpa, .{ .id = leaf, .first_key = first_key }); + lit = lit[n..]; + } + self.nodes.items[prev].next = 0; + self.first_leaf = level.items[0].id; + + // Interior levels: group the level below into internal nodes whose + // separators are the first keys of its members. + while (level.items.len > 1) { + var next: std.ArrayListUnmanaged(Level) = .empty; + defer next.deinit(gpa); + var i: usize = 0; + while (i < level.items.len) { + const node = self.alloc_node(); + self.nodes.items[node].first_child = level.items[i].id; + self.nodes.items[level.items[i].id].parent = node; + const first_key = level.items[i].first_key; + var used: usize = 0; + var slots: usize = 0; + var j = i + 1; + while (j < level.items.len) : (j += 1) { + const rec_len: u64 = level.items[j].first_key.len; + const inline_bytes: usize = if (rec_len > inline_limit) 0 else @intCast(rec_len); + if ((slots + 1) * slot_size + used + inline_bytes > page_data) break; + self.store_record(node, @intCast(slots), .{ + .key = level.items[j].first_key, + .child = level.items[j].id, + .spill_off = 0, + }); + self.nodes.items[level.items[j].id].parent = node; + slots += 1; + used += inline_bytes; + } + try next.append(gpa, .{ .id = node, .first_key = first_key }); + i = j; + } + level.deinit(gpa); + level = next; + next = .empty; // ownership moved to `level`; the defer frees nothing + self.depth += 1; + } + self.root = level.items[0].id; + self.nodes.items[self.root].parent = 0; + self.entry_count = self.staging.items.len; + } }; /// The index whose key pattern is exactly `key_pairs` (same paths, same @@ -598,7 +1311,7 @@ fn default_name(gpa: std.mem.Allocator, key_pairs: []const bson.Pair) ![]u8 { // --------------------------------------------------------------------------- /// Encoded-key byte order, tie-broken by the id bytes. This is the total -/// order the entry array is kept in. +/// order batches are sorted into before packing. /// /// bson.encode_key guarantees this reproduces component-wise bson.compare /// exactly, so ordering is a memcmp. Key direction is deliberately not @@ -614,19 +1327,20 @@ fn entry_less(_: void, a: Entry, b: Entry) bool { return compare_entries(a, b) == .lt; } -/// Order of an encoded `prefix` against an entry key — `.eq` when the key -/// starts with it. The search-side counterpart of compare_entries: no id -/// tie-break, so a partial key matches a whole range. +/// Order of an encoded search prefix against a stored key, the tree's +/// lookup comparator: `.eq` when the key starts with the prefix, so a +/// partial key matches a whole band; `.lt` when the key sorts strictly +/// before the prefix; `.gt` when after. /// /// Comparing raw bytes is sound because every column encoding is /// self-delimiting, so a prefix of the encoded key is exactly the encoding /// of its leading columns. -fn prefix_order(prefix: []const u8, e: Entry) std.math.Order { - const n = @min(prefix.len, e.key.len); - const o = std.mem.order(u8, prefix[0..n], e.key[0..n]); +fn cmp_prefix(prefix: []const u8, key: []const u8) std.math.Order { + const n = @min(prefix.len, key.len); + const o = std.mem.order(u8, prefix[0..n], key[0..n]); if (o != .eq) return o; // The key ran out first, so it sorts below the prefix. - if (prefix.len > e.key.len) return .gt; + if (prefix.len > key.len) return .gt; return .eq; } @@ -783,7 +1497,7 @@ pub const Plan = struct { try self.index.lookup_range(gpa, key, self.lo, self.lo_incl, self.hi, self.hi_incl, out); } } - // Entries come out of the array in key order, so a forward scan is + // Entries come out of the leaves in key order, so a forward scan is // already sorted; a backward one just reads it the other way. This // must happen before the dedupe pass below, which sorts by id and // would destroy the order — the conditions that set provides_sort @@ -843,13 +1557,13 @@ fn plan_better(a: *const Plan, b: *const Plan) bool { } /// Whether scanning `ix` in order satisfies `sort`, and if so whether that -/// means reading the array backwards. +/// means reading the tree backwards. /// /// The sort keys must line up with the index components that follow the /// equality-pinned prefix: those components are fixed to one value each, so /// they do not affect the order of what remains. Directions must agree /// uniformly — every key the same way round, or every key opposite — since -/// the array can only be read forwards or backwards. +/// the leaves can only be read forwards or backwards. /// /// A multikey index is excluded: it emits a document once per indexed /// value, so its order is not an order on documents. @@ -1101,6 +1815,19 @@ fn expect_ids(gpa: std.mem.Allocator, ix: *const Index, key: []const bson.Value, for (out.items, 0..) |id, i| try testing.expectEqualStrings(expected[i], id); } +/// The tree's entries as (key, id) references, in iteration order. +fn refs_of(gpa: std.mem.Allocator, ix: *const Index, out: *std.ArrayListUnmanaged(EntryRef)) !void { + var it = ix.iter(); + while (it.next()) |e| try out.append(gpa, e); +} + +/// EntryRef ordering, mirroring compare_entries: key bytes, then id bytes. +fn ref_lt(a: EntryRef, b: EntryRef) bool { + const o = std.mem.order(u8, a.key, b.key); + if (o != .eq) return o == .lt; + return std.mem.order(u8, a.id, b.id) == .lt; +} + test "entries sort across numeric types and string/null/objectid" { const gpa = testing.allocator; var ix = try simple_index(gpa, &.{"a"}, false, false); @@ -1123,10 +1850,15 @@ test "entries sort across numeric types and string/null/objectid" { try expect_ids(gpa, &ix, &.{.{ .string = "b" }}, &.{"i3"}); try expect_ids(gpa, &ix, &.{.{ .object_id = [_]u8{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 } }}, &.{"i5"}); - // Full order: null, 5, "b", oid — sorted entries respect it. - try testing.expect(compare_entries(ix.entries.items[0], ix.entries.items[1]) == .lt); - try testing.expect(compare_entries(ix.entries.items[1], ix.entries.items[2]) == .lt); - try testing.expect(compare_entries(ix.entries.items[2], ix.entries.items[3]) == .lt); + // Full order: null, 5, 5, "b", oid — iteration respects it (equal + // keys tie-break by id, as the leaf position rule does). + var refs: std.ArrayListUnmanaged(EntryRef) = .empty; + defer refs.deinit(gpa); + try refs_of(gpa, &ix, &refs); + try testing.expectEqual(@as(usize, 5), refs.items.len); + try testing.expect(ref_lt(refs.items[0], refs.items[1])); + try testing.expect(ref_lt(refs.items[1], refs.items[2])); + try testing.expect(ref_lt(refs.items[2], refs.items[3])); } test "missing field is indexed as null; sparse skips the document" { @@ -1135,13 +1867,13 @@ test "missing field is indexed as null; sparse skips the document" { defer ix.deinit(gpa); const d = doc_of(&.{.{ .key = "_id", .value = .{ .int32 = 1 } }}); _ = try ix.add_doc(gpa, &d, "m1", true); - try testing.expectEqual(@as(usize, 1), ix.entries.items.len); + try testing.expectEqual(@as(usize, 1), ix.count()); try expect_ids(gpa, &ix, &.{.null}, &.{"m1"}); var sp = try simple_index(gpa, &.{"a"}, false, true); defer sp.deinit(gpa); _ = try sp.add_doc(gpa, &d, "m2", true); - try testing.expectEqual(@as(usize, 0), sp.entries.items.len); + try testing.expectEqual(@as(usize, 0), sp.count()); } test "multikey expansion indexes the array and its elements" { @@ -1153,7 +1885,7 @@ test "multikey expansion indexes the array and its elements" { _ = try ix.add_doc(gpa, &d, "mk1", true); // 3 entries: the array itself, "a", "b". - try testing.expectEqual(@as(usize, 3), ix.entries.items.len); + try testing.expectEqual(@as(usize, 3), ix.count()); // Element query. try expect_ids(gpa, &ix, &.{.{ .string = "a" }}, &.{"mk1"}); try expect_ids(gpa, &ix, &.{.{ .string = "b" }}, &.{"mk1"}); @@ -1169,7 +1901,7 @@ test "per-document dedup keeps {a: [1,1]} under a unique index" { const d = doc_of(&.{.{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 1 } } } }}); _ = try ix.add_doc(gpa, &d, "d1", true); // Entries after dedup: the array itself and one element. - try testing.expectEqual(@as(usize, 2), ix.entries.items.len); + try testing.expectEqual(@as(usize, 2), ix.count()); } test "parallel arrays are rejected" { @@ -1190,7 +1922,7 @@ test "parallel arrays are rejected" { .{ .key = "b", .value = .{ .int32 = 3 } }, }); _ = try ix.add_doc(gpa, &ok, "p2", true); - try testing.expectEqual(@as(usize, 3), ix.entries.items.len); + try testing.expectEqual(@as(usize, 3), ix.count()); } test "unique conflict across documents, replace of own entries allowed" { @@ -1259,7 +1991,7 @@ test "empty index and remove_id" { const d = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .int32 = 1 } } }); _ = try ix.add_doc(gpa, &d, "e1", true); ix.remove_id(gpa, "e1"); - try testing.expectEqual(@as(usize, 0), ix.entries.items.len); + try testing.expectEqual(@as(usize, 0), ix.count()); try ix.lookup_eq(gpa, &.{.{ .int32 = 1 }}, &out); try testing.expectEqual(@as(usize, 0), out.items.len); } @@ -1358,42 +2090,139 @@ test "remove_doc leaves the index identical to a full scan removal" { try by_doc.append_doc_entries(gpa, &docs[i], id); try by_scan.append_doc_entries(gpa, &docs[i], id); } - _ = try by_doc.finish_bulk(false); - _ = try by_scan.finish_bulk(false); - try testing.expectEqual(by_scan.entries.items.len, by_doc.entries.items.len); + _ = try by_doc.finish_bulk(gpa, false); + _ = try by_scan.finish_bulk(gpa, false); + try testing.expectEqual(by_scan.count(), by_doc.count()); // Remove in a shuffled order, so removals interleave rather than - // peeling the array from one end. + // peeling the leaves from one end. var order: [n]usize = undefined; for (0..n) |i| order[i] = i; rand.shuffle(usize, &order); + var doc_refs: std.ArrayListUnmanaged(EntryRef) = .empty; + var scan_refs: std.ArrayListUnmanaged(EntryRef) = .empty; + defer doc_refs.deinit(gpa); + defer scan_refs.deinit(gpa); for (order) |i| { by_doc.remove_doc(gpa, &docs[i], ids.items[i]); by_scan.remove_id(gpa, ids.items[i]); - testing.expectEqual(by_scan.entries.items.len, by_doc.entries.items.len) catch |err| { + doc_refs.clearRetainingCapacity(); + scan_refs.clearRetainingCapacity(); + try refs_of(gpa, &by_doc, &doc_refs); + try refs_of(gpa, &by_scan, &scan_refs); + + testing.expectEqual(scan_refs.items.len, doc_refs.items.len) catch |err| { std.debug.print("sparse={} after removing doc {d}: scan left {d}, remove_doc left {d}\n", .{ - sparse, i, by_scan.entries.items.len, by_doc.entries.items.len, + sparse, i, scan_refs.items.len, doc_refs.items.len, }); return err; }; - for (by_doc.entries.items, by_scan.entries.items) |x, y| { + for (doc_refs.items, scan_refs.items) |x, y| { try testing.expect(std.mem.eql(u8, x.key, y.key)); try testing.expect(std.mem.eql(u8, x.id, y.id)); } } - try testing.expectEqual(@as(usize, 0), by_doc.entries.items.len); + try testing.expectEqual(@as(usize, 0), by_doc.count()); } } +test "incremental inserts and removals stay identical to a brute-force model" { + // The bulk-path differential covers packing; this one covers the + // incremental path: one insert at a time (leaf splits), interleaved + // with removals (empty-leaf unlinks, one-child internal nodes), all + // checked against a reference model at every step. + const gpa = testing.allocator; + var prng = std.Random.DefaultPrng.init(0x0dd_ba11); + const rand = prng.random(); + + var ix = try simple_index(gpa, &.{ "a", "b" }, false, false); + defer ix.deinit(gpa); + + var model: std.ArrayListUnmanaged(ModelFact) = .empty; + defer { + for (model.items) |m| gpa.free(m.id); + model.deinit(gpa); + } + var live: std.ArrayListUnmanaged(bool) = .empty; + defer live.deinit(gpa); + + const n = 600; + for (0..n) |i| { + // Insert. + const a = rand.intRangeAtMost(i32, 0, 30); + const b = rand.intRangeAtMost(i32, 0, 30); + const id = try std.fmt.allocPrint(gpa, "id{d}", .{i}); + const d = doc_of(&.{ + .{ .key = "_id", .value = .{ .int32 = @intCast(i) } }, + .{ .key = "a", .value = .{ .int32 = a } }, + .{ .key = "b", .value = .{ .int32 = b } }, + }); + _ = try ix.add_doc(gpa, &d, id, false); + try model.append(gpa, .{ .a = a, .b = b, .id = id }); + try live.append(gpa, true); + + // Verify equality and one-sided range against the model. + try verify_model(gpa, &ix, model.items, live.items, rand); + } + // Remove a random half in random order. + var order: std.ArrayListUnmanaged(usize) = .empty; + defer order.deinit(gpa); + for (0..n) |i| if (rand.boolean()) try order.append(gpa, i); + rand.shuffle(usize, order.items); + for (order.items) |i| { + const m = model.items[i]; + const d = doc_of(&.{ + .{ .key = "_id", .value = .{ .int32 = @intCast(i) } }, + .{ .key = "a", .value = .{ .int32 = m.a } }, + .{ .key = "b", .value = .{ .int32 = m.b } }, + }); + ix.remove_doc(gpa, &d, m.id); + live.items[i] = false; + try verify_model(gpa, &ix, model.items, live.items, rand); + } +} + +/// One document's facts in the incremental-mutation differential. +const ModelFact = struct { a: i32, b: i32, id: []const u8 }; + +fn verify_model(gpa: std.mem.Allocator, ix: *const Index, model: []const ModelFact, live: []const bool, rand: std.Random) !void { + // lookup_eq over a random value. + const a = rand.intRangeAtMost(i32, 0, 30); + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .int32 = a }}, &out); + var expected: usize = 0; + for (model, 0..) |m, mi| { + if (live[mi] and m.a == a) expected += 1; + } + try testing.expectEqual(expected, out.items.len); + + // lookup_range over a random one-sided range. + out.clearRetainingCapacity(); + const lo_v = rand.intRangeAtMost(i32, -5, 35); + const lo_incl = rand.boolean(); + try ix.lookup_range(gpa, &.{.{ .int32 = a }}, .{ .int32 = lo_v }, lo_incl, null, false, &out); + expected = 0; + for (model, 0..) |m, mi| { + if (!live[mi]) continue; + if (m.a != a) continue; + if (m.b < lo_v) continue; + if (m.b == lo_v and !lo_incl) continue; + expected += 1; + } + try testing.expectEqual(expected, out.items.len); +} + test "lookup_range matches a brute-force filter over random data" { const gpa = testing.allocator; - // lookup_range binary-searches both ends instead of scanning the - // equality band. Bounds like that are easy to get subtly wrong at the - // inclusive/exclusive edges and where the band ends, so check the whole - // result set against the definition rather than spot-checking. + // lookup_range seeks both ends of the range inside the tree instead of + // scanning the whole equality band. Bounds like that are easy to get + // subtly wrong at the inclusive/exclusive edges and where the band + // ends, so check the whole result set against the definition rather + // than spot-checking. var prng = std.Random.DefaultPrng.init(0x5eed_1234); const rand = prng.random(); @@ -1423,7 +2252,7 @@ test "lookup_range matches a brute-force filter over random data" { }); try ix.append_doc_entries(gpa, &d, id); } - _ = try ix.finish_bulk(false); + _ = try ix.finish_bulk(gpa, false); var out: std.ArrayListUnmanaged([]const u8) = .empty; defer out.deinit(gpa); diff --git a/src/spill.zig b/src/spill.zig new file mode 100644 index 0000000..8c377d5 --- /dev/null +++ b/src/spill.zig @@ -0,0 +1,74 @@ +// Dev stress test for the overflow slab (records > 1024 bytes): +// zig run -O ReleaseFast src/spill.zig +// Keys straddling the spill threshold (10 B .. 100 KB) through insert, +// lookup, delete and iteration. +const std = @import("std"); +const index = @import("index.zig"); +const bson = @import("bson.zig"); + +fn doc_of(pairs: []const bson.Pair) bson.Document { + return .{ .arena = undefined, .pairs = pairs }; +} + +pub fn main() !void { + const gpa = std.heap.page_allocator; + var prng = std.Random.DefaultPrng.init(0x1234_5678); + const rand = prng.random(); + + var keys = [_]index.IndexKey{.{ .path = "tag", .descending = false }}; + var ix = try index.Index.init(gpa, "tag", &keys, false, false, null); + + // Keys straddling the spill threshold: inline, exactly at the limit, + // just over, and one very long. Each id is a short static string. + const lens = [_]usize{ 10, 1023, 1024, 1025, 2000, 100_000 }; + var strings: [lens.len][]u8 = undefined; + var docs: [lens.len]bson.Document = undefined; + var pairs: [2]bson.Pair = undefined; + for (lens, 0..) |len, i| { + strings[i] = try gpa.alloc(u8, len); + for (strings[i]) |*c| c.* = 'a' + @as(u8, @intCast(rand.intRangeAtMost(u8, 0, 25))); + // add a distinguishing suffix so keys are unique + std.mem.copyForwards(u8, strings[i][len - 4 ..], &[_]u8{ @intCast(i), 0xff, 0x00, 0x00 }); + pairs[0] = .{ .key = "_id", .value = .{ .int32 = @intCast(i) } }; + pairs[1] = .{ .key = "tag", .value = .{ .string = strings[i] } }; + docs[i] = doc_of(&pairs); + _ = try ix.add_doc(gpa, &docs[i], &[_]u8{ 'i', 'd', @intCast(i + 1) }, true); + } + std.debug.print("count={d} leaves={d} depth={d} overflow={d}\n", .{ ix.count(), ix.leaf_count, ix.depth, ix.overflow.items.len }); + if (ix.overflow.items.len < 100_000) return error.NoSpill; + + // Every entry is found by exact key. + for (lens, 0..) |_, i| { + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .string = strings[i] }}, &out); + if (out.items.len != 1) { + std.debug.print("lookup {d} got {d}\n", .{ i, out.items.len }); + return error.Bad; + } + if (!std.mem.eql(u8, out.items[0], &[_]u8{ 'i', 'd', @intCast(i + 1) })) return error.Bad; + } + + // Delete the spilled ones and the inline ones alternately. + for (lens, 0..) |_, i| { + if (i % 2 == 0) continue; + ix.remove_doc(gpa, &docs[i], &[_]u8{ 'i', 'd', @intCast(i + 1) }); + } + if (ix.count() != 3) return error.Bad; + for (lens, 0..) |_, i| { + if (i % 2 == 0) continue; + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .string = strings[i] }}, &out); + if (out.items.len != 0) return error.Bad; + } + // Iteration still sees the survivors in order. + var it = ix.iter(); + var seen: usize = 0; + while (it.next()) |e| { + seen += 1; + _ = e; + } + if (seen != 3) return error.Bad; + std.debug.print("SPILL OK (seen={d})\n", .{seen}); +} diff --git a/src/spill2.zig b/src/spill2.zig new file mode 100644 index 0000000..3fa167b --- /dev/null +++ b/src/spill2.zig @@ -0,0 +1,88 @@ +// Dev stress test: spilled records through leaf splits and internal levels. +// zig run -O ReleaseFast src/spill2.zig +// 5000 docs with 2 KiB keys: every record spills; the tree still answers +// exact lookups and deletes half of them. +const std = @import("std"); +const index = @import("index.zig"); +const bson = @import("bson.zig"); + +fn doc_of(pairs: []const bson.Pair) bson.Document { + return .{ .arena = undefined, .pairs = pairs }; +} + +pub fn main() !void { + const gpa = std.heap.page_allocator; + + var keys = [_]index.IndexKey{.{ .path = "tag", .descending = false }}; + var ix = try index.Index.init(gpa, "tag", &keys, false, false, null); + + // 5000 docs, each with a 2 KiB key: spills on every record, forcing + // leaves and internal nodes to hold overflow references. + const N = 5000; + var ids: std.ArrayListUnmanaged([]u8) = .empty; + defer { + for (ids.items) |x| gpa.free(x); + ids.deinit(gpa); + } + var pairs: [2]bson.Pair = undefined; + var buf = try gpa.alloc(u8, 2000); + defer gpa.free(buf); + var facts: std.ArrayListUnmanaged(struct { key: []u8 }) = .empty; + defer { + for (facts.items) |f| gpa.free(f.key); + facts.deinit(gpa); + } + for (0..N) |i| { + for (buf) |*c| c.* = 'x'; + // unique suffix + std.mem.writeInt(u32, buf[0..4], @intCast(i), .little); + const key = try gpa.dupe(u8, buf); + try facts.append(gpa, .{ .key = key }); + const id = try std.fmt.allocPrint(gpa, "id{d}", .{i}); + try ids.append(gpa, id); + pairs[0] = .{ .key = "_id", .value = .{ .int32 = @intCast(i) } }; + pairs[1] = .{ .key = "tag", .value = .{ .string = key } }; + const d = doc_of(&pairs); + _ = try ix.add_doc(gpa, &d, id, false); + } + std.debug.print("count={d} leaves={d} depth={d} overflow={d}\n", .{ ix.count(), ix.leaf_count, ix.depth, ix.overflow.items.len }); + if (ix.count() != N) return error.Bad; + + // Spot-check exact lookups. + var prng2 = std.Random.DefaultPrng.init(0xabc); + const rand2 = prng2.random(); + for (0..300) |_| { + const i = rand2.intRangeAtMost(usize, 0, N - 1); + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .string = facts.items[i].key }}, &out); + if (out.items.len != 1 or !std.mem.eql(u8, out.items[0], ids.items[i])) { + std.debug.print("lookup mismatch at {d}\n", .{i}); + return error.Bad; + } + } + // Delete half (random), verify count and no leftover. + var order: std.ArrayListUnmanaged(usize) = .empty; + defer order.deinit(gpa); + for (0..N) |i| if (i % 2 == 0) try order.append(gpa, i); + rand2.shuffle(usize, order.items); + var removed: usize = 0; + for (order.items) |i| { + pairs[0] = .{ .key = "_id", .value = .{ .int32 = @intCast(i) } }; + pairs[1] = .{ .key = "tag", .value = .{ .string = facts.items[i].key } }; + const d = doc_of(&pairs); + ix.remove_doc(gpa, &d, ids.items[i]); + removed += 1; + if (ix.count() != N - removed) { + std.debug.print("count mismatch at {d}: {d} != {d}\n", .{ i, ix.count(), N - removed }); + return error.Bad; + } + } + for (order.items) |i| { + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .string = facts.items[i].key }}, &out); + if (out.items.len != 0) return error.Bad; + } + std.debug.print("SPILL2 OK (leaves={d} depth={d})\n", .{ ix.leaf_count, ix.depth }); +} diff --git a/src/stress.zig b/src/stress.zig new file mode 100644 index 0000000..f877560 --- /dev/null +++ b/src/stress.zig @@ -0,0 +1,174 @@ +// Dev stress test for the index B+tree (not part of the build): +// zig run -O ReleaseFast src/stress.zig +// Bulk-builds 30k entries, inserts 20k more one at a time (splits at depth 2), +// deletes 16.6k randomly (empty-leaf cascades), drains everything, and checks +// lookups against a brute-force model throughout. +const std = @import("std"); +const index = @import("index.zig"); +const bson = @import("bson.zig"); + +fn doc_of(pairs: []const bson.Pair) bson.Document { + return .{ .arena = undefined, .pairs = pairs }; +} + +const Fact = struct { a: i32, b: i32 }; + +fn check_range(gpa: std.mem.Allocator, ix: *const index.Index, prefix: bson.Value, lo: ?bson.Value, hi: ?bson.Value, facts: []const Fact, alive: []const bool) !void { + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_range(gpa, &.{prefix}, lo, true, hi, false, &out); + var expected: usize = 0; + for (facts, 0..) |f, fi| { + if (!alive[fi]) continue; + if (f.a != prefix.int32) continue; + if (lo) |l| if (f.b < l.int32) continue; + if (hi) |h| if (f.b >= h.int32) continue; + expected += 1; + } + if (out.items.len != expected) { + std.debug.print("MISMATCH: prefix={d} lo={?d} hi={?d}: got {d}, want {d}\n", .{ prefix.int32, if (lo) |l| l.int32 else null, if (hi) |h| h.int32 else null, out.items.len, expected }); + std.process.exit(1); + } +} + +pub fn main() !void { + const gpa = std.heap.page_allocator; + var prng = std.Random.DefaultPrng.init(0xBADCAFE); + const rand = prng.random(); + + const N = 30_000; + var ids: std.ArrayListUnmanaged([]u8) = .empty; + var facts: std.ArrayListUnmanaged(Fact) = .empty; + var alive: std.ArrayListUnmanaged(bool) = .empty; + var pairs: [2]bson.Pair = undefined; + + // 1. Bulk build an index over N docs. + var keys = [_]index.IndexKey{ .{ .path = "a", .descending = false }, .{ .path = "b", .descending = false } }; + var ix = try index.Index.init(gpa, "ab", &keys, false, false, null); + for (0..N) |i| { + const id = try std.fmt.allocPrint(gpa, "id{d}", .{i}); + try ids.append(gpa, id); + const a = rand.intRangeAtMost(i32, 0, 99); + const b = rand.intRangeAtMost(i32, 0, 999); + try facts.append(gpa, .{ .a = a, .b = b }); + try alive.append(gpa, true); + pairs[0] = .{ .key = "a", .value = .{ .int32 = a } }; + pairs[1] = .{ .key = "b", .value = .{ .int32 = b } }; + const d = doc_of(&pairs); + try ix.append_doc_entries(gpa, &d, id); + } + _ = try ix.finish_bulk(gpa, false); + std.debug.print("bulk: count={d} leaves={d} depth={d} nodes={d}\n", .{ ix.count(), ix.leaf_count, ix.depth, ix.nodes.items.len }); + if (ix.count() != N) return error.BadCount; + + // Random range checks against brute force. + for (0..500) |_| { + const a = rand.intRangeAtMost(i32, 0, 99); + const lo_v = rand.intRangeAtMost(i32, -10, 1009); + const hi_v = rand.intRangeAtMost(i32, -10, 1009); + const use_lo = rand.boolean(); + const use_hi = rand.boolean(); + try check_range(gpa, &ix, .{ .int32 = a }, if (use_lo) .{ .int32 = lo_v } else null, if (use_hi) .{ .int32 = hi_v } else null, facts.items, alive.items); + } + + // 2. Incremental inserts, random order (splits + rebalancing-free path). + const M = 20_000; + for (0..M) |i| { + const id = try std.fmt.allocPrint(gpa, "new{d}", .{i}); + try ids.append(gpa, id); + const a = rand.intRangeAtMost(i32, 0, 99); + const b = rand.intRangeAtMost(i32, 0, 999); + try facts.append(gpa, .{ .a = a, .b = b }); + try alive.append(gpa, true); + pairs[0] = .{ .key = "a", .value = .{ .int32 = a } }; + pairs[1] = .{ .key = "b", .value = .{ .int32 = b } }; + const d = doc_of(&pairs); + _ = try ix.add_doc(gpa, &d, id, false); + } + std.debug.print("after inserts: count={d} leaves={d} depth={d} nodes={d}\n", .{ ix.count(), ix.leaf_count, ix.depth, ix.nodes.items.len }); + if (ix.count() != N + M) return error.BadCount; + for (0..500) |_| { + const a = rand.intRangeAtMost(i32, 0, 99); + const lo_v = rand.intRangeAtMost(i32, -10, 1009); + const hi_v = rand.intRangeAtMost(i32, -10, 1009); + const use_lo = rand.boolean(); + const use_hi = rand.boolean(); + try check_range(gpa, &ix, .{ .int32 = a }, if (use_lo) .{ .int32 = lo_v } else null, if (use_hi) .{ .int32 = hi_v } else null, facts.items, alive.items); + } + + // 3. Delete every 3rd doc in random order (empty-leaf cascades, + // one-child internals). + var order: std.ArrayListUnmanaged(usize) = .empty; + for (0..N + M) |i| if (i % 3 == 0) try order.append(gpa, i); + rand.shuffle(usize, order.items); + for (order.items) |i| { + const a = facts.items[i].a; + const b = facts.items[i].b; + pairs[0] = .{ .key = "a", .value = .{ .int32 = a } }; + pairs[1] = .{ .key = "b", .value = .{ .int32 = b } }; + const d = doc_of(&pairs); + ix.remove_doc(gpa, &d, ids.items[i]); + alive.items[i] = false; + } + std.debug.print("after deletes: count={d} leaves={d} depth={d} nodes={d}\n", .{ ix.count(), ix.leaf_count, ix.depth, ix.nodes.items.len }); + if (ix.count() != (N + M) - (N + M) / 3 - 1) return error.BadCount; + for (0..500) |_| { + const a = rand.intRangeAtMost(i32, 0, 99); + const lo_v = rand.intRangeAtMost(i32, -10, 1009); + const hi_v = rand.intRangeAtMost(i32, -10, 1009); + const use_lo = rand.boolean(); + const use_hi = rand.boolean(); + try check_range(gpa, &ix, .{ .int32 = a }, if (use_lo) .{ .int32 = lo_v } else null, if (use_hi) .{ .int32 = hi_v } else null, facts.items, alive.items); + } + + // 4. Equality lookups still exact. + for (0..300) |_| { + const a = rand.intRangeAtMost(i32, 0, 99); + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .int32 = a }}, &out); + var expected: usize = 0; + for (facts.items, 0..) |f, fi| { if (alive.items[fi] and f.a == a) expected += 1; } + if (out.items.len != expected) { + std.debug.print("EQ MISMATCH a={d}: got {d} want {d}\n", .{ a, out.items.len, expected }); + return error.BadCount; + } + } + + // 5. Delete everything (empty-leaf cascades, one-child internals), + // then verify the tree still works for fresh inserts. + var live: std.ArrayListUnmanaged(usize) = .empty; + defer live.deinit(gpa); + for (facts.items, 0..) |_, i| if (alive.items[i]) try live.append(gpa, i); + rand.shuffle(usize, live.items); + var remaining = live.items.len; + for (live.items) |i| { + const a = facts.items[i].a; + const b = facts.items[i].b; + pairs[0] = .{ .key = "a", .value = .{ .int32 = a } }; + pairs[1] = .{ .key = "b", .value = .{ .int32 = b } }; + const d = doc_of(&pairs); + ix.remove_doc(gpa, &d, ids.items[i]); + remaining -= 1; + if (ix.count() != remaining) { + std.debug.print("COUNT MISMATCH during drain: {d} != {d}\n", .{ ix.count(), remaining }); + return error.BadCount; + } + } + std.debug.print("after full drain: count={d} leaves={d} depth={d} nodes={d}\n", .{ ix.count(), ix.leaf_count, ix.depth, ix.nodes.items.len }); + if (ix.count() != 0) return error.BadCount; + // The drained tree still accepts and finds entries. + pairs[0] = .{ .key = "a", .value = .{ .int32 = 7 } }; + pairs[1] = .{ .key = "b", .value = .{ .int32 = 42 } }; + const d2 = doc_of(&pairs); + _ = try ix.add_doc(gpa, &d2, "final", false); + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .int32 = 7 }}, &out); + if (out.items.len != 1) return error.BadCount; + var it = ix.iter(); + if (it.next() == null) return error.BadCount; + if (it.next() != null) return error.BadCount; + + std.debug.print("STRESS OK\n", .{}); +} diff --git a/tests/e2e/results/phase2.txt b/tests/e2e/results/phase2.txt new file mode 100644 index 0000000..737d74b --- /dev/null +++ b/tests/e2e/results/phase2.txt @@ -0,0 +1,38 @@ +# Phase 2 gate — mongo-lite vs MongoDB 8.3.7, 1g dataset / ~16k docs +# Ratio < 1.0 = mongo-lite faster. Reproduce: bash tests/e2e/compare-run.sh 1g 16k +# This run includes roadmap item 1 (B+tree over the encoded keys). +# Compare: tests/e2e/results/phase1.txt (pre-tree baseline). + +benchmark mongo-lite mongodb ratio +insertOne (sequential) ×200 0.19 ms 4.1 ms 0.0x +bulk insert throughput 810.4 MB/s 714.4 MB/s 1.1x +docs loaded 65,536 65,536 1.0x +createIndex({k: 1}) 50.8 ms 82.4 ms 0.6x +countDocuments({}) 1.5 ms 13.8 ms 0.1x +findOne({_id: }) 0.57 ms 0.67 ms 0.9x +findOne({k: 500}) (indexed) 0.57 ms 1.8 ms 0.3x +find({p: {$gte,$lt}}).count() (scan) 20.3 ms 12.9 ms 1.6x +find({}).sort({_id:-1}).limit(20) 6.2 ms 2.7 ms 2.3x +find({}, {proj}).limit(1000) 3.7 ms 4.5 ms 0.8x +aggregate $group by k 11.5 ms 15.5 ms 0.7x +updateOne({_id}) ×50 0.17 ms 0.19 ms 0.9x +updateMany({k: 7}, {$inc}) 1.8 ms 6.7 ms 0.3x +deleteOne({_id}) + insertOne 0.50 ms 5.0 ms 0.1x +node client RSS 152 MB 156 MB 1.0x +server RSS 1974 MB 1474 MB +kill -9 reopen 0.8s 1.3s +db on disk 1028MB 96MB + +# Item 1 (B+tree over encoded keys) deltas vs phase1: +# updateMany 17.3 -> 1.8 ms (2.8x slower than mongod -> 3.7x faster): +# entry removal was a per-entry binary search into a sorted +# array with an orderedRemove memmove behind it; now it is a +# descent plus a leaf-local slot removal. +# createIndex 62.4 -> 50.8 ms (bulk packing replaces append+sort) +# findOne(k:500) indexed 0.64 -> 0.57 ms (unchanged shape, tree search) +# +# Remaining gaps and where they are addressed: +# db on disk 11x -> Phase 3 (block-compressed log) +# sort+limit 2.3x -> Phase 2 (ordered _id index) +# range-scan 1.6x -> Phase 4 (contiguous byte storage, not the matcher) +# server RSS 1.3x -> Phase 4 (per-document arena -> byte storage)