Skip to main content

reddb_engine/storage/
btree.rs

1//! B+ Tree Implementation for Page-Based Storage
2//!
3//! A B+ tree optimized for disk I/O with the following properties:
4//! - All values stored in leaf nodes
5//! - Interior nodes only contain keys and child pointers
6//! - Leaf nodes form a doubly-linked list for range scans
7//! - Each node fits in a single 16 KiB page
8//!
9//! # Page Layout
10//!
11//! Interior Node:
12//! ```text
13//! ┌────────────────────────────────────────────────────────────┐
14//! │ PageHeader (32 bytes)                                      │
15//! ├────────────────────────────────────────────────────────────┤
16//! │ right_child: u32 - Rightmost child pointer                 │
17//! │ cells: [key_len: u16, key: [u8], child: u32]...           │
18//! └────────────────────────────────────────────────────────────┘
19//! ```
20//!
21//! Leaf Node:
22//! ```text
23//! ┌────────────────────────────────────────────────────────────┐
24//! │ PageHeader (32 bytes)                                      │
25//! ├────────────────────────────────────────────────────────────┤
26//! │ prev_leaf: u32 - Previous leaf (0 = none)                  │
27//! │ next_leaf: u32 - Next leaf (0 = none)                      │
28//! │ cells: [key_len: u16, val_len: u16, key: [u8], val: [u8]] │
29//! └────────────────────────────────────────────────────────────┘
30//! ```
31
32use std::cmp::Ordering;
33use std::sync::{Arc, RwLock};
34
35use super::page::{Page, PageType, HEADER_SIZE, PAGE_SIZE};
36use super::page_store::PageStore;
37use super::pager::PagerError;
38
39/// Maximum key size (to ensure at least 2 keys per node)
40pub const MAX_KEY_SIZE: usize = 1024;
41
42/// Hard upper bound on logical value size.
43///
44/// Per ADR 0023 (large-value overflow), values up to 256 MiB are stored
45/// transparently — inline below [`value_layout::OVERFLOW_THRESHOLD`],
46/// inline+compressed when LZ4 shrinks them under the same threshold, or
47/// spilled into an [`overflow::OverflowChain`] otherwise. Values above
48/// this ceiling are rejected with [`BTreeError::ValueTooLarge`] before
49/// any compression or page allocation runs.
50///
51/// Re-exported from [`value_layout::MAX_VALUE_SIZE`] so callers that
52/// historically read this constant (`unified::store::impl_pages`,
53/// `query::planner::stats_catalog`) continue to compile.
54pub const MAX_VALUE_SIZE: usize = value_layout::MAX_VALUE_SIZE;
55
56/// Minimum fill factor before merge (as percentage)
57pub const MIN_FILL_FACTOR: usize = 25;
58
59/// Offset of prev_leaf in leaf page
60const LEAF_PREV_OFFSET: usize = HEADER_SIZE;
61
62/// Offset of next_leaf in leaf page
63const LEAF_NEXT_OFFSET: usize = HEADER_SIZE + 4;
64
65/// Start of the slot array in leaf pages.
66///
67/// Slotted-page layout:
68/// ```text
69/// [PageHeader 32B][prev u32 | next u32 | 8B]
70/// [slot 0: u16][slot 1: u16]...[slot N-1: u16]  ← grows forward
71/// ... free space ...
72/// [cell N-1][cell N-2]...[cell 0]               ← grows backward from free_end
73/// ```
74/// Each cell is laid out as `[key_len:u16][val_len:u16][key][val]`.
75/// `page.cell_count()` is N; `page.free_end()` is the offset of the
76/// lowest (most recently written) cell. The slot array lives right
77/// after the leaf-chain links and each u16 slot is the absolute page
78/// offset of its cell.
79const LEAF_SLOT_ARRAY_OFFSET: usize = HEADER_SIZE + 8;
80
81// A split must still fit two maximum-key inline cells, including the value
82// flag, cell headers and slot pointers, after the native page-size change.
83const _: () = assert!(
84    2 * (MAX_KEY_SIZE + value_layout::OVERFLOW_THRESHOLD + 1 + 4 + 2)
85        <= PAGE_SIZE - LEAF_SLOT_ARRAY_OFFSET
86);
87
88/// Kept for source-compat with older code paths (e.g. interior-node
89/// helpers); equivalent to `LEAF_SLOT_ARRAY_OFFSET` for leaves.
90const LEAF_DATA_OFFSET: usize = LEAF_SLOT_ARRAY_OFFSET;
91
92/// Start of cell data in interior pages (right_child is in header)
93const INTERIOR_DATA_OFFSET: usize = HEADER_SIZE;
94
95/// B+ Tree error types
96#[derive(Debug, Clone)]
97pub enum BTreeError {
98    /// Key not found
99    NotFound,
100    /// Key already exists
101    DuplicateKey,
102    /// Key too large
103    KeyTooLarge(usize),
104    /// Value too large
105    ValueTooLarge(usize),
106    /// Tree is corrupted
107    Corrupted(String),
108    /// Pager error
109    Pager(String),
110    /// Internal lock was poisoned by a panicked thread
111    LockPoisoned(String),
112    /// Large-value layout encode/decode failed (overflow chain or LZ4).
113    ValueLayout(String),
114    /// Materializing a value would exceed the caller's read memory allowance.
115    ReadLimitExceeded {
116        required_bytes: usize,
117        limit_bytes: usize,
118    },
119}
120
121impl std::fmt::Display for BTreeError {
122    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
123        match self {
124            Self::NotFound => write!(f, "Key not found"),
125            Self::DuplicateKey => write!(f, "Key already exists"),
126            Self::KeyTooLarge(size) => {
127                write!(f, "Key too large: {} bytes (max {})", size, MAX_KEY_SIZE)
128            }
129            Self::ValueTooLarge(size) => write!(
130                f,
131                "Value too large: {} bytes (max {})",
132                size, MAX_VALUE_SIZE
133            ),
134            Self::Corrupted(msg) => write!(f, "B-tree corrupted: {}", msg),
135            Self::Pager(msg) => write!(f, "Pager error: {}", msg),
136            Self::LockPoisoned(msg) => write!(f, "Lock poisoned: {}", msg),
137            Self::ValueLayout(msg) => write!(f, "Value layout: {}", msg),
138            Self::ReadLimitExceeded {
139                required_bytes,
140                limit_bytes,
141            } => write!(
142                f,
143                "B-tree read needs {required_bytes} bytes over read memory limit {limit_bytes}"
144            ),
145        }
146    }
147}
148
149impl From<value_layout::ValueLayoutError> for BTreeError {
150    fn from(e: value_layout::ValueLayoutError) -> Self {
151        match e {
152            value_layout::ValueLayoutError::ValueTooLarge(n) => Self::ValueTooLarge(n),
153            value_layout::ValueLayoutError::ReadLimitExceeded {
154                required_bytes,
155                limit_bytes,
156            } => Self::ReadLimitExceeded {
157                required_bytes,
158                limit_bytes,
159            },
160            other => Self::ValueLayout(other.to_string()),
161        }
162    }
163}
164
165impl std::error::Error for BTreeError {}
166
167impl From<PagerError> for BTreeError {
168    fn from(e: PagerError) -> Self {
169        Self::Pager(e.to_string())
170    }
171}
172
173impl From<super::page::PageError> for BTreeError {
174    fn from(e: super::page::PageError) -> Self {
175        Self::Corrupted(e.to_string())
176    }
177}
178
179/// Result type for B+ tree operations
180pub type BTreeResult<T> = Result<T, BTreeError>;
181
182/// B+ Tree cursor for iteration
183pub struct BTreeCursor<P: PageStore> {
184    /// A cursor can outlive the tree handle that created it.
185    _retirement_pin: Arc<()>,
186    /// Current leaf page ID
187    leaf_page_id: u32,
188    /// Current position within leaf
189    position: usize,
190    /// Pager reference
191    pager: Arc<P>,
192    /// A6 — true once we've issued a prefetch hint for the next leaf on
193    /// this leaf. Reset to false every time we move to a new leaf so the
194    /// hint fires once per leaf, at the halfway point.
195    prefetched_next: bool,
196    /// Published collection generations are immutable. Retain at most one
197    /// leaf while walking them; mutable-tree cursors keep re-reading pages.
198    immutable: bool,
199    immutable_leaf: Option<Box<Page>>,
200}
201
202impl<P: PageStore> BTreeCursor<P> {
203    /// Copy the position of an immutable cursor without advancing the original.
204    pub(crate) fn fork_immutable(&self) -> Self {
205        assert!(self.immutable, "only immutable cursors may be forked");
206        Self {
207            _retirement_pin: Arc::clone(&self._retirement_pin),
208            leaf_page_id: self.leaf_page_id,
209            position: self.position,
210            pager: Arc::clone(&self.pager),
211            prefetched_next: self.prefetched_next,
212            immutable: true,
213            immutable_leaf: self.immutable_leaf.clone(),
214        }
215    }
216
217    /// Only collection readers whose host pins an immutable published root
218    /// may retain a leaf across calls. This is not a mutable-tree snapshot.
219    pub(crate) fn with_immutable_leaves(mut self) -> Self {
220        self.immutable = true;
221        self
222    }
223
224    fn prepare_immutable_leaf(&mut self) -> BTreeResult<bool> {
225        if self.leaf_page_id == 0 {
226            return Ok(false);
227        }
228        if self.immutable_leaf.is_none() {
229            let page = self.pager.read_page(self.leaf_page_id)?;
230            if page.page_type()? != PageType::BTreeLeaf {
231                return Err(BTreeError::Corrupted(
232                    "cursor reached a non-leaf page".into(),
233                ));
234            }
235            self.immutable_leaf = Some(Box::new(page));
236        }
237        let page = self.immutable_leaf.as_ref().expect("loaded immutable leaf");
238        if self.position < usize::from(page.cell_count()) {
239            return Ok(true);
240        }
241        let next = read_next_leaf(page);
242        if next == 0 {
243            self.leaf_page_id = 0;
244            self.immutable_leaf = None;
245            return Ok(false);
246        }
247        let next_page = self.pager.read_page(next)?;
248        if next_page.page_type()? != PageType::BTreeLeaf || next_page.cell_count() == 0 {
249            return Err(BTreeError::Corrupted(
250                "leaf chain reached an invalid leaf".into(),
251            ));
252        }
253        let first_key = read_leaf_cell_borrowed(&next_page, 0)?.0;
254        if leaf_high_key(page)?
255            .as_ref()
256            .is_some_and(|key| first_key <= key.as_slice())
257        {
258            return Err(BTreeError::Corrupted(
259                "leaf chain does not advance its key range".into(),
260            ));
261        }
262        **self.immutable_leaf.as_mut().expect("loaded immutable leaf") = next_page;
263        self.leaf_page_id = next;
264        self.position = 0;
265        self.prefetched_next = false;
266        Ok(true)
267    }
268
269    /// Inspect the next key without materializing its value or advancing.
270    #[cfg(feature = "full")]
271    pub(crate) fn peek_key_with<T>(
272        &mut self,
273        visit: impl FnOnce(&[u8]) -> T,
274    ) -> BTreeResult<Option<T>> {
275        if self.immutable {
276            if !self.prepare_immutable_leaf()? {
277                return Ok(None);
278            }
279            let page = self.immutable_leaf.as_ref().expect("loaded immutable leaf");
280            return Ok(Some(visit(read_leaf_cell_borrowed(page, self.position)?.0)));
281        }
282        let Some(page) = self.prepare_next_leaf()? else {
283            return Ok(None);
284        };
285        Ok(Some(visit(
286            read_leaf_cell_borrowed(&page, self.position)?.0,
287        )))
288    }
289
290    /// Inspect the next key without materializing its value or advancing.
291    pub fn peek_key(&mut self) -> BTreeResult<Option<Vec<u8>>> {
292        if self.immutable {
293            if !self.prepare_immutable_leaf()? {
294                return Ok(None);
295            }
296            let page = self.immutable_leaf.as_ref().expect("loaded immutable leaf");
297            return Ok(Some(
298                read_leaf_cell_borrowed(page, self.position)?.0.to_vec(),
299            ));
300        }
301        let Some(page) = self.prepare_next_leaf()? else {
302            return Ok(None);
303        };
304        Ok(Some(read_leaf_cell(&page, self.position)?.0))
305    }
306
307    /// Skip a shadowed or deleted base record without reading its overflow value.
308    pub fn skip_next(&mut self) -> BTreeResult<bool> {
309        if self.immutable {
310            if !self.prepare_immutable_leaf()? {
311                return Ok(false);
312            }
313            self.position += 1;
314            return Ok(true);
315        }
316        if self.prepare_next_leaf()?.is_none() {
317            return Ok(false);
318        }
319        self.position += 1;
320        Ok(true)
321    }
322
323    fn prepare_next_leaf(&mut self) -> BTreeResult<Option<Page>> {
324        if self.leaf_page_id == 0 {
325            return Ok(None);
326        }
327        let page = self.pager.read_page(self.leaf_page_id)?;
328        if page.page_type()? != PageType::BTreeLeaf {
329            return Err(BTreeError::Corrupted(
330                "cursor reached a non-leaf page".into(),
331            ));
332        }
333        if self.position < usize::from(page.cell_count()) {
334            return Ok(Some(page));
335        }
336        let next = read_next_leaf(&page);
337        if next == 0 {
338            self.leaf_page_id = 0;
339            return Ok(None);
340        }
341        let next_page = self.pager.read_page(next)?;
342        if next_page.page_type()? != PageType::BTreeLeaf || next_page.cell_count() == 0 {
343            return Err(BTreeError::Corrupted(
344                "leaf chain reached an invalid leaf".into(),
345            ));
346        }
347        let first_key = read_leaf_cell(&next_page, 0)?.0;
348        if leaf_high_key(&page)?
349            .as_ref()
350            .is_some_and(|key| first_key <= *key)
351        {
352            return Err(BTreeError::Corrupted(
353                "leaf chain does not advance its key range".into(),
354            ));
355        }
356        self.leaf_page_id = next;
357        self.position = 0;
358        self.prefetched_next = false;
359        Ok(Some(next_page))
360    }
361
362    /// Read one value within the caller's materialization allowance.
363    pub fn next_with_budget(
364        &mut self,
365        limit_bytes: usize,
366    ) -> BTreeResult<Option<(Vec<u8>, Vec<u8>)>> {
367        self.next_impl(limit_bytes)
368    }
369
370    /// Move to next entry
371    pub fn next(&mut self) -> BTreeResult<Option<(Vec<u8>, Vec<u8>)>> {
372        self.next_impl(usize::MAX)
373    }
374
375    fn next_impl(&mut self, limit_bytes: usize) -> BTreeResult<Option<(Vec<u8>, Vec<u8>)>> {
376        if self.immutable {
377            return self.next_immutable(limit_bytes);
378        }
379        if self.leaf_page_id == 0 {
380            return Ok(None);
381        }
382
383        let page = self.pager.read_page(self.leaf_page_id)?;
384        if page.page_type()? != PageType::BTreeLeaf {
385            return Err(BTreeError::Corrupted(
386                "cursor reached a non-leaf page".into(),
387            ));
388        }
389        let cell_count = page.cell_count() as usize;
390
391        // A6 — prefetch the next leaf when we cross the halfway point of
392        // the current one. The hint fires at most once per leaf (guarded
393        // by `prefetched_next`) so the syscall overhead is amortized over
394        // ~half a leaf's worth of entries. The kernel uses the hint to
395        // start DMA-ing the next page while we finish the current one.
396        if !self.prefetched_next && self.position >= cell_count / 2 && cell_count > 0 {
397            let next_leaf = read_next_leaf(&page);
398            if next_leaf != 0 {
399                self.pager.prefetch_hint(next_leaf);
400            }
401            self.prefetched_next = true;
402        }
403
404        // Check if we have more cells in current page
405        if self.position < cell_count {
406            let (key, stored) = read_leaf_cell(&page, self.position)?;
407            let value =
408                value_layout::decode_with_budget(self.pager.as_ref(), &stored, limit_bytes)?;
409            self.position += 1;
410            return Ok(Some((key, value)));
411        }
412
413        // Move to next leaf
414        let next_leaf = read_next_leaf(&page);
415        if next_leaf == 0 {
416            self.leaf_page_id = 0;
417            return Ok(None);
418        }
419        let previous_high_key = leaf_high_key(&page)?;
420
421        self.leaf_page_id = next_leaf;
422        self.position = 0;
423        self.prefetched_next = false; // reset for the new leaf
424
425        // Read from new leaf
426        let page = self.pager.read_page(self.leaf_page_id)?;
427        if page.page_type()? != PageType::BTreeLeaf {
428            return Err(BTreeError::Corrupted(
429                "leaf chain reached a non-leaf page".into(),
430            ));
431        }
432        if page.cell_count() == 0 {
433            return Ok(None);
434        }
435
436        let (key, stored) = read_leaf_cell(&page, 0)?;
437        // Every leaf transition must advance the key range. This bounds a
438        // corrupt cyclic chain without growing a visited-page set or an
439        // unbounded checkpoint buffer (#2376).
440        if previous_high_key
441            .as_ref()
442            .is_some_and(|previous| key <= *previous)
443        {
444            self.leaf_page_id = 0;
445            return Err(BTreeError::Corrupted(
446                "leaf chain does not advance its key range".into(),
447            ));
448        }
449        let value = value_layout::decode_with_budget(self.pager.as_ref(), &stored, limit_bytes)?;
450        self.position = 1;
451        Ok(Some((key, value)))
452    }
453
454    fn next_immutable(&mut self, limit_bytes: usize) -> BTreeResult<Option<(Vec<u8>, Vec<u8>)>> {
455        if !self.prepare_immutable_leaf()? {
456            return Ok(None);
457        }
458        let page = self.immutable_leaf.as_ref().expect("loaded immutable leaf");
459        if !self.prefetched_next && self.position >= usize::from(page.cell_count()) / 2 {
460            let next = read_next_leaf(page);
461            if next != 0 {
462                self.pager.prefetch_hint(next);
463            }
464            self.prefetched_next = true;
465        }
466        let (key, stored) = read_leaf_cell_borrowed(page, self.position)?;
467        let value = value_layout::decode_with_budget(self.pager.as_ref(), stored, limit_bytes)?;
468        self.position += 1;
469        Ok(Some((key.to_vec(), value)))
470    }
471
472    /// Peek at current entry without advancing
473    pub fn peek(&self) -> BTreeResult<Option<(Vec<u8>, Vec<u8>)>> {
474        if self.leaf_page_id == 0 {
475            return Ok(None);
476        }
477
478        let page = self.pager.read_page(self.leaf_page_id)?;
479        let cell_count = page.cell_count() as usize;
480
481        if self.position >= cell_count {
482            return Ok(None);
483        }
484
485        let (key, stored) = read_leaf_cell(&page, self.position)?;
486        let value = value_layout::decode(self.pager.as_ref(), &stored)?;
487        Ok(Some((key, value)))
488    }
489}
490
491/// B+ Tree implementation
492pub struct BTree<P: PageStore> {
493    retirement_pin: Arc<()>,
494    /// Pager for page I/O
495    pager: Arc<P>,
496    /// Root page ID (0 = empty tree)
497    root_page_id: RwLock<u32>,
498    /// D2 fastpath: cached rightmost leaf (page_id, high_key).
499    ///
500    /// When the next key to insert is strictly greater than `high_key`,
501    /// skip `find_leaf` (tree descent from root) and go directly to the
502    /// cached page. Invalidated on splits and deletes.
503    ///
504    /// Mirrors `BTPageCacheIsValid` / `BTREE_FASTPATH_MIN_LEVEL` in
505    /// `nbtinsert.c:31`.
506    rightmost_leaf: RwLock<Option<(u32, Vec<u8>)>>,
507    /// Serialises structural writers (insert / upsert / delete / bulk).
508    ///
509    /// Writers copy a page, modify the copy and write it back, and there
510    /// are no page latches, so two writers on one tree would lose each
511    /// other's leaf updates or tear a split (#2376). Reentrant because the
512    /// public mutators call each other (`upsert` -> `insert`, the bulk
513    /// paths -> `insert_encoded`). Readers do not take it.
514    writer: parking_lot::ReentrantMutex<()>,
515}
516
517#[path = "btree/impl.rs"]
518mod btree_impl;
519
520#[path = "btree/reclaim.rs"]
521mod reclaim;
522
523#[cfg(test)]
524#[path = "btree/read_budget_tests.rs"]
525mod read_budget_tests;
526
527#[path = "btree/lehman_yao.rs"]
528pub mod lehman_yao;
529
530#[path = "btree/value_layout.rs"]
531pub mod value_layout;
532// ==================== Search Helpers ====================
533
534enum SearchResult {
535    Found(usize),
536    NotFound(usize),
537}
538
539fn search_leaf(page: &Page, key: &[u8]) -> BTreeResult<SearchResult> {
540    let cell_count = page.cell_count() as usize;
541
542    // Binary search
543    let mut low = 0;
544    let mut high = cell_count;
545
546    while low < high {
547        let mid = (low + high) / 2;
548        let (cell_key, _) = read_leaf_cell(page, mid)?;
549
550        match cell_key.as_slice().cmp(key) {
551            Ordering::Less => low = mid + 1,
552            Ordering::Greater => high = mid,
553            Ordering::Equal => return Ok(SearchResult::Found(mid)),
554        }
555    }
556
557    Ok(SearchResult::NotFound(low))
558}
559
560fn find_child(page: &Page, key: &[u8]) -> BTreeResult<u32> {
561    let cell_count = page.cell_count() as usize;
562
563    // Binary search for the correct child
564    for i in 0..cell_count {
565        let (cell_key, child) = read_interior_cell(page, i)?;
566        if key < cell_key.as_slice() {
567            return Ok(child);
568        }
569    }
570
571    // Key is >= all keys, use right child
572    Ok(page.right_child())
573}
574
575fn find_first_child(page: &Page) -> BTreeResult<u32> {
576    if page.cell_count() == 0 {
577        return Ok(page.right_child());
578    }
579    let (_, child) = read_interior_cell(page, 0)?;
580    Ok(child)
581}
582
583fn leaf_min_bytes() -> usize {
584    (PAGE_SIZE - LEAF_DATA_OFFSET) * MIN_FILL_FACTOR / 100
585}
586
587fn interior_min_bytes() -> usize {
588    (PAGE_SIZE - INTERIOR_DATA_OFFSET) * MIN_FILL_FACTOR / 100
589}
590
591/// Bytes consumed by one leaf entry in the slotted layout: one u16
592/// slot in the pointer array plus the cell itself (`[key_len:u16]
593/// [val_len:u16][key][val]`).
594fn leaf_entry_size(entry: &(Vec<u8>, Vec<u8>)) -> usize {
595    2 + 4 + entry.0.len() + entry.1.len()
596}
597
598fn leaf_entries_size(entries: &[(Vec<u8>, Vec<u8>)]) -> usize {
599    entries.iter().map(leaf_entry_size).sum()
600}
601
602#[inline]
603fn leaf_slot_offset_for(index: usize) -> usize {
604    LEAF_SLOT_ARRAY_OFFSET + index * 2
605}
606
607#[inline]
608fn leaf_read_slot(page: &Page, index: usize) -> BTreeResult<usize> {
609    let data = page.as_bytes();
610    let slot_pos = leaf_slot_offset_for(index);
611    if slot_pos + 2 > PAGE_SIZE {
612        return Err(BTreeError::Corrupted("slot array overflows page".into()));
613    }
614    Ok(u16::from_le_bytes([data[slot_pos], data[slot_pos + 1]]) as usize)
615}
616
617#[inline]
618fn leaf_write_slot(page: &mut Page, index: usize, cell_offset: u16) -> BTreeResult<()> {
619    let data = page.as_bytes_mut();
620    let slot_pos = leaf_slot_offset_for(index);
621    if slot_pos + 2 > PAGE_SIZE {
622        return Err(BTreeError::Corrupted("slot array overflows page".into()));
623    }
624    data[slot_pos..slot_pos + 2].copy_from_slice(&cell_offset.to_le_bytes());
625    Ok(())
626}
627
628#[inline]
629fn leaf_slots_end(page: &Page) -> usize {
630    LEAF_SLOT_ARRAY_OFFSET + (page.cell_count() as usize) * 2
631}
632
633#[inline]
634fn leaf_cells_start(page: &Page) -> usize {
635    let end = page.free_end() as usize;
636    if end == 0 {
637        PAGE_SIZE
638    } else {
639        end
640    }
641}
642
643#[inline]
644fn leaf_free_bytes(page: &Page) -> usize {
645    let slot_end = leaf_slots_end(page);
646    let cells = leaf_cells_start(page);
647    cells.saturating_sub(slot_end)
648}
649
650fn interior_key_size(key: &[u8]) -> usize {
651    2 + key.len() + 4
652}
653
654fn interior_entries_size(keys: &[Vec<u8>]) -> usize {
655    keys.iter().map(|k| interior_key_size(k)).sum()
656}
657
658// ==================== Leaf Page Helpers ====================
659
660/// Read the cell at slot `index` in O(1). Follows the u16 slot pointer
661/// into the cell data area; the cell header is `[key_len:u16][val_len:u16]`
662/// followed by the raw key and value bytes.
663fn read_leaf_cell(page: &Page, index: usize) -> BTreeResult<(Vec<u8>, Vec<u8>)> {
664    let (key, value) = read_leaf_cell_borrowed(page, index)?;
665    Ok((key.to_vec(), value.to_vec()))
666}
667
668fn read_leaf_cell_borrowed(page: &Page, index: usize) -> BTreeResult<(&[u8], &[u8])> {
669    let cell_count = page.cell_count() as usize;
670    if index >= cell_count {
671        return Err(BTreeError::Corrupted("Cell index out of range".into()));
672    }
673    let offset = leaf_read_slot(page, index)?;
674    let data = page.as_bytes();
675    if offset + 4 > PAGE_SIZE {
676        return Err(BTreeError::Corrupted("cell header out of range".into()));
677    }
678    let key_len = u16::from_le_bytes([data[offset], data[offset + 1]]) as usize;
679    let val_len = u16::from_le_bytes([data[offset + 2], data[offset + 3]]) as usize;
680    let end = offset + 4 + key_len + val_len;
681    if end > PAGE_SIZE {
682        return Err(BTreeError::Corrupted("cell body out of range".into()));
683    }
684    let key = &data[offset + 4..offset + 4 + key_len];
685    let value = &data[offset + 4 + key_len..end];
686    Ok((key, value))
687}
688
689fn read_leaf_entries(page: &Page) -> BTreeResult<Vec<(Vec<u8>, Vec<u8>)>> {
690    let cell_count = page.cell_count() as usize;
691    let mut entries = Vec::with_capacity(cell_count);
692    for i in 0..cell_count {
693        entries.push(read_leaf_cell(page, i)?);
694    }
695    Ok(entries)
696}
697
698/// Wipe and re-lay an entire leaf in slotted form. Used by the split
699/// path (which hands us a pre-sorted `entries` vector) and by
700/// `clear_leaf_cells` under the hood. Entries are appended at the end
701/// of the free area so the highest-index slot ends up at the lowest
702/// offset — same shape the single-key insert path produces.
703fn write_leaf_entries(page: &mut Page, entries: &[(Vec<u8>, Vec<u8>)]) -> BTreeResult<()> {
704    clear_leaf_cells(page);
705    for (i, (k, v)) in entries.iter().enumerate() {
706        let cell_size = 4 + k.len() + v.len();
707        let cells_start = leaf_cells_start(page);
708        let slot_end = LEAF_SLOT_ARRAY_OFFSET + (i + 1) * 2;
709        if slot_end + cell_size > cells_start {
710            return Err(BTreeError::Corrupted("leaf rebuild overflowed page".into()));
711        }
712        let cell_offset = cells_start - cell_size;
713        {
714            let data = page.as_bytes_mut();
715            data[cell_offset..cell_offset + 2].copy_from_slice(&(k.len() as u16).to_le_bytes());
716            data[cell_offset + 2..cell_offset + 4].copy_from_slice(&(v.len() as u16).to_le_bytes());
717            data[cell_offset + 4..cell_offset + 4 + k.len()].copy_from_slice(k);
718            data[cell_offset + 4 + k.len()..cell_offset + 4 + k.len() + v.len()].copy_from_slice(v);
719        }
720        page.set_free_end(cell_offset as u16);
721        leaf_write_slot(page, i, cell_offset as u16)?;
722    }
723    page.set_cell_count(entries.len() as u16);
724    page.set_free_start((LEAF_SLOT_ARRAY_OFFSET + entries.len() * 2) as u16);
725    Ok(())
726}
727
728fn can_insert_leaf(page: &Page, key: &[u8], value: &[u8]) -> bool {
729    let needed = 2 + 4 + key.len() + value.len();
730    leaf_free_bytes(page) >= needed
731}
732
733/// Insert `(key, value)` into the slotted leaf in O(log M) search +
734/// O(M) slot-array memmove. Cell data is appended at the tail of the
735/// free area (backward from `free_end`); the slot pointer is inserted
736/// Overwrite the value portion of cell at `index` in-place. Valid when
737/// `new_value.len() <= existing_value.len()` — the caller
738/// (typically `BTree::upsert`) checks this before calling.
739///
740/// Same-length writes are trivial (byte copy). Shrinks update the cell's
741/// `v_len` header and leave the trailing bytes as dead space until the
742/// page is naturally rewritten; the slotted layout keeps cells
743/// addressable by their own offset+len so the gap is harmless.
744///
745/// Fixed-size column updates (SET score = X on a typed schema) hit this
746/// path, avoiding the delete + re-insert + rebalance cycle that
747/// dominates bulk_update.
748pub(crate) fn overwrite_leaf_value(
749    page: &mut Page,
750    index: usize,
751    new_value: &[u8],
752) -> BTreeResult<()> {
753    let cell_offset = leaf_read_slot(page, index)?;
754    let data = page.as_bytes();
755    if cell_offset + 4 > data.len() {
756        return Err(BTreeError::Corrupted(
757            "leaf cell header out of bounds".into(),
758        ));
759    }
760    let k_len = u16::from_le_bytes([data[cell_offset], data[cell_offset + 1]]) as usize;
761    let old_v_len = u16::from_le_bytes([data[cell_offset + 2], data[cell_offset + 3]]) as usize;
762    if new_value.len() > old_v_len {
763        return Err(BTreeError::Corrupted(
764            "overwrite_leaf_value: new value larger than slot".into(),
765        ));
766    }
767    let value_start = cell_offset + 4 + k_len;
768    if value_start + new_value.len() > data.len() {
769        return Err(BTreeError::Corrupted(
770            "leaf cell value out of bounds".into(),
771        ));
772    }
773    let data = page.as_bytes_mut();
774    // Update the v_len header if the value shrank. read_leaf_cell reads
775    // v_len from this header so the shortened value deserializes correctly.
776    if new_value.len() != old_v_len {
777        let new_len = new_value.len() as u16;
778        data[cell_offset + 2..cell_offset + 4].copy_from_slice(&new_len.to_le_bytes());
779    }
780    data[value_start..value_start + new_value.len()].copy_from_slice(new_value);
781    Ok(())
782}
783
784/// at the sorted position.
785fn insert_into_leaf(page: &mut Page, key: &[u8], value: &[u8]) -> BTreeResult<()> {
786    // 1. Binary search the slot array to find the insertion index.
787    let cell_count = page.cell_count() as usize;
788    let mut low = 0;
789    let mut high = cell_count;
790    while low < high {
791        let mid = (low + high) / 2;
792        let (cell_key, _) = read_leaf_cell(page, mid)?;
793        match cell_key.as_slice().cmp(key) {
794            Ordering::Less => low = mid + 1,
795            Ordering::Greater => high = mid,
796            Ordering::Equal => {
797                // Duplicate keys are tolerated by the B-tree (caller
798                // decides semantics); append after the existing run.
799                low = mid + 1;
800                break;
801            }
802        }
803    }
804    let insert_pos = low;
805
806    // 2. Reserve the cell at the tail of the free area.
807    let cell_size = 4 + key.len() + value.len();
808    let slot_end_after = LEAF_SLOT_ARRAY_OFFSET + (cell_count + 1) * 2;
809    let cells_start = leaf_cells_start(page);
810    if slot_end_after + cell_size > cells_start {
811        return Err(BTreeError::Corrupted("leaf page full".into()));
812    }
813    let cell_offset = cells_start - cell_size;
814    {
815        let data = page.as_bytes_mut();
816        data[cell_offset..cell_offset + 2].copy_from_slice(&(key.len() as u16).to_le_bytes());
817        data[cell_offset + 2..cell_offset + 4].copy_from_slice(&(value.len() as u16).to_le_bytes());
818        data[cell_offset + 4..cell_offset + 4 + key.len()].copy_from_slice(key);
819        data[cell_offset + 4 + key.len()..cell_offset + 4 + key.len() + value.len()]
820            .copy_from_slice(value);
821    }
822    page.set_free_end(cell_offset as u16);
823
824    // 3. Shift the slot-array tail right by one slot, then write the
825    //    new pointer into the freed slot. This is a single memmove on
826    //    a couple dozen u16s — far cheaper than the O(M²) rebuild.
827    {
828        let shift_from = leaf_slot_offset_for(insert_pos);
829        let shift_to = shift_from + 2;
830        let shift_len = (cell_count - insert_pos) * 2;
831        if shift_len > 0 {
832            let data = page.as_bytes_mut();
833            data.copy_within(shift_from..shift_from + shift_len, shift_to);
834        }
835    }
836    leaf_write_slot(page, insert_pos, cell_offset as u16)?;
837
838    // 4. Bump counters.
839    page.set_cell_count((cell_count + 1) as u16);
840    page.set_free_start((LEAF_SLOT_ARRAY_OFFSET + (cell_count + 1) * 2) as u16);
841    Ok(())
842}
843
844/// Remove the slot at `index`. Cell bytes are left in place (lazy
845/// compaction); the slot-array tail is shifted left to close the gap.
846/// The caller is expected to call `clear_leaf_cells` + rebuild if the
847/// page wants its free space reclaimed.
848fn delete_from_leaf(page: &mut Page, index: usize) -> BTreeResult<()> {
849    let cell_count = page.cell_count() as usize;
850    if index >= cell_count {
851        return Err(BTreeError::Corrupted("delete index out of range".into()));
852    }
853    if index + 1 < cell_count {
854        let shift_from = leaf_slot_offset_for(index + 1);
855        let shift_to = leaf_slot_offset_for(index);
856        let shift_len = (cell_count - index - 1) * 2;
857        let data = page.as_bytes_mut();
858        data.copy_within(shift_from..shift_from + shift_len, shift_to);
859    }
860    page.set_cell_count((cell_count - 1) as u16);
861    page.set_free_start((LEAF_SLOT_ARRAY_OFFSET + (cell_count - 1) * 2) as u16);
862    Ok(())
863}
864
865/// Reset the leaf to an empty slotted state — cell count 0, free
866/// cursors at the extremes. Zeroes the slot array + cell data area so
867/// stale bytes never leak through a corrupted read.
868fn clear_leaf_cells(page: &mut Page) {
869    {
870        let data = page.as_bytes_mut();
871        for byte in &mut data[LEAF_SLOT_ARRAY_OFFSET..] {
872            *byte = 0;
873        }
874    }
875    page.set_cell_count(0);
876    page.set_free_start(LEAF_SLOT_ARRAY_OFFSET as u16);
877    page.set_free_end(PAGE_SIZE as u16);
878}
879
880fn init_leaf_links(page: &mut Page, prev: u32, next: u32) {
881    let data = page.as_bytes_mut();
882    data[LEAF_PREV_OFFSET..LEAF_PREV_OFFSET + 4].copy_from_slice(&prev.to_le_bytes());
883    data[LEAF_NEXT_OFFSET..LEAF_NEXT_OFFSET + 4].copy_from_slice(&next.to_le_bytes());
884}
885
886fn read_next_leaf(page: &Page) -> u32 {
887    let data = page.as_bytes();
888    u32::from_le_bytes([
889        data[LEAF_NEXT_OFFSET],
890        data[LEAF_NEXT_OFFSET + 1],
891        data[LEAF_NEXT_OFFSET + 2],
892        data[LEAF_NEXT_OFFSET + 3],
893    ])
894}
895
896/// Returns the right-sibling page ID for a leaf page.
897/// 0 means this is the rightmost leaf — used by the D2 fastpath cache.
898#[inline]
899fn leaf_right_sibling(page: &Page) -> u32 {
900    read_next_leaf(page)
901}
902
903/// Returns the maximum (last) key stored in a leaf page, or `None` if empty.
904/// This is the Lehman-Yao "high_key" — any key > this value belongs on a
905/// right-sibling page. Reads only the last slot so it is O(1).
906fn leaf_high_key(page: &Page) -> BTreeResult<Option<Vec<u8>>> {
907    let cell_count = page.cell_count() as usize;
908    if cell_count == 0 {
909        return Ok(None);
910    }
911    let (key, _) = read_leaf_cell(page, cell_count - 1)?;
912    Ok(Some(key))
913}
914
915fn set_prev_leaf(page: &mut Page, prev: u32) {
916    let data = page.as_bytes_mut();
917    data[LEAF_PREV_OFFSET..LEAF_PREV_OFFSET + 4].copy_from_slice(&prev.to_le_bytes());
918}
919
920fn set_next_leaf(page: &mut Page, next: u32) {
921    let data = page.as_bytes_mut();
922    data[LEAF_NEXT_OFFSET..LEAF_NEXT_OFFSET + 4].copy_from_slice(&next.to_le_bytes());
923}
924
925// ==================== Interior Page Helpers ====================
926//
927// Slotted-page layout mirroring the leaf format:
928//
929// ```text
930// [PageHeader 32B]
931// [slot 0: u16][slot 1: u16]...[slot N-1: u16]  ← grows forward
932// ... free space ...
933// [cell N-1][cell N-2]...[cell 0]               ← grows backward from free_end
934// ```
935//
936// Each cell is `[key_len:u16][key][child:u32]`. The right-most child
937// pointer lives in `PageHeader.right_child` exactly as before.
938
939#[inline]
940fn interior_slot_offset_for(index: usize) -> usize {
941    INTERIOR_DATA_OFFSET + index * 2
942}
943
944#[inline]
945fn interior_read_slot(page: &Page, index: usize) -> BTreeResult<usize> {
946    let data = page.as_bytes();
947    let slot_pos = interior_slot_offset_for(index);
948    if slot_pos + 2 > PAGE_SIZE {
949        return Err(BTreeError::Corrupted(
950            "interior slot array overflows page".into(),
951        ));
952    }
953    Ok(u16::from_le_bytes([data[slot_pos], data[slot_pos + 1]]) as usize)
954}
955
956#[inline]
957fn interior_write_slot(page: &mut Page, index: usize, cell_offset: u16) -> BTreeResult<()> {
958    let data = page.as_bytes_mut();
959    let slot_pos = interior_slot_offset_for(index);
960    if slot_pos + 2 > PAGE_SIZE {
961        return Err(BTreeError::Corrupted(
962            "interior slot array overflows page".into(),
963        ));
964    }
965    data[slot_pos..slot_pos + 2].copy_from_slice(&cell_offset.to_le_bytes());
966    Ok(())
967}
968
969#[inline]
970fn interior_cells_start(page: &Page) -> usize {
971    let end = page.free_end() as usize;
972    if end == 0 {
973        PAGE_SIZE
974    } else {
975        end
976    }
977}
978
979#[inline]
980fn interior_free_bytes(page: &Page) -> usize {
981    let slot_end = INTERIOR_DATA_OFFSET + (page.cell_count() as usize) * 2;
982    interior_cells_start(page).saturating_sub(slot_end)
983}
984
985fn read_interior_cell(page: &Page, index: usize) -> BTreeResult<(Vec<u8>, u32)> {
986    let cell_count = page.cell_count() as usize;
987    if index >= cell_count {
988        return Err(BTreeError::Corrupted("Cell index out of range".into()));
989    }
990    let offset = interior_read_slot(page, index)?;
991    let data = page.as_bytes();
992    if offset + 2 > PAGE_SIZE {
993        return Err(BTreeError::Corrupted(
994            "interior cell header out of range".into(),
995        ));
996    }
997    let key_len = u16::from_le_bytes([data[offset], data[offset + 1]]) as usize;
998    let end = offset + 2 + key_len + 4;
999    if end > PAGE_SIZE {
1000        return Err(BTreeError::Corrupted(
1001            "interior cell body out of range".into(),
1002        ));
1003    }
1004    let key = data[offset + 2..offset + 2 + key_len].to_vec();
1005    let child = u32::from_le_bytes([
1006        data[offset + 2 + key_len],
1007        data[offset + 2 + key_len + 1],
1008        data[offset + 2 + key_len + 2],
1009        data[offset + 2 + key_len + 3],
1010    ]);
1011    Ok((key, child))
1012}
1013
1014fn read_interior_keys_children(page: &Page) -> BTreeResult<(Vec<Vec<u8>>, Vec<u32>)> {
1015    let cell_count = page.cell_count() as usize;
1016    let mut keys = Vec::with_capacity(cell_count);
1017    let mut children = Vec::with_capacity(cell_count + 1);
1018
1019    for i in 0..cell_count {
1020        let (key, child) = read_interior_cell(page, i)?;
1021        keys.push(key);
1022        children.push(child);
1023    }
1024
1025    if cell_count == 0 {
1026        let right_child = page.right_child();
1027        if right_child != 0 {
1028            children.push(right_child);
1029        }
1030    } else {
1031        children.push(page.right_child());
1032    }
1033
1034    Ok((keys, children))
1035}
1036
1037/// Bulk-write an interior page from scratch in slotted form. Used by
1038/// the split path. `keys` and `children` follow the B+ tree contract:
1039/// `children.len() == keys.len() + 1`, with the last child landing in
1040/// `PageHeader.right_child` (as before).
1041fn write_interior_entries(page: &mut Page, keys: &[Vec<u8>], children: &[u32]) -> BTreeResult<()> {
1042    if !keys.is_empty() && children.len() != keys.len() + 1 {
1043        return Err(BTreeError::Corrupted(
1044            "Interior keys/children length mismatch".into(),
1045        ));
1046    }
1047
1048    clear_interior_cells(page);
1049    if keys.is_empty() {
1050        let right_child = children.first().copied().unwrap_or(0);
1051        page.set_right_child(right_child);
1052        return Ok(());
1053    }
1054
1055    for (i, key) in keys.iter().enumerate() {
1056        let cell_size = 2 + key.len() + 4;
1057        let cells_start = interior_cells_start(page);
1058        let slot_end = INTERIOR_DATA_OFFSET + (i + 1) * 2;
1059        if slot_end + cell_size > cells_start {
1060            return Err(BTreeError::Corrupted(
1061                "interior rebuild overflowed page".into(),
1062            ));
1063        }
1064        let cell_offset = cells_start - cell_size;
1065        {
1066            let data = page.as_bytes_mut();
1067            data[cell_offset..cell_offset + 2].copy_from_slice(&(key.len() as u16).to_le_bytes());
1068            data[cell_offset + 2..cell_offset + 2 + key.len()].copy_from_slice(key);
1069            data[cell_offset + 2 + key.len()..cell_offset + 2 + key.len() + 4]
1070                .copy_from_slice(&children[i].to_le_bytes());
1071        }
1072        page.set_free_end(cell_offset as u16);
1073        interior_write_slot(page, i, cell_offset as u16)?;
1074    }
1075    page.set_cell_count(keys.len() as u16);
1076    page.set_free_start((INTERIOR_DATA_OFFSET + keys.len() * 2) as u16);
1077    page.set_right_child(*children.last().ok_or_else(|| {
1078        BTreeError::Corrupted("write_interior_entries: children empty with non-empty keys".into())
1079    })?);
1080    Ok(())
1081}
1082
1083fn can_insert_interior(page: &Page, key: &[u8]) -> bool {
1084    let needed = 2 + 2 + key.len() + 4;
1085    interior_free_bytes(page) >= needed
1086}
1087
1088/// Insert a `(key, left_child)` separator into the interior page.
1089/// `right_child` replaces whatever child used to sit to the right of
1090/// the inserted key (either a middle child's pointer or the page's
1091/// `right_child` when the key is the new maximum).
1092fn insert_into_interior(
1093    page: &mut Page,
1094    key: &[u8],
1095    left_child: u32,
1096    right_child: u32,
1097) -> BTreeResult<()> {
1098    // 1. Binary search the slot array for the insert position.
1099    let cell_count = page.cell_count() as usize;
1100    let mut low = 0;
1101    let mut high = cell_count;
1102    while low < high {
1103        let mid = (low + high) / 2;
1104        let (cell_key, _) = read_interior_cell(page, mid)?;
1105        match cell_key.as_slice().cmp(key) {
1106            Ordering::Less => low = mid + 1,
1107            Ordering::Greater | Ordering::Equal => high = mid,
1108        }
1109    }
1110    let insert_pos = low;
1111
1112    // 2. Redirect the previous owner of the split to `right_child`.
1113    //    If the insertion is at the tail, the previous owner was the
1114    //    page's right_child pointer; otherwise it was the child slot
1115    //    of the cell currently at `insert_pos`.
1116    if insert_pos < cell_count {
1117        let old_offset = interior_read_slot(page, insert_pos)?;
1118        let data = page.as_bytes();
1119        let key_len = u16::from_le_bytes([data[old_offset], data[old_offset + 1]]) as usize;
1120        let child_pos = old_offset + 2 + key_len;
1121        let data = page.as_bytes_mut();
1122        data[child_pos..child_pos + 4].copy_from_slice(&right_child.to_le_bytes());
1123    } else {
1124        page.set_right_child(right_child);
1125    }
1126
1127    // 3. Reserve the new cell at the tail of the free area.
1128    let cell_size = 2 + key.len() + 4;
1129    let slot_end_after = INTERIOR_DATA_OFFSET + (cell_count + 1) * 2;
1130    let cells_start = interior_cells_start(page);
1131    if slot_end_after + cell_size > cells_start {
1132        return Err(BTreeError::Corrupted("interior page full".into()));
1133    }
1134    let cell_offset = cells_start - cell_size;
1135    {
1136        let data = page.as_bytes_mut();
1137        data[cell_offset..cell_offset + 2].copy_from_slice(&(key.len() as u16).to_le_bytes());
1138        data[cell_offset + 2..cell_offset + 2 + key.len()].copy_from_slice(key);
1139        data[cell_offset + 2 + key.len()..cell_offset + 2 + key.len() + 4]
1140            .copy_from_slice(&left_child.to_le_bytes());
1141    }
1142    page.set_free_end(cell_offset as u16);
1143
1144    // 4. Shift the slot-array tail right by one slot, write the new
1145    //    pointer into the freed slot, bump counters.
1146    {
1147        let shift_from = interior_slot_offset_for(insert_pos);
1148        let shift_to = shift_from + 2;
1149        let shift_len = (cell_count - insert_pos) * 2;
1150        if shift_len > 0 {
1151            let data = page.as_bytes_mut();
1152            data.copy_within(shift_from..shift_from + shift_len, shift_to);
1153        }
1154    }
1155    interior_write_slot(page, insert_pos, cell_offset as u16)?;
1156    page.set_cell_count((cell_count + 1) as u16);
1157    page.set_free_start((INTERIOR_DATA_OFFSET + (cell_count + 1) * 2) as u16);
1158    Ok(())
1159}
1160
1161fn clear_interior_cells(page: &mut Page) {
1162    {
1163        let data = page.as_bytes_mut();
1164        for byte in &mut data[INTERIOR_DATA_OFFSET..] {
1165            *byte = 0;
1166        }
1167    }
1168    page.set_cell_count(0);
1169    page.set_free_start(INTERIOR_DATA_OFFSET as u16);
1170    page.set_free_end(PAGE_SIZE as u16);
1171}