Electra epoch processing

consensus/state_processing/src/common/initiate_validator_exit.rs

@@ -8,12 +8,12 @@ pub fn initiate_validator_exit<E: EthSpec>(
     index: usize,
     spec: &ChainSpec,
 ) -> Result<(), Error> {
-    // We do things in a slightly different order to the spec here. Instead of immediately checking
-    // whether the validator has already exited, we instead prepare the exit cache and compute the
-    // cheap-to-calculate values from that. *Then* we look up the validator a single time in the
-    // validator tree (expensive), make the check and mutate as appropriate. Compared to the spec
-    // ordering, this saves us from looking up the validator in the validator registry multiple
-    // times.
+    let validator = state.get_validator_cow(index)?;
+
+    // Return if the validator already initiated exit
+    if validator.exit_epoch != spec.far_future_epoch {
+        return Ok(());
+    }
 
     // Ensure the exit cache is built.
     state.build_exit_cache(spec)?;
@@ -36,14 +36,7 @@ pub fn initiate_validator_exit<E: EthSpec>(
         exit_queue_epoch
     };
 
-    let validator = state.get_validator_cow(index)?;
-
-    // Return if the validator already initiated exit
-    if validator.exit_epoch != spec.far_future_epoch {
-        return Ok(());
-    }
-
-    let validator = validator.into_mut()?;
+    let validator = state.get_validator_mut(index)?;
     validator.exit_epoch = exit_queue_epoch;
     validator.withdrawable_epoch =
         exit_queue_epoch.safe_add(spec.min_validator_withdrawability_delay)?;

consensus/state_processing/src/epoch_cache.rs

@@ -9,6 +9,7 @@ use types::{ActivationQueue, BeaconState, ChainSpec, EthSpec, ForkName, Hash256}
 pub struct PreEpochCache {
     epoch_key: EpochCacheKey,
     effective_balances: Vec<u64>,
+    total_active_balance: u64,
 }
 
 impl PreEpochCache {
@@ -36,27 +37,59 @@ impl PreEpochCache {
         Ok(Self {
             epoch_key,
             effective_balances: Vec::with_capacity(state.validators().len()),
+            total_active_balance: 0,
         })
     }
 
-    pub fn push_effective_balance(&mut self, effective_balance: u64) {
-        self.effective_balances.push(effective_balance);
+    pub fn update_effective_balance(
+        &mut self,
+        validator_index: usize,
+        effective_balance: u64,
+        is_active_next_epoch: bool,
+    ) -> Result<(), EpochCacheError> {
+        if validator_index == self.effective_balances.len() {
+            self.effective_balances.push(effective_balance);
+            if is_active_next_epoch {
+                self.total_active_balance
+                    .safe_add_assign(effective_balance)?;
+            }
+
+            Ok(())
+        } else if let Some(existing_balance) = self.effective_balances.get_mut(validator_index) {
+            // Update total active balance for a late change in effective balance. This happens when
+            // processing consolidations.
+            if is_active_next_epoch {
+                self.total_active_balance
+                    .safe_add_assign(effective_balance)?;
+                self.total_active_balance
+                    .safe_sub_assign(*existing_balance)?;
+            }
+            *existing_balance = effective_balance;
+            Ok(())
+        } else {
+            Err(EpochCacheError::ValidatorIndexOutOfBounds { validator_index })
+        }
+    }
+
+    pub fn get_total_active_balance(&self) -> u64 {
+        self.total_active_balance
     }
 
     pub fn into_epoch_cache(
         self,
-        total_active_balance: u64,
         activation_queue: ActivationQueue,
         spec: &ChainSpec,
     ) -> Result<EpochCache, EpochCacheError> {
         let epoch = self.epoch_key.epoch;
+        let total_active_balance = self.total_active_balance;
         let sqrt_total_active_balance = SqrtTotalActiveBalance::new(total_active_balance);
         let base_reward_per_increment = BaseRewardPerIncrement::new(total_active_balance, spec)?;
 
         let effective_balance_increment = spec.effective_balance_increment;
-        let max_effective_balance_eth = spec
-            .max_effective_balance
-            .safe_div(effective_balance_increment)?;
+        let max_effective_balance =
+            spec.max_effective_balance_for_fork(spec.fork_name_at_epoch(epoch));
+        let max_effective_balance_eth =
+            max_effective_balance.safe_div(effective_balance_increment)?;
 
         let mut base_rewards = Vec::with_capacity(max_effective_balance_eth.safe_add(1)? as usize);
 
@@ -131,9 +164,9 @@ pub fn initialize_epoch_cache<E: EthSpec>(
             decision_block_root,
         },
         effective_balances,
+        total_active_balance,
     };
-    *state.epoch_cache_mut() =
-        pre_epoch_cache.into_epoch_cache(total_active_balance, activation_queue, spec)?;
+    *state.epoch_cache_mut() = pre_epoch_cache.into_epoch_cache(activation_queue, spec)?;
 
     Ok(())
 }

consensus/state_processing/src/per_epoch_processing/errors.rs

@@ -25,6 +25,9 @@ pub enum EpochProcessingError {
     InvalidFlagIndex(usize),
     MilhouseError(milhouse::Error),
     EpochCache(EpochCacheError),
+    SinglePassMissingActivationQueue,
+    MissingEarliestExitEpoch,
+    MissingExitBalanceToConsume,
 }
 
 impl From<InclusionError> for EpochProcessingError {