Skip to content

Commit

Permalink
Add new helpers
Browse files Browse the repository at this point in the history
  • Loading branch information
pawanjay176 committed Apr 25, 2024
1 parent d2de5f3 commit 71defde
Show file tree
Hide file tree
Showing 2 changed files with 164 additions and 5 deletions.
127 changes: 127 additions & 0 deletions consensus/types/src/beacon_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2092,6 +2092,40 @@ impl<E: EthSpec> BeaconState<E> {
.map_err(Into::into)
}

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))?;
// 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<u64, Error> {
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(
Expand Down Expand Up @@ -2142,6 +2176,99 @@ impl<E: EthSpec> BeaconState<E> {
.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<Epoch, Error> {
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<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
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).
Expand Down
42 changes: 37 additions & 5 deletions consensus/types/src/validator.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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
}
}
}

Expand Down

0 comments on commit 71defde

Please sign in to comment.