508 lines
16 KiB
Rust
508 lines
16 KiB
Rust
use std::io;
|
|
use std::path::Path;
|
|
|
|
use crate::page::{Page, PageFile, PageId};
|
|
|
|
// 磁盘 B+ 树：键是变长字节串，允许重复键，叶子链支持范围扫描。
|
|
|
|
const META_PAGE: PageId = 0;
|
|
const META_MAGIC: &[u8; 4] = b"CBK2";
|
|
const OFF_PAGE_SIZE: usize = 4;
|
|
const OFF_MAX_KEY: usize = 8;
|
|
const OFF_ROOT: usize = 10;
|
|
const OFF_ORDER: usize = 18;
|
|
|
|
const NODE_LEAF: u8 = 0;
|
|
const NODE_INTERNAL: u8 = 1;
|
|
const OFF_KIND: usize = 0;
|
|
const OFF_COUNT: usize = 1;
|
|
const OFF_NEXT: usize = 3;
|
|
const HEADER: usize = 11;
|
|
const SLOT: usize = 2;
|
|
const CELL_FIXED: usize = 2 + 8;
|
|
|
|
/// 键长上限：节点里的格子必须放得进一页
|
|
pub const MAX_KEY_BYTES: usize = 64;
|
|
|
|
/// 一个节点最多几个键
|
|
fn max_keys(order: usize) -> usize {
|
|
order - 1
|
|
}
|
|
|
|
/// 一个格子最多占多少字节
|
|
fn max_cell_bytes() -> usize {
|
|
CELL_FIXED + MAX_KEY_BYTES
|
|
}
|
|
|
|
/// 这个 order 放得进一页吗
|
|
pub fn fits_in_page(page_size: usize, order: usize) -> bool {
|
|
order >= 2 && HEADER + max_keys(order) * (SLOT + max_cell_bytes()) <= page_size
|
|
}
|
|
|
|
/// 一页最多几个孩子
|
|
pub fn max_order(page_size: usize) -> usize {
|
|
page_size.saturating_sub(HEADER) / (SLOT + max_cell_bytes()) + 1
|
|
}
|
|
|
|
/// 读一个 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_PAGE_SIZE, page_size as u32);
|
|
write_u16(page, OFF_MAX_KEY, MAX_KEY_BYTES as u16);
|
|
write_u64(page, OFF_ROOT, root);
|
|
write_u16(page, OFF_ORDER, order as u16);
|
|
}
|
|
|
|
/// 校验文件头与配置
|
|
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_PAGE_SIZE) as usize;
|
|
let stored_order = read_u16(page, OFF_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(())
|
|
}
|
|
|
|
/// 一个节点：叶子存行位置，内部节点存孩子页号
|
|
struct Node {
|
|
kind: u8,
|
|
next: PageId,
|
|
keys: Vec<Vec<u8>>,
|
|
values: Vec<u64>,
|
|
}
|
|
|
|
/// 一个格子占多少字节
|
|
fn cell_bytes(key: &[u8]) -> usize {
|
|
CELL_FIXED + key.len()
|
|
}
|
|
|
|
/// 节点占多少字节
|
|
fn node_bytes(node: &Node) -> usize {
|
|
HEADER + node.keys.len() * SLOT + node.keys.iter().map(|k| cell_bytes(k)).sum::<usize>()
|
|
}
|
|
|
|
/// 把节点写回一页（格子从页尾往前摆，槽位从头部往后长）
|
|
fn put_node(p: &mut Page, node: &Node) -> io::Result<()> {
|
|
if node.keys.len() > u16::MAX as usize || node_bytes(node) > p.data.len() {
|
|
return Err(io::Error::new(
|
|
io::ErrorKind::InvalidInput,
|
|
"index node does not fit in a page",
|
|
));
|
|
}
|
|
p.zero();
|
|
write_u8(p, OFF_KIND, node.kind);
|
|
write_u16(p, OFF_COUNT, node.keys.len() as u16);
|
|
// 内部节点比键多一个孩子，最后一个放在 next 槽位里
|
|
let trailer = if node.kind == NODE_INTERNAL {
|
|
node.values.last().copied().unwrap_or(0)
|
|
} else {
|
|
node.next
|
|
};
|
|
write_u64(p, OFF_NEXT, trailer);
|
|
let mut cursor = p.data.len();
|
|
for (i, key) in node.keys.iter().enumerate() {
|
|
let size = cell_bytes(key);
|
|
cursor -= size;
|
|
write_u16(p, cursor, key.len() as u16);
|
|
p.data[cursor + 2..cursor + 2 + key.len()].copy_from_slice(key);
|
|
write_u64(p, cursor + 2 + key.len(), node.values[i]);
|
|
write_u16(p, HEADER + i * SLOT, cursor as u16);
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 从一页读出节点
|
|
fn get_node(p: &Page) -> Node {
|
|
let kind = read_u8(p, OFF_KIND);
|
|
let count = read_u16(p, OFF_COUNT) as usize;
|
|
let trailer = read_u64(p, OFF_NEXT);
|
|
let mut keys = Vec::with_capacity(count);
|
|
let mut values = Vec::with_capacity(count + 1);
|
|
for i in 0..count {
|
|
let off = read_u16(p, HEADER + i * SLOT) as usize;
|
|
let klen = read_u16(p, off) as usize;
|
|
keys.push(p.data[off + 2..off + 2 + klen].to_vec());
|
|
values.push(read_u64(p, off + 2 + klen));
|
|
}
|
|
if kind == NODE_INTERNAL {
|
|
values.push(trailer);
|
|
Node { kind, next: 0, keys, values }
|
|
} else {
|
|
Node { kind, next: trailer, keys, values }
|
|
}
|
|
}
|
|
|
|
/// 叶子里的插入位置：同键按 value 排，保证顺序稳定
|
|
fn leaf_slot(keys: &[Vec<u8>], values: &[u64], key: &[u8], value: u64) -> usize {
|
|
let mut i = keys.partition_point(|k| k.as_slice() < key);
|
|
while i < keys.len() && keys[i].as_slice() == key && values[i] < value {
|
|
i += 1;
|
|
}
|
|
i
|
|
}
|
|
|
|
/// 内部节点里选哪个孩子
|
|
fn child_slot(keys: &[Vec<u8>], key: &[u8]) -> usize {
|
|
keys.partition_point(|k| k.as_slice() <= key)
|
|
}
|
|
|
|
/// 切开一个节点：返回上提的键与右半节点
|
|
fn split_node(node: &mut Node) -> (Vec<u8>, Node) {
|
|
let mid = node.keys.len() / 2;
|
|
if node.kind == NODE_LEAF {
|
|
let sep = node.keys[mid].clone();
|
|
let right = Node {
|
|
kind: NODE_LEAF,
|
|
next: node.next,
|
|
keys: node.keys.split_off(mid),
|
|
values: node.values.split_off(mid),
|
|
};
|
|
(sep, right)
|
|
} else {
|
|
let sep = node.keys[mid].clone();
|
|
let right_keys = node.keys.split_off(mid + 1);
|
|
let right_values = node.values.split_off(mid + 1);
|
|
node.keys.pop();
|
|
(
|
|
sep,
|
|
Node { kind: NODE_INTERNAL, next: 0, keys: right_keys, values: right_values },
|
|
)
|
|
}
|
|
}
|
|
|
|
/// 磁盘 B+ 树
|
|
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 check_key(&self, key: &[u8]) -> io::Result<()> {
|
|
if key.len() > MAX_KEY_BYTES {
|
|
return Err(io::Error::new(
|
|
io::ErrorKind::InvalidInput,
|
|
format!(
|
|
"index key is {} bytes, at most {} are supported",
|
|
key.len(),
|
|
MAX_KEY_BYTES
|
|
),
|
|
));
|
|
}
|
|
Ok(())
|
|
}
|
|
|
|
/// 读根页页号；空树是 None
|
|
fn root(&mut self) -> io::Result<Option<PageId>> {
|
|
let r = self.file.with_page(META_PAGE, |p| read_u64(p, OFF_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_ROOT, r.unwrap_or(0)))
|
|
}
|
|
|
|
/// 追加一个新节点页
|
|
fn alloc(&mut self, node: &Node) -> io::Result<PageId> {
|
|
let page = self.file.append_page()?;
|
|
let id = page.id;
|
|
self.file.update_page(id, |p| put_node(p, node))??;
|
|
Ok(id)
|
|
}
|
|
|
|
/// 按页号读节点
|
|
fn load(&mut self, id: PageId) -> io::Result<Node> {
|
|
self.file.with_page(id, get_node)
|
|
}
|
|
|
|
/// 按页号写节点
|
|
fn store(&mut self, id: PageId, node: &Node) -> io::Result<()> {
|
|
self.file.update_page(id, |p| put_node(p, node))?
|
|
}
|
|
|
|
/// 清空整棵树
|
|
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(())
|
|
}
|
|
|
|
/// 树是空的吗（只看最左叶子，不用整树遍历）
|
|
pub fn is_empty(&mut self) -> io::Result<bool> {
|
|
let Some(mut id) = self.root()? else {
|
|
return Ok(true);
|
|
};
|
|
loop {
|
|
let node = self.load(id)?;
|
|
if node.kind == NODE_LEAF {
|
|
return Ok(node.keys.is_empty());
|
|
}
|
|
id = node.values[0];
|
|
}
|
|
}
|
|
|
|
/// 按 key 取第一个位置
|
|
pub fn get(&mut self, key: &[u8]) -> io::Result<Option<u64>> {
|
|
Ok(self.get_all(key)?.into_iter().next())
|
|
}
|
|
|
|
/// 按 key 取全部位置（重复键按值升序）
|
|
pub fn get_all(&mut self, key: &[u8]) -> io::Result<Vec<u64>> {
|
|
let Some(mut id) = self.root()? else {
|
|
return Ok(Vec::new());
|
|
};
|
|
loop {
|
|
let node = self.load(id)?;
|
|
if node.kind == NODE_LEAF {
|
|
let start = node.keys.partition_point(|k| k.as_slice() < key);
|
|
let mut out = Vec::new();
|
|
for i in start..node.keys.len() {
|
|
if node.keys[i].as_slice() != key {
|
|
break;
|
|
}
|
|
out.push(node.values[i]);
|
|
}
|
|
return Ok(out);
|
|
}
|
|
id = node.values[child_slot(&node.keys, key)];
|
|
}
|
|
}
|
|
|
|
/// 插入一个键值对（键可重复）
|
|
pub fn insert(&mut self, key: &[u8], value: u64) -> io::Result<()> {
|
|
self.check_key(key)?;
|
|
match self.root()? {
|
|
None => {
|
|
let node = Node {
|
|
kind: NODE_LEAF,
|
|
next: 0,
|
|
keys: vec![key.to_vec()],
|
|
values: vec![value],
|
|
};
|
|
let id = self.alloc(&node)?;
|
|
self.set_root(Some(id))
|
|
}
|
|
Some(root) => match self.insert_rec(root, key, value)? {
|
|
None => Ok(()),
|
|
Some((sep, right_id)) => {
|
|
let node = Node {
|
|
kind: NODE_INTERNAL,
|
|
next: 0,
|
|
keys: vec![sep],
|
|
values: vec![root, right_id],
|
|
};
|
|
let id = self.alloc(&node)?;
|
|
self.set_root(Some(id))
|
|
}
|
|
},
|
|
}
|
|
}
|
|
|
|
/// 递归插入，溢出就把右半节点与上提键交给上层
|
|
fn insert_rec(
|
|
&mut self,
|
|
id: PageId,
|
|
key: &[u8],
|
|
value: u64,
|
|
) -> io::Result<Option<(Vec<u8>, PageId)>> {
|
|
let mut node = self.load(id)?;
|
|
if node.kind == NODE_LEAF {
|
|
let at = leaf_slot(&node.keys, &node.values, key, value);
|
|
node.keys.insert(at, key.to_vec());
|
|
node.values.insert(at, value);
|
|
} else {
|
|
let child = node.values[child_slot(&node.keys, key)];
|
|
if let Some((sep, right_id)) = self.insert_rec(child, key, value)? {
|
|
let at = node.keys.partition_point(|k| k.as_slice() <= sep.as_slice());
|
|
node.keys.insert(at, sep);
|
|
node.values.insert(at + 1, right_id);
|
|
}
|
|
}
|
|
if node.keys.len() <= max_keys(self.order) {
|
|
self.store(id, &node)?;
|
|
return Ok(None);
|
|
}
|
|
let (sep, mut right) = split_node(&mut node);
|
|
let right_id = self.alloc(&right)?;
|
|
if node.kind == NODE_LEAF {
|
|
right.next = node.next;
|
|
node.next = right_id;
|
|
}
|
|
self.store(id, &node)?;
|
|
Ok(Some((sep, right_id)))
|
|
}
|
|
|
|
/// 删掉一个 (key, value)
|
|
pub fn delete(&mut self, key: &[u8], value: u64) -> io::Result<bool> {
|
|
let Some(root) = self.root()? else {
|
|
return Ok(false);
|
|
};
|
|
self.delete_rec(root, key, value)
|
|
}
|
|
|
|
/// 递归删一个 (key, value)
|
|
fn delete_rec(&mut self, id: PageId, key: &[u8], value: u64) -> io::Result<bool> {
|
|
let mut node = self.load(id)?;
|
|
if node.kind == NODE_LEAF {
|
|
let at = leaf_slot(&node.keys, &node.values, key, value);
|
|
if at < node.keys.len() && node.keys[at].as_slice() == key && node.values[at] == value {
|
|
node.keys.remove(at);
|
|
node.values.remove(at);
|
|
self.store(id, &node)?;
|
|
return Ok(true);
|
|
}
|
|
return Ok(false);
|
|
}
|
|
let child = node.values[child_slot(&node.keys, key)];
|
|
self.delete_rec(child, key, value)
|
|
}
|
|
|
|
/// 叶子链上按顺序取全部 (键, 值)
|
|
pub fn iter_all(&mut self) -> io::Result<Vec<(Vec<u8>, u64)>> {
|
|
self.scan(None, None)
|
|
}
|
|
|
|
/// 范围扫描：lo / hi 都是闭区间，None 表示不限
|
|
pub fn scan(
|
|
&mut self,
|
|
lo: Option<&[u8]>,
|
|
hi: Option<&[u8]>,
|
|
) -> io::Result<Vec<(Vec<u8>, u64)>> {
|
|
let Some(mut id) = self.root()? else {
|
|
return Ok(Vec::new());
|
|
};
|
|
loop {
|
|
let node = self.load(id)?;
|
|
if node.kind == NODE_LEAF {
|
|
id = match lo {
|
|
None => id,
|
|
Some(start) => {
|
|
let at = node.keys.partition_point(|k| k.as_slice() < start);
|
|
if at < node.keys.len() { id } else { node.next }
|
|
}
|
|
};
|
|
break;
|
|
}
|
|
id = match lo {
|
|
None => node.values[0],
|
|
Some(start) => node.values[child_slot(&node.keys, start)],
|
|
};
|
|
}
|
|
let mut out = Vec::new();
|
|
while id != 0 {
|
|
let node = self.load(id)?;
|
|
for i in 0..node.keys.len() {
|
|
if lo.is_some_and(|start| node.keys[i].as_slice() < start) {
|
|
continue;
|
|
}
|
|
if hi.is_some_and(|end| node.keys[i].as_slice() > end) {
|
|
return Ok(out);
|
|
}
|
|
out.push((node.keys[i].clone(), node.values[i]));
|
|
}
|
|
id = node.next;
|
|
}
|
|
Ok(out)
|
|
}
|
|
|
|
/// 脏页写回
|
|
pub fn flush(&mut self) -> io::Result<()> {
|
|
self.file.flush()
|
|
}
|
|
|
|
/// 当前缓存页数
|
|
pub fn cached_pages(&self) -> usize {
|
|
self.file.cached_pages()
|
|
}
|
|
}
|