Fix partial withdrawals count

presets/minimal/electra.yaml

@@ -42,4 +42,4 @@ MAX_WITHDRAWAL_REQUESTS_PER_PAYLOAD: 2
 # Withdrawals processing
 # ---------------------------------------------------------------
 # 2**0 ( = 1) pending withdrawals
-MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP: 1
+MAX_PENDING_PARTIALS_PER_WITHDRAWALS_SWEEP: 2

specs/electra/beacon-chain.md

@@ -1007,6 +1007,7 @@ def get_expected_withdrawals(state: BeaconState) -> Tuple[Sequence[Withdrawal],
     withdrawal_index = state.next_withdrawal_index
     validator_index = state.next_withdrawal_validator_index
     withdrawals: List[Withdrawal] = []
+    partial_withdrawals_count = 0
 
     # [New in Electra:EIP7251] Consume pending partial withdrawals
     for withdrawal in state.pending_partial_withdrawals:
@@ -1026,7 +1027,7 @@ def get_expected_withdrawals(state: BeaconState) -> Tuple[Sequence[Withdrawal],
             ))
             withdrawal_index += WithdrawalIndex(1)
 
-    partial_withdrawals_count = len(withdrawals)
+        partial_withdrawals_count += 1
 
     # Sweep for remaining.
     bound = min(len(state.validators), MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP)

tests/core/pyspec/eth2spec/test/capella/block_processing/test_process_withdrawals.py

@@ -24,80 +24,10 @@
     set_eth1_withdrawal_credential_with_balance,
     set_validator_fully_withdrawable,
     set_validator_partially_withdrawable,
+    run_withdrawals_processing,
 )
 
 
-def verify_post_state(state, spec, expected_withdrawals,
-                      fully_withdrawable_indices, partial_withdrawals_indices):
-    # Consider verifying also the condition when no withdrawals are expected.
-    if len(expected_withdrawals) == 0:
-        return
-
-    expected_withdrawals_validator_indices = [withdrawal.validator_index for withdrawal in expected_withdrawals]
-    assert state.next_withdrawal_index == expected_withdrawals[-1].index + 1
-
-    if len(expected_withdrawals) == spec.MAX_WITHDRAWALS_PER_PAYLOAD:
-        # NOTE: ideally we would also check in the case with
-        # fewer than maximum withdrawals but that requires the pre-state info
-        next_withdrawal_validator_index = (expected_withdrawals_validator_indices[-1] + 1) % len(state.validators)
-        assert state.next_withdrawal_validator_index == next_withdrawal_validator_index
-
-    for index in fully_withdrawable_indices:
-        if index in expected_withdrawals_validator_indices:
-            assert state.balances[index] == 0
-        else:
-            assert state.balances[index] > 0
-    for index in partial_withdrawals_indices:
-        if index in expected_withdrawals_validator_indices:
-            assert state.balances[index] == spec.MAX_EFFECTIVE_BALANCE
-        else:
-            assert state.balances[index] > spec.MAX_EFFECTIVE_BALANCE
-
-
-def run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=None,
-                               fully_withdrawable_indices=None, partial_withdrawals_indices=None, valid=True):
-    """
-    Run ``process_withdrawals``, yielding:
-      - pre-state ('pre')
-      - execution payload ('execution_payload')
-      - post-state ('post').
-    If ``valid == False``, run expecting ``AssertionError``
-    """
-    expected_withdrawals = get_expected_withdrawals(spec, state)
-    assert len(expected_withdrawals) <= spec.MAX_WITHDRAWALS_PER_PAYLOAD
-    if num_expected_withdrawals is not None:
-        assert len(expected_withdrawals) == num_expected_withdrawals
-
-    pre_state = state.copy()
-    yield 'pre', state
-    yield 'execution_payload', execution_payload
-
-    if not valid:
-        expect_assertion_error(lambda: spec.process_withdrawals(state, execution_payload))
-        yield 'post', None
-        return
-
-    spec.process_withdrawals(state, execution_payload)
-
-    yield 'post', state
-
-    if len(expected_withdrawals) == 0:
-        next_withdrawal_validator_index = (
-            pre_state.next_withdrawal_validator_index + spec.MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP
-        )
-        assert state.next_withdrawal_validator_index == next_withdrawal_validator_index % len(state.validators)
-    elif len(expected_withdrawals) <= spec.MAX_WITHDRAWALS_PER_PAYLOAD:
-        bound = min(spec.MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP, spec.MAX_WITHDRAWALS_PER_PAYLOAD)
-        assert len(get_expected_withdrawals(spec, state)) <= bound
-    elif len(expected_withdrawals) > spec.MAX_WITHDRAWALS_PER_PAYLOAD:
-        raise ValueError('len(expected_withdrawals) should not be greater than MAX_WITHDRAWALS_PER_PAYLOAD')
-
-    if fully_withdrawable_indices is not None or partial_withdrawals_indices is not None:
-        verify_post_state(state, spec, expected_withdrawals, fully_withdrawable_indices, partial_withdrawals_indices)
-
-    return expected_withdrawals
-
-
 @with_capella_and_later
 @spec_state_test
 def test_success_zero_expected_withdrawals(spec, state):

