From 354f823c1628e56de6627961dc3746a1f09c7a2b Mon Sep 17 00:00:00 2001 From: Paul Hauner Date: Mon, 15 Apr 2019 15:13:02 +1000 Subject: [PATCH] Tidy tree hash cache, add new trait --- eth2/utils/tree_hash/src/btree_overlay.rs | 2 +- eth2/utils/tree_hash/src/cached_tree_hash.rs | 13 +- eth2/utils/tree_hash/src/impls.rs | 2 +- eth2/utils/tree_hash/src/impls/vec.rs | 6 +- eth2/utils/tree_hash/src/lib.rs | 12 +- eth2/utils/tree_hash/tests/tests.rs | 121 ++++++++++++++++++- 6 files changed, 138 insertions(+), 18 deletions(-) diff --git a/eth2/utils/tree_hash/src/btree_overlay.rs b/eth2/utils/tree_hash/src/btree_overlay.rs index 7d5602c0b..8c859d046 100644 --- a/eth2/utils/tree_hash/src/btree_overlay.rs +++ b/eth2/utils/tree_hash/src/btree_overlay.rs @@ -12,7 +12,7 @@ pub struct BTreeOverlay { impl BTreeOverlay { pub fn new(item: &T, initial_offset: usize) -> Result where - T: CachedTreeHash, + T: CachedTreeHashSubtree, { item.btree_overlay(initial_offset) } diff --git a/eth2/utils/tree_hash/src/cached_tree_hash.rs b/eth2/utils/tree_hash/src/cached_tree_hash.rs index 4022f6b7b..97f9388a1 100644 --- a/eth2/utils/tree_hash/src/cached_tree_hash.rs +++ b/eth2/utils/tree_hash/src/cached_tree_hash.rs @@ -15,7 +15,7 @@ impl Into> for TreeHashCache { impl TreeHashCache { pub fn new(item: &T) -> Result where - T: CachedTreeHash, + T: CachedTreeHashSubtree, { item.new_cache() } @@ -32,7 +32,7 @@ impl TreeHashCache { leaves_and_subtrees: Vec, ) -> Result where - T: CachedTreeHash, + T: CachedTreeHashSubtree, { let offset_handler = BTreeOverlay::new(item, 0)?; @@ -55,7 +55,7 @@ impl TreeHashCache { // 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()?); + leaves.append(&mut t.root().ok_or_else(|| Error::NoBytesForRoot)?.to_vec()); cache.append(&mut t.into_merkle_tree()); } @@ -89,11 +89,8 @@ impl TreeHashCache { 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 root(&self) -> Option<&[u8]> { + self.cache.get(0..HASHSIZE) } pub fn splice(&mut self, chunk_range: Range, replace_with: Self) { diff --git a/eth2/utils/tree_hash/src/impls.rs b/eth2/utils/tree_hash/src/impls.rs index 6849fd55c..bd5c352c9 100644 --- a/eth2/utils/tree_hash/src/impls.rs +++ b/eth2/utils/tree_hash/src/impls.rs @@ -4,7 +4,7 @@ use ssz::ssz_encode; mod vec; -impl CachedTreeHash for u64 { +impl CachedTreeHashSubtree for u64 { fn item_type() -> ItemType { ItemType::Basic } diff --git a/eth2/utils/tree_hash/src/impls/vec.rs b/eth2/utils/tree_hash/src/impls/vec.rs index 7c0993f43..c02460cf3 100644 --- a/eth2/utils/tree_hash/src/impls/vec.rs +++ b/eth2/utils/tree_hash/src/impls/vec.rs @@ -1,8 +1,8 @@ use super::*; -impl CachedTreeHash> for Vec +impl CachedTreeHashSubtree> for Vec where - T: CachedTreeHash, + T: CachedTreeHashSubtree, { fn item_type() -> ItemType { ItemType::List @@ -168,7 +168,7 @@ where fn get_packed_leaves(vec: &Vec) -> Result, Error> where - T: CachedTreeHash, + T: CachedTreeHashSubtree, { let num_packed_bytes = (BYTES_PER_CHUNK / T::packing_factor()) * vec.len(); let num_leaves = num_sanitized_leaves(num_packed_bytes); diff --git a/eth2/utils/tree_hash/src/lib.rs b/eth2/utils/tree_hash/src/lib.rs index b3167a37d..179f557ce 100644 --- a/eth2/utils/tree_hash/src/lib.rs +++ b/eth2/utils/tree_hash/src/lib.rs @@ -1,6 +1,5 @@ use hashing::hash; use int_to_bytes::int_to_bytes32; -use std::fmt::Debug; use std::ops::Range; mod btree_overlay; @@ -36,8 +35,15 @@ pub enum ItemType { Composite, } -// TODO: remove debug requirement. -pub trait CachedTreeHash: Debug { +pub trait CachedTreeHash: CachedTreeHashSubtree + Sized { + fn update_internal_tree_hash_cache(self, old: T) -> Result<(Self, Self), Error>; + + fn cached_tree_hash_root(&self) -> Option>; + + fn clone_without_tree_hash_cache(&self) -> Self; +} + +pub trait CachedTreeHashSubtree { fn item_type() -> ItemType; fn btree_overlay(&self, chunk_offset: usize) -> Result; diff --git a/eth2/utils/tree_hash/tests/tests.rs b/eth2/utils/tree_hash/tests/tests.rs index 22780bcac..c61a010ca 100644 --- a/eth2/utils/tree_hash/tests/tests.rs +++ b/eth2/utils/tree_hash/tests/tests.rs @@ -2,6 +2,123 @@ use hashing::hash; use int_to_bytes::{int_to_bytes32, int_to_bytes8}; use tree_hash::*; +#[derive(Clone, Debug)] +pub struct InternalCache { + pub a: u64, + pub b: u64, + pub cache: Option, +} + +impl CachedTreeHash for InternalCache { + fn update_internal_tree_hash_cache(mut self, mut old: Self) -> Result<(Self, Self), Error> { + let mut local_cache = old.cache; + old.cache = None; + + if let Some(ref mut local_cache) = local_cache { + self.update_cache(&old, local_cache, 0)?; + } else { + local_cache = Some(self.new_cache()?) + } + + self.cache = local_cache; + + Ok((old, self)) + } + + fn cached_tree_hash_root(&self) -> Option> { + match &self.cache { + None => None, + Some(c) => Some(c.root()?.to_vec()), + } + } + + fn clone_without_tree_hash_cache(&self) -> Self { + Self { + a: self.a, + b: self.b, + cache: None, + } + } +} + +#[test] +fn works_when_embedded() { + let old = InternalCache { + a: 99, + b: 99, + cache: None, + }; + + let mut new = old.clone_without_tree_hash_cache(); + new.a = 1; + new.b = 2; + + let (_old, new) = new.update_internal_tree_hash_cache(old).unwrap(); + + let root = new.cached_tree_hash_root().unwrap(); + + let leaves = vec![int_to_bytes32(1), int_to_bytes32(2)]; + let merkle = merkleize(join(leaves)); + + assert_eq!(&merkle[0..32], &root[..]); +} + +impl CachedTreeHashSubtree for InternalCache { + fn item_type() -> ItemType { + ItemType::Composite + } + + fn new_cache(&self) -> Result { + let tree = TreeHashCache::from_leaves_and_subtrees( + self, + vec![self.a.new_cache()?, self.b.new_cache()?], + )?; + + Ok(tree) + } + + fn btree_overlay(&self, chunk_offset: usize) -> Result { + let mut lengths = vec![]; + + lengths.push(BTreeOverlay::new(&self.a, 0)?.total_nodes()); + lengths.push(BTreeOverlay::new(&self.b, 0)?.total_nodes()); + + BTreeOverlay::from_lengths(chunk_offset, lengths) + } + + fn packed_encoding(&self) -> Result, Error> { + Err(Error::ShouldNeverBePacked(Self::item_type())) + } + + fn packing_factor() -> usize { + 1 + } + + fn update_cache( + &self, + other: &Self, + cache: &mut TreeHashCache, + chunk: usize, + ) -> Result { + let offset_handler = BTreeOverlay::new(self, chunk)?; + + // Skip past the internal nodes and update any changed leaf nodes. + { + let chunk = offset_handler.first_leaf_node()?; + let chunk = self.a.update_cache(&other.a, cache, chunk)?; + let _chunk = self.b.update_cache(&other.b, cache, chunk)?; + } + + for (&parent, children) in offset_handler.iter_internal_nodes().rev() { + if cache.either_modified(children)? { + cache.modify_chunk(parent, &cache.hash_children(children)?)?; + } + } + + Ok(offset_handler.next_node) + } +} + fn num_nodes(num_leaves: usize) -> usize { 2 * num_leaves - 1 } @@ -14,7 +131,7 @@ pub struct Inner { pub d: u64, } -impl CachedTreeHash for Inner { +impl CachedTreeHashSubtree for Inner { fn item_type() -> ItemType { ItemType::Composite } @@ -86,7 +203,7 @@ pub struct Outer { pub c: u64, } -impl CachedTreeHash for Outer { +impl CachedTreeHashSubtree for Outer { fn item_type() -> ItemType { ItemType::Composite }