Add trait.

common/append_merkle/src/lib.rs

@@ -1,4 +1,5 @@
 mod merkle_tree;
+mod node_manager;
 mod proof;
 mod sha3;
 
@@ -7,17 +8,20 @@ use std::cmp::Ordering;
 use std::collections::{BTreeMap, HashMap};
 use std::fmt::Debug;
 use std::marker::PhantomData;
+use std::sync::Arc;
 use tracing::{trace, warn};
 
 pub use crate::merkle_tree::{
     Algorithm, HashElement, MerkleTreeInitialData, MerkleTreeRead, ZERO_HASHES,
 };
+pub use crate::node_manager::{EmptyNodeDatabase, NodeDatabase, NodeManager};
 pub use proof::{Proof, RangeProof};
 pub use sha3::Sha3Algorithm;
 
 pub struct AppendMerkleTree<E: HashElement, A: Algorithm<E>> {
     /// Keep all the nodes in the latest version. `layers[0]` is the layer of leaves.
     layers: Vec<Vec<E>>,
+    node_manager: NodeManager<E>,
     /// Keep the delta nodes that can be used to construct a history tree.
     /// The key is the root node of that version.
     delta_nodes_map: BTreeMap<u64, DeltaNodes<E>>,
@@ -33,9 +37,15 @@ pub struct AppendMerkleTree<E: HashElement, A: Algorithm<E>> {
 }
 
 impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
-    pub fn new(leaves: Vec<E>, leaf_height: usize, start_tx_seq: Option<u64>) -> Self {
+    pub fn new(
+        node_db: Arc<dyn NodeDatabase<E>>,
+        leaves: Vec<E>,
+        leaf_height: usize,
+        start_tx_seq: Option<u64>,
+    ) -> Self {
         let mut merkle = Self {
             layers: vec![leaves],
+            node_manager: NodeManager::new(node_db),
             delta_nodes_map: BTreeMap::new(),
             root_to_tx_seq_map: HashMap::new(),
             min_depth: None,
@@ -62,12 +72,14 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
     }
 
     pub fn new_with_subtrees(
+        node_db: Arc<dyn NodeDatabase<E>>,
         initial_data: MerkleTreeInitialData<E>,
         leaf_height: usize,
         start_tx_seq: Option<u64>,
     ) -> Result<Self> {
         let mut merkle = Self {
             layers: vec![vec![]],
+            node_manager: NodeManager::new(node_db),
             delta_nodes_map: BTreeMap::new(),
             root_to_tx_seq_map: HashMap::new(),
             min_depth: None,
@@ -103,6 +115,8 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
             // Create an empty merkle tree with `depth`.
             let mut merkle = Self {
                 layers: vec![vec![]; depth],
+                // dummy node manager for the last chunk.
+                node_manager: NodeManager::new(Arc::new(EmptyNodeDatabase {})),
                 delta_nodes_map: BTreeMap::new(),
                 root_to_tx_seq_map: HashMap::new(),
                 min_depth: Some(depth),
@@ -123,6 +137,8 @@ impl<E: HashElement, A: Algorithm<E>> AppendMerkleTree<E, A> {
             layers[0] = leaves;
             let mut merkle = Self {
                 layers,
+                // dummy node manager for the last chunk.
+                node_manager: NodeManager::new(Arc::new(EmptyNodeDatabase {})),
                 delta_nodes_map: BTreeMap::new(),
                 root_to_tx_seq_map: HashMap::new(),
                 min_depth: Some(depth),
@@ -699,9 +715,11 @@ macro_rules! ensure_eq {
 #[cfg(test)]
 mod tests {
     use crate::merkle_tree::MerkleTreeRead;
+    use crate::node_manager::EmptyNodeDatabase;
     use crate::sha3::Sha3Algorithm;
     use crate::AppendMerkleTree;
     use ethereum_types::H256;
+    use std::sync::Arc;
 
     #[test]
     fn test_proof() {
@@ -711,8 +729,12 @@ mod tests {
             for _ in 0..entry_len {
                 data.push(H256::random());
             }
-            let mut merkle =
+            let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(
-                AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
+                Arc::new(EmptyNodeDatabase {}),
+                vec![H256::zero()],
+                0,
+                None,
+            );
             merkle.append_list(data.clone());
             merkle.commit(Some(0));
             verify(&data, &mut merkle);
@@ -739,8 +761,12 @@ mod tests {
             for _ in 0..entry_len {
                 data.push(H256::random());
             }
-            let mut merkle =
+            let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(
-                AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
+                Arc::new(EmptyNodeDatabase {}),
+                vec![H256::zero()],
+                0,
+                None,
+            );
             merkle.append_list(data.clone());
             merkle.commit(Some(0));
 
@@ -764,7 +790,12 @@ mod tests {
     #[test]
     fn test_proof_at_version() {
         let n = [2, 255, 256, 257];
-        let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
+        let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(
+            Arc::new(EmptyNodeDatabase {}),
+            vec![H256::zero()],
+            0,
+            None,
+        );
         let mut start_pos = 0;
 
         for (tx_seq, &entry_len) in n.iter().enumerate() {

common/append_merkle/src/node_manager.rs

@@ -0,0 +1,53 @@
+use crate::HashElement;
+use anyhow::Result;
+use std::collections::HashMap;
+use std::sync::Arc;
+use tracing::error;
+
+pub struct NodeManager<E: HashElement> {
+    cache: HashMap<(usize, usize), E>,
+    db: Arc<dyn NodeDatabase<E>>,
+}
+
+impl<E: HashElement> NodeManager<E> {
+    pub fn new(db: Arc<dyn NodeDatabase<E>>) -> Self {
+        Self {
+            cache: HashMap::new(),
+            db,
+        }
+    }
+
+    pub fn get_node(&self, layer: usize, pos: usize) -> Option<E> {
+        match self.cache.get(&(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 add_node(&mut self, layer: usize, pos: usize, node: E) {
+        if let Err(e) = self.db.save_node(layer, pos, &node) {
+            error!("Failed to save node: {}", e);
+        }
+        self.cache.insert((layer, pos), node);
+    }
+}
+
+pub trait NodeDatabase<E: HashElement>: Send + Sync {
+    fn get_node(&self, layer: usize, pos: usize) -> Result<Option<E>>;
+    fn save_node(&self, layer: usize, pos: usize, node: &E) -> Result<()>;
+}
+
+/// A dummy database structure for in-memory merkle tree that will not read/write db.
+pub struct EmptyNodeDatabase {}
+impl<E: HashElement> NodeDatabase<E> for EmptyNodeDatabase {
+    fn get_node(&self, _layer: usize, _pos: usize) -> Result<Option<E>> {
+        Ok(None)
+    }
+
+    fn save_node(&self, _layer: usize, _pos: usize, _node: &E) -> Result<()> {
+        Ok(())
+    }
+}

node/storage/src/log_store/flow_store.rs

@@ -10,7 +10,7 @@ use crate::log_store::log_manager::{
 use crate::log_store::{FlowRead, FlowSeal, FlowWrite};
 use crate::{try_option, ZgsKeyValueDB};
 use anyhow::{anyhow, bail, Result};
-use append_merkle::{MerkleTreeInitialData, MerkleTreeRead};
+use append_merkle::{EmptyNodeDatabase, MerkleTreeInitialData, MerkleTreeRead, NodeDatabase};
 use itertools::Itertools;
 use parking_lot::RwLock;
 use shared_types::{ChunkArray, DataRoot, FlowProof, Merkle};
@@ -25,15 +25,15 @@ use tracing::{debug, error, trace};
 use zgs_spec::{BYTES_PER_SECTOR, SEALS_PER_LOAD, SECTORS_PER_LOAD, SECTORS_PER_SEAL};
 
 pub struct FlowStore {
-    db: FlowDBStore,
+    db: Arc<FlowDBStore>,
     seal_manager: SealTaskManager,
     config: FlowConfig,
 }
 
 impl FlowStore {
-    pub fn new(db: Arc<dyn ZgsKeyValueDB>, config: FlowConfig) -> Self {
+    pub fn new(db: Arc<FlowDBStore>, config: FlowConfig) -> Self {
         Self {
-            db: FlowDBStore::new(db),
+            db,
             seal_manager: Default::default(),
             config,
         }
@@ -436,7 +436,13 @@ impl FlowDBStore {
         let mut expected_index = 0;
 
         let empty_data = vec![0; PORA_CHUNK_SIZE * ENTRY_SIZE];
-        let empty_root = *Merkle::new(data_to_merkle_leaves(&empty_data)?, 0, None).root();
+        let empty_root = *Merkle::new(
+            Arc::new(EmptyNodeDatabase {}),
+            data_to_merkle_leaves(&empty_data)?,
+            0,
+            None,
+        )
+        .root();
 
         for r in self.kvdb.iter(COL_ENTRY_BATCH_ROOT) {
             let (index_bytes, root_bytes) = r?;
@@ -666,3 +672,22 @@ fn decode_mpt_node_key(data: &[u8]) -> Result<(usize, usize)> {
     let position = try_decode_usize(&data[mem::size_of::<u64>()..])?;
     Ok((layer_index, position))
 }
+
+impl NodeDatabase<DataRoot> for FlowDBStore {
+    fn get_node(&self, layer: usize, pos: usize) -> Result<Option<DataRoot>> {
+        Ok(self
+            .kvdb
+            .get(COL_FLOW_MPT_NODES, &encode_mpt_node_key(layer, pos))?
+            .map(|v| DataRoot::from_slice(&v)))
+    }
+
+    fn save_node(&self, layer: usize, pos: usize, node: &DataRoot) -> Result<()> {
+        let mut tx = self.kvdb.transaction();
+        tx.put(
+            COL_FLOW_MPT_NODES,
+            &encode_mpt_node_key(layer, pos),
+            node.as_bytes(),
+        );
+        Ok(self.kvdb.write(tx)?)
+    }
+}

node/storage/src/log_store/load_chunk/mod.rs

@@ -4,15 +4,15 @@ mod seal;
 mod serde;
 
 use ::serde::{Deserialize, Serialize};
-use std::cmp::min;
-
 use anyhow::Result;
 use ethereum_types::H256;
 use ssz_derive::{Decode, Encode};
+use std::cmp::min;
+use std::sync::Arc;
 
 use crate::log_store::log_manager::data_to_merkle_leaves;
 use crate::try_option;
-use append_merkle::{Algorithm, MerkleTreeRead, Sha3Algorithm};
+use append_merkle::{Algorithm, EmptyNodeDatabase, MerkleTreeRead, Sha3Algorithm};
 use shared_types::{ChunkArray, DataRoot, Merkle};
 use tracing::trace;
 use zgs_spec::{
@@ -248,7 +248,7 @@ impl EntryBatch {
         } else {
             vec![]
         };
-        let mut merkle = Merkle::new(initial_leaves, 0, None);
+        let mut merkle = Merkle::new(Arc::new(EmptyNodeDatabase {}), initial_leaves, 0, None);
         for subtree in self.data.get_subtree_list() {
             trace!(?subtree, "get subtree, leaves={}", merkle.leaves());
             if subtree.start_sector != merkle.leaves() {

node/storage/src/log_store/log_manager.rs

@@ -1,12 +1,12 @@
 use crate::config::ShardConfig;
-use crate::log_store::flow_store::{batch_iter_sharded, FlowConfig, FlowStore};
+use crate::log_store::flow_store::{batch_iter_sharded, FlowConfig, FlowDBStore, FlowStore};
 use crate::log_store::tx_store::TransactionStore;
 use crate::log_store::{
     FlowRead, FlowWrite, LogStoreChunkRead, LogStoreChunkWrite, LogStoreRead, LogStoreWrite,
 };
 use crate::{try_option, ZgsKeyValueDB};
 use anyhow::{anyhow, bail, Result};
-use append_merkle::{Algorithm, MerkleTreeRead, Sha3Algorithm};
+use append_merkle::{Algorithm, EmptyNodeDatabase, MerkleTreeRead, Sha3Algorithm};
 use ethereum_types::H256;
 use kvdb_rocksdb::{Database, DatabaseConfig};
 use merkle_light::merkle::{log2_pow2, MerkleTree};
@@ -626,7 +626,8 @@ impl LogManager {
         executor: task_executor::TaskExecutor,
     ) -> Result<Self> {
         let tx_store = TransactionStore::new(db.clone())?;
-        let flow_store = Arc::new(FlowStore::new(db.clone(), config.flow));
+        let flow_db = Arc::new(FlowDBStore::new(db.clone()));
+        let flow_store = Arc::new(FlowStore::new(flow_db.clone(), config.flow));
         let mut initial_data = flow_store.get_chunk_root_list()?;
         // If the last tx `put_tx` does not complete, we will revert it in `initial_data.subtree_list`
         // first and call `put_tx` later. The known leaves in its data will be saved in `extra_leaves`
@@ -705,8 +706,12 @@ impl LogManager {
             }
         }
 
-        let mut pora_chunks_merkle =
+        let mut pora_chunks_merkle = Merkle::new_with_subtrees(
-            Merkle::new_with_subtrees(initial_data, log2_pow2(PORA_CHUNK_SIZE), start_tx_seq)?;
+            flow_db,
+            initial_data,
+            log2_pow2(PORA_CHUNK_SIZE),
+            start_tx_seq,
+        )?;
         let last_chunk_merkle = match start_tx_seq {
             Some(tx_seq) => {
                 tx_store.rebuild_last_chunk_merkle(pora_chunks_merkle.leaves(), tx_seq)?
@@ -977,7 +982,13 @@ impl LogManager {
                         let data = pad_data[start_index * ENTRY_SIZE
                             ..(start_index + PORA_CHUNK_SIZE) * ENTRY_SIZE]
                             .to_vec();
-                        let root = *Merkle::new(data_to_merkle_leaves(&data)?, 0, None).root();
+                        let root = *Merkle::new(
+                            Arc::new(EmptyNodeDatabase {}),
+                            data_to_merkle_leaves(&data)?,
+                            0,
+                            None,
+                        )
+                        .root();
                         merkle.pora_chunks_merkle.append(root);
                         root_map.insert(merkle.pora_chunks_merkle.leaves() - 1, (root, 1));
                         start_index += PORA_CHUNK_SIZE;

node/storage/src/log_store/tests.rs

@@ -3,11 +3,14 @@ use crate::log_store::log_manager::{
     PORA_CHUNK_SIZE,
 };
 use crate::log_store::{LogStoreChunkRead, LogStoreChunkWrite, LogStoreRead, LogStoreWrite};
-use append_merkle::{Algorithm, AppendMerkleTree, MerkleTreeRead, Sha3Algorithm};
+use append_merkle::{
+    Algorithm, AppendMerkleTree, EmptyNodeDatabase, MerkleTreeRead, Sha3Algorithm,
+};
 use ethereum_types::H256;
 use rand::random;
 use shared_types::{compute_padded_chunk_size, ChunkArray, Transaction, CHUNK_SIZE};
 use std::cmp;
+use std::sync::Arc;
 use task_executor::test_utils::TestRuntime;
 
 #[test]
@@ -26,7 +29,12 @@ fn test_put_get() {
         data[i * CHUNK_SIZE] = random();
     }
     let (padded_chunks, _) = compute_padded_chunk_size(data_size);
-    let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(vec![H256::zero()], 0, None);
+    let mut merkle = AppendMerkleTree::<H256, Sha3Algorithm>::new(
+        Arc::new(EmptyNodeDatabase {}),
+        vec![H256::zero()],
+        0,
+        None,
+    );
     merkle.append_list(data_to_merkle_leaves(&LogManager::padding_raw(start_offset - 1)).unwrap());
     let mut data_padded = data.clone();
     data_padded.append(&mut vec![0u8; CHUNK_SIZE]);
@@ -124,6 +132,7 @@ fn test_root() {
         let mt = sub_merkle_tree(&data).unwrap();
         println!("{:?} {}", mt.root(), hex::encode(mt.root()));
         let append_mt = AppendMerkleTree::<H256, Sha3Algorithm>::new(
+            Arc::new(EmptyNodeDatabase {}),
             data_to_merkle_leaves(&data).unwrap(),
             0,
             None,