From d418f8dae79692bf28db1d6db9475992c6a78140 Mon Sep 17 00:00:00 2001
From: Lucas Saldanha <lucas.saldanha@consensys.net>
Date: Fri, 20 Oct 2023 10:18:04 +1300
Subject: [PATCH] Implemented isFinalizedRootOrDescendant (#7606)

---
 .../teku/storage/protoarray/ProtoArray.java   | 35 +++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

diff --git a/storage/src/main/java/tech/pegasys/teku/storage/protoarray/ProtoArray.java b/storage/src/main/java/tech/pegasys/teku/storage/protoarray/ProtoArray.java
index bc5e9d17abb..e0c1066b3f8 100644
--- a/storage/src/main/java/tech/pegasys/teku/storage/protoarray/ProtoArray.java
+++ b/storage/src/main/java/tech/pegasys/teku/storage/protoarray/ProtoArray.java
@@ -639,13 +639,44 @@ public boolean nodeIsViableForHead(ProtoNode node) {
                   .isGreaterThanOrEqualTo(currentEpoch);
     }
 
-    final UInt64 finalizedSlot = spec.computeStartSlotAtEpoch(finalizedCheckpoint.getEpoch());
     final boolean correctFinalized =
         node.getFinalizedCheckpoint().getEpoch().equals(initialEpoch)
-            || hasAncestorAtSlot(node, finalizedSlot, finalizedCheckpoint.getRoot());
+            || isFinalizedRootOrDescendant(node);
     return correctJustified && correctFinalized;
   }
 
+  private boolean isFinalizedRootOrDescendant(final ProtoNode node) {
+    final UInt64 finalizedEpoch = finalizedCheckpoint.getEpoch();
+    final Bytes32 finalizedRoot = finalizedCheckpoint.getRoot();
+
+    final Checkpoint nodeFinalizedCheckpoint = node.getFinalizedCheckpoint();
+    if (nodeFinalizedCheckpoint.getEpoch().equals(finalizedEpoch)
+        && nodeFinalizedCheckpoint.getRoot().equals(finalizedRoot)) {
+      return true;
+    }
+
+    final Checkpoint nodeJustifiedCheckpoint = node.getJustifiedCheckpoint();
+    if (nodeJustifiedCheckpoint.getEpoch().equals(finalizedEpoch)
+        && nodeJustifiedCheckpoint.getRoot().equals(finalizedRoot)) {
+      return true;
+    }
+
+    final Checkpoint nodeUnrealizedFinalizedCheckpoint = node.getUnrealizedFinalizedCheckpoint();
+    if (nodeUnrealizedFinalizedCheckpoint.getEpoch().equals(finalizedEpoch)
+        && nodeUnrealizedFinalizedCheckpoint.getRoot().equals(finalizedRoot)) {
+      return true;
+    }
+
+    final Checkpoint nodeUnrealizedJustifiedCheckpoint = node.getUnrealizedJustifiedCheckpoint();
+    if (nodeUnrealizedJustifiedCheckpoint.getEpoch().equals(finalizedEpoch)
+        && nodeUnrealizedJustifiedCheckpoint.getRoot().equals(finalizedRoot)) {
+      return true;
+    }
+
+    final UInt64 finalizedSlot = spec.computeStartSlotAtEpoch(finalizedCheckpoint.getEpoch());
+    return hasAncestorAtSlot(node, finalizedSlot, finalizedRoot);
+  }
+
   /**
    * This is similar to the <a
    * href="https://github.com/ethereum/consensus-specs/blob/dev/specs/phase0/fork-choice.md#get_ancestor">get_ancestor</a>