refactor attestations filtering

Consensys · Jul 24, 2024 · 89c0843 · 89c0843
1 parent c797ec0
commit 89c0843
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 9 deletions.
diff --git a/ethereum/spec/src/testFixtures/java/tech/pegasys/teku/spec/util/DataStructureUtil.java b/ethereum/spec/src/testFixtures/java/tech/pegasys/teku/spec/util/DataStructureUtil.java
@@ -790,7 +790,11 @@ public AttestationData randomAttestationData() {
 
   public AttestationData randomAttestationData(final UInt64 slot) {
     return new AttestationData(
-        slot, randomUInt64(), randomBytes32(), randomCheckpoint(), randomCheckpoint());
+        slot,
+        randomUInt64(getMaxCommitteesPerSlot()),
+        randomBytes32(),
+        randomCheckpoint(),
+        randomCheckpoint());
   }
 
   public AttestationData randomAttestationData(final UInt64 slot, final Bytes32 blockRoot) {

diff --git a/...c/main/java/tech/pegasys/teku/statetransition/attestation/AggregatingAttestationPool.java b/...c/main/java/tech/pegasys/teku/statetransition/attestation/AggregatingAttestationPool.java
@@ -25,7 +25,9 @@
 import java.util.Set;
 import java.util.TreeMap;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.Function;
 import java.util.function.Predicate;
+import java.util.stream.Stream;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.tuweni.bytes.Bytes;
@@ -273,32 +275,50 @@ public synchronized List<Attestation> getAttestations(
     final Predicate<Map.Entry<UInt64, Set<Bytes>>> filterForSlot =
         (entry) -> maybeSlot.map(slot -> entry.getKey().equals(slot)).orElse(true);
 
-    final Predicate<MatchingDataAttestationGroup> filterForCommitteeIndex =
-        (group) ->
-            maybeCommitteeIndex
-                .map(index -> group.getAttestationData().getIndex().equals(index))
-                .orElse(true);
     final UInt64 slot = maybeSlot.orElse(recentChainData.getCurrentSlot().orElse(UInt64.ZERO));
     final SchemaDefinitions schemaDefinitions = spec.atSlot(slot).getSchemaDefinitions();
 
     final boolean requestRequiresAttestationWithCommitteeBits =
         schemaDefinitions.getAttestationSchema().requiresCommitteeBits();
 
+    // Committee index filter predicate is used for pre Electra attestations only (the filter is
+    // applied to the index at the attestation data level)
+    // Post Electra the index is always set to 0 at the attestation data level (the filter
+    // is rather applied to the committee bits)
+    final Predicate<MatchingDataAttestationGroup> filterForCommitteeIndex =
+        (group) ->
+            requestRequiresAttestationWithCommitteeBits
+                || maybeCommitteeIndex
+                    .map(index -> group.getAttestationData().getIndex().equals(index))
+                    .orElse(true);
+
     return dataHashBySlot.descendingMap().entrySet().stream()
         .filter(filterForSlot)
         .map(Map.Entry::getValue)
         .flatMap(Collection::stream)
         .map(attestationGroupByDataHash::get)
         .filter(Objects::nonNull)
         .filter(filterForCommitteeIndex)
-        .flatMap(MatchingDataAttestationGroup::stream)
+        .flatMap(
+            streamMatchingAttestations(
+                maybeCommitteeIndex, requestRequiresAttestationWithCommitteeBits))
         .map(ValidatableAttestation::getAttestation)
         .filter(
             attestation ->
                 attestation.requiresCommitteeBits() == requestRequiresAttestationWithCommitteeBits)
         .toList();
   }
 
+  private Function<MatchingDataAttestationGroup, Stream<ValidatableAttestation>>
+      streamMatchingAttestations(
+          final Optional<UInt64> maybeCommitteeIndex,
+          final boolean requestRequiresAttestationWithCommitteeBits) {
+    return matchingDataAttestationGroup ->
+        requestRequiresAttestationWithCommitteeBits
+            ? matchingDataAttestationGroup.stream(maybeCommitteeIndex)
+            : matchingDataAttestationGroup.stream();
+  }
+
   private boolean isValid(
       final BeaconState stateAtBlockSlot, final AttestationData attestationData) {
     return spec.validateAttestation(stateAtBlockSlot, attestationData).isEmpty();

diff --git a/...st/java/tech/pegasys/teku/statetransition/attestation/AggregatingAttestationPoolTest.java b/...st/java/tech/pegasys/teku/statetransition/attestation/AggregatingAttestationPoolTest.java
@@ -522,7 +522,10 @@ public void getAttestations_shouldReturnAllAttestations() {
   }
 
   @TestTemplate
-  public void getAttestations_shouldReturnAttestationsForGivenCommitteeIndexOnly() {
+  public void getAttestations_shouldReturnAttestationsForGivenCommitteeIndexOnly_PreElectra() {
+    assumeThat(specMilestone).isLessThan(ELECTRA);
+    // Pre Electra the committee index filter is applied to the index set at the attestation data
+    // level
     final AttestationData attestationData1 = dataStructureUtil.randomAttestationData();
     final AttestationData attestationData2 =
         new AttestationData(
@@ -531,14 +534,28 @@ public void getAttestations_shouldReturnAttestationsForGivenCommitteeIndexOnly()
             attestationData1.getBeaconBlockRoot(),
             attestationData1.getSource(),
             attestationData1.getTarget());
-    Attestation attestation1 = addAttestationFromValidators(attestationData1, 1, 2, 3);
+    final Attestation attestation1 = addAttestationFromValidators(attestationData1, 1, 2, 3);
     addAttestationFromValidators(attestationData2, 4, 5, 6);
     assertThat(
             aggregatingPool.getAttestations(
                 Optional.empty(), Optional.of(attestationData1.getIndex())))
         .containsExactly(attestation1);
   }
 
+  @TestTemplate
+  public void getAttestations_shouldReturnAttestationsForGivenCommitteeIndexOnly_PostElectra() {
+    assumeThat(specMilestone).isGreaterThanOrEqualTo(ELECTRA);
+    // Post Electra the committee index filter is applied to the committee bits
+    final AttestationData attestationData1 = dataStructureUtil.randomAttestationData();
+    final AttestationData attestationData2 = dataStructureUtil.randomAttestationData();
+    final Attestation attestation1 = addAttestationFromValidators(attestationData1, 1, 2, 3);
+    final Optional<UInt64> committeeIndexFilter = committeeIndex;
+    committeeIndex = Optional.of(committeeIndex.get().plus(1));
+    addAttestationFromValidators(attestationData2, 4, 5, 6);
+    assertThat(aggregatingPool.getAttestations(Optional.empty(), committeeIndexFilter))
+        .containsExactly(attestation1);
+  }
+
   @TestTemplate
   public void getAttestations_shouldReturnAttestationsForGivenSlotOnly() {
     final AttestationData attestationData1 = dataStructureUtil.randomAttestationData();