Add batch operations to imt.rs

crates/imt/src/imt.rs

@@ -1,3 +1,4 @@
+use std::collections::HashSet;
 pub struct IMT {
     nodes: Vec<Vec<IMTNode>>,
     zeroes: Vec<IMTNode>,
@@ -138,6 +139,11 @@ impl IMT {
         self.update(index, self.zeroes[0].clone())
     }
 
+    pub fn batch_delete(&mut self, indices: Vec<usize>) -> Result<(), &'static str> {
+        let updates: Vec<(usize, IMTNode)> = indices.into_iter().map(|index| (index, self.zeroes[0].clone())).collect();
+        self.batch_update(updates)
+    }
+
     pub fn create_proof(&self, index: usize) -> Result<IMTMerkleProof, &'static str> {
         if index >= self.nodes[0].len() {
             return Err("The leaf does not exist in this tree");
@@ -190,6 +196,89 @@ impl IMT {
 
         node == proof.root
     }
+
+    pub fn batch_insert(&mut self, leaves: Vec<IMTNode>) -> Result<(), &'static str> {
+        let old_len = self.nodes[0].len();
+        if old_len + leaves.len() > self.arity.pow(self.depth as u32) {
+            return Err("The tree cannot contain more than arity^depth leaves");
+        }
+
+        // Append all new leaves to the leaf level
+        self.nodes[0].extend(leaves);
+
+        // Update the tree level by level, starting from the first new leaf
+        let mut start_idx = old_len;
+        for level in 0..self.depth {
+            let parent_start_idx = start_idx / self.arity;
+            let parent_end_idx = ((self.nodes[level].len() as f64) / (self.arity as f64)).ceil() as usize;
+            for idx in parent_start_idx..parent_end_idx {
+                let position = idx * self.arity;
+                let children: Vec<_> = (0..self.arity)
+                    .map(|i| {
+                        self.nodes[level]
+                            .get(position + i)
+                            .cloned()
+                            .unwrap_or_else(|| self.zeroes[level].clone())
+                    })
+                    .collect();
+                let node = (self.hash)(children);
+                if idx < self.nodes[level + 1].len() {
+                    self.nodes[level + 1][idx] = node;
+                } else {
+                    self.nodes[level + 1].push(node);
+                }
+            }
+            start_idx = parent_start_idx;
+        }
+
+        Ok(())
+    }
+
+    /// Updates multiple leaves at specified indices with new values.
+    pub fn batch_update(&mut self, updates: Vec<(usize, IMTNode)>) -> Result<(), &'static str> {
+        // Collect indices and validate them in one pass
+        let mut updated_indices = HashSet::new();
+        for &(index, _) in &updates {
+            if index >= self.nodes[0].len() {
+                return Err("Index out of range");
+            }
+            updated_indices.insert(index);
+        }
+
+        // Apply updates to the leaf level
+        for (index, new_value) in updates {
+            self.nodes[0][index] = new_value;
+        }
+
+        // Update the tree level by level using the set of affected indices
+        for level in 0..self.depth {
+            let mut parent_updated_indices: HashSet<usize> = HashSet::new();
+            for &idx in &updated_indices {
+                let parent_idx = idx / self.arity;
+                if !parent_updated_indices.contains(&parent_idx) {
+                    let position = parent_idx * self.arity;
+                    let children: Vec<_> = (0..self.arity)
+                        .map(|i| {
+                            self.nodes[level]
+                                .get(position + i)
+                                .cloned()
+                                .unwrap_or_else(|| self.zeroes[level].clone())
+                        })
+                        .collect();
+                    let node = (self.hash)(children);
+                    if parent_idx < self.nodes[level + 1].len() {
+                        self.nodes[level + 1][parent_idx] = node;
+                    } else {
+                        self.nodes[level + 1].push(node);
+                    }
+                    parent_updated_indices.insert(parent_idx);
+                }
+            }
+            updated_indices = parent_updated_indices;
+        }
+
+        Ok(())
+    }
 }
 
 #[cfg(test)]
@@ -216,14 +305,61 @@ mod tests {
         assert!(imt.insert("leaf1".to_string()).is_ok());
     }
 
