Skip to main content

bplus_index/
lib.rs

1//! Arena-backed B+Tree for in-memory sorted indexes.
2//!
3//! `BPlusIndex<K, V>` is a sorted multimap where entries are `(K, V)` pairs.
4//! Multiple values per key are supported natively -- the `(K, V)` pair is the
5//! unique identity of each entry. This makes it suitable for secondary indexes
6//! in data-oriented designs where one key maps to many row identifiers.
7//!
8//! # Design
9//!
10//! - **Arena storage**: all nodes live in contiguous `Vec`s, indexed by `u32` ids.
11//!   No `Box`, no pointer chasing between siblings during range scans.
12//! - **Cache-friendly leaves**: each leaf holds up to 64 keys and 64 values in
13//!   inline `ArrayVec`s. Binary search within a leaf touches 1-2 cache lines.
14//! - **Linked leaves**: leaves form a doubly-linked list for O(log n + k) range
15//!   scans without climbing back up the tree.
16//! - **Zero `unsafe`**: built on top of `arrayvec` which encapsulates all the
17//!   `MaybeUninit` machinery.
18//! - **Non-unique key support**: inner node separators store `(K, V)` tuples for
19//!   correct routing even when thousands of entries share the same key.
20//!
21//! # Complexity
22//!
23//! | Operation      | Time           | Notes                                   |
24//! |----------------|----------------|-----------------------------------------|
25//! | `insert`       | O(log n)       | Amortized, with rare leaf/inner splits  |
26//! | `delete`       | O(log n)       | With borrow-or-merge rebalancing        |
27//! | `update`       | O(log n)       | Delete + insert                         |
28//! | `lookup_unique`| O(log n)       | Returns first value for a given key     |
29//! | `lookup_range` | O(log n + k)   | All values for a given key              |
30//! | `range`        | O(log n + k)   | Arbitrary key range with bounds         |
31//! | `iter`         | O(n)           | Full scan via leaf chain                |
32//!
33//! # Example
34//!
35//! ```
36//! use bplus_index::BPlusIndex;
37//!
38//! let mut idx = BPlusIndex::<i32, u64>::new();
39//!
40//! // Non-unique: multiple values for the same key
41//! idx.insert(10, 100);
42//! idx.insert(10, 200);
43//! idx.insert(20, 300);
44//!
45//! // Lookup all values for key 10
46//! let vals: Vec<_> = idx.lookup_range(&10).copied().collect();
47//! assert_eq!(vals, vec![100, 200]);
48//!
49//! // Range scan
50//! let entries: Vec<_> = idx.range(5..15).collect();
51//! assert_eq!(entries.len(), 2);
52//! ```
53
54use arrayvec::ArrayVec;
55use std::borrow::Borrow;
56use std::cmp::Ordering;
57use std::ops::{Bound, RangeBounds};
58
59// ---------------------------------------------------------------------------
60// Constants
61// ---------------------------------------------------------------------------
62
63/// Maximum number of keys per leaf node.
64const LEAF_CAP: usize = 64;
65
66/// Maximum number of keys (separators) per inner node.
67const INNER_CAP: usize = 64;
68
69/// Maximum number of children per inner node (INNER_CAP + 1).
70const INNER_CHILDREN: usize = INNER_CAP + 1;
71
72/// Minimum occupancy for a leaf before rebalancing is triggered.
73const MIN_LEAF: usize = LEAF_CAP / 4;
74
75/// Minimum occupancy for an inner node before rebalancing is triggered.
76const MIN_INNER: usize = INNER_CAP / 4;
77
78/// Sentinel value indicating "no node" in prev/next links.
79const NONE: u32 = u32::MAX;
80
81/// Maximum tree depth. Supports up to 64^20 entries (effectively unlimited).
82const MAX_DEPTH: usize = 20;
83
84// ---------------------------------------------------------------------------
85// ArrayVec helpers
86// ---------------------------------------------------------------------------
87
88/// Insert `val` at position `idx`, shifting elements to the right.
89fn av_insert<T, const N: usize>(av: &mut ArrayVec<T, N>, idx: usize, val: T) {
90    av.insert(idx, val);
91}
92
93/// Split `av` at position `at`, returning a new `ArrayVec` with `[at..]`.
94/// The original keeps `[..at]`.
95fn av_split_off<T, const N: usize>(av: &mut ArrayVec<T, N>, at: usize) -> ArrayVec<T, N> {
96    let mut right = ArrayVec::new();
97    for item in av.drain(at..) {
98        right.push(item);
99    }
100    right
101}
102
103/// Move all elements from `src` into `dst`.
104fn av_append<T, const N: usize>(dst: &mut ArrayVec<T, N>, src: &mut ArrayVec<T, N>) {
105    for item in src.drain(..) {
106        dst.push(item);
107    }
108}
109
110// ---------------------------------------------------------------------------
111// Leaf node
112// ---------------------------------------------------------------------------
113
114/// A leaf node in the B+Tree.
115///
116/// Stores up to [`LEAF_CAP`] key-value pairs in sorted order. Leaves are
117/// linked via `prev`/`next` indices for efficient range scans.
118struct LeafNode<K, V> {
119    keys: ArrayVec<K, LEAF_CAP>,
120    vals: ArrayVec<V, LEAF_CAP>,
121    prev: u32,
122    next: u32,
123}
124
125impl<K, V> LeafNode<K, V> {
126    fn new() -> Self {
127        Self {
128            keys: ArrayVec::new(),
129            vals: ArrayVec::new(),
130            prev: NONE,
131            next: NONE,
132        }
133    }
134
135    #[inline]
136    fn len(&self) -> usize {
137        self.keys.len()
138    }
139
140    #[inline]
141    fn is_full(&self) -> bool {
142        self.keys.is_full()
143    }
144
145    /// Returns `true` if this leaf has enough entries to lend one to a sibling.
146    fn can_lend(&self) -> bool {
147        self.keys.len() > MIN_LEAF
148    }
149}
150
151impl<K: Ord, V> LeafNode<K, V> {
152    /// Binary search on the composite `(K, V)` ordering.
153    ///
154    /// Returns `Ok(idx)` on exact match, `Err(idx)` for the insertion point.
155    fn search_kv_pos<Q>(&self, k: &Q, v: &V) -> Result<usize, usize>
156    where
157        K: Borrow<Q> + Ord,
158        Q: Ord + ?Sized,
159        V: Ord,
160    {
161        let mut lo = 0usize;
162        let mut hi = self.len();
163        while lo < hi {
164            let mid = lo + (hi - lo) / 2;
165            match self.keys[mid]
166                .borrow()
167                .cmp(k)
168                .then_with(|| self.vals[mid].cmp(v))
169            {
170                Ordering::Less => lo = mid + 1,
171                Ordering::Greater => hi = mid,
172                Ordering::Equal => return Ok(mid),
173            }
174        }
175        Err(lo)
176    }
177
178    /// First position where `key >= k` (lower bound).
179    fn lower_bound_k<Q>(&self, k: &Q) -> usize
180    where
181        K: Borrow<Q> + Ord,
182        Q: Ord + ?Sized,
183    {
184        self.keys.partition_point(|sk| sk.borrow() < k)
185    }
186
187    /// First position where `key > k` (upper bound).
188    fn upper_bound_k<Q>(&self, k: &Q) -> usize
189    where
190        K: Borrow<Q> + Ord,
191        Q: Ord + ?Sized,
192    {
193        self.keys.partition_point(|sk| sk.borrow() <= k)
194    }
195}
196
197// ---------------------------------------------------------------------------
198// Inner node
199// ---------------------------------------------------------------------------
200
201/// An inner (branch) node in the B+Tree.
202///
203/// Stores separators as `(K, V)` pairs rather than just `K`. This is required
204/// for correct child routing when keys are not unique: if many entries share
205/// the same K, the V component of the separator disambiguates which subtree
206/// contains a given `(K, V)` pair.
207struct InnerNode<K, V> {
208    keys: ArrayVec<K, INNER_CAP>,
209    vals: ArrayVec<V, INNER_CAP>,
210    children: ArrayVec<u32, INNER_CHILDREN>,
211}
212
213impl<K, V> InnerNode<K, V> {
214    fn new() -> Self {
215        Self {
216            keys: ArrayVec::new(),
217            vals: ArrayVec::new(),
218            children: ArrayVec::new(),
219        }
220    }
221
222    fn len(&self) -> usize {
223        self.keys.len()
224    }
225
226    fn is_full(&self) -> bool {
227        self.keys.is_full()
228    }
229
230    /// Returns `true` if this inner node can lend a child to a sibling.
231    fn can_lend(&self) -> bool {
232        self.keys.len() > MIN_INNER
233    }
234}
235
236impl<K: Ord, V: Ord> InnerNode<K, V> {
237    /// Find the child index for a given `(K, V)` using upper-bound semantics.
238    ///
239    /// Returns the index of the first separator strictly greater than `(k, v)`.
240    /// The child at that index is the correct subtree for descent.
241    fn find_child_kv<Q>(&self, k: &Q, v: &V) -> usize
242    where
243        K: Borrow<Q> + Ord,
244        Q: Ord + ?Sized,
245    {
246        let mut lo = 0usize;
247        let mut hi = self.len();
248        while lo < hi {
249            let mid = lo + (hi - lo) / 2;
250            match self.keys[mid]
251                .borrow()
252                .cmp(k)
253                .then_with(|| self.vals[mid].cmp(v))
254            {
255                Ordering::Less | Ordering::Equal => lo = mid + 1,
256                Ordering::Greater => hi = mid,
257            }
258        }
259        lo
260    }
261
262    /// Find the child index using K-only lower-bound semantics.
263    ///
264    /// Used for `lookup_unique`, `lookup_range`, and `range` queries where we
265    /// want to find the leftmost subtree that might contain a given key.
266    fn find_child_k_lower<Q>(&self, k: &Q) -> usize
267    where
268        K: Borrow<Q> + Ord,
269        Q: Ord + ?Sized,
270    {
271        self.keys.partition_point(|sk| sk.borrow() < k)
272    }
273}
274
275// ---------------------------------------------------------------------------
276// Descent path (stack-allocated)
277// ---------------------------------------------------------------------------
278
279#[derive(Clone, Copy)]
280struct PathEntry {
281    node: u32,
282    child_pos: usize,
283}
284
285/// Records the path from root to leaf during descent, so that insert/delete
286/// can propagate splits and merges back up without re-descending.
287struct Path {
288    entries: [PathEntry; MAX_DEPTH],
289    len: usize,
290}
291
292impl Path {
293    fn new() -> Self {
294        Self {
295            entries: [PathEntry {
296                node: 0,
297                child_pos: 0,
298            }; MAX_DEPTH],
299            len: 0,
300        }
301    }
302
303    fn push(&mut self, node: u32, child_pos: usize) {
304        self.entries[self.len] = PathEntry { node, child_pos };
305        self.len += 1;
306    }
307}
308
309// ---------------------------------------------------------------------------
310// BPlusIndex
311// ---------------------------------------------------------------------------
312
313/// An arena-backed B+Tree sorted multimap.
314///
315/// Entries are `(K, V)` pairs sorted lexicographically. The pair is the unique
316/// identity: inserting the same `(K, V)` twice is a no-op, but inserting
317/// `(K, V1)` and `(K, V2)` stores both.
318///
319/// Nodes are stored in contiguous `Vec`s and referenced by `u32` indices.
320/// Freed nodes are recycled via an internal free list.
321pub struct BPlusIndex<K, V> {
322    leaves: Vec<LeafNode<K, V>>,
323    inners: Vec<InnerNode<K, V>>,
324    free_leaves: Vec<u32>,
325    free_inners: Vec<u32>,
326    root: u32,
327    height: u32,
328    len: usize,
329    first_leaf: u32,
330}
331
332impl<K: Ord + Clone, V: Ord + Clone> BPlusIndex<K, V> {
333    /// Creates a new empty index.
334    pub fn new() -> Self {
335        let mut idx = Self {
336            leaves: Vec::new(),
337            inners: Vec::new(),
338            free_leaves: Vec::new(),
339            free_inners: Vec::new(),
340            root: 0,
341            height: 0,
342            len: 0,
343            first_leaf: 0,
344        };
345        idx.alloc_leaf(); // leaf 0 = empty root
346        idx
347    }
348
349    /// Returns the number of entries in the index.
350    pub fn len(&self) -> usize {
351        self.len
352    }
353
354    /// Returns `true` if the index contains no entries.
355    pub fn is_empty(&self) -> bool {
356        self.len == 0
357    }
358
359    // -- Arena helpers ------------------------------------------------------
360
361    fn alloc_leaf(&mut self) -> u32 {
362        if let Some(id) = self.free_leaves.pop() {
363            self.leaves[id as usize] = LeafNode::new();
364            id
365        } else {
366            let id = self.leaves.len() as u32;
367            self.leaves.push(LeafNode::new());
368            id
369        }
370    }
371
372    fn free_leaf(&mut self, id: u32) {
373        self.free_leaves.push(id);
374    }
375
376    fn alloc_inner(&mut self) -> u32 {
377        if let Some(id) = self.free_inners.pop() {
378            self.inners[id as usize] = InnerNode::new();
379            id
380        } else {
381            let id = self.inners.len() as u32;
382            self.inners.push(InnerNode::new());
383            id
384        }
385    }
386
387    fn free_inner(&mut self, id: u32) {
388        self.free_inners.push(id);
389    }
390
391    #[inline]
392    fn leaf(&self, id: u32) -> &LeafNode<K, V> {
393        &self.leaves[id as usize]
394    }
395
396    #[inline]
397    fn leaf_mut(&mut self, id: u32) -> &mut LeafNode<K, V> {
398        &mut self.leaves[id as usize]
399    }
400
401    #[inline]
402    fn inner(&self, id: u32) -> &InnerNode<K, V> {
403        &self.inners[id as usize]
404    }
405
406    #[inline]
407    fn inner_mut(&mut self, id: u32) -> &mut InnerNode<K, V> {
408        &mut self.inners[id as usize]
409    }
410
411    // -- Descent ------------------------------------------------------------
412
413    /// Descend using full `(K, V)` comparison. Used by insert and delete to
414    /// find the exact leaf and record the path for split/merge propagation.
415    fn descend_kv<Q>(&self, k: &Q, v: &V) -> (u32, Path)
416    where
417        K: Borrow<Q> + Ord,
418        Q: Ord + ?Sized,
419    {
420        let mut path = Path::new();
421        if self.height == 0 {
422            return (self.root, path);
423        }
424        let mut node = self.root;
425        for level in (1..=self.height).rev() {
426            let inner = self.inner(node);
427            let child_pos = inner.find_child_kv(k, v);
428            path.push(node, child_pos);
429            node = inner.children[child_pos];
430            if level == 1 {
431                return (node, path);
432            }
433        }
434        unreachable!()
435    }
436
437    /// Descend using K-only lower-bound. Used by lookups and range queries to
438    /// find the leftmost leaf that could contain the key.
439    fn descend_k<Q>(&self, k: &Q) -> u32
440    where
441        K: Borrow<Q> + Ord,
442        Q: Ord + ?Sized,
443    {
444        if self.height == 0 {
445            return self.root;
446        }
447        let mut node = self.root;
448        for level in (1..=self.height).rev() {
449            let inner = self.inner(node);
450            let child_pos = inner.find_child_k_lower(k);
451            node = inner.children[child_pos];
452            if level == 1 {
453                return node;
454            }
455        }
456        unreachable!()
457    }
458
459    // -- Insert -------------------------------------------------------------
460
461    /// Inserts a `(key, val)` pair into the index.
462    ///
463    /// If the exact `(key, val)` pair already exists, this is a no-op.
464    /// Multiple values for the same key are allowed.
465    pub fn insert(&mut self, key: K, val: V) {
466        let (leaf_id, path) = self.descend_kv(&key, &val);
467        let leaf = self.leaf(leaf_id);
468
469        let pos = match leaf.search_kv_pos(&key, &val) {
470            Ok(_) => return, // duplicate (K, V)
471            Err(p) => p,
472        };
473
474        if !self.leaf(leaf_id).is_full() {
475            let leaf = self.leaf_mut(leaf_id);
476            av_insert(&mut leaf.keys, pos, key);
477            av_insert(&mut leaf.vals, pos, val);
478            self.len += 1;
479            return;
480        }
481
482        // Leaf is full: split then propagate upward.
483        let (right_id, sep_k, sep_v) = self.split_leaf(leaf_id, pos, key, val);
484        self.propagate_split(path, sep_k, sep_v, right_id);
485        self.len += 1;
486    }
487
488    /// Split a full leaf at the midpoint, inserting `(key, val)` at `pos`.
489    /// Returns the new right leaf id and the separator to promote.
490    fn split_leaf(&mut self, leaf_id: u32, pos: usize, key: K, val: V) -> (u32, K, V) {
491        let split = LEAF_CAP / 2;
492
493        let mut right_keys = av_split_off(&mut self.leaf_mut(leaf_id).keys, split);
494        let mut right_vals = av_split_off(&mut self.leaf_mut(leaf_id).vals, split);
495
496        if pos <= split {
497            let leaf = self.leaf_mut(leaf_id);
498            av_insert(&mut leaf.keys, pos, key);
499            av_insert(&mut leaf.vals, pos, val);
500        } else {
501            av_insert(&mut right_keys, pos - split, key);
502            av_insert(&mut right_vals, pos - split, val);
503        }
504
505        let sep_k = right_keys[0].clone();
506        let sep_v = right_vals[0].clone();
507
508        let right_id = self.alloc_leaf();
509        let old_next = self.leaf(leaf_id).next;
510
511        let right = self.leaf_mut(right_id);
512        right.keys = right_keys;
513        right.vals = right_vals;
514        right.prev = leaf_id;
515        right.next = old_next;
516
517        self.leaf_mut(leaf_id).next = right_id;
518
519        if old_next != NONE {
520            self.leaf_mut(old_next).prev = right_id;
521        }
522
523        (right_id, sep_k, sep_v)
524    }
525
526    /// Propagate a split upward through inner nodes. If the root overflows,
527    /// a new root is created (increasing tree height by 1).
528    fn propagate_split(&mut self, path: Path, mut sep_k: K, mut sep_v: V, mut right_child: u32) {
529        for i in (0..path.len).rev() {
530            let inner_id = path.entries[i].node;
531            let pos = self.inner(inner_id).find_child_kv(&sep_k, &sep_v);
532
533            if !self.inner(inner_id).is_full() {
534                let inner = self.inner_mut(inner_id);
535                av_insert(&mut inner.keys, pos, sep_k);
536                av_insert(&mut inner.vals, pos, sep_v);
537                av_insert(&mut inner.children, pos + 1, right_child);
538                return;
539            }
540
541            let (new_right, med_k, med_v) =
542                self.split_inner(inner_id, pos, sep_k, sep_v, right_child);
543            sep_k = med_k;
544            sep_v = med_v;
545            right_child = new_right;
546        }
547
548        self.grow_root(sep_k, sep_v, right_child);
549    }
550
551    /// Split a full inner node, returning the new right node and the median
552    /// separator to promote to the parent.
553    fn split_inner(
554        &mut self,
555        inner_id: u32,
556        pos: usize,
557        key: K,
558        val: V,
559        child: u32,
560    ) -> (u32, K, V) {
561        let split = INNER_CAP / 2;
562
563        let mut right_keys = av_split_off(&mut self.inner_mut(inner_id).keys, split + 1);
564        let mut right_vals = av_split_off(&mut self.inner_mut(inner_id).vals, split + 1);
565        let mut right_children = av_split_off(&mut self.inner_mut(inner_id).children, split + 1);
566
567        let med_k = self.inner_mut(inner_id).keys.remove(split);
568        let med_v = self.inner_mut(inner_id).vals.remove(split);
569
570        if pos <= split {
571            let inner = self.inner_mut(inner_id);
572            av_insert(&mut inner.keys, pos, key);
573            av_insert(&mut inner.vals, pos, val);
574            av_insert(&mut inner.children, pos + 1, child);
575        } else {
576            let rp = pos - split - 1;
577            av_insert(&mut right_keys, rp, key);
578            av_insert(&mut right_vals, rp, val);
579            av_insert(&mut right_children, rp + 1, child);
580        }
581
582        let right_id = self.alloc_inner();
583        let right = self.inner_mut(right_id);
584        right.keys = right_keys;
585        right.vals = right_vals;
586        right.children = right_children;
587
588        (right_id, med_k, med_v)
589    }
590
591    /// Create a new root containing one separator and two children.
592    fn grow_root(&mut self, sep_k: K, sep_v: V, right_child: u32) {
593        let old_root = self.root;
594        let new_root = self.alloc_inner();
595        let root = self.inner_mut(new_root);
596        root.keys.push(sep_k);
597        root.vals.push(sep_v);
598        root.children.push(old_root);
599        root.children.push(right_child);
600        self.root = new_root;
601        self.height += 1;
602    }
603
604    // -- Delete -------------------------------------------------------------
605
606    /// Removes the exact `(key, val)` pair from the index.
607    ///
608    /// Returns `true` if the entry was found and removed, `false` otherwise.
609    pub fn delete<Q>(&mut self, key: &Q, val: &V) -> bool
610    where
611        K: Borrow<Q> + Ord,
612        Q: Ord + ?Sized,
613    {
614        let (leaf_id, path) = self.descend_kv(key, val);
615
616        let pos = match self.leaf(leaf_id).search_kv_pos(key, val) {
617            Ok(p) => p,
618            Err(_) => return false,
619        };
620
621        self.leaf_mut(leaf_id).keys.remove(pos);
622        self.leaf_mut(leaf_id).vals.remove(pos);
623        self.len -= 1;
624
625        // Root leaf never needs rebalancing.
626        if self.height == 0 {
627            return true;
628        }
629
630        if self.leaf(leaf_id).len() >= MIN_LEAF {
631            return true;
632        }
633
634        self.rebalance_leaf(leaf_id, &path);
635        true
636    }
637
638    /// Rebalance a leaf that has fallen below minimum occupancy.
639    ///
640    /// Tries in order: borrow from right sibling, borrow from left sibling,
641    /// merge with a sibling. Merges may cascade up through inner nodes.
642    fn rebalance_leaf(&mut self, leaf_id: u32, path: &Path) {
643        let parent_entry = path.entries[path.len - 1];
644        let parent_id = parent_entry.node;
645        let child_pos = parent_entry.child_pos;
646
647        // Try borrow from right sibling.
648        if child_pos + 1 < self.inner(parent_id).children.len() {
649            let right_id = self.inner(parent_id).children[child_pos + 1];
650            if self.leaf(right_id).can_lend() {
651                let rk = self.leaf_mut(right_id).keys.remove(0);
652                let rv = self.leaf_mut(right_id).vals.remove(0);
653                self.leaf_mut(leaf_id).keys.push(rk);
654                self.leaf_mut(leaf_id).vals.push(rv);
655
656                let new_sep_k = self.leaf(right_id).keys[0].clone();
657                let new_sep_v = self.leaf(right_id).vals[0].clone();
658                self.inner_mut(parent_id).keys[child_pos] = new_sep_k;
659                self.inner_mut(parent_id).vals[child_pos] = new_sep_v;
660                return;
661            }
662        }
663
664        // Try borrow from left sibling.
665        if child_pos > 0 {
666            let left_id = self.inner(parent_id).children[child_pos - 1];
667            if self.leaf(left_id).can_lend() {
668                let lk = self.leaf_mut(left_id).keys.pop().unwrap();
669                let lv = self.leaf_mut(left_id).vals.pop().unwrap();
670                av_insert(&mut self.leaf_mut(leaf_id).keys, 0, lk);
671                av_insert(&mut self.leaf_mut(leaf_id).vals, 0, lv);
672
673                let new_sep_k = self.leaf(leaf_id).keys[0].clone();
674                let new_sep_v = self.leaf(leaf_id).vals[0].clone();
675                self.inner_mut(parent_id).keys[child_pos - 1] = new_sep_k;
676                self.inner_mut(parent_id).vals[child_pos - 1] = new_sep_v;
677                return;
678            }
679        }
680
681        // Merge with a sibling.
682        if child_pos + 1 < self.inner(parent_id).children.len() {
683            let right_id = self.inner(parent_id).children[child_pos + 1];
684            self.merge_leaves(leaf_id, right_id);
685            self.inner_mut(parent_id).keys.remove(child_pos);
686            self.inner_mut(parent_id).vals.remove(child_pos);
687            self.inner_mut(parent_id).children.remove(child_pos + 1);
688            self.free_leaf(right_id);
689        } else {
690            let left_id = self.inner(parent_id).children[child_pos - 1];
691            self.merge_leaves(left_id, leaf_id);
692            self.inner_mut(parent_id).keys.remove(child_pos - 1);
693            self.inner_mut(parent_id).vals.remove(child_pos - 1);
694            self.inner_mut(parent_id).children.remove(child_pos);
695            self.free_leaf(leaf_id);
696        }
697
698        self.maybe_shrink_or_rebalance_inner(parent_id, path, path.len - 1);
699    }
700
701    /// Merge the right leaf into the left leaf and fix the linked list.
702    fn merge_leaves(&mut self, left_id: u32, right_id: u32) {
703        let right = self.leaf_mut(right_id);
704        let mut rk = std::mem::replace(&mut right.keys, ArrayVec::new());
705        let mut rv = std::mem::replace(&mut right.vals, ArrayVec::new());
706        let right_next = right.next;
707
708        let left = self.leaf_mut(left_id);
709        av_append(&mut left.keys, &mut rk);
710        av_append(&mut left.vals, &mut rv);
711        left.next = right_next;
712
713        if right_next != NONE {
714            self.leaf_mut(right_next).prev = left_id;
715        }
716    }
717
718    /// After a merge, check if the parent needs shrinking (root with one child)
719    /// or rebalancing (inner node below minimum occupancy).
720    fn maybe_shrink_or_rebalance_inner(&mut self, inner_id: u32, path: &Path, path_idx: usize) {
721        // Root with zero separators -> shrink the tree by one level.
722        if inner_id == self.root {
723            if self.inner(inner_id).keys.is_empty() {
724                let only_child = self.inner(inner_id).children[0];
725                self.free_inner(inner_id);
726                self.root = only_child;
727                self.height -= 1;
728                if self.height == 0 {
729                    self.first_leaf = self.root;
730                }
731            }
732            return;
733        }
734
735        if self.inner(inner_id).len() >= MIN_INNER {
736            return;
737        }
738
739        self.rebalance_inner(inner_id, path, path_idx);
740    }
741
742    /// Rebalance an inner node: borrow from a sibling or merge.
743    fn rebalance_inner(&mut self, inner_id: u32, path: &Path, path_idx: usize) {
744        if path_idx == 0 {
745            return;
746        }
747
748        let parent_entry = path.entries[path_idx - 1];
749        let parent_id = parent_entry.node;
750        let child_pos = parent_entry.child_pos;
751
752        // Try borrow from right sibling.
753        if child_pos + 1 < self.inner(parent_id).children.len() {
754            let right_id = self.inner(parent_id).children[child_pos + 1];
755            if self.inner(right_id).can_lend() {
756                let sep_k = self.inner_mut(parent_id).keys.remove(child_pos);
757                let sep_v = self.inner_mut(parent_id).vals.remove(child_pos);
758
759                let rk = self.inner_mut(right_id).keys.remove(0);
760                let rv = self.inner_mut(right_id).vals.remove(0);
761                let rc = self.inner_mut(right_id).children.remove(0);
762
763                av_insert(&mut self.inner_mut(parent_id).keys, child_pos, rk);
764                av_insert(&mut self.inner_mut(parent_id).vals, child_pos, rv);
765
766                self.inner_mut(inner_id).keys.push(sep_k);
767                self.inner_mut(inner_id).vals.push(sep_v);
768                self.inner_mut(inner_id).children.push(rc);
769                return;
770            }
771        }
772
773        // Try borrow from left sibling.
774        if child_pos > 0 {
775            let left_id = self.inner(parent_id).children[child_pos - 1];
776            if self.inner(left_id).can_lend() {
777                let sep_k = self.inner_mut(parent_id).keys.remove(child_pos - 1);
778                let sep_v = self.inner_mut(parent_id).vals.remove(child_pos - 1);
779
780                let lk = self.inner_mut(left_id).keys.pop().unwrap();
781                let lv = self.inner_mut(left_id).vals.pop().unwrap();
782                let lc = self.inner_mut(left_id).children.pop().unwrap();
783
784                av_insert(&mut self.inner_mut(parent_id).keys, child_pos - 1, lk);
785                av_insert(&mut self.inner_mut(parent_id).vals, child_pos - 1, lv);
786
787                av_insert(&mut self.inner_mut(inner_id).keys, 0, sep_k);
788                av_insert(&mut self.inner_mut(inner_id).vals, 0, sep_v);
789                av_insert(&mut self.inner_mut(inner_id).children, 0, lc);
790                return;
791            }
792        }
793
794        // Merge with right sibling.
795        if child_pos + 1 < self.inner(parent_id).children.len() {
796            let right_id = self.inner(parent_id).children[child_pos + 1];
797
798            let sep_k = self.inner_mut(parent_id).keys.remove(child_pos);
799            let sep_v = self.inner_mut(parent_id).vals.remove(child_pos);
800            self.inner_mut(parent_id).children.remove(child_pos + 1);
801
802            self.inner_mut(inner_id).keys.push(sep_k);
803            self.inner_mut(inner_id).vals.push(sep_v);
804
805            let right = self.inner_mut(right_id);
806            let mut rk = std::mem::replace(&mut right.keys, ArrayVec::new());
807            let mut rv = std::mem::replace(&mut right.vals, ArrayVec::new());
808            let mut rc = std::mem::replace(&mut right.children, ArrayVec::new());
809
810            av_append(&mut self.inner_mut(inner_id).keys, &mut rk);
811            av_append(&mut self.inner_mut(inner_id).vals, &mut rv);
812            av_append(&mut self.inner_mut(inner_id).children, &mut rc);
813
814            self.free_inner(right_id);
815        } else {
816            // Merge into left sibling.
817            let left_id = self.inner(parent_id).children[child_pos - 1];
818
819            let sep_k = self.inner_mut(parent_id).keys.remove(child_pos - 1);
820            let sep_v = self.inner_mut(parent_id).vals.remove(child_pos - 1);
821            self.inner_mut(parent_id).children.remove(child_pos);
822
823            self.inner_mut(left_id).keys.push(sep_k);
824            self.inner_mut(left_id).vals.push(sep_v);
825
826            let cur = self.inner_mut(inner_id);
827            let mut ck = std::mem::replace(&mut cur.keys, ArrayVec::new());
828            let mut cv = std::mem::replace(&mut cur.vals, ArrayVec::new());
829            let mut cc = std::mem::replace(&mut cur.children, ArrayVec::new());
830
831            av_append(&mut self.inner_mut(left_id).keys, &mut ck);
832            av_append(&mut self.inner_mut(left_id).vals, &mut cv);
833            av_append(&mut self.inner_mut(left_id).children, &mut cc);
834
835            self.free_inner(inner_id);
836        }
837
838        self.maybe_shrink_or_rebalance_inner(parent_id, path, path_idx - 1);
839    }
840
841    // -- Update -------------------------------------------------------------
842
843    /// Re-indexes an entry by deleting `(old_key, val)` and inserting `(new_key, val)`.
844    ///
845    /// This is a convenience for the common case where the indexed field of a
846    /// row changes but the row identifier (the value) stays the same.
847    pub fn update<Q>(&mut self, old_key: &Q, new_key: K, val: V)
848    where
849        K: Borrow<Q> + Ord,
850        Q: Ord + ?Sized,
851    {
852        self.delete(old_key, &val);
853        self.insert(new_key, val);
854    }
855
856    // -- Lookups ------------------------------------------------------------
857
858    /// Returns a reference to the first value associated with `key`, or `None`.
859    ///
860    /// For unique indexes this is the only value. For non-unique indexes this
861    /// returns the smallest `V` for the given `K`.
862    ///
863    /// Time: O(log n).
864    pub fn lookup_unique<Q>(&self, key: &Q) -> Option<&V>
865    where
866        K: Borrow<Q> + Ord,
867        Q: Ord + ?Sized,
868    {
869        let leaf_id = self.descend_k(key);
870        let leaf = self.leaf(leaf_id);
871        let pos = leaf.lower_bound_k(key);
872        if pos < leaf.len() && leaf.keys[pos].borrow() == key {
873            return Some(&leaf.vals[pos]);
874        }
875        // Key might be at the start of the next leaf (boundary case).
876        if leaf.next != NONE {
877            let next = self.leaf(leaf.next);
878            if !next.keys.is_empty() && next.keys[0].borrow() == key {
879                return Some(&next.vals[0]);
880            }
881        }
882        None
883    }
884
885    /// Returns an iterator over all values associated with `key`.
886    ///
887    /// Values are yielded in ascending `V` order. The iterator follows the
888    /// leaf chain, so it works correctly even when entries for one key span
889    /// multiple leaves.
890    ///
891    /// Time: O(log n) for the initial descent, then O(k) for k results.
892    pub fn lookup_range<'a, 'b, Q>(&'a self, key: &'b Q) -> LookupRange<'a, 'b, Q, K, V>
893    where
894        K: Borrow<Q> + Ord,
895        Q: Ord + ?Sized,
896    {
897        let leaf_id = self.descend_k(key);
898        let leaf = self.leaf(leaf_id);
899        let pos = leaf.lower_bound_k(key);
900        LookupRange {
901            tree: self,
902            key,
903            leaf: leaf_id,
904            pos,
905        }
906    }
907
908    /// Returns an iterator over all entries `(&K, &V)` whose keys fall within
909    /// the given range bounds.
910    ///
911    /// Supports all standard range syntax: `a..b`, `a..=b`, `a..`, `..b`, `..`.
912    ///
913    /// Time: O(log n) for the initial descent, then O(k) for k results.
914    pub fn range<R: RangeBounds<K>>(&self, bounds: R) -> RangeIter<'_, K, V> {
915        let (leaf, pos) = match bounds.start_bound() {
916            Bound::Unbounded => (self.first_leaf, 0),
917            Bound::Included(k) => {
918                let lid = self.descend_k(k);
919                let l = self.leaf(lid);
920                (lid, l.lower_bound_k(k))
921            }
922            Bound::Excluded(k) => {
923                let lid = self.descend_k(k);
924                let l = self.leaf(lid);
925                (lid, l.upper_bound_k(k))
926            }
927        };
928
929        let end = match bounds.end_bound() {
930            Bound::Unbounded => RangeEnd::Unbounded,
931            Bound::Included(k) => RangeEnd::Included(k.clone()),
932            Bound::Excluded(k) => RangeEnd::Excluded(k.clone()),
933        };
934
935        RangeIter {
936            tree: self,
937            leaf,
938            pos,
939            end,
940        }
941    }
942
943    /// Returns an iterator over all entries in ascending order.
944    ///
945    /// Time: O(n).
946    pub fn iter(&self) -> RangeIter<'_, K, V> {
947        RangeIter {
948            tree: self,
949            leaf: self.first_leaf,
950            pos: 0,
951            end: RangeEnd::Unbounded,
952        }
953    }
954}
955
956impl<K: Ord + Clone, V: Ord + Clone> Default for BPlusIndex<K, V> {
957    fn default() -> Self {
958        Self::new()
959    }
960}
961
962// ---------------------------------------------------------------------------
963// Iterators
964// ---------------------------------------------------------------------------
965
966/// Iterator returned by [`BPlusIndex::lookup_range`].
967///
968/// Yields `&V` for all entries matching a specific key.
969pub struct LookupRange<'a, 'b, Q, K, V>
970where
971    K: Borrow<Q> + Ord,
972    Q: Ord + ?Sized,
973{
974    tree: &'a BPlusIndex<K, V>,
975    key: &'b Q,
976    leaf: u32,
977    pos: usize,
978}
979
980impl<'a, 'b, Q, K, V> Iterator for LookupRange<'a, 'b, Q, K, V>
981where
982    K: Borrow<Q> + Ord + Clone,
983    Q: Ord + ?Sized,
984    V: Ord + Clone,
985{
986    type Item = &'a V;
987
988    fn next(&mut self) -> Option<Self::Item> {
989        loop {
990            if self.leaf == NONE {
991                return None;
992            }
993            let leaf = self.tree.leaf(self.leaf);
994            if self.pos < leaf.len() {
995                if leaf.keys[self.pos].borrow() == self.key {
996                    let v = &leaf.vals[self.pos];
997                    self.pos += 1;
998                    return Some(v);
999                }
1000                return None;
1001            }
1002            self.leaf = leaf.next;
1003            self.pos = 0;
1004        }
1005    }
1006}
1007
1008enum RangeEnd<K> {
1009    Unbounded,
1010    Included(K),
1011    Excluded(K),
1012}
1013
1014/// Iterator returned by [`BPlusIndex::range`] and [`BPlusIndex::iter`].
1015///
1016/// Yields `(&K, &V)` pairs in ascending order. Follows the leaf chain
1017/// for sequential memory access.
1018pub struct RangeIter<'a, K, V> {
1019    tree: &'a BPlusIndex<K, V>,
1020    leaf: u32,
1021    pos: usize,
1022    end: RangeEnd<K>,
1023}
1024
1025impl<'a, K: Ord + Clone, V: Ord + Clone> Iterator for RangeIter<'a, K, V> {
1026    type Item = (&'a K, &'a V);
1027
1028    fn next(&mut self) -> Option<Self::Item> {
1029        loop {
1030            if self.leaf == NONE {
1031                return None;
1032            }
1033            let leaf = self.tree.leaf(self.leaf);
1034            if self.pos < leaf.len() {
1035                let k = &leaf.keys[self.pos];
1036                match &self.end {
1037                    RangeEnd::Excluded(end) if k >= end => return None,
1038                    RangeEnd::Included(end) if k > end => return None,
1039                    _ => {}
1040                }
1041                let v = &leaf.vals[self.pos];
1042                self.pos += 1;
1043                return Some((k, v));
1044            }
1045            self.leaf = leaf.next;
1046            self.pos = 0;
1047        }
1048    }
1049}
1050
1051#[cfg(test)]
1052mod tests {
1053    use super::*;
1054    use std::collections::HashSet;
1055    use std::fmt::Debug;
1056
1057    // -----------------------------------------------------------------------
1058    // Structural invariant validator
1059    // -----------------------------------------------------------------------
1060
1061    impl<K: Ord + Clone + Debug, V: Ord + Clone + Debug> BPlusIndex<K, V> {
1062        /// Validate all B+Tree structural invariants. Panics with a descriptive
1063        /// message on any violation.
1064        fn validate(&self) {
1065            self.validate_len();
1066            if self.height == 0 {
1067                self.validate_root_leaf();
1068            } else {
1069                self.validate_inner_recursive(self.root, self.height, None, None);
1070            }
1071            self.validate_leaf_chain();
1072            self.validate_first_leaf();
1073        }
1074
1075        fn validate_len(&self) {
1076            let actual: usize = self.iter().count();
1077            assert_eq!(
1078                actual, self.len,
1079                "len mismatch: stored {} but iter yields {}",
1080                self.len, actual
1081            );
1082        }
1083
1084        fn validate_root_leaf(&self) {
1085            let leaf = self.leaf(self.root);
1086            assert!(
1087                leaf.keys.len() == leaf.vals.len(),
1088                "root leaf: keys.len != vals.len"
1089            );
1090            // Root leaf has no minimum occupancy constraint.
1091            self.validate_leaf_sorted(self.root);
1092        }
1093
1094        fn validate_leaf_sorted(&self, id: u32) {
1095            let leaf = self.leaf(id);
1096            for i in 1..leaf.len() {
1097                let cmp = leaf.keys[i - 1]
1098                    .cmp(&leaf.keys[i])
1099                    .then_with(|| leaf.vals[i - 1].cmp(&leaf.vals[i]));
1100                assert!(
1101                    cmp == Ordering::Less,
1102                    "leaf {}: entries not strictly sorted at position {}: ({:?},{:?}) vs ({:?},{:?})",
1103                    id,
1104                    i,
1105                    leaf.keys[i - 1],
1106                    leaf.vals[i - 1],
1107                    leaf.keys[i],
1108                    leaf.vals[i]
1109                );
1110            }
1111        }
1112
1113        fn validate_inner_recursive(
1114            &self,
1115            node_id: u32,
1116            level: u32,
1117            lower: Option<(&K, &V)>,
1118            upper: Option<(&K, &V)>,
1119        ) {
1120            let inner = self.inner(node_id);
1121
1122            // children.len == keys.len + 1
1123            assert_eq!(
1124                inner.children.len(),
1125                inner.keys.len() + 1,
1126                "inner {}: children.len ({}) != keys.len ({}) + 1",
1127                node_id,
1128                inner.children.len(),
1129                inner.keys.len()
1130            );
1131
1132            // keys/vals sorted
1133            for i in 1..inner.len() {
1134                let cmp = inner.keys[i - 1]
1135                    .cmp(&inner.keys[i])
1136                    .then_with(|| inner.vals[i - 1].cmp(&inner.vals[i]));
1137                assert!(
1138                    cmp == Ordering::Less,
1139                    "inner {}: separators not sorted at position {}",
1140                    node_id,
1141                    i
1142                );
1143            }
1144
1145            // Minimum occupancy (except root)
1146            if node_id != self.root {
1147                assert!(
1148                    inner.len() >= MIN_INNER,
1149                    "inner {}: underflow, len {} < MIN_INNER {}",
1150                    node_id,
1151                    inner.len(),
1152                    MIN_INNER
1153                );
1154            }
1155
1156            // Recurse into children
1157            for ci in 0..inner.children.len() {
1158                let child_lower = if ci == 0 {
1159                    lower
1160                } else {
1161                    Some((&inner.keys[ci - 1], &inner.vals[ci - 1]))
1162                };
1163                let child_upper = if ci == inner.keys.len() {
1164                    upper
1165                } else {
1166                    Some((&inner.keys[ci], &inner.vals[ci]))
1167                };
1168
1169                if level == 1 {
1170                    // Child is a leaf
1171                    let leaf_id = inner.children[ci];
1172                    let leaf = self.leaf(leaf_id);
1173
1174                    // Leaf occupancy (non-root leaves)
1175                    if self.len > 0 {
1176                        // Allow underflow only if tree is very small
1177                        let total_leaves = self.iter().count();
1178                        if total_leaves > LEAF_CAP {
1179                            assert!(
1180                                leaf.len() >= MIN_LEAF,
1181                                "leaf {}: underflow, len {} < MIN_LEAF {}",
1182                                leaf_id,
1183                                leaf.len(),
1184                                MIN_LEAF
1185                            );
1186                        }
1187                    }
1188
1189                    self.validate_leaf_sorted(leaf_id);
1190
1191                    // All entries in leaf must be within [child_lower, child_upper)
1192                    for i in 0..leaf.len() {
1193                        if let Some((lk, lv)) = child_lower {
1194                            let cmp = leaf.keys[i].cmp(lk).then_with(|| leaf.vals[i].cmp(lv));
1195                            assert!(
1196                                cmp != Ordering::Less,
1197                                "leaf {}: entry ({:?},{:?}) below lower bound ({:?},{:?})",
1198                                leaf_id,
1199                                leaf.keys[i],
1200                                leaf.vals[i],
1201                                lk,
1202                                lv
1203                            );
1204                        }
1205                        if let Some((uk, uv)) = child_upper {
1206                            let cmp = leaf.keys[i].cmp(uk).then_with(|| leaf.vals[i].cmp(uv));
1207                            assert!(
1208                                cmp == Ordering::Less,
1209                                "leaf {}: entry ({:?},{:?}) not below upper bound ({:?},{:?})",
1210                                leaf_id,
1211                                leaf.keys[i],
1212                                leaf.vals[i],
1213                                uk,
1214                                uv
1215                            );
1216                        }
1217                    }
1218                } else {
1219                    // Child is an inner node
1220                    self.validate_inner_recursive(
1221                        inner.children[ci],
1222                        level - 1,
1223                        child_lower,
1224                        child_upper,
1225                    );
1226                }
1227            }
1228        }
1229
1230        fn validate_leaf_chain(&self) {
1231            // Walk the forward chain and check prev/next consistency.
1232            let mut visited = HashSet::new();
1233            let mut leaf_id = self.first_leaf;
1234            let mut prev_id = NONE;
1235            let mut total_entries = 0usize;
1236
1237            while leaf_id != NONE {
1238                assert!(
1239                    visited.insert(leaf_id),
1240                    "leaf chain: cycle detected at leaf {}",
1241                    leaf_id
1242                );
1243
1244                let leaf = self.leaf(leaf_id);
1245                assert_eq!(
1246                    leaf.prev, prev_id,
1247                    "leaf {}: prev is {} but expected {}",
1248                    leaf_id, leaf.prev, prev_id
1249                );
1250
1251                // Cross-leaf ordering: last of prev < first of current
1252                if prev_id != NONE && leaf.len() > 0 {
1253                    let prev_leaf = self.leaf(prev_id);
1254                    if prev_leaf.len() > 0 {
1255                        let prev_last_k = &prev_leaf.keys[prev_leaf.len() - 1];
1256                        let prev_last_v = &prev_leaf.vals[prev_leaf.len() - 1];
1257                        let cur_first_k = &leaf.keys[0];
1258                        let cur_first_v = &leaf.vals[0];
1259                        let cmp = prev_last_k
1260                            .cmp(cur_first_k)
1261                            .then_with(|| prev_last_v.cmp(cur_first_v));
1262                        assert!(
1263                            cmp == Ordering::Less,
1264                            "leaf chain: last entry of leaf {} ({:?},{:?}) >= first entry of leaf {} ({:?},{:?})",
1265                            prev_id,
1266                            prev_last_k,
1267                            prev_last_v,
1268                            leaf_id,
1269                            cur_first_k,
1270                            cur_first_v
1271                        );
1272                    }
1273                }
1274
1275                total_entries += leaf.len();
1276                prev_id = leaf_id;
1277                leaf_id = leaf.next;
1278            }
1279
1280            assert_eq!(
1281                total_entries, self.len,
1282                "leaf chain: total entries {} != stored len {}",
1283                total_entries, self.len
1284            );
1285        }
1286
1287        fn validate_first_leaf(&self) {
1288            if self.height == 0 {
1289                assert_eq!(
1290                    self.first_leaf, self.root,
1291                    "height 0 but first_leaf != root"
1292                );
1293            } else {
1294                // first_leaf should be reachable by always going to child 0
1295                let mut node = self.root;
1296                for level in (1..=self.height).rev() {
1297                    let inner = self.inner(node);
1298                    node = inner.children[0];
1299                    if level == 1 {
1300                        assert_eq!(
1301                            self.first_leaf, node,
1302                            "first_leaf {} != leftmost leaf {}",
1303                            self.first_leaf, node
1304                        );
1305                        return;
1306                    }
1307                }
1308            }
1309        }
1310    }
1311
1312    // -----------------------------------------------------------------------
1313    // Helper: validate after every operation
1314    // -----------------------------------------------------------------------
1315
1316    fn insert_and_validate(idx: &mut BPlusIndex<i32, u64>, k: i32, v: u64) {
1317        idx.insert(k, v);
1318        idx.validate();
1319    }
1320
1321    fn delete_and_validate(idx: &mut BPlusIndex<i32, u64>, k: &i32, v: &u64) -> bool {
1322        let result = idx.delete(k, v);
1323        idx.validate();
1324        result
1325    }
1326
1327    // -----------------------------------------------------------------------
1328    // Empty tree
1329    // -----------------------------------------------------------------------
1330
1331    #[test]
1332    fn empty_tree() {
1333        let idx = BPlusIndex::<i32, u64>::new();
1334        assert_eq!(idx.len(), 0);
1335        assert!(idx.is_empty());
1336        assert_eq!(idx.lookup_unique(&0), None);
1337        assert_eq!(idx.lookup_range(&0).count(), 0);
1338        assert_eq!(idx.range(..).count(), 0);
1339        assert_eq!(idx.range(0..10).count(), 0);
1340        assert_eq!(idx.iter().count(), 0);
1341        idx.validate();
1342    }
1343
1344    #[test]
1345    fn empty_tree_delete() {
1346        let mut idx = BPlusIndex::<i32, u64>::new();
1347        assert!(!idx.delete(&1, &1));
1348        idx.validate();
1349    }
1350
1351    // -----------------------------------------------------------------------
1352    // Single element
1353    // -----------------------------------------------------------------------
1354
1355    #[test]
1356    fn single_insert_delete() {
1357        let mut idx = BPlusIndex::<i32, u64>::new();
1358        insert_and_validate(&mut idx, 42, 1);
1359        assert_eq!(idx.len(), 1);
1360        assert_eq!(idx.lookup_unique(&42), Some(&1));
1361        assert!(delete_and_validate(&mut idx, &42, &1));
1362        assert_eq!(idx.len(), 0);
1363        assert_eq!(idx.lookup_unique(&42), None);
1364    }
1365
1366    // -----------------------------------------------------------------------
1367    // Duplicates
1368    // -----------------------------------------------------------------------
1369
1370    #[test]
1371    fn insert_duplicate_kv_is_noop() {
1372        let mut idx = BPlusIndex::<i32, u64>::new();
1373        insert_and_validate(&mut idx, 5, 100);
1374        insert_and_validate(&mut idx, 5, 100);
1375        assert_eq!(idx.len(), 1);
1376    }
1377
1378    // -----------------------------------------------------------------------
1379    // Non-unique keys
1380    // -----------------------------------------------------------------------
1381
1382    #[test]
1383    fn non_unique_keys_small() {
1384        let mut idx = BPlusIndex::<i32, u64>::new();
1385        insert_and_validate(&mut idx, 5, 100);
1386        insert_and_validate(&mut idx, 5, 200);
1387        insert_and_validate(&mut idx, 5, 300);
1388        insert_and_validate(&mut idx, 10, 400);
1389
1390        let vals: Vec<_> = idx.lookup_range(&5).copied().collect();
1391        assert_eq!(vals, vec![100, 200, 300]);
1392        assert_eq!(idx.lookup_unique(&5), Some(&100)); // smallest V
1393    }
1394
1395    #[test]
1396    fn non_unique_keys_spanning_leaves() {
1397        // 200 entries with the same key forces them across multiple leaves
1398        let mut idx = BPlusIndex::<i32, u64>::new();
1399        for v in 0..200u64 {
1400            insert_and_validate(&mut idx, 42, v);
1401        }
1402        assert_eq!(idx.len(), 200);
1403
1404        let vals: Vec<_> = idx.lookup_range(&42).copied().collect();
1405        assert_eq!(vals, (0..200u64).collect::<Vec<_>>());
1406
1407        // Delete every other one
1408        for v in (0..200u64).step_by(2) {
1409            assert!(delete_and_validate(&mut idx, &42, &v));
1410        }
1411        assert_eq!(idx.len(), 100);
1412
1413        let vals: Vec<_> = idx.lookup_range(&42).copied().collect();
1414        let expected: Vec<u64> = (0..200u64).filter(|v| v % 2 != 0).collect();
1415        assert_eq!(vals, expected);
1416    }
1417
1418    // -----------------------------------------------------------------------
1419    // Leaf split boundary (LEAF_CAP = 64)
1420    // -----------------------------------------------------------------------
1421
1422    #[test]
1423    fn insert_exactly_leaf_cap() {
1424        let mut idx = BPlusIndex::<i32, u64>::new();
1425        for i in 0..LEAF_CAP as i32 {
1426            insert_and_validate(&mut idx, i, i as u64);
1427        }
1428        assert_eq!(idx.len(), LEAF_CAP);
1429        assert_eq!(idx.height, 0); // still fits in root leaf
1430    }
1431
1432    #[test]
1433    fn insert_triggers_first_split() {
1434        let mut idx = BPlusIndex::<i32, u64>::new();
1435        for i in 0..=LEAF_CAP as i32 {
1436            insert_and_validate(&mut idx, i, i as u64);
1437        }
1438        assert_eq!(idx.len(), LEAF_CAP + 1);
1439        assert_eq!(idx.height, 1); // split created an inner root
1440    }
1441
1442    #[test]
1443    fn insert_triggers_inner_split() {
1444        // LEAF_CAP=64, INNER_CAP=64
1445        // Fill enough to overflow inner nodes: ~64*64 = 4096+ entries
1446        let mut idx = BPlusIndex::<i32, u64>::new();
1447        let n = LEAF_CAP * INNER_CAP + 100;
1448        for i in 0..n as i32 {
1449            idx.insert(i, i as u64);
1450        }
1451        idx.validate();
1452        assert!(idx.height >= 2, "expected height >= 2, got {}", idx.height);
1453    }
1454
1455    // -----------------------------------------------------------------------
1456    // Delete boundary cases
1457    // -----------------------------------------------------------------------
1458
1459    #[test]
1460    fn delete_nonexistent_key() {
1461        let mut idx = BPlusIndex::<i32, u64>::new();
1462        insert_and_validate(&mut idx, 1, 10);
1463        assert!(!idx.delete(&999, &10));
1464        assert!(!idx.delete(&1, &999));
1465        assert_eq!(idx.len(), 1);
1466        idx.validate();
1467    }
1468
1469    #[test]
1470    fn delete_triggers_borrow_from_right() {
1471        let mut idx = BPlusIndex::<i32, u64>::new();
1472        // Build a tree with 2 leaves
1473        for i in 0..LEAF_CAP as i32 + 10 {
1474            idx.insert(i, i as u64);
1475        }
1476        idx.validate();
1477        // Delete from left leaf until it borrows from right
1478        for i in 0..LEAF_CAP as i32 / 2 {
1479            delete_and_validate(&mut idx, &i, &(i as u64));
1480        }
1481    }
1482
1483    #[test]
1484    fn delete_triggers_borrow_from_left() {
1485        let mut idx = BPlusIndex::<i32, u64>::new();
1486        for i in 0..LEAF_CAP as i32 + 10 {
1487            idx.insert(i, i as u64);
1488        }
1489        idx.validate();
1490        // Delete from right side
1491        for i in (LEAF_CAP as i32..LEAF_CAP as i32 + 10).rev() {
1492            delete_and_validate(&mut idx, &i, &(i as u64));
1493        }
1494    }
1495
1496    #[test]
1497    fn delete_triggers_leaf_merge() {
1498        let mut idx = BPlusIndex::<i32, u64>::new();
1499        // Build a 2-leaf tree
1500        for i in 0..LEAF_CAP as i32 + 2 {
1501            idx.insert(i, i as u64);
1502        }
1503        idx.validate();
1504        assert_eq!(idx.height, 1);
1505        // Delete enough to force underflow + merge
1506        for i in (0..LEAF_CAP as i32 + 2).rev() {
1507            delete_and_validate(&mut idx, &i, &(i as u64));
1508            if idx.len() <= MIN_LEAF {
1509                break;
1510            }
1511        }
1512        // Once only MIN_LEAF entries remain, tree should have shrunk to height 0
1513        assert_eq!(idx.height, 0, "tree should have shrunk back to single leaf");
1514    }
1515
1516    #[test]
1517    fn delete_all_forward() {
1518        let mut idx = BPlusIndex::<i32, u64>::new();
1519        let n = 500;
1520        for i in 0..n {
1521            idx.insert(i, i as u64);
1522        }
1523        for i in 0..n {
1524            assert!(delete_and_validate(&mut idx, &i, &(i as u64)));
1525        }
1526        assert_eq!(idx.len(), 0);
1527    }
1528
1529    #[test]
1530    fn delete_all_reverse() {
1531        let mut idx = BPlusIndex::<i32, u64>::new();
1532        let n = 500;
1533        for i in 0..n {
1534            idx.insert(i, i as u64);
1535        }
1536        for i in (0..n).rev() {
1537            assert!(delete_and_validate(&mut idx, &i, &(i as u64)));
1538        }
1539        assert_eq!(idx.len(), 0);
1540    }
1541
1542    #[test]
1543    fn delete_all_from_middle_out() {
1544        let mut idx = BPlusIndex::<i32, u64>::new();
1545        let n = 500i32;
1546        for i in 0..n {
1547            idx.insert(i, i as u64);
1548        }
1549        // Delete from the center outward
1550        let mid = n / 2;
1551        for d in 0..mid {
1552            delete_and_validate(&mut idx, &(mid + d), &((mid + d) as u64));
1553            if mid - d > 0 {
1554                delete_and_validate(&mut idx, &(mid - d - 1), &((mid - d - 1) as u64));
1555            }
1556        }
1557        assert_eq!(idx.len(), 0);
1558    }
1559
1560    #[test]
1561    fn delete_cascade_inner_rebalance() {
1562        // Build a large enough tree for inner node rebalancing
1563        let mut idx = BPlusIndex::<i32, u64>::new();
1564        let n = 5000;
1565        for i in 0..n {
1566            idx.insert(i, i as u64);
1567        }
1568        idx.validate();
1569        let initial_height = idx.height;
1570        assert!(initial_height >= 2);
1571
1572        // Delete ~90% of entries
1573        for i in 0..(n * 9 / 10) {
1574            assert!(delete_and_validate(&mut idx, &i, &(i as u64)));
1575        }
1576        // Height should have decreased
1577        assert!(idx.height < initial_height);
1578    }
1579
1580    // -----------------------------------------------------------------------
1581    // Lookup
1582    // -----------------------------------------------------------------------
1583
1584    #[test]
1585    fn lookup_unique_all() {
1586        let mut idx = BPlusIndex::<i32, u64>::new();
1587        for i in 0..200 {
1588            idx.insert(i, i as u64 * 10);
1589        }
1590        for i in 0..200 {
1591            assert_eq!(idx.lookup_unique(&i), Some(&(i as u64 * 10)));
1592        }
1593        assert_eq!(idx.lookup_unique(&999), None);
1594        assert_eq!(idx.lookup_unique(&-1), None);
1595    }
1596
1597    #[test]
1598    fn lookup_unique_boundary_between_leaves() {
1599        // Key that falls exactly at a leaf split point
1600        let mut idx = BPlusIndex::<i32, u64>::new();
1601        for i in 0..200 {
1602            idx.insert(i, i as u64);
1603        }
1604        // Verify every key, including those at leaf boundaries
1605        for i in 0..200 {
1606            assert_eq!(
1607                idx.lookup_unique(&i),
1608                Some(&(i as u64)),
1609                "lookup_unique failed for key {}",
1610                i
1611            );
1612        }
1613    }
1614
1615    #[test]
1616    fn lookup_range_empty_result() {
1617        let mut idx = BPlusIndex::<i32, u64>::new();
1618        idx.insert(5, 100);
1619        idx.insert(10, 200);
1620        assert_eq!(idx.lookup_range(&7).count(), 0);
1621    }
1622
1623    #[test]
1624    fn lookup_range_key_not_in_tree() {
1625        let idx = BPlusIndex::<i32, u64>::new();
1626        assert_eq!(idx.lookup_range(&42).count(), 0);
1627    }
1628
1629    // -----------------------------------------------------------------------
1630    // Range queries
1631    // -----------------------------------------------------------------------
1632
1633    #[test]
1634    fn range_exclusive_both_ends() {
1635        let mut idx = BPlusIndex::<i32, u64>::new();
1636        for i in 0..100 {
1637            idx.insert(i, i as u64);
1638        }
1639        let vals: Vec<_> = idx.range(10..20).map(|(k, _)| *k).collect();
1640        assert_eq!(vals, (10..20).collect::<Vec<_>>());
1641    }
1642
1643    #[test]
1644    fn range_inclusive_both_ends() {
1645        let mut idx = BPlusIndex::<i32, u64>::new();
1646        for i in 0..100 {
1647            idx.insert(i, i as u64);
1648        }
1649        let vals: Vec<_> = idx.range(10..=20).map(|(k, _)| *k).collect();
1650        assert_eq!(vals, (10..=20).collect::<Vec<_>>());
1651    }
1652
1653    #[test]
1654    fn range_unbounded_start() {
1655        let mut idx = BPlusIndex::<i32, u64>::new();
1656        for i in 0..50 {
1657            idx.insert(i, i as u64);
1658        }
1659        let vals: Vec<_> = idx.range(..10).map(|(k, _)| *k).collect();
1660        assert_eq!(vals, (0..10).collect::<Vec<_>>());
1661    }
1662
1663    #[test]
1664    fn range_unbounded_end() {
1665        let mut idx = BPlusIndex::<i32, u64>::new();
1666        for i in 0..50 {
1667            idx.insert(i, i as u64);
1668        }
1669        let vals: Vec<_> = idx.range(40..).map(|(k, _)| *k).collect();
1670        assert_eq!(vals, (40..50).collect::<Vec<_>>());
1671    }
1672
1673    #[test]
1674    fn range_fully_unbounded() {
1675        let mut idx = BPlusIndex::<i32, u64>::new();
1676        for i in 0..50 {
1677            idx.insert(i, i as u64);
1678        }
1679        let vals: Vec<_> = idx.range(..).map(|(k, _)| *k).collect();
1680        assert_eq!(vals, (0..50).collect::<Vec<_>>());
1681    }
1682
1683    #[test]
1684    fn range_empty_result() {
1685        let mut idx = BPlusIndex::<i32, u64>::new();
1686        for i in 0..50 {
1687            idx.insert(i, i as u64);
1688        }
1689        assert_eq!(idx.range(100..200).count(), 0);
1690        assert_eq!(idx.range(-10..-5).count(), 0);
1691    }
1692
1693    #[test]
1694    fn range_on_empty_tree() {
1695        let idx = BPlusIndex::<i32, u64>::new();
1696        assert_eq!(idx.range(0..100).count(), 0);
1697        assert_eq!(idx.range(..).count(), 0);
1698    }
1699
1700    #[test]
1701    fn range_spanning_multiple_leaves() {
1702        let mut idx = BPlusIndex::<i32, u64>::new();
1703        for i in 0..500 {
1704            idx.insert(i, i as u64);
1705        }
1706        let vals: Vec<_> = idx.range(100..400).map(|(k, _)| *k).collect();
1707        assert_eq!(vals, (100..400).collect::<Vec<_>>());
1708    }
1709
1710    #[test]
1711    fn range_with_non_unique_keys() {
1712        let mut idx = BPlusIndex::<i32, u64>::new();
1713        for k in 0..10 {
1714            for v in 0..5u64 {
1715                idx.insert(k, v);
1716            }
1717        }
1718        // Range 3..7 should return all V for keys 3, 4, 5, 6
1719        let entries: Vec<_> = idx.range(3..7).collect();
1720        assert_eq!(entries.len(), 4 * 5);
1721        for (k, _) in &entries {
1722            assert!(**k >= 3 && **k < 7);
1723        }
1724    }
1725
1726    // -----------------------------------------------------------------------
1727    // Iteration
1728    // -----------------------------------------------------------------------
1729
1730    #[test]
1731    fn iter_empty() {
1732        let idx = BPlusIndex::<i32, u64>::new();
1733        assert_eq!(idx.iter().count(), 0);
1734    }
1735
1736    #[test]
1737    fn iter_sorted_after_reverse_insert() {
1738        let mut idx = BPlusIndex::<i32, u64>::new();
1739        for i in (0..300).rev() {
1740            idx.insert(i, i as u64);
1741        }
1742        idx.validate();
1743        let keys: Vec<_> = idx.iter().map(|(k, _)| *k).collect();
1744        assert_eq!(keys, (0..300).collect::<Vec<_>>());
1745    }
1746
1747    #[test]
1748    fn iter_sorted_after_random_order_insert() {
1749        let mut idx = BPlusIndex::<i32, u64>::new();
1750        // Insert in a pseudo-random order
1751        let mut keys: Vec<i32> = (0..500).collect();
1752        // Simple deterministic shuffle
1753        for i in (1..keys.len()).rev() {
1754            let j = (i * 2654435761) % (i + 1);
1755            keys.swap(i, j);
1756        }
1757        for k in &keys {
1758            idx.insert(*k, *k as u64);
1759        }
1760        idx.validate();
1761        let result: Vec<_> = idx.iter().map(|(k, _)| *k).collect();
1762        assert_eq!(result, (0..500).collect::<Vec<_>>());
1763    }
1764
1765    // -----------------------------------------------------------------------
1766    // Update
1767    // -----------------------------------------------------------------------
1768
1769    #[test]
1770    fn update_moves_entry() {
1771        let mut idx = BPlusIndex::<i32, u64>::new();
1772        idx.insert(5, 42);
1773        idx.update(&5, 10, 42);
1774        assert_eq!(idx.lookup_unique(&5), None);
1775        assert_eq!(idx.lookup_unique(&10), Some(&42));
1776        idx.validate();
1777    }
1778
1779    #[test]
1780    fn update_nonexistent_key_still_inserts() {
1781        let mut idx = BPlusIndex::<i32, u64>::new();
1782        idx.insert(5, 42);
1783        idx.update(&999, 10, 42); // delete of 999 is a no-op
1784        // 42 is now under both key 5 and key 10 (update doesn't remove from 5)
1785        // Actually: update deletes (999, 42) which doesn't exist, then inserts (10, 42)
1786        assert_eq!(idx.lookup_unique(&5), Some(&42));
1787        assert_eq!(idx.lookup_unique(&10), Some(&42));
1788        assert_eq!(idx.len(), 2);
1789        idx.validate();
1790    }
1791
1792    // -----------------------------------------------------------------------
1793    // Free list / node recycling
1794    // -----------------------------------------------------------------------
1795
1796    #[test]
1797    fn node_recycling_after_delete_insert() {
1798        let mut idx = BPlusIndex::<i32, u64>::new();
1799        // Build up
1800        for i in 0..500 {
1801            idx.insert(i, i as u64);
1802        }
1803        let leaves_after_insert = idx.leaves.len();
1804
1805        // Tear down
1806        for i in 0..500 {
1807            idx.delete(&i, &(i as u64));
1808        }
1809        idx.validate();
1810        let free_after_delete = idx.free_leaves.len();
1811        assert!(free_after_delete > 0, "expected freed leaves");
1812
1813        // Build up again -- should reuse freed nodes
1814        for i in 0..500 {
1815            idx.insert(i, i as u64);
1816        }
1817        idx.validate();
1818        assert_eq!(
1819            idx.leaves.len(),
1820            leaves_after_insert,
1821            "arena grew instead of recycling"
1822        );
1823    }
1824
1825    // -----------------------------------------------------------------------
1826    // Stress / combined operations
1827    // -----------------------------------------------------------------------
1828
1829    #[test]
1830    fn stress_insert_delete_interleaved() {
1831        let mut idx = BPlusIndex::<i32, u64>::new();
1832        for i in 0..1000 {
1833            idx.insert(i, i as u64);
1834        }
1835        idx.validate();
1836
1837        // Delete even
1838        for i in (0..1000).step_by(2) {
1839            assert!(idx.delete(&i, &(i as u64)));
1840        }
1841        idx.validate();
1842        assert_eq!(idx.len(), 500);
1843
1844        // Verify odd remain
1845        for i in (1..1000).step_by(2) {
1846            assert_eq!(idx.lookup_unique(&i), Some(&(i as u64)));
1847        }
1848
1849        // Re-insert even
1850        for i in (0..1000).step_by(2) {
1851            idx.insert(i, i as u64);
1852        }
1853        idx.validate();
1854        assert_eq!(idx.len(), 1000);
1855
1856        let keys: Vec<_> = idx.iter().map(|(k, _)| *k).collect();
1857        assert_eq!(keys, (0..1000).collect::<Vec<_>>());
1858    }
1859
1860    #[test]
1861    fn stress_insert_delete_reinsert_cycles() {
1862        let mut idx = BPlusIndex::<i32, u64>::new();
1863        for cycle in 0..5 {
1864            let base = cycle * 100;
1865            for i in base..base + 100 {
1866                insert_and_validate(&mut idx, i, i as u64);
1867            }
1868            // Delete first half of this cycle
1869            for i in base..base + 50 {
1870                assert!(delete_and_validate(&mut idx, &i, &(i as u64)));
1871            }
1872        }
1873        // Should have 50 entries from each of 5 cycles = 250
1874        assert_eq!(idx.len(), 250);
1875    }
1876
1877    #[test]
1878    fn stress_large_tree() {
1879        let mut idx = BPlusIndex::<i32, u64>::new();
1880        let n = 10_000i32;
1881        for i in 0..n {
1882            idx.insert(i, i as u64);
1883        }
1884        idx.validate();
1885        assert_eq!(idx.len(), n as usize);
1886
1887        // Random-ish deletions
1888        for i in (0..n).step_by(3) {
1889            assert!(idx.delete(&i, &(i as u64)));
1890        }
1891        idx.validate();
1892
1893        // Verify remaining
1894        for i in 0..n {
1895            if i % 3 == 0 {
1896                assert_eq!(idx.lookup_unique(&i), None);
1897            } else {
1898                assert_eq!(idx.lookup_unique(&i), Some(&(i as u64)));
1899            }
1900        }
1901    }
1902
1903    // -----------------------------------------------------------------------
1904    // Composite key (simulating real use case)
1905    // -----------------------------------------------------------------------
1906
1907    #[test]
1908    fn composite_key_range() {
1909        // (game_mode, mmr) -> ticket_id
1910        let mut idx = BPlusIndex::<(u8, i32), u64>::new();
1911        // Mode 1: various MMRs
1912        for i in 0..100 {
1913            idx.insert((1, 1000 + i), i as u64);
1914        }
1915        // Mode 2: various MMRs
1916        for i in 0..50 {
1917            idx.insert((2, 1000 + i), 100 + i as u64);
1918        }
1919        idx.validate();
1920
1921        // Range: mode 1, mmr 1020..1030
1922        let tickets: Vec<_> = idx.range((1, 1020)..(1, 1030)).map(|(_, v)| *v).collect();
1923        assert_eq!(tickets, (20..30u64).collect::<Vec<_>>());
1924
1925        // Range: mode 2 only
1926        let mode2: Vec<_> = idx.range((2, 0)..(3, 0)).map(|(_, v)| *v).collect();
1927        assert_eq!(mode2.len(), 50);
1928    }
1929
1930    // -----------------------------------------------------------------------
1931    // Edge: all same key
1932    // -----------------------------------------------------------------------
1933
1934    #[test]
1935    fn all_same_key() {
1936        let mut idx = BPlusIndex::<i32, u64>::new();
1937        for v in 0..500u64 {
1938            idx.insert(0, v);
1939        }
1940        idx.validate();
1941        assert_eq!(idx.len(), 500);
1942
1943        let vals: Vec<_> = idx.lookup_range(&0).copied().collect();
1944        assert_eq!(vals, (0..500u64).collect::<Vec<_>>());
1945
1946        // Delete all
1947        for v in 0..500u64 {
1948            assert!(delete_and_validate(&mut idx, &0, &v));
1949        }
1950        assert_eq!(idx.len(), 0);
1951    }
1952}