diff --git a/eth2/utils/ssz/src/cached_tree_hash.rs b/eth2/utils/ssz/src/cached_tree_hash.rs index be3fe98de..3b900e503 100644 --- a/eth2/utils/ssz/src/cached_tree_hash.rs +++ b/eth2/utils/ssz/src/cached_tree_hash.rs @@ -11,16 +11,27 @@ const BYTES_PER_CHUNK: usize = 32; const HASHSIZE: usize = 32; const MERKLE_HASH_CHUNCK: usize = 2 * BYTES_PER_CHUNK; +#[derive(Debug, PartialEq, Clone)] +pub enum Error { + LeavesAndSubtreesIncomplete(usize), + ShouldNotProduceOffsetHandler, + NoFirstNode, + BytesAreNotEvenChunks(usize), + NoModifiedFieldForChunk(usize), + NoBytesForChunk(usize), + NoChildrenForHashing((usize, usize)), +} + pub trait CachedTreeHash { type Item: CachedTreeHash; - fn build_cache_bytes(&self) -> Vec; + fn build_cache(&self) -> Result; /// 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) -> Option; + fn offset_handler(&self, initial_offset: usize) -> Result; fn num_child_nodes(&self) -> usize; @@ -29,9 +40,10 @@ pub trait CachedTreeHash { other: &Self::Item, cache: &mut TreeHashCache, chunk: usize, - ) -> Option; + ) -> Result; } +#[derive(Debug, PartialEq, Clone)] pub struct TreeHashCache { cache: Vec, chunk_modified: Vec, @@ -44,11 +56,17 @@ impl Into> for TreeHashCache { } impl TreeHashCache { - pub fn new(mut leaves_and_subtrees: Vec, offset_handler: OffsetHandler) -> Option { - if leaves_and_subtrees.len() % BYTES_PER_CHUNK != 0 { - return None; - } + pub fn new(item: &T) -> Result + where + T: CachedTreeHash, + { + item.build_cache() + } + pub fn from_leaves_and_subtrees( + mut leaves_and_subtrees: Vec, + offset_handler: OffsetHandler, + ) -> Result { // 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; @@ -56,13 +74,22 @@ impl TreeHashCache { cache.resize(internal_node_bytes, 0); cache.append(&mut leaves_and_subtrees); + dbg!(cache.len() / BYTES_PER_CHUNK); + // 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; - leaves.extend_from_slice(cache.get(start..end)?); + dbg!(end); + dbg!(cache.len()); + + leaves.extend_from_slice( + cache + .get(start..end) + .ok_or_else(|| Error::LeavesAndSubtreesIncomplete(*leaf_chunk))?, + ); } // Merkleize the leaves, then split the leaf nodes off them. Then, replace all-zeros @@ -71,18 +98,18 @@ impl TreeHashCache { merkleized.split_off(internal_node_bytes); cache.splice(0..internal_node_bytes, merkleized); - Some(Self { + Ok(Self { chunk_modified: vec![false; cache.len() / BYTES_PER_CHUNK], cache, }) } - pub fn from_bytes(bytes: Vec) -> Option { + pub fn from_bytes(bytes: Vec) -> Result { if bytes.len() % BYTES_PER_CHUNK > 0 { - return None; + return Err(Error::BytesAreNotEvenChunks(bytes.len())); } - Some(Self { + Ok(Self { chunk_modified: vec![false; bytes.len() / BYTES_PER_CHUNK], cache: bytes, }) @@ -121,15 +148,18 @@ impl TreeHashCache { Some(()) } - pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Option<()> { + pub fn modify_chunk(&mut self, chunk: usize, to: &[u8]) -> Result<(), Error> { let start = chunk * BYTES_PER_CHUNK; let end = start + BYTES_PER_CHUNK; - self.cache.get_mut(start..end)?.copy_from_slice(to); + self.cache + .get_mut(start..end) + .ok_or_else(|| Error::NoBytesForChunk(chunk))? + .copy_from_slice(to); self.chunk_modified[chunk] = true; - Some(()) + Ok(()) } pub fn chunk_equals(&mut self, chunk: usize, other: &[u8]) -> Option { @@ -139,57 +169,30 @@ impl TreeHashCache { Some(self.cache.get(start..end)? == other) } - pub fn changed(&self, chunk: usize) -> Option { - self.chunk_modified.get(chunk).cloned() + pub fn changed(&self, chunk: usize) -> Result { + self.chunk_modified + .get(chunk) + .cloned() + .ok_or_else(|| Error::NoModifiedFieldForChunk(chunk)) } - pub fn either_modified(&self, children: (&usize, &usize)) -> Option { - Some(self.changed(*children.0)? | self.changed(*children.1)?) + pub fn either_modified(&self, children: (&usize, &usize)) -> Result { + Ok(self.changed(*children.0)? | self.changed(*children.1)?) } - /* - pub fn children_modified(&self, parent_chunk: usize, child_offsets: &[usize]) -> Option { - let children = children(parent_chunk); - - let a = *child_offsets.get(children.0)?; - let b = *child_offsets.get(children.1)?; - - Some(self.changed(a)? | self.changed(b)?) - } - */ - - pub fn hash_children(&self, children: (&usize, &usize)) -> Option> { + pub fn hash_children(&self, children: (&usize, &usize)) -> Result, Error> { let start = children.0 * BYTES_PER_CHUNK; let end = start + BYTES_PER_CHUNK * 2; - Some(hash(&self.cache.get(start..end)?)) + let children = &self + .cache + .get(start..end) + .ok_or_else(|| Error::NoChildrenForHashing((*children.0, *children.1)))?; + + Ok(hash(children)) } } -/* -pub struct LocalCache { - offsets: Vec, -} - -impl LocalCache { - -} - -pub struct OffsetBTree { - offsets: Vec, -} - -impl From> for OffsetBTree { - fn from(offsets: Vec) -> Self { - Self { offsets } - } -} - -impl OffsetBTree { - fn -} -*/ - fn children(parent: usize) -> (usize, usize) { ((2 * parent + 1), (2 * parent + 2)) } @@ -206,7 +209,7 @@ pub struct OffsetHandler { } impl OffsetHandler { - fn from_lengths(offset: usize, mut lengths: Vec) -> Self { + 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() { lengths.len() @@ -228,20 +231,23 @@ impl OffsetHandler { next_node += lengths[i]; } - Self { + Ok(Self { num_internal_nodes, num_leaf_nodes, offsets, next_node, - } + }) } pub fn total_nodes(&self) -> usize { self.num_internal_nodes + self.num_leaf_nodes } - pub fn first_leaf_node(&self) -> Option { - self.offsets.get(self.num_internal_nodes).cloned() + pub fn first_leaf_node(&self) -> Result { + self.offsets + .get(self.num_internal_nodes) + .cloned() + .ok_or_else(|| Error::NoFirstNode) } pub fn next_node(&self) -> usize { @@ -314,6 +320,15 @@ pub fn sanitise_bytes(mut bytes: Vec) -> Vec { bytes } +fn pad_for_leaf_count(num_leaves: usize, bytes: &mut Vec) { + let required_leaves = num_leaves.next_power_of_two(); + + bytes.resize( + bytes.len() + (required_leaves - num_leaves) * BYTES_PER_CHUNK, + 0, + ); +} + fn last_leaf_needs_padding(num_bytes: usize) -> bool { num_bytes % HASHSIZE != 0 } diff --git a/eth2/utils/ssz/src/cached_tree_hash/impls.rs b/eth2/utils/ssz/src/cached_tree_hash/impls.rs index a4d1c7d1a..9c0a8ec6d 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/impls.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/impls.rs @@ -4,16 +4,16 @@ use crate::{ssz_encode, Encodable}; impl CachedTreeHash for u64 { type Item = Self; - fn build_cache_bytes(&self) -> Vec { - merkleize(ssz_encode(self)) + fn build_cache(&self) -> Result { + TreeHashCache::from_bytes(merkleize(ssz_encode(self))) } fn num_bytes(&self) -> usize { 8 } - fn offset_handler(&self, _initial_offset: usize) -> Option { - None + fn offset_handler(&self, _initial_offset: usize) -> Result { + Err(Error::ShouldNotProduceOffsetHandler) } fn num_child_nodes(&self) -> usize { @@ -25,13 +25,13 @@ impl CachedTreeHash for u64 { other: &Self, cache: &mut TreeHashCache, chunk: usize, - ) -> Option { + ) -> Result { if self != other { let leaf = merkleize(ssz_encode(self)); cache.modify_chunk(chunk, &leaf)?; } - Some(chunk + 1) + Ok(chunk + 1) } } diff --git a/eth2/utils/ssz/src/cached_tree_hash/tests.rs b/eth2/utils/ssz/src/cached_tree_hash/tests.rs index 9cb012c79..a6be7f9ae 100644 --- a/eth2/utils/ssz/src/cached_tree_hash/tests.rs +++ b/eth2/utils/ssz/src/cached_tree_hash/tests.rs @@ -1,5 +1,5 @@ use super::*; -use int_to_bytes::{int_to_bytes32, int_to_bytes8}; +use int_to_bytes::int_to_bytes32; #[derive(Clone)] pub struct Inner { @@ -12,21 +12,19 @@ pub struct Inner { impl CachedTreeHash for Inner { type Item = Self; - fn build_cache_bytes(&self) -> Vec { + fn build_cache(&self) -> Result { + let offset_handler = self.offset_handler(0)?; + let mut leaves_and_subtrees = vec![]; - leaves_and_subtrees.append(&mut self.a.build_cache_bytes()); - leaves_and_subtrees.append(&mut self.b.build_cache_bytes()); - leaves_and_subtrees.append(&mut self.c.build_cache_bytes()); - leaves_and_subtrees.append(&mut self.d.build_cache_bytes()); + 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()); - // TODO: fix unwrap - let offset_handler = self.offset_handler(0).unwrap(); + pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees); - // TODO: fix unwrap - let cache = TreeHashCache::new(leaves_and_subtrees, offset_handler).unwrap(); - - cache.into() + TreeHashCache::from_leaves_and_subtrees(leaves_and_subtrees, self.offset_handler(0)?) } fn num_bytes(&self) -> usize { @@ -40,7 +38,7 @@ impl CachedTreeHash for Inner { bytes } - fn offset_handler(&self, initial_offset: usize) -> Option { + fn offset_handler(&self, initial_offset: usize) -> Result { let mut offsets = vec![]; offsets.push(self.a.num_child_nodes() + 1); @@ -48,7 +46,7 @@ impl CachedTreeHash for Inner { offsets.push(self.c.num_child_nodes() + 1); offsets.push(self.d.num_child_nodes() + 1); - Some(OffsetHandler::from_lengths(initial_offset, offsets)) + OffsetHandler::from_lengths(initial_offset, offsets) } fn num_child_nodes(&self) -> usize { @@ -68,7 +66,7 @@ impl CachedTreeHash for Inner { other: &Self, cache: &mut TreeHashCache, chunk: usize, - ) -> Option { + ) -> Result { let offset_handler = self.offset_handler(chunk)?; // Skip past the internal nodes and update any changed leaf nodes. @@ -86,7 +84,7 @@ impl CachedTreeHash for Inner { } } - Some(offset_handler.next_node()) + Ok(offset_handler.next_node()) } } @@ -100,20 +98,18 @@ pub struct Outer { impl CachedTreeHash for Outer { type Item = Self; - fn build_cache_bytes(&self) -> Vec { + fn build_cache(&self) -> Result { + let offset_handler = self.offset_handler(0)?; + let mut leaves_and_subtrees = vec![]; - leaves_and_subtrees.append(&mut self.a.build_cache_bytes()); - leaves_and_subtrees.append(&mut self.b.build_cache_bytes()); - leaves_and_subtrees.append(&mut self.c.build_cache_bytes()); + 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()); - // TODO: fix unwrap - let offset_handler = self.offset_handler(0).unwrap(); + pad_for_leaf_count(offset_handler.num_leaf_nodes, &mut leaves_and_subtrees); - // TODO: fix unwrap - let cache = TreeHashCache::new(leaves_and_subtrees, offset_handler).unwrap(); - - cache.into() + TreeHashCache::from_leaves_and_subtrees(leaves_and_subtrees, self.offset_handler(0)?) } fn num_bytes(&self) -> usize { @@ -135,14 +131,14 @@ impl CachedTreeHash for Outer { num_nodes(leaves) + children - 1 } - fn offset_handler(&self, initial_offset: usize) -> Option { + fn offset_handler(&self, initial_offset: usize) -> Result { 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); - Some(OffsetHandler::from_lengths(initial_offset, offsets)) + OffsetHandler::from_lengths(initial_offset, offsets) } fn cached_hash_tree_root( @@ -150,7 +146,7 @@ impl CachedTreeHash for Outer { other: &Self, cache: &mut TreeHashCache, chunk: usize, - ) -> Option { + ) -> Result { let offset_handler = self.offset_handler(chunk)?; // Skip past the internal nodes and update any changed leaf nodes. @@ -167,7 +163,7 @@ impl CachedTreeHash for Outer { } } - Some(offset_handler.next_node()) + Ok(offset_handler.next_node()) } } @@ -199,17 +195,16 @@ fn partial_modification_to_inner_struct() { ..original_inner.clone() }; - // Build the initial cache. - let original_cache = original_outer.build_cache_bytes(); - // Modify outer let modified_outer = Outer { b: modified_inner.clone(), ..original_outer.clone() }; + println!("AAAAAAAAA"); // Perform a differential hash - let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap(); + let mut cache_struct = TreeHashCache::new(&original_outer).unwrap(); + println!("BBBBBBBBBB"); modified_outer .cached_hash_tree_root(&original_outer, &mut cache_struct, 0) @@ -220,7 +215,7 @@ fn partial_modification_to_inner_struct() { // Generate reference data. let mut data = vec![]; data.append(&mut int_to_bytes32(0)); - let inner_bytes = modified_inner.build_cache_bytes(); + let inner_bytes: Vec = TreeHashCache::new(&modified_inner).unwrap().into(); data.append(&mut int_to_bytes32(5)); let leaves = vec![ @@ -254,7 +249,7 @@ fn partial_modification_to_outer() { }; // Build the initial cache. - let original_cache = original_outer.build_cache_bytes(); + // let original_cache = original_outer.build_cache_bytes(); // Modify outer let modified_outer = Outer { @@ -263,7 +258,7 @@ fn partial_modification_to_outer() { }; // Perform a differential hash - let mut cache_struct = TreeHashCache::from_bytes(original_cache.clone()).unwrap(); + let mut cache_struct = TreeHashCache::new(&original_outer).unwrap(); modified_outer .cached_hash_tree_root(&original_outer, &mut cache_struct, 0) @@ -274,7 +269,7 @@ fn partial_modification_to_outer() { // Generate reference data. let mut data = vec![]; data.append(&mut int_to_bytes32(0)); - let inner_bytes = inner.build_cache_bytes(); + let inner_bytes: Vec = TreeHashCache::new(&inner).unwrap().into(); data.append(&mut int_to_bytes32(5)); let leaves = vec![ @@ -308,12 +303,12 @@ fn outer_builds() { }; // Build the function output. - let cache = outer.build_cache_bytes(); + let cache: Vec = TreeHashCache::new(&outer).unwrap().into(); // Generate reference data. let mut data = vec![]; data.append(&mut int_to_bytes32(0)); - let inner_bytes = inner.build_cache_bytes(); + let inner_bytes: Vec = inner.build_cache().unwrap().into(); data.append(&mut int_to_bytes32(5)); let leaves = vec![ @@ -427,7 +422,7 @@ fn generic_test(index: usize) { d: 4, }; - let cache = inner.build_cache_bytes(); + let cache: Vec = TreeHashCache::new(&inner).unwrap().into(); let changed_inner = match index { 0 => Inner { @@ -498,7 +493,7 @@ fn inner_builds() { d: 4, }; - let cache = inner.build_cache_bytes(); + let cache: Vec = TreeHashCache::new(&inner).unwrap().into(); assert_eq!(expected, cache); }