From c1de1c3bca81c7259e23ffbc64517dcdb444b811 Mon Sep 17 00:00:00 2001
From: rymnc <43716372+rymnc@users.noreply.github.com>
Date: Tue, 9 May 2023 22:38:04 +0530
Subject: [PATCH] fix: fn works as expected now

---
 src/tree.rs            | 11 ++++++-----
 tests/memory_keccak.rs | 42 ++++++++++++++++++++++++++++++++++++++++++
 tests/sled_keccak.rs   |  3 ++-
 3 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/tree.rs b/src/tree.rs
index d3b5a59..a2f0b32 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -227,7 +227,7 @@ where
         let end = start + leaves.len();
 
         // check if the leaves are in the correct range
-        if leaves.len() + start > self.capacity() {
+        if leaves.len() + start - to_remove_indices.len() > self.capacity() {
             return Err(PmtreeErrorKind::TreeError(TreeErrorKind::IndexOutOfBounds));
         }
 
@@ -238,7 +238,9 @@ where
 
         // check if the indices are unique
         if to_remove_indices.len() != to_remove_indices.iter().collect::<HashSet<_>>().len() {
-            return Err(PmtreeErrorKind::TreeError(TreeErrorKind::CustomError("Indices are not unique".to_string())));
+            return Err(PmtreeErrorKind::TreeError(TreeErrorKind::CustomError(
+                "Indices are not unique".to_string(),
+            )));
         }
 
         let mut subtree = HashMap::<Key, H::Fr>::new();
@@ -246,13 +248,12 @@ where
         let root_key = Key(0, 0);
 
         subtree.insert(root_key, self.root);
-        
+
         self.fill_nodes(root_key, start, end, &mut subtree, &leaves, start)?;
         for i in to_remove_indices {
-            subtree.insert(Key(self.depth, i - leaves.len()), H::default_leaf());
+            subtree.insert(Key(self.depth, i - leaves.len() + 1), H::default_leaf());
         }
 
-
         let subtree = Arc::new(RwLock::new(subtree));
 
         let root_val = rayon::ThreadPoolBuilder::new()
diff --git a/tests/memory_keccak.rs b/tests/memory_keccak.rs
index 7005920..1ba2e4c 100644
--- a/tests/memory_keccak.rs
+++ b/tests/memory_keccak.rs
@@ -150,3 +150,45 @@ fn set_range() -> PmtreeResult<()> {
 
     Ok(())
 }
+
+
+#[test]
+fn batch_operations() -> PmtreeResult<()> {
+    let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, MemoryDBConfig)?;
+
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000001"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000002"),
+    ];
+
+    mt.batch_insert(None, leaves)?;
+
+    assert_eq!(
+        mt.root(),
+        hex!("893760ec5b5bee236f29e85aef64f17139c3c1b7ff24ce64eb6315fca0f2485b")
+    );
+
+    let leaves = [
+        hex!("0000000000000000000000000000000000000000000000000000000000000003"),
+        hex!("0000000000000000000000000000000000000000000000000000000000000004"),
+    ];
+
+    mt.batch_operations(None, leaves, [])?;
+
+    assert_eq!(
+        mt.root(),
+        hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36")
+    );
+
+    // this should be a no-op
+    mt.batch_operations(None, [leaves[1]], [3])?;
+
+    assert_eq!(
+        mt.root(),
+        hex!("222ff5e0b5877792c2bc1670e2ccd0c2c97cd7bb1672a57d598db05092d3d72c")
+    );
+
+
+    Ok(())
+}
diff --git a/tests/sled_keccak.rs b/tests/sled_keccak.rs
index 8737810..96dd647 100644
--- a/tests/sled_keccak.rs
+++ b/tests/sled_keccak.rs
@@ -230,7 +230,8 @@ fn batch_operations() -> PmtreeResult<()> {
         hex!("a9bb8c3f1f12e9aa903a50c47f314b57610a3ab32f2d463293f58836def38d36")
     );
 
-    mt.batch_operations(None, [], [3])?;
+    // this should be a no-op
+    mt.batch_operations(None, [leaves[1]], [3])?;
 
     assert_eq!(
         mt.root(),