fix(rln): atomic operation edge case

rln/src/pm_tree_adapter.rs

@@ -193,32 +193,60 @@ impl ZerokitMerkleTree for PmTree {
     ) -> Result<()> {
         let leaves = leaves.into_iter().collect::<Vec<_>>();
         let indices = indices.into_iter().collect::<HashSet<_>>();
-        let end = start + leaves.len();
-
-        if leaves.len() + start - indices.len() > self.capacity() {
-            return Err(Report::msg("index out of bounds"));
-        }
-
-        // extend the range to include indices to be removed
-        let min_index = indices.iter().min().unwrap_or(&start);
-        let max_index = indices.iter().max().unwrap_or(&end);
-
-        let mut new_leaves = Vec::new();
-
-        // insert leaves into new_leaves
-        for i in *min_index..*max_index {
-            if indices.contains(&i) {
-                // insert 0
-                new_leaves.push(Self::Hasher::default_leaf());
-            } else {
-                // insert leaf
-                new_leaves.push(leaves[i - start]);
+        let end = start + leaves.len() + indices.len();
+
+        // handle each case appropriately -
+        // case 1: both leaves and indices to be removed are passed in
+        // case 2: only leaves are passed in
+        // case 3: only indices are passed in
+        // case 4: neither leaves nor indices are passed in
+        match (leaves.len(), indices.len()) {
+            (0, 0) => {
+                Err(Report::msg("no leaves or indices to be removed"))
+            }
+            (0, _) => {
+                // case 3
+                // remove indices
+                let mut new_leaves = Vec::new();
+                for i in start..end {
+                    if indices.contains(&i) {
+                        // insert 0
+                        new_leaves.push(Self::Hasher::default_leaf());
+                    } else {
+                        // insert leaf
+                        new_leaves.push(self.tree.get(i)?);
+                    }
+                }
+                self.tree
+                    .set_range(start, new_leaves)
+                    .map_err(|e| Report::msg(e.to_string()))
+            }
+            (_, 0) => {
+                // case 2
+                // insert leaves
+                self.tree
+                    .set_range(start, leaves)
+                    .map_err(|e| Report::msg(e.to_string()))
+            }
+            (_, _) => {
+                // case 1
+                // remove indices
+                let mut new_leaves = Vec::new();
+                let new_start = start + leaves.len();
+                for i in new_start..=end {
+                    if indices.contains(&i) {
+                        // Insert 0
+                        new_leaves.push(Self::Hasher::default_leaf());
+                    } else if let Some(leaf) = leaves.get(i - new_start) {
+                        // Insert leaf
+                        new_leaves.push(*leaf);
+                    }
+                }
+                self.tree
+                    .set_range(start, new_leaves)
+                    .map_err(|e| Report::msg(e.to_string()))
             }
         }
-
-        self.tree
-            .set_range(start, new_leaves)
-            .map_err(|e| Report::msg(e.to_string()))
     }
 
     fn update_next(&mut self, leaf: FrOf<Self::Hasher>) -> Result<()> {

rln/src/public.rs

@@ -362,7 +362,7 @@ impl RLN<'_> {
         // We set the leaves
         self.tree
             .override_range(index, leaves, indices)
-            .map_err(|_| Report::msg("Could not perform the batch operation"))?;
+            .map_err(|e| Report::msg(format!("Could not perform the batch operation: {e}")))?;
         Ok(())
     }
 
@@ -1387,7 +1387,7 @@ mod test {
 
         assert_eq!(root_batch_with_init, root_single_additions);
 
-        rln.flush();
+        rln.flush().unwrap();
     }
 
     #[test]
@@ -1430,7 +1430,7 @@ mod test {
         let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices).unwrap());
         let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&last_leaf).unwrap());
 
-        rln.atomic_operation(no_of_leaves, leaves_buffer, indices_buffer)
+        rln.atomic_operation(last_leaf_index, leaves_buffer, indices_buffer)
             .unwrap();
 
         // We get the root of the tree obtained after a no-op
@@ -1441,6 +1441,43 @@ mod test {
         assert_eq!(root_after_insertion, root_after_noop);
     }
 
+    #[test]
+    fn test_atomic_operation_zero_indexed() {
+        // Test duplicated from https://github.com/waku-org/go-zerokit-rln/pull/12/files
+        let tree_height = TEST_TREE_HEIGHT;
+        let no_of_leaves = 256;
+
+        // We generate a vector of random leaves
+        let mut leaves: Vec<Fr> = Vec::new();
+        let mut rng = thread_rng();
+        for _ in 0..no_of_leaves {
+            leaves.push(Fr::rand(&mut rng));
+        }
+
+        // We create a new tree
+        let input_buffer =
+            Cursor::new(json!({ "resources_folder": TEST_RESOURCES_FOLDER }).to_string());
+        let mut rln = RLN::new(tree_height, input_buffer).unwrap();
+
+        // We add leaves in a batch into the tree
+        let mut buffer = Cursor::new(vec_fr_to_bytes_le(&leaves).unwrap());
+        rln.init_tree_with_leaves(&mut buffer).unwrap();
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        // We check if number of leaves set is consistent
+        assert_eq!(rln.tree.leaves_set(), no_of_leaves);
+
+        let zero_index = 0;
+        let indices = vec![zero_index as u8];
+        let zero_leaf: Vec<Fr> = vec![];
+        let indices_buffer = Cursor::new(vec_u8_to_bytes_le(&indices).unwrap());
+        let leaves_buffer = Cursor::new(vec_fr_to_bytes_le(&zero_leaf).unwrap());
+        rln.atomic_operation(0, leaves_buffer, indices_buffer)
+            .unwrap();
+    }
+
     #[allow(unused_must_use)]
     #[test]
     // This test checks if `set_leaves_from` throws an error when the index is out of bounds

rln/tests/ffi.rs

@@ -265,7 +265,7 @@ mod test {
 
         let success = atomic_operation(
             rln_pointer,
-            no_of_leaves as usize,
+            last_leaf_index as usize,
             leaves_buffer,
             indices_buffer,
         );