sochdb-storage 2.0.4

// SPDX-License-Identifier: AGPL-3.0-or-later
// SochDB - LLM-Optimized Embedded Database
// Copyright (C) 2026 Sushanth Reddy Vanagala (https://github.com/sushanthpy)
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Copy-on-Write B-Tree Index (Recommendation 5)
//!
//! ## Problem
//!
//! DashMap provides O(1) average but:
//! - Hash computation: ~30ns
//! - Bucket traversal: ~15ns (collision resolution)
//! - Cache miss on data access: ~100ns
//! - Worst case (many collisions): ~230ns
//!
//! ## Solution
//!
//! Copy-on-Write B-tree with:
//! - O(log N) lookups with ~3-4 comparisons for 1M rows
//! - Cache-friendly page layout (256 entries per node)
//! - Predictable worst-case performance
//! - Persistent structure for crash recovery
//!
//! ## Performance Analysis
//!
//! B-tree depth for N keys, branching factor B:
//! ```text
//! depth = ceil(log_B(N))
//! ```
//!
//! For N = 1,000,000 and B = 256:
//! ```text
//! depth = ceil(log_256(1,000,000)) = ceil(2.5) = 3 levels
//! ```
//!
//! Per-lookup cost:
//! ```text
//! T_btree = depth × T_node_search
//!         = 3 × (T_cache_miss + T_binary_search)
//!         = 3 × (100ns + 50ns) = 450ns
//! ```
//!
//! B-tree provides consistent 450ns vs DashMap's 45-500ns variance.

use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering as AtomicOrdering};

use parking_lot::RwLock;

/// Default node size (256 entries fits in 4KB page with some metadata)
pub const DEFAULT_NODE_SIZE: usize = 256;

/// Minimum node size (for small trees)
pub const MIN_NODE_SIZE: usize = 4;

/// Maximum tree depth (safety limit)
pub const MAX_DEPTH: usize = 32;

// =============================================================================
// B-Tree Node
// =============================================================================

/// A key-value pair in the B-tree
#[derive(Debug, Clone)]
pub struct BTreeEntry<K, V> {
    pub key: K,
    pub value: V,
}

impl<K, V> BTreeEntry<K, V> {
    pub fn new(key: K, value: V) -> Self {
        Self { key, value }
    }
}

/// Interior node: keys and child pointers
#[derive(Debug, Clone)]
pub struct InteriorNode<K, V> {
    /// Keys (n-1 keys for n children)
    keys: Vec<K>,
    /// Child node pointers
    children: Vec<Arc<Node<K, V>>>,
}

/// Leaf node: key-value pairs
#[derive(Debug, Clone)]
pub struct LeafNode<K, V> {
    /// Key-value entries (sorted by key)
    entries: Vec<BTreeEntry<K, V>>,
    /// Optional next leaf pointer for range scans
    next: Option<Arc<Node<K, V>>>,
}

/// B-tree node (either interior or leaf)
#[derive(Debug, Clone)]
pub enum Node<K, V> {
    Interior(InteriorNode<K, V>),
    Leaf(LeafNode<K, V>),
}

impl<K: Clone + Ord, V: Clone> Node<K, V> {
    /// Create new empty leaf node
    pub fn new_leaf() -> Self {
        Node::Leaf(LeafNode {
            entries: Vec::new(),
            next: None,
        })
    }

    /// Create interior node from children
    pub fn new_interior(keys: Vec<K>, children: Vec<Arc<Node<K, V>>>) -> Self {
        Node::Interior(InteriorNode { keys, children })
    }

    /// Check if this is a leaf node
    pub fn is_leaf(&self) -> bool {
        matches!(self, Node::Leaf(_))
    }

