301 lines
8.8 KiB
Rust
301 lines
8.8 KiB
Rust
use crate::config::MAX_KEYS;
|
|
use crate::node::*;
|
|
use crate::record::Record;
|
|
use crate::storage::Storage;
|
|
|
|
pub struct BPlusTree<S> {
|
|
storage: S,
|
|
root_loc: usize,
|
|
}
|
|
|
|
impl<S> BPlusTree<S>
|
|
where
|
|
S: Storage,
|
|
{
|
|
pub fn open(mut storage: S) -> Self {
|
|
let root = Node::Leaf(LeafNode::new());
|
|
storage.write_node(0, &root);
|
|
BPlusTree {
|
|
storage,
|
|
root_loc: 0,
|
|
}
|
|
}
|
|
|
|
pub fn find(&mut self, key: i32) -> Option<Record> {
|
|
let mut current_loc = self.root_loc;
|
|
loop {
|
|
let node = self.storage.read_node(current_loc)?;
|
|
match node {
|
|
Node::Internal(internal) => {
|
|
let mut i = 0;
|
|
while i < internal.keys.len() && key >= internal.keys[i] {
|
|
i += 1;
|
|
}
|
|
current_loc = internal.children[i];
|
|
}
|
|
Node::Leaf(leaf) => {
|
|
for (i, k) in leaf.keys.iter().enumerate() {
|
|
if *k == key {
|
|
return Some(leaf.values[i]);
|
|
}
|
|
}
|
|
return None;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
pub fn insert(&mut self, value: Record) {
|
|
let key = value[0];
|
|
let mut path = Vec::new();
|
|
let mut current_loc = self.root_loc;
|
|
let mut current_node = self.storage.read_node(current_loc).unwrap();
|
|
|
|
while let Node::Internal(internal) = current_node {
|
|
path.push((current_loc, internal.clone()));
|
|
let mut i = 0;
|
|
while i < internal.keys.len() && key >= internal.keys[i] {
|
|
i += 1;
|
|
}
|
|
current_loc = internal.children[i];
|
|
current_node = self.storage.read_node(current_loc).unwrap();
|
|
}
|
|
|
|
// Insert into the leaf node
|
|
if let Node::Leaf(mut leaf) = current_node {
|
|
// Insert or update the key-value pair
|
|
let mut inserted = false;
|
|
for (i, k) in leaf.keys.iter_mut().enumerate() {
|
|
if *k == key {
|
|
leaf.values[i] = value;
|
|
inserted = true;
|
|
break;
|
|
} else if key < *k {
|
|
leaf.keys.insert(i, key);
|
|
leaf.values.insert(i, value);
|
|
inserted = true;
|
|
break;
|
|
}
|
|
}
|
|
if !inserted {
|
|
leaf.keys.push(key);
|
|
leaf.values.push(value);
|
|
}
|
|
|
|
self.storage
|
|
.write_node(current_loc, &Node::Leaf(leaf.clone()));
|
|
|
|
if leaf.keys.len() > MAX_KEYS {
|
|
self.split_leaf(current_loc, leaf, &mut path);
|
|
}
|
|
}
|
|
}
|
|
|
|
fn split_leaf(&mut self, loc: usize, leaf: LeafNode, path: &mut Vec<(usize, InternalNode)>) {
|
|
let mid = leaf.keys.len() / 2;
|
|
let new_leaf = LeafNode {
|
|
keys: leaf.keys[mid..].to_vec(),
|
|
values: leaf.values[mid..].to_vec(),
|
|
next: leaf.next,
|
|
};
|
|
let new_leaf_loc = self.storage.total_nodes();
|
|
|
|
let original_leaf = LeafNode {
|
|
keys: leaf.keys[..mid].to_vec(),
|
|
values: leaf.values[..mid].to_vec(),
|
|
next: Some(new_leaf_loc),
|
|
};
|
|
|
|
self.storage.write_node(loc, &Node::Leaf(original_leaf));
|
|
self.storage
|
|
.write_node(new_leaf_loc, &Node::Leaf(new_leaf.clone()));
|
|
|
|
let new_key = new_leaf.keys[0];
|
|
if path.is_empty() {
|
|
// Create a new root
|
|
let new_root = Node::Internal(InternalNode {
|
|
keys: vec![new_key],
|
|
children: vec![loc, new_leaf_loc],
|
|
});
|
|
let new_root_loc = self.storage.total_nodes();
|
|
self.storage.write_node(new_root_loc, &new_root);
|
|
self.root_loc = new_root_loc;
|
|
} else {
|
|
self.insert_into_parent(new_key, new_leaf_loc, path);
|
|
}
|
|
}
|
|
|
|
fn insert_into_parent(
|
|
&mut self,
|
|
key: i32,
|
|
new_child_loc: usize,
|
|
path: &mut Vec<(usize, InternalNode)>,
|
|
) {
|
|
let (parent_loc, mut parent) = path.pop().unwrap();
|
|
|
|
// Insert the new key and child into the parent
|
|
let mut i = 0;
|
|
while i < parent.keys.len() && key >= parent.keys[i] {
|
|
i += 1;
|
|
}
|
|
parent.keys.insert(i, key);
|
|
parent.children.insert(i + 1, new_child_loc);
|
|
|
|
self.storage
|
|
.write_node(parent_loc, &Node::Internal(parent.clone()));
|
|
|
|
if parent.keys.len() > MAX_KEYS {
|
|
self.split_internal(parent_loc, parent, path);
|
|
}
|
|
}
|
|
|
|
fn split_internal(
|
|
&mut self,
|
|
loc: usize,
|
|
internal: InternalNode,
|
|
path: &mut Vec<(usize, InternalNode)>,
|
|
) {
|
|
let mid = internal.keys.len() / 2;
|
|
let new_internal = InternalNode {
|
|
keys: internal.keys[mid + 1..].to_vec(),
|
|
children: internal.children[mid + 1..].to_vec(),
|
|
};
|
|
let new_internal_loc = self.storage.total_nodes();
|
|
|
|
let original_internal = InternalNode {
|
|
keys: internal.keys[..mid].to_vec(),
|
|
children: internal.children[..mid + 1].to_vec(),
|
|
};
|
|
|
|
self.storage
|
|
.write_node(loc, &Node::Internal(original_internal));
|
|
self.storage
|
|
.write_node(new_internal_loc, &Node::Internal(new_internal));
|
|
|
|
let new_key = internal.keys[mid];
|
|
if path.is_empty() {
|
|
// Create a new root
|
|
let new_root = Node::Internal(InternalNode {
|
|
keys: vec![new_key],
|
|
children: vec![loc, new_internal_loc],
|
|
});
|
|
let new_root_loc = self.storage.total_nodes();
|
|
self.storage.write_node(new_root_loc, &new_root);
|
|
self.root_loc = new_root_loc;
|
|
} else {
|
|
self.insert_into_parent(new_key, new_internal_loc, path);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use crate::storage::InMemoryStorage;
|
|
|
|
use proptest::prelude::*;
|
|
use proptest::proptest;
|
|
|
|
use std::collections::BTreeMap;
|
|
|
|
fn arb_record() -> impl Strategy<Value = Record> {
|
|
prop::array::uniform7(any::<i32>())
|
|
}
|
|
|
|
fn arb_key() -> impl Strategy<Value = i32> {
|
|
any::<i32>()
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
enum Op {
|
|
Insert(Record),
|
|
Find(i32),
|
|
}
|
|
|
|
fn arb_op() -> impl Strategy<Value = Op> {
|
|
prop_oneof![
|
|
arb_record().prop_map(Op::Insert),
|
|
arb_key().prop_map(Op::Find),
|
|
]
|
|
}
|
|
|
|
fn arb_ops() -> impl Strategy<Value = Vec<Op>> {
|
|
// Long enough to trigger many splits, but still fast
|
|
prop::collection::vec(arb_op(), 1..500)
|
|
}
|
|
|
|
proptest! {
|
|
#![proptest_config(ProptestConfig {
|
|
cases: 200,
|
|
.. ProptestConfig::default()
|
|
})]
|
|
|
|
#[test]
|
|
fn prop_bplustree_matches_btreemap(ops in arb_ops()) {
|
|
let storage = InMemoryStorage::new();
|
|
let mut tree = BPlusTree::open(storage);
|
|
|
|
let mut model = BTreeMap::<i32, Record>::new();
|
|
|
|
for op in ops {
|
|
match op {
|
|
Op::Insert(record) => {
|
|
let key = record[0];
|
|
tree.insert(record);
|
|
model.insert(key, record);
|
|
}
|
|
|
|
Op::Find(key) => {
|
|
let tree_res = tree.find(key);
|
|
let model_res = model.get(&key).copied();
|
|
prop_assert_eq!(tree_res, model_res);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
proptest! {
|
|
#[test]
|
|
fn prop_all_inserted_keys_are_findable(records in prop::collection::vec(arb_record(), 0..1000)) {
|
|
let storage = InMemoryStorage::new();
|
|
let mut tree = BPlusTree::open(storage);
|
|
let mut model = BTreeMap::<i32, Record>::new();
|
|
|
|
for rec in records {
|
|
tree.insert(rec);
|
|
model.insert(rec[0], rec);
|
|
}
|
|
|
|
for (key, value) in model {
|
|
prop_assert_eq!(tree.find(key), Some(value));
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn sanity_single_insert() {
|
|
let storage = InMemoryStorage::new();
|
|
let mut tree = BPlusTree::open(storage);
|
|
|
|
let rec = [42, 1, 2, 3, 4, 5, 6];
|
|
tree.insert(rec);
|
|
|
|
assert_eq!(tree.find(42), Some(rec));
|
|
assert_eq!(tree.find(7), None);
|
|
}
|
|
|
|
#[test]
|
|
fn sanity_update_overwrites_value() {
|
|
let storage = InMemoryStorage::new();
|
|
let mut tree = BPlusTree::open(storage);
|
|
|
|
let r1 = [1, 1, 1, 1, 1, 1, 1];
|
|
let r2 = [1, 9, 9, 9, 9, 9, 9];
|
|
|
|
tree.insert(r1);
|
|
tree.insert(r2);
|
|
|
|
assert_eq!(tree.find(1), Some(r2));
|
|
}
|
|
}
|