From b26f1f8e1cfd357f2991d1e8230a52305f491a44 Mon Sep 17 00:00:00 2001
From: Paul Hauner <paul@paulhauner.com>
Date: Sun, 31 Mar 2019 13:40:12 +1100
Subject: [PATCH] Add `build_all_caches` method to `BeaconState`

Also adds a few more cache builds in BeaconChain.
---
 beacon_node/beacon_chain/src/beacon_chain.rs | 23 ++++++++++++--------
 eth2/types/src/beacon_state.rs               | 11 ++++++++++
 2 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/beacon_node/beacon_chain/src/beacon_chain.rs b/beacon_node/beacon_chain/src/beacon_chain.rs
index efa83bdd3..046f37a81 100644
--- a/beacon_node/beacon_chain/src/beacon_chain.rs
+++ b/beacon_node/beacon_chain/src/beacon_chain.rs
@@ -130,10 +130,7 @@ where
             state_root,
         ));
 
-        genesis_state.build_epoch_cache(RelativeEpoch::Previous, &spec)?;
-        genesis_state.build_epoch_cache(RelativeEpoch::Current, &spec)?;
-        genesis_state.build_epoch_cache(RelativeEpoch::NextWithoutRegistryChange, &spec)?;
-        genesis_state.build_epoch_cache(RelativeEpoch::NextWithRegistryChange, &spec)?;
+        genesis_state.build_all_caches(&spec)?;
 
         Ok(Self {
             block_store,
@@ -318,6 +315,8 @@ where
             per_slot_processing(&mut state, &latest_block_header, &self.spec)?;
         }
 
+        state.build_all_caches(&self.spec)?;
+
         *self.state.write() = state;
 
         Ok(())
@@ -342,11 +341,17 @@ where
 
             per_slot_processing(&mut *state, &latest_block_header, &self.spec)?;
         }
-        state.build_epoch_cache(RelativeEpoch::Previous, &self.spec)?;
-        state.build_epoch_cache(RelativeEpoch::Current, &self.spec)?;
-        state.build_epoch_cache(RelativeEpoch::NextWithoutRegistryChange, &self.spec)?;
-        state.build_epoch_cache(RelativeEpoch::NextWithRegistryChange, &self.spec)?;
-        state.update_pubkey_cache()?;
+
+        state.build_all_caches(&self.spec)?;
+
+        Ok(())
+    }
+
+    /// Build all of the caches on the current state.
+    ///
+    /// Ideally this shouldn't be required, however we leave it here for testing.
+    pub fn ensure_state_caches_are_built(&self) -> Result<(), Error> {
+        self.state.write().build_all_caches(&self.spec)?;
 
         Ok(())
     }
diff --git a/eth2/types/src/beacon_state.rs b/eth2/types/src/beacon_state.rs
index 1e5278124..774e8eb76 100644
--- a/eth2/types/src/beacon_state.rs
+++ b/eth2/types/src/beacon_state.rs
@@ -661,6 +661,17 @@ impl BeaconState {
         })
     }
 
+    /// Build all the caches, if they need to be built.
+    pub fn build_all_caches(&mut self, spec: &ChainSpec) -> Result<(), Error> {
+        self.build_epoch_cache(RelativeEpoch::Previous, spec)?;
+        self.build_epoch_cache(RelativeEpoch::Current, spec)?;
+        self.build_epoch_cache(RelativeEpoch::NextWithoutRegistryChange, spec)?;
+        self.build_epoch_cache(RelativeEpoch::NextWithRegistryChange, spec)?;
+        self.update_pubkey_cache()?;
+
+        Ok(())
+    }
+
     /// Build an epoch cache, unless it is has already been built.
     pub fn build_epoch_cache(
         &mut self,