obj_core/btree/

insert.rs

//! B+tree insert path with copy-on-write splits.
//!
//! See `docs/format.md` § B+tree write semantics. Every node touched
//! along the insert path is rewritten to a freshly-allocated page;
//! the displaced pages enter the freelist only after the new root is
//! staged in the pager's WAL transaction buffer. The caller is
//! responsible for calling `Pager::commit` to make the insert
//! durable.
//!
//! # Power-of-ten posture
//!
//! - **Rule 1.** The walk root → leaf uses an explicit
//!   `heapless::Vec<_, MAX_BTREE_DEPTH>` path stack. The bubble-up
//!   loop is a `while let Some(...)` over that stack — no
//!   recursion.
//! - **Rule 2.** Every loop is bounded: by the path length (≤
//!   `MAX_BTREE_DEPTH`), by `key_count` (a node's slot count), or
//!   by `LEAF_SLOT_CAP` / `INTERNAL_SLOT_CAP`.
//! - **Rule 5.** `validate_node` runs on every encoded node behind
//!   a `debug_assert!`; release builds surface invariant breaks as
//!   `Error::BTreeInvariantViolated`.

#![forbid(unsafe_code)]

use heapless::Vec as HeaplessVec;

use crate::btree::node::{
    decode_node, encode_node, max_inline_value, max_key_len, DecodedNode, InternalEntry, LeafEntry,
    NodeKind, INTERNAL_LEFTMOST_CHILD_BYTES, INTERNAL_SLOT_BYTES, LEAF_SLOT_BYTES,
};
use crate::btree::{BTree, MAX_BTREE_DEPTH};
use crate::error::{Error, Result};
use crate::pager::page::{Page, PageId, PAGE_SIZE, PAGE_TRAILER_SIZE};
use crate::pager::Pager;
use crate::platform::FileBackend;

/// One step in a descending path: the node's id, its decoded
/// representation, and the index of the child we followed.
struct PathFrame {
    page_id: PageId,
    node: DecodedNode,
    /// For internal-node frames: the child index we recursed into.
    /// Unused for the leaf frame at the bottom of the stack.
    child_index: usize,
}

/// Outcome of replacing one node along the insert path: either the
/// node still fits in a single page (we just COW it), or it had to
/// split into two siblings with a promoted pivot key.
enum ReplaceOutcome {
    /// Node fits in one page; new page-id replaces the old one in
    /// its parent's child slot.
    Fits { new_id: PageId },
    /// Node split. `left_id` is the new left half (same role as the
    /// pre-insert node), `right_id` is the brand-new right sibling,
    /// `promoted_key` is the smallest key of the right half (for
    /// leaves) or the pivot moved up (for internals).
    Split {
        left_id: PageId,
        right_id: PageId,
        promoted_key: Vec<u8>,
    },
}

/// Total bytes available for the slot directory + heap in a single
/// node (excludes the node header and the page trailer).
const PAYLOAD_BYTES: usize = PAGE_SIZE - PAGE_TRAILER_SIZE - crate::btree::node::NODE_HEADER_SIZE;

impl<F: FileBackend> BTree<F> {
    /// Insert `key → value` into the tree. Returns
    /// [`Error::BTreeKeyExists`] if `key` is already present.
    ///
    /// Copy-on-write: every node touched on the path is allocated as
    /// a fresh page through the pager. The pre-insert pages enter the
    /// freelist before this function returns, but only after every
    /// new page has been staged in the WAL transaction. The caller
    /// must call [`Pager::commit`] to make the insert durable.
    ///
    /// # Errors
    ///
    /// - [`Error::BTreeKeyTooLarge`] if the key exceeds `PAGE_SIZE / 4`.
    /// - [`Error::BTreeValueTooLarge`] if the (key, value) pair will
    ///   not fit inline in a leaf.
    /// - [`Error::BTreeKeyExists`] if `key` is already present.
    /// - [`Error::BTreeDepthExceeded`] if the tree height would
    ///   exceed `MAX_BTREE_DEPTH`.
    /// - [`Error::Corruption`] / [`Error::Io`] propagated from the
    ///   pager.
    pub fn insert(&mut self, pager: &mut Pager<F>, key: &[u8], value: &[u8]) -> Result<()> {
        check_key_value_size(key, value)?;
        let path = self.descend_with_path(pager, key)?;
        self.apply_insert(pager, path, key, value)
    }

