refactor: update the structure of state regen operation

packages/beacon-node/src/chain/archiveStore/differentialState/apply.ts

@@ -0,0 +1,95 @@
+import {PubkeyIndexMap} from "@chainsafe/pubkey-index-map";
+import {BeaconConfig} from "@lodestar/config";
+import {Logger} from "@lodestar/logger";
+import {IBeaconDb} from "../../../db/index.ts";
+import {IStateDiffCodec} from "../interface.ts";
+import {replayBlocks} from "../utils/replayBlocks.ts";
+import {StateRegenArtifacts} from "./fetch.ts";
+import {StateRegenPlan} from "./plan.ts";
+import {BeaconStateSnapshot} from "./ssz.ts";
+import {replayStateDifferentials} from "./stateDifferential.ts";
+import {beaconStateBytesToSnapshot, snapshotToBeaconStateBytes} from "./stateSnapshot.ts";
+
+export type StateRegenContext = {
+  codec: IStateDiffCodec;
+  config: BeaconConfig;
+  logger?: Logger;
+  pubkey2index: PubkeyIndexMap;
+  db: IBeaconDb;
+};
+
+export async function applyStateRegenPlan(
+  ctx: StateRegenContext,
+  plan: StateRegenPlan,
+  artifacts: StateRegenArtifacts
+): Promise<BeaconStateSnapshot> {
+  // When we start a node from a certain checkpoint which is usually
+  // not the snapshot epoch but we fetch it because of the fallback settings
+  if (plan.snapshotSlot !== artifacts.snapshot.slot) {
+    ctx.logger?.warn("Expected snapshot not found", {
+      expectedSnapshotSlot: plan.snapshotSlot,
+      availableSnapshotSlot: artifacts.snapshot.slot,
+    });
+  }
+
+  // TODO: Need to do further thinking if we fail here with fatal error
+  if (artifacts.missingDiffs.length) {
+    ctx.logger?.warn("Missing some diff states", {
+      snapshotSlot: plan.snapshotSlot,
+      diffPath: plan.diffSlots.join(","),
+      missingDiffs: artifacts.missingDiffs.join(","),
+    });
+  }
+
+  const orderedDiffs = [];
+  for (const diffSlot of plan.diffSlots) {
+    const diff = artifacts.diffs.find((d) => d.slot === diffSlot);
+    if (diff) {
+      orderedDiffs.push(diff);
+    }
+  }
+
+  if (orderedDiffs.length + artifacts.missingDiffs.length !== plan.diffSlots.length) {
+    throw new Error(`Can not find required state diffs ${plan.diffSlots.join(",")}`);
+  }
+
+  if (plan.blockReplay && orderedDiffs.at(-1)?.slot !== plan.blockReplay.fromSlot - 1) {
+    throw new Error(`Can not replay blocks due to missing state diffs ${artifacts.missingDiffs.join(",")}`);
+  }
+
+  ctx.logger?.verbose("Replaying state diffs", {
+    snapshotSlot: plan.snapshotSlot,
+    diffPath: plan.diffSlots.join(","),
+    availableDiffs: orderedDiffs.map((d) => d.slot).join(","),
+  });
+
+  const stateWithDiffApplied = await replayStateDifferentials(
+    {codec: ctx.codec, logger: ctx.logger},
+    {stateDifferentials: orderedDiffs, stateSnapshot: artifacts.snapshot}
+  );
+
+  if (stateWithDiffApplied.stateBytes.byteLength === 0 || stateWithDiffApplied.balancesBytes.byteLength === 0) {
+    throw new Error(
+      `Invalid state after applying diffs: 
+      stateBytesSize=${stateWithDiffApplied.stateBytes.byteLength},
+      balancesBytesSize=${stateWithDiffApplied.balancesBytes.byteLength}`
+    );
+  }
+
+  if (!plan.blockReplay) return stateWithDiffApplied;
+
+  const stateBytes = snapshotToBeaconStateBytes({config: ctx.config}, stateWithDiffApplied);
+
+  ctx.logger?.verbose("Replaying blocks", {
+    fromSlot: plan.blockReplay.fromSlot,
+    tillSlot: plan.blockReplay.tillSlot,
+  });
+
+  const replayed = await replayBlocks(ctx, {
+    stateBytes,
+    fromSlot: plan.blockReplay.fromSlot,
+    toSlot: plan.blockReplay.tillSlot,
+  });
+
+  return beaconStateBytesToSnapshot({config: ctx.config}, plan.blockReplay.tillSlot, replayed);
+}

packages/beacon-node/src/chain/archiveStore/differentialState/execute.ts

@@ -0,0 +1,22 @@
+import {BeaconState, Slot} from "@lodestar/types";
+import {StateRegenContext, applyStateRegenPlan} from "./apply.ts";
+import {fetchStateRegenArtifacts} from "./fetch.ts";
+import {HierarchicalLayers} from "./hierarchicalLayers.ts";
+import {buildStateRegenPlan} from "./plan.ts";
+import {snapshotToBeaconState} from "./stateSnapshot.ts";
+
+export async function regenerateState(
+  ctx: StateRegenContext & {layers: HierarchicalLayers},
+  target: Slot,
+  opts?: {fallbackSnapshot?: boolean}
+): Promise<BeaconState | null> {
+  ctx.logger?.verbose("Regenerating state via state differential", {
+    slot: target,
+  });
+
+  const plan = buildStateRegenPlan(ctx.layers, target);
+  const artifacts = await fetchStateRegenArtifacts(ctx.db, plan, opts);
+  const finalState = await applyStateRegenPlan(ctx, plan, artifacts);
+
+  return snapshotToBeaconState(ctx, finalState);
+}

packages/beacon-node/src/chain/archiveStore/differentialState/fetch.ts

