From e434ca3aef36d39a895f21cb89357c20a2335e8c Mon Sep 17 00:00:00 2001
From: Paul Harris <paul.harris@consensys.net>
Date: Tue, 13 Jun 2023 10:59:56 +1000
Subject: [PATCH] Moved the chain heads iteration below unknown in
 AttestationStateSelector (#7251)

Signed-off-by: Paul Harris <paul.harris@consensys.net>
---
 .../validation/AttestationStateSelector.java  | 64 ++++++++++---------
 1 file changed, 33 insertions(+), 31 deletions(-)

diff --git a/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/validation/AttestationStateSelector.java b/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/validation/AttestationStateSelector.java
index 1e0c0749e8b..2b15faef639 100644
--- a/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/validation/AttestationStateSelector.java
+++ b/ethereum/statetransition/src/main/java/tech/pegasys/teku/statetransition/validation/AttestationStateSelector.java
@@ -74,33 +74,13 @@ public SafeFuture<Optional<BeaconState>> getStateToValidate(
     }
 
     final UInt64 headEpoch = spec.computeEpochAtSlot(chainHead.getSlot());
-    if (attestationEpoch
-        .plus(spec.getSpecConfig(headEpoch).getEpochsPerHistoricalVector())
-        .isGreaterThan(headEpoch)) {
-      // If it's an ancestor of head and within historic slots, use the chain head
-      if (isAncestorOfChainHead(chainHead.getRoot(), targetBlockRoot)) {
-        appliedSelectorRule.labels("ancestor_of_head").inc();
-        return chainHead.getState().thenApply(Optional::of);
-      }
-      // if it's an ancestor of any chain head within historic slots, use that chain head.
-      final Optional<BeaconState> maybeChainHeadData =
-          recentChainData.getChainHeads().stream()
-              .filter(head -> isAncestorOfChainHead(head.getRoot(), targetBlockRoot))
-              .findFirst()
-              .flatMap(
-                  protoNodeData -> {
-                    LOG.debug(
-                        "Found chain for target {}, chain head {}",
-                        () -> attestationData.getTarget().getRoot(),
-                        protoNodeData::getStateRoot);
-                    return recentChainData
-                        .getStore()
-                        .getBlockStateIfAvailable(protoNodeData.getRoot());
-                  });
-      if (maybeChainHeadData.isPresent()) {
-        appliedSelectorRule.labels("ancestor_of_fork").inc();
-        return completedFuture(maybeChainHeadData);
-      }
+    final boolean isWithinHistoricalEpochs =
+        attestationEpoch
+            .plus(spec.getSpecConfig(headEpoch).getEpochsPerHistoricalVector())
+            .isGreaterThan(headEpoch);
+    if (isWithinHistoricalEpochs && isAncestorOfChainHead(chainHead.getRoot(), targetBlockRoot)) {
+      appliedSelectorRule.labels("ancestor_of_head").inc();
+      return chainHead.getState().thenApply(Optional::of);
     }
 
     final UInt64 earliestSlot =
@@ -113,10 +93,6 @@ public SafeFuture<Optional<BeaconState>> getStateToValidate(
       return completedFuture(Optional.of(finalizedState));
     }
 
-    // Otherwise, use the state from the earliest allowed slot.
-    // This maximises the chance that the state we get will be on the canonical fork and so useful
-    // for other requests, and means all attestations for that epoch refer to the same slot,
-    // minimising the number of states we need
     final Optional<UInt64> targetBlockSlot = recentChainData.getSlotForBlockRoot(targetBlockRoot);
     if (targetBlockSlot.isEmpty()) {
       // Block became unknown, so ignore it
@@ -124,6 +100,32 @@ public SafeFuture<Optional<BeaconState>> getStateToValidate(
       return completedFuture(Optional.empty());
     }
 
+    // We generally have the chain head states, so attempt to locate which chain head is a
+    // descendent and use that
+    if (isWithinHistoricalEpochs) {
+      // if it's an ancestor of any chain head within historic slots, use that chain head.
+      final Optional<BeaconState> maybeChainHeadData =
+          recentChainData.getChainHeads().stream()
+              .filter(
+                  head ->
+                      isAncestorOfChainHead(head.getRoot(), targetBlockRoot, targetBlockSlot.get()))
+              .findFirst()
+              .flatMap(
+                  protoNodeData -> {
+                    LOG.trace(
+                        "Found chain for target {}, chain head {}",
+                        () -> attestationData.getTarget().getRoot(),
+                        protoNodeData::getStateRoot);
+                    return recentChainData
+                        .getStore()
+                        .getBlockStateIfAvailable(protoNodeData.getRoot());
+                  });
+      if (maybeChainHeadData.isPresent()) {
+        appliedSelectorRule.labels("ancestor_of_fork").inc();
+        return completedFuture(maybeChainHeadData);
+      }
+    }
+
     // Check if the attestations head block state is already available in cache and usable
     if (targetBlockSlot.get().isGreaterThanOrEqualTo(earliestSlot)) {
       final Optional<BeaconState> maybeState =