sql_rs/storage/

btree.rs

use crate::{Result, SqlRsError};
use std::path::Path;
use std::sync::Arc;
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};

use super::file::FileManager;
use super::page::{Page, PageId, PageType, PageCache};

const MAX_KEYS_PER_NODE: usize = 64;

#[derive(Debug, Serialize, Deserialize)]
struct BTreeNode {
    page_id: PageId,
    is_leaf: bool,
    keys: Vec<Vec<u8>>,
    values: Vec<Vec<u8>>,
    children: Vec<PageId>,
}

impl BTreeNode {
    fn new(page_id: PageId, is_leaf: bool) -> Self {
        Self {
            page_id,
            is_leaf,
            keys: Vec::new(),
            values: Vec::new(),
            children: Vec::new(),
        }
    }
    
    fn serialize(&self) -> Vec<u8> {
        bincode::serialize(self).unwrap_or_default()
    }
    
    fn deserialize(data: &[u8]) -> Option<Self> {
        bincode::deserialize(data).ok()
    }
}

pub struct BTree {
    file_manager: Arc<RwLock<FileManager>>,
    cache: RwLock<PageCache>,
    root_page_id: PageId,
}

impl BTree {
    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
        let file_manager = FileManager::open(path)?;

        if file_manager.page_count() == 0 {
            return Err(SqlRsError::Storage("Empty database file".to_string()));
        }

