Make modified helpers in electra fork aware (#5698)

* Make modified helpers in electra fork aware

* Make more functions fork aware

* formatting fixes only.

consensus/state_processing/src/per_block_processing.rs

@@ -508,6 +508,7 @@ pub fn get_expected_withdrawals<E: EthSpec>(
     let mut withdrawal_index = state.next_withdrawal_index()?;
     let mut validator_index = state.next_withdrawal_validator_index()?;
     let mut withdrawals = vec![];
+    let fork_name = state.fork_name_unchecked();
 
     let bound = std::cmp::min(
         state.validators().len() as u64,
@@ -518,7 +519,7 @@ pub fn get_expected_withdrawals<E: EthSpec>(
         let balance = *state.balances().get(validator_index as usize).ok_or(
             BeaconStateError::BalancesOutOfBounds(validator_index as usize),
         )?;
-        if validator.is_fully_withdrawable_at(balance, epoch, spec) {
+        if validator.is_fully_withdrawable_at(balance, epoch, spec, fork_name) {
             withdrawals.push(Withdrawal {
                 index: withdrawal_index,
                 validator_index,
@@ -528,7 +529,7 @@ pub fn get_expected_withdrawals<E: EthSpec>(
                 amount: balance,
             });
             withdrawal_index.safe_add_assign(1)?;
-        } else if validator.is_partially_withdrawable_validator(balance, spec) {
+        } else if validator.is_partially_withdrawable_validator(balance, spec, fork_name) {
             withdrawals.push(Withdrawal {
                 index: withdrawal_index,
                 validator_index,

consensus/state_processing/src/per_epoch_processing/registry_updates.rs

@@ -19,19 +19,20 @@ pub fn process_registry_updates<E: EthSpec>(
         validator.is_active_at(current_epoch)
             && validator.effective_balance <= spec.ejection_balance
     };
+    let fork_name = state.fork_name_unchecked();
     let indices_to_update: Vec<_> = state
         .validators()
         .iter()
         .enumerate()
         .filter(|(_, validator)| {
-            validator.is_eligible_for_activation_queue(spec) || is_ejectable(validator)
+            validator.is_eligible_for_activation_queue(spec, fork_name) || is_ejectable(validator)
         })
         .map(|(idx, _)| idx)
         .collect();
 
     for index in indices_to_update {
         let validator = state.get_validator_mut(index)?;
-        if validator.is_eligible_for_activation_queue(spec) {
+        if validator.is_eligible_for_activation_queue(spec, fork_name) {
             validator.activation_eligibility_epoch = current_epoch.safe_add(1)?;
         }
         if is_ejectable(validator) {

consensus/state_processing/src/per_epoch_processing/single_pass.rs

@@ -466,7 +466,7 @@ fn process_single_registry_update(
 ) -> Result<(), Error> {
     let current_epoch = state_ctxt.current_epoch;
 
-    if validator.is_eligible_for_activation_queue(spec) {
+    if validator.is_eligible_for_activation_queue(spec, state_ctxt.fork_name) {
         validator.make_mut()?.activation_eligibility_epoch = current_epoch.safe_add(1)?;
     }
 

consensus/types/src/validator.rs

@@ -1,5 +1,5 @@
 use crate::{
-    test_utils::TestRandom, Address, BeaconState, ChainSpec, Epoch, EthSpec, Hash256,
+    test_utils::TestRandom, Address, BeaconState, ChainSpec, Epoch, EthSpec, ForkName, Hash256,
     PublicKeyBytes,
 };
 use serde::{Deserialize, Serialize};
@@ -57,8 +57,31 @@ impl Validator {
 
     /// Returns `true` if the validator is eligible to join the activation queue.
     ///
-    /// Modified in electra
+    /// Calls the correct function depending on the provided `fork_name`.
-    pub fn is_eligible_for_activation_queue(&self, spec: &ChainSpec) -> bool {
+    pub fn is_eligible_for_activation_queue(
+        &self,
+        spec: &ChainSpec,
+        current_fork: ForkName,
+    ) -> bool {
+        if current_fork >= ForkName::Electra {
+            self.is_eligible_for_activation_queue_electra(spec)
+        } else {
+            self.is_eligible_for_activation_queue_base(spec)
+        }
+    }
+
+    /// Returns `true` if the validator is eligible to join the activation queue.
+    ///
+    /// Spec v0.12.1
+    fn is_eligible_for_activation_queue_base(&self, spec: &ChainSpec) -> bool {
+        self.activation_eligibility_epoch == spec.far_future_epoch
+            && self.effective_balance == spec.max_effective_balance
+    }
+
+    /// Returns `true` if the validator is eligible to join the activation queue.
+    ///
+    /// Modified in electra as part of EIP 7251.
+    fn is_eligible_for_activation_queue_electra(&self, spec: &ChainSpec) -> bool {
         self.activation_eligibility_epoch == spec.far_future_epoch
             && self.effective_balance >= spec.min_activation_balance
     }
@@ -131,8 +154,40 @@ impl Validator {
 
     /// Returns `true` if the validator is fully withdrawable at some epoch.
     ///
-    /// Note: Modified in electra.
+    /// Calls the correct function depending on the provided `fork_name`.
-    pub fn is_fully_withdrawable_at(&self, balance: u64, epoch: Epoch, spec: &ChainSpec) -> bool {
+    pub fn is_fully_withdrawable_at(
+        &self,
+        balance: u64,
+        epoch: Epoch,
+        spec: &ChainSpec,
+        current_fork: ForkName,
+    ) -> bool {
+        if current_fork >= ForkName::Electra {
+            self.is_fully_withdrawable_at_electra(balance, epoch, spec)
+        } else {
+            self.is_fully_withdrawable_at_capella(balance, epoch, spec)
+        }
+    }
+
+    /// Returns `true` if the validator is fully withdrawable at some epoch.
+    fn is_fully_withdrawable_at_capella(
+        &self,
+        balance: u64,
+        epoch: Epoch,
+        spec: &ChainSpec,
+    ) -> bool {
+        self.has_eth1_withdrawal_credential(spec) && self.withdrawable_epoch <= epoch && balance > 0
+    }
+
+    /// Returns `true` if the validator is fully withdrawable at some epoch.
+    ///
+    /// Modified in electra as part of EIP 7251.
+    fn is_fully_withdrawable_at_electra(
+        &self,
+        balance: u64,
+        epoch: Epoch,
+        spec: &ChainSpec,
+    ) -> bool {
         self.has_execution_withdrawal_credential(spec)
             && self.withdrawable_epoch <= epoch
             && balance > 0
@@ -140,8 +195,35 @@ impl Validator {
 
     /// Returns `true` if the validator is partially withdrawable.
     ///
-    /// Note: Modified in electra.
+    /// Calls the correct function depending on the provided `fork_name`.
-    pub fn is_partially_withdrawable_validator(&self, balance: u64, spec: &ChainSpec) -> bool {
+    pub fn is_partially_withdrawable_validator(
+        &self,
+        balance: u64,
+        spec: &ChainSpec,
+        current_fork: ForkName,
+    ) -> bool {
+        if current_fork >= ForkName::Electra {
+            self.is_partially_withdrawable_validator_electra(balance, spec)
+        } else {
+            self.is_partially_withdrawable_validator_capella(balance, spec)
+        }
+    }
+
+    /// Returns `true` if the validator is partially withdrawable.
+    fn is_partially_withdrawable_validator_capella(&self, balance: u64, spec: &ChainSpec) -> bool {
+        self.has_eth1_withdrawal_credential(spec)
+            && self.effective_balance == spec.max_effective_balance
+            && balance > spec.max_effective_balance
+    }
+
+    /// Returns `true` if the validator is partially withdrawable.
+    ///
+    /// Modified in electra as part of EIP 7251.
+    pub fn is_partially_withdrawable_validator_electra(
+        &self,
+        balance: u64,
+        spec: &ChainSpec,
+    ) -> bool {
         let max_effective_balance = self.get_validator_max_effective_balance(spec);
         let has_max_effective_balance = self.effective_balance == max_effective_balance;
         let has_excess_balance = balance > max_effective_balance;