From 0632a00a48f6da3cd9a3e7ac1df332e6b12aee92 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sun, 14 Apr 2019 18:50:12 +1000
Subject: [PATCH] Fix failing test for shrinking vec of structs

---
 eth2/utils/ssz/src/cached_tree_hash.rs        |  7 ++-
 eth2/utils/ssz/src/cached_tree_hash/impls.rs  | 14 +++---
 eth2/utils/ssz/src/cached_tree_hash/resize.rs |  9 ++--
 eth2/utils/ssz/src/cached_tree_hash/tests.rs  | 43 +++++++++++++++++++
 4 files changed, 59 insertions(+), 14 deletions(-)

diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs
index 0e6bdf986f..7a722766ce 100644
--- a/eth2/utils/ssz/src/cached_tree_hash.rs
+++ b/eth2/utils/ssz/src/cached_tree_hash.rs
@@ -18,6 +18,7 @@ pub enum Error {
     NoBytesForRoot,
     UnableToObtainSlices,
     UnableToGrowMerkleTree,
+    UnableToShrinkMerkleTree,
     BytesAreNotEvenChunks(usize),
     NoModifiedFieldForChunk(usize),
     NoBytesForChunk(usize),
@@ -303,10 +304,14 @@ impl OffsetHandler {
         self.num_leaf_nodes.trailing_zeros() as usize
     }
 
-    pub fn node_range(&self) -> Range<usize> {
+    pub fn chunk_range(&self) -> Range<usize> {
         self.first_node..self.next_node
     }
 
+    pub fn total_chunks(&self) -> usize {
+        self.next_node - self.first_node
+    }
+
     pub fn total_nodes(&self) -> usize {
         self.num_internal_nodes + self.num_leaf_nodes
     }
diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs
index f16e6a62b1..0377649cbe 100644
--- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs
+++ b/eth2/utils/ssz/src/cached_tree_hash/impls.rs
@@ -119,7 +119,7 @@ where
 
             // Get slices of the exsiting tree from the cache.
             let (old_bytes, old_flags) = cache
-                .slices(old_offset_handler.node_range())
+                .slices(old_offset_handler.chunk_range())
                 .ok_or_else(|| Error::UnableToObtainSlices)?;
 
             let (new_bytes, new_flags) =
@@ -137,16 +137,16 @@ where
                         old_flags,
                         old_offset_handler.height(),
                         offset_handler.height(),
-                        offset_handler.total_nodes(),
+                        offset_handler.total_chunks(),
                     )
-                    .ok_or_else(|| Error::UnableToGrowMerkleTree)?
+                    .ok_or_else(|| Error::UnableToShrinkMerkleTree)?
                 };
 
             // Create a `TreeHashCache` from the raw elements.
-            let expanded_cache = TreeHashCache::from_elems(new_bytes, new_flags);
+            let modified_cache = TreeHashCache::from_elems(new_bytes, new_flags);
 
-            // Splice the newly created `TreeHashCache` over the existing, smaller elements.
-            cache.splice(old_offset_handler.node_range(), expanded_cache);
+            // Splice the newly created `TreeHashCache` over the existing elements.
+            cache.splice(old_offset_handler.chunk_range(), modified_cache);
         }
 
         match T::item_type() {
@@ -208,8 +208,6 @@ where
 
         for (&parent, children) in offset_handler.iter_internal_nodes().rev() {
             if cache.either_modified(children)? {
-                dbg!(parent);
-                dbg!(children);
                 cache.modify_chunk(parent, &cache.hash_children(children)?)?;
             }
         }
diff --git a/eth2/utils/ssz/src/cached_tree_hash/resize.rs b/eth2/utils/ssz/src/cached_tree_hash/resize.rs
index 0b492770fd..44b3f0ea59 100644
--- a/eth2/utils/ssz/src/cached_tree_hash/resize.rs
+++ b/eth2/utils/ssz/src/cached_tree_hash/resize.rs
@@ -73,11 +73,10 @@ pub fn shrink_merkle_cache(
     let mut bytes = vec![0; to_nodes * HASHSIZE];
     let mut flags = vec![true; to_nodes];
 
-    let leaf_level = to_height;
-
-    for i in 0..=leaf_level as usize {
+    for i in 0..=to_height as usize {
         let from_i = i + from_height - to_height;
-        let (from_byte_slice, from_flag_slice) = if from_i == leaf_level {
+
+        let (from_byte_slice, from_flag_slice) = if from_i == from_height {
             (
                 from_bytes.get(first_byte_at_height(from_i)..)?,
                 from_flags.get(first_node_at_height(from_i)..)?,
@@ -89,7 +88,7 @@ pub fn shrink_merkle_cache(
             )
         };
 
-        let (to_byte_slice, to_flag_slice) = if i == leaf_level {
+        let (to_byte_slice, to_flag_slice) = if i == to_height {
             (
                 bytes.get_mut(first_byte_at_height(i)..)?,
                 flags.get_mut(first_node_at_height(i)..)?,
diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs
index 22d01ec1a1..f09fac4193 100644
--- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs
+++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs
@@ -563,6 +563,49 @@ fn shortened_vec_of_inner_within_power_of_two_boundary() {
     test_inner_vec_modifications(original, modified, reference_vec);
 }
 
+#[test]
+fn shortened_vec_of_inner_outside_power_of_two_boundary() {
+    let original = vec![
+        Inner {
+            a: 0,
+            b: 1,
+            c: 2,
+            d: 3,
+        },
+        Inner {
+            a: 4,
+            b: 5,
+            c: 6,
+            d: 7,
+        },
+        Inner {
+            a: 8,
+            b: 9,
+            c: 10,
+            d: 11,
+        },
+        Inner {
+            a: 12,
+            b: 13,
+            c: 14,
+            d: 15,
+        },
+        Inner {
+            a: 16,
+            b: 17,
+            c: 18,
+            d: 19,
+        },
+    ];
+
+    let mut modified = original.clone();
+    modified.pop(); // remove the last element from the list.
+
+    let reference_vec: Vec<u64> = (0..16).collect();
+
+    test_inner_vec_modifications(original, modified, reference_vec);
+}
+
 #[test]
 fn lengthened_vec_of_inner_within_power_of_two_boundary() {
     let original = vec![