diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs index 0588ab772..b676ececc 100644 --- a/eth2/utils/ssz/src/cached_tree_hash.rs +++ b/eth2/utils/ssz/src/cached_tree_hash.rs @@ -124,6 +124,17 @@ impl TreeHashCache { }) } + pub fn from_bytes(bytes: Vec, initial_modified_state: bool) -> Result { + 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) -> Result { - if bytes.len() % BYTES_PER_CHUNK > 0 { - return Err(Error::BytesAreNotEvenChunks(bytes.len())); - } + pub fn splice(&mut self, chunk_range: Range, 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(&mut self, chunk: usize, replace_with: Vec) { - self.chunk_splice(chunk..chunk + 1, replace_with); - } - - pub fn chunk_splice(&mut self, chunk_range: Range, replace_with: Vec) { + /* + pub fn byte_splice(&mut self, chunk_range: Range, replace_with: Vec) { 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 { self.cache } + + pub fn into_components(self) -> (Vec, Vec) { + (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) -> Range { + node_range.start * HASHSIZE..node_range.end * HASHSIZE +} + #[derive(Debug)] pub struct OffsetHandler { num_internal_nodes: usize, diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs index 2d0ab5059..f598de79a 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/impls.rs @@ -7,7 +7,10 @@ impl CachedTreeHash for u64 { } fn build_tree_hash_cache(&self) -> Result { - 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 { 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 = 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. diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs index d48ed9eb8..d784a0889 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs @@ -355,7 +355,7 @@ fn test_u64_vec_modifications(original: Vec, modified: Vec) { let original_cache: Vec = 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)