test cleanup and bug fix

worldcoin · Apr 17, 2024 · 3aacf09 · 3aacf09
1 parent e4b7db2
commit 3aacf09
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 67 deletions.
diff --git a/src/cascading_merkle_tree.rs b/src/cascading_merkle_tree.rs
@@ -925,75 +925,19 @@ mod tests {
         tree.validate().unwrap();
     }
 
-    #[test]
-    fn test_extend_one_from_slice() {
-        let mut tree = CascadingMerkleTree::<TestHasher>::new(vec![], 10, &1);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-
-        tree.extend_from_slice(&[]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        let expected: Vec<usize> = vec![];
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), expected);
-
-        tree.extend_from_slice(&[2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![2]);
-
-        tree.extend_from_slice(&[2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![2, 2]);
-
-        tree.extend_from_slice(&[2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![2, 2, 2]);
-
-        tree.extend_from_slice(&[2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![2, 2, 2, 2]);
-
-        tree.extend_from_slice(&[2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![2, 2, 2, 2, 2]);
-    }
-
     #[test]
     fn test_extend_from_slice() {
-        let mut tree = CascadingMerkleTree::<TestHasher>::new(vec![], 10, &1);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-
-        tree.extend_from_slice(&[2, 2, 2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![2, 2, 2]);
-
-        tree.extend_from_slice(&[2, 2, 2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![
-            2, 2, 2, 2, 2, 2
-        ]);
-
-        tree.extend_from_slice(&[2, 2, 2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![
-            2, 2, 2, 2, 2, 2, 2, 2, 2
-        ]);
-
-        tree.extend_from_slice(&[2, 2, 2]);
-        debug_tree(&tree);
-        tree.validate().unwrap();
-        assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec![
-            2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2
-        ]);
+        for increment in 1..20 {
+            let mut tree = CascadingMerkleTree::<TestHasher>::new(vec![], 30, &1);
+            let mut vec = vec![];
+            for _ in 0..20 {
+                tree.extend_from_slice(&vec![2; increment]);
+                vec.extend_from_slice(&vec![2; increment]);
+                debug_tree(&tree);
+                tree.validate().unwrap();
+                assert_eq!(tree.leaves().collect::<Vec<usize>>(), vec);
+            }
+        }
     }
 
     #[test]

diff --git a/src/cascading_merkle_tree/storage_ops.rs b/src/cascading_merkle_tree/storage_ops.rs
@@ -371,6 +371,7 @@ pub fn propogate_partial_subtree<H: Hasher>(
             .par_iter_mut()
             .enumerate()
             .for_each(|(i, value)| {
+                let i = i + range.start;
                 let left = &child_layer[2 * i];
                 let right = &child_layer[2 * i + 1];
                 *value = H::hash_node(left, right);