    /// Walk root → leaf, recording each ancestor's page-id, decoded
    /// representation, and the child index that was followed.
    fn descend_with_path(
        &self,
        pager: &mut Pager<F>,
        key: &[u8],
    ) -> Result<HeaplessVec<PathFrame, MAX_BTREE_DEPTH>> {
        let mut path: HeaplessVec<PathFrame, MAX_BTREE_DEPTH> = HeaplessVec::new();
        let mut current = self.root;
        loop {
            let decoded = {
                let page_ref = pager.read_page(current)?;
                decode_node(page_ref.as_bytes())?
            };
            match decoded.kind {
                NodeKind::Leaf => {
                    let frame = PathFrame {
                        page_id: current,
                        node: decoded,
                        child_index: 0,
                    };
                    if path.push(frame).is_err() {
                        return Err(Error::BTreeDepthExceeded {
                            limit: MAX_BTREE_DEPTH,
                        });
                    }
                    return Ok(path);
                }
                NodeKind::Internal => {
                    let child_index = pivot_index(&decoded, key);
                    let raw = decoded.children[child_index];
                    let next = PageId::new(raw).ok_or(Error::BTreeInvariantViolated {
                        reason: "internal node had zero child page-id",
                    })?;
                    let frame = PathFrame {
                        page_id: current,
                        node: decoded,
                        child_index,
                    };
                    if path.push(frame).is_err() {
                        return Err(Error::BTreeDepthExceeded {
                            limit: MAX_BTREE_DEPTH,
                        });
                    }
                    current = next;
                }
            }
        }
    }

    /// Apply the (key, value) insertion to the leaf at the bottom of
    /// the path, then bubble any split up to the root.
    fn apply_insert(
        &mut self,
        pager: &mut Pager<F>,
        mut path: HeaplessVec<PathFrame, MAX_BTREE_DEPTH>,
        key: &[u8],
        value: &[u8],
    ) -> Result<()> {
        let mut freed: HeaplessVec<PageId, { MAX_BTREE_DEPTH * 2 }> = HeaplessVec::new();
        let Some(leaf_frame) = path.pop() else {
            return Err(Error::BTreeInvariantViolated {
                reason: "insert: descend returned empty path",
            });
        };
        let mut outcome = replace_leaf(pager, leaf_frame, key, value, &mut freed)?;
        while let Some(parent_frame) = path.pop() {
            outcome = replace_internal(pager, parent_frame, outcome, &mut freed)?;
        }
        let new_root = build_new_root(pager, outcome)?;
        self.root = new_root;
        // Free displaced pages AFTER every new page is staged in the
        // WAL transaction buffer. This is the COW contract: pre-
        // insert pages must remain readable through the previous
        // root until the new root is committed.
        for old_id in freed.iter().copied() {
            pager.free_page(old_id)?;
        }
        Ok(())
    }
}

fn replace_leaf<F: FileBackend>(
    pager: &mut Pager<F>,
    frame: PathFrame,
    key: &[u8],
    value: &[u8],
    freed: &mut HeaplessVec<PageId, { MAX_BTREE_DEPTH * 2 }>,
) -> Result<ReplaceOutcome> {
    let mut leaf = frame.node;
    if leaf.leaves.iter().any(|e| e.key.as_slice() == key) {
        return Err(Error::BTreeKeyExists);
    }
    let insert_at = leaf
        .leaves
        .iter()
        .position(|e| e.key.as_slice() > key)
        .unwrap_or(leaf.leaves.len());
    leaf.leaves.insert(
        insert_at,
        LeafEntry {
            key: key.to_vec(),
            value: value.to_vec(),
        },
    );
    push_freed(freed, frame.page_id)?;
    if leaf.occupied_bytes() <= PAYLOAD_BYTES {
        let new_id = write_new_node(pager, &leaf)?;
        return Ok(ReplaceOutcome::Fits { new_id });
    }
    split_leaf(pager, leaf)
}