tests/core/pyspec/eth2spec/test/electra/block_processing/test_process_withdrawals.py

@@ -0,0 +1,121 @@
+import random
+
+from eth2spec.test.context import (
+    spec_state_test,
+    with_presets,
+    with_electra_and_later,
+)
+from eth2spec.test.helpers.constants import MAINNET, MINIMAL
+from eth2spec.test.helpers.execution_payload import (
+    build_empty_execution_payload,
+    compute_el_block_hash,
+)
+from eth2spec.test.helpers.random import (
+    randomize_state,
+)
+from eth2spec.test.helpers.state import (
+    next_epoch,
+    next_slot,
+)
+from eth2spec.test.helpers.withdrawals import (
+    get_expected_withdrawals,
+    prepare_expected_withdrawals,
+    set_eth1_withdrawal_credential_with_balance,
+    set_validator_fully_withdrawable,
+    set_validator_partially_withdrawable,
+    prepare_expected_withdrawals_compounding,
+    run_withdrawals_processing,
+    set_compounding_withdrawal_credential,
+    prepare_pending_withdrawal,
+)
+
+
+@with_electra_and_later
+@spec_state_test
+def test_success_mixed_fully_and_partial_withdrawable_compounding(spec, state):
+    num_full_withdrawals = spec.MAX_WITHDRAWALS_PER_PAYLOAD // 2
+    num_partial_withdrawals = spec.MAX_WITHDRAWALS_PER_PAYLOAD - num_full_withdrawals
+    fully_withdrawable_indices, partial_withdrawals_indices = prepare_expected_withdrawals_compounding(
+        spec, state,
+        rng=random.Random(42),
+        num_full_withdrawals=num_full_withdrawals,
+        num_partial_withdrawals_sweep=num_partial_withdrawals,
+    )
+
+    next_slot(spec, state)
+    execution_payload = build_empty_execution_payload(spec, state)
+
+    yield from run_withdrawals_processing(
+        spec, state, execution_payload,
+        fully_withdrawable_indices=fully_withdrawable_indices,
+        partial_withdrawals_indices=partial_withdrawals_indices)
+
+
+@with_electra_and_later
+@spec_state_test
+def test_success_no_max_effective_balance_compounding(spec, state):
+    validator_index = len(state.validators) // 2
+    # To be partially withdrawable, the validator's effective balance must be maxed out
+    set_compounding_withdrawal_credential(spec, state, validator_index)
+    validator = state.validators[validator_index]
+    validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA - spec.EFFECTIVE_BALANCE_INCREMENT
+    state.balances[validator_index] = validator.effective_balance
+
+    assert not spec.is_partially_withdrawable_validator(validator, state.balances[validator_index])
+
+    execution_payload = build_empty_execution_payload(spec, state)
+
+    yield from run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=0)
+
+
+@with_electra_and_later
+@spec_state_test
+def test_success_no_excess_balance_compounding(spec, state):
+    validator_index = len(state.validators) // 2
+    # To be partially withdrawable, the validator needs an excess balance
+    set_compounding_withdrawal_credential(spec, state, validator_index)
+    validator = state.validators[validator_index]
+    validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+    state.balances[validator_index] = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+
+    assert not spec.is_partially_withdrawable_validator(validator, state.balances[validator_index])
+
+    execution_payload = build_empty_execution_payload(spec, state)
+
+    yield from run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=0)
+
+
+@with_electra_and_later
+@spec_state_test
+def test_success_excess_balance_but_no_max_effective_balance_compounding(spec, state):
+    validator_index = len(state.validators) // 2
+    # To be partially withdrawable, the validator needs both a maxed out effective balance and an excess balance
+    set_compounding_withdrawal_credential(spec, state, validator_index)
+    validator = state.validators[validator_index]
+    validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA - spec.EFFECTIVE_BALANCE_INCREMENT
+    state.balances[validator_index] = spec.MAX_EFFECTIVE_BALANCE_ELECTRA + spec.EFFECTIVE_BALANCE_INCREMENT
+
+    assert not spec.is_partially_withdrawable_validator(validator, state.balances[validator_index])
+
+    execution_payload = build_empty_execution_payload(spec, state)
+
+    yield from run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=0)
+
+
+@with_electra_and_later
+@spec_state_test
+def test_pending_withdrawals_one_skipped_one_effective(spec, state):
+    index_0 = 3
+    index_1 = 5
+
+    withdrawal_0 = prepare_pending_withdrawal(spec, state, index_0)
+    withdrawal_1 = prepare_pending_withdrawal(spec, state, index_1)
+
+    # If validator doesn't have an excess balance pending withdrawal is skipped
+    state.balances[index_0] = spec.MIN_ACTIVATION_BALANCE
+
+    execution_payload = build_empty_execution_payload(spec, state)
+
+    yield from run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=1)
+
+    assert state.pending_partial_withdrawals == []

