From 3d6d19badc9447ed74ce4cf5fbbdbf806ff4a2de Mon Sep 17 00:00:00 2001
From: Tuyen Nguyen <vutuyen2636@gmail.com>
Date: Wed, 23 Aug 2023 17:15:38 +0700
Subject: [PATCH] fix: improve forkchoice getBlock hasBlock

---
 .../fork-choice/src/forkChoice/forkChoice.ts  | 31 +++++++++----------
 .../fork-choice/src/forkChoice/interface.ts   |  4 ---
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/packages/fork-choice/src/forkChoice/forkChoice.ts b/packages/fork-choice/src/forkChoice/forkChoice.ts
index 6dda2c639df2..3cb2f152de6c 100644
--- a/packages/fork-choice/src/forkChoice/forkChoice.ts
+++ b/packages/fork-choice/src/forkChoice/forkChoice.ts
@@ -592,7 +592,12 @@ export class ForkChoice implements IForkChoice {
    * Returns `true` if the block is known **and** a descendant of the finalized root.
    */
   hasBlockHex(blockRoot: RootHex): boolean {
-    return this.protoArray.hasBlock(blockRoot) && this.isDescendantOfFinalized(blockRoot);
+    const node = this.protoArray.getNode(blockRoot);
+    if (node === undefined) {
+      return false;
+    }
+
+    return this.protoArray.isFinalizedRootOrDescendant(node);
   }
 
   /**
@@ -610,20 +615,21 @@ export class ForkChoice implements IForkChoice {
   }
 
   /**
-   * Returns a `ProtoBlock` if the block is known **and** a descendant of the finalized root.
+   * Returns a MUTABLE `ProtoBlock` if the block is known **and** a descendant of the finalized root.
    */
   getBlockHex(blockRoot: RootHex): ProtoBlock | null {
-    const block = this.protoArray.getBlock(blockRoot);
-    if (!block) {
+    const node = this.protoArray.getNode(blockRoot);
+    if (!node) {
       return null;
     }
-    // If available, use the parent_root to perform the lookup since it will involve one
-    // less lookup. This involves making the assumption that the finalized block will
-    // always have `block.parent_root` of `None`.
-    if (!this.isDescendantOfFinalized(blockRoot)) {
+
+    if (!this.protoArray.isFinalizedRootOrDescendant(node)) {
       return null;
     }
-    return block;
+
+    return {
+      ...node,
+    };
   }
 
   getJustifiedBlock(): ProtoBlock {
@@ -648,13 +654,6 @@ export class ForkChoice implements IForkChoice {
     return block;
   }
 
-  /**
-   * Return `true` if `block_root` is equal to the finalized root, or a known descendant of it.
-   */
-  isDescendantOfFinalized(blockRoot: RootHex): boolean {
-    return this.protoArray.isDescendant(this.fcStore.finalizedCheckpoint.rootHex, blockRoot);
-  }
-
   /**
    * Returns true if the `descendantRoot` has an ancestor with `ancestorRoot`.
    *
diff --git a/packages/fork-choice/src/forkChoice/interface.ts b/packages/fork-choice/src/forkChoice/interface.ts
index adedd4cd3c9d..6d26f17d493b 100644
--- a/packages/fork-choice/src/forkChoice/interface.ts
+++ b/packages/fork-choice/src/forkChoice/interface.ts
@@ -161,10 +161,6 @@ export interface IForkChoice {
   getBlockHex(blockRoot: RootHex): ProtoBlock | null;
   getFinalizedBlock(): ProtoBlock;
   getJustifiedBlock(): ProtoBlock;
-  /**
-   * Return `true` if `block_root` is equal to the finalized root, or a known descendant of it.
-   */
-  isDescendantOfFinalized(blockRoot: RootHex): boolean;
   /**
    * Returns true if the `descendantRoot` has an ancestor with `ancestorRoot`.
    *