From 3aacf0928838ae689e9634c0f3bedcb8a948f617 Mon Sep 17 00:00:00 2001
From: Eric Woolsey <ewoolsey@ualberta.ca>
Date: Tue, 16 Apr 2024 17:03:20 -0700
Subject: [PATCH] test cleanup and bug fix

---
 src/cascading_merkle_tree.rs             | 78 ++++--------------------
 src/cascading_merkle_tree/storage_ops.rs |  1 +
 2 files changed, 12 insertions(+), 67 deletions(-)

diff --git a/src/cascading_merkle_tree.rs b/src/cascading_merkle_tree.rs
index 18e54a1..a4063e4 100644
--- 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
index 1e5fbae..35b862f 100644
--- 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);