        Ok(Self {
            file_manager: Arc::new(RwLock::new(file_manager)),
            cache: RwLock::new(PageCache::new(100)),
            root_page_id: 0,
        })
    }

    pub fn create<P: AsRef<Path>>(path: P) -> Result<Self> {
        let mut file_manager = FileManager::create(path)?;

        let root_page_id = file_manager.allocate_page();
        let mut root_page = Page::new(root_page_id, PageType::BTreeLeaf);
        root_page.write_header();

        let root_node = BTreeNode::new(root_page_id, true);
        let serialized = root_node.serialize();
        root_page.data[1..serialized.len() + 1].copy_from_slice(&serialized);

        file_manager.write_page(&root_page)?;
        file_manager.flush()?;

        Ok(Self {
            file_manager: Arc::new(RwLock::new(file_manager)),
            cache: RwLock::new(PageCache::new(100)),
            root_page_id,
        })
    }
    
    fn read_node(&self, page_id: PageId) -> Result<BTreeNode> {
        {
            let cache = self.cache.read();
            if let Some(page) = cache.get(page_id) {
                if let Some(node) = BTreeNode::deserialize(&page.data[1..]) {
                    return Ok(node);
                }
            }
        }
        
        let mut file_manager = self.file_manager.write();
        let page = file_manager.read_page(page_id)?;
        
        let node = BTreeNode::deserialize(&page.data[1..])
            .ok_or_else(|| SqlRsError::Storage("Failed to deserialize node".to_string()))?;
        
        self.cache.write().insert(page);
        
        Ok(node)
    }
    
    fn write_node(&self, node: &BTreeNode) -> Result<()> {
        let mut page = Page::new(node.page_id, if node.is_leaf { PageType::BTreeLeaf } else { PageType::BTreeInternal });
        page.write_header();
        
        let serialized = node.serialize();
        if serialized.len() + 1 > page.data.len() {
            return Err(SqlRsError::Storage("Node too large for page".to_string()));
        }
        
        page.data[1..serialized.len() + 1].copy_from_slice(&serialized);
        
        self.cache.write().insert(page.clone());
        self.file_manager.write().write_page(&page)?;
        
        Ok(())
    }
    
    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        self.search_recursive(self.root_page_id, key)
    }
    
    fn search_recursive(&self, page_id: PageId, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let node = self.read_node(page_id)?;
        
        if node.is_leaf {
            for (i, k) in node.keys.iter().enumerate() {
                if k.as_slice() == key {
                    return Ok(Some(node.values[i].clone()));
                }
            }
            Ok(None)
        } else {
            let mut child_idx = 0;
            for (i, k) in node.keys.iter().enumerate() {
                if key < k.as_slice() {
                    child_idx = i;
                    break;
                }
                child_idx = i + 1;
            }
            
            if child_idx < node.children.len() {
                self.search_recursive(node.children[child_idx], key)
            } else {
                Ok(None)
            }
        }
    }
    
    pub fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<()> {
        let mut root = self.read_node(self.root_page_id)?;
        
        if root.keys.len() >= MAX_KEYS_PER_NODE {
            let new_root_id = self.file_manager.write().allocate_page();
            let mut new_root = BTreeNode::new(new_root_id, false);
            new_root.children.push(self.root_page_id);
            
            self.split_child(&mut new_root, 0)?;
            self.root_page_id = new_root_id;
            self.write_node(&new_root)?;
            
            root = new_root;
        }
        
        self.insert_non_full(root, key, value)?;
        Ok(())
    }
    
    fn insert_non_full(&self, mut node: BTreeNode, key: &[u8], value: &[u8]) -> Result<()> {
        if node.is_leaf {
            let pos = node.keys.iter().position(|k| k.as_slice() >= key).unwrap_or(node.keys.len());
            
            if pos < node.keys.len() && node.keys[pos] == key {
                node.values[pos] = value.to_vec();
            } else {
                node.keys.insert(pos, key.to_vec());
                node.values.insert(pos, value.to_vec());
            }
            
            self.write_node(&node)?;
        } else {
            let mut pos = node.keys.iter().position(|k| k.as_slice() > key).unwrap_or(node.keys.len());
            
            let child_id = node.children[pos];
            let child = self.read_node(child_id)?;
            
            if child.keys.len() >= MAX_KEYS_PER_NODE {
                self.split_child(&mut node, pos)?;
                
                if key > node.keys[pos].as_slice() {
                    pos += 1;
                }
            }
            
            let child_id = node.children[pos];
            let child = self.read_node(child_id)?;
            self.insert_non_full(child, key, value)?;
        }
        
        Ok(())
    }
    
    fn split_child(&self, parent: &mut BTreeNode, child_idx: usize) -> Result<()> {
        let child_id = parent.children[child_idx];
        let mut child = self.read_node(child_id)?;
        
        let mid = child.keys.len() / 2;
        
        let new_node_id = self.file_manager.write().allocate_page();
        let mut new_node = BTreeNode::new(new_node_id, child.is_leaf);
        
        if child.is_leaf {
            new_node.keys = child.keys.split_off(mid);
            new_node.values = child.values.split_off(mid);
            
            let mid_key = new_node.keys[0].clone();
            parent.keys.insert(child_idx, mid_key);
        } else {
            new_node.keys = child.keys.split_off(mid + 1);
            new_node.values = child.values.split_off(mid + 1);
            new_node.children = child.children.split_off(mid + 1);
            
            let mid_key = child.keys.pop().unwrap();
            child.values.pop();
            
            parent.keys.insert(child_idx, mid_key);
        }
        
        parent.children.insert(child_idx + 1, new_node_id);
        
        self.write_node(&child)?;
        self.write_node(&new_node)?;
        self.write_node(parent)?;
        
        Ok(())
    }
    
    pub fn delete(&mut self, key: &[u8]) -> Result<()> {
        let mut root = self.read_node(self.root_page_id)?;
        self.delete_recursive(&mut root, key)?;

        // If root has no keys but has a child, make the child the new root
        if !root.is_leaf && root.keys.is_empty() {
            if let Some(&new_root_id) = root.children.first() {
                self.root_page_id = new_root_id;
            }
        }

        Ok(())
    }

    fn delete_recursive(&self, node: &mut BTreeNode, key: &[u8]) -> Result<bool> {
        let idx = node.keys.iter().position(|k| k.as_slice() >= key).unwrap_or(node.keys.len());

        if node.is_leaf {
            if idx < node.keys.len() && node.keys[idx] == key {
                node.keys.remove(idx);
                node.values.remove(idx);
                self.write_node(node)?;
                return Ok(true);
            }
            return Ok(false);
        }

        // Key is in internal node
        if idx < node.keys.len() && node.keys[idx] == key {
            return self.delete_from_internal_node(node, idx);
        }

        // Key is in child
        let child_id = node.children[idx];
        let mut child = self.read_node(child_id)?;

        // Ensure child has enough keys
        if child.keys.len() <= MAX_KEYS_PER_NODE / 2 {
            self.ensure_min_keys(node, idx)?;
            // Reload child and potentially new node index
            let child_id = node.children[idx];
            let mut child = self.read_node(child_id)?;
            return self.delete_recursive(&mut child, key);
        }

        self.delete_recursive(&mut child, key)
    }

    fn delete_from_internal_node(&self, node: &mut BTreeNode, idx: usize) -> Result<bool> {
        let predecessor_child_id = node.children[idx];
        let predecessor_child = self.read_node(predecessor_child_id)?;

        if predecessor_child.keys.len() > MAX_KEYS_PER_NODE / 2 {
            // Get predecessor (rightmost key of left child)
            let pred_key = self.get_predecessor(predecessor_child_id)?;

            node.keys[idx] = pred_key.clone();
            self.write_node(node)?;

            let mut child = self.read_node(predecessor_child_id)?;
            return self.delete_recursive(&mut child, &pred_key);
        }

        let successor_child_id = node.children[idx + 1];
        let successor_child = self.read_node(successor_child_id)?;

        if successor_child.keys.len() > MAX_KEYS_PER_NODE / 2 {
            // Get successor (leftmost key of right child)
            let succ_key = self.get_successor(successor_child_id)?;

            node.keys[idx] = succ_key.clone();
            self.write_node(node)?;

            let mut child = self.read_node(successor_child_id)?;
            return self.delete_recursive(&mut child, &succ_key);
        }

        // Merge both children and delete
        self.merge_children(node, idx)?;
        let merged_child_id = node.children[idx];
        let mut merged_child = self.read_node(merged_child_id)?;
        let key_to_delete = node.keys[idx].clone();
        self.delete_recursive(&mut merged_child, &key_to_delete)
    }

    fn get_predecessor(&self, mut page_id: PageId) -> Result<Vec<u8>> {
        loop {
            let node = self.read_node(page_id)?;
            if node.is_leaf {
                return Ok(node.keys.last().cloned().unwrap());
            }
            page_id = *node.children.last().unwrap();
        }
    }

    fn get_successor(&self, mut page_id: PageId) -> Result<Vec<u8>> {
        loop {
            let node = self.read_node(page_id)?;
            if node.is_leaf {
                return Ok(node.keys.first().cloned().unwrap());
            }
            page_id = node.children[0];
        }
    }

    fn ensure_min_keys(&self, parent: &mut BTreeNode, idx: usize) -> Result<()> {
        let child_id = parent.children[idx];
        let child = self.read_node(child_id)?;

        if child.keys.len() > MAX_KEYS_PER_NODE / 2 {
            return Ok(());
        }

        // Try to borrow from left sibling
        if idx > 0 {
            let left_sibling_id = parent.children[idx - 1];
            let left_sibling = self.read_node(left_sibling_id)?;

            if left_sibling.keys.len() > MAX_KEYS_PER_NODE / 2 {
                return self.borrow_from_left(parent, idx);
            }
        }

        // Try to borrow from right sibling
        if idx < parent.children.len() - 1 {
            let right_sibling_id = parent.children[idx + 1];
            let right_sibling = self.read_node(right_sibling_id)?;

            if right_sibling.keys.len() > MAX_KEYS_PER_NODE / 2 {
                return self.borrow_from_right(parent, idx);
            }
        }

        // Merge with a sibling
        if idx > 0 {
            self.merge_children(parent, idx - 1)?;
        } else {
            self.merge_children(parent, idx)?;
        }

        Ok(())
    }

    fn borrow_from_left(&self, parent: &mut BTreeNode, idx: usize) -> Result<()> {
        let child_id = parent.children[idx];
        let left_sibling_id = parent.children[idx - 1];

        let mut child = self.read_node(child_id)?;
        let mut left_sibling = self.read_node(left_sibling_id)?;

        // Move key from parent to child
        child.keys.insert(0, parent.keys[idx - 1].clone());
        if child.is_leaf {
            child.values.insert(0, left_sibling.values.pop().unwrap());
        } else {
            child.children.insert(0, left_sibling.children.pop().unwrap());
            child.values.insert(0, left_sibling.values.pop().unwrap());
        }

        // Move key from left sibling to parent
        parent.keys[idx - 1] = left_sibling.keys.pop().unwrap();

        self.write_node(&child)?;
        self.write_node(&left_sibling)?;
        self.write_node(parent)?;

        Ok(())
    }

    fn borrow_from_right(&self, parent: &mut BTreeNode, idx: usize) -> Result<()> {
        let child_id = parent.children[idx];
        let right_sibling_id = parent.children[idx + 1];

        let mut child = self.read_node(child_id)?;
        let mut right_sibling = self.read_node(right_sibling_id)?;

        // Move key from parent to child
        child.keys.push(parent.keys[idx].clone());
        if child.is_leaf {
            child.values.push(right_sibling.values.remove(0));
        } else {
            child.children.push(right_sibling.children.remove(0));
            child.values.push(right_sibling.values.remove(0));
        }

        // Move key from right sibling to parent
        parent.keys[idx] = right_sibling.keys.remove(0);

        self.write_node(&child)?;
        self.write_node(&right_sibling)?;
        self.write_node(parent)?;

        Ok(())
    }

    fn merge_children(&self, parent: &mut BTreeNode, idx: usize) -> Result<()> {
        let left_child_id = parent.children[idx];
        let right_child_id = parent.children[idx + 1];

        let mut left_child = self.read_node(left_child_id)?;
        let right_child = self.read_node(right_child_id)?;

        // Move key from parent to left child
        left_child.keys.push(parent.keys.remove(idx));
        parent.children.remove(idx + 1);

        // Merge keys, values, and children
        left_child.keys.extend(right_child.keys);
        left_child.values.extend(right_child.values);
        if !left_child.is_leaf {
            left_child.children.extend(right_child.children);
        }

        self.write_node(&left_child)?;
        self.write_node(parent)?;

        Ok(())
    }
    
    pub fn scan(&self, start: &[u8], end: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
        let mut results = Vec::new();
        self.scan_recursive(self.root_page_id, start, end, &mut results)?;
        Ok(results)
    }
    
    fn scan_recursive(&self, page_id: PageId, start: &[u8], end: &[u8], results: &mut Vec<(Vec<u8>, Vec<u8>)>) -> Result<()> {
        let node = self.read_node(page_id)?;
        
        if node.is_leaf {
            for (i, key) in node.keys.iter().enumerate() {
                if key.as_slice() >= start && key.as_slice() < end {
                    results.push((key.clone(), node.values[i].clone()));
                }
            }
        } else {
            for (i, child_id) in node.children.iter().enumerate() {
                if i < node.keys.len() {
                    if node.keys[i].as_slice() >= start {
                        self.scan_recursive(*child_id, start, end, results)?;
                    }
                } else {
                    self.scan_recursive(*child_id, start, end, results)?;
                }
            }
        }
        
        Ok(())
    }
    
    pub fn flush(&mut self) -> Result<()> {
        self.cache.write().clear();
        self.file_manager.write().flush()
    }

    pub fn get_file_manager(&self) -> Arc<RwLock<FileManager>> {
        Arc::clone(&self.file_manager)
    }
}