refactor attestation filtering

Consensys · Jul 25, 2024 · b9604a8 · b9604a8
1 parent 4883f6a
commit b9604a8
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 35 deletions.
diff --git a/.../main/java/tech/pegasys/teku/beaconrestapi/handlers/v2/beacon/GetBlockAttestationsV2.java b/.../main/java/tech/pegasys/teku/beaconrestapi/handlers/v2/beacon/GetBlockAttestationsV2.java
@@ -39,11 +39,8 @@
 import tech.pegasys.teku.infrastructure.restapi.endpoints.EndpointMetadata;
 import tech.pegasys.teku.infrastructure.restapi.endpoints.RestApiEndpoint;
 import tech.pegasys.teku.infrastructure.restapi.endpoints.RestApiRequest;
-import tech.pegasys.teku.spec.Spec;
-import tech.pegasys.teku.spec.SpecMilestone;
 import tech.pegasys.teku.spec.datastructures.metadata.ObjectAndMetaData;
 import tech.pegasys.teku.spec.datastructures.operations.Attestation;
-import tech.pegasys.teku.spec.datastructures.operations.AttesterSlashing;
 import tech.pegasys.teku.spec.schemas.SchemaDefinitionCache;
 
 public class GetBlockAttestationsV2 extends RestApiEndpoint {
@@ -119,20 +116,6 @@ private static Predicate<List<Attestation>> phase0AttestationsPredicate() {
     return attestations -> attestations.isEmpty() || !attestations.get(0).requiresCommitteeBits();
   }
 
-  private static Predicate<AttesterSlashing> electraAttesterSlashing(final Spec spec) {
-    return attesterSlashing ->
-        spec.atSlot(attesterSlashing.getAttestation1().getData().getSlot())
-            .getMilestone()
-            .isGreaterThanOrEqualTo(SpecMilestone.ELECTRA);
-  }
-
-  private static Predicate<AttesterSlashing> phase0AttesterSlashing(final Spec spec) {
-    return attesterSlashing ->
-        !spec.atSlot(attesterSlashing.getAttestation1().getData().getSlot())
-            .getMilestone()
-            .isGreaterThanOrEqualTo(SpecMilestone.ELECTRA);
-  }
-
   private static Predicate<List<Attestation>> electraAttestationsPredicate() {
     // Only once we are in Electra attestations will have committee bits
     return attestations -> !attestations.isEmpty() && attestations.get(0).requiresCommitteeBits();

diff --git a/...c/main/java/tech/pegasys/teku/statetransition/attestation/AggregatingAttestationPool.java b/...c/main/java/tech/pegasys/teku/statetransition/attestation/AggregatingAttestationPool.java
@@ -276,7 +276,7 @@ public synchronized List<Attestation> getAttestations(
     final UInt64 slot = maybeSlot.orElse(recentChainData.getCurrentSlot().orElse(UInt64.ZERO));
     final SchemaDefinitions schemaDefinitions = spec.atSlot(slot).getSchemaDefinitions();
 
-    final boolean requestRequiresAttestationWithCommitteeBits =
+    final boolean requiresCommitteeBits =
         schemaDefinitions.getAttestationSchema().requiresCommitteeBits();
 
     return dataHashBySlot.descendingMap().entrySet().stream()
@@ -287,11 +287,8 @@ public synchronized List<Attestation> getAttestations(
         .filter(Objects::nonNull)
         .flatMap(
             matchingDataAttestationGroup ->
-                matchingDataAttestationGroup.stream(maybeCommitteeIndex))
+                matchingDataAttestationGroup.stream(maybeCommitteeIndex, requiresCommitteeBits))
         .map(ValidatableAttestation::getAttestation)
-        .filter(
-            attestation ->
-                attestation.requiresCommitteeBits() == requestRequiresAttestationWithCommitteeBits)
         .toList();
   }
 

diff --git a/...main/java/tech/pegasys/teku/statetransition/attestation/MatchingDataAttestationGroup.java b/...main/java/tech/pegasys/teku/statetransition/attestation/MatchingDataAttestationGroup.java
@@ -147,6 +147,22 @@ public Stream<ValidatableAttestation> stream(final Optional<UInt64> committeeInd
     return StreamSupport.stream(spliterator(committeeIndex), false);
   }
 
+  public Stream<ValidatableAttestation> stream(
+      final Optional<UInt64> committeeIndex, final boolean requiresCommitteeBits) {
+    // Group Attestation type mismatch
+    if (requiresCommitteeBits != includedValidators.requiresCommitteeBits()) {
+      return Stream.empty();
+    }
+    // Pre electra attestation and committee index not matching
+    if (committeeIndex.isPresent()
+        && !includedValidators.requiresCommitteeBits()
+        && !attestationData.getIndex().equals(committeeIndex.get())) {
+      return Stream.empty();
+    }
+    // Filter based on the committee index
+    return StreamSupport.stream(spliterator(committeeIndex), false);
+  }
+
   public Spliterator<ValidatableAttestation> spliterator(final Optional<UInt64> committeeIndex) {
     return Spliterators.spliteratorUnknownSize(iterator(committeeIndex), 0);
   }
@@ -270,23 +286,21 @@ public Iterator<ValidatableAttestation> getRemainingAttestations() {
           .iterator();
     }
 
+    /*
+    If we have attestations with committeeBits (Electra) then, if maybeCommitteeIndex is specified, we will consider attestation related to that committee only
+     */
     private boolean maybeFilterOnCommitteeIndex(final ValidatableAttestation candidate) {
-      final Attestation attestation = candidate.getAttestation();
-      final Optional<SszBitvector> maybeCommitteeBits = attestation.getCommitteeBits();
-      if (maybeCommitteeIndex.isEmpty()) {
+      final Optional<SszBitvector> maybeCommitteeBits =
+          candidate.getAttestation().getCommitteeBits();
+      if (maybeCommitteeBits.isEmpty() || maybeCommitteeIndex.isEmpty()) {
         return true;
       }
-      // Pre Electra attestation, filter based on the attestation data index
-      if (maybeCommitteeBits.isEmpty()) {
-        return attestation.getData().getIndex().equals(maybeCommitteeIndex.get());
-      } else {
-        // Post Electra attestation, filter on committee bits
-        final SszBitvector committeeBits = maybeCommitteeBits.get();
-        if (committeeBits.getBitCount() != 1) {
-          return false;
-        }
-        return committeeBits.isSet(maybeCommitteeIndex.get().intValue());
+
+      final SszBitvector committeeBits = maybeCommitteeBits.get();
+      if (committeeBits.getBitCount() != 1) {
+        return false;
       }
+      return committeeBits.isSet(maybeCommitteeIndex.get().intValue());
     }
   }
 }