diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs index 3b900e503..83b516ac7 100644 --- a/eth2/utils/ssz/src/cached_tree_hash.rs +++ b/eth2/utils/ssz/src/cached_tree_hash.rs @@ -25,13 +25,13 @@ pub enum Error { pub trait CachedTreeHash { type Item: CachedTreeHash; - fn build_cache(&self) -> Result; + fn leaves_and_subtrees(&self) -> Vec; /// Return the number of bytes when this element is encoded as raw SSZ _without_ length /// prefixes. fn num_bytes(&self) -> usize; - fn offset_handler(&self, initial_offset: usize) -> Result; + fn offsets(&self) -> Result, Error>; fn num_child_nodes(&self) -> usize; @@ -60,13 +60,17 @@ impl TreeHashCache { where T: CachedTreeHash, { - item.build_cache() + Self::from_leaves_and_subtrees(item.leaves_and_subtrees(), OffsetHandler::new(item, 0)?) } pub fn from_leaves_and_subtrees( mut leaves_and_subtrees: Vec, offset_handler: OffsetHandler, ) -> Result { + // Pad the leaves with zeros if the number of immediate leaf-nodes (without recursing into + // sub-trees) is not an even power-of-two. + pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees); + // Allocate enough bytes to store the internal nodes and the leaves and subtrees, then fill // all the to-be-built internal nodes with zeros and append the leaves and subtrees. let internal_node_bytes = offset_handler.num_internal_nodes * BYTES_PER_CHUNK; @@ -209,6 +213,13 @@ pub struct OffsetHandler { } impl OffsetHandler { + pub fn new(item: &T, initial_offset: usize) -> Result + where + T: CachedTreeHash, + { + Self::from_lengths(initial_offset, item.offsets()?) + } + fn from_lengths(offset: usize, mut lengths: Vec) -> Result { // Extend it to the next power-of-two, if it is not already. let num_leaf_nodes = if lengths.len().is_power_of_two() { diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs index 9c0a8ec6d..54a690c6d 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/impls.rs @@ -4,15 +4,15 @@ use crate::{ssz_encode, Encodable}; impl CachedTreeHash for u64 { type Item = Self; - fn build_cache(&self) -> Result { - TreeHashCache::from_bytes(merkleize(ssz_encode(self))) + fn leaves_and_subtrees(&self) -> Vec { + merkleize(ssz_encode(self)) } fn num_bytes(&self) -> usize { 8 } - fn offset_handler(&self, _initial_offset: usize) -> Result { + fn offsets(&self) -> Result, Error> { Err(Error::ShouldNotProduceOffsetHandler) } diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs index a6be7f9ae..f6c52ef8f 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs @@ -12,19 +12,15 @@ pub struct Inner { impl CachedTreeHash for Inner { type Item = Self; - fn build_cache(&self) -> Result { - let offset_handler = self.offset_handler(0)?; - + fn leaves_and_subtrees(&self) -> Vec { let mut leaves_and_subtrees = vec![]; - leaves_and_subtrees.append(&mut self.a.build_cache()?.into()); - leaves_and_subtrees.append(&mut self.b.build_cache()?.into()); - leaves_and_subtrees.append(&mut self.c.build_cache()?.into()); - leaves_and_subtrees.append(&mut self.d.build_cache()?.into()); + leaves_and_subtrees.append(&mut self.a.leaves_and_subtrees()); + leaves_and_subtrees.append(&mut self.b.leaves_and_subtrees()); + leaves_and_subtrees.append(&mut self.c.leaves_and_subtrees()); + leaves_and_subtrees.append(&mut self.d.leaves_and_subtrees()); - pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees); - - TreeHashCache::from_leaves_and_subtrees(leaves_and_subtrees, self.offset_handler(0)?) + leaves_and_subtrees } fn num_bytes(&self) -> usize { @@ -38,7 +34,7 @@ impl CachedTreeHash for Inner { bytes } - fn offset_handler(&self, initial_offset: usize) -> Result { + fn offsets(&self) -> Result, Error> { let mut offsets = vec![]; offsets.push(self.a.num_child_nodes() + 1); @@ -46,7 +42,7 @@ impl CachedTreeHash for Inner { offsets.push(self.c.num_child_nodes() + 1); offsets.push(self.d.num_child_nodes() + 1); - OffsetHandler::from_lengths(initial_offset, offsets) + Ok(offsets) } fn num_child_nodes(&self) -> usize { @@ -67,7 +63,7 @@ impl CachedTreeHash for Inner { cache: &mut TreeHashCache, chunk: usize, ) -> Result { - let offset_handler = self.offset_handler(chunk)?; + let offset_handler = OffsetHandler::new(self, chunk)?; // Skip past the internal nodes and update any changed leaf nodes. { @@ -98,18 +94,14 @@ pub struct Outer { impl CachedTreeHash for Outer { type Item = Self; - fn build_cache(&self) -> Result { - let offset_handler = self.offset_handler(0)?; - + fn leaves_and_subtrees(&self) -> Vec { let mut leaves_and_subtrees = vec![]; - leaves_and_subtrees.append(&mut self.a.build_cache()?.into()); - leaves_and_subtrees.append(&mut self.b.build_cache()?.into()); - leaves_and_subtrees.append(&mut self.c.build_cache()?.into()); + leaves_and_subtrees.append(&mut self.a.leaves_and_subtrees()); + leaves_and_subtrees.append(&mut self.b.leaves_and_subtrees()); + leaves_and_subtrees.append(&mut self.c.leaves_and_subtrees()); - pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees); - - TreeHashCache::from_leaves_and_subtrees(leaves_and_subtrees, self.offset_handler(0)?) + leaves_and_subtrees } fn num_bytes(&self) -> usize { @@ -131,14 +123,14 @@ impl CachedTreeHash for Outer { num_nodes(leaves) + children - 1 } - fn offset_handler(&self, initial_offset: usize) -> Result { + fn offsets(&self) -> Result, Error> { let mut offsets = vec![]; offsets.push(self.a.num_child_nodes() + 1); offsets.push(self.b.num_child_nodes() + 1); offsets.push(self.c.num_child_nodes() + 1); - OffsetHandler::from_lengths(initial_offset, offsets) + Ok(offsets) } fn cached_hash_tree_root( @@ -147,7 +139,7 @@ impl CachedTreeHash for Outer { cache: &mut TreeHashCache, chunk: usize, ) -> Result { - let offset_handler = self.offset_handler(chunk)?; + let offset_handler = OffsetHandler::new(self, chunk)?; // Skip past the internal nodes and update any changed leaf nodes. { @@ -308,7 +300,7 @@ fn outer_builds() { // Generate reference data. let mut data = vec![]; data.append(&mut int_to_bytes32(0)); - let inner_bytes: Vec = inner.build_cache().unwrap().into(); + let inner_bytes: Vec = TreeHashCache::new(&inner).unwrap().into(); data.append(&mut int_to_bytes32(5)); let leaves = vec![