ChuBlog/chusql-storage/src/btree.rs

use std::io;
use std::path::Path;

use crate::page::{Page, PageFile, PageId};

// 磁盘 B+ 树：节点占一页，文件头记页大小与 order。

const META_PAGE: PageId = 0;
const META_MAGIC: &[u8; 4] = b"CBTR";
const OFF_META_PAGE_SIZE: usize = 4;
const OFF_META_ORDER: usize = 8;
const OFF_META_ROOT: usize = 10;

const NODE_LEAF: u8 = 0;
const NODE_INTERNAL: u8 = 1;
const OFF_TYPE: usize = 0;
const OFF_COUNT: usize = 1;
const OFF_NEXT: usize = 3;
const OFF_KEYS: usize = 12;

/// 一个节点占多少字节
fn node_bytes(order: usize) -> usize {
    OFF_KEYS + max_keys(order) * 8 + order * 8
}

/// 最多几个键
fn max_keys(order: usize) -> usize {
    order - 1
}

/// 值区域的起点
fn off_payload(order: usize) -> usize {
    OFF_KEYS + max_keys(order) * 8
}

/// 这个 order 放得进一页吗
pub fn fits_in_page(page_size: usize, order: usize) -> bool {
    order >= 2 && node_bytes(order) <= page_size
}

/// 一页最多几个孩子
pub fn max_order(page_size: usize) -> usize {
    page_size.saturating_sub(4) / 16
}

/// 读一个 u8
fn read_u8(p: &Page, off: usize) -> u8 {
    p.data[off]
}

/// 写一个 u8
fn write_u8(p: &mut Page, off: usize, v: u8) {
    p.data[off] = v;
}

/// 读一个小端 u16
fn read_u16(p: &Page, off: usize) -> u16 {
    u16::from_le_bytes([p.data[off], p.data[off + 1]])
}

/// 写一个小端 u16
fn write_u16(p: &mut Page, off: usize, v: u16) {
    p.data[off..off + 2].copy_from_slice(&v.to_le_bytes());
}

/// 读一个小端 u32
fn read_u32(p: &Page, off: usize) -> u32 {
    u32::from_le_bytes([p.data[off], p.data[off + 1], p.data[off + 2], p.data[off + 3]])
}

/// 写一个小端 u32
fn write_u32(p: &mut Page, off: usize, v: u32) {
    p.data[off..off + 4].copy_from_slice(&v.to_le_bytes());
}

/// 读一个小端 u64
fn read_u64(p: &Page, off: usize) -> u64 {
    let mut b = [0u8; 8];
    b.copy_from_slice(&p.data[off..off + 8]);
    u64::from_le_bytes(b)
}

/// 写一个小端 u64
fn write_u64(p: &mut Page, off: usize, v: u64) {
    p.data[off..off + 8].copy_from_slice(&v.to_le_bytes());
}

/// 写文件头
fn write_meta(page: &mut Page, page_size: usize, order: usize, root: PageId) {
    page.zero();
    page.data[0..4].copy_from_slice(META_MAGIC);
    write_u32(page, OFF_META_PAGE_SIZE, page_size as u32);
    write_u16(page, OFF_META_ORDER, order as u16);
    write_u64(page, OFF_META_ROOT, root);
}

/// 校验文件头与配置
fn check_meta(page: &Page, page_size: usize, order: usize) -> io::Result<()> {
    if page.data[0..4] != *META_MAGIC {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            "not a chusql index file (bad magic)",
        ));
    }
    let stored_page = read_u32(page, OFF_META_PAGE_SIZE) as usize;
    let stored_order = read_u16(page, OFF_META_ORDER) as usize;
    if stored_page != page_size || stored_order != order {
        return Err(io::Error::new(
            io::ErrorKind::InvalidData,
            format!(
                "index file was built with page.size={} btree.order={}, config says page.size={} btree.order={}",
                stored_page, stored_order, page_size, order
            ),
        ));
    }
    Ok(())
}

enum NodeData {
    Leaf {
        keys: Vec<i64>,
        values: Vec<u64>,
        next: PageId,
    },
    Internal {
        keys: Vec<i64>,
        children: Vec<PageId>,
    },
}

/// 从一页读出节点
fn read_node(p: &Page, order: usize) -> NodeData {
    let ty = read_u8(p, OFF_TYPE);
    let n = read_u16(p, OFF_COUNT) as usize;
    let payload = off_payload(order);
    if ty == NODE_LEAF {
        let next = read_u64(p, OFF_NEXT);
        let mut keys = Vec::with_capacity(n);
        let mut values = Vec::with_capacity(n);
        for i in 0..n {
            keys.push(read_u64(p, OFF_KEYS + i * 8) as i64);
            values.push(read_u64(p, payload + i * 8));
        }
        NodeData::Leaf { keys, values, next }
    } else {
        let mut keys = Vec::with_capacity(n);
        let mut children = Vec::with_capacity(n + 1);
        for i in 0..n {
            keys.push(read_u64(p, OFF_KEYS + i * 8) as i64);
        }
        for i in 0..(n + 1) {
            children.push(read_u64(p, payload + i * 8));
        }
        NodeData::Internal { keys, children }
    }
}

