use anyhow::{bail, Result};
use super::{
cell::{encode_cell, fetch_cell, free_cell, Cell},
encoding::{
decode_page, interior_too_empty, interior_too_full, leaf_too_empty, leaf_too_full, Page,
},
};
use crate::{
encoding::{d_decode_row, d_encode_row},
kv::encoding::{occupancy_interior, occupancy_leaf},
value::Value,
Db,
};
pub fn write_op(store: &mut Db, root_page: u64, op: BTreeWriteOp) -> Result<Option<u64>> {
let new_root = match sub_write(store, root_page, op.clone())? {
WriteOpResult::WasFine(None) if op == BTreeWriteOp::DeleteTree => return Ok(None),
WriteOpResult::WasFine(new_root) => new_root,
WriteOpResult::WasTooFull(le, cell, gt) => {
Some(store.append_interior((le, vec![(cell, gt)]))?)
}
WriteOpResult::WasTooEmptyLeaf(cells) => {
store.append_or_overwrite_leaf(root_page, cells)?
}
WriteOpResult::WasTooEmptyInterior((left, cells)) => {
if cells.is_empty() {
store.free_page(root_page)?;
Some(left)
} else {
store.append_or_overwrite_interior(root_page, (left, cells))?
}
}
};
#[cfg(debug_assertions)]
are_keys_malformed(store, new_root.unwrap_or(root_page)).unwrap();
Ok(new_root)
}
#[derive(Clone, PartialEq, Eq, Debug)]
pub enum BTreeWriteOp {
DeleteTree,
Set(Vec<Value>, Vec<Value>),
Delete(Vec<Value>),
MergeLeaf(Vec<(Cell, Cell)>),
}
#[derive(Debug)]
enum WriteOpResult {
WasFine(Option<u64>),
WasTooFull(u64, Cell, u64),
WasTooEmptyLeaf(Vec<(Cell, Cell)>),
WasTooEmptyInterior((u64, Vec<(Cell, u64)>)),
}
fn sub_write(db: &mut Db, root_page: u64, op: BTreeWriteOp) -> Result<WriteOpResult> {
match decode_page(&db.get_page(root_page)?)? {
Page::BTreeLeaf(_tx_id, mut cells) => {
match op.clone() {
BTreeWriteOp::DeleteTree => {
for (k, v) in cells {
free_cell(db, &k)?;
free_cell(db, &v)?;
}
db.free_page(root_page)?;
return Ok(WriteOpResult::WasFine(None));
}
BTreeWriteOp::Set(key, _) | BTreeWriteOp::Delete(key) => {
let binsearch = cells.binary_search_by(|probe| {
let k = d_decode_row(&fetch_cell(db, &probe.0).unwrap()).unwrap();
k.cmp(&key)
});
match &op {
BTreeWriteOp::Set(_, value) => {
let new_pair = (
encode_cell(db, d_encode_row(&key))?,
encode_cell(db, d_encode_row(&value))?,
);
match binsearch {
Ok(old_idx) => {
free_cell(db, &cells[old_idx].0)?;
free_cell(db, &cells[old_idx].1)?;
cells[old_idx] = new_pair
}
Err(new_idx) => cells.insert(new_idx, new_pair),
};
}
BTreeWriteOp::Delete(_) => {
match binsearch {
Ok(idx) => {
let (k, v) = cells.remove(idx);
free_cell(db, &k)?;
free_cell(db, &v)?;
}
Err(_) => {}
};
}
_ => bail!("b+ tree is malformed"),
}
}
BTreeWriteOp::MergeLeaf(mut new_cells) => {
cells.append(&mut new_cells);
}
}
finish_leaf(db, root_page, cells)
}
Page::BTreeInterior(_tx_id, (mut left, mut cells)) => {
match op.clone() {
BTreeWriteOp::DeleteTree => {
sub_write(db, left, op.clone())?;
for (k, ptr) in cells {
free_cell(db, &k)?;
sub_write(db, ptr, op.clone())?;
}
db.free_page(root_page)?;
return Ok(WriteOpResult::WasFine(None));
}
BTreeWriteOp::Set(key, _) | BTreeWriteOp::Delete(key) => {
let (Ok(id) | Err(id)) = cells.binary_search_by(|probe| {
let k = d_decode_row(&fetch_cell(db, &probe.0).unwrap()).unwrap();
k.cmp(&key)
});
let child_page_no = if id == 0 { left } else { cells[id - 1].1 };
{
let mut need_rewrite = false;
let mut id = id;
let mut war_crime = Some(sub_write(db, child_page_no, op.clone())?);
while let Some(write_op) = war_crime.take() {
match write_op {
WriteOpResult::WasFine(None) => {}
_ => {
need_rewrite = true;
}
}
war_crime = match write_op {
WriteOpResult::WasFine(ptr) => {
if let Some(ptr) = ptr {
if id == 0 {
left = ptr;
} else {
cells[id - 1].1 = ptr;
}
}
None
}
WriteOpResult::WasTooFull(le, cell, gt) => {
if id == 0 {
left = le;
} else {
cells[id - 1].1 = le;
}
cells.insert(id, (cell, gt));
None
}
WriteOpResult::WasTooEmptyLeaf(mut leaf_cells) => {
if cells.is_empty() {
Some(WriteOpResult::WasFine(db.append_or_overwrite_leaf(
if id == 0 { left } else { cells[id - 1].1 },
leaf_cells,
)?))
} else {
Some(if id == 0 {
let (k, ptr) = cells.remove(0);
let Page::BTreeLeaf(_, mut right_cells) =
decode_page(&db.get_page(ptr)?)?
else {
bail!("b+ tree is malformed")
};
db.free_page(ptr)?;
free_cell(db, &k)?;
leaf_cells.append(&mut right_cells);
finish_leaf(db, left, leaf_cells)?
} else {
let (k, ptr) = cells.remove(id - 1);
db.free_page(ptr)?;
free_cell(db, &k)?;
id -= 1;
sub_write(
db,
if id == 0 { left } else { cells[id - 1].1 },
BTreeWriteOp::MergeLeaf(leaf_cells),
)?
})
}
}
WriteOpResult::WasTooEmptyInterior((new_left, mut new_cells)) => {
if cells.is_empty() {
db.free_page(if id == 0 { left } else { cells[id - 1].1 })?;
Some(WriteOpResult::WasFine(Some(new_left)))
} else {
Some(if id == 0 {
let (k, ptr) = cells.remove(0);
let Page::BTreeInterior(_, mut right_interior) =
decode_page(&db.get_page(ptr)?)?
else {
bail!("b+ tree is malformed")
};
db.free_page(ptr)?;
new_cells.push((k, right_interior.0));
new_cells.append(&mut right_interior.1);
finish_interior(db, left, new_left, new_cells)?
} else {
let (k, ptr) = cells.remove(id - 1);
db.free_page(ptr)?;
id -= 1;
let Page::BTreeInterior(_, (left_left, mut left_cells)) =
decode_page(&db.get_page(if id == 0 {
left
} else {
cells[id - 1].1
})?)?
else {
bail!("b+ tree is malformed")
};
left_cells.push((k, new_left));
left_cells.append(&mut new_cells);
finish_interior(
db,
if id == 0 { left } else { cells[id - 1].1 },
left_left,
left_cells,
)?
})
}
}
};
}
if !need_rewrite {
return Ok(WriteOpResult::WasFine(None));
}
}
}
BTreeWriteOp::MergeLeaf(..) => bail!("b+ tree is malformed"),
};
finish_interior(db, root_page, left, cells)
}
_ => bail!("b+ tree is malformed"),
}
}
fn finish_leaf(db: &mut Db, root_page: u64, cells: Vec<(Cell, Cell)>) -> Result<WriteOpResult> {
Ok(if leaf_too_full(&cells) {
let mut left_idx = 0;
let mut right_idx = cells.len() - 1;
while left_idx < right_idx {
let left_occ = occupancy_leaf(&cells[..left_idx]);
let right_occ = occupancy_leaf(&cells[right_idx..]);
if left_occ < right_occ {
left_idx += 1;
} else {
right_idx -= 1;
}
}
db.free_page(root_page)?;
let left_leaf = cells[0..left_idx].to_vec();
let right_leaf = cells[left_idx..].to_vec();
assert!(!leaf_too_full(&left_leaf));
assert!(!leaf_too_full(&right_leaf));
let cells_l = db.append_leaf(left_leaf)?;
let cells_r = db.append_leaf(right_leaf)?;
let cell_content = fetch_cell(db, &cells[0..left_idx].last().unwrap().0)?;
WriteOpResult::WasTooFull(cells_l, encode_cell(db, cell_content)?, cells_r)
} else if leaf_too_empty(&cells) {
WriteOpResult::WasTooEmptyLeaf(cells)
} else {
let id = db.append_or_overwrite_leaf(root_page, cells);
WriteOpResult::WasFine(id?)
})
}
fn finish_interior(
db: &mut Db,
root_page: u64,
left: u64,
cells: Vec<(Cell, u64)>,
) -> Result<WriteOpResult> {
Ok(if interior_too_full(&cells) {
let mut left_idx = 0;
let mut right_idx = cells.len() - 1;
while left_idx < right_idx {
let left_occ = occupancy_interior(&cells[..left_idx]);
let right_occ = occupancy_interior(&cells[right_idx..]);
if left_occ < right_occ {
left_idx += 1;
} else {
right_idx -= 1;
}
}
let middle_one = cells[left_idx].clone();
db.free_page(root_page)?;
let left_interior = (left, cells[0..left_idx].to_vec());
let right_interior = (middle_one.1, cells[left_idx + 1..].to_vec());
assert!(!interior_too_full(&left_interior.1));
assert!(!interior_too_full(&right_interior.1));
let cells_l = db.append_interior(left_interior)?;
let cells_r = db.append_interior(right_interior)?;
let cell_content = fetch_cell(db, &middle_one.0)?;
WriteOpResult::WasTooFull(cells_l, encode_cell(db, cell_content)?, cells_r)
} else if interior_too_empty(&cells) {
WriteOpResult::WasTooEmptyInterior((left, cells))
} else {
let id = db.append_or_overwrite_interior(root_page, (left, cells))?;
WriteOpResult::WasFine(id)
})
}
#[cfg(debug_assertions)]
fn are_keys_malformed(db: &mut Db, root_page: u64) -> Result<Option<(Vec<Value>, Vec<Value>)>> {
use anyhow::bail;
match decode_page(&db.get_page(root_page)?)? {
Page::BTreeLeaf(_tx_id, cells) => {
let mut keys = cells
.into_iter()
.map(|(k, _v)| d_decode_row(&fetch_cell(db, &k)?))
.collect::<Result<Vec<Vec<Value>>>>()?;
keys.sort();
if keys.is_empty() {
return Ok(None);
}
Ok(Some((
keys.first().unwrap().clone(),
keys.last().unwrap().clone(),
)))
}
Page::BTreeInterior(_tx_id, (left, cells)) => {
match (
are_keys_malformed(db, left)?,
cells
.get(0)
.map(|(k, _)| d_decode_row(&fetch_cell(db, &k)?))
.transpose()?,
) {
(Some((lessest, _)), Some(k)) if lessest > k => {
bail!(
"lowest key left of {} is {}",
k.iter()
.map(|e| e.to_string())
.collect::<Vec<_>>()
.join(","),
lessest
.iter()
.map(|e| e.to_string())
.collect::<Vec<_>>()
.join(",")
)
}
_ => {}
}
for (k, ptr) in &cells {
match (
d_decode_row(&fetch_cell(db, &k)?)?,
are_keys_malformed(db, *ptr)?,
) {
(k, Some((lessest, _))) if !(lessest > k) => {
bail!(
"lowest key left of {} is {}",
k.iter()
.map(|e| e.to_string())
.collect::<Vec<_>>()
.join(","),
lessest
.iter()
.map(|e| e.to_string())
.collect::<Vec<_>>()
.join(",")
)
}
_ => {}
}
}
let greatest = cells
.iter()
.flat_map(|(_, ptr)| are_keys_malformed(db, *ptr).unwrap())
.map(|(_, gt)| gt)
.max();
let lessest = cells
.iter()
.flat_map(|(_, ptr)| are_keys_malformed(db, *ptr).unwrap())
.map(|(lt, _)| lt)
.min();
if let (Some(lessest), Some(greatest)) = (greatest, lessest) {
Ok(Some((lessest, greatest)))
} else {
Ok(None)
}
}
_ => unreachable!(),
}
}
#[cfg(debug_assertions)]
use std::collections::BTreeSet;
#[cfg(debug_assertions)]
pub fn did_we_leak_memory(db: &mut Db) -> Result<(), BTreeSet<u64>> {
fn explore_all_branches(db: &mut Db, inaccessible_pages: &mut BTreeSet<u64>, root_page: u64) {
inaccessible_pages.remove(&root_page);
match decode_page(&db.get_page(root_page).unwrap()).unwrap() {
Page::BTreeLeaf(_, cells) => {
for (k, v) in cells {
access_overflow_pages(db, inaccessible_pages, &k);
access_overflow_pages(db, inaccessible_pages, &v);
}
}
Page::BTreeInterior(_, (left, cells)) => {
explore_all_branches(db, inaccessible_pages, left);
for (k, ptr) in cells {
access_overflow_pages(db, inaccessible_pages, &k);
explore_all_branches(db, inaccessible_pages, ptr);
}
}
_ => unreachable!(),
}
}
fn access_overflow_pages(db: &mut Db, inaccessible_pages: &mut BTreeSet<u64>, cell: &Cell) {
match cell {
Cell::Full(_) => {}
Cell::Partial(_, page, _) => {
let mut next = Some(*page);
while let Some(ptr) = next {
inaccessible_pages.remove(&ptr);
let Page::Overflow(_, next_next_ptr, _) =
decode_page(&db.get_page(ptr).unwrap()).unwrap()
else {
unreachable!()
};
next = next_next_ptr;
}
}
}
}
use std::collections::BTreeSet;
use crate::{
catalog::MATDB_TABLES,
kv::{Range, RangeIterKV},
};
let mut inaccessible_pages = BTreeSet::new();
for i in 0..db.page_count {
inaccessible_pages.insert(i);
}
let mut officially_free_pages: BTreeSet<u64> = BTreeSet::new();
officially_free_pages.extend(db.free_list.iter());
officially_free_pages.extend(db.next_free_list.iter());
let mut next_free_list_ptr = db.next_freelist_ptr;
while let Some(ptr) = next_free_list_ptr {
inaccessible_pages.remove(&ptr);
let Page::Freelist(next_ptr, free_pages) = decode_page(&db.get_page(ptr).unwrap()).unwrap()
else {
unreachable!()
};
next_free_list_ptr = next_ptr;
officially_free_pages.extend(free_pages);
}
let mut all_the_table_roots = Vec::new();
let mut iter = RangeIterKV::new(
Range {
backwards: false,
prefix: vec![],
start: None,
end: None,
},
db.this_tx_id,
MATDB_TABLES.into(),
);
while let Some(root) = iter.next(db).unwrap() {
all_the_table_roots.push(root.1[1].clone().into_i64().unwrap() as u64);
}
all_the_table_roots.push(db.main_tree_root);
for root in all_the_table_roots {
explore_all_branches(db, &mut inaccessible_pages, root);
}
let mut leaked = inaccessible_pages
.difference(&officially_free_pages)
.cloned()
.collect::<BTreeSet<u64>>();
for page in leaked.clone() {
let decoded_page = decode_page(&db.get_page(page).unwrap()).unwrap();
let does_count = match decoded_page {
Page::Header | Page::Footer(_, _, _, _) => false,
Page::BTreeLeaf(tx_id, _)
| Page::Overflow(tx_id, _, _)
| Page::BTreeInterior(tx_id, _) => db.this_tx_id == tx_id,
Page::Freelist(_, _) => true,
};
if !does_count {
leaked.remove(&page);
} else {
println!("{decoded_page:?}");
}
}
if leaked.is_empty() {
Ok(())
} else {
Err(leaked)
}
}