Fix failing tree hash tests

eth2/utils/ssz/src/cached_tree_hash.rs

@@ -13,9 +13,9 @@ const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK;
 
 #[derive(Debug, PartialEq, Clone)]
 pub enum Error {
-    LeavesAndSubtreesIncomplete(usize),
     ShouldNotProduceOffsetHandler,
     NoFirstNode,
+    NoBytesForRoot,
     BytesAreNotEvenChunks(usize),
     NoModifiedFieldForChunk(usize),
     NoBytesForChunk(usize),
@@ -25,7 +25,7 @@ pub enum Error {
 pub trait CachedTreeHash {
     type Item: CachedTreeHash;
 
-    fn leaves_and_subtrees(&self) -> Vec<u8>;
+    fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error>;
 
     /// Return the number of bytes when this element is encoded as raw SSZ _without_ length
     /// prefixes.
@@ -60,42 +60,44 @@ impl TreeHashCache {
     where
         T: CachedTreeHash,
     {
-        Self::from_leaves_and_subtrees(item.leaves_and_subtrees(), OffsetHandler::new(item, 0)?)
+        item.build_tree_hash_cache()
     }
 
-    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);
+    pub fn from_leaves_and_subtrees<T>(
+        item: &T,
+        leaves_and_subtrees: Vec<Self>,
+    ) -> Result<Self, Error>
+    where
+        T: CachedTreeHash,
+    {
+        let offset_handler = OffsetHandler::new(item, 0)?;
+
+        // Note how many leaves were provided. If is not a power-of-two, we'll need to pad it out
+        // later.
+        let num_provided_leaf_nodes = leaves_and_subtrees.len();
 
         // 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;
-        let mut cache = Vec::with_capacity(internal_node_bytes + leaves_and_subtrees.len());
+        let leaves_and_subtrees_bytes = leaves_and_subtrees
+            .iter()
+            .fold(0, |acc, t| acc + t.bytes_len());
+        let mut cache = Vec::with_capacity(leaves_and_subtrees_bytes + internal_node_bytes);
         cache.resize(internal_node_bytes, 0);
-        cache.append(&mut leaves_and_subtrees);
 
-        dbg!(cache.len() / BYTES_PER_CHUNK);
+        // Allocate enough bytes to store all the leaves.
+        let mut leaves = Vec::with_capacity(offset_handler.num_leaf_nodes * HASHSIZE);
 
-        // Concat all the leaves into one big byte array, ready for `merkleize`.
-        let mut leaves = vec![];
-        for leaf_chunk in offset_handler.iter_leaf_nodes() {
-            let start = leaf_chunk * BYTES_PER_CHUNK;
-            let end = start + BYTES_PER_CHUNK;
-
-            dbg!(end);
-            dbg!(cache.len());
-
-            leaves.extend_from_slice(
-                cache
-                    .get(start..end)
-                    .ok_or_else(|| Error::LeavesAndSubtreesIncomplete(*leaf_chunk))?,
-            );
+        // Iterate through all of the leaves/subtrees, adding their root as a leaf node and then
+        // concatenating their merkle trees.
+        for t in leaves_and_subtrees {
+            leaves.append(&mut t.root()?);
+            cache.append(&mut t.into_merkle_tree());
         }
 
+        // Pad the leaves to an even power-of-two, using zeros.
+        pad_for_leaf_count(num_provided_leaf_nodes, &mut cache);
+
         // Merkleize the leaves, then split the leaf nodes off them. Then, replace all-zeros
         // internal nodes created earlier with the internal nodes generated by `merkleize`.
         let mut merkleized = merkleize(leaves);
@@ -108,6 +110,17 @@ impl TreeHashCache {
         })
     }
 
+    pub fn bytes_len(&self) -> usize {
+        self.cache.len()
+    }
+
+    pub fn root(&self) -> Result<Vec<u8>, Error> {
+        self.cache
+            .get(0..HASHSIZE)
+            .ok_or_else(|| Error::NoBytesForRoot)
+            .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()));
@@ -195,6 +208,10 @@ impl TreeHashCache {
 
         Ok(hash(children))
     }
+
+    pub fn into_merkle_tree(self) -> Vec<u8> {
+        self.cache
+    }
 }
 
 fn children(parent: usize) -> (usize, usize) {
@@ -205,6 +222,7 @@ fn num_nodes(num_leaves: usize) -> usize {
     2 * num_leaves - 1
 }
 
+#[derive(Debug)]
 pub struct OffsetHandler {
     num_internal_nodes: usize,
     num_leaf_nodes: usize,

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

@@ -4,8 +4,8 @@ use crate::{ssz_encode, Encodable};
 impl CachedTreeHash for u64 {
     type Item = Self;
 
-    fn leaves_and_subtrees(&self) -> Vec<u8> {
-        merkleize(ssz_encode(self))
+    fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
+        Ok(TreeHashCache::from_bytes(merkleize(ssz_encode(self)))?)
     }
 
     fn num_bytes(&self) -> usize {

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

@@ -12,15 +12,18 @@ pub struct Inner {
 impl CachedTreeHash for Inner {
     type Item = Self;
 
-    fn leaves_and_subtrees(&self) -> Vec<u8> {
-        let mut leaves_and_subtrees = vec![];
+    fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
+        let tree = TreeHashCache::from_leaves_and_subtrees(
+            self,
+            vec![
+                self.a.build_tree_hash_cache()?,
+                self.b.build_tree_hash_cache()?,
+                self.c.build_tree_hash_cache()?,
+                self.d.build_tree_hash_cache()?,
+            ],
+        )?;
 
-        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());
-
-        leaves_and_subtrees
+        Ok(tree)
     }
 
     fn num_bytes(&self) -> usize {
@@ -94,14 +97,17 @@ pub struct Outer {
 impl CachedTreeHash for Outer {
     type Item = Self;
 
-    fn leaves_and_subtrees(&self) -> Vec<u8> {
-        let mut leaves_and_subtrees = vec![];
+    fn build_tree_hash_cache(&self) -> Result<TreeHashCache, Error> {
+        let tree = TreeHashCache::from_leaves_and_subtrees(
+            self,
+            vec![
+                self.a.build_tree_hash_cache()?,
+                self.b.build_tree_hash_cache()?,
+                self.c.build_tree_hash_cache()?,
+            ],
+        )?;
 
-        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
+        Ok(tree)
     }
 
     fn num_bytes(&self) -> usize {
@@ -193,10 +199,8 @@ fn partial_modification_to_inner_struct() {
         ..original_outer.clone()
     };
 
-    println!("AAAAAAAAA");
     // Perform a differential hash
     let mut cache_struct = TreeHashCache::new(&original_outer).unwrap();
-    println!("BBBBBBBBBB");
 
     modified_outer
         .cached_hash_tree_root(&original_outer, &mut cache_struct, 0)