tests/core/pyspec/eth2spec/test/helpers/withdrawals.py

@@ -41,21 +41,33 @@ def set_eth1_withdrawal_credential_with_balance(spec, state, index, balance=None
 
 
 def set_validator_partially_withdrawable(spec, state, index, excess_balance=1000000000):
-    set_eth1_withdrawal_credential_with_balance(spec, state, index, spec.MAX_EFFECTIVE_BALANCE + excess_balance)
     validator = state.validators[index]
+    if is_post_electra(spec) and spec.has_compounding_withdrawal_credential(validator):
+        validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+        state.balances[index] = validator.effective_balance + excess_balance
+    else:
+        set_eth1_withdrawal_credential_with_balance(spec, state, index, spec.MAX_EFFECTIVE_BALANCE + excess_balance)
 
-    assert spec.is_partially_withdrawable_validator(validator, state.balances[index])
+    assert spec.is_partially_withdrawable_validator(state.validators[index], state.balances[index])
 
 
-def prepare_expected_withdrawals(spec, state, rng,
-                                 num_full_withdrawals=0, num_partial_withdrawals=0):
+def sample_withdrawal_indices(spec, state, rng, num_full_withdrawals, num_partial_withdrawals):
     bound = min(len(state.validators), spec.MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP)
     assert num_full_withdrawals + num_partial_withdrawals <= bound
     eligible_validator_indices = list(range(bound))
     sampled_indices = rng.sample(eligible_validator_indices, num_full_withdrawals + num_partial_withdrawals)
     fully_withdrawable_indices = rng.sample(sampled_indices, num_full_withdrawals)
     partial_withdrawals_indices = list(set(sampled_indices).difference(set(fully_withdrawable_indices)))
 
+    return fully_withdrawable_indices, partial_withdrawals_indices
+
+
+def prepare_expected_withdrawals(spec, state, rng,
+                                 num_full_withdrawals=0, num_partial_withdrawals=0):
+    fully_withdrawable_indices, partial_withdrawals_indices = sample_withdrawal_indices(
+        spec, state, rng, num_full_withdrawals, num_partial_withdrawals
+    )
+
     for index in fully_withdrawable_indices:
         set_validator_fully_withdrawable(spec, state, index)
     for index in partial_withdrawals_indices:
@@ -70,3 +82,130 @@ def set_compounding_withdrawal_credential(spec, state, index, address=None):
 
     validator = state.validators[index]
     validator.withdrawal_credentials = spec.COMPOUNDING_WITHDRAWAL_PREFIX + b'\x00' * 11 + address
+
+
+def prepare_expected_withdrawals_compounding(spec, state, rng,
+                                             num_full_withdrawals=0,
+                                             num_partial_withdrawals_sweep=0,
+                                             excess_balance=1000000000):
+    assert is_post_electra(spec)
+
+    fully_withdrawable_indices, partial_withdrawals_sweep_indices = sample_withdrawal_indices(
+        spec, state, rng, num_full_withdrawals, num_partial_withdrawals_sweep
+    )
+
+    for index in fully_withdrawable_indices + partial_withdrawals_sweep_indices:
+        validator = state.validators[index]
+        validator.effective_balance = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+        state.balances[index] = spec.MAX_EFFECTIVE_BALANCE_ELECTRA
+        address = state.validators[index].withdrawal_credentials[12:]
+        set_compounding_withdrawal_credential(spec, state, index, address=address)
+
+    for index in fully_withdrawable_indices:
+        set_validator_fully_withdrawable(spec, state, index)
+    for index in partial_withdrawals_sweep_indices:
+        set_validator_partially_withdrawable(spec, state, index)
+
+    return fully_withdrawable_indices, partial_withdrawals_sweep_indices
+
+
+def prepare_pending_withdrawal(spec, state, validator_index,
+                               effective_balance=32_000_000_000, amount=1_000_000_000):
+    assert is_post_electra(spec)
+
+    set_compounding_withdrawal_credential(spec, state, validator_index)
+    state.validators[validator_index].effective_balance = effective_balance
+    state.balances[validator_index] = effective_balance + amount
+
+    withdrawal = spec.PendingPartialWithdrawal(
+        index=validator_index,
+        amount=amount,
+        withdrawable_epoch=spec.get_current_epoch(state),
+    )
+    state.pending_partial_withdrawals.append(withdrawal)
+
+    return withdrawal
+
+#
+# Run processing
+#
+
+def verify_post_state(state, spec, expected_withdrawals,
+                      fully_withdrawable_indices, partial_withdrawals_indices):
+    # Consider verifying also the condition when no withdrawals are expected.
+    if len(expected_withdrawals) == 0:
+        return
+
+    expected_withdrawals_validator_indices = [withdrawal.validator_index for withdrawal in expected_withdrawals]
+    assert state.next_withdrawal_index == expected_withdrawals[-1].index + 1
+
+    if len(expected_withdrawals) == spec.MAX_WITHDRAWALS_PER_PAYLOAD:
+        # NOTE: ideally we would also check in the case with
+        # fewer than maximum withdrawals but that requires the pre-state info
+        next_withdrawal_validator_index = (expected_withdrawals_validator_indices[-1] + 1) % len(state.validators)
+        assert state.next_withdrawal_validator_index == next_withdrawal_validator_index
+
+    for index in fully_withdrawable_indices:
+        if index in expected_withdrawals_validator_indices:
+            assert state.balances[index] == 0
+        else:
+            assert state.balances[index] > 0
+    for index in partial_withdrawals_indices:
+        if is_post_electra(spec):
+            max_effective_balance = spec.get_max_effective_balance(state.validators[index])
+        else:
+            max_effective_balance = spec.MAX_EFFECTIVE_BALANCE
+
+        if index in expected_withdrawals_validator_indices:
+            assert state.balances[index] == max_effective_balance
+        else:
+            assert state.balances[index] > max_effective_balance
+
+
+def run_withdrawals_processing(spec, state, execution_payload, num_expected_withdrawals=None,
+                               fully_withdrawable_indices=None, partial_withdrawals_indices=None, valid=True):
+    """
+    Run ``process_withdrawals``, yielding:
+      - pre-state ('pre')
+      - execution payload ('execution_payload')
+      - post-state ('post').
+    If ``valid == False``, run expecting ``AssertionError``
+    """
+    expected_withdrawals = get_expected_withdrawals(spec, state)
+    assert len(expected_withdrawals) <= spec.MAX_WITHDRAWALS_PER_PAYLOAD
+    if num_expected_withdrawals is not None:
+        assert len(expected_withdrawals) == num_expected_withdrawals
+
+    pre_state = state.copy()
+    yield 'pre', state
+    yield 'execution_payload', execution_payload
+
+    if not valid:
+        try:
+            spec.process_withdrawals(state, execution_payload)
+            raise AssertionError('expected an assertion error, but got none.')
+        except AssertionError:
+            pass
+
+        yield 'post', None
+        return
+
+    spec.process_withdrawals(state, execution_payload)
+
+    yield 'post', state
+
+    if len(expected_withdrawals) == 0:
+        next_withdrawal_validator_index = (
+            pre_state.next_withdrawal_validator_index + spec.MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP
+        )
+        assert state.next_withdrawal_validator_index == next_withdrawal_validator_index % len(state.validators)
+    elif len(expected_withdrawals) <= spec.MAX_WITHDRAWALS_PER_PAYLOAD:
+        bound = min(spec.MAX_VALIDATORS_PER_WITHDRAWALS_SWEEP, spec.MAX_WITHDRAWALS_PER_PAYLOAD)
+        assert len(get_expected_withdrawals(spec, state)) <= bound
+    elif len(expected_withdrawals) > spec.MAX_WITHDRAWALS_PER_PAYLOAD:
+        raise ValueError('len(expected_withdrawals) should not be greater than MAX_WITHDRAWALS_PER_PAYLOAD')
+
+    if fully_withdrawable_indices is not None or partial_withdrawals_indices is not None:
+        verify_post_state(state, spec, expected_withdrawals, fully_withdrawable_indices, partial_withdrawals_indices)
+
+    return expected_withdrawals