    /// Get number of entries (for leaf) or children (for interior)
    pub fn len(&self) -> usize {
        match self {
            Node::Interior(n) => n.children.len(),
            Node::Leaf(n) => n.entries.len(),
        }
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Check if node is full
    pub fn is_full(&self, max_size: usize) -> bool {
        self.len() >= max_size
    }

    /// Search for key in leaf node - O(log n) binary search
    pub fn search(&self, key: &K) -> SearchResult<K, V> {
        match self {
            Node::Leaf(leaf) => match leaf.entries.binary_search_by(|e| e.key.cmp(key)) {
                Ok(idx) => SearchResult::Found(leaf.entries[idx].value.clone()),
                Err(idx) => SearchResult::NotFound(idx),
            },
            Node::Interior(interior) => {
                // Find child to descend into
                // For keys [k1, k2, ..., kn] and children [c0, c1, ..., cn]:
                // c_i contains keys in range [k_{i-1}, k_i) where k_0 = -∞
                // Use <= to ensure keys equal to separator go to right subtree
                let idx = interior.keys.partition_point(|k| k <= key);
                SearchResult::Child(idx, interior.children[idx].clone())
            }
        }
    }

    /// Get entry at index (for leaf only)
    pub fn get_entry(&self, idx: usize) -> Option<&BTreeEntry<K, V>> {
        match self {
            Node::Leaf(leaf) => leaf.entries.get(idx),
            _ => None,
        }
    }
}

/// Result of searching a node
pub enum SearchResult<K, V> {
    /// Key found with value
    Found(V),
    /// Key not found, would be at index
    NotFound(usize),
    /// Need to search child at index
    Child(usize, Arc<Node<K, V>>),
}

// =============================================================================
// Copy-on-Write B-Tree
// =============================================================================

/// Copy-on-Write B-tree for ordered key-value storage
///
/// ## Structural Sharing
///
/// When modifying the tree, only the path from root to the modified leaf
/// is copied. All other nodes are shared between versions.
///
/// ```text
/// Before insert:           After insert (key X):
///       [A]                      [A']  (new root)
///      /   \                    /    \
///    [B]   [C]               [B']    [C]  (B copied, C shared)
///   /  \   /  \             /  \    /  \
/// [D] [E] [F] [G]         [D] [E'] [F] [G]  (E copied, others shared)
/// ```
///
/// This enables:
/// - Lock-free reads (immutable snapshots)
/// - O(log N) modifications
/// - Crash recovery (old version preserved until commit)
pub struct CowBTree<K, V> {
    /// Root node (atomically swapped on updates)
    root: RwLock<Arc<Node<K, V>>>,
    /// Node size (branching factor)
    node_size: usize,
    /// Current version (for MVCC)
    version: AtomicU64,
    /// Entry count
    count: AtomicU64,
}

impl<K: Clone + Ord + std::fmt::Debug, V: Clone + std::fmt::Debug> CowBTree<K, V> {
    /// Create new empty B-tree
    pub fn new() -> Self {
        Self::with_node_size(DEFAULT_NODE_SIZE)
    }

    /// Create with specific node size
    pub fn with_node_size(node_size: usize) -> Self {
        let node_size = node_size.max(MIN_NODE_SIZE);
        Self {
            root: RwLock::new(Arc::new(Node::new_leaf())),
            node_size,
            version: AtomicU64::new(0),
            count: AtomicU64::new(0),
        }
    }

    /// Get value for key - O(log N)
    pub fn get(&self, key: &K) -> Option<V> {
        let root = self.root.read().clone();
        self.search_node(&root, key)
    }

    /// Search recursively through nodes
    fn search_node(&self, node: &Arc<Node<K, V>>, key: &K) -> Option<V> {
        match node.search(key) {
            SearchResult::Found(value) => Some(value),
            SearchResult::NotFound(_) => None,
            SearchResult::Child(_, child) => self.search_node(&child, key),
        }
    }

    /// Insert key-value pair - O(log N)
    ///
    /// Returns previous value if key existed
    pub fn insert(&self, key: K, value: V) -> Option<V> {
        let mut root = self.root.write();
        let (new_root, old_value, split) = self.insert_recursive(&root, key.clone(), value);

        if let Some((split_key, right_child)) = split {
            // Root was split, create new root
            let new_interior = Node::new_interior(vec![split_key], vec![new_root, right_child]);
            *root = Arc::new(new_interior);
        } else {
            *root = new_root;
        }

        if old_value.is_none() {
            self.count.fetch_add(1, AtomicOrdering::Relaxed);
        }
        self.version.fetch_add(1, AtomicOrdering::Relaxed);

        old_value
    }

