freash start
This commit is contained in:
bronku
parent
410e26cea8
commit
d730a57706
6 changed files with 248 additions and 1028 deletions
595
src/btree.rs
595
src/btree.rs
|
|
@ -1,595 +0,0 @@
|
|||
use crate::Record;
|
||||
use crate::consts::MAX_KEYS;
|
||||
use crate::node::Node;
|
||||
use crate::node_storage::NodeStorage;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct BTree {
|
||||
storage: NodeStorage,
|
||||
}
|
||||
|
||||
enum FindResult {
|
||||
EmptyTree,
|
||||
|
||||
Found {
|
||||
page: usize,
|
||||
index: usize,
|
||||
node: Node,
|
||||
record: Record,
|
||||
},
|
||||
|
||||
NotFound {
|
||||
node: (Node, usize), // leaf node where insertion must happen
|
||||
parent: Option<(Node, usize)>, // optional parent for reducing disk reads on split/compensate
|
||||
},
|
||||
}
|
||||
|
||||
impl BTree {
|
||||
pub fn new(filename: &str) -> Self {
|
||||
BTree {
|
||||
storage: NodeStorage::open(filename),
|
||||
}
|
||||
}
|
||||
|
||||
fn find(&mut self, key: i32) -> FindResult {
|
||||
if self.storage.num_nodes() == 0 {
|
||||
return FindResult::EmptyTree;
|
||||
}
|
||||
|
||||
let mut current = 0;
|
||||
let mut parent: Option<(Node, usize)> = None;
|
||||
|
||||
'outer: loop {
|
||||
let node = self.storage.read_node(current);
|
||||
|
||||
if node.is_leaf {
|
||||
for i in 0..node.num_keys {
|
||||
if let Some(rec) = node.keys[i] {
|
||||
if rec.key == key {
|
||||
return FindResult::Found {
|
||||
page: current,
|
||||
index: i,
|
||||
node,
|
||||
record: rec,
|
||||
};
|
||||
}
|
||||
}
|
||||
}
|
||||
return FindResult::NotFound {
|
||||
node: (node, current),
|
||||
parent,
|
||||
};
|
||||
}
|
||||
|
||||
let mut prev_key = i32::MIN;
|
||||
|
||||
for i in 0..node.num_keys {
|
||||
let node_key = node.keys[i].unwrap().key;
|
||||
|
||||
if key == node_key {
|
||||
let rec = node.keys[i].unwrap();
|
||||
return FindResult::Found {
|
||||
page: current,
|
||||
index: i,
|
||||
node,
|
||||
record: rec,
|
||||
};
|
||||
}
|
||||
|
||||
if prev_key < key && key < node_key {
|
||||
parent = Some((node, current));
|
||||
current = node.children[i].unwrap();
|
||||
continue 'outer;
|
||||
}
|
||||
|
||||
prev_key = node_key;
|
||||
}
|
||||
|
||||
parent = Some((node, current));
|
||||
current = node.children[node.num_keys].unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
pub fn search(&mut self, key: i32) -> Option<Record> {
|
||||
match self.find(key) {
|
||||
FindResult::Found { record, .. } => Some(record),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
|
||||
fn try_insert_without_split(mut node: Node, key: Record) -> Option<Node> {
|
||||
if node.num_keys >= MAX_KEYS {
|
||||
return None;
|
||||
}
|
||||
|
||||
let mut pos = 0;
|
||||
while pos < node.num_keys && node.keys[pos].unwrap().key < key.key {
|
||||
pos += 1;
|
||||
}
|
||||
|
||||
for i in (pos..node.num_keys).rev() {
|
||||
node.keys[i + 1] = node.keys[i];
|
||||
}
|
||||
|
||||
node.keys[pos] = Some(key);
|
||||
node.num_keys += 1;
|
||||
Some(node)
|
||||
}
|
||||
|
||||
fn redistribute(
|
||||
mut left: Node,
|
||||
left_page: usize,
|
||||
mut parent: Node,
|
||||
parent_page: usize,
|
||||
mut right: Node,
|
||||
right_page: usize,
|
||||
record: Record,
|
||||
) -> Vec<(Node, usize)> {
|
||||
let mut all_keys: Vec<Record> = vec![];
|
||||
|
||||
// Collect left keys
|
||||
for i in 0..left.num_keys {
|
||||
if let Some(k) = left.keys[i] {
|
||||
all_keys.push(k);
|
||||
}
|
||||
}
|
||||
|
||||
// Add parent separator key
|
||||
let last_left_key = all_keys.last().unwrap().key;
|
||||
let separator_idx = parent
|
||||
.keys
|
||||
.iter()
|
||||
.position(|k| match k {
|
||||
Some(r) => r.key > last_left_key,
|
||||
None => false,
|
||||
})
|
||||
.unwrap();
|
||||
let separating_key = parent.keys[separator_idx].unwrap();
|
||||
all_keys.push(separating_key);
|
||||
|
||||
// Collect right keys
|
||||
for i in 0..right.num_keys {
|
||||
if let Some(k) = right.keys[i] {
|
||||
all_keys.push(k);
|
||||
}
|
||||
}
|
||||
|
||||
// Add new record
|
||||
all_keys.push(record);
|
||||
all_keys.sort_by_key(|r| r.key);
|
||||
|
||||
// Split keys evenly
|
||||
let total_keys = all_keys.len();
|
||||
let left_num = total_keys / 2;
|
||||
let right_num = total_keys - left_num - 1;
|
||||
|
||||
left.keys.fill(None);
|
||||
for i in 0..left_num {
|
||||
left.keys[i] = Some(all_keys[i]);
|
||||
}
|
||||
left.num_keys = left_num;
|
||||
|
||||
let middle_key = all_keys[left_num];
|
||||
let right_start = left_num + 1;
|
||||
|
||||
// Update parent key safely
|
||||
if separator_idx < parent.keys.len() {
|
||||
parent.keys[separator_idx] = Some(middle_key);
|
||||
} else {
|
||||
parent.keys[parent.keys.iter().position(|k| k.is_none()).unwrap()] = Some(middle_key);
|
||||
}
|
||||
|
||||
right.keys.fill(None);
|
||||
for i in 0..right_num {
|
||||
right.keys[i] = Some(all_keys[right_start + i]);
|
||||
}
|
||||
right.num_keys = right_num;
|
||||
|
||||
// Update parent pointers
|
||||
left.parent = Some(parent_page);
|
||||
right.parent = Some(parent_page);
|
||||
|
||||
vec![
|
||||
(left, left_page),
|
||||
(parent, parent_page),
|
||||
(right, right_page),
|
||||
]
|
||||
}
|
||||
|
||||
fn try_compensate(
|
||||
&mut self,
|
||||
node: (Node, usize),
|
||||
parent: Option<(Node, usize)>,
|
||||
input: Record,
|
||||
) -> Option<Vec<(Node, usize)>> {
|
||||
let (parent, parent_idx) = parent?;
|
||||
let (node, node_idx) = node;
|
||||
|
||||
let node_idx_in_parent = parent
|
||||
.children
|
||||
.iter()
|
||||
.position(|&child_opt| child_opt == Some(node_idx))?;
|
||||
|
||||
if node_idx_in_parent > 0 {
|
||||
let left_sibling_idx = parent.children[node_idx_in_parent - 1]?;
|
||||
let left_sibling = self.storage.read_node(left_sibling_idx);
|
||||
|
||||
if left_sibling.is_leaf && left_sibling.num_keys < MAX_KEYS {
|
||||
return Some(BTree::redistribute(
|
||||
left_sibling,
|
||||
left_sibling_idx,
|
||||
parent,
|
||||
parent_idx,
|
||||
node,
|
||||
node_idx,
|
||||
input,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
// 4. Attempt right sibling compensation
|
||||
if node_idx_in_parent + 1 < parent.children.len() {
|
||||
let right_sibling_idx = parent.children[node_idx_in_parent + 1]?;
|
||||
let right_sibling = self.storage.read_node(right_sibling_idx);
|
||||
|
||||
if right_sibling.is_leaf && right_sibling.num_keys < MAX_KEYS {
|
||||
return Some(BTree::redistribute(
|
||||
node,
|
||||
node_idx,
|
||||
parent,
|
||||
parent_idx,
|
||||
right_sibling,
|
||||
right_sibling_idx,
|
||||
input,
|
||||
));
|
||||
}
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
fn split_recursive(&mut self, node_page: usize, input: Record) -> Option<usize> {
|
||||
let node = self.storage.read_node(node_page);
|
||||
let mut keys: Vec<Record> = node
|
||||
.keys
|
||||
.iter()
|
||||
.filter_map(|k| k.as_ref())
|
||||
.cloned()
|
||||
.collect();
|
||||
keys.push(input);
|
||||
keys.sort_by_key(|r| r.key);
|
||||
|
||||
let mid = keys.len() / 2;
|
||||
let middle_key = keys[mid];
|
||||
|
||||
// Create left node
|
||||
let mut left_node = Node::new(node.is_leaf);
|
||||
left_node.num_keys = mid;
|
||||
for i in 0..mid {
|
||||
left_node.keys[i] = Some(keys[i]);
|
||||
}
|
||||
if !node.is_leaf {
|
||||
for i in 0..=mid {
|
||||
left_node.children[i] = node.children[i];
|
||||
}
|
||||
}
|
||||
|
||||
// Create right node
|
||||
let mut right_node = Node::new(node.is_leaf);
|
||||
right_node.num_keys = keys.len() - mid - 1;
|
||||
for i in 0..right_node.num_keys as usize {
|
||||
right_node.keys[i] = Some(keys[mid + 1 + i]);
|
||||
}
|
||||
if !node.is_leaf {
|
||||
for i in 0..=right_node.num_keys as usize {
|
||||
right_node.children[i] = node.children[mid + 1 + i];
|
||||
}
|
||||
}
|
||||
|
||||
// Update parent pointers for children
|
||||
if !node.is_leaf {
|
||||
for i in 0..=left_node.num_keys as usize {
|
||||
if let Some(child_page) = left_node.children[i] {
|
||||
let mut child = self.storage.read_node(child_page);
|
||||
child.parent = Some(node_page);
|
||||
self.storage.write_node(child_page, &child);
|
||||
}
|
||||
}
|
||||
for i in 0..=right_node.num_keys as usize {
|
||||
if let Some(child_page) = right_node.children[i] {
|
||||
let mut child = self.storage.read_node(child_page);
|
||||
child.parent = Some(self.storage.num_nodes()); // Placeholder, will be updated
|
||||
self.storage.write_node(child_page, &child);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Write left and right nodes to storage
|
||||
let left_page = if node_page == 0 {
|
||||
// If splitting the root, left node becomes new root
|
||||
self.storage.write_node(node_page, &left_node);
|
||||
node_page
|
||||
} else {
|
||||
self.storage.append_node(&left_node)
|
||||
};
|
||||
let right_page = self.storage.append_node(&right_node);
|
||||
|
||||
// Update parent
|
||||
if node_page == 0 {
|
||||
// Root was split, create new root
|
||||
let mut new_root = Node::new(false);
|
||||
new_root.keys[0] = Some(middle_key);
|
||||
new_root.children[0] = Some(left_page);
|
||||
new_root.children[1] = Some(right_page);
|
||||
new_root.num_keys = 1;
|
||||
new_root.parent = None;
|
||||
left_node.parent = Some(self.storage.num_nodes());
|
||||
right_node.parent = Some(self.storage.num_nodes());
|
||||
self.storage.write_node(left_page, &left_node);
|
||||
self.storage.write_node(right_page, &right_node);
|
||||
let new_root_page = self.storage.append_node(&new_root);
|
||||
Some(new_root_page)
|
||||
} else {
|
||||
// Update parent with middle key and new child
|
||||
let parent_page = node.parent.unwrap();
|
||||
let mut parent = self.storage.read_node(parent_page);
|
||||
|
||||
// Insert middle key into parent
|
||||
let mut key_pos = 0;
|
||||
while key_pos < parent.num_keys && parent.keys[key_pos].unwrap().key < middle_key.key {
|
||||
key_pos += 1;
|
||||
}
|
||||
for i in (key_pos..parent.num_keys as usize).rev() {
|
||||
parent.keys[i + 1] = parent.keys[i];
|
||||
parent.children[i + 2] = parent.children[i + 1];
|
||||
}
|
||||
parent.keys[key_pos] = Some(middle_key);
|
||||
parent.children[key_pos + 1] = Some(right_page);
|
||||
parent.num_keys += 1;
|
||||
self.storage.write_node(parent_page, &parent);
|
||||
// Recursively split parent if full
|
||||
if parent.num_keys as usize == MAX_KEYS {
|
||||
self.split_recursive(parent_page, middle_key)
|
||||
} else {
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn insert(&mut self, input: Record) {
|
||||
use FindResult::*;
|
||||
|
||||
match self.find(input.key) {
|
||||
Found {
|
||||
page,
|
||||
mut node,
|
||||
index,
|
||||
..
|
||||
} => {
|
||||
// Key exists ─ update record
|
||||
node.keys[index] = Some(input);
|
||||
self.storage.write_node(page, &node);
|
||||
return;
|
||||
}
|
||||
|
||||
EmptyTree => {
|
||||
// Create first root
|
||||
let mut root = Node::new(true);
|
||||
root.keys[0] = Some(input);
|
||||
root.num_keys = 1;
|
||||
root.parent = None;
|
||||
self.storage.append_node(&root);
|
||||
return;
|
||||
}
|
||||
|
||||
NotFound { node, parent } => {
|
||||
// 1) Try normal insertion
|
||||
if let Some(updated_node) = BTree::try_insert_without_split(node.0, input) {
|
||||
self.storage.write_node(node.1, &updated_node);
|
||||
return;
|
||||
}
|
||||
|
||||
// 2) Try compensation
|
||||
if let Some(updated_nodes) = self.try_compensate(node, parent, input) {
|
||||
for (n, idx) in updated_nodes {
|
||||
self.storage.write_node(idx, &n);
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
// 3) Must split and possibly recurse upward #todo
|
||||
let _ = self.split_recursive(node.1, input);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use tempfile::tempfile;
|
||||
|
||||
// Helper function to create a BTree with a temporary file
|
||||
fn create_btree() -> BTree {
|
||||
let file = tempfile().unwrap();
|
||||
BTree {
|
||||
storage: NodeStorage {
|
||||
file,
|
||||
page_reads: 0,
|
||||
page_writes: 0,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
fn create_record(key: i32) -> Record {
|
||||
let mut out = Record::random();
|
||||
out.key = key;
|
||||
out
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_empty_tree() {
|
||||
let mut btree = create_btree();
|
||||
assert!(btree.search(1).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_single_insert_and_search() {
|
||||
let mut btree = create_btree();
|
||||
let key = 42;
|
||||
let record = create_record(key);
|
||||
btree.insert(record);
|
||||
assert!(btree.search(key).is_some());
|
||||
assert_eq!(btree.search(key).unwrap().key, key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_and_search_multiple() {
|
||||
let mut btree = create_btree();
|
||||
let keys = vec![10, 20, 30, 40, 50];
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key));
|
||||
}
|
||||
for key in &keys {
|
||||
assert!(btree.search(*key).is_some());
|
||||
assert_eq!(btree.search(*key).unwrap().key, *key);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_duplicate_key() {
|
||||
let mut btree = create_btree();
|
||||
let key = 42;
|
||||
let record1 = create_record(key);
|
||||
let record2 = create_record(key);
|
||||
btree.insert(record1);
|
||||
btree.insert(record2);
|
||||
assert!(btree.search(key).is_some());
|
||||
// Ensure the data was updated (assuming insert updates existing keys)
|
||||
assert_eq!(btree.search(key).unwrap().key, record2.key);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_and_search_min_max_keys() {
|
||||
let mut btree = create_btree();
|
||||
let min_key = i32::MIN;
|
||||
let max_key = i32::MAX;
|
||||
btree.insert(create_record(min_key));
|
||||
btree.insert(create_record(max_key));
|
||||
assert!(btree.search(min_key).is_some());
|
||||
assert!(btree.search(max_key).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_search_nonexistent_key() {
|
||||
let mut btree = create_btree();
|
||||
btree.insert(create_record(10));
|
||||
btree.insert(create_record(20));
|
||||
btree.insert(create_record(30));
|
||||
assert!(btree.search(15).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_split_and_search() {
|
||||
let mut btree = create_btree();
|
||||
// Insert enough keys to force a split
|
||||
let keys = (1..=crate::consts::MAX_KEYS + 1).collect::<Vec<_>>();
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key as i32));
|
||||
}
|
||||
for key in &keys {
|
||||
assert!(btree.search(*key as i32).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_recursive_split() {
|
||||
let mut btree = create_btree();
|
||||
// Insert enough keys to force multiple splits
|
||||
let keys = (1..=crate::consts::MAX_KEYS * 2).collect::<Vec<_>>();
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key as i32));
|
||||
}
|
||||
for key in &keys {
|
||||
assert!(btree.search(*key as i32).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_stress_insert_and_search() {
|
||||
let mut btree = create_btree();
|
||||
let keys = (1..=1000).collect::<Vec<_>>();
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key as i32));
|
||||
}
|
||||
for key in &keys {
|
||||
assert!(btree.search(*key as i32).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_at_boundary() {
|
||||
let mut btree = create_btree();
|
||||
// Insert keys at the boundary of the node capacity
|
||||
let keys = (1..=crate::consts::MAX_KEYS).collect::<Vec<_>>();
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key as i32));
|
||||
}
|
||||
// Insert one more to force a split
|
||||
btree.insert(create_record((crate::consts::MAX_KEYS + 1) as i32));
|
||||
for key in 1..=crate::consts::MAX_KEYS + 1 {
|
||||
assert!(btree.search(key as i32).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_sorted_keys() {
|
||||
let mut btree = create_btree();
|
||||
// Insert keys in sorted order to test rightmost insertion
|
||||
let keys = (1..=crate::consts::MAX_KEYS * 2).collect::<Vec<_>>();
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key as i32));
|
||||
}
|
||||
for key in 1..=crate::consts::MAX_KEYS * 2 {
|
||||
assert!(btree.search(key as i32).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_reverse_sorted_keys() {
|
||||
let mut btree = create_btree();
|
||||
// Insert keys in reverse order to test leftmost insertion
|
||||
let keys = (1..=crate::consts::MAX_KEYS * 2).rev().collect::<Vec<_>>();
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key as i32));
|
||||
}
|
||||
for key in 1..=crate::consts::MAX_KEYS * 2 {
|
||||
assert!(btree.search(key as i32).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_and_search_large_keys() {
|
||||
let mut btree = create_btree();
|
||||
// Insert large keys to test edge cases
|
||||
let keys = vec![i32::MAX - 1, i32::MAX - 2, i32::MAX - 3];
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key));
|
||||
}
|
||||
for key in &keys {
|
||||
assert!(btree.search(*key).is_some());
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_insert_and_search_negative_keys() {
|
||||
let mut btree = create_btree();
|
||||
// Insert negative keys
|
||||
let keys = vec![-1, -2, -3, -4, -5];
|
||||
for key in &keys {
|
||||
btree.insert(create_record(*key));
|
||||
}
|
||||
for key in &keys {
|
||||
assert!(btree.search(*key).is_some());
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -1,2 +0,0 @@
|
|||
pub const MIN_DEGREE: usize = 3; // Example minimum degree
|
||||
pub const MAX_KEYS: usize = 2 * MIN_DEGREE; // Max keys per node
|
||||
287
src/main.rs
287
src/main.rs
|
|
@ -1,41 +1,258 @@
|
|||
mod btree;
|
||||
mod consts;
|
||||
mod node;
|
||||
mod node_storage;
|
||||
mod record;
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
|
||||
struct ValueType([u32; 2]); // Example implementation
|
||||
|
||||
use crate::btree::BTree;
|
||||
use crate::record::Record;
|
||||
type Value = u64;
|
||||
|
||||
const MAX_KEYS: usize = 4;
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
enum Node {
|
||||
Leaf(LeafNode),
|
||||
Internal(InternalNode),
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct LeafNode {
|
||||
keys: Vec<ValueType>,
|
||||
values: Vec<Value>,
|
||||
next: Option<usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct InternalNode {
|
||||
keys: Vec<ValueType>,
|
||||
children: Vec<usize>,
|
||||
}
|
||||
|
||||
trait Storage {
|
||||
fn read_node(&self, loc: usize) -> Option<Node>;
|
||||
fn write_node(&mut self, loc: usize, node: &Node);
|
||||
fn total_nodes(&self) -> usize;
|
||||
}
|
||||
|
||||
struct BPlusTree<S> {
|
||||
storage: S,
|
||||
root_loc: usize,
|
||||
}
|
||||
|
||||
impl<S> BPlusTree<S>
|
||||
where
|
||||
S: Storage,
|
||||
{
|
||||
pub fn open(mut storage: S) -> Self {
|
||||
// Initialize with an empty root node (leaf)
|
||||
let root = Node::Leaf(LeafNode {
|
||||
keys: Vec::new(),
|
||||
values: Vec::new(),
|
||||
next: None,
|
||||
});
|
||||
storage.write_node(0, &root);
|
||||
BPlusTree {
|
||||
storage,
|
||||
root_loc: 0,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn find(&self, key: ValueType) -> Option<Value> {
|
||||
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, key: ValueType, value: Value) {
|
||||
let mut path = Vec::new();
|
||||
let mut current_loc = self.root_loc;
|
||||
let mut current_node = self.storage.read_node(current_loc).unwrap();
|
||||
|
||||
// Traverse to the leaf node, recording the path
|
||||
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);
|
||||
}
|
||||
|
||||
// Write the updated leaf back to storage
|
||||
self.storage
|
||||
.write_node(current_loc, &Node::Leaf(leaf.clone()));
|
||||
|
||||
// Check if the leaf needs to be split
|
||||
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: ValueType,
|
||||
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);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Example in-memory storage implementation for testing
|
||||
struct InMemoryStorage {
|
||||
nodes: Vec<Option<Node>>,
|
||||
}
|
||||
|
||||
impl Storage for InMemoryStorage {
|
||||
fn read_node(&self, loc: usize) -> Option<Node> {
|
||||
self.nodes.get(loc)?.clone()
|
||||
}
|
||||
|
||||
fn write_node(&mut self, loc: usize, node: &Node) {
|
||||
if loc >= self.nodes.len() {
|
||||
self.nodes.resize(loc + 1, None);
|
||||
}
|
||||
self.nodes[loc] = Some(node.clone());
|
||||
}
|
||||
|
||||
fn total_nodes(&self) -> usize {
|
||||
self.nodes.len()
|
||||
}
|
||||
}
|
||||
|
||||
fn main() {
|
||||
let mut tree = BTree::new("btree.dat");
|
||||
tree.insert(Record {
|
||||
key: 1,
|
||||
a: [0, 0, 0, 0, 0],
|
||||
x: 0,
|
||||
});
|
||||
let storage = InMemoryStorage { nodes: Vec::new() };
|
||||
let mut tree = BPlusTree::open(storage);
|
||||
|
||||
let rec = tree.search(0);
|
||||
println!("{:?}", rec);
|
||||
// Insert some key-value pairs
|
||||
tree.insert(ValueType([1, 2]), 100);
|
||||
tree.insert(ValueType([3, 4]), 200);
|
||||
|
||||
/*
|
||||
let mut btree = BTree::new();
|
||||
|
||||
let mut values: Vec<i32> = vec![];
|
||||
|
||||
println!("Inserting values");
|
||||
for _ in 0..10 {
|
||||
let rec = Record::random();
|
||||
values.push(rec.key);
|
||||
btree.insert(rec);
|
||||
}
|
||||
|
||||
println!("\nB-Tree structure:");
|
||||
btree.print_tree();
|
||||
|
||||
println!("\nSearch results:");
|
||||
println!("Search for 6: {:?}", btree.search(values[0]));
|
||||
println!("Search for 15: {:?}", btree.search(values[4]));
|
||||
println!("Search for 21: {:?}", btree.search(100));
|
||||
*/
|
||||
// Find a key
|
||||
let value = tree.find(ValueType([1, 2]));
|
||||
println!("Found value: {:?}", value);
|
||||
}
|
||||
|
|
|
|||
91
src/node.rs
91
src/node.rs
|
|
@ -1,91 +0,0 @@
|
|||
use crate::consts::*;
|
||||
use crate::record::Record;
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
pub struct Node {
|
||||
pub keys: [Option<Record>; MAX_KEYS],
|
||||
pub children: [Option<usize>; MAX_KEYS + 1],
|
||||
pub num_keys: usize,
|
||||
pub is_leaf: bool,
|
||||
pub parent: Option<usize>,
|
||||
}
|
||||
|
||||
impl Node {
|
||||
pub fn new(is_leaf: bool) -> Self {
|
||||
Self {
|
||||
keys: [(); MAX_KEYS].map(|_| None),
|
||||
children: [(); MAX_KEYS + 1].map(|_| None),
|
||||
num_keys: 0,
|
||||
is_leaf,
|
||||
parent: None,
|
||||
}
|
||||
}
|
||||
|
||||
/*
|
||||
|
||||
pub fn is_full(&self) -> bool {
|
||||
self.keys.len() >= 2 * MIN_DEGREE - 1
|
||||
}
|
||||
|
||||
pub fn search(&self, key: i32) -> Option<Record> {
|
||||
let mut i = 0;
|
||||
while i < self.keys.len() && key > self.keys[i].key {
|
||||
i += 1;
|
||||
}
|
||||
|
||||
if i < self.keys.len() && key == self.keys[i].key {
|
||||
return Some(self.keys[i]);
|
||||
}
|
||||
|
||||
if self.is_leaf {
|
||||
return None;
|
||||
}
|
||||
|
||||
self.children[i].search(key)
|
||||
}
|
||||
|
||||
pub fn insert_non_full(&mut self, rec: Record) {
|
||||
let mut i = self.keys.len();
|
||||
|
||||
if self.is_leaf {
|
||||
self.keys.push(rec);
|
||||
i = self.keys.len() - 1;
|
||||
while i > 0 && self.keys[i] < self.keys[i - 1] {
|
||||
self.keys.swap(i, i - 1);
|
||||
i -= 1;
|
||||
}
|
||||
} else {
|
||||
while i > 0 && rec < self.keys[i - 1] {
|
||||
i -= 1;
|
||||
}
|
||||
|
||||
if self.children[i].is_full() {
|
||||
self.split_child(i);
|
||||
if rec > self.keys[i] {
|
||||
i += 1;
|
||||
}
|
||||
}
|
||||
self.children[i].insert_non_full(rec);
|
||||
}
|
||||
}
|
||||
|
||||
pub fn split_child(&mut self, i: usize) {
|
||||
let t = MIN_DEGREE;
|
||||
let full_child = &mut self.children[i];
|
||||
let mut new_child = Box::new(Node::new(full_child.is_leaf));
|
||||
|
||||
// Move the second half of keys to new child
|
||||
new_child.keys = full_child.keys.split_off(t);
|
||||
let median = full_child.keys.pop().unwrap();
|
||||
|
||||
// Move the second half of children if not a leaf
|
||||
if !full_child.is_leaf {
|
||||
new_child.children = full_child.children.split_off(t);
|
||||
}
|
||||
|
||||
// Insert median key and new child into parent
|
||||
self.keys.insert(i, median);
|
||||
self.children.insert(i + 1, new_child);
|
||||
}
|
||||
*/
|
||||
}
|
||||
|
|
@ -1,263 +0,0 @@
|
|||
use std::fs::{File, OpenOptions};
|
||||
use std::io::{Read, Seek, SeekFrom, Write};
|
||||
|
||||
use crate::node::Node;
|
||||
use crate::record::Record;
|
||||
|
||||
pub const PAGE_SIZE: usize = 512;
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct NodeStorage {
|
||||
pub file: File,
|
||||
pub page_reads: usize,
|
||||
pub page_writes: usize,
|
||||
}
|
||||
|
||||
impl NodeStorage {
|
||||
pub fn open(path: &str) -> Self {
|
||||
let file = OpenOptions::new()
|
||||
.read(true)
|
||||
.write(true)
|
||||
.create(true)
|
||||
.open(path)
|
||||
.unwrap();
|
||||
|
||||
Self {
|
||||
file,
|
||||
page_reads: 0,
|
||||
page_writes: 0,
|
||||
}
|
||||
}
|
||||
fn serialize_node(node: &Node) -> [u8; PAGE_SIZE] {
|
||||
let mut out = String::new();
|
||||
|
||||
// Format:
|
||||
// L|numkeys|key0;key1;...|child0,child1,...
|
||||
out.push(if node.is_leaf { 'L' } else { 'I' });
|
||||
out.push('|');
|
||||
out.push_str(&node.num_keys.to_string());
|
||||
out.push('|');
|
||||
|
||||
// Keys
|
||||
for i in 0..node.num_keys {
|
||||
let rec = node.keys[i].unwrap();
|
||||
out.push_str(&rec.to_text());
|
||||
if i + 1 < node.num_keys {
|
||||
out.push(';');
|
||||
}
|
||||
}
|
||||
|
||||
out.push('|');
|
||||
|
||||
// Children (m = num_keys, m+1 children)
|
||||
for i in 0..=node.num_keys {
|
||||
match node.children[i] {
|
||||
Some(idx) => out.push_str(&idx.to_string()),
|
||||
None => out.push('.'),
|
||||
}
|
||||
|
||||
if i < node.num_keys {
|
||||
out.push(',');
|
||||
}
|
||||
}
|
||||
|
||||
out.push('|');
|
||||
|
||||
// parent
|
||||
match node.parent {
|
||||
Some(p) => out.push_str(&p.to_string()),
|
||||
None => out.push('.'),
|
||||
}
|
||||
|
||||
// Convert to fixed-size block
|
||||
let mut block = [b' '; PAGE_SIZE];
|
||||
let bytes = out.as_bytes();
|
||||
assert!(bytes.len() <= PAGE_SIZE, "Node too large to serialize");
|
||||
|
||||
block[..bytes.len()].copy_from_slice(bytes);
|
||||
block
|
||||
}
|
||||
|
||||
fn deserialize_node(block: &[u8; PAGE_SIZE]) -> Node {
|
||||
let text = std::str::from_utf8(block).unwrap().trim_end();
|
||||
let parts: Vec<&str> = text.split('|').collect();
|
||||
|
||||
let mut node = Node::new(parts[0] == "L");
|
||||
node.num_keys = parts[1].parse().unwrap();
|
||||
|
||||
if !parts[2].is_empty() {
|
||||
for (i, ks) in parts[2].split(';').enumerate() {
|
||||
let rec = Record::from_text(ks);
|
||||
node.keys[i] = Some(rec);
|
||||
}
|
||||
}
|
||||
|
||||
if !parts[3].is_empty() {
|
||||
for (i, cs) in parts[3].split(',').enumerate() {
|
||||
node.children[i] = if cs == "." {
|
||||
None
|
||||
} else {
|
||||
Some(cs.parse().unwrap())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
node.parent = if parts.len() > 4 && parts[4] != "." {
|
||||
Some(parts[4].parse().unwrap())
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
node
|
||||
}
|
||||
|
||||
pub fn read_node(&mut self, index: usize) -> Node {
|
||||
let offset = (index * PAGE_SIZE) as u64;
|
||||
|
||||
self.file.seek(SeekFrom::Start(offset)).unwrap();
|
||||
let mut block = [0u8; PAGE_SIZE];
|
||||
self.file.read_exact(&mut block).unwrap();
|
||||
|
||||
self.page_reads += 1;
|
||||
|
||||
Self::deserialize_node(&block)
|
||||
}
|
||||
|
||||
pub fn write_node(&mut self, index: usize, node: &Node) {
|
||||
let offset = (index * PAGE_SIZE) as u64;
|
||||
|
||||
let block = Self::serialize_node(node);
|
||||
|
||||
self.file.seek(SeekFrom::Start(offset)).unwrap();
|
||||
self.file.write_all(&block).unwrap();
|
||||
|
||||
self.page_writes += 1;
|
||||
}
|
||||
|
||||
pub fn append_node(&mut self, node: &Node) -> usize {
|
||||
let index = self.num_nodes();
|
||||
self.write_node(index, node);
|
||||
index
|
||||
}
|
||||
|
||||
pub fn num_nodes(&self) -> usize {
|
||||
let len = self.file.metadata().unwrap().len() as usize;
|
||||
len / PAGE_SIZE
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::node::Node;
|
||||
use crate::record::Record;
|
||||
use rand::Rng;
|
||||
use tempfile::tempfile;
|
||||
|
||||
fn random_node() -> Node {
|
||||
let mut node = Node::new(true);
|
||||
let mut rng = rand::rng();
|
||||
let n = rng.random_range(1..=crate::consts::MAX_KEYS);
|
||||
node.num_keys = n;
|
||||
for i in 0..n {
|
||||
node.keys[i] = Some(Record::random());
|
||||
}
|
||||
for i in 0..=n {
|
||||
node.children[i] = Some(rng.random_range(0..100));
|
||||
}
|
||||
|
||||
node.parent = Some(rng.random_range(0..200));
|
||||
|
||||
node
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_append_and_read_single_node() {
|
||||
let file = tempfile().unwrap();
|
||||
let mut storage = NodeStorage {
|
||||
file,
|
||||
page_reads: 0,
|
||||
page_writes: 0,
|
||||
};
|
||||
|
||||
let node = random_node();
|
||||
let index = storage.append_node(&node);
|
||||
|
||||
assert_eq!(index, 0);
|
||||
assert_eq!(storage.page_writes, 1);
|
||||
assert_eq!(storage.num_nodes(), 1);
|
||||
|
||||
let read_node = storage.read_node(0);
|
||||
assert_eq!(read_node.num_keys, node.num_keys);
|
||||
|
||||
for i in 0..node.num_keys {
|
||||
assert_eq!(read_node.keys[i].unwrap().key, node.keys[i].unwrap().key);
|
||||
}
|
||||
assert_eq!(storage.page_reads, 1);
|
||||
assert_eq!(read_node.parent, node.parent);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_overwrite_node() {
|
||||
let file = tempfile().unwrap();
|
||||
let mut storage = NodeStorage {
|
||||
file,
|
||||
page_reads: 0,
|
||||
page_writes: 0,
|
||||
};
|
||||
|
||||
let node1 = random_node();
|
||||
let node2 = random_node();
|
||||
|
||||
let index = storage.append_node(&node1);
|
||||
storage.write_node(index, &node2);
|
||||
|
||||
let read_node = storage.read_node(index);
|
||||
assert_eq!(read_node.num_keys, node2.num_keys);
|
||||
|
||||
for i in 0..node2.num_keys {
|
||||
assert_eq!(read_node.keys[i].unwrap().key, node2.keys[i].unwrap().key);
|
||||
}
|
||||
|
||||
assert_eq!(read_node.parent, node2.parent);
|
||||
|
||||
assert_eq!(storage.page_writes, 2);
|
||||
assert_eq!(storage.page_reads, 1);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_multiple_nodes() {
|
||||
let file = tempfile().unwrap();
|
||||
let mut storage = NodeStorage {
|
||||
file,
|
||||
page_reads: 0,
|
||||
page_writes: 0,
|
||||
};
|
||||
|
||||
let mut nodes = vec![];
|
||||
|
||||
// append 10 random nodes
|
||||
for _ in 0..10 {
|
||||
let node = random_node();
|
||||
storage.append_node(&node);
|
||||
nodes.push(node);
|
||||
}
|
||||
|
||||
assert_eq!(storage.num_nodes(), 10);
|
||||
assert_eq!(storage.page_writes, 10);
|
||||
|
||||
// read them back and verify
|
||||
for i in 0..10 {
|
||||
let read_node = storage.read_node(i);
|
||||
let orig = &nodes[i];
|
||||
|
||||
assert_eq!(read_node.parent, orig.parent);
|
||||
assert_eq!(read_node.num_keys, orig.num_keys);
|
||||
for j in 0..orig.num_keys {
|
||||
assert_eq!(read_node.keys[j].unwrap().key, orig.keys[j].unwrap().key);
|
||||
}
|
||||
}
|
||||
|
||||
assert_eq!(storage.page_reads, 10);
|
||||
}
|
||||
}
|
||||
|
|
@ -1,46 +0,0 @@
|
|||
use rand::Rng;
|
||||
|
||||
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Clone, Copy)]
|
||||
pub struct Record {
|
||||
pub key: i32,
|
||||
pub a: [i32; 5],
|
||||
pub x: i32,
|
||||
}
|
||||
|
||||
impl Record {
|
||||
pub fn new(a: [i32; 5], x: i32) -> Self {
|
||||
let mut key: i32 = 0;
|
||||
let mut x_n: i32 = 1;
|
||||
for i in 0..5 {
|
||||
// rust is annyoing like that, and panics on overflow in arithmetic operations
|
||||
key = key.wrapping_add(a[i].wrapping_mul(x_n));
|
||||
x_n = x_n.wrapping_mul(x);
|
||||
}
|
||||
|
||||
Self { key, a, x }
|
||||
}
|
||||
|
||||
pub fn random() -> Self {
|
||||
let mut rng = rand::rng();
|
||||
let a = rng.random::<[i32; 5]>();
|
||||
let x = rng.random::<i32>();
|
||||
Self::new(a, x)
|
||||
}
|
||||
|
||||
pub fn to_text(&self) -> String {
|
||||
format!(
|
||||
"{},{},{},{},{},{},{}",
|
||||
self.key, self.a[0], self.a[1], self.a[2], self.a[3], self.a[4], self.x
|
||||
)
|
||||
}
|
||||
|
||||
pub fn from_text(s: &str) -> Self {
|
||||
let parts: Vec<i32> = s.split(',').map(|p| p.parse().unwrap()).collect();
|
||||
|
||||
Self {
|
||||
key: parts[0],
|
||||
a: [parts[1], parts[2], parts[3], parts[4], parts[5]],
|
||||
x: parts[6],
|
||||
}
|
||||
}
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue