From fc17d5fea4c46ae35ab11dc9c6b1ae7a93dfaccb Mon Sep 17 00:00:00 2001 From: Paul Hauner Date: Fri, 29 Mar 2019 14:37:27 +1100 Subject: [PATCH] Fix failing tree hash tests --- eth2/utils/ssz/src/cached_tree_hash.rs | 72 ++++++++++++-------- eth2/utils/ssz/src/cached_tree_hash/impls.rs | 4 +- eth2/utils/ssz/src/cached_tree_hash/tests.rs | 38 ++++++----- 3 files changed, 68 insertions(+), 46 deletions(-) diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs index 83b516ac7..510185b40 100644 --- a/eth2/utils/ssz/src/cached_tree_hash.rs +++ b/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; + fn build_tree_hash_cache(&self) -> Result; /// 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, - 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); + pub fn from_leaves_and_subtrees( + item: &T, + leaves_and_subtrees: Vec, + ) -> Result + 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, Error> { + self.cache + .get(0..HASHSIZE) + .ok_or_else(|| Error::NoBytesForRoot) + .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())); @@ -195,6 +208,10 @@ impl TreeHashCache { Ok(hash(children)) } + + pub fn into_merkle_tree(self) -> Vec { + 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, diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs index 54a690c6d..012a4a8be 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs +++ b/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 { - merkleize(ssz_encode(self)) + fn build_tree_hash_cache(&self) -> Result { + Ok(TreeHashCache::from_bytes(merkleize(ssz_encode(self)))?) } fn num_bytes(&self) -> usize { diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs index f6c52ef8f..0593b2bae 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs +++ b/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 { - let mut leaves_and_subtrees = vec![]; + fn build_tree_hash_cache(&self) -> Result { + 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 { - let mut leaves_and_subtrees = vec![]; + fn build_tree_hash_cache(&self) -> Result { + 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)