refactor committee index filtering

Consensys · Jul 25, 2024 · ae2b2d7 · ae2b2d7
1 parent f375c94
commit ae2b2d7
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 37 deletions.
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,9 +25,7 @@
 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;
@@ -281,44 +279,22 @@ public synchronized List<Attestation> getAttestations(
     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(
-            streamMatchingAttestations(
-                maybeCommitteeIndex, requestRequiresAttestationWithCommitteeBits))
+            matchingDataAttestationGroup ->
+                matchingDataAttestationGroup.stream(maybeCommitteeIndex))
         .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/...main/java/tech/pegasys/teku/statetransition/attestation/MatchingDataAttestationGroup.java b/...main/java/tech/pegasys/teku/statetransition/attestation/MatchingDataAttestationGroup.java
@@ -270,21 +270,23 @@ 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 Optional<SszBitvector> maybeCommitteeBits =
-          candidate.getAttestation().getCommitteeBits();
-      if (maybeCommitteeBits.isEmpty() || maybeCommitteeIndex.isEmpty()) {
+      final Attestation attestation = candidate.getAttestation();
+      final Optional<SszBitvector> maybeCommitteeBits = attestation.getCommitteeBits();
+      if (maybeCommitteeIndex.isEmpty()) {
         return true;
       }
-
-      final SszBitvector committeeBits = maybeCommitteeBits.get();
-      if (committeeBits.getBitCount() != 1) {
-        return false;
+      // 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());
       }
-      return committeeBits.isSet(maybeCommitteeIndex.get().intValue());
     }
   }
 }