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::cmp::Ordering;
56use std::ops::{Bound, RangeBounds};
57
58// ---------------------------------------------------------------------------
59// Constants
60// ---------------------------------------------------------------------------
61
62/// Maximum number of keys per leaf node.
63const LEAF_CAP: usize = 64;
64
65/// Maximum number of keys (separators) per inner node.
66const INNER_CAP: usize = 64;
67
68/// Maximum number of children per inner node (INNER_CAP + 1).
69const INNER_CHILDREN: usize = INNER_CAP + 1;
70
71/// Minimum occupancy for a leaf before rebalancing is triggered.
72const MIN_LEAF: usize = LEAF_CAP / 4;
73
74/// Minimum occupancy for an inner node before rebalancing is triggered.
75const MIN_INNER: usize = INNER_CAP / 4;
76
77/// Sentinel value indicating "no node" in prev/next links.
78const NONE: u32 = u32::MAX;
79
80/// Maximum tree depth. Supports up to 64^20 entries (effectively unlimited).
81const MAX_DEPTH: usize = 20;
82
83// ---------------------------------------------------------------------------
84// ArrayVec helpers
85// ---------------------------------------------------------------------------
86
87/// Insert `val` at position `idx`, shifting elements to the right.
88fn av_insert<T, const N: usize>(av: &mut ArrayVec<T, N>, idx: usize, val: T) {
89    av.insert(idx, val);
90}
91
92/// Split `av` at position `at`, returning a new `ArrayVec` with `[at..]`.
93/// The original keeps `[..at]`.
94fn av_split_off<T, const N: usize>(av: &mut ArrayVec<T, N>, at: usize) -> ArrayVec<T, N> {
95    let mut right = ArrayVec::new();
96    for item in av.drain(at..) {
97        right.push(item);
98    }
99    right
100}
101
102/// Move all elements from `src` into `dst`.
103fn av_append<T, const N: usize>(dst: &mut ArrayVec<T, N>, src: &mut ArrayVec<T, N>) {
104    for item in src.drain(..) {
105        dst.push(item);
106    }
107}
108
109// ---------------------------------------------------------------------------
110// Leaf node
111// ---------------------------------------------------------------------------
112
113/// A leaf node in the B+Tree.
114///
115/// Stores up to [`LEAF_CAP`] key-value pairs in sorted order. Leaves are
116/// linked via `prev`/`next` indices for efficient range scans.
117struct LeafNode<K, V> {
118    keys: ArrayVec<K, LEAF_CAP>,
119    vals: ArrayVec<V, LEAF_CAP>,
120    prev: u32,
121    next: u32,
122}
123
124impl<K, V> LeafNode<K, V> {
125    fn new() -> Self {
126        Self {
127            keys: ArrayVec::new(),
128            vals: ArrayVec::new(),
129            prev: NONE,
130            next: NONE,
131        }
132    }
133
134    #[inline]
135    fn len(&self) -> usize {
136        self.keys.len()
137    }
138
139    #[inline]
140    fn is_full(&self) -> bool {
141        self.keys.is_full()
142    }
143
144    /// Returns `true` if this leaf has enough entries to lend one to a sibling.
145    fn can_lend(&self) -> bool {
146        self.keys.len() > MIN_LEAF
147    }
148}
149
150impl<K: Ord, V> LeafNode<K, V> {
151    /// Binary search on the composite `(K, V)` ordering.
152    ///
153    /// Returns `Ok(idx)` on exact match, `Err(idx)` for the insertion point.
154    fn search_kv_pos(&self, k: &K, v: &V) -> Result<usize, usize>
155    where
156        V: Ord,
157    {
158        let mut lo = 0usize;
159        let mut hi = self.len();
160        while lo < hi {
161            let mid = lo + (hi - lo) / 2;
162            match self.keys[mid].cmp(k).then_with(|| self.vals[mid].cmp(v)) {
163                Ordering::Less => lo = mid + 1,
164                Ordering::Greater => hi = mid,
165                Ordering::Equal => return Ok(mid),
166            }
167        }
168        Err(lo)
169    }
170
171    /// First position where `key >= k` (lower bound).
172    fn lower_bound_k(&self, k: &K) -> usize {
173        self.keys.partition_point(|sk| sk < k)
174    }
175
176    /// First position where `key > k` (upper bound).
177    fn upper_bound_k(&self, k: &K) -> usize {
178        self.keys.partition_point(|sk| sk <= k)
179    }
180}
181
182// ---------------------------------------------------------------------------
183// Inner node
184// ---------------------------------------------------------------------------
185
186/// An inner (branch) node in the B+Tree.
187///
188/// Stores separators as `(K, V)` pairs rather than just `K`. This is required
189/// for correct child routing when keys are not unique: if many entries share
190/// the same K, the V component of the separator disambiguates which subtree
191/// contains a given `(K, V)` pair.
192struct InnerNode<K, V> {
193    keys: ArrayVec<K, INNER_CAP>,
194    vals: ArrayVec<V, INNER_CAP>,
195    children: ArrayVec<u32, INNER_CHILDREN>,
196}
197
198impl<K, V> InnerNode<K, V> {
199    fn new() -> Self {
200        Self {
201            keys: ArrayVec::new(),
202            vals: ArrayVec::new(),
203            children: ArrayVec::new(),
204        }
205    }
206
207    fn len(&self) -> usize {
208        self.keys.len()
209    }
210
211    fn is_full(&self) -> bool {
212        self.keys.is_full()
213    }
214
215    /// Returns `true` if this inner node can lend a child to a sibling.
216    fn can_lend(&self) -> bool {
217        self.keys.len() > MIN_INNER
218    }
219}
220
221impl<K: Ord, V: Ord> InnerNode<K, V> {
222    /// Find the child index for a given `(K, V)` using upper-bound semantics.
223    ///
224    /// Returns the index of the first separator strictly greater than `(k, v)`.
225    /// The child at that index is the correct subtree for descent.
226    fn find_child_kv(&self, k: &K, v: &V) -> usize {
227        let mut lo = 0usize;
228        let mut hi = self.len();
229        while lo < hi {
230            let mid = lo + (hi - lo) / 2;
231            match self.keys[mid].cmp(k).then_with(|| self.vals[mid].cmp(v)) {
232                Ordering::Less | Ordering::Equal => lo = mid + 1,
233                Ordering::Greater => hi = mid,
234            }
235        }
236        lo
237    }
238
239    /// Find the child index using K-only lower-bound semantics.
240    ///
241    /// Used for `lookup_unique`, `lookup_range`, and `range` queries where we
242    /// want to find the leftmost subtree that might contain a given key.
243    fn find_child_k_lower(&self, k: &K) -> usize {
244        self.keys.partition_point(|sk| sk < k)
245    }
246}
247
248// ---------------------------------------------------------------------------
249// Descent path (stack-allocated)
250// ---------------------------------------------------------------------------
251
252#[derive(Clone, Copy)]
253struct PathEntry {
254    node: u32,
255    child_pos: usize,
256}
257
258/// Records the path from root to leaf during descent, so that insert/delete
259/// can propagate splits and merges back up without re-descending.
260struct Path {
261    entries: [PathEntry; MAX_DEPTH],
262    len: usize,
263}
264
265impl Path {
266    fn new() -> Self {
267        Self {
268            entries: [PathEntry {
269                node: 0,
270                child_pos: 0,
271            }; MAX_DEPTH],
272            len: 0,
273        }
274    }
275
276    fn push(&mut self, node: u32, child_pos: usize) {
277        self.entries[self.len] = PathEntry { node, child_pos };
278        self.len += 1;
279    }
280}
281
282// ---------------------------------------------------------------------------
283// BPlusIndex
284// ---------------------------------------------------------------------------
285
286/// An arena-backed B+Tree sorted multimap.
287///
288/// Entries are `(K, V)` pairs sorted lexicographically. The pair is the unique
289/// identity: inserting the same `(K, V)` twice is a no-op, but inserting
290/// `(K, V1)` and `(K, V2)` stores both.
291///
292/// Nodes are stored in contiguous `Vec`s and referenced by `u32` indices.
293/// Freed nodes are recycled via an internal free list.
294pub struct BPlusIndex<K, V> {
295    leaves: Vec<LeafNode<K, V>>,
296    inners: Vec<InnerNode<K, V>>,
297    free_leaves: Vec<u32>,
298    free_inners: Vec<u32>,
299    root: u32,
300    height: u32,
301    len: usize,
302    first_leaf: u32,
303}
304
305impl<K: Ord + Clone, V: Ord + Clone> BPlusIndex<K, V> {
306    /// Creates a new empty index.
307    pub fn new() -> Self {
308        let mut idx = Self {
309            leaves: Vec::new(),
310            inners: Vec::new(),
311            free_leaves: Vec::new(),
312            free_inners: Vec::new(),
313            root: 0,
314            height: 0,
315            len: 0,
316            first_leaf: 0,
317        };
318        idx.alloc_leaf(); // leaf 0 = empty root
319        idx
320    }
321
322    /// Returns the number of entries in the index.
323    pub fn len(&self) -> usize {
324        self.len
325    }
326
327    /// Returns `true` if the index contains no entries.
328    pub fn is_empty(&self) -> bool {
329        self.len == 0
330    }
331
332    // -- Arena helpers ------------------------------------------------------
333
334    fn alloc_leaf(&mut self) -> u32 {
335        if let Some(id) = self.free_leaves.pop() {
336            self.leaves[id as usize] = LeafNode::new();
337            id
338        } else {
339            let id = self.leaves.len() as u32;
340            self.leaves.push(LeafNode::new());
341            id
342        }
343    }
344
345    fn free_leaf(&mut self, id: u32) {
346        self.free_leaves.push(id);
347    }
348
349    fn alloc_inner(&mut self) -> u32 {
350        if let Some(id) = self.free_inners.pop() {
351            self.inners[id as usize] = InnerNode::new();
352            id
353        } else {
354            let id = self.inners.len() as u32;
355            self.inners.push(InnerNode::new());
356            id
357        }
358    }
359
360    fn free_inner(&mut self, id: u32) {
361        self.free_inners.push(id);
362    }
363
364    #[inline]
365    fn leaf(&self, id: u32) -> &LeafNode<K, V> {
366        &self.leaves[id as usize]
367    }
368
369    #[inline]
370    fn leaf_mut(&mut self, id: u32) -> &mut LeafNode<K, V> {
371        &mut self.leaves[id as usize]
372    }
373
374    #[inline]
375    fn inner(&self, id: u32) -> &InnerNode<K, V> {
376        &self.inners[id as usize]
377    }
378
379    #[inline]
380    fn inner_mut(&mut self, id: u32) -> &mut InnerNode<K, V> {
381        &mut self.inners[id as usize]
382    }
383
384    // -- Descent ------------------------------------------------------------
385
386    /// Descend using full `(K, V)` comparison. Used by insert and delete to
387    /// find the exact leaf and record the path for split/merge propagation.
388    fn descend_kv(&self, k: &K, v: &V) -> (u32, Path) {
389        let mut path = Path::new();
390        if self.height == 0 {
391            return (self.root, path);
392        }
393        let mut node = self.root;
394        for level in (1..=self.height).rev() {
395            let inner = self.inner(node);
396            let child_pos = inner.find_child_kv(k, v);
397            path.push(node, child_pos);
398            node = inner.children[child_pos];
399            if level == 1 {
400                return (node, path);
401            }
402        }
403        unreachable!()
404    }
405
406    /// Descend using K-only lower-bound. Used by lookups and range queries to
407    /// find the leftmost leaf that could contain the key.
408    fn descend_k(&self, k: &K) -> u32 {
409        if self.height == 0 {
410            return self.root;
411        }
412        let mut node = self.root;
413        for level in (1..=self.height).rev() {
414            let inner = self.inner(node);
415            let child_pos = inner.find_child_k_lower(k);
416            node = inner.children[child_pos];
417            if level == 1 {
418                return node;
419            }
420        }
421        unreachable!()
422    }
423
424    // -- Insert -------------------------------------------------------------
425
426    /// Inserts a `(key, val)` pair into the index.
427    ///
428    /// If the exact `(key, val)` pair already exists, this is a no-op.
429    /// Multiple values for the same key are allowed.
430    pub fn insert(&mut self, key: K, val: V) {
431        let (leaf_id, path) = self.descend_kv(&key, &val);
432        let leaf = self.leaf(leaf_id);
433
434        let pos = match leaf.search_kv_pos(&key, &val) {
435            Ok(_) => return, // duplicate (K, V)
436            Err(p) => p,
437        };
438
439        if !self.leaf(leaf_id).is_full() {
440            let leaf = self.leaf_mut(leaf_id);
441            av_insert(&mut leaf.keys, pos, key);
442            av_insert(&mut leaf.vals, pos, val);
443            self.len += 1;
444            return;
445        }
446
447        // Leaf is full: split then propagate upward.
448        let (right_id, sep_k, sep_v) = self.split_leaf(leaf_id, pos, key, val);
449        self.propagate_split(path, sep_k, sep_v, right_id);
450        self.len += 1;
451    }
452
453    /// Split a full leaf at the midpoint, inserting `(key, val)` at `pos`.
454    /// Returns the new right leaf id and the separator to promote.
455    fn split_leaf(&mut self, leaf_id: u32, pos: usize, key: K, val: V) -> (u32, K, V) {
456        let split = LEAF_CAP / 2;
457
458        let mut right_keys = av_split_off(&mut self.leaf_mut(leaf_id).keys, split);
459        let mut right_vals = av_split_off(&mut self.leaf_mut(leaf_id).vals, split);
460
461        if pos <= split {
462            let leaf = self.leaf_mut(leaf_id);
463            av_insert(&mut leaf.keys, pos, key);
464            av_insert(&mut leaf.vals, pos, val);
465        } else {
466            av_insert(&mut right_keys, pos - split, key);
467            av_insert(&mut right_vals, pos - split, val);
468        }
469
470        let sep_k = right_keys[0].clone();
471        let sep_v = right_vals[0].clone();
472
473        let right_id = self.alloc_leaf();
474        let old_next = self.leaf(leaf_id).next;
475
476        let right = self.leaf_mut(right_id);
477        right.keys = right_keys;
478        right.vals = right_vals;
479        right.prev = leaf_id;
480        right.next = old_next;
481
482        self.leaf_mut(leaf_id).next = right_id;
483
484        if old_next != NONE {
485            self.leaf_mut(old_next).prev = right_id;
486        }
487
488        (right_id, sep_k, sep_v)
489    }
490
491    /// Propagate a split upward through inner nodes. If the root overflows,
492    /// a new root is created (increasing tree height by 1).
493    fn propagate_split(&mut self, path: Path, mut sep_k: K, mut sep_v: V, mut right_child: u32) {
494        for i in (0..path.len).rev() {
495            let inner_id = path.entries[i].node;
496            let pos = self.inner(inner_id).find_child_kv(&sep_k, &sep_v);
497
498            if !self.inner(inner_id).is_full() {
499                let inner = self.inner_mut(inner_id);
500                av_insert(&mut inner.keys, pos, sep_k);
501                av_insert(&mut inner.vals, pos, sep_v);
502                av_insert(&mut inner.children, pos + 1, right_child);
503                return;
504            }
505
506            let (new_right, med_k, med_v) =
507                self.split_inner(inner_id, pos, sep_k, sep_v, right_child);
508            sep_k = med_k;
509            sep_v = med_v;
510            right_child = new_right;
511        }
512
513        self.grow_root(sep_k, sep_v, right_child);
514    }
515
516    /// Split a full inner node, returning the new right node and the median
517    /// separator to promote to the parent.
518    fn split_inner(
519        &mut self,
520        inner_id: u32,
521        pos: usize,
522        key: K,
523        val: V,
524        child: u32,
525    ) -> (u32, K, V) {
526        let split = INNER_CAP / 2;
527
528        let mut right_keys = av_split_off(&mut self.inner_mut(inner_id).keys, split + 1);
529        let mut right_vals = av_split_off(&mut self.inner_mut(inner_id).vals, split + 1);
530        let mut right_children = av_split_off(&mut self.inner_mut(inner_id).children, split + 1);
531
532        let med_k = self.inner_mut(inner_id).keys.remove(split);
533        let med_v = self.inner_mut(inner_id).vals.remove(split);
534
535        if pos <= split {
536            let inner = self.inner_mut(inner_id);
537            av_insert(&mut inner.keys, pos, key);
538            av_insert(&mut inner.vals, pos, val);
539            av_insert(&mut inner.children, pos + 1, child);
540        } else {
541            let rp = pos - split - 1;
542            av_insert(&mut right_keys, rp, key);
543            av_insert(&mut right_vals, rp, val);
544            av_insert(&mut right_children, rp + 1, child);
545        }
546
547        let right_id = self.alloc_inner();
548        let right = self.inner_mut(right_id);
549        right.keys = right_keys;
550        right.vals = right_vals;
551        right.children = right_children;
552
553        (right_id, med_k, med_v)
554    }
555
556    /// Create a new root containing one separator and two children.
557    fn grow_root(&mut self, sep_k: K, sep_v: V, right_child: u32) {
558        let old_root = self.root;
559        let new_root = self.alloc_inner();
560        let root = self.inner_mut(new_root);
561        root.keys.push(sep_k);
562        root.vals.push(sep_v);
563        root.children.push(old_root);
564        root.children.push(right_child);
565        self.root = new_root;
566        self.height += 1;
567    }
568
569    // -- Delete -------------------------------------------------------------
570
571    /// Removes the exact `(key, val)` pair from the index.
572    ///
573    /// Returns `true` if the entry was found and removed, `false` otherwise.
574    pub fn delete(&mut self, key: &K, val: &V) -> bool {
575        let (leaf_id, path) = self.descend_kv(key, val);
576
577        let pos = match self.leaf(leaf_id).search_kv_pos(key, val) {
578            Ok(p) => p,
579            Err(_) => return false,
580        };
581
582        self.leaf_mut(leaf_id).keys.remove(pos);
583        self.leaf_mut(leaf_id).vals.remove(pos);
584        self.len -= 1;
585
586        // Root leaf never needs rebalancing.
587        if self.height == 0 {
588            return true;
589        }
590
591        if self.leaf(leaf_id).len() >= MIN_LEAF {
592            return true;
593        }
594
595        self.rebalance_leaf(leaf_id, &path);
596        true
597    }
598
599    /// Rebalance a leaf that has fallen below minimum occupancy.
600    ///
601    /// Tries in order: borrow from right sibling, borrow from left sibling,
602    /// merge with a sibling. Merges may cascade up through inner nodes.
603    fn rebalance_leaf(&mut self, leaf_id: u32, path: &Path) {
604        let parent_entry = path.entries[path.len - 1];
605        let parent_id = parent_entry.node;
606        let child_pos = parent_entry.child_pos;
607
608        // Try borrow from right sibling.
609        if child_pos + 1 < self.inner(parent_id).children.len() {
610            let right_id = self.inner(parent_id).children[child_pos + 1];
611            if self.leaf(right_id).can_lend() {
612                let rk = self.leaf_mut(right_id).keys.remove(0);
613                let rv = self.leaf_mut(right_id).vals.remove(0);
614                self.leaf_mut(leaf_id).keys.push(rk);
615                self.leaf_mut(leaf_id).vals.push(rv);
616
617                let new_sep_k = self.leaf(right_id).keys[0].clone();
618                let new_sep_v = self.leaf(right_id).vals[0].clone();
619                self.inner_mut(parent_id).keys[child_pos] = new_sep_k;
620                self.inner_mut(parent_id).vals[child_pos] = new_sep_v;
621                return;
622            }
623        }
624
625        // Try borrow from left sibling.
626        if child_pos > 0 {
627            let left_id = self.inner(parent_id).children[child_pos - 1];
628            if self.leaf(left_id).can_lend() {
629                let lk = self.leaf_mut(left_id).keys.pop().unwrap();
630                let lv = self.leaf_mut(left_id).vals.pop().unwrap();
631                av_insert(&mut self.leaf_mut(leaf_id).keys, 0, lk);
632                av_insert(&mut self.leaf_mut(leaf_id).vals, 0, lv);
633
634                let new_sep_k = self.leaf(leaf_id).keys[0].clone();
635                let new_sep_v = self.leaf(leaf_id).vals[0].clone();
636                self.inner_mut(parent_id).keys[child_pos - 1] = new_sep_k;
637                self.inner_mut(parent_id).vals[child_pos - 1] = new_sep_v;
638                return;
639            }
640        }
641
642        // Merge with a sibling.
643        if child_pos + 1 < self.inner(parent_id).children.len() {
644            let right_id = self.inner(parent_id).children[child_pos + 1];
645            self.merge_leaves(leaf_id, right_id);
646            self.inner_mut(parent_id).keys.remove(child_pos);
647            self.inner_mut(parent_id).vals.remove(child_pos);
648            self.inner_mut(parent_id).children.remove(child_pos + 1);
649            self.free_leaf(right_id);
650        } else {
651            let left_id = self.inner(parent_id).children[child_pos - 1];
652            self.merge_leaves(left_id, leaf_id);
653            self.inner_mut(parent_id).keys.remove(child_pos - 1);
654            self.inner_mut(parent_id).vals.remove(child_pos - 1);
655            self.inner_mut(parent_id).children.remove(child_pos);
656            self.free_leaf(leaf_id);
657        }
658
659        self.maybe_shrink_or_rebalance_inner(parent_id, path, path.len - 1);
660    }
661
662    /// Merge the right leaf into the left leaf and fix the linked list.
663    fn merge_leaves(&mut self, left_id: u32, right_id: u32) {
664        let right = self.leaf_mut(right_id);
665        let mut rk = std::mem::replace(&mut right.keys, ArrayVec::new());
666        let mut rv = std::mem::replace(&mut right.vals, ArrayVec::new());
667        let right_next = right.next;
668
669        let left = self.leaf_mut(left_id);
670        av_append(&mut left.keys, &mut rk);
671        av_append(&mut left.vals, &mut rv);
672        left.next = right_next;
673
674        if right_next != NONE {
675            self.leaf_mut(right_next).prev = left_id;
676        }
677    }
678
679    /// After a merge, check if the parent needs shrinking (root with one child)
680    /// or rebalancing (inner node below minimum occupancy).
681    fn maybe_shrink_or_rebalance_inner(&mut self, inner_id: u32, path: &Path, path_idx: usize) {
682        // Root with zero separators -> shrink the tree by one level.
683        if inner_id == self.root {
684            if self.inner(inner_id).keys.is_empty() {
685                let only_child = self.inner(inner_id).children[0];
686                self.free_inner(inner_id);
687                self.root = only_child;
688                self.height -= 1;
689                if self.height == 0 {
690                    self.first_leaf = self.root;
691                }
692            }
693            return;
694        }
695
696        if self.inner(inner_id).len() >= MIN_INNER {
697            return;
698        }
699
700        self.rebalance_inner(inner_id, path, path_idx);
701    }
702
703    /// Rebalance an inner node: borrow from a sibling or merge.
704    fn rebalance_inner(&mut self, inner_id: u32, path: &Path, path_idx: usize) {
705        if path_idx == 0 {
706            return;
707        }
708
709        let parent_entry = path.entries[path_idx - 1];
710        let parent_id = parent_entry.node;
711        let child_pos = parent_entry.child_pos;
712
713        // Try borrow from right sibling.
714        if child_pos + 1 < self.inner(parent_id).children.len() {
715            let right_id = self.inner(parent_id).children[child_pos + 1];
716            if self.inner(right_id).can_lend() {
717                let sep_k = self.inner_mut(parent_id).keys.remove(child_pos);
718                let sep_v = self.inner_mut(parent_id).vals.remove(child_pos);
719
720                let rk = self.inner_mut(right_id).keys.remove(0);
721                let rv = self.inner_mut(right_id).vals.remove(0);
722                let rc = self.inner_mut(right_id).children.remove(0);
723
724                av_insert(&mut self.inner_mut(parent_id).keys, child_pos, rk);
725                av_insert(&mut self.inner_mut(parent_id).vals, child_pos, rv);
726
727                self.inner_mut(inner_id).keys.push(sep_k);
728                self.inner_mut(inner_id).vals.push(sep_v);
729                self.inner_mut(inner_id).children.push(rc);
730                return;
731            }
732        }
733
734        // Try borrow from left sibling.
735        if child_pos > 0 {
736            let left_id = self.inner(parent_id).children[child_pos - 1];
737            if self.inner(left_id).can_lend() {
738                let sep_k = self.inner_mut(parent_id).keys.remove(child_pos - 1);
739                let sep_v = self.inner_mut(parent_id).vals.remove(child_pos - 1);
740
741                let lk = self.inner_mut(left_id).keys.pop().unwrap();
742                let lv = self.inner_mut(left_id).vals.pop().unwrap();
743                let lc = self.inner_mut(left_id).children.pop().unwrap();
744
745                av_insert(&mut self.inner_mut(parent_id).keys, child_pos - 1, lk);
746                av_insert(&mut self.inner_mut(parent_id).vals, child_pos - 1, lv);
747
748                av_insert(&mut self.inner_mut(inner_id).keys, 0, sep_k);
749                av_insert(&mut self.inner_mut(inner_id).vals, 0, sep_v);
750                av_insert(&mut self.inner_mut(inner_id).children, 0, lc);
751                return;
752            }
753        }
754
755        // Merge with right sibling.
756        if child_pos + 1 < self.inner(parent_id).children.len() {
757            let right_id = self.inner(parent_id).children[child_pos + 1];
758
759            let sep_k = self.inner_mut(parent_id).keys.remove(child_pos);
760            let sep_v = self.inner_mut(parent_id).vals.remove(child_pos);
761            self.inner_mut(parent_id).children.remove(child_pos + 1);
762
763            self.inner_mut(inner_id).keys.push(sep_k);
764            self.inner_mut(inner_id).vals.push(sep_v);
765
766            let right = self.inner_mut(right_id);
767            let mut rk = std::mem::replace(&mut right.keys, ArrayVec::new());
768            let mut rv = std::mem::replace(&mut right.vals, ArrayVec::new());
769            let mut rc = std::mem::replace(&mut right.children, ArrayVec::new());
770
771            av_append(&mut self.inner_mut(inner_id).keys, &mut rk);
772            av_append(&mut self.inner_mut(inner_id).vals, &mut rv);
773            av_append(&mut self.inner_mut(inner_id).children, &mut rc);
774
775            self.free_inner(right_id);
776        } else {
777            // Merge into left sibling.
778            let left_id = self.inner(parent_id).children[child_pos - 1];
779
780            let sep_k = self.inner_mut(parent_id).keys.remove(child_pos - 1);
781            let sep_v = self.inner_mut(parent_id).vals.remove(child_pos - 1);
782            self.inner_mut(parent_id).children.remove(child_pos);
783
784            self.inner_mut(left_id).keys.push(sep_k);
785            self.inner_mut(left_id).vals.push(sep_v);
786
787            let cur = self.inner_mut(inner_id);
788            let mut ck = std::mem::replace(&mut cur.keys, ArrayVec::new());
789            let mut cv = std::mem::replace(&mut cur.vals, ArrayVec::new());
790            let mut cc = std::mem::replace(&mut cur.children, ArrayVec::new());
791
792            av_append(&mut self.inner_mut(left_id).keys, &mut ck);
793            av_append(&mut self.inner_mut(left_id).vals, &mut cv);
794            av_append(&mut self.inner_mut(left_id).children, &mut cc);
795
796            self.free_inner(inner_id);
797        }
798
799        self.maybe_shrink_or_rebalance_inner(parent_id, path, path_idx - 1);
800    }
801
802    // -- Update -------------------------------------------------------------
803
804    /// Re-indexes an entry by deleting `(old_key, val)` and inserting `(new_key, val)`.
805    ///
806    /// This is a convenience for the common case where the indexed field of a
807    /// row changes but the row identifier (the value) stays the same.
808    pub fn update(&mut self, old_key: &K, new_key: K, val: V) {
809        self.delete(old_key, &val);
810        self.insert(new_key, val);
811    }
812
813    // -- Lookups ------------------------------------------------------------
814
815    /// Returns a reference to the first value associated with `key`, or `None`.
816    ///
817    /// For unique indexes this is the only value. For non-unique indexes this
818    /// returns the smallest `V` for the given `K`.
819    ///
820    /// Time: O(log n).
821    pub fn lookup_unique(&self, key: &K) -> Option<&V> {
822        let leaf_id = self.descend_k(key);
823        let leaf = self.leaf(leaf_id);
824        let pos = leaf.lower_bound_k(key);
825        if pos < leaf.len() && leaf.keys[pos] == *key {
826            return Some(&leaf.vals[pos]);
827        }
828        // Key might be at the start of the next leaf (boundary case).
829        if leaf.next != NONE {
830            let next = self.leaf(leaf.next);
831            if !next.keys.is_empty() && next.keys[0] == *key {
832                return Some(&next.vals[0]);
833            }
834        }
835        None
836    }
837
838    /// Returns an iterator over all values associated with `key`.
839    ///
840    /// Values are yielded in ascending `V` order. The iterator follows the
841    /// leaf chain, so it works correctly even when entries for one key span
842    /// multiple leaves.
843    ///
844    /// Time: O(log n) for the initial descent, then O(k) for k results.
845    pub fn lookup_range<'a>(&'a self, key: &'a K) -> LookupRange<'a, K, V> {
846        let leaf_id = self.descend_k(key);
847        let leaf = self.leaf(leaf_id);
848        let pos = leaf.lower_bound_k(key);
849        LookupRange {
850            tree: self,
851            key,
852            leaf: leaf_id,
853            pos,
854        }
855    }
856
857    /// Returns an iterator over all entries `(&K, &V)` whose keys fall within
858    /// the given range bounds.
859    ///
860    /// Supports all standard range syntax: `a..b`, `a..=b`, `a..`, `..b`, `..`.
861    ///
862    /// Time: O(log n) for the initial descent, then O(k) for k results.
863    pub fn range<R: RangeBounds<K>>(&self, bounds: R) -> RangeIter<'_, K, V> {
864        let (leaf, pos) = match bounds.start_bound() {
865            Bound::Unbounded => (self.first_leaf, 0),
866            Bound::Included(k) => {
867                let lid = self.descend_k(k);
868                let l = self.leaf(lid);
869                (lid, l.lower_bound_k(k))
870            }
871            Bound::Excluded(k) => {
872                let lid = self.descend_k(k);
873                let l = self.leaf(lid);
874                (lid, l.upper_bound_k(k))
875            }
876        };
877
878        let end = match bounds.end_bound() {
879            Bound::Unbounded => RangeEnd::Unbounded,
880            Bound::Included(k) => RangeEnd::Included(k.clone()),
881            Bound::Excluded(k) => RangeEnd::Excluded(k.clone()),
882        };
883
884        RangeIter {
885            tree: self,
886            leaf,
887            pos,
888            end,
889        }
890    }
891
892    /// Returns an iterator over all entries in ascending order.
893    ///
894    /// Time: O(n).
895    pub fn iter(&self) -> RangeIter<'_, K, V> {
896        RangeIter {
897            tree: self,
898            leaf: self.first_leaf,
899            pos: 0,
900            end: RangeEnd::Unbounded,
901        }
902    }
903}
904
905impl<K: Ord + Clone, V: Ord + Clone> Default for BPlusIndex<K, V> {
906    fn default() -> Self {
907        Self::new()
908    }
909}
910
911// ---------------------------------------------------------------------------
912// Iterators
913// ---------------------------------------------------------------------------
914
915/// Iterator returned by [`BPlusIndex::lookup_range`].
916///
917/// Yields `&V` for all entries matching a specific key.
918pub struct LookupRange<'a, K, V> {
919    tree: &'a BPlusIndex<K, V>,
920    key: &'a K,
921    leaf: u32,
922    pos: usize,
923}
924
925impl<'a, K: Ord + Clone, V: Ord + Clone> Iterator for LookupRange<'a, K, V> {
926    type Item = &'a V;
927
928    fn next(&mut self) -> Option<Self::Item> {
929        loop {
930            if self.leaf == NONE {
931                return None;
932            }
933            let leaf = self.tree.leaf(self.leaf);
934            if self.pos < leaf.len() {
935                if leaf.keys[self.pos] == *self.key {
936                    let v = &leaf.vals[self.pos];
937                    self.pos += 1;
938                    return Some(v);
939                }
940                return None;
941            }
942            self.leaf = leaf.next;
943            self.pos = 0;
944        }
945    }
946}
947
948enum RangeEnd<K> {
949    Unbounded,
950    Included(K),
951    Excluded(K),
952}
953
954/// Iterator returned by [`BPlusIndex::range`] and [`BPlusIndex::iter`].
955///
956/// Yields `(&K, &V)` pairs in ascending order. Follows the leaf chain
957/// for sequential memory access.
958pub struct RangeIter<'a, K, V> {
959    tree: &'a BPlusIndex<K, V>,
960    leaf: u32,
961    pos: usize,
962    end: RangeEnd<K>,
963}
964
965impl<'a, K: Ord + Clone, V: Ord + Clone> Iterator for RangeIter<'a, K, V> {
966    type Item = (&'a K, &'a V);
967
968    fn next(&mut self) -> Option<Self::Item> {
969        loop {
970            if self.leaf == NONE {
971                return None;
972            }
973            let leaf = self.tree.leaf(self.leaf);
974            if self.pos < leaf.len() {
975                let k = &leaf.keys[self.pos];
976                match &self.end {
977                    RangeEnd::Excluded(end) if k >= end => return None,
978                    RangeEnd::Included(end) if k > end => return None,
979                    _ => {}
980                }
981                let v = &leaf.vals[self.pos];
982                self.pos += 1;
983                return Some((k, v));
984            }
985            self.leaf = leaf.next;
986            self.pos = 0;
987        }
988    }
989}
990
991// ---------------------------------------------------------------------------
992// Tests
993// ---------------------------------------------------------------------------
994
995#[cfg(test)]
996mod tests {
997    use super::*;
998
999    #[test]
1000    fn insert_and_lookup_unique() {
1001        let mut idx = BPlusIndex::<i32, u64>::new();
1002        for i in 0..200 {
1003            idx.insert(i, i as u64 * 10);
1004        }
1005        assert_eq!(idx.len(), 200);
1006        for i in 0..200 {
1007            assert_eq!(idx.lookup_unique(&i), Some(&(i as u64 * 10)));
1008        }
1009        assert_eq!(idx.lookup_unique(&999), None);
1010    }
1011
1012    #[test]
1013    fn insert_duplicates_ignored() {
1014        let mut idx = BPlusIndex::<i32, u64>::new();
1015        idx.insert(5, 100);
1016        idx.insert(5, 100);
1017        assert_eq!(idx.len(), 1);
1018    }
1019
1020    #[test]
1021    fn non_unique_keys() {
1022        let mut idx = BPlusIndex::<i32, u64>::new();
1023        idx.insert(5, 100);
1024        idx.insert(5, 200);
1025        idx.insert(5, 300);
1026        idx.insert(10, 400);
1027
1028        let vals: Vec<_> = idx.lookup_range(&5).copied().collect();
1029        assert_eq!(vals, vec![100, 200, 300]);
1030    }
1031
1032    #[test]
1033    fn delete_and_rebalance() {
1034        let mut idx = BPlusIndex::<i32, u64>::new();
1035        let n = 500;
1036        for i in 0..n {
1037            idx.insert(i, i as u64);
1038        }
1039        for i in 0..n {
1040            assert!(idx.delete(&i, &(i as u64)));
1041        }
1042        assert_eq!(idx.len(), 0);
1043        assert!(idx.is_empty());
1044    }
1045
1046    #[test]
1047    fn delete_nonexistent() {
1048        let mut idx = BPlusIndex::<i32, u64>::new();
1049        idx.insert(1, 10);
1050        assert!(!idx.delete(&1, &999));
1051        assert!(!idx.delete(&999, &10));
1052        assert_eq!(idx.len(), 1);
1053    }
1054
1055    #[test]
1056    fn range_query() {
1057        let mut idx = BPlusIndex::<i32, u64>::new();
1058        for i in 0..100 {
1059            idx.insert(i, i as u64);
1060        }
1061        let vals: Vec<_> = idx.range(10..20).map(|(k, _)| *k).collect();
1062        assert_eq!(vals, (10..20).collect::<Vec<_>>());
1063    }
1064
1065    #[test]
1066    fn range_inclusive() {
1067        let mut idx = BPlusIndex::<i32, u64>::new();
1068        for i in 0..10 {
1069            idx.insert(i, i as u64);
1070        }
1071        let vals: Vec<_> = idx.range(3..=7).map(|(k, _)| *k).collect();
1072        assert_eq!(vals, vec![3, 4, 5, 6, 7]);
1073    }
1074
1075    #[test]
1076    fn iter_full() {
1077        let mut idx = BPlusIndex::<i32, u64>::new();
1078        for i in (0..300).rev() {
1079            idx.insert(i, i as u64);
1080        }
1081        let keys: Vec<_> = idx.iter().map(|(k, _)| *k).collect();
1082        assert_eq!(keys, (0..300).collect::<Vec<_>>());
1083    }
1084
1085    #[test]
1086    fn update_moves_entry() {
1087        let mut idx = BPlusIndex::<i32, u64>::new();
1088        idx.insert(5, 42);
1089        idx.update(&5, 10, 42);
1090        assert_eq!(idx.lookup_unique(&5), None);
1091        assert_eq!(idx.lookup_unique(&10), Some(&42));
1092    }
1093
1094    #[test]
1095    fn stress_insert_delete_interleaved() {
1096        let mut idx = BPlusIndex::<i32, u64>::new();
1097        for i in 0..1000 {
1098            idx.insert(i, i as u64);
1099        }
1100        for i in (0..1000).step_by(2) {
1101            assert!(idx.delete(&i, &(i as u64)));
1102        }
1103        assert_eq!(idx.len(), 500);
1104        for i in (1..1000).step_by(2) {
1105            assert_eq!(idx.lookup_unique(&i), Some(&(i as u64)));
1106        }
1107        for i in (0..1000).step_by(2) {
1108            idx.insert(i, i as u64);
1109        }
1110        assert_eq!(idx.len(), 1000);
1111        let keys: Vec<_> = idx.iter().map(|(k, _)| *k).collect();
1112        assert_eq!(keys, (0..1000).collect::<Vec<_>>());
1113    }
1114
1115    #[test]
1116    fn many_duplicates_same_key() {
1117        let mut idx = BPlusIndex::<i32, u64>::new();
1118        for v in 0..200u64 {
1119            idx.insert(42, v);
1120        }
1121        assert_eq!(idx.len(), 200);
1122        let vals: Vec<_> = idx.lookup_range(&42).copied().collect();
1123        assert_eq!(vals, (0..200u64).collect::<Vec<_>>());
1124
1125        for v in (0..200u64).step_by(2) {
1126            assert!(idx.delete(&42, &v));
1127        }
1128        assert_eq!(idx.len(), 100);
1129        let vals: Vec<_> = idx.lookup_range(&42).copied().collect();
1130        assert_eq!(
1131            vals,
1132            (0..200u64).step_by(2).map(|v| v + 1).collect::<Vec<_>>()
1133        );
1134    }
1135
1136    #[test]
1137    fn reverse_insertion_order() {
1138        let mut idx = BPlusIndex::<i32, u64>::new();
1139        for i in (0..500).rev() {
1140            idx.insert(i, i as u64);
1141        }
1142        let keys: Vec<_> = idx.iter().map(|(k, _)| *k).collect();
1143        assert_eq!(keys, (0..500).collect::<Vec<_>>());
1144    }
1145}