Add more electra helpers (#5653)

* Add new helpers

* Fix some stuff

* Fix compilation errors

* lint

* Address review

Author: pawanjay176

consensus/types/src/beacon_state.rs

@@ -2092,6 +2092,38 @@ impl<E: EthSpec> BeaconState<E> {
             .map_err(Into::into)
     }
 
+    /// Get active balance for the given `validator_index`.
+    pub fn get_active_balance(
+        &self,
+        validator_index: usize,
+        spec: &ChainSpec,
+    ) -> Result<u64, Error> {
+        let max_effective_balance = self
+            .validators()
+            .get(validator_index)
+            .map(|validator| validator.get_validator_max_effective_balance(spec))
+            .ok_or(Error::UnknownValidator(validator_index))?;
+        Ok(std::cmp::min(
+            *self
+                .balances()
+                .get(validator_index)
+                .ok_or(Error::UnknownValidator(validator_index))?,
+            max_effective_balance,
+        ))
+    }
+
+    pub fn get_pending_balance_to_withdraw(&self, validator_index: usize) -> Result<u64, Error> {
+        let mut pending_balance = 0;
+        for withdrawal in self
+            .pending_partial_withdrawals()?
+            .iter()
+            .filter(|withdrawal| withdrawal.index as usize == validator_index)
+        {
+            pending_balance.safe_add_assign(withdrawal.amount)?;
+        }
+        Ok(pending_balance)
+    }
+
     // ******* Electra mutators *******
 
     pub fn queue_excess_active_balance(
@@ -2142,6 +2174,101 @@ impl<E: EthSpec> BeaconState<E> {
             .map_err(Into::into)
     }
 
+    /// Change the withdrawal prefix of the given `validator_index` to the compounding withdrawal validator prefix.
+    pub fn switch_to_compounding_validator(
+        &mut self,
+        validator_index: usize,
+        spec: &ChainSpec,
+    ) -> Result<(), Error> {
+        let validator = self
+            .validators_mut()
+            .get_mut(validator_index)
+            .ok_or(Error::UnknownValidator(validator_index))?;
+        if validator.has_eth1_withdrawal_credential(spec) {
+            validator.withdrawal_credentials.as_fixed_bytes_mut()[0] =
+                spec.compounding_withdrawal_prefix_byte;
+            self.queue_excess_active_balance(validator_index, spec)?;
+        }
+        Ok(())
+    }
+
+    pub fn compute_exit_epoch_and_update_churn(
+        &mut self,
+        exit_balance: u64,
+        spec: &ChainSpec,
+    ) -> Result<Epoch, Error> {
+        let mut earliest_exit_epoch = std::cmp::max(
+            self.earliest_exit_epoch()?,
+            self.compute_activation_exit_epoch(self.current_epoch(), spec)?,
+        );
+
+        let per_epoch_churn = self.get_activation_exit_churn_limit(spec)?;
+        // New epoch for exits
+        let mut exit_balance_to_consume = if self.earliest_exit_epoch()? < earliest_exit_epoch {
+            per_epoch_churn
+        } else {
+            self.exit_balance_to_consume()?
+        };
+
+        // Exit doesn't fit in the current earliest epoch
+        if exit_balance > exit_balance_to_consume {
+            let balance_to_process = exit_balance.safe_sub(exit_balance_to_consume)?;
+            let additional_epochs = balance_to_process
+                .safe_sub(1)?
+                .safe_div(per_epoch_churn)?
+                .safe_add(1)?;
+            earliest_exit_epoch.safe_add_assign(additional_epochs)?;
+            exit_balance_to_consume
+                .safe_add_assign(additional_epochs.safe_mul(per_epoch_churn)?)?;
+        }
+        let state = self.as_electra_mut()?;
+        // Consume the balance and update state variables
+        state.exit_balance_to_consume = exit_balance_to_consume.safe_sub(exit_balance)?;
+        state.earliest_exit_epoch = earliest_exit_epoch;
+
+        Ok(state.earliest_exit_epoch)
+    }
+
+    pub fn compute_consolidation_epoch_and_update_churn(
+        &mut self,
+        consolidation_balance: u64,
+        spec: &ChainSpec,
+    ) -> Result<Epoch, Error> {
+        let mut earliest_consolidation_epoch = std::cmp::max(
+            self.earliest_consolidation_epoch()?,
+            self.compute_activation_exit_epoch(self.current_epoch(), spec)?,
+        );
+
+        let per_epoch_consolidation_churn = self.get_consolidation_churn_limit(spec)?;
+
+        // New epoch for consolidations
+        let mut consolidation_balance_to_consume =
+            if self.earliest_consolidation_epoch()? < earliest_consolidation_epoch {
+                per_epoch_consolidation_churn
+            } else {
+                self.consolidation_balance_to_consume()?
+            };
+        // Consolidation doesn't fit in the current earliest epoch
+        if consolidation_balance > consolidation_balance_to_consume {
+            let balance_to_process =
+                consolidation_balance.safe_sub(consolidation_balance_to_consume)?;
+            let additional_epochs = balance_to_process
+                .safe_sub(1)?
+                .safe_div(per_epoch_consolidation_churn)?
+                .safe_add(1)?;
+            earliest_consolidation_epoch.safe_add_assign(additional_epochs)?;
+            consolidation_balance_to_consume
+                .safe_add_assign(additional_epochs.safe_mul(per_epoch_consolidation_churn)?)?;
+        }
+        // Consume the balance and update state variables
+        let state = self.as_electra_mut()?;
+        state.consolidation_balance_to_consume =
+            consolidation_balance_to_consume.safe_sub(consolidation_balance)?;
+        state.earliest_consolidation_epoch = earliest_consolidation_epoch;
+
+        Ok(state.earliest_consolidation_epoch)
+    }
+
     #[allow(clippy::arithmetic_side_effects)]
     pub fn rebase_on(&mut self, base: &Self, spec: &ChainSpec) -> Result<(), Error> {
         // Required for macros (which use type-hints internally).

consensus/types/src/validator.rs

@@ -57,10 +57,10 @@ impl Validator {
 
     /// Returns `true` if the validator is eligible to join the activation queue.
     ///
-    /// Spec v0.12.1
+    /// Modified in electra
     pub fn is_eligible_for_activation_queue(&self, spec: &ChainSpec) -> bool {
         self.activation_eligibility_epoch == spec.far_future_epoch
-            && self.effective_balance == spec.max_effective_balance
+            && self.effective_balance >= spec.min_activation_balance
     }
 
     /// Returns `true` if the validator is eligible to be activated.
@@ -130,15 +130,39 @@ impl Validator {
     }
 
     /// Returns `true` if the validator is fully withdrawable at some epoch.
+    ///
+    /// Note: Modified in electra.
     pub fn is_fully_withdrawable_at(&self, balance: u64, epoch: Epoch, spec: &ChainSpec) -> bool {
-        self.has_eth1_withdrawal_credential(spec) && self.withdrawable_epoch <= epoch && balance > 0
+        self.has_execution_withdrawal_credential(spec)
+            && self.withdrawable_epoch <= epoch
+            && balance > 0
     }
 
     /// Returns `true` if the validator is partially withdrawable.
+    ///
+    /// Note: Modified in electra.
     pub fn is_partially_withdrawable_validator(&self, balance: u64, spec: &ChainSpec) -> bool {
-        self.has_eth1_withdrawal_credential(spec)
-            && self.effective_balance == spec.max_effective_balance
-            && balance > spec.max_effective_balance
+        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;
+        self.has_execution_withdrawal_credential(spec)
+            && has_max_effective_balance
+            && has_excess_balance
+    }
+
+    /// Returns `true` if the validator has a 0x01 or 0x02 prefixed withdrawal credential.
+    pub fn has_execution_withdrawal_credential(&self, spec: &ChainSpec) -> bool {
+        self.has_compounding_withdrawal_credential(spec)
+            || self.has_eth1_withdrawal_credential(spec)
+    }
+
+    /// Returns the max effective balance for a validator in gwei.
+    pub fn get_validator_max_effective_balance(&self, spec: &ChainSpec) -> u64 {
+        if self.has_compounding_withdrawal_credential(spec) {
+            spec.max_effective_balance_electra
+        } else {
+            spec.min_activation_balance
+        }
     }
 }