1use super::*;
7use crate::{
8 btree::BTreeTable,
9 column::Column,
10 db::{RcKey, RcValue},
11 error::Result,
12 log::{LogQuery, LogWriter},
13 table::key::TableKeyQuery,
14 Operation,
15};
16
17#[derive(Debug)]
18pub struct BTree {
19 pub(super) depth: u32,
20 pub(super) root_index: Option<Address>,
21 pub(super) record_id: u64,
22}
23
24impl BTree {
25 pub fn new(root_index: Option<Address>, depth: u32, record_id: u64) -> Self {
26 BTree { root_index, depth, record_id }
27 }
28
29 pub fn open(values: TablesRef, log: &impl LogQuery, record_id: u64) -> Result<Self> {
30 let btree_header = BTreeTable::btree_header(log, values)?;
31
32 let root_index =
33 if btree_header.root == NULL_ADDRESS { None } else { Some(btree_header.root) };
34 Ok(BTree::new(root_index, btree_header.depth, record_id))
35 }
36
37 pub fn write_sorted_changes(
38 &mut self,
39 mut changes: &[Operation<RcKey, RcValue>],
40 btree: TablesRef,
41 log: &mut LogWriter,
42 ) -> Result<()> {
43 let mut root = BTree::fetch_root(self.root_index.unwrap_or(NULL_ADDRESS), btree, log)?;
44 let changes = &mut changes;
45
46 while !changes.is_empty() {
47 match root.change(None, self.depth, changes, btree, log)? {
48 (Some((sep, right)), _) => {
49 self.depth += 1;
51 let left = std::mem::take(&mut root);
52 let left_index = self.root_index.take();
53 let new_index = BTreeTable::write_node_plan(btree, left, log, left_index)?;
54 let new_index = if new_index.is_some() { new_index } else { left_index };
55 root.set_child(0, Node::new_child(new_index));
56 root.set_child(1, right);
57 root.set_separator(0, sep);
58 },
59 (_, true) =>
60 if let Some((node_index, node)) = root.need_remove_root(btree, log)? {
61 self.depth -= 1;
62 if let Some(index) = self.root_index.take() {
63 BTreeTable::write_plan_remove_node(btree, log, index)?;
64 }
65 self.root_index = node_index;
66 root = node;
67 },
68 _ => (),
69 }
70 *changes = &changes[1..];
71 }
72
73 if root.changed {
74 let new_index = BTreeTable::write_node_plan(btree, root, log, self.root_index)?;
75
76 if new_index.is_some() {
77 self.root_index = new_index;
78 }
79 }
80 Ok(())
81 }
82
83 #[cfg(test)]
84 pub fn is_balanced(&self, tables: TablesRef, log: &impl LogQuery) -> Result<bool> {
85 let root = BTree::fetch_root(self.root_index.unwrap_or(NULL_ADDRESS), tables, log)?;
86 root.is_balanced(tables, log, 0)
87 }
88
89 pub fn get(
90 &self,
91 key: &[u8],
92 values: TablesRef,
93 log: &impl LogQuery,
94 ) -> Result<Option<Vec<u8>>> {
95 let root = BTree::fetch_root(self.root_index.unwrap_or(NULL_ADDRESS), values, log)?;
96 if let Some(address) = root.get(key, values, log)? {
97 let key_query = TableKeyQuery::Fetch(None);
98 let r = Column::get_value(key_query, address, values, log)?;
99 Ok(r.map(|r| r.1))
100 } else {
101 Ok(None)
102 }
103 }
104
105 pub fn fetch_root(root: Address, tables: TablesRef, log: &impl LogQuery) -> Result<Node> {
106 if root == NULL_ADDRESS {
107 Ok(Node::default())
108 } else {
109 let root = BTreeTable::get_encoded_entry(root, log, tables)?;
110 Node::from_encoded(root)
111 }
112 }
113}