use std::io::{Cursor, Write};
use anyhow::{bail, Result};
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
use crate::{Db, PAGE_SIZE};
use super::cell::{read_cell, write_cell, Cell};
pub const FREE_PAGES_IN_FOOTER_COUNT: usize = (PAGE_SIZE - 1 - 8 - 8 - 8) / 8;
pub const FREE_PAGES_IN_FREELIST_COUNT: usize = (PAGE_SIZE - 1 - 8) / 8;
const OVERFLOW_PAGE_HEADER_SIZE: usize = 1 + 8 + 8;
#[repr(u8)]
enum PageType {
Header = 0x4D,
Footer = 0x4E,
BTreeInterior = 0x01,
BTreeLeaf = 0x02,
PayloadOverflow = 0x53,
Freelist = 0x54,
}
#[derive(Debug)]
pub enum Page {
Header,
Footer(u64, u64, Option<u64>, Vec<u64>),
BTreeLeaf(u64, Vec<(Cell, Cell)>),
BTreeInterior(u64, (u64, Vec<(Cell, u64)>)),
Overflow(u64, Option<u64>, Vec<u8>),
Freelist(Option<u64>, Vec<u64>),
}
impl Db {
pub(super) fn add_header(&mut self) -> Result<()> {
self.set_first_page(&encode_page(&Page::Header))
}
pub(super) fn append_footer(&mut self) -> Result<u64> {
self.append_page_at_end(&encode_page(&Page::Footer(
self.main_tree_root,
self.this_tx_id,
self.next_freelist_ptr,
self.free_list
.iter()
.chain(self.next_free_list.iter())
.copied()
.collect(),
)))
}
pub(super) fn append_leaf(&mut self, cells: Vec<(Cell, Cell)>) -> Result<u64> {
self.append_page(&encode_page(&Page::BTreeLeaf(self.this_tx_id, cells)))
}
pub(super) fn append_or_overwrite_leaf(
&mut self,
prev_page_no: u64,
cells: Vec<(Cell, Cell)>,
) -> Result<Option<u64>> {
self.append_or_overwrite_page(
prev_page_no,
&encode_page(&Page::BTreeLeaf(self.this_tx_id, cells)),
)
}
pub(super) fn append_interior(&mut self, interior: (u64, Vec<(Cell, u64)>)) -> Result<u64> {
self.append_page(&encode_page(&Page::BTreeInterior(
self.this_tx_id,
interior,
)))
}
pub(super) fn append_or_overwrite_interior(
&mut self,
prev_page_no: u64,
interior: (u64, Vec<(Cell, u64)>),
) -> Result<Option<u64>> {
self.append_or_overwrite_page(
prev_page_no,
&encode_page(&Page::BTreeInterior(self.this_tx_id, interior)),
)
}
pub(super) fn append_freelist(
&mut self,
next: Option<u64>,
free_pages: Vec<u64>,
) -> Result<u64> {
self.append_page(&encode_page(&Page::Freelist(next, free_pages)))
}
pub(super) fn append_overflow_pages(&mut self, bytes: Vec<u8>) -> Result<u64> {
let mut prev_page = 0;
for chunk in bytes.chunks(PAGE_SIZE - OVERFLOW_PAGE_HEADER_SIZE).rev() {
prev_page = self.append_page(&encode_page(&Page::Overflow(
self.this_tx_id,
if prev_page == 0 {
None
} else {
Some(prev_page)
},
chunk.to_vec(),
)))?;
}
Ok(prev_page)
}
pub(super) fn read_overflow_pages(&mut self, first_page: u64, length: u64) -> Result<Vec<u8>> {
let mut buf = Cursor::new(vec![0; length as usize]);
let mut page = Some(first_page);
while let Some(page_no) = page {
match decode_page(&self.get_page(page_no)?)? {
Page::Overflow(_tx_id, next, bytes) => {
page = next;
buf.write_all(if page.is_some() {
&bytes
} else {
&bytes[..length as usize % (PAGE_SIZE - OVERFLOW_PAGE_HEADER_SIZE)]
})
.unwrap();
}
_ => bail!("overflow page next-linked to a non-overflow page"),
}
}
Ok(buf.into_inner())
}
}
pub fn leaf_too_empty(cells: &[(Cell, Cell)]) -> bool {
occupancy_leaf(cells) < PAGE_SIZE / 8
}
pub fn leaf_too_full(cells: &[(Cell, Cell)]) -> bool {
PAGE_SIZE < occupancy_leaf(cells)
}
pub fn interior_too_empty(cells: &[(Cell, u64)]) -> bool {
occupancy_interior(cells) < PAGE_SIZE / 8
}
pub fn interior_too_full(cells: &[(Cell, u64)]) -> bool {
PAGE_SIZE < occupancy_interior(cells)
}
pub fn occupancy_leaf(cells: &[(Cell, Cell)]) -> usize {
1 + 8 + 8 + cells
.iter()
.map(|(k, v)| {
[k, v]
.iter()
.map(|c| match c {
Cell::Full(vec) => 2 + vec.len(),
Cell::Partial(vec, _, _) => 2 + vec.len() + 8 + 8,
})
.sum::<usize>()
})
.sum::<usize>()
}
pub fn occupancy_interior(cells: &[(Cell, u64)]) -> usize {
1 + 8 + 8 + 8 + cells
.iter()
.map(|(k, _)| match k {
Cell::Full(vec) => 2 + vec.len(),
Cell::Partial(vec, _, _) => 2 + vec.len() + 8 + 8,
} + 8 )
.sum::<usize>()
}
pub fn decode_page(page: &[u8]) -> Result<Page> {
let mut c = Cursor::new(page);
let ty_byte = c.read_u8().unwrap();
if ty_byte == PageType::Header as u8 {
Ok(Page::Header)
} else if ty_byte == PageType::Footer as u8 {
let main_tree_root = c.read_u64::<BigEndian>().unwrap();
let tx_id = c.read_u64::<BigEndian>().unwrap();
let next_freelist = c.read_u64::<BigEndian>().unwrap();
let mut free_pages = Vec::new();
for _ in 0..FREE_PAGES_IN_FOOTER_COUNT {
let free_page = c.read_u64::<BigEndian>().unwrap();
if free_page == 0 {
break;
}
free_pages.push(free_page);
}
Ok(Page::Footer(
main_tree_root,
tx_id,
if next_freelist == 0 {
None
} else {
Some(next_freelist)
},
free_pages,
))
} else if ty_byte == PageType::BTreeLeaf as u8 {
let tx_id = c.read_u64::<BigEndian>().unwrap();
let len = c.read_u64::<BigEndian>().unwrap();
let mut cells = Vec::new();
for _ in 0..len {
let k = read_cell(&mut c).unwrap();
let v = read_cell(&mut c).unwrap();
cells.push((k, v));
}
Ok(Page::BTreeLeaf(tx_id, cells))
} else if ty_byte == PageType::BTreeInterior as u8 {
let tx_id = c.read_u64::<BigEndian>().unwrap();
let len = c.read_u64::<BigEndian>().unwrap();
let mut cells = Vec::new();
let left = c.read_u64::<BigEndian>().unwrap();
for _ in 0..len {
let k = read_cell(&mut c).unwrap();
let right = c.read_u64::<BigEndian>().unwrap();
cells.push((k, right));
}
Ok(Page::BTreeInterior(tx_id, (left, cells)))
} else if ty_byte == PageType::PayloadOverflow as u8 {
let tx_id = c.read_u64::<BigEndian>().unwrap();
let next = c.read_u64::<BigEndian>().unwrap();
Ok(Page::Overflow(
tx_id,
if next == 0 { None } else { Some(next) },
page[OVERFLOW_PAGE_HEADER_SIZE..].to_vec(),
))
} else if ty_byte == PageType::Freelist as u8 {
let next = c.read_u64::<BigEndian>().unwrap();
let mut free_pages = Vec::new();
for _ in 0..FREE_PAGES_IN_FREELIST_COUNT {
let free_page = c.read_u64::<BigEndian>().unwrap();
if free_page == 0 {
break;
}
free_pages.push(free_page);
}
Ok(Page::Freelist(
if next == 0 { None } else { Some(next) },
free_pages,
))
} else {
bail!("invalid page type byte: {ty_byte:x?}")
}
}
fn encode_page(page: &Page) -> Vec<u8> {
let mut c = Cursor::new(vec![0u8; PAGE_SIZE]);
match page {
Page::Header => {
c.write_u8(PageType::Header as u8).unwrap();
c.into_inner()
}
Page::Footer(main_tree_root, tx_id, next_freelist, free_pages) => {
c.write_u8(PageType::Footer as u8).unwrap();
c.write_u64::<BigEndian>(*main_tree_root).unwrap();
c.write_u64::<BigEndian>(*tx_id).unwrap();
c.write_u64::<BigEndian>(next_freelist.unwrap_or(0))
.unwrap();
for page in free_pages {
c.write_u64::<BigEndian>(*page).unwrap();
}
c.into_inner()
}
Page::BTreeLeaf(tx_id, cells) => {
c.write_u8(PageType::BTreeLeaf as u8).unwrap();
c.write_u64::<BigEndian>(*tx_id).unwrap();
c.write_u64::<BigEndian>(cells.len() as u64).unwrap();
for (k, v) in cells {
write_cell(&mut c, k);
write_cell(&mut c, v);
}
c.into_inner()
}
Page::BTreeInterior(tx_id, (left, cells)) => {
c.write_u8(PageType::BTreeInterior as u8).unwrap();
c.write_u64::<BigEndian>(*tx_id).unwrap();
c.write_u64::<BigEndian>(cells.len() as u64).unwrap();
c.write_u64::<BigEndian>(*left).unwrap();
for (k, right) in cells {
write_cell(&mut c, k);
c.write_u64::<BigEndian>(*right).unwrap();
}
c.into_inner()
}
Page::Overflow(tx_id, next, bytes) => {
c.write_u8(PageType::PayloadOverflow as u8).unwrap();
c.write_u64::<BigEndian>(*tx_id).unwrap();
c.write_u64::<BigEndian>(next.unwrap_or(0)).unwrap();
c.write_all(&bytes).unwrap();
c.into_inner()
}
Page::Freelist(next, free_pages) => {
c.write_u8(PageType::Freelist as u8).unwrap();
c.write_u64::<BigEndian>(next.unwrap_or(0)).unwrap();
for page in free_pages {
c.write_u64::<BigEndian>(*page).unwrap();
}
c.into_inner()
}
}
}