diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs index 227e6a4f95..015816403c 100644 --- a/eth2/types/src/beacon_state.rs +++ b/eth2/types/src/beacon_state.rs @@ -1,6 +1,7 @@ use self::epoch_cache::{get_active_validator_indices, EpochCache, Error as EpochCacheError}; use crate::test_utils::TestRandom; use crate::*; +use cached_tree_hash::{TreeHashCache, Error as TreeHashCacheError}; use int_to_bytes::int_to_bytes32; use pubkey_cache::PubkeyCache; use rand::RngCore; @@ -42,6 +43,7 @@ pub enum Error { EpochCacheUninitialized(RelativeEpoch), RelativeEpochError(RelativeEpochError), EpochCacheError(EpochCacheError), + TreeHashCacheError(TreeHashCacheError), } /// The state of the `BeaconChain` at some slot. @@ -123,6 +125,12 @@ pub struct BeaconState { #[tree_hash(skip_hashing)] #[test_random(default)] pub pubkey_cache: PubkeyCache, + #[serde(skip_serializing, skip_deserializing)] + #[ssz(skip_serializing)] + #[ssz(skip_deserializing)] + #[tree_hash(skip_hashing)] + #[test_random(default)] + pub tree_hash_cache: TreeHashCache, } impl BeaconState { @@ -198,6 +206,7 @@ impl BeaconState { EpochCache::default(), ], pubkey_cache: PubkeyCache::default(), + tree_hash_cache: TreeHashCache::default(), } } @@ -683,6 +692,7 @@ impl BeaconState { self.build_epoch_cache(RelativeEpoch::NextWithoutRegistryChange, spec)?; self.build_epoch_cache(RelativeEpoch::NextWithRegistryChange, spec)?; self.update_pubkey_cache()?; + self.update_tree_hash_cache()?; Ok(()) } @@ -789,6 +799,38 @@ impl BeaconState { pub fn drop_pubkey_cache(&mut self) { self.pubkey_cache = PubkeyCache::default() } + + /// Update the tree hash cache, building it for the first time if it is empty. + /// + /// Returns the `tree_hash_root` resulting from the update. This root can be considered the + /// canonical root of `self`. + pub fn update_tree_hash_cache(&mut self) -> Result { + if self.tree_hash_cache.is_empty() { + self.tree_hash_cache = TreeHashCache::new(self, 0)?; + } else { + // Move the cache outside of `self` to satisfy the borrow checker. + let mut cache = std::mem::replace(&mut self.tree_hash_cache, TreeHashCache::default()); + + cache.update(self)?; + + // Move the updated cache back into `self`. + self.tree_hash_cache = cache + } + + self.cached_tree_hash_root() + } + + /// Returns the tree hash root determined by the last execution of `self.update_tree_hash_cache(..)`. + /// + /// Note: does _not_ update the cache and may return an outdated root. + /// + /// Returns an error if the cache is not initialized or if an error is encountered during the + /// cache update. + pub fn cached_tree_hash_root(&self) -> Result { + self.tree_hash_cache.root() + .and_then(|b| Ok(Hash256::from_slice(b))) + .map_err(|e| e.into()) + } } impl From for Error { @@ -802,3 +844,9 @@ impl From for Error { Error::EpochCacheError(e) } } + +impl From for Error { + fn from(e: TreeHashCacheError) -> Error { + Error::TreeHashCacheError(e) + } +} diff --git a/eth2/types/src/beacon_state/tests.rs b/eth2/types/src/beacon_state/tests.rs index f417dd5551..d5862559ac 100644 --- a/eth2/types/src/beacon_state/tests.rs +++ b/eth2/types/src/beacon_state/tests.rs @@ -56,3 +56,22 @@ fn cache_initialization() { test_cache_initialization(&mut state, RelativeEpoch::NextWithRegistryChange, &spec); test_cache_initialization(&mut state, RelativeEpoch::NextWithoutRegistryChange, &spec); } + +#[test] +fn tree_hash_cache() { + use crate::test_utils::{SeedableRng, TestRandom, XorShiftRng}; + use tree_hash::TreeHash; + + let mut rng = XorShiftRng::from_seed([42; 16]); + + let mut state = BeaconState::random_for_test(&mut rng); + + let root = state.update_tree_hash_cache().unwrap(); + + assert_eq!(root.as_bytes(), &state.tree_hash_root()[..]); + + state.slot = state.slot + 1; + + let root = state.update_tree_hash_cache().unwrap(); + assert_eq!(root.as_bytes(), &state.tree_hash_root()[..]); +} diff --git a/eth2/utils/cached_tree_hash/src/errors.rs b/eth2/utils/cached_tree_hash/src/errors.rs index cd387fa47e..d9ac02913f 100644 --- a/eth2/utils/cached_tree_hash/src/errors.rs +++ b/eth2/utils/cached_tree_hash/src/errors.rs @@ -9,6 +9,7 @@ pub enum Error { UnableToGrowMerkleTree, UnableToShrinkMerkleTree, TreeCannotHaveZeroNodes, + CacheNotInitialized, ShouldNeverBePacked(TreeHashType), BytesAreNotEvenChunks(usize), NoModifiedFieldForChunk(usize), diff --git a/eth2/utils/cached_tree_hash/src/lib.rs b/eth2/utils/cached_tree_hash/src/lib.rs index b1397b0f4c..ee5f982753 100644 --- a/eth2/utils/cached_tree_hash/src/lib.rs +++ b/eth2/utils/cached_tree_hash/src/lib.rs @@ -45,17 +45,7 @@ impl CachedTreeHasher { where T: CachedTreeHash, { - // Reset the per-hash counters. - self.cache.chunk_index = 0; - self.cache.schema_index = 0; - - // Reset the "modified" flags for the cache. - self.cache.reset_modifications(); - - // Update the cache with the (maybe) changed object. - item.update_tree_hash_cache(&mut self.cache)?; - - Ok(()) + self.cache.update(item) } pub fn tree_hash_root(&self) -> Result, Error> { diff --git a/eth2/utils/cached_tree_hash/src/tree_hash_cache.rs b/eth2/utils/cached_tree_hash/src/tree_hash_cache.rs index efdd9b65b9..089e38469d 100644 --- a/eth2/utils/cached_tree_hash/src/tree_hash_cache.rs +++ b/eth2/utils/cached_tree_hash/src/tree_hash_cache.rs @@ -12,6 +12,18 @@ pub struct TreeHashCache { pub schema_index: usize, } +impl Default for TreeHashCache { + fn default() -> TreeHashCache { + TreeHashCache { + cache: vec![], + chunk_modified: vec![], + schemas: vec![], + chunk_index: 0, + schema_index: 0, + } + } +} + impl Into> for TreeHashCache { fn into(self) -> Vec { self.cache @@ -26,6 +38,24 @@ impl TreeHashCache { item.new_tree_hash_cache(depth) } + pub fn update(&mut self, item: &T) -> Result<(), Error> + where + T: CachedTreeHash, + { + if self.is_empty() { + Err(Error::CacheNotInitialized) + } else { + // Reset the per-hash counters. + self.chunk_index = 0; + self.schema_index = 0; + + // Reset the "modified" flags for the cache. + self.reset_modifications(); + + item.update_tree_hash_cache(self) + } + } + pub fn from_leaves_and_subtrees( item: &T, leaves_and_subtrees: Vec, @@ -108,6 +138,10 @@ impl TreeHashCache { }) } + pub fn is_empty(&self) -> bool { + self.chunk_modified.is_empty() + } + pub fn get_overlay( &self, schema_index: usize, @@ -210,9 +244,13 @@ impl TreeHashCache { } pub fn root(&self) -> Result<&[u8], Error> { - self.cache - .get(0..HASHSIZE) - .ok_or_else(|| Error::NoBytesForRoot) + if self.is_empty() { + Err(Error::CacheNotInitialized) + } else { + self.cache + .get(0..HASHSIZE) + .ok_or_else(|| Error::NoBytesForRoot) + } } pub fn splice(&mut self, chunk_range: Range, bytes: Vec, bools: Vec) {