From 719ab1d6cefa9807ded9044b6130f03b1bae5038 Mon Sep 17 00:00:00 2001
From: NC <adrninistrator1@protonmail.com>
Date: Tue, 30 Jan 2024 18:25:02 +0800
Subject: [PATCH] Add `getPreStateSync()`

---
 packages/beacon-node/src/chain/chain.ts       | 10 ++-
 .../beacon-node/src/chain/regen/interface.ts  |  1 +
 .../beacon-node/src/chain/regen/queued.ts     | 64 +++++++++++--------
 3 files changed, 45 insertions(+), 30 deletions(-)

diff --git a/packages/beacon-node/src/chain/chain.ts b/packages/beacon-node/src/chain/chain.ts
index 016e5b5cb836..92da02ddab82 100644
--- a/packages/beacon-node/src/chain/chain.ts
+++ b/packages/beacon-node/src/chain/chain.ts
@@ -994,9 +994,13 @@ export class BeaconChain implements IBeaconChain {
   }
 
   async getBlockRewards(block: allForks.FullOrBlindedBeaconBlock): Promise<BlockRewards> {
-    const preState = (await this.regen.getPreState(block, {dontTransferCache: true}, RegenCaller.restApi)).clone();
+    const preState = this.regen.getPreStateSync(block);
     const postState = this.regen.getStateSync(toHexString(block.stateRoot)) ?? undefined;
-    const result = computeBlockRewards(block, preState, postState);
-    return result;
+
+    if (preState === null) {
+      throw Error(`Pre-state is unavailable given block's parent root ${toHexString(block.parentRoot)}`);
+    }
+
+    return computeBlockRewards(block, preState, postState);
   }
 }
diff --git a/packages/beacon-node/src/chain/regen/interface.ts b/packages/beacon-node/src/chain/regen/interface.ts
index be481de9abc8..b861378ff440 100644
--- a/packages/beacon-node/src/chain/regen/interface.ts
+++ b/packages/beacon-node/src/chain/regen/interface.ts
@@ -35,6 +35,7 @@ export interface IStateRegenerator extends IStateRegeneratorInternal {
   dropCache(): void;
   dumpCacheSummary(): routes.lodestar.StateCacheItem[];
   getStateSync(stateRoot: RootHex): CachedBeaconStateAllForks | null;
+  getPreStateSync(block: allForks.BeaconBlock): CachedBeaconStateAllForks | null;
   getCheckpointStateSync(cp: CheckpointHex): CachedBeaconStateAllForks | null;
   getClosestHeadState(head: ProtoBlock): CachedBeaconStateAllForks | null;
   pruneOnCheckpoint(finalizedEpoch: Epoch, justifiedEpoch: Epoch, headStateRoot: RootHex): void;
diff --git a/packages/beacon-node/src/chain/regen/queued.ts b/packages/beacon-node/src/chain/regen/queued.ts
index 928c2e399b9a..a65c462227f3 100644
--- a/packages/beacon-node/src/chain/regen/queued.ts
+++ b/packages/beacon-node/src/chain/regen/queued.ts
@@ -71,6 +71,40 @@ export class QueuedStateRegenerator implements IStateRegenerator {
     return this.stateCache.get(stateRoot);
   }
 
+  getPreStateSync(block: allForks.BeaconBlock): CachedBeaconStateAllForks | null {
+    const parentRoot = toHexString(block.parentRoot);
+    const parentBlock = this.forkChoice.getBlockHex(parentRoot);
+    if (!parentBlock) {
+      throw new RegenError({
+        code: RegenErrorCode.BLOCK_NOT_IN_FORKCHOICE,
+        blockRoot: block.parentRoot,
+      });
+    }
+
+    const parentEpoch = computeEpochAtSlot(parentBlock.slot);
+    const blockEpoch = computeEpochAtSlot(block.slot);
+
+    // Check the checkpoint cache (if the pre-state is a checkpoint state)
+    if (parentEpoch < blockEpoch) {
+      const checkpointState = this.checkpointStateCache.getLatest(parentRoot, blockEpoch);
+      if (checkpointState && computeEpochAtSlot(checkpointState.slot) === blockEpoch) {
+        return checkpointState;
+      }
+    }
+
+    // Check the state cache, only if the state doesn't need to go through an epoch transition.
+    // Otherwise the state transition may not be cached and wasted. Queue for regen since the
+    // work required will still be significant.
+    if (parentEpoch === blockEpoch) {
+      const state = this.stateCache.get(parentBlock.stateRoot);
+      if (state) {
+        return state;
+      }
+    }
+
+    return null;
+  }
+
   getCheckpointStateSync(cp: CheckpointHex): CachedBeaconStateAllForks | null {
     return this.checkpointStateCache.get(cp);
   }
@@ -137,34 +171,10 @@ export class QueuedStateRegenerator implements IStateRegenerator {
     this.metrics?.regenFnCallTotal.inc({caller: rCaller, entrypoint: RegenFnName.getPreState});
 
     // First attempt to fetch the state from caches before queueing
-    const parentRoot = toHexString(block.parentRoot);
-    const parentBlock = this.forkChoice.getBlockHex(parentRoot);
-    if (!parentBlock) {
-      throw new RegenError({
-        code: RegenErrorCode.BLOCK_NOT_IN_FORKCHOICE,
-        blockRoot: block.parentRoot,
-      });
-    }
+    const cachedState = this.getPreStateSync(block);
 
-    const parentEpoch = computeEpochAtSlot(parentBlock.slot);
-    const blockEpoch = computeEpochAtSlot(block.slot);
-
-    // Check the checkpoint cache (if the pre-state is a checkpoint state)
-    if (parentEpoch < blockEpoch) {
-      const checkpointState = this.checkpointStateCache.getLatest(parentRoot, blockEpoch);
-      if (checkpointState && computeEpochAtSlot(checkpointState.slot) === blockEpoch) {
-        return checkpointState;
-      }
-    }
-
-    // Check the state cache, only if the state doesn't need to go through an epoch transition.
-    // Otherwise the state transition may not be cached and wasted. Queue for regen since the
-    // work required will still be significant.
-    if (parentEpoch === blockEpoch) {
-      const state = this.stateCache.get(parentBlock.stateRoot);
-      if (state) {
-        return state;
-      }
+    if (cachedState !== null) {
+      return cachedState;
     }
 
     // The state is not immediately available in the caches, enqueue the job