From 71defde858d157727a7ccecfbbe7ca3e9011714c Mon Sep 17 00:00:00 2001 From: Pawan Dhananjay Date: Thu, 25 Apr 2024 16:28:34 -0700 Subject: [PATCH] Add new helpers --- consensus/types/src/beacon_state.rs | 127 ++++++++++++++++++++++++++++ consensus/types/src/validator.rs | 42 +++++++-- 2 files changed, 164 insertions(+), 5 deletions(-) diff --git a/consensus/types/src/beacon_state.rs b/consensus/types/src/beacon_state.rs index 1504f1f6eee..523f652aea1 100644 --- a/consensus/types/src/beacon_state.rs +++ b/consensus/types/src/beacon_state.rs @@ -2092,6 +2092,40 @@ impl BeaconState { .map_err(Into::into) } + pub fn get_active_balance( + &self, + validator_index: usize, + spec: &ChainSpec, + ) -> Result { + let max_effective_balance = self + .validators() + .get(validator_index) + .map(|validator| validator.get_validator_max_effective_balance(spec)) + .ok_or(Error::UnknownValidator(validator_index))?; + // TODO(pawan): this is assuming balances and validat + 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: u64) -> Result { + Ok(self + .pending_partial_withdrawals()? + .iter() + .filter_map(|withdrawal| { + if withdrawal.index == validator_index { + Some(withdrawal.amount) + } else { + None + } + }) + .sum()) + } + // ******* Electra mutators ******* pub fn queue_excess_active_balance( @@ -2142,6 +2176,99 @@ impl BeaconState { .map_err(Into::into) } + pub fn switch_to_compounding_validator( + &mut self, + validator_index: u64, + 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 = + spec.compounding_withdrawal_prefix_byte + validator.withdrawal_credentials[1..]; + self.queue_excess_active_balance(validator_index, spec) + } else { + Ok(()) + } + } + + pub fn compute_exit_epoch_and_update_churn( + &self, + exit_balance: u64, + spec: &ChainSpec, + ) -> Result { + let mut earliest_exit_epoch = std::cmp::max( + self.earliest_exit_epoch()?, + self.compute_activation_exit_epoch(self.current_epoch()), + ); + + 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 - 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)?)?; + } + // Consume the balance and update state variables + self.exit_balance_to_consume_mut() = exit_balance_to_consume.safe_sub(exit_balance)?; + self.earliest_exit_epoch_mut() = earliest_exit_epoch; + + Ok(self.earliest_exit_epoch()) + } + + pub fn compute_consolidation_epoch_and_update_churn( + &mut self, + consolidation_balance: u64, + spec: &ChainSpec, + ) -> Result { + 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 + self.consolidation_balance_to_consume_mut() = + consolidation_balance_to_consume.safe_sub(consolidation_balance)?; + self.earliest_consolidation_epoch_mut() = earliest_consolidation_epoch; + + Ok(self.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). diff --git a/consensus/types/src/validator.rs b/consensus/types/src/validator.rs index 5fd18552f6b..c70c08ec937 100644 --- a/consensus/types/src/validator.rs +++ b/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. @@ -135,10 +135,42 @@ impl Validator { } /// 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 `true` if the validator if fully withdrawable. + pub fn is_fully_withdrawable_validator( + &self, + balance: u64, + epoch: Epoch, + spec: &ChainSpec, + ) -> bool { + self.has_execution_withdrawal_credential(spec) + && self.withdrawable_epoch <= epoch + && balance > 0 + } + + /// 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 + } } }