Avoid building caches during block replay (#783)

Also, make the ExitCache safe.

eth2/types/src/beacon_state.rs

@@ -58,6 +58,7 @@ pub enum Error {
     PreviousCommitteeCacheUninitialized,
     CurrentCommitteeCacheUninitialized,
     RelativeEpochError(RelativeEpochError),
+    ExitCacheUninitialized,
     CommitteeCacheUninitialized(Option<RelativeEpoch>),
     SszTypesError(ssz_types::Error),
     CachedTreeHashError(cached_tree_hash::Error),
@@ -779,13 +780,19 @@ impl<T: EthSpec> BeaconState<T> {
 
     /// Build all the caches, if they need to be built.
     pub fn build_all_caches(&mut self, spec: &ChainSpec) -> Result<(), Error> {
+        self.build_all_committee_caches(spec)?;
+        self.update_pubkey_cache()?;
+        self.build_tree_hash_cache()?;
+        self.exit_cache.build(&self.validators, spec)?;
+
+        Ok(())
+    }
+
+    /// Build all committee caches, if they need to be built.
+    pub fn build_all_committee_caches(&mut self, spec: &ChainSpec) -> Result<(), Error> {
         self.build_committee_cache(RelativeEpoch::Previous, spec)?;
         self.build_committee_cache(RelativeEpoch::Current, spec)?;
         self.build_committee_cache(RelativeEpoch::Next, spec)?;
-        self.update_pubkey_cache()?;
-        self.build_tree_hash_cache()?;
-        self.exit_cache.build_from_registry(&self.validators, spec);
-
         Ok(())
     }
 

eth2/types/src/beacon_state/exit_cache.rs

@@ -1,35 +1,68 @@
-use super::{ChainSpec, Epoch, Validator};
+use super::{BeaconStateError, ChainSpec, Epoch, Validator};
 use serde_derive::{Deserialize, Serialize};
 use std::collections::HashMap;
 
 /// Map from exit epoch to the number of validators known to be exiting/exited at that epoch.
 #[derive(Debug, Default, Clone, PartialEq, Serialize, Deserialize)]
-pub struct ExitCache(HashMap<Epoch, u64>);
+pub struct ExitCache {
+    initialized: bool,
+    exits_per_epoch: HashMap<Epoch, u64>,
+}
 
 impl ExitCache {
+    /// Ensure the cache is built, and do nothing if it's already initialized.
+    pub fn build(
+        &mut self,
+        validators: &[Validator],
+        spec: &ChainSpec,
+    ) -> Result<(), BeaconStateError> {
+        if self.initialized {
+            Ok(())
+        } else {
+            self.force_build(validators, spec)
+        }
+    }
+
     /// Add all validators with a non-trivial exit epoch to the cache.
-    pub fn build_from_registry(&mut self, validators: &[Validator], spec: &ChainSpec) {
+    pub fn force_build(
+        &mut self,
+        validators: &[Validator],
+        spec: &ChainSpec,
+    ) -> Result<(), BeaconStateError> {
+        self.initialized = true;
         validators
             .iter()
             .filter(|validator| validator.exit_epoch != spec.far_future_epoch)
-            .for_each(|validator| self.record_validator_exit(validator.exit_epoch));
+            .try_for_each(|validator| self.record_validator_exit(validator.exit_epoch))
+    }
+
+    /// Check that the cache is initialized and return an error if it isn't.
+    pub fn check_initialized(&self) -> Result<(), BeaconStateError> {
+        if self.initialized {
+            Ok(())
+        } else {
+            Err(BeaconStateError::ExitCacheUninitialized)
+        }
     }
 
     /// Record the exit of a single validator in the cache.
     ///
     /// Must only be called once per exiting validator.
-    pub fn record_validator_exit(&mut self, exit_epoch: Epoch) {
-        *self.0.entry(exit_epoch).or_insert(0) += 1;
+    pub fn record_validator_exit(&mut self, exit_epoch: Epoch) -> Result<(), BeaconStateError> {
+        self.check_initialized()?;
+        *self.exits_per_epoch.entry(exit_epoch).or_insert(0) += 1;
+        Ok(())
     }
 
     /// Get the greatest epoch for which validator exits are known.
-    pub fn max_epoch(&self) -> Option<Epoch> {
-        // This could probably be made even faster by caching the maximum.
-        self.0.keys().max().cloned()
+    pub fn max_epoch(&self) -> Result<Option<Epoch>, BeaconStateError> {
+        self.check_initialized()?;
+        Ok(self.exits_per_epoch.keys().max().cloned())
     }
 
     /// Get the number of validators exiting/exited at a given epoch, or zero if not known.
-    pub fn get_churn_at(&self, epoch: Epoch) -> u64 {
-        self.0.get(&epoch).cloned().unwrap_or(0)
+    pub fn get_churn_at(&self, epoch: Epoch) -> Result<u64, BeaconStateError> {
+        self.check_initialized()?;
+        Ok(self.exits_per_epoch.get(&epoch).cloned().unwrap_or(0))
     }
 }