From 0d3f2e065744cf78d95dd5b4f0b41ea5a2c8a887 Mon Sep 17 00:00:00 2001
From: Magamedrasul Ibragimov <curryrasul@gmail.com>
Date: Thu, 2 Feb 2023 15:47:19 +0400
Subject: [PATCH] test: implement batch_insert for memory_keccak

---
 src/database.rs        |  2 +-
 src/tree.rs            |  7 ++++++-
 tests/memory_keccak.rs | 20 +++++++++++++++++++-
 tests/poseidon.rs      |  6 ++++--
 tests/sled_keccak.rs   |  2 +-
 5 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/database.rs b/src/database.rs
index 6b82857..60f5ee4 100644
--- a/src/database.rs
+++ b/src/database.rs
@@ -21,5 +21,5 @@ pub trait Database {
     fn put(&mut self, key: DBKey, value: Value) -> Result<()>;
 
     /// Batc
-    fn put_batch(&mut self, subtree: &HashMap<DBKey, Value>) -> Result<()>;
+    fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()>;
 }
diff --git a/src/tree.rs b/src/tree.rs
index 48581b9..3bec22e 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -221,7 +221,12 @@ where
 
         let subtree = RwLock::into_inner(Arc::try_unwrap(subtree).unwrap()).unwrap();
 
-        self.db.put_batch(&subtree.into_iter().)?;
+        self.db.put_batch(
+            subtree
+                .into_iter()
+                .map(|(key, value)| (key.into(), H::serialize(value)))
+                .collect(),
+        )?;
 
         // Update root value and next_index in memory
         self.root = root_val;
diff --git a/tests/memory_keccak.rs b/tests/memory_keccak.rs
index 342e275..829bb1c 100644
--- a/tests/memory_keccak.rs
+++ b/tests/memory_keccak.rs
@@ -25,7 +25,9 @@ impl Database for MemoryDB {
         Ok(())
     }
 
-    fn put_batch(&mut self, subtree: &HashMap<DBKey, Value>) -> Result<()> {
+    fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()> {
+        self.0.extend(subtree.into_iter());
+
         Ok(())
     }
 }
@@ -101,3 +103,19 @@ fn insert_delete() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn batch_insertions() -> Result<()> {
+    let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, "abacaba")?;
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000004"),
+    ];
+
+    mt.batch_insert(&leaves)?;
+
+    Ok(())
+}
diff --git a/tests/poseidon.rs b/tests/poseidon.rs
index cab941e..598f531 100644
--- a/tests/poseidon.rs
+++ b/tests/poseidon.rs
@@ -62,7 +62,9 @@ impl Database for MemoryDB {
         Ok(())
     }
 
-    fn put_batch(&mut self, subtree: &HashMap<DBKey, Value>) -> Result<()> {
+    fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()> {
+        self.0.extend(subtree.into_iter());
+
         Ok(())
     }
 }
@@ -100,7 +102,7 @@ impl Database for MySled {
         Ok(())
     }
 
-    fn put_batch(&mut self, subtree: &HashMap<DBKey, Value>) -> Result<()> {
+    fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()> {
         Ok(())
     }
 }
diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs
index fbe49aa..4f93af7 100644
--- a/tests/sled_keccak.rs
+++ b/tests/sled_keccak.rs
@@ -40,7 +40,7 @@ impl Database for MySled {
         Ok(())
     }
 
-    fn put_batch(&mut self, subtree: &HashMap<DBKey, Value>) -> Result<()> {
+    fn put_batch(&mut self, subtree: HashMap<DBKey, Value>) -> Result<()> {
         Ok(())
     }
 }