use crate::HashElement; use anyhow::Result; use lru::LruCache; use std::any::Any; use std::num::NonZeroUsize; use std::sync::Arc; use tracing::error; pub struct NodeManager { cache: LruCache<(usize, usize), E>, layer_size: Vec, db: Arc>, db_tx: Option>>, } impl NodeManager { pub fn new(db: Arc>, capacity: usize) -> Result { let mut layer = 0; let mut layer_size = Vec::new(); while let Some(size) = db.get_layer_size(layer)? { layer_size.push(size); layer += 1; } Ok(Self { cache: LruCache::new(NonZeroUsize::new(capacity).expect("capacity should be non-zero")), layer_size, db, db_tx: None, }) } pub fn new_dummy() -> Self { Self { cache: LruCache::unbounded(), layer_size: vec![], db: Arc::new(EmptyNodeDatabase {}), db_tx: None, } } pub fn push_node(&mut self, layer: usize, node: E) { self.add_node(layer, self.layer_size[layer], node); self.set_layer_size(layer, self.layer_size[layer] + 1); } pub fn append_nodes(&mut self, layer: usize, nodes: &[E]) { let mut pos = self.layer_size[layer]; let mut saved_nodes = Vec::with_capacity(nodes.len()); for node in nodes { self.cache.put((layer, pos), node.clone()); saved_nodes.push((layer, pos, node)); pos += 1; } self.set_layer_size(layer, pos); self.db_tx().save_node_list(&saved_nodes); } pub fn get_node(&self, layer: usize, pos: usize) -> Option { match self.cache.peek(&(layer, pos)) { Some(node) => Some(node.clone()), None => self.db.get_node(layer, pos).unwrap_or_else(|e| { error!("Failed to get node: {}", e); None }), } } pub fn get_nodes(&self, layer: usize, start_pos: usize, end_pos: usize) -> NodeIterator { NodeIterator { node_manager: self, layer, start_pos, end_pos, } } pub fn add_node(&mut self, layer: usize, pos: usize, node: E) { // No need to insert if the value is unchanged. if self.cache.get(&(layer, pos)) != Some(&node) { self.db_tx().save_node(layer, pos, &node); self.cache.put((layer, pos), node); } } pub fn add_layer(&mut self) { self.layer_size.push(0); let layer = self.layer_size.len() - 1; self.db_tx().save_layer_size(layer, 0); } pub fn layer_size(&self, layer: usize) -> usize { self.layer_size[layer] } pub fn num_layers(&self) -> usize { self.layer_size.len() } pub fn truncate_nodes(&mut self, layer: usize, pos_end: usize) { let mut removed_nodes = Vec::new(); for pos in pos_end..self.layer_size[layer] { self.cache.pop(&(layer, pos)); removed_nodes.push((layer, pos)); } self.db_tx().remove_node_list(&removed_nodes); self.set_layer_size(layer, pos_end); } pub fn truncate_layer(&mut self, layer: usize) { self.truncate_nodes(layer, 0); if layer == self.num_layers() - 1 { self.layer_size.pop(); self.db_tx().remove_layer_size(layer); } } pub fn start_transaction(&mut self) { if self.db_tx.is_some() { error!("start new tx before commit"); panic!("start new tx before commit"); } self.db_tx = Some(self.db.start_transaction()); } pub fn commit(&mut self) { let tx = match self.db_tx.take() { Some(tx) => tx, None => { error!("db_tx is None"); return; } }; if let Err(e) = self.db.commit(tx) { error!("Failed to commit db transaction: {}", e); } } fn db_tx(&mut self) -> &mut dyn NodeTransaction { (*self.db_tx.as_mut().expect("tx checked")).as_mut() } fn set_layer_size(&mut self, layer: usize, size: usize) { self.layer_size[layer] = size; self.db_tx().save_layer_size(layer, size); } } pub struct NodeIterator<'a, E: HashElement> { node_manager: &'a NodeManager, layer: usize, start_pos: usize, end_pos: usize, } impl<'a, E: HashElement> Iterator for NodeIterator<'a, E> { type Item = E; fn next(&mut self) -> Option { if self.start_pos < self.end_pos { let r = self.node_manager.get_node(self.layer, self.start_pos); self.start_pos += 1; r } else { None } } } pub trait NodeDatabase: Send + Sync { fn get_node(&self, layer: usize, pos: usize) -> Result>; fn get_layer_size(&self, layer: usize) -> Result>; fn start_transaction(&self) -> Box>; fn commit(&self, tx: Box>) -> Result<()>; } pub trait NodeTransaction: Send + Sync { fn save_node(&mut self, layer: usize, pos: usize, node: &E); /// `nodes` are a list of tuples `(layer, pos, node)`. fn save_node_list(&mut self, nodes: &[(usize, usize, &E)]); fn remove_node_list(&mut self, nodes: &[(usize, usize)]); fn save_layer_size(&mut self, layer: usize, size: usize); fn remove_layer_size(&mut self, layer: usize); fn into_any(self: Box) -> Box; } /// A dummy database structure for in-memory merkle tree that will not read/write db. pub struct EmptyNodeDatabase {} pub struct EmptyNodeTransaction {} impl NodeDatabase for EmptyNodeDatabase { fn get_node(&self, _layer: usize, _pos: usize) -> Result> { Ok(None) } fn get_layer_size(&self, _layer: usize) -> Result> { Ok(None) } fn start_transaction(&self) -> Box> { Box::new(EmptyNodeTransaction {}) } fn commit(&self, _tx: Box>) -> Result<()> { Ok(()) } } impl NodeTransaction for EmptyNodeTransaction { fn save_node(&mut self, _layer: usize, _pos: usize, _node: &E) {} fn save_node_list(&mut self, _nodes: &[(usize, usize, &E)]) {} fn remove_node_list(&mut self, _nodes: &[(usize, usize)]) {} fn save_layer_size(&mut self, _layer: usize, _size: usize) {} fn remove_layer_size(&mut self, _layer: usize) {} fn into_any(self: Box) -> Box { self } }