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