From 0d8d3385bef62d4bcb2c094fd7b6ad9d7b6e0801 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Thu, 28 Mar 2019 14:17:25 +1100
Subject: [PATCH] Pass tree hash caching tests

---
 eth2/utils/ssz/src/cached_tree_hash/impls.rs |  8 +++
 eth2/utils/ssz/src/cached_tree_hash/tests.rs | 53 +++++++++++++++++++-
 2 files changed, 59 insertions(+), 2 deletions(-)

diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs
index b27d28c4b..1c2bf342e 100644
--- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs
+++ b/eth2/utils/ssz/src/cached_tree_hash/impls.rs
@@ -85,6 +85,14 @@ where
             }
         }
 
+        // Iterate backwards through the internal nodes, rehashing any node where it's children
+        // have changed.
+        for chunk in (chunk..chunk + num_internal_nodes).into_iter().rev() {
+            if cache.children_modified(chunk)? {
+                cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
+            }
+        }
+
         Some(chunk + num_nodes)
     }
 }
diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs
index f4a4b1d46..6c7567250 100644
--- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs
+++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs
@@ -58,7 +58,7 @@ impl CachedTreeHash for Inner {
 
         // Iterate backwards through the internal nodes, rehashing any node where it's children
         // have changed.
-        for chunk in (0..num_internal_nodes).into_iter().rev() {
+        for chunk in (chunk..chunk + num_internal_nodes).into_iter().rev() {
             if cache.children_modified(chunk)? {
                 cache.modify_chunk(chunk, &cache.hash_children(chunk)?)?;
             }
@@ -77,7 +77,56 @@ fn join(many: Vec<Vec<u8>>) -> Vec<u8> {
 }
 
 #[test]
-fn vec_of_u64() {
+fn partial_modification_u64_vec() {
+    let n: u64 = 50;
+
+    let original_vec: Vec<u64> = (0..n).collect();
+
+    // Generate initial cache.
+    let original_cache = original_vec.build_cache_bytes();
+
+    // Modify the vec
+    let mut modified_vec = original_vec.clone();
+    modified_vec[n as usize - 1] = 42;
+
+    // Perform a differential hash
+    let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap();
+    modified_vec.cached_hash_tree_root(&original_vec, &mut cache_struct, 0);
+    let modified_cache: Vec<u8> = cache_struct.into();
+
+    // Generate reference data.
+    let mut data = vec![];
+    for i in &modified_vec {
+        data.append(&mut int_to_bytes8(*i));
+    }
+    let data = sanitise_bytes(data);
+    let expected = merkleize(data);
+
+    assert_eq!(expected, modified_cache);
+}
+
+#[test]
+fn large_vec_of_u64_builds() {
+    let n: u64 = 50;
+
+    let my_vec: Vec<u64> = (0..n).collect();
+
+    // Generate function output.
+    let cache = my_vec.build_cache_bytes();
+
+    // Generate reference data.
+    let mut data = vec![];
+    for i in &my_vec {
+        data.append(&mut int_to_bytes8(*i));
+    }
+    let data = sanitise_bytes(data);
+    let expected = merkleize(data);
+
+    assert_eq!(expected, cache);
+}
+
+#[test]
+fn vec_of_u64_builds() {
     let data = join(vec![
         int_to_bytes8(1),
         int_to_bytes8(2),