fn replace_internal<F: FileBackend>(
    pager: &mut Pager<F>,
    frame: PathFrame,
    child_outcome: ReplaceOutcome,
    freed: &mut HeaplessVec<PageId, { MAX_BTREE_DEPTH * 2 }>,
) -> Result<ReplaceOutcome> {
    let mut internal = frame.node;
    let idx = frame.child_index;
    match child_outcome {
        ReplaceOutcome::Fits { new_id } => {
            internal.children[idx] = new_id.get();
        }
        ReplaceOutcome::Split {
            left_id,
            right_id,
            promoted_key,
        } => {
            internal.children[idx] = left_id.get();
            internal
                .internals
                .insert(idx, InternalEntry { key: promoted_key });
            internal.children.insert(idx + 1, right_id.get());
        }
    }
    push_freed(freed, frame.page_id)?;
    if internal.occupied_bytes() <= PAYLOAD_BYTES {
        let new_id = write_new_node(pager, &internal)?;
        return Ok(ReplaceOutcome::Fits { new_id });
    }
    split_internal(pager, internal)
}

/// Build the new root from the outcome of `apply_insert`'s final
/// bubble. If no split, the root is just the new id; if a split, we
/// allocate a fresh internal node holding (left, promoted, right).
fn build_new_root<F: FileBackend>(pager: &mut Pager<F>, outcome: ReplaceOutcome) -> Result<PageId> {
    let (left_id, right_id, promoted_key) = match outcome {
        ReplaceOutcome::Fits { new_id } => return Ok(new_id),
        ReplaceOutcome::Split {
            left_id,
            right_id,
            promoted_key,
        } => (left_id, right_id, promoted_key),
    };
    let level = node_level_after_split(pager, left_id)?;
    let next_level = level.checked_add(1).ok_or(Error::BTreeDepthExceeded {
        limit: MAX_BTREE_DEPTH,
    })?;
    let root_node = DecodedNode {
        kind: NodeKind::Internal,
        level: next_level,
        next_sibling: 0,
        children: vec![left_id.get(), right_id.get()],
        leaves: Vec::new(),
        internals: vec![InternalEntry { key: promoted_key }],
    };
    write_new_node(pager, &root_node)
}

fn push_freed(freed: &mut HeaplessVec<PageId, { MAX_BTREE_DEPTH * 2 }>, id: PageId) -> Result<()> {
    freed.push(id).map_err(|_| Error::BTreeInvariantViolated {
        reason: "insert: too many displaced pages to track",
    })
}

/// Pick the child index in `node` whose subtree contains `key`.
fn pivot_index(node: &DecodedNode, key: &[u8]) -> usize {
    let mut idx = node.internals.len();
    for (i, pivot) in node.internals.iter().enumerate() {
        if pivot.key.as_slice() > key {
            idx = i;
            break;
        }
    }
    idx
}

/// Validate that `key` / `value` fit the format-version-0 bounds.
fn check_key_value_size(key: &[u8], value: &[u8]) -> Result<()> {
    if key.len() > max_key_len() {
        return Err(Error::BTreeKeyTooLarge {
            key_len: key.len(),
            max: max_key_len(),
        });
    }
    let v_max = max_inline_value(key.len());
    if value.len() > v_max {
        return Err(Error::BTreeValueTooLarge {
            value_len: value.len(),
            max: v_max,
        });
    }
    Ok(())
}

/// Encode `node` into a fresh page and write it through the pager.
/// Returns the newly-allocated `PageId`.
pub(crate) fn write_new_node<F: FileBackend>(
    pager: &mut Pager<F>,
    node: &DecodedNode,
) -> Result<PageId> {
    let new_id = pager.alloc_page()?;
    let mut page = Page::zeroed();
    encode_node(node, &mut page)?;
    pager.write_page(new_id, &page)?;
    Ok(new_id)
}

/// Split an overflowing leaf into two. Allocates new ids for both
/// halves; the right half's `next_sibling` is inherited, the left
/// half points at the new right.
///
/// Takes the leaf by value because the function relocates the
/// internal `Vec`s into two new nodes (no allocation of a clone).
fn split_leaf<F: FileBackend>(
    pager: &mut Pager<F>,
    mut leaf: DecodedNode,
) -> Result<ReplaceOutcome> {
    let mid = leaf.leaves.len() / 2;
    let original_sibling = leaf.next_sibling;
    // Split off the right half by draining; the left half stays in
    // `leaf.leaves`. `Vec::split_off` returns a new `Vec` and
    // truncates the receiver — no clone, no extra alloc.
    let right_entries: Vec<LeafEntry> = leaf.leaves.split_off(mid);
    let promoted_key = right_entries[0].key.clone();
    let right_node = DecodedNode {
        kind: NodeKind::Leaf,
        level: 0,
        next_sibling: original_sibling,
        children: Vec::new(),
        leaves: right_entries,
        internals: Vec::new(),
    };
    let right_id = write_new_node(pager, &right_node)?;
    let left_node = DecodedNode {
        kind: NodeKind::Leaf,
        level: 0,
        next_sibling: right_id.get(),
        children: Vec::new(),
        leaves: leaf.leaves,
        internals: Vec::new(),
    };
    let left_id = write_new_node(pager, &left_node)?;
    Ok(ReplaceOutcome::Split {
        left_id,
        right_id,
        promoted_key,
    })
}

