Simplify total balance cache building

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -4953,7 +4953,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             metrics::start_timer(&metrics::BLOCK_PRODUCTION_ATTESTATION_TIMES);
 
         // Epoch cache and total balance cache are required for op pool packing.
-        state.build_total_active_balance_cache_at(state.current_epoch(), &self.spec)?;
+        state.build_total_active_balance_cache(&self.spec)?;
         initialize_epoch_cache(&mut state, &self.spec)?;
 
         let mut prev_filter_cache = HashMap::new();

consensus/state_processing/src/epoch_cache.rs

@@ -104,7 +104,7 @@ pub fn initialize_epoch_cache<E: EthSpec>(
         .proposer_shuffling_decision_root(Hash256::zero())
         .map_err(EpochCacheError::BeaconState)?;
 
-    state.build_total_active_balance_cache_at(current_epoch, spec)?;
+    state.build_total_active_balance_cache(spec)?;
     let total_active_balance = state.get_total_active_balance_at_epoch(current_epoch)?;
 
     // Collect effective balances and compute activation queue.

consensus/state_processing/src/per_epoch_processing/altair.rs

@@ -30,7 +30,7 @@ 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)?;
+    state.build_total_active_balance_cache(spec)?;
     initialize_epoch_cache(state, spec)?;
     initialize_progressive_balances_cache::<T>(state, spec)?;
 

consensus/state_processing/src/per_epoch_processing/base.rs

@@ -24,7 +24,7 @@ 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)?;
+    state.build_total_active_balance_cache(spec)?;
     initialize_epoch_cache(state, spec)?;
 
     // Load the struct we use to assign validators into sets based on their participation.

consensus/types/src/beacon_state.rs

@@ -1480,19 +1480,13 @@ impl<T: EthSpec> BeaconState<T> {
     ///
     /// This method should rarely be invoked because single-pass epoch processing keeps the total
     /// active balance cache up to date.
-    pub fn compute_total_active_balance_slow(
-        &self,
-        epoch: Epoch,
-        spec: &ChainSpec,
-    ) -> Result<u64, Error> {
-        if epoch != self.current_epoch() && epoch != self.next_epoch()? {
-            return Err(Error::EpochOutOfBounds);
-        }
+    pub fn compute_total_active_balance_slow(&self, spec: &ChainSpec) -> Result<u64, Error> {
+        let current_epoch = self.current_epoch();
 
         let mut total_active_balance = 0;
 
         for validator in self.validators() {
-            if validator.is_active_at(epoch) {
+            if validator.is_active_at(current_epoch) {
                 total_active_balance.safe_add_assign(validator.effective_balance)?;
             }
         }
@@ -1536,26 +1530,24 @@ impl<T: EthSpec> BeaconState<T> {
         *self.total_active_balance_mut() = Some((epoch, balance));
     }
 
-    /// Build the total active balance cache.
-    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)?;
+    /// Build the total active balance cache for the current epoch if it is not already built.
+    pub fn build_total_active_balance_cache(&mut self, spec: &ChainSpec) -> Result<(), Error> {
+        if self
+            .get_total_active_balance_at_epoch(self.current_epoch())
+            .is_err()
+        {
+            self.force_build_total_active_balance_cache(spec)?;
         }
         Ok(())
     }
 
     /// Build the total active balance cache, even if it is already built.
-    pub fn force_build_total_active_balance_cache_at(
+    pub fn force_build_total_active_balance_cache(
         &mut self,
-        epoch: Epoch,
         spec: &ChainSpec,
     ) -> Result<(), Error> {
-        let total_active_balance = self.compute_total_active_balance_slow(epoch, spec)?;
-        *self.total_active_balance_mut() = Some((epoch, total_active_balance));
+        let total_active_balance = self.compute_total_active_balance_slow(spec)?;
+        *self.total_active_balance_mut() = Some((self.current_epoch(), total_active_balance));
         Ok(())
     }
 
@@ -1685,7 +1677,7 @@ impl<T: EthSpec> BeaconState<T> {
         }
 
         if self.total_active_balance().is_none() && relative_epoch == RelativeEpoch::Current {
-            self.build_total_active_balance_cache_at(self.current_epoch(), spec)?;
+            self.build_total_active_balance_cache(spec)?;
         }
         Ok(())
     }