Refactor TreeHashCache splice method

eth2/utils/ssz/src/cached_tree_hash.rs

@@ -124,6 +124,17 @@ impl TreeHashCache {
         })
     }
 
+    pub fn from_bytes(bytes: Vec<u8>, initial_modified_state: bool) -> Result<Self, Error> {
+        if bytes.len() % BYTES_PER_CHUNK > 0 {
+            return Err(Error::BytesAreNotEvenChunks(bytes.len()));
+        }
+
+        Ok(Self {
+            chunk_modified: vec![initial_modified_state; bytes.len() / BYTES_PER_CHUNK],
+            cache: bytes,
+        })
+    }
+
     pub fn bytes_len(&self) -> usize {
         self.cache.len()
     }
@@ -135,22 +146,19 @@ impl TreeHashCache {
             .and_then(|slice| Ok(slice.to_vec()))
     }
 
-    pub fn from_bytes(bytes: Vec<u8>) -> Result<Self, Error> {
-        if bytes.len() % BYTES_PER_CHUNK > 0 {
-            return Err(Error::BytesAreNotEvenChunks(bytes.len()));
-        }
+    pub fn splice(&mut self, chunk_range: Range<usize>, replace_with: Self) {
+        let (bytes, bools) = replace_with.into_components();
 
-        Ok(Self {
-            chunk_modified: vec![false; bytes.len() / BYTES_PER_CHUNK],
-            cache: bytes,
-        })
+        // Update the `chunk_modified` vec, marking all spliced-in nodes as changed.
+        self.chunk_modified.splice(
+            chunk_range.clone(),
+            bools,
+        );
+        self.cache.splice(node_range_to_byte_range(chunk_range), bytes);
     }
 
-    pub fn single_chunk_splice<I>(&mut self, chunk: usize, replace_with: Vec<u8>) {
-        self.chunk_splice(chunk..chunk + 1, replace_with);
-    }
-
-    pub fn chunk_splice(&mut self, chunk_range: Range<usize>, replace_with: Vec<u8>) {
+    /*
+    pub fn byte_splice(&mut self, chunk_range: Range<usize>, replace_with: Vec<u8>) {
         let byte_start = chunk_range.start * BYTES_PER_CHUNK;
         let byte_end = chunk_range.end * BYTES_PER_CHUNK;
 
@@ -161,6 +169,7 @@ impl TreeHashCache {
         );
         self.cache.splice(byte_start..byte_end, replace_with);
     }
+    */
 
     pub fn maybe_update_chunk(&mut self, chunk: usize, to: &[u8]) -> Result<(), Error> {
         let start = chunk * BYTES_PER_CHUNK;
@@ -227,6 +236,10 @@ impl TreeHashCache {
     pub fn into_merkle_tree(self) -> Vec<u8> {
         self.cache
     }
+
+    pub fn into_components(self) -> (Vec<u8>, Vec<bool>) {
+        (self.cache, self.chunk_modified)
+    }
 }
 
 fn children(parent: usize) -> (usize, usize) {
@@ -237,6 +250,10 @@ fn num_nodes(num_leaves: usize) -> usize {
     2 * num_leaves - 1
 }
 
+fn node_range_to_byte_range(node_range: Range<usize>) -> Range<usize> {
+    node_range.start * HASHSIZE..node_range.end * HASHSIZE
+}
+
 #[derive(Debug)]
 pub struct OffsetHandler {
     num_internal_nodes: usize,

eth2/utils/ssz/src/cached_tree_hash/impls.rs

@@ -7,7 +7,10 @@ impl CachedTreeHash<u64> for u64 {
     }
 
     fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
-        Ok(TreeHashCache::from_bytes(merkleize(ssz_encode(self)))?)
+        Ok(TreeHashCache::from_bytes(
+            merkleize(ssz_encode(self)),
+            false,
+        )?)
     }
 
     fn num_bytes(&self) -> usize {
@@ -55,7 +58,7 @@ where
 
     fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
         match T::item_type() {
-            ItemType::Basic => TreeHashCache::from_bytes(merkleize(get_packed_leaves(self))),
+            ItemType::Basic => TreeHashCache::from_bytes(merkleize(get_packed_leaves(self)), false),
             ItemType::Composite | ItemType::List => {
                 let subtrees = self
                     .iter()
@@ -123,7 +126,11 @@ where
                     }
                 }
                 let first_leaf_chunk = offset_handler.first_leaf_node()?;
-                cache.chunk_splice(first_leaf_chunk..offset_handler.next_node, leaves);
+
+                cache.splice(
+                    first_leaf_chunk..offset_handler.next_node,
+                    TreeHashCache::from_bytes(leaves, true)?,
+                );
             }
             ItemType::Composite | ItemType::List => {
                 let mut i = offset_handler.num_leaf_nodes;
@@ -141,14 +148,22 @@ where
                             // Splice out the entire tree of the removed node, replacing it with a
                             // single padding node.
                             let end_chunk = OffsetHandler::new(old, start_chunk)?.next_node();
-                            cache.chunk_splice(start_chunk..end_chunk, vec![0; HASHSIZE]);
+
+                            cache.splice(
+                                start_chunk..end_chunk,
+                                TreeHashCache::from_bytes(vec![0; HASHSIZE], true)?,
+                            );
                         }
                         // The item existed in the previous list but does exist in this list.
                         //
                         // I.e., the list has been lengthened.
                         (None, Some(new)) => {
                             let bytes: Vec<u8> = TreeHashCache::new(new)?.into();
-                            cache.chunk_splice(start_chunk..start_chunk + 1, bytes);
+
+                            cache.splice(
+                                start_chunk..start_chunk + 1,
+                                TreeHashCache::from_bytes(bytes, true)?,
+                            );
                         }
                         // The item didn't exist in the old list and doesn't exist in the new list,
                         // nothing to do.

eth2/utils/ssz/src/cached_tree_hash/tests.rs

@@ -355,7 +355,7 @@ fn test_u64_vec_modifications(original: Vec<u64>, modified: Vec<u64>) {
     let original_cache: Vec<u8> = TreeHashCache::new(&original).unwrap().into();
 
     // Perform a differential hash
-    let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap();
+    let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone(), false).unwrap();
     modified
         .cached_hash_tree_root(&original, &mut cache_struct, 0)
         .unwrap();
@@ -723,7 +723,7 @@ fn generic_test(index: usize) {
         _ => panic!("bad index"),
     };
 
-    let mut cache_struct = TreeHashCache::from_bytes(cache.clone()).unwrap();
+    let mut cache_struct = TreeHashCache::from_bytes(cache.clone(), false).unwrap();
 
     changed_inner
         .cached_hash_tree_root(&inner, &mut cache_struct, 0)