diff --git a/consensus/state_processing/src/per_epoch_processing/altair.rs b/consensus/state_processing/src/per_epoch_processing/altair.rs index 5c635609b9b..a52580c7bfb 100644 --- a/consensus/state_processing/src/per_epoch_processing/altair.rs +++ b/consensus/state_processing/src/per_epoch_processing/altair.rs @@ -28,6 +28,8 @@ pub fn process_epoch( 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)?; diff --git a/consensus/state_processing/src/per_epoch_processing/base.rs b/consensus/state_processing/src/per_epoch_processing/base.rs index 7363ffd07b2..478d620cc32 100644 --- a/consensus/state_processing/src/per_epoch_processing/base.rs +++ b/consensus/state_processing/src/per_epoch_processing/base.rs @@ -24,6 +24,8 @@ pub fn process_epoch( 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. // diff --git a/consensus/state_processing/src/per_epoch_processing/capella.rs b/consensus/state_processing/src/per_epoch_processing/capella.rs index 34a1b5cfda2..656cfb67fa0 100644 --- a/consensus/state_processing/src/per_epoch_processing/capella.rs +++ b/consensus/state_processing/src/per_epoch_processing/capella.rs @@ -24,6 +24,8 @@ pub fn process_epoch( 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)?; diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs index 336c906e58d..c44430edf82 100644 --- a/consensus/types/src/beacon_state.rs +++ b/consensus/types/src/beacon_state.rs @@ -1468,10 +1468,24 @@ impl BeaconState { } /// 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 BeaconState { 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 BeaconState { } 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 BeaconState { // 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)?; } } } diff --git a/testing/ef_tests/src/cases/epoch_processing.rs b/testing/ef_tests/src/cases/epoch_processing.rs index 62ccb792c6e..40c75cfbaf5 100644 --- a/testing/ef_tests/src/cases/epoch_processing.rs +++ b/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 EpochTransition for RewardsAndPenalties { impl EpochTransition for RegistryUpdates { fn run(state: &mut BeaconState, spec: &ChainSpec) -> Result<(), EpochProcessingError> { + initialize_epoch_cache(state, state.current_epoch(), spec)?; process_registry_updates(state, spec) } } diff --git a/testing/ef_tests/src/cases/operations.rs b/testing/ef_tests/src/cases/operations.rs index d8b890b8018..313cc6780e5 100644 --- a/testing/ef_tests/src/cases/operations.rs +++ b/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 Operation for Attestation { spec: &ChainSpec, _: &Operations, ) -> Result<(), BlockProcessingError> { + initialize_epoch_cache(state, state.current_epoch(), spec)?; let mut ctxt = ConsensusContext::new(state.slot()); match state { BeaconState::Base(_) => base::process_attestations(