feat: make Merkle map thread safe

Author: jpraynaud

mithril-common/src/crypto_helper/merkle_map.rs

@@ -5,7 +5,7 @@ use serde::{Deserialize, Serialize};
 use std::{
     collections::{BTreeMap, HashMap},
     hash::Hash,
-    rc::Rc,
+    sync::Arc,
 };
 
 use crate::{StdError, StdResult};
@@ -266,10 +266,10 @@ impl<K: MKMapKey> From<MKProof> for MKMapProof<K> {
 #[derive(Clone)]
 pub enum MKMapNode<K: MKMapKey> {
     /// A Merkle map
-    Map(Rc<MKMap<K, Self>>),
+    Map(Arc<MKMap<K, Self>>),
 
     /// A full Merkle tree
-    Tree(Rc<MKTree>),
+    Tree(Arc<MKTree>),
 
     /// A Merkle tree node
     TreeNode(MKTreeNode),
@@ -327,13 +327,13 @@ impl<K: MKMapKey> MKMapValue<K> for MKMapNode<K> {
 
 impl<K: MKMapKey> From<MKMap<K, MKMapNode<K>>> for MKMapNode<K> {
     fn from(other: MKMap<K, MKMapNode<K>>) -> Self {
-        MKMapNode::Map(Rc::new(other))
+        MKMapNode::Map(Arc::new(other))
     }
 }
 
 impl<K: MKMapKey> From<MKTree> for MKMapNode<K> {
     fn from(other: MKTree) -> Self {
-        MKMapNode::Tree(Rc::new(other))
+        MKMapNode::Tree(Arc::new(other))
     }
 }
 
@@ -658,7 +658,7 @@ mod tests {
         ];
         let merkle_tree_node_entries = &entries
             .into_iter()
-            .map(|(range, mktree)| (range.to_owned(), MKMapNode::Tree(Rc::new(mktree))))
+            .map(|(range, mktree)| (range.to_owned(), MKMapNode::Tree(Arc::new(mktree))))
             .collect::<Vec<_>>()
             .chunks(10)
             .map(|entries| {
@@ -668,7 +668,7 @@ mod tests {
                         .fold(BlockRange::new(0, 0), |acc, (range, _)| {
                             acc.try_add(range).unwrap()
                         }),
-                    MKMapNode::Map(Rc::new(MKMap::new(entries).unwrap())),
+                    MKMapNode::Map(Arc::new(MKMap::new(entries).unwrap())),
                 )
             })
             .collect::<Vec<_>>();

mithril-common/src/crypto_helper/merkle_tree.rs

@@ -3,14 +3,14 @@ use anyhow::anyhow;
 use anyhow::Context;
 use blake2::{Blake2s256, Digest};
 use ckb_merkle_mountain_range::{
-    util::MemStore, MMRStoreReadOps, MMRStoreWriteOps, Merge, MerkleProof, Result as MMRResult, MMR,
+    MMRStoreReadOps, MMRStoreWriteOps, Merge, MerkleProof, Result as MMRResult, MMR,
 };
 use serde::{Deserialize, Serialize};
 use std::{
     collections::{BTreeMap, HashMap},
     fmt::Display,
     ops::{Add, Deref},
-    sync::Arc,
+    sync::{Arc, RwLock},
 };
 
 use crate::{StdError, StdResult};
@@ -194,28 +194,33 @@ impl From<MKProof> for MKTreeNode {
 
 /// A Merkle tree store
 pub struct MKTreeStore<T> {
-    inner_store: MemStore<T>,
+    inner_store: RwLock<HashMap<u64, T>>,
 }
 
 impl<T> MKTreeStore<T> {
     fn new() -> Self {
         Self {
-            inner_store: MemStore::<T>::default(),
+            inner_store: RwLock::new(HashMap::new()),
         }
     }
 }
 
 impl<T: Clone> MMRStoreReadOps<T> for MKTreeStore<T> {
     fn get_elem(&self, pos: u64) -> MMRResult<Option<T>> {
-        let inner_store = &self.inner_store;
-        inner_store.get_elem(pos)
+        let inner_store = self.inner_store.read().unwrap();
+
+        Ok((*inner_store).get(&pos).cloned())
     }
 }
 
 impl<T> MMRStoreWriteOps<T> for MKTreeStore<T> {
     fn append(&mut self, pos: u64, elems: Vec<T>) -> MMRResult<()> {
-        let mut inner_store = &self.inner_store;
-        inner_store.append(pos, elems)
+        let mut inner_store = self.inner_store.write().unwrap();
+        for (i, elem) in elems.into_iter().enumerate() {
+            (*inner_store).insert(pos + i as u64, elem);
+        }
+
+        Ok(())
     }
 }