/// 把节点写回一页
fn write_node(p: &mut Page, node: &NodeData, order: usize) {
    p.zero();
    let payload = off_payload(order);
    match node {
        NodeData::Leaf { keys, values, next } => {
            write_u8(p, OFF_TYPE, NODE_LEAF);
            write_u16(p, OFF_COUNT, keys.len() as u16);
            write_u64(p, OFF_NEXT, *next);
            for (i, k) in keys.iter().enumerate() {
                write_u64(p, OFF_KEYS + i * 8, *k as u64);
            }
            for (i, v) in values.iter().enumerate() {
                write_u64(p, payload + i * 8, *v);
            }
        }
        NodeData::Internal { keys, children } => {
            write_u8(p, OFF_TYPE, NODE_INTERNAL);
            write_u16(p, OFF_COUNT, keys.len() as u16);
            for (i, k) in keys.iter().enumerate() {
                write_u64(p, OFF_KEYS + i * 8, *k as u64);
            }
            for (i, c) in children.iter().enumerate() {
                write_u64(p, payload + i * 8, *c);
            }
        }
    }
}

/// 二分选孩子
fn child_index(keys: &[i64], key: i64) -> usize {
    let mut lo = 0;
    let mut hi = keys.len();
    while lo < hi {
        let mid = (lo + hi) / 2;
        if keys[mid] <= key {
            lo = mid + 1;
        } else {
            hi = mid;
        }
    }
    lo
}

pub struct DiskBTree {
    file: PageFile,
    order: usize,
}

