fix: resolve state id as checkpoint

packages/beacon-node/src/api/impl/beacon/state/utils.ts

@@ -3,11 +3,14 @@ import {FAR_FUTURE_EPOCH, GENESIS_SLOT} from "@lodestar/params";
 import {BeaconStateAllForks, PubkeyIndexMap} from "@lodestar/state-transition";
 import {BLSPubkey, Epoch, phase0, RootHex, Slot, ValidatorIndex} from "@lodestar/types";
 import {fromHex} from "@lodestar/utils";
-import {IForkChoice} from "@lodestar/fork-choice";
+import {CheckpointWithHex, IForkChoice} from "@lodestar/fork-choice";
 import {IBeaconChain} from "../../../../chain/index.js";
 import {ApiError, ValidationError} from "../../errors.js";
 
-export function resolveStateId(forkChoice: IForkChoice, stateId: routes.beacon.StateId): RootHex | Slot {
+export function resolveStateId(
+  forkChoice: IForkChoice,
+  stateId: routes.beacon.StateId
+): RootHex | Slot | CheckpointWithHex {
   if (stateId === "head") {
     return forkChoice.getHead().stateRoot;
   }
@@ -17,11 +20,11 @@ export function resolveStateId(forkChoice: IForkChoice, stateId: routes.beacon.S
   }
 
   if (stateId === "finalized") {
-    return forkChoice.getFinalizedBlock().stateRoot;
+    return forkChoice.getFinalizedCheckpoint();
   }
 
   if (stateId === "justified") {
-    return forkChoice.getJustifiedBlock().stateRoot;
+    return forkChoice.getJustifiedCheckpoint();
   }
 
   if (typeof stateId === "string" && stateId.startsWith("0x")) {
@@ -39,37 +42,41 @@ export function resolveStateId(forkChoice: IForkChoice, stateId: routes.beacon.S
 
 export async function getStateResponse(
   chain: IBeaconChain,
-  stateId: routes.beacon.StateId
+  inStateId: routes.beacon.StateId
 ): Promise<{state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean}> {
-  const rootOrSlot = resolveStateId(chain.forkChoice, stateId);
+  const stateId = resolveStateId(chain.forkChoice, inStateId);
 
   const res =
-    typeof rootOrSlot === "string"
-      ? await chain.getStateByStateRoot(rootOrSlot)
-      : await chain.getStateBySlot(rootOrSlot);
+    typeof stateId === "string"
+      ? await chain.getStateByStateRoot(stateId)
+      : typeof stateId === "number"
+        ? await chain.getStateBySlot(stateId)
+        : chain.getStateByCheckpoint(stateId);
 
   if (!res) {
-    throw new ApiError(404, `No state found for id '${stateId}'`);
+    throw new ApiError(404, `No state found for id '${inStateId}'`);
   }
 
   return res;
 }
 
 export async function getStateResponseWithRegen(
   chain: IBeaconChain,
-  stateId: routes.beacon.StateId
+  inStateId: routes.beacon.StateId
 ): Promise<{state: BeaconStateAllForks | Uint8Array; executionOptimistic: boolean; finalized: boolean}> {
-  const rootOrSlot = resolveStateId(chain.forkChoice, stateId);
+  const stateId = resolveStateId(chain.forkChoice, inStateId);
 
   const res =
-    typeof rootOrSlot === "string"
-      ? await chain.getStateByStateRoot(rootOrSlot, {allowRegen: true})
-      : rootOrSlot >= chain.forkChoice.getFinalizedBlock().slot
-        ? await chain.getStateBySlot(rootOrSlot, {allowRegen: true})
-        : await chain.getHistoricalStateBySlot(rootOrSlot);
+    typeof stateId === "string"
+      ? await chain.getStateByStateRoot(stateId, {allowRegen: true})
+      : typeof stateId === "number"
+        ? stateId >= chain.forkChoice.getFinalizedBlock().slot
+          ? await chain.getStateBySlot(stateId, {allowRegen: true})
+          : await chain.getHistoricalStateBySlot(stateId)
+        : await chain.getStateOrBytesByCheckpoint(stateId);
 
   if (!res) {
-    throw new ApiError(404, `No state found for id '${stateId}'`);
+    throw new ApiError(404, `No state found for id '${inStateId}'`);
   }
 
   return res;

packages/beacon-node/src/chain/chain.ts

@@ -518,7 +518,7 @@ export class BeaconChain implements IBeaconChain {
   getStateByCheckpoint(
     checkpoint: CheckpointWithHex
   ): {state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean} | null {
-    // TODO: this is not guaranteed to work with new state caches, should work on this before we turn n-historical state on
+    // finalized or justified checkpoint states maynot be available with PersistentCheckpointStateCache, use getCheckpointStateOrBytes() api to get Uint8Array
     const cachedStateCtx = this.regen.getCheckpointStateSync(checkpoint);
     if (cachedStateCtx) {
       const block = this.forkChoice.getBlock(cachedStateCtx.latestBlockHeader.hashTreeRoot());
@@ -533,6 +533,23 @@ export class BeaconChain implements IBeaconChain {
     return null;
   }
 
+  async getStateOrBytesByCheckpoint(
+    checkpoint: CheckpointWithHex
+  ): Promise<{state: CachedBeaconStateAllForks | Uint8Array; executionOptimistic: boolean; finalized: boolean} | null> {
+    const cachedStateCtx = await this.regen.getCheckpointStateOrBytes(checkpoint);
+    if (cachedStateCtx) {
+      const block = this.forkChoice.getBlock(checkpoint.root);
+      const finalizedEpoch = this.forkChoice.getFinalizedCheckpoint().epoch;
+      return {
+        state: cachedStateCtx,
+        executionOptimistic: block != null && isOptimisticBlock(block),
+        finalized: checkpoint.epoch <= finalizedEpoch && finalizedEpoch !== GENESIS_EPOCH,
+      };
+    }
+
+    return null;
+  }
+
   async getCanonicalBlockAtSlot(
     slot: Slot
   ): Promise<{block: SignedBeaconBlock; executionOptimistic: boolean; finalized: boolean} | null> {

packages/beacon-node/src/chain/interface.ts

@@ -160,6 +160,10 @@ export interface IBeaconChain {
   getStateByCheckpoint(
     checkpoint: CheckpointWithHex
   ): {state: BeaconStateAllForks; executionOptimistic: boolean; finalized: boolean} | null;
+  /** Return state bytes by checkpoint */
+  getStateOrBytesByCheckpoint(
+    checkpoint: CheckpointWithHex
+  ): Promise<{state: CachedBeaconStateAllForks | Uint8Array; executionOptimistic: boolean; finalized: boolean} | null>;
 
   /**
    * Since we can have multiple parallel chains,