@@ -0,0 +1,34 @@
+import {Slot} from "@lodestar/types";
+import {IBeaconDb} from "../../../db/index.ts";
+import {StateRegenPlan} from "./plan.ts";
+import {BeaconStateDifferential, BeaconStateSnapshot} from "./ssz.ts";
+import {getStateDifferential} from "./stateDifferential.ts";
+import {getStateSnapshot} from "./stateSnapshot.ts";
+
+export type StateRegenArtifacts = {
+  snapshot: BeaconStateSnapshot;
+  diffs: BeaconStateDifferential[];
+  missingDiffs: Slot[];
+};
+
+export async function fetchStateRegenArtifacts(
+  db: IBeaconDb,
+  plan: StateRegenPlan,
+  opts: {fallbackSnapshot?: boolean} = {}
+): Promise<StateRegenArtifacts> {
+  const snapshot = await getStateSnapshot({db}, {slot: plan.snapshotSlot, fallback: opts.fallbackSnapshot ?? false});
+
+  if (!snapshot) {
+    throw new Error(`Can not find state snapshot for slot=${plan.snapshotSlot}`);
+  }
+
+  const diffs: BeaconStateDifferential[] = [];
+  const missingDiffs: Slot[] = [];
+
+  for (const edge of plan.diffSlots) {
+    const diff = await getStateDifferential({db}, {slot: edge});
+    diff ? diffs.push(diff) : missingDiffs.push(edge);
+  }
+
+  return {snapshot, diffs, missingDiffs};
+}

packages/beacon-node/src/chain/archiveStore/differentialState/plan.ts

@@ -0,0 +1,34 @@
+import {Slot} from "@lodestar/types";
+import {HierarchicalLayers} from "./hierarchicalLayers.ts";
+
+export type StateRegenPlan = {
+  targetSlot: Slot;
+  snapshotSlot: Slot;
+  diffSlots: Slot[];
+  blockReplay?: {fromSlot: Slot; tillSlot: Slot};
+};
+
+export function buildStateRegenPlan(layers: HierarchicalLayers, target: Slot): StateRegenPlan {
+  const path = layers.computeSlotPath(target);
+  const [snapshotSlot, ...diffSlots] = path;
+  const lastDiffSlot = diffSlots.at(-1);
+
+  if (target === lastDiffSlot || target === snapshotSlot) {
+    return {
+      snapshotSlot,
+      diffSlots,
+      blockReplay: undefined,
+      targetSlot: target,
+    };
+  }
+
+  return {
+    snapshotSlot,
+    diffSlots,
+    blockReplay: {
+      fromSlot: lastDiffSlot ? lastDiffSlot + 1 : snapshotSlot + 1,
+      tillSlot: target,
+    },
+    targetSlot: target,
+  };
+}

packages/beacon-node/src/chain/archiveStore/utils/replayBlocks.ts

@@ -47,8 +47,7 @@ export async function replayBlocks(
     }
   );
 
-  // Will use this for metrics
-  // biome-ignore lint/correctness/noUnusedVariables: <explanation>
+  // biome-ignore lint/correctness/noUnusedVariables: Will use this for metrics
   let blockCount = 0;
 
   for await (const block of db.blockArchive.valuesStream({gt: fromSlot, lte: toSlot})) {
@@ -61,8 +60,7 @@ export async function replayBlocks(
         dataAvailabilityStatus: DataAvailabilityStatus.Available,
       });
     } catch (e) {
-      // Add metrics for error
-      // biome-ignore lint/complexity/noUselessCatch: <explanation>
+      // biome-ignore lint/complexity/noUselessCatch: Add metrics for error
       throw e;
     }
     blockCount++;

packages/beacon-node/src/db/repositories/index.ts

@@ -1,11 +1,11 @@
 export {AttesterSlashingRepository} from "./attesterSlashing.js";
 export {BackfilledRanges} from "./backfilledRanges.js";
+export {BeaconStateDifferentialArchiveRepository} from "./beaconStateDifferentialArchive.js";
+export {BeaconStateSnapshotArchiveRepository} from "./beaconStateSnapshotArchive.js";
 export {BlobSidecarsRepository} from "./blobSidecars.js";
 export {BlobSidecarsArchiveRepository} from "./blobSidecarsArchive.js";
 export {BlockRepository} from "./block.js";
 export type {BlockArchiveBatchPutBinaryItem, BlockFilterOptions} from "./blockArchive.js";
-export {BeaconStateSnapshotArchiveRepository} from "./beaconStateSnapshotArchive.js";
-export {BeaconStateDifferentialArchiveRepository} from "./beaconStateDifferentialArchive.js";
 export {BlockArchiveRepository} from "./blockArchive.js";
 export {BLSToExecutionChangeRepository} from "./blsToExecutionChange.js";
 export {DataColumnSidecarRepository} from "./dataColumnSidecar.js";

packages/beacon-node/test/fixtures/differentialState/hierarchicalLayers.ts

@@ -323,7 +323,7 @@ export const nonOverlappingLayersData: LayersTest[] = [
     path: [0, computeStartSlotAtEpoch(3)],
   },
   {
-    title: "after slot of first diff layer",
+    title: "one slot after first diff layer",
     slot: computeStartSlotAtEpoch(3) + 1,
     path: [0, computeStartSlotAtEpoch(3)],
     blockReplay: {
@@ -346,7 +346,7 @@ export const nonOverlappingLayersData: LayersTest[] = [
     path: [0, computeStartSlotAtEpoch(5)],
   },
   {
-    title: "after slot of second diff layer",
+    title: "one slot after second diff layer",
     slot: computeStartSlotAtEpoch(5) + 1,
     path: [0, computeStartSlotAtEpoch(5)],
     blockReplay: {