/// Split an overflowing internal node. The pivot at the split point
/// moves *up* (not duplicated, unlike leaves); each half keeps
/// `mid` / `K - mid - 1` pivots respectively.
fn split_internal<F: FileBackend>(
    pager: &mut Pager<F>,
    mut internal: DecodedNode,
) -> Result<ReplaceOutcome> {
    let k = internal.internals.len();
    debug_assert!(k >= 2, "internal split needs ≥ 2 pivots");
    let mid = k / 2;
    // Right half: pivots [mid+1..K) and children [mid+1..K+1).
    // The pivot at `mid` is promoted (not retained in either half).
    let right_pivots: Vec<InternalEntry> = internal.internals.split_off(mid + 1);
    let right_children: Vec<u64> = internal.children.split_off(mid + 1);
    // Remove the promoted pivot from the left half (now at the end of
    // the truncated `internals` vec).
    let promoted_pivot = internal
        .internals
        .pop()
        .ok_or(Error::BTreeInvariantViolated {
            reason: "internal split: missing promoted pivot",
        })?;
    let level = internal.level;
    let right_node = DecodedNode {
        kind: NodeKind::Internal,
        level,
        next_sibling: 0,
        children: right_children,
        leaves: Vec::new(),
        internals: right_pivots,
    };
    let right_id = write_new_node(pager, &right_node)?;
    let left_node = DecodedNode {
        kind: NodeKind::Internal,
        level,
        next_sibling: 0,
        children: internal.children,
        leaves: Vec::new(),
        internals: internal.internals,
    };
    let left_id = write_new_node(pager, &left_node)?;
    Ok(ReplaceOutcome::Split {
        left_id,
        right_id,
        promoted_key: promoted_pivot.key,
    })
}

/// Determine the level of a node that just emerged from a split. We
/// read the left half's page back through the pager (cheap — it was
/// just staged in the WAL view) to learn its level.
fn node_level_after_split<F: FileBackend>(pager: &mut Pager<F>, id: PageId) -> Result<u8> {
    let page_ref = pager.read_page(id)?;
    let decoded = decode_node(page_ref.as_bytes())?;
    Ok(decoded.level)
}

// `_PAYLOAD_BYTES` reference to keep the linker happy if the const
// is not otherwise visible to insert callers; used in
// `occupied_bytes() <= PAYLOAD_BYTES`.
const _: usize = PAYLOAD_BYTES;

/// Force-link the unused private slot/header byte constants so any
/// future refactor that drops them gets a hard error here rather
/// than silently changing the encode layout.
const _UNUSED_CHECKS: () = {
    let _ = INTERNAL_LEFTMOST_CHILD_BYTES;
    let _ = INTERNAL_SLOT_BYTES;
    let _ = LEAF_SLOT_BYTES;
};

#[cfg(test)]
mod tests {
    use super::*;
    use crate::pager::{Config, Pager};
    use crate::platform::FileHandle;

    use proptest::prelude::*;
    use rand::prelude::IndexedRandom;
    use rand::SeedableRng;
    use rand_chacha::ChaCha8Rng;
    use std::collections::BTreeMap;

    fn config() -> Config {
        Config::default()
    }

    #[test]
    fn insert_single_key_round_trip() {
        let mut pager = Pager::<FileHandle>::memory(config()).expect("pager");
        let mut tree = BTree::<FileHandle>::empty(&mut pager).expect("empty");
        tree.insert(&mut pager, b"hello", b"world").expect("ins");
        assert_eq!(
            tree.get(&mut pager, b"hello").expect("get"),
            Some(b"world".to_vec())
        );
    }