impl DiskBTree {
    /// 打开或新建索引文件
    pub fn open<P: AsRef<Path>>(
            path: P,
            page_size: usize,
            order: usize,
            pool_size: usize,
        ) -> io::Result<Self> {
        if !fits_in_page(page_size, order) {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!(
                    "btree order {} does not fit in a {}-byte page (max {})",
                    order,
                    page_size,
                    max_order(page_size)
                ),
            ));
        }
        let mut file = PageFile::with_options(path, page_size, pool_size)?;
        if file.num_pages()? == 0 {
            let mut meta = Page::new(META_PAGE, page_size);
            write_meta(&mut meta, page_size, order, 0);
            file.write_page(&meta)?;
        } else {
            let meta = file.read_page(META_PAGE)?;
            check_meta(&meta, page_size, order)?;
        }
        Ok(DiskBTree { file, order })
    }

    /// 读根页号
    fn root(&mut self) -> io::Result<Option<PageId>> {
        let r = self.file.with_page(META_PAGE, |p| read_u64(p, OFF_META_ROOT))?;
        Ok(if r == 0 { None } else { Some(r) })
    }

    /// 写根页号
    fn set_root(&mut self, r: Option<PageId>) -> io::Result<()> {
        self.file
            .update_page(META_PAGE, |p| write_u64(p, OFF_META_ROOT, r.unwrap_or(0)))
    }

    /// 分配新页
    fn alloc(&mut self) -> io::Result<PageId> {
        Ok(self.file.append_page()?.id)
    }

    /// 清空整棵树
    pub fn clear(&mut self) -> io::Result<()> {
        let page_size = self.file.page_size();
        self.file.truncate()?;
        let mut meta = Page::new(META_PAGE, page_size);
        write_meta(&mut meta, page_size, self.order, 0);
        self.file.write_page(&meta)?;
        Ok(())
    }

    /// 空页，用来拼节点
    fn new_page(&self, id: PageId) -> Page {
        Page::new(id, self.file.page_size())
    }

    /// 按 key 找 value
    pub fn get(&mut self, key: i64) -> io::Result<Option<u64>> {
        match self.root()? {
            None => Ok(None),
            Some(r) => self.get_rec(r, key),
        }
    }

    /// 递归查找
    fn get_rec(&mut self, page_id: PageId, key: i64) -> io::Result<Option<u64>> {
        let order = self.order;
        match self.file.with_page(page_id, |p| read_node(p, order))? {
            NodeData::Leaf { keys, values, .. } => {
                Ok(keys.binary_search(&key).ok().map(|i| values[i]))
            }
            NodeData::Internal { keys, children } => {
                let i = child_index(&keys, key);
                self.get_rec(children[i], key)
            }
        }
    }

    /// 插入或覆盖
    pub fn insert(&mut self, key: i64, value: u64) -> io::Result<()> {
        match self.root()? {
            None => {
                let id = self.alloc()?;
                let node = NodeData::Leaf {
                    keys: vec![key],
                    values: vec![value],
                    next: 0,
                };
                let mut p = self.new_page(id);
                write_node(&mut p, &node, self.order);
                self.file.write_page(&p)?;
                self.set_root(Some(id))
            }
            Some(r) => {
                let (_, split) = self.insert_rec(r, key, value)?;
                if let Some((sep, right_id)) = split {
                    let root_id = self.alloc()?;
                    let node = NodeData::Internal {
                        keys: vec![sep],
                        children: vec![r, right_id],
                    };
                    let mut p = self.new_page(root_id);
                    write_node(&mut p, &node, self.order);
                    self.file.write_page(&p)?;
                    self.set_root(Some(root_id))
                } else {
                    Ok(())
                }
            }
        }
    }

    /// 递归插入，可能上提
    fn insert_rec(
        &mut self,
        page_id: PageId,
        key: i64,
        value: u64,
    ) -> io::Result<(PageId, Option<(i64, PageId)>)> {
        let order = self.order;
        let node = self.file.with_page(page_id, |p| read_node(p, order))?;
        let limit = max_keys(order);
        match node {
            NodeData::Leaf { mut keys, mut values, next } => {
                match keys.binary_search(&key) {
                    Ok(i) => values[i] = value,
                    Err(i) => {
                        keys.insert(i, key);
                        values.insert(i, value);
                    }
                }
                if keys.len() <= limit {
                    self.file.update_page(page_id, |p| {
                        write_node(p, &NodeData::Leaf { keys, values, next }, order)
                    })?;
                    Ok((page_id, None))
                } else {
                    let mid = keys.len() / 2;
                    let rkeys = keys.split_off(mid);
                    let rvalues = values.split_off(mid);
                    let sep = rkeys[0];

                    let right_id = self.alloc()?;
                    let mut rp = self.new_page(right_id);
                    write_node(
                        &mut rp,
                        &NodeData::Leaf { keys: rkeys, values: rvalues, next },
                        order,
                    );
                    self.file.write_page(&rp)?;

                    self.file.update_page(page_id, |p| {
                        write_node(p, &NodeData::Leaf { keys, values, next: right_id }, order)
                    })?;

                    Ok((page_id, Some((sep, right_id))))
                }
            }

            NodeData::Internal { mut keys, mut children } => {
                let i = child_index(&keys, key);
                let (new_child, split) = self.insert_rec(children[i], key, value)?;
                children[i] = new_child;
                match split {
                    None => {
                        self.file.update_page(page_id, |p| {
                            write_node(p, &NodeData::Internal { keys, children }, order)
                        })?;
                        Ok((page_id, None))
                    }
                    Some((sep, right_id)) => {
                        keys.insert(i, sep);
                        children.insert(i + 1, right_id);
                        if keys.len() <= limit {
                            self.file.update_page(page_id, |p| {
                                write_node(p, &NodeData::Internal { keys, children }, order)
                            })?;
                            Ok((page_id, None))
                        } else {
                            let mid = keys.len() / 2;
                            let up = keys[mid];
                            let rkeys = keys.split_off(mid + 1);
                            keys.truncate(mid);
                            let rchildren = children.split_off(mid + 1);

                            let right_id = self.alloc()?;
                            let mut rp = self.new_page(right_id);
                            write_node(
                                &mut rp,
                                &NodeData::Internal { keys: rkeys, children: rchildren },
                                order,
                            );
                            self.file.write_page(&rp)?;

                            self.file.update_page(page_id, |p| {
                                write_node(p, &NodeData::Internal { keys, children }, order)
                            })?;

                            Ok((page_id, Some((up, right_id))))
                        }
                    }
                }
            }
        }
    }

    /// 删掉一个键，键不存在返回 false
    pub fn delete(&mut self, key: i64) -> io::Result<bool> {
        match self.root()? {
            None => Ok(false),
            Some(r) => self.delete_rec(r, key),
        }
    }

    /// 递归删除：一路走到叶子，删掉那一对
    fn delete_rec(&mut self, page_id: PageId, key: i64) -> io::Result<bool> {
        let order = self.order;
        match self.file.with_page(page_id, |p| read_node(p, order))? {
            NodeData::Leaf { mut keys, mut values, next } => match keys.binary_search(&key) {
                Err(_) => Ok(false),
                Ok(i) => {
                    keys.remove(i);
                    values.remove(i);
                    self.file.update_page(page_id, |p| {
                        write_node(p, &NodeData::Leaf { keys, values, next }, order)
                    })?;
                    Ok(true)
                }
            },
            NodeData::Internal { keys, children } => {
                let i = child_index(&keys, key);
                self.delete_rec(children[i], key)
            }
        }
    }

    /// 按键升序返回全部
    pub fn iter_all(&mut self) -> io::Result<Vec<(i64, u64)>> {
        let mut out = Vec::new();
        if let Some(r) = self.root()? {
            self.walk(r, &mut out)?;
        }
        Ok(out)
    }

    /// 中序遍历
    fn walk(&mut self, page_id: PageId, out: &mut Vec<(i64, u64)>) -> io::Result<()> {
        let order = self.order;
        match self.file.with_page(page_id, |p| read_node(p, order))? {
            NodeData::Leaf { keys, values, .. } => {
                for i in 0..keys.len() {
                    out.push((keys[i], values[i]));
                }
            }
            NodeData::Internal { children, .. } => {
                for c in children {
                    self.walk(c, out)?;
                }
            }
        }
        Ok(())
    }

    /// 把脏页写回
    pub fn flush(&mut self) -> io::Result<()> {
        self.file.flush()
    }

}