Fix EpochCache handling in ef-tests (#4454)

consensus/state_processing/src/epoch_cache.rs

@@ -23,6 +23,7 @@ pub fn initialize_epoch_cache<E: EthSpec>(
     }
 
     // Compute base rewards.
+    state.build_total_active_balance_cache_at(epoch, spec)?;
     let total_active_balance = state.get_total_active_balance_at_epoch(epoch)?;
     let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_active_balance);
     let base_reward_per_increment = BaseRewardPerIncrement::new(total_active_balance, spec)?;

consensus/state_processing/src/per_epoch_processing/altair.rs

@@ -28,6 +28,8 @@ pub fn process_epoch<T: EthSpec>(
     state.build_committee_cache(RelativeEpoch::Previous, spec)?;
     state.build_committee_cache(RelativeEpoch::Current, spec)?;
     state.build_committee_cache(RelativeEpoch::Next, spec)?;
+    state.build_total_active_balance_cache_at(state.current_epoch(), spec)?;
+    initialize_epoch_cache(state, state.current_epoch(), spec)?;
 
     // Pre-compute participating indices and total balances.
     let mut participation_cache = ParticipationCache::new(state, spec)?;

consensus/state_processing/src/per_epoch_processing/base.rs

@@ -24,6 +24,8 @@ pub fn process_epoch<T: EthSpec>(
     state.build_committee_cache(RelativeEpoch::Previous, spec)?;
     state.build_committee_cache(RelativeEpoch::Current, spec)?;
     state.build_committee_cache(RelativeEpoch::Next, spec)?;
+    state.build_total_active_balance_cache_at(state.current_epoch(), spec)?;
+    initialize_epoch_cache(state, state.current_epoch(), spec)?;
 
     // Load the struct we use to assign validators into sets based on their participation.
     //

consensus/state_processing/src/per_epoch_processing/capella.rs

@@ -24,6 +24,8 @@ pub fn process_epoch<T: EthSpec>(
     state.build_committee_cache(RelativeEpoch::Previous, spec)?;
     state.build_committee_cache(RelativeEpoch::Current, spec)?;
     state.build_committee_cache(RelativeEpoch::Next, spec)?;
+    state.build_total_active_balance_cache_at(state.current_epoch(), spec)?;
+    initialize_epoch_cache(state, state.current_epoch(), spec)?;
 
     // Pre-compute participating indices and total balances.
     let mut participation_cache = ParticipationCache::new(state, spec)?;

consensus/types/src/beacon_state.rs

@@ -1468,10 +1468,24 @@ impl<T: EthSpec> BeaconState<T> {
     }
 
     /// Build the total active balance cache.
-    fn build_total_active_balance_cache(&mut self, spec: &ChainSpec) -> Result<(), Error> {
-        let current_epoch = self.current_epoch();
-        let total_active_balance = self.compute_total_active_balance(current_epoch, spec)?;
-        *self.total_active_balance_mut() = Some((current_epoch, total_active_balance));
+    pub fn build_total_active_balance_cache_at(
+        &mut self,
+        epoch: Epoch,
+        spec: &ChainSpec,
+    ) -> Result<(), Error> {
+        if self.get_total_active_balance_at_epoch(epoch).is_err() {
+            self.force_build_total_active_balance_cache_at(epoch, spec)?;
+        }
+        Ok(())
+    }
+
+    pub fn force_build_total_active_balance_cache_at(
+        &mut self,
+        epoch: Epoch,
+        spec: &ChainSpec,
+    ) -> Result<(), Error> {
+        let total_active_balance = self.compute_total_active_balance(epoch, spec)?;
+        *self.total_active_balance_mut() = Some((epoch, total_active_balance));
         Ok(())
     }
 
@@ -1552,6 +1566,7 @@ impl<T: EthSpec> BeaconState<T> {
         self.drop_committee_cache(RelativeEpoch::Next)?;
         self.drop_pubkey_cache();
         *self.exit_cache_mut() = ExitCache::default();
+        *self.epoch_cache_mut() = EpochCache::default();
         Ok(())
     }
 
@@ -1580,7 +1595,7 @@ impl<T: EthSpec> BeaconState<T> {
         }
 
         if self.total_active_balance().is_none() && relative_epoch == RelativeEpoch::Current {
-            self.build_total_active_balance_cache(spec)?;
+            self.build_total_active_balance_cache_at(self.current_epoch(), spec)?;
         }
         Ok(())
     }
@@ -1874,7 +1889,7 @@ impl<T: EthSpec> BeaconState<T> {
                 // Ensure total active balance cache remains built whenever current committee
                 // cache is built.
                 if epoch == self.current_epoch() {
-                    self.build_total_active_balance_cache(spec)?;
+                    self.build_total_active_balance_cache_at(self.current_epoch(), spec)?;
                 }
             }
         }

testing/ef_tests/src/cases/epoch_processing.rs

@@ -5,6 +5,7 @@ use crate::decode::{ssz_decode_state, yaml_decode_file};
 use crate::type_name;
 use crate::type_name::TypeName;
 use serde_derive::Deserialize;
+use state_processing::epoch_cache::initialize_epoch_cache;
 use state_processing::per_epoch_processing::capella::process_historical_summaries_update;
 use state_processing::per_epoch_processing::{
     altair, base,
@@ -135,6 +136,7 @@ impl<E: EthSpec> EpochTransition<E> for RewardsAndPenalties {
 
 impl<E: EthSpec> EpochTransition<E> for RegistryUpdates {
     fn run(state: &mut BeaconState<E>, spec: &ChainSpec) -> Result<(), EpochProcessingError> {
+        initialize_epoch_cache(state, state.current_epoch(), spec)?;
         process_registry_updates(state, spec)
     }
 }

testing/ef_tests/src/cases/operations.rs

@@ -4,6 +4,7 @@ use crate::case_result::{check_state_diff, compare_beacon_state_results_without_
 use crate::decode::{ssz_decode_file, ssz_decode_file_with, ssz_decode_state, yaml_decode_file};
 use crate::testing_spec;
 use serde_derive::Deserialize;
+use state_processing::epoch_cache::initialize_epoch_cache;
 use state_processing::{
     per_block_processing::{
         errors::BlockProcessingError,
@@ -86,6 +87,7 @@ impl<E: EthSpec> Operation<E> for Attestation<E> {
         spec: &ChainSpec,
         _: &Operations<E, Self>,
     ) -> Result<(), BlockProcessingError> {
+        initialize_epoch_cache(state, state.current_epoch(), spec)?;
         let mut ctxt = ConsensusContext::new(state.slot());
         match state {
             BeaconState::Base(_) => base::process_attestations(