    #[test]
    fn duplicate_key_errors() {
        let mut pager = Pager::<FileHandle>::memory(config()).expect("pager");
        let mut tree = BTree::<FileHandle>::empty(&mut pager).expect("empty");
        tree.insert(&mut pager, b"k", b"v1").expect("ins");
        let err = tree
            .insert(&mut pager, b"k", b"v2")
            .expect_err("dup must fail");
        assert!(matches!(err, Error::BTreeKeyExists));
    }

    #[test]
    fn insert_growth_splits_root() {
        let mut pager = Pager::<FileHandle>::memory(config()).expect("pager");
        let mut tree = BTree::<FileHandle>::empty(&mut pager).expect("empty");
        // 256-byte values guarantee splits after ~15 inserts.
        let value = vec![0xABu8; 256];
        for i in 0..200u32 {
            let key = format!("key-{i:08}");
            tree.insert(&mut pager, key.as_bytes(), &value)
                .expect("ins");
        }
        for i in 0..200u32 {
            let key = format!("key-{i:08}");
            assert_eq!(
                tree.get(&mut pager, key.as_bytes()).expect("get"),
                Some(value.clone()),
                "key {key}"
            );
        }
        // Tree should have at least 2 levels now.
        let root = tree.root();
        let page_ref = pager.read_page(root).expect("read root");
        let decoded = decode_node(page_ref.as_bytes()).expect("decode root");
        assert!(
            decoded.level >= 1,
            "expected internal root, got {decoded:?}"
        );
    }

    proptest! {
        #![proptest_config(ProptestConfig {
            cases: 16,
            max_shrink_iters: 32,
            .. ProptestConfig::default()
        })]

        #[test]
        fn insert_oracle_property(seed in any::<u64>()) {
            run_insert_oracle(seed, 200);
        }
    }

    /// Run 10k random insert operations against a `BTreeMap` oracle.
    /// This is the in-module sanity check; the full 1M-op oracle
    /// lives in `tests/btree_oracle.rs` (issue #29).
    #[test]
    fn insert_oracle_10k() {
        for seed in 0..3u64 {
            run_insert_oracle(seed, 10_000);
        }
    }

    fn run_insert_oracle(seed: u64, ops: usize) {
        let mut rng = ChaCha8Rng::seed_from_u64(seed);
        let mut pager = Pager::<FileHandle>::memory(config()).expect("pager");
        let mut tree = BTree::<FileHandle>::empty(&mut pager).expect("empty");
        let mut oracle: BTreeMap<Vec<u8>, Vec<u8>> = BTreeMap::new();
        for op in 0..ops {
            let key = random_key(&mut rng);
            let value = random_value(&mut rng);
            let key_already = oracle.contains_key(&key);
            let res = tree.insert(&mut pager, &key, &value);
            if key_already {
                assert!(
                    matches!(res, Err(Error::BTreeKeyExists)),
                    "seed {seed} op {op}: expected BTreeKeyExists, got {res:?}"
                );
            } else {
                res.unwrap_or_else(|e| panic!("seed {seed} op {op}: insert err {e:?}"));
                oracle.insert(key.clone(), value.clone());
            }
            if op.is_multiple_of(127) {
                // Spot-check a few random oracle keys.
                let keys: Vec<&Vec<u8>> = oracle.keys().collect();
                if !keys.is_empty() {
                    let sample: Vec<&Vec<u8>> =
                        keys.choose_multiple(&mut rng, 4).copied().collect();
                    for k in sample {
                        assert_eq!(
                            tree.get(&mut pager, k).expect("get").as_ref(),
                            oracle.get(k),
                            "seed {seed} op {op}: key {k:?}"
                        );
                    }
                }
            }
        }
        // Final pass: every oracle key resolves correctly.
        for (k, v) in &oracle {
            assert_eq!(
                tree.get(&mut pager, k).expect("get").as_ref(),
                Some(v),
                "seed {seed} final: key {k:?}"
            );
        }
    }

    fn random_key(rng: &mut ChaCha8Rng) -> Vec<u8> {
        use rand::Rng;
        let len = rng.random_range(1..16);
        (0..len).map(|_| rng.random_range(b'a'..=b'z')).collect()
    }

    fn random_value(rng: &mut ChaCha8Rng) -> Vec<u8> {
        use rand::Rng;
        let len = rng.random_range(0..64);
        (0..len).map(|_| rng.random()).collect()
    }
}