diff --git a/crates/imt/src/imt.rs b/crates/imt/src/imt.rs index 741c384..ea89465 100644 --- a/crates/imt/src/imt.rs +++ b/crates/imt/src/imt.rs @@ -1,3 +1,4 @@ +use std::collections::HashSet; pub struct IMT { nodes: Vec>, zeroes: Vec, @@ -138,6 +139,11 @@ impl IMT { self.update(index, self.zeroes[0].clone()) } + pub fn batch_delete(&mut self, indices: Vec) -> 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 { 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) -> 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 = 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,6 +305,26 @@ 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; @@ -223,7 +332,34 @@ mod tests { 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)); + } + }