+    #[test] 
+    fn test_batch_insert() {
+        let hash: IMTHashFunction = simple_hash_function;
+        let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec![]).unwrap();
+
+        // Insert multiple leaves at once
+        let leaves = vec!["leaf1".to_string(), "leaf2".to_string(), "leaf3".to_string()];
+        let result = imt.batch_insert(leaves);
+        assert!(result.is_ok());
+
+        // Verify the leaves are correctly inserted
+        assert_eq!(
+            imt.leaves(),
+            vec!["leaf1".to_string(), "leaf2".to_string(), "leaf3".to_string()]
+        );
+
+        // Check that the root is computed (specific value depends on hash function)
+        assert!(imt.root().is_some());
+    }
+
     #[test]
     fn test_delete() {
         let hash: IMTHashFunction = simple_hash_function;
         let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec!["leaf1".to_string()]).unwrap();
 
         assert!(imt.delete(0).is_ok());
     }
+
+    #[test]
+    fn test_batch_delete() {
+        let hash: IMTHashFunction = simple_hash_function;
+        let mut imt = IMT::new(
+            hash,
+            3,
+            "zero".to_string(),
+            2,
+            vec!["leaf1".to_string(), "leaf2".to_string(), "leaf3".to_string()],
+        )
+        .unwrap();
+
+        // Delete multiple leaves
+        let indices = vec![0, 2];
+        let result = imt.batch_delete(indices);
+        assert!(result.is_ok());
 
+        // Verify deleted leaves are set to zero
+        assert_eq!(
+            imt.leaves(),
+            vec!["zero".to_string(), "leaf2".to_string(), "zero".to_string()]
+        );
+
+        // Check that the root is updated
+        assert!(imt.root().is_some());
+    }
+
     #[test]
     fn test_update() {
         let hash: IMTHashFunction = simple_hash_function;
@@ -232,6 +368,35 @@ mod tests {
         assert!(imt.update(0, "new_leaf".to_string()).is_ok());
     }
 
+    #[test]
+    fn test_batch_update() {
+        let hash: IMTHashFunction = simple_hash_function;
+        let mut imt = IMT::new(
+            hash,
+            3,
+            "zero".to_string(),
+            2,
+            vec!["leaf1".to_string(), "leaf2".to_string(), "leaf3".to_string()],
+        )
+        .unwrap();
+
+        // Update multiple leaves
+        let updates = vec![
+            (0, "new_leaf1".to_string()),
+            (2, "new_leaf3".to_string()),
+        ];
+        let result = imt.batch_update(updates);
+        assert!(result.is_ok());
+
+        // Verify the updated leaves
+        assert_eq!(
+            imt.leaves(),
+            vec!["new_leaf1".to_string(), "leaf2".to_string(), "new_leaf3".to_string()]
+        );
+
+        // Check that the root is updated
+        assert!(imt.root().is_some());
+    }
     #[test]
     fn test_create_and_verify_proof() {
         let hash: IMTHashFunction = simple_hash_function;
@@ -275,6 +440,24 @@ mod tests {
         assert!(result.is_err());
     }
 
+    #[test]
+    fn test_delete_many_invalid_index() {
+        let hash: IMTHashFunction = simple_hash_function;
+        let mut imt = IMT::new(
+            hash,
+            3,
+            "zero".to_string(),
+            2,
+            vec!["leaf1".to_string(), "leaf2".to_string()],
+        )
+        .unwrap();
+
+        // Attempt to delete an invalid index
+        let indices = vec![0, 2];
+        let result = imt.batch_delete(indices);
+        assert!(result.is_err());
+    }
+
     #[test]
     fn should_not_delete_nonexistent_leaf() {
         let hash: IMTHashFunction = simple_hash_function;
@@ -322,4 +505,27 @@ mod tests {
         assert_eq!(imt.depth(), 3);
         assert_eq!(imt.arity(), 2);
     }
+
+    #[test]
+    fn test_batched_operations_with_proof() {
+        let hash: IMTHashFunction = simple_hash_function;
+        let mut imt = IMT::new(hash, 3, "zero".to_string(), 2, vec![]).unwrap();
+
+        // Perform batched insertions
+        let leaves = vec!["leaf1".to_string(), "leaf2".to_string(), "leaf3".to_string()];
+        imt.batch_insert(leaves).unwrap();
+
+        // Perform batched updates
+        let updates = vec![(1, "new_leaf2".to_string())];
+        imt.batch_update(updates).unwrap();
+
+        // Perform batched deletions
+        let indices = vec![2];
+        imt.batch_delete(indices).unwrap();
+
+        // Create and verify a proof for an existing leaf
+        let proof = imt.create_proof(0).unwrap();
+        assert!(imt.verify_proof(&proof));
+    }
+
 }