Update get_expected_withdrawals to return processed_validators_sweep_count.

Author: jimmygchen

beacon_node/beacon_chain/src/beacon_chain.rs

@@ -4794,7 +4794,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
         let proposal_epoch = proposal_slot.epoch(T::EthSpec::slots_per_epoch());
         if head_state.current_epoch() == proposal_epoch {
             return get_expected_withdrawals(&unadvanced_state, &self.spec)
-                .map(|(withdrawals, _, _)| withdrawals)
+                .map(|(withdrawals, _, _, _)| withdrawals)
                 .map_err(Error::PrepareProposerFailed);
         }
 
@@ -4812,7 +4812,7 @@ impl<T: BeaconChainTypes> BeaconChain<T> {
             &self.spec,
         )?;
         get_expected_withdrawals(&advanced_state, &self.spec)
-            .map(|(withdrawals, _, _)| withdrawals)
+            .map(|(withdrawals, _, _, _)| withdrawals)
             .map_err(Error::PrepareProposerFailed)
     }
 

beacon_node/http_api/src/builder_states.rs

@@ -32,7 +32,7 @@ pub fn get_next_withdrawals<T: BeaconChainTypes>(
     }
 
     match get_expected_withdrawals(&state, &chain.spec) {
-        Ok((withdrawals, _, _)) => Ok(withdrawals),
+        Ok((withdrawals, _, _, _)) => Ok(withdrawals),
         Err(e) => Err(warp_utils::reject::custom_server_error(format!(
             "failed to get expected withdrawal: {:?}",
             e

consensus/state_processing/src/per_block_processing.rs

@@ -659,85 +659,109 @@ pub fn get_pending_partial_withdrawals<E: EthSpec>(
     Ok(Some(processed_count))
 }
 
-/// Compute the next batch of withdrawals which should be included in a block.
+/// Get withdrawals from the validator sweep.
 ///
-/// https://ethereum.github.io/consensus-specs/specs/gloas/beacon-chain/#modified-get_expected_withdrawals
-#[allow(clippy::type_complexity)]
-pub fn get_expected_withdrawals<E: EthSpec>(
+/// This function iterates through validators starting from `next_withdrawal_validator_index`
+/// and adds full or partial withdrawals for eligible validators.
+///
+/// https://ethereum.github.io/consensus-specs/specs/capella/beacon-chain/#new-get_validators_sweep_withdrawals
+pub fn get_validators_sweep_withdrawals<E: EthSpec>(
     state: &BeaconState<E>,
+    withdrawal_index: &mut u64,
+    withdrawals: &mut Vec<Withdrawal>,
     spec: &ChainSpec,
-) -> Result<(Withdrawals<E>, Option<usize>, Option<usize>), BlockProcessingError> {
+) -> Result<u64, BlockProcessingError> {
     let epoch = state.current_epoch();
-    let mut withdrawal_index = state.next_withdrawal_index()?;
-    let mut withdrawals = Vec::<Withdrawal>::with_capacity(E::max_withdrawals_per_payload());
     let fork_name = state.fork_name_unchecked();
-
     let mut validator_index = state.next_withdrawal_validator_index()?;
-
-    // [New in Gloas:EIP7732]
-    // Get builder withdrawals
-    let processed_builder_withdrawals_count =
-        get_builder_withdrawals(state, &mut withdrawal_index, &mut withdrawals)?;
-
-    // [New in Electra:EIP7251]
-    // Get partial withdrawals.
-    let processed_partial_withdrawals_count =
-        get_pending_partial_withdrawals(state, &mut withdrawal_index, &mut withdrawals, spec)?;
-
-    let bound = std::cmp::min(
+    let validators_limit = std::cmp::min(
         state.validators().len() as u64,
         spec.max_validators_per_withdrawals_sweep,
     );
-    for _ in 0..bound {
+    let withdrawals_limit = E::max_withdrawals_per_payload();
+
+    // There must be at least one space reserved for validator sweep withdrawals
+    block_verify!(
+        withdrawals.len() < withdrawals_limit,
+        BlockProcessingError::WithdrawalsLimitExceeded {
+            limit: withdrawals_limit,
+            prior_withdrawals: withdrawals.len()
+        }
+    );
+
+    let mut processed_count: u64 = 0;
+
+    for _ in 0..validators_limit {
+        if withdrawals.len() >= withdrawals_limit {
+            break;
+        }
+
         let validator = state.get_validator(validator_index as usize)?;
-        let partially_withdrawn_balance = withdrawals
-            .iter()
-            .filter_map(|withdrawal| {
-                (withdrawal.validator_index == validator_index).then_some(withdrawal.amount)
-            })
-            .safe_sum()?;
-        let balance = state
-            .balances()
-            .get(validator_index as usize)
-            .ok_or(BeaconStateError::BalancesOutOfBounds(
-                validator_index as usize,
-            ))?
-            .safe_sub(partially_withdrawn_balance)?;
+        let balance = get_balance_after_withdrawals(state, validator_index, withdrawals)?;
+
         if validator.is_fully_withdrawable_validator(balance, epoch, spec, fork_name) {
             withdrawals.push(Withdrawal {
-                index: withdrawal_index,
+                index: *withdrawal_index,
                 validator_index,
                 address: validator
-                    .get_execution_withdrawal_address(spec, state.fork_name_unchecked())
+                    .get_execution_withdrawal_address(spec, fork_name)
                     .ok_or(BlockProcessingError::WithdrawalCredentialsInvalid)?,
                 amount: balance,
             });
             withdrawal_index.safe_add_assign(1)?;
         } else if validator.is_partially_withdrawable_validator(balance, spec, fork_name) {
             withdrawals.push(Withdrawal {
-                index: withdrawal_index,
+                index: *withdrawal_index,
                 validator_index,
                 address: validator
-                    .get_execution_withdrawal_address(spec, state.fork_name_unchecked())
+                    .get_execution_withdrawal_address(spec, fork_name)
                     .ok_or(BlockProcessingError::WithdrawalCredentialsInvalid)?,
                 amount: balance.safe_sub(validator.get_max_effective_balance(spec, fork_name))?,
             });
             withdrawal_index.safe_add_assign(1)?;
         }
-        if withdrawals.len() == E::max_withdrawals_per_payload() {
-            break;
-        }
+
         validator_index = validator_index
             .safe_add(1)?
             .safe_rem(state.validators().len() as u64)?;
+        processed_count.safe_add_assign(1)?;
     }
 
+    Ok(processed_count)
+}
+
+/// Compute the next batch of withdrawals which should be included in a block.
+///
+/// https://ethereum.github.io/consensus-specs/specs/gloas/beacon-chain/#modified-get_expected_withdrawals
+#[allow(clippy::type_complexity)]
+pub fn get_expected_withdrawals<E: EthSpec>(
+    state: &BeaconState<E>,
+    spec: &ChainSpec,
+) -> Result<(Withdrawals<E>, Option<usize>, Option<usize>, u64), BlockProcessingError> {
+    let mut withdrawal_index = state.next_withdrawal_index()?;
+    let mut withdrawals = Vec::<Withdrawal>::with_capacity(E::max_withdrawals_per_payload());
+
+    // [New in Gloas:EIP7732]
+    // Get builder withdrawals
+    let processed_builder_withdrawals_count =
+        get_builder_withdrawals(state, &mut withdrawal_index, &mut withdrawals)?;
+
+    // [New in Electra:EIP7251]
+    // Get partial withdrawals.
+    let processed_partial_withdrawals_count =
+        get_pending_partial_withdrawals(state, &mut withdrawal_index, &mut withdrawals, spec)?;
+
+    // Get validators sweep withdrawals
+    let processed_validators_sweep_count =
+        get_validators_sweep_withdrawals(state, &mut withdrawal_index, &mut withdrawals, spec)?;
+
     Ok((
         withdrawals
             .try_into()
             .map_err(BlockProcessingError::SszTypesError)?,
         processed_builder_withdrawals_count,
         processed_partial_withdrawals_count,
+        processed_validators_sweep_count,
     ))
 }
 

consensus/state_processing/src/per_block_processing/process_withdrawals.rs

@@ -82,7 +82,7 @@ pub mod capella {
     ) -> Result<(), BlockProcessingError> {
         // check if capella enabled because this function will run on the merge block where the fork is technically still Bellatrix
         if state.fork_name_unchecked().capella_enabled() {
-            let (expected_withdrawals, _, partial_withdrawals_count) =
+            let (expected_withdrawals, _, partial_withdrawals_count, _) =
                 get_expected_withdrawals(state, spec)?;
 
             let expected_root = expected_withdrawals.tree_hash_root();
@@ -127,7 +127,7 @@ pub mod gloas {
             return Ok(());
         }
 
-        let (expected_withdrawals, builder_withdrawals_count, partial_withdrawals_count) =
+        let (expected_withdrawals, builder_withdrawals_count, partial_withdrawals_count, _) =
             get_expected_withdrawals(state, spec)?;
 
         for withdrawal in expected_withdrawals.iter() {

testing/ef_tests/src/cases/operations.rs

@@ -28,8 +28,8 @@ use types::{
     Attestation, AttesterSlashing, BeaconBlock, BeaconBlockBody, BeaconBlockBodyBellatrix,
     BeaconBlockBodyCapella, BeaconBlockBodyDeneb, BeaconBlockBodyElectra, BeaconBlockBodyFulu,
     BeaconState, BlindedPayload, ConsolidationRequest, Deposit, DepositRequest, ExecutionPayload,
-    ForkVersionDecode, FullPayload, PayloadAttestation, ProposerSlashing, SignedBlsToExecutionChange,
-    SignedVoluntaryExit, SyncAggregate, WithdrawalRequest,
+    ForkVersionDecode, FullPayload, PayloadAttestation, ProposerSlashing,
+    SignedBlsToExecutionChange, SignedVoluntaryExit, SyncAggregate, WithdrawalRequest,
 };
 
 #[derive(Debug, Clone, Default, Deserialize)]

testing/ef_tests/tests/tests.rs

@@ -378,8 +378,10 @@ mod ssz_static {
             .run();
         SszStaticHandler::<BeaconBlockBodyFulu<MinimalEthSpec>, MinimalEthSpec>::fulu_only().run();
         SszStaticHandler::<BeaconBlockBodyFulu<MainnetEthSpec>, MainnetEthSpec>::fulu_only().run();
-        SszStaticHandler::<BeaconBlockBodyGloas<MinimalEthSpec>, MinimalEthSpec>::gloas_only().run();
-        SszStaticHandler::<BeaconBlockBodyGloas<MainnetEthSpec>, MainnetEthSpec>::gloas_only().run();
+        SszStaticHandler::<BeaconBlockBodyGloas<MinimalEthSpec>, MinimalEthSpec>::gloas_only()
+            .run();
+        SszStaticHandler::<BeaconBlockBodyGloas<MainnetEthSpec>, MainnetEthSpec>::gloas_only()
+            .run();
     }
 
     // Altair and later