Updated get_expected_withdrawals and process_withdrawals to use MAX_V…

…ALIDATORS_PER_WITHDRAWALS_SWEEP

ethereum/spec/src/main/java/tech/pegasys/teku/spec/config/builder/CapellaBuilder.java

@@ -67,7 +67,8 @@ public CapellaBuilder maxWithdrawalsPerPayload(final int maxWithdrawalsPerPayloa
     return this;
   }
 
-  public CapellaBuilder maxValidatorsPerWithdrawalsSweep(final int maxValidatorsPerWithdrawalSweep) {
+  public CapellaBuilder maxValidatorsPerWithdrawalsSweep(
+      final int maxValidatorsPerWithdrawalSweep) {
     this.maxValidatorsPerWithdrawalSweep = maxValidatorsPerWithdrawalSweep;
     return this;
   }

ethereum/spec/src/main/java/tech/pegasys/teku/spec/logic/versions/capella/block/BlockProcessorCapella.java

@@ -181,6 +181,7 @@ public void processBlsToExecutionChangesNoValidation(
     }
   }
 
+  // process_withdrawals
   @Override
   public void processWithdrawals(
       final MutableBeaconState genericState, final ExecutionPayloadSummary payloadSummary)
@@ -192,34 +193,57 @@ public void processWithdrawals(
             .getWithdrawalsSchemaRequired()
             .createFromElements(getExpectedWithdrawals(state));
 
-    if (payloadSummary.getOptionalWithdrawalsRoot().isEmpty()
-        || !expectedWithdrawals
-            .hashTreeRoot()
-            .equals(payloadSummary.getOptionalWithdrawalsRoot().get())) {
-      throw new BlockProcessingException(
-          "Expected "
-              + expectedWithdrawals.hashTreeRoot()
-              + " withdrawals root, but withdrawals root was "
-              + (payloadSummary.getOptionalWithdrawalsRoot().isPresent()
-                  ? payloadSummary.getOptionalWithdrawalsRoot().get()
-                  : "MISSING"));
-    }
+    assertWithdrawalsInExecutionPayloadMatchExpected(payloadSummary, expectedWithdrawals);
+
     for (int i = 0; i < expectedWithdrawals.size(); i++) {
       final Withdrawal withdrawal = expectedWithdrawals.get(i);
       beaconStateMutators.decreaseBalance(
           state, withdrawal.getValidatorIndex().intValue(), withdrawal.getAmount());
     }
 
-    if (expectedWithdrawals.size() > 0) {
+    if (expectedWithdrawals.size() != 0) {
+      final int validatorCount = genericState.getValidators().size();
+      final int maxWithdrawalsPerPayload = specConfigCapella.getMaxWithdrawalsPerPayload();
+      final int maxValidatorsPerWithdrawalsSweep =
+          specConfigCapella.getMaxValidatorsPerWithdrawalSweep();
+
       final Withdrawal latestWithdrawal = expectedWithdrawals.get(expectedWithdrawals.size() - 1);
-      final int nextWithdrawalValidatorIndex =
-          incrementValidatorIndex(
-              latestWithdrawal.getValidatorIndex().intValue(), genericState.getValidators().size());
       state.setNextWithdrawalIndex(latestWithdrawal.getIndex().increment());
-      state.setNextWithdrawalValidatorIndex(UInt64.valueOf(nextWithdrawalValidatorIndex));
+
+      final int nextWithdrawalValidatorIndex;
+      if (expectedWithdrawals.size() == maxWithdrawalsPerPayload) {
+        // Update the next validator index to start the next withdrawal sweep
+        nextWithdrawalValidatorIndex = latestWithdrawal.getValidatorIndex().intValue() + 1;
+      } else {
+        // Advance sweep by the max length of the sweep if there was not a full set of withdrawals
+        nextWithdrawalValidatorIndex =
+            state.getNextWithdrawalValidatorIndex().intValue() + maxValidatorsPerWithdrawalsSweep;
+      }
+      state.setNextWithdrawalValidatorIndex(
+          UInt64.valueOf(nextWithdrawalValidatorIndex % validatorCount));
+    }
+  }
+
+  private static void assertWithdrawalsInExecutionPayloadMatchExpected(
+      final ExecutionPayloadSummary payloadSummary, final SszList<Withdrawal> expectedWithdrawals)
+      throws BlockProcessingException {
+    // the spec does a element-to-element comparison but Teku is comparing the hash of the tree
+    if (payloadSummary.getOptionalWithdrawalsRoot().isEmpty()
+        || !expectedWithdrawals
+            .hashTreeRoot()
+            .equals(payloadSummary.getOptionalWithdrawalsRoot().get())) {
+      final String msg =
+          String.format(
+              "Withdrawals in execution payload are different from expected (expected withdrawals root is %s but was "
+                  + "%s)",
+              expectedWithdrawals.hashTreeRoot(),
+              payloadSummary
+                  .getOptionalWithdrawalsRoot()
+                  .map(Bytes::toHexString)
+                  .orElse("MISSING"));
+      throw new BlockProcessingException(msg);
     }
   }
-  // process_withdrawals
 
   @Override
   public Optional<List<Withdrawal>> getExpectedWithdrawals(final BeaconState preState) {
@@ -235,12 +259,14 @@ private List<Withdrawal> getExpectedWithdrawals(final BeaconStateCapella preStat
     final SszUInt64List balances = preState.getBalances();
     final int validatorCount = validators.size();
     final int maxWithdrawalsPerPayload = specConfigCapella.getMaxWithdrawalsPerPayload();
+    final int maxValidatorsPerWithdrawalsSweep =
+        specConfigCapella.getMaxValidatorsPerWithdrawalSweep();
+    final int bound = Math.min(validatorCount, maxValidatorsPerWithdrawalsSweep);
+
     UInt64 withdrawalIndex = preState.getNextWithdrawalIndex();
     int validatorIndex = preState.getNextWithdrawalValidatorIndex().intValue();
-    for (int i = 0;
-        i < validatorCount && expectedWithdrawals.size() < maxWithdrawalsPerPayload;
-        i++) {
 
+    for (int i = 0; i < bound; i++) {
       final Validator validator = validators.get(validatorIndex);
       if (predicates.hasEth1WithdrawalCredential(validator)) {
         final UInt64 balance = balances.get(validatorIndex).get();
@@ -264,9 +290,13 @@ private List<Withdrawal> getExpectedWithdrawals(final BeaconStateCapella preStat
                   balance.minus(specConfig.getMaxEffectiveBalance())));
           withdrawalIndex = withdrawalIndex.increment();
         }
+
+        if (expectedWithdrawals.size() == maxWithdrawalsPerPayload) {
+          break;
+        }
       }
 
-      validatorIndex = incrementValidatorIndex(validatorIndex, validatorCount);
+      validatorIndex = (validatorIndex + 1) % validatorCount;
     }
 
     return expectedWithdrawals;
@@ -278,11 +308,6 @@ public static Bytes32 getWithdrawalAddressFromEth1Address(final Bytes20 toExecut
         Bytes.concatenate(ETH1_WITHDRAWAL_KEY_PREFIX, toExecutionAddress.getWrappedBytes()));
   }
 
-  @VisibleForTesting
-  static int incrementValidatorIndex(final int validatorIndex, final int validatorCount) {
-    return (validatorIndex + 1) % validatorCount;
-  }
-
   @VisibleForTesting
   BlockValidationResult verifyBlsToExecutionChangesPreProcessing(
       final BeaconState genericState,

ethereum/spec/src/test/java/tech/pegasys/teku/spec/logic/versions/capella/block/BlockProcessorCapellaTest.java

@@ -16,6 +16,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
 import org.apache.tuweni.bytes.Bytes32;
@@ -27,10 +28,12 @@
 import tech.pegasys.teku.infrastructure.unsigned.UInt64;
 import tech.pegasys.teku.spec.Spec;
 import tech.pegasys.teku.spec.TestSpecFactory;
+import tech.pegasys.teku.spec.config.SpecConfig;
 import tech.pegasys.teku.spec.datastructures.blocks.SignedBeaconBlock;
 import tech.pegasys.teku.spec.datastructures.execution.versions.capella.Withdrawal;
 import tech.pegasys.teku.spec.datastructures.operations.BlsToExecutionChange;
 import tech.pegasys.teku.spec.datastructures.operations.SignedBlsToExecutionChange;
+import tech.pegasys.teku.spec.datastructures.state.Fork;
 import tech.pegasys.teku.spec.datastructures.state.Validator;
 import tech.pegasys.teku.spec.datastructures.state.beaconstate.BeaconState;
 import tech.pegasys.teku.spec.generator.ChainBuilder;
@@ -78,6 +81,45 @@ void shouldCreateExpectedWithdrawalAddress() {
     assertThat(bytes32.toHexString()).startsWith("0x010000000000000000000000");
   }
 
+  @Test
+  void shouldNotSweepMoreValidatorsThanLimit() {
+    final int maxValidatorsPerWithdrawalSweep =
+        spec.getGenesisSpecConfig().toVersionCapella().get().getMaxValidatorsPerWithdrawalSweep();
+    final DataStructureUtil data = new DataStructureUtil(TestSpecFactory.createMinimalCapella());
+
+    final Validator withdrawableValidator1 =
+        makeValidator(data.randomPublicKey(), data.randomEth1WithdrawalCredentials());
+    final UInt64 withdrawableValidator1Balance =
+        spec.getGenesisSpecConfig().getMaxEffectiveBalance().plus(1024000);
+
+    final Validator withdrawableValidator2 =
+        makeValidator(data.randomPublicKey(), data.randomEth1WithdrawalCredentials());
+    final UInt64 withdrawableValidator2Balance =
+        spec.getGenesisSpecConfig().getMaxEffectiveBalance().plus(1024000);
+
+    final List<Validator> validators = new ArrayList<>();
+    final List<UInt64> balances = new ArrayList<>();
+    for (int i = 0; i < maxValidatorsPerWithdrawalSweep - 1; i++) {
+      validators.add(dataStructureUtil.randomValidator());
+      balances.add(spec.getGenesisSpecConfig().getMaxEffectiveBalance());
+    }
+
+    // Withdrawable validator within limit
+    validators.add(withdrawableValidator1);
+    balances.add(withdrawableValidator1Balance);
+
+    // Withdrawable validator outside limit
+    validators.add(withdrawableValidator2);
+    balances.add(withdrawableValidator2Balance);
+
+    final BeaconState preState = createBeaconStateWithValidatorsAndBalances(validators, balances);
+
+    final Optional<List<Withdrawal>> withdrawals =
+        spec.getBlockProcessor(preState.getSlot()).getExpectedWithdrawals(preState);
+    assertThat(withdrawals).isPresent();
+    assertThat(withdrawals.get()).hasSize(1);
+  }
+
   @Test
   void shouldFindPartialWithdrawals() {
     final DataStructureUtil data = new DataStructureUtil(TestSpecFactory.createMinimalCapella());
@@ -107,13 +149,6 @@ void shouldFindFullWithdrawals() {
     assertThat(withdrawals.get().get(0).getAmount()).isEqualTo(balance);
   }
 
-  @Test
-  void shouldIncrementIndexCorrectly() {
-    assertThat(BlockProcessorCapella.incrementValidatorIndex(9, 10)).isEqualTo(0);
-    assertThat(BlockProcessorCapella.incrementValidatorIndex(8, 10)).isEqualTo(9);
-    assertThat(BlockProcessorCapella.incrementValidatorIndex(0, 10)).isEqualTo(1);
-  }
-
   @Test
   public void shouldRejectBlockWithMoreThanOneBlsExecutionChangesForSameValidator() {
     final int validatorIndex = 0;
@@ -171,4 +206,23 @@ private SszList<SignedBlsToExecutionChange> createBlsToExecutionChangeList(
         .getBlsToExecutionChangesSchema()
         .of(signedBlsToExecutionChanges);
   }
+
+  protected BeaconState createBeaconStateWithValidatorsAndBalances(
+      final List<Validator> validators, final List<UInt64> balances) {
+    return spec.getGenesisSpec()
+        .getSchemaDefinitions()
+        .getBeaconStateSchema()
+        .createEmpty()
+        .updated(
+            beaconState -> {
+              beaconState.setSlot(dataStructureUtil.randomUInt64());
+              beaconState.setFork(
+                  new Fork(
+                      spec.getGenesisSpecConfig().getGenesisForkVersion(),
+                      spec.getGenesisSpecConfig().getGenesisForkVersion(),
+                      SpecConfig.GENESIS_EPOCH));
+              beaconState.getValidators().appendAll(validators);
+              beaconState.getBalances().appendAllElements(balances);
+            });
+  }
 }