Move offset_handler construction into self

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<TreeHashCache, Error>;
+    fn leaves_and_subtrees(&self) -> Vec<u8>;
 
     /// 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<OffsetHandler, Error>;
+    fn offsets(&self) -> Result<Vec<usize>, 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<u8>,
         offset_handler: OffsetHandler,
     ) -> Result<Self, Error> {
+        // 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<T>(item: &T, initial_offset: usize) -> Result<Self, Error>
+    where
+        T: CachedTreeHash,
+    {
+        Self::from_lengths(initial_offset, item.offsets()?)
+    }
+
     fn from_lengths(offset: usize, mut lengths: Vec<usize>) -> Result<Self, Error> {
         // Extend it to the next power-of-two, if it is not already.
         let num_leaf_nodes = if lengths.len().is_power_of_two() {

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, Error> {
+    fn leaves_and_subtrees(&self) -> Vec<u8> {
-        TreeHashCache::from_bytes(merkleize(ssz_encode(self)))
+        merkleize(ssz_encode(self))
     }
 
     fn num_bytes(&self) -> usize {
         8
     }
 
-    fn offset_handler(&self, _initial_offset: usize) -> Result<OffsetHandler, Error> {
+    fn offsets(&self) -> Result<Vec<usize>, Error> {
         Err(Error::ShouldNotProduceOffsetHandler)
     }
 

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<TreeHashCache, Error> {
+    fn leaves_and_subtrees(&self) -> Vec<u8> {
-        let offset_handler = self.offset_handler(0)?;
-
         let mut leaves_and_subtrees = vec![];
 
-        leaves_and_subtrees.append(&mut self.a.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.a.leaves_and_subtrees());
-        leaves_and_subtrees.append(&mut self.b.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.b.leaves_and_subtrees());
-        leaves_and_subtrees.append(&mut self.c.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.c.leaves_and_subtrees());
-        leaves_and_subtrees.append(&mut self.d.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.d.leaves_and_subtrees());
 
-        pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees);
+        leaves_and_subtrees
-
-        TreeHashCache::from_leaves_and_subtrees(leaves_and_subtrees, self.offset_handler(0)?)
     }
 
     fn num_bytes(&self) -> usize {
@@ -38,7 +34,7 @@ impl CachedTreeHash for Inner {
         bytes
     }
 
-    fn offset_handler(&self, initial_offset: usize) -> Result<OffsetHandler, Error> {
+    fn offsets(&self) -> Result<Vec<usize>, 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<usize, Error> {
-        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<TreeHashCache, Error> {
+    fn leaves_and_subtrees(&self) -> Vec<u8> {
-        let offset_handler = self.offset_handler(0)?;
-
         let mut leaves_and_subtrees = vec![];
 
-        leaves_and_subtrees.append(&mut self.a.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.a.leaves_and_subtrees());
-        leaves_and_subtrees.append(&mut self.b.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.b.leaves_and_subtrees());
-        leaves_and_subtrees.append(&mut self.c.build_cache()?.into());
+        leaves_and_subtrees.append(&mut self.c.leaves_and_subtrees());
 
-        pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees);
+        leaves_and_subtrees
-
-        TreeHashCache::from_leaves_and_subtrees(leaves_and_subtrees, self.offset_handler(0)?)
     }
 
     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<OffsetHandler, Error> {
+    fn offsets(&self) -> Result<Vec<usize>, 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<usize, Error> {
-        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<u8> = inner.build_cache().unwrap().into();
+    let inner_bytes: Vec<u8> = TreeHashCache::new(&inner).unwrap().into();
     data.append(&mut int_to_bytes32(5));
 
     let leaves = vec![