Add tree hash cache as field to BeaconState.

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<Hash256, Error> {
+        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<Hash256, Error> {
+        self.tree_hash_cache.root()
+            .and_then(|b| Ok(Hash256::from_slice(b)))
+            .map_err(|e| e.into())
+    }
 }
 
 impl From<RelativeEpochError> for Error {
@@ -802,3 +844,9 @@ impl From<EpochCacheError> for Error {
         Error::EpochCacheError(e)
     }
 }
+
+impl From<TreeHashCacheError> for Error {
+    fn from(e: TreeHashCacheError) -> Error {
+        Error::TreeHashCacheError(e)
+    }
+}

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()[..]);
+}

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),

eth2/utils/cached_tree_hash/src/lib.rs

@@ -45,17 +45,7 @@ impl CachedTreeHasher {
     where
         T: CachedTreeHash<T>,
     {
-        // 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<Vec<u8>, Error> {

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<Vec<u8>> for TreeHashCache {
     fn into(self) -> Vec<u8> {
         self.cache
@@ -26,6 +38,24 @@ impl TreeHashCache {
         item.new_tree_hash_cache(depth)
     }
 
+    pub fn update<T>(&mut self, item: &T) -> Result<(), Error>
+    where
+        T: CachedTreeHash<T>,
+    {
+        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<T>(
         item: &T,
         leaves_and_subtrees: Vec<Self>,
@@ -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<usize>, bytes: Vec<u8>, bools: Vec<bool>) {