Skip to content

Commit

Permalink
feat: block io alloc size
Browse files Browse the repository at this point in the history
  • Loading branch information
@peeeuzin
peeeuzin committed Dec 5, 2024
1 parent cc1a137 commit 4b29fc5
Show file tree
Hide file tree
Showing 3 changed files with 170 additions and 73 deletions.
147 changes: 99 additions & 48 deletions src/btree/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,17 @@ use std::{
cmp::{self, Ordering},
fmt::Debug,
marker::PhantomData,
path::Path,
};

pub mod node;
pub mod spec;
use spec::{BTreeCell, BTreePageHeader, PageType, BTREE_PAGE_HEADER_SIZE};
use spec::{
BTreeBlockHeader, BTreeCell, BTreePageHeader, PageType, BTREE_BLOCK_ALLOC_SIZE,
BTREE_PAGE_HEADER_SIZE,
};

pub const MAX_BRANCHING_FACTOR: u16 = 100;

#[derive(Debug)]
pub struct Search {
Expand All @@ -24,26 +30,67 @@ pub struct Search {
pub index: Either<u16, u16>,
}

pub struct BTree<'p, K, V> {
pub struct BTree<K, V> {
root: PageNumber,
b: u16,
io: &'p mut BlockIO,
io: BlockIO,

_k: PhantomData<K>,
_v: PhantomData<V>,
}

impl<
'p,
K: Serialize + DeserializeOwned + PartialOrd + Ord + Clone + Debug,
V: Serialize + DeserializeOwned + PartialOrd + Ord + Clone + Debug,
> BTree<'p, K, V>
> BTree<K, V>
{
pub fn new(io: &'p mut BlockIO, root: PageNumber, b: u16) -> Result<Self> {
pub fn new<P>(block_path: P) -> Result<Self>
where
P: AsRef<Path>,
{
let io: BlockIO =
BlockIO::new(block_path, BTREE_BLOCK_ALLOC_SIZE as u16).map_err(Error::IoError)?;

if io.exists().map_err(Error::IoError)? {
Self::open_block(io)
} else {
Self::create_block(io)
}
}

fn open_block(mut io: BlockIO) -> Result<Self> {
let bytes = io.read_alloc_data().map_err(Error::IoError)?;
let block_header: BTreeBlockHeader =
bincode::deserialize(&bytes).map_err(Error::SerializeError)?;

Ok(Self {
root,
io,
b,
root: block_header.root,
b: block_header.b,
_k: PhantomData,
_v: PhantomData,
})
}

fn create_block(mut io: BlockIO) -> Result<Self> {
let mut root: Page<BTreeCell<K, V>> = Page::create(BTREE_PAGE_HEADER_SIZE).unwrap();
let metadata = BTreePageHeader::new(PageType::Leaf, None);
root.write_special(&metadata.to_bytes()).unwrap();

io.write_page(0, &root).unwrap();

let block_header = BTreeBlockHeader {
b: MAX_BRANCHING_FACTOR,
root: 0,
};

io.alloc_data(&bincode::serialize(&block_header).map_err(Error::SerializeError)?)
.map_err(Error::IoError)?;

Ok(Self {
io,
root: 0,
b: MAX_BRANCHING_FACTOR,
_k: PhantomData,
_v: PhantomData,
})
Expand Down Expand Up @@ -122,15 +169,15 @@ impl<
if root.is_full(self.b) {
let mut splited = root.split(self.b)?;

let sibling_page_index = self
let sibling_page_number = self
.io
.write_new_page(&splited.sibling_node.page)
.map_err(Error::IoError)?;

// create new root
let mut new_root = Page::<BTreeCell<K, V>>::create(BTREE_PAGE_HEADER_SIZE)?;
// set sibling page index to new root right child
let new_root_metadata = BTreePageHeader::new(PageType::Root, Some(sibling_page_index));
let new_root_metadata = BTreePageHeader::new(PageType::Root, Some(sibling_page_number));
new_root.write_special(&new_root_metadata.to_bytes())?;

splited.median_cell.left_child = Some(self.root);
Expand All @@ -143,7 +190,8 @@ impl<
.map_err(Error::IoError)?;

// write new root to disk and set new root index
self.root = self.io.write_new_page(&new_root).map_err(Error::IoError)?;
let root_page = self.io.write_new_page(&new_root).map_err(Error::IoError)?;
self.set_root(root_page)?;
}

self.insert_non_full(self.root, key, value)
Expand All @@ -165,43 +213,43 @@ impl<

match node.header.kind {
PageType::Internal | PageType::Root => {
let page_child_index = node.child(index)?.unwrap();
let page_child_page_number = node.child(index)?.unwrap();

let page_child: BTreeNode<K, V> = self
.io
.read_page(page_child_index.into())
.read_page(page_child_page_number.into())
.map_err(Error::IoError)?
.try_into()?;

if !page_child.is_full(self.b) {
return self.insert_non_full(page_child_index, key, value);
return self.insert_non_full(page_child_page_number, key, value);
}

let mut splited = page_child.split(self.b)?;

let sibling_page_index = self
let sibling_page_number = self
.io
.write_new_page(&splited.sibling_node.page)
.map_err(Error::IoError)?;

node.set_child(index + 1, sibling_page_index)?;
node.set_child(index + 1, sibling_page_number)?;

splited.median_cell.left_child = Some(page_child_index);
splited.median_cell.left_child = Some(page_child_page_number);

let median_cell_key = splited.median_cell.key.clone();
node.page.insert(index, splited.median_cell)?;

self.io
.write_page(page_child_index.into(), &splited.node.page)
.write_page(page_child_page_number.into(), &splited.node.page)
.map_err(Error::IoError)?;

self.io
.write_page(page_number.into(), &node.page)
.map_err(Error::IoError)?;

let insert_page_offset = match key.cmp(&median_cell_key) {
Ordering::Less | Ordering::Equal => page_child_index,
Ordering::Greater => sibling_page_index,
Ordering::Less | Ordering::Equal => page_child_page_number,
Ordering::Greater => sibling_page_number,
};

self.insert_non_full(insert_page_offset, key, value)
Expand Down Expand Up @@ -276,18 +324,22 @@ impl<
key: &K,
parents: &mut Vec<PageNumber>,
) -> Result<()> {
// search key
let search = self.search_from_subtree(key, page_number, parents)?;

// get searched page node
let mut node: BTreeNode<K, V> = self
.io
.read_page(search.page.into())
.map_err(Error::IoError)?
.try_into()?;

// return if not found
if search.index.is_right() {
return Err(Error::NotFound(format!("key {:?}", key)));
}

// delete cell
let index = *search.index.left().unwrap();
node.page.delete(index)?;
let page = node.page.compact()?;
Expand All @@ -296,6 +348,7 @@ impl<
.write_page(search.page.into(), &page)
.map_err(Error::IoError)?;

// check for underflow
self.borrow_if_needed(search.page, parents, key)
}

Expand Down Expand Up @@ -334,11 +387,11 @@ impl<
true => index - 1,
};

let sibling_page_index = parent.child(sibling_index)?.unwrap();
let sibling_page_number = parent.child(sibling_index)?.unwrap();

let sibling: BTreeNode<K, V> = self
.io
.read_page(sibling_page_index.into())
.read_page(sibling_page_number.into())
.map_err(Error::IoError)?
.try_into()?;

Expand All @@ -350,7 +403,7 @@ impl<

let mut merged_node = node.merge(sibling, cell)?;

let merged_page_index = self
let merged_page_number = self
.io
.write_new_page(&merged_node.page)
.map_err(Error::IoError)?;
Expand All @@ -362,12 +415,12 @@ impl<
.write_special(&merged_node.header.to_bytes())?;

self.io
.write_page(merged_page_index.into(), &merged_node.page)
.write_page(merged_page_number.into(), &merged_node.page)
.map_err(Error::IoError)?;

self.root = merged_page_index
self.set_root(merged_page_number)?;
} else {
parent.set_child(merged_node_idx, merged_page_index)?;
parent.set_child(merged_node_idx, merged_page_number)?;
}

self.io
Expand All @@ -380,30 +433,28 @@ impl<

Ok(())
}
}

#[cfg(test)]
mod btree_tests {
use super::*;

#[test]
fn create_btree() {
let mut io = BlockIO::new("test_data/btree/create_btree").unwrap();

let mut root: Page<BTreeCell<u32, String>> = Page::create(BTREE_PAGE_HEADER_SIZE).unwrap();
let metadata = BTreePageHeader::new(PageType::Leaf, None);
root.write_special(&metadata.to_bytes()).unwrap();

io.write_page(0, &root).unwrap();
fn set_root(&mut self, page_number: PageNumber) -> Result<()> {
self.root = page_number;

let mut btree = BTree::<u32, String>::new(&mut io, 0, 10).unwrap();

for i in 0..100 {
btree.insert(i, i.to_string()).unwrap();
}

btree.delete(&99).unwrap();
let block_header = BTreeBlockHeader {
b: MAX_BRANCHING_FACTOR,
root: self.root,
};

println!("{:?}", btree.keys())
self.io
.alloc_data(&bincode::serialize(&block_header).map_err(Error::SerializeError)?)
.map_err(Error::IoError)?;
Ok(())
}
}

// #[cfg(test)]
// mod btree_tests {
// use super::*;

// #[test]
// fn create_btree() {
// let mut btree = BTree::<u32, String>::new("test_data/btree").unwrap();
// }
// }
9 changes: 9 additions & 0 deletions src/btree/spec.rs
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,16 @@
use serde::{Deserialize, Serialize};
use std::mem;

use crate::page::spec::PageNumber;

pub const BTREE_BLOCK_ALLOC_SIZE: usize = mem::size_of::<BTreeBlockHeader>();

#[derive(Clone, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize, Debug)]
pub struct BTreeBlockHeader {
pub root: PageNumber,
pub b: u16,
}

#[derive(Clone, Eq, Ord, PartialEq, PartialOrd, Serialize, Deserialize, Debug)]
pub struct BTreeCell<K, V> {
pub left_child: Option<PageNumber>,
Expand Down
Loading

0 comments on commit 4b29fc5

Please sign in to comment.