    /// Recursive insert with copy-on-write
    /// Returns: (new_node, old_value, optional_split)
    fn insert_recursive(
        &self,
        node: &Arc<Node<K, V>>,
        key: K,
        value: V,
    ) -> (Arc<Node<K, V>>, Option<V>, Option<(K, Arc<Node<K, V>>)>) {
        match node.as_ref() {
            Node::Leaf(leaf) => {
                let mut new_entries = leaf.entries.clone();

                // Binary search for position
                match new_entries.binary_search_by(|e| e.key.cmp(&key)) {
                    Ok(idx) => {
                        // Key exists, update value
                        let old_value = new_entries[idx].value.clone();
                        new_entries[idx].value = value;
                        (
                            Arc::new(Node::Leaf(LeafNode {
                                entries: new_entries,
                                next: leaf.next.clone(),
                            })),
                            Some(old_value),
                            None,
                        )
                    }
                    Err(idx) => {
                        // Insert new entry
                        new_entries.insert(idx, BTreeEntry::new(key, value));

                        // Check if split needed
                        if new_entries.len() > self.node_size {
                            let mid = new_entries.len() / 2;
                            let split_key = new_entries[mid].key.clone();

                            let right_entries = new_entries.split_off(mid);
                            let right_node = Arc::new(Node::Leaf(LeafNode {
                                entries: right_entries,
                                next: leaf.next.clone(),
                            }));

                            let left_node = Arc::new(Node::Leaf(LeafNode {
                                entries: new_entries,
                                next: Some(right_node.clone()),
                            }));

                            (left_node, None, Some((split_key, right_node)))
                        } else {
                            (
                                Arc::new(Node::Leaf(LeafNode {
                                    entries: new_entries,
                                    next: leaf.next.clone(),
                                })),
                                None,
                                None,
                            )
                        }
                    }
                }
            }
            Node::Interior(interior) => {
                // Find child to insert into
                // Use <= to be consistent with search
                let idx = interior.keys.partition_point(|k| k <= &key);

                let (new_child, old_value, child_split) =
                    self.insert_recursive(&interior.children[idx], key, value);

                let mut new_keys = interior.keys.clone();
                let mut new_children = interior.children.clone();
                new_children[idx] = new_child;

                if let Some((split_key, right_child)) = child_split {
                    // Child was split, insert new key and child
                    new_keys.insert(idx, split_key);
                    new_children.insert(idx + 1, right_child);

                    // Check if this node needs to split
                    if new_children.len() > self.node_size {
                        let mid = new_keys.len() / 2;
                        let up_key = new_keys[mid].clone();

                        let right_keys = new_keys.split_off(mid + 1);
                        new_keys.pop(); // Remove the key that goes up

                        let right_children = new_children.split_off(mid + 1);

                        let right_node = Arc::new(Node::Interior(InteriorNode {
                            keys: right_keys,
                            children: right_children,
                        }));

                        let left_node = Arc::new(Node::Interior(InteriorNode {
                            keys: new_keys,
                            children: new_children,
                        }));

                        (left_node, old_value, Some((up_key, right_node)))
                    } else {
                        (
                            Arc::new(Node::Interior(InteriorNode {
                                keys: new_keys,
                                children: new_children,
                            })),
                            old_value,
                            None,
                        )
                    }
                } else {
                    (
                        Arc::new(Node::Interior(InteriorNode {
                            keys: new_keys,
                            children: new_children,
                        })),
                        old_value,
                        None,
                    )
                }
            }
        }
    }

    /// Check if key exists - O(log N)
    pub fn contains_key(&self, key: &K) -> bool {
        self.get(key).is_some()
    }

    /// Get entry count
    pub fn len(&self) -> usize {
        self.count.load(AtomicOrdering::Relaxed) as usize
    }

    /// Check if empty
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Get current version
    pub fn version(&self) -> u64 {
        self.version.load(AtomicOrdering::Relaxed)
    }

    /// Get tree depth
    pub fn depth(&self) -> usize {
        let root = self.root.read().clone();
        self.node_depth(&root)
    }

    fn node_depth(&self, node: &Arc<Node<K, V>>) -> usize {
        match node.as_ref() {
            Node::Leaf(_) => 1,
            Node::Interior(interior) => {
                1 + interior
                    .children
                    .first()
                    .map(|c| self.node_depth(c))
                    .unwrap_or(0)
            }
        }
    }

    /// Range scan - O(log N + K) where K is result count
    pub fn range(&self, start: &K, end: &K) -> Vec<(K, V)> {
        let root = self.root.read().clone();
        let mut results = Vec::new();
        self.range_search(&root, start, end, &mut results);
        results
    }

    fn range_search(&self, node: &Arc<Node<K, V>>, start: &K, end: &K, results: &mut Vec<(K, V)>) {
        match node.as_ref() {
            Node::Leaf(leaf) => {
                for entry in &leaf.entries {
                    if &entry.key >= start && &entry.key < end {
                        results.push((entry.key.clone(), entry.value.clone()));
                    } else if &entry.key >= end {
                        break;
                    }
                }
                // Follow next pointer for linked leaves
                if let Some(ref next) = leaf.next {
                    if let Some(last_entry) = leaf.entries.last() {
                        if &last_entry.key < end {
                            self.range_search(next, start, end, results);
                        }
                    }
                }
            }
            Node::Interior(interior) => {
                // Find first child that might contain start
                let start_idx = interior.keys.partition_point(|k| k < start);
                // Find last child that might contain end
                let end_idx = interior.keys.partition_point(|k| k < end);

                for idx in start_idx..=end_idx.min(interior.children.len() - 1) {
                    self.range_search(&interior.children[idx], start, end, results);
                }
            }
        }
    }

    /// Get snapshot for read-only access
    pub fn snapshot(&self) -> BTreeSnapshot<K, V> {
        BTreeSnapshot {
            root: self.root.read().clone(),
            version: self.version.load(AtomicOrdering::Relaxed),
        }
    }
}

impl<K: Clone + Ord + std::fmt::Debug, V: Clone + std::fmt::Debug> Default for CowBTree<K, V> {
    fn default() -> Self {
        Self::new()
    }
}

/// Read-only snapshot of B-tree
#[derive(Clone)]
pub struct BTreeSnapshot<K, V> {
    root: Arc<Node<K, V>>,
    version: u64,
}

impl<K: Clone + Ord, V: Clone> BTreeSnapshot<K, V> {
    /// Get value for key
    pub fn get(&self, key: &K) -> Option<V> {
        self.search_node(&self.root, key)
    }

    fn search_node(&self, node: &Arc<Node<K, V>>, key: &K) -> Option<V> {
        match node.search(key) {
            SearchResult::Found(value) => Some(value),
            SearchResult::NotFound(_) => None,
            SearchResult::Child(_, child) => self.search_node(&child, key),
        }
    }

    /// Get snapshot version
    pub fn version(&self) -> u64 {
        self.version
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_btree_insert_get() {
        let tree: CowBTree<i32, String> = CowBTree::new();

        tree.insert(5, "five".to_string());
        tree.insert(3, "three".to_string());
        tree.insert(7, "seven".to_string());

        assert_eq!(tree.get(&5), Some("five".to_string()));
        assert_eq!(tree.get(&3), Some("three".to_string()));
        assert_eq!(tree.get(&7), Some("seven".to_string()));
        assert_eq!(tree.get(&1), None);
    }

    #[test]
    fn test_btree_update() {
        let tree: CowBTree<i32, String> = CowBTree::new();

        tree.insert(1, "one".to_string());
        assert_eq!(tree.insert(1, "ONE".to_string()), Some("one".to_string()));
        assert_eq!(tree.get(&1), Some("ONE".to_string()));
    }

    #[test]
    fn test_btree_many_inserts() {
        let tree: CowBTree<i32, i32> = CowBTree::with_node_size(4);

        for i in 0..1000 {
            tree.insert(i, i * 10);
        }

        assert_eq!(tree.len(), 1000);

        for i in 0..1000 {
            assert_eq!(tree.get(&i), Some(i * 10));
        }
    }

    #[test]
    fn test_btree_range() {
        let tree: CowBTree<i32, i32> = CowBTree::new();

        for i in 0..100 {
            tree.insert(i, i);
        }

        let range = tree.range(&10, &20);
        assert_eq!(range.len(), 10);
        assert_eq!(range[0], (10, 10));
        assert_eq!(range[9], (19, 19));
    }

    #[test]
    fn test_btree_snapshot() {
        let tree: CowBTree<i32, String> = CowBTree::new();

        tree.insert(1, "one".to_string());
        let snap1 = tree.snapshot();

        tree.insert(2, "two".to_string());
        let snap2 = tree.snapshot();

        // snap1 should not see key 2
        assert_eq!(snap1.get(&1), Some("one".to_string()));
        // Note: In a full MVCC implementation, snap1 wouldn't see key 2
        // but our simple CoW only tracks version numbers

        // snap2 should see both
        assert_eq!(snap2.get(&1), Some("one".to_string()));
        assert_eq!(snap2.get(&2), Some("two".to_string()));
    }

    #[test]
    fn test_btree_depth() {
        let tree: CowBTree<i32, i32> = CowBTree::with_node_size(4);

        // Empty tree has depth 1 (just root leaf)
        assert_eq!(tree.depth(), 1);

        // Insert enough to cause splits
        for i in 0..100 {
            tree.insert(i, i);
        }

        // With node size 4 and 100 entries, depth should be reasonable
        let depth = tree.depth();
        assert!(depth >= 2 && depth <= 6, "Unexpected depth: {}", depth);
    }

    #[test]
    fn test_btree_reverse_order() {
        let tree: CowBTree<i32, i32> = CowBTree::with_node_size(4);

        // Insert in reverse order
        for i in (0..100).rev() {
            tree.insert(i, i);
        }

        assert_eq!(tree.len(), 100);

        // Should still be searchable
        for i in 0..100 {
            assert_eq!(tree.get(&i), Some(i));
        }
    }

    #[test]
    fn test_btree_string_keys() {
        let tree: CowBTree<String, i32> = CowBTree::new();

        tree.insert("apple".to_string(), 1);
        tree.insert("banana".to_string(), 2);
        tree.insert("cherry".to_string(), 3);

        assert_eq!(tree.get(&"banana".to_string()), Some(2));

        let range = tree.range(&"a".to_string(), &"c".to_string());
        assert_eq!(range.len(), 2); // apple, banana
    }
}