Fix bug with cached tree hash, passes tests

eth2/utils/cached_tree_hash/tests/tests.rs

@@ -10,7 +10,7 @@ pub struct NestedStruct {
 
 fn test_routine<T>(original: T, modified: Vec<T>)
 where
-    T: CachedTreeHash<T>,
+    T: CachedTreeHash<T> + std::fmt::Debug,
 {
     let mut hasher = CachedTreeHasher::new(&original).unwrap();
 
@@ -20,10 +20,23 @@ where
 
     for (i, modified) in modified.iter().enumerate() {
         println!("-- Start of modification {} --", i);
-        // Test after a modification
+
+        // Update the existing hasher.
         hasher
             .update(modified)
             .expect(&format!("Modification {}", i));
+
+        // Create a new hasher from the "modified" struct.
+        let modified_hasher = CachedTreeHasher::new(modified).unwrap();
+
+        // Test that the modified hasher has the same number of chunks as a newly built hasher.
+        assert_eq!(
+            hasher.cache.chunk_modified.len(),
+            modified_hasher.cache.chunk_modified.len(),
+            "Number of chunks is different"
+        );
+
+        // Test the root generated by the updated hasher matches a non-cached tree hash root.
         let standard_root = modified.tree_hash_root();
         let cached_root = hasher
             .tree_hash_root()
@@ -73,7 +86,7 @@ fn test_inner() {
 }
 
 #[test]
-fn test_vec() {
+fn test_vec_of_u64() {
     let original: Vec<u64> = vec![1, 2, 3, 4, 5];
 
     let modified: Vec<Vec<u64>> = vec![
@@ -113,6 +126,29 @@ fn test_nested_list_of_u64() {
     test_routine(original, modified);
 }
 
+#[test]
+fn test_shrinking_vec_of_vec() {
+    let original: Vec<Vec<u64>> = vec![vec![1], vec![2], vec![3], vec![4], vec![5]];
+    let modified: Vec<Vec<u64>> = original[0..3].to_vec();
+
+    let new_hasher = CachedTreeHasher::new(&modified).unwrap();
+
+    let mut modified_hasher = CachedTreeHasher::new(&original).unwrap();
+    modified_hasher.update(&modified).unwrap();
+
+    assert_eq!(
+        new_hasher.cache.schemas.len(),
+        modified_hasher.cache.schemas.len(),
+        "Schema count is different"
+    );
+
+    assert_eq!(
+        new_hasher.cache.chunk_modified.len(),
+        modified_hasher.cache.chunk_modified.len(),
+        "Chunk count is different"
+    );
+}
+
 #[derive(Clone, Debug, TreeHash, CachedTreeHash)]
 pub struct StructWithVec {
     pub a: u64,