parity_db/btree/
btree.rs

1// Copyright 2021-2022 Parity Technologies (UK) Ltd.
2// This file is dual-licensed as Apache-2.0 or MIT.
3
4//! Btree overlay definition and methods.
5
6use 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					// add one level
50					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}