Add merge fork transition process (#3197)

* Add upgrade state logic

* Initialize TransitionStore from db

* Update stub type

packages/beacon-state-transition/src/allForks/stateTransition.ts

@@ -1,6 +1,6 @@
 /* eslint-disable import/namespace */
-import {allForks, phase0 as phase0Types, Slot, ssz} from "@chainsafe/lodestar-types";
-import {ForkName, GENESIS_EPOCH, SLOTS_PER_EPOCH} from "@chainsafe/lodestar-params";
+import {allForks, Slot, ssz} from "@chainsafe/lodestar-types";
+import {ForkName, SLOTS_PER_EPOCH} from "@chainsafe/lodestar-params";
 import * as phase0 from "../phase0";
 import * as altair from "../altair";
 import * as merge from "../merge";
@@ -12,7 +12,8 @@ import {computeEpochAtSlot} from "../util";
 import {toHexString} from "@chainsafe/ssz";
 
 type StateAllForks = CachedBeaconState<allForks.BeaconState>;
-type StatePhase0 = CachedBeaconState<phase0Types.BeaconState>;
+type StatePhase0 = CachedBeaconState<phase0.BeaconState>;
+type StateAltair = CachedBeaconState<altair.BeaconState>;
 
 type ProcessBlockFn = (state: StateAllForks, block: allForks.BeaconBlock, verifySignatures: boolean) => void;
 type ProcessEpochFn = (state: StateAllForks, epochProcess: IEpochProcess) => void;
@@ -159,12 +160,12 @@ function processSlotsWithTransientCache(
       }
 
       // Upgrade state if exactly at epoch boundary
-      switch (computeEpochAtSlot(postState.slot)) {
-        case GENESIS_EPOCH:
-          break; // Don't do any upgrades at genesis epoch
-        case config.ALTAIR_FORK_EPOCH:
-          postState = altair.upgradeState(postState as StatePhase0) as StateAllForks;
-          break;
+      const stateSlot = computeEpochAtSlot(postState.slot);
+      if (stateSlot === config.ALTAIR_FORK_EPOCH) {
+        postState = altair.upgradeState(postState as StatePhase0) as StateAllForks;
+      }
+      if (stateSlot === config.MERGE_FORK_EPOCH) {
+        postState = merge.upgradeState(postState as StateAltair) as StateAllForks;
       }
     } else {
       postState.slot++;

packages/beacon-state-transition/src/merge/index.ts

@@ -1,4 +1,5 @@
 export {processBlock} from "./block";
+export {upgradeState} from "./upgradeState";
 export * from "./utils";
 
 // re-export merge lodestar types for ergonomic usage downstream

packages/beacon-state-transition/src/merge/upgradeState.ts

@@ -0,0 +1,34 @@
+import {altair, merge, ssz} from "@chainsafe/lodestar-types";
+import {CachedBeaconState, createCachedBeaconState} from "../allForks/util";
+import {getCurrentEpoch} from "../util";
+import {TreeBacked} from "@chainsafe/ssz";
+import {IBeaconConfig} from "@chainsafe/lodestar-config";
+
+/**
+ * Upgrade a state from altair to merge.
+ */
+export function upgradeState(state: CachedBeaconState<altair.BeaconState>): CachedBeaconState<merge.BeaconState> {
+  const {config} = state;
+  const postTreeBackedState = upgradeTreeBackedState(config, state);
+  // TODO: This seems very sub-optimal, review
+  return createCachedBeaconState(config, postTreeBackedState);
+}
+
+function upgradeTreeBackedState(
+  config: IBeaconConfig,
+  state: CachedBeaconState<altair.BeaconState>
+): TreeBacked<merge.BeaconState> {
+  const stateTB = ssz.phase0.BeaconState.createTreeBacked(state.tree);
+
+  // TODO: Does this preserve the hashing cache? In altair devnets memory spikes on the fork transition
+  const postState = ssz.merge.BeaconState.createTreeBacked(stateTB.tree);
+  postState.fork = {
+    previousVersion: stateTB.fork.currentVersion,
+    currentVersion: config.MERGE_FORK_VERSION,
+    epoch: getCurrentEpoch(stateTB),
+  };
+  // Execution-layer
+  postState.latestExecutionPayloadHeader = ssz.merge.ExecutionPayloadHeader.defaultTreeBacked();
+
+  return postState;
+}

packages/db/src/schema.ts

@@ -58,7 +58,9 @@ export enum Bucket {
   altair_lightclientFinalizedCheckpoint = 33, // Epoch -> FinalizedCheckpointData
   // Note: this is the state root for the checkpoint block, NOT necessarily the state root at the epoch boundary
   altair_lightClientInitProof = 34, // Epoch -> Proof
-  altair_lightClientSyncCommitteeProof = 35, // SyncPeriod -> Proof
+  altair_lightClientSyncCommitteeProof = 35, // SyncPeriod ->
+
+  merge_totalTerminalDifficulty = 40, // Single item -> Uint256
 }
 
 export enum Key {

packages/fork-choice/src/forkChoice/forkChoice.ts

@@ -75,7 +75,7 @@ export class ForkChoice implements IForkChoice {
     private readonly config: IChainForkConfig,
     private readonly fcStore: IForkChoiceStore,
     /** Nullable until merge time comes */
-    private readonly transitionStore: ITransitionStore | null,
+    private transitionStore: ITransitionStore | null,
     /** The underlying representation of the block DAG. */
     private readonly protoArray: ProtoArray,
     /**
@@ -91,6 +91,12 @@ export class ForkChoice implements IForkChoice {
     this.head = this.updateHead();
   }
 
+  /** For merge transition. Initialize transition store when merge fork condition is met */
+  initializeTransitionStore(transitionStore: ITransitionStore): void {
+    if (this.transitionStore !== null) throw Error("transitionStore already initialized");
+    this.transitionStore = transitionStore;
+  }
+
   /**
    * Returns the block root of an ancestor of `blockRoot` at the given `slot`.
    * (Note: `slot` refers to the block that is *returned*, not the one that is supplied.)

packages/fork-choice/src/forkChoice/interface.ts

@@ -1,13 +1,16 @@
 import {Epoch, Slot, ValidatorIndex, phase0, allForks, Root, RootHex} from "@chainsafe/lodestar-types";
 import {IProtoBlock} from "../protoArray/interface";
 import {CheckpointWithHex} from "./store";
+import {ITransitionStore} from "./transitionStore";
 
 export type CheckpointHex = {
   epoch: Epoch;
   root: RootHex;
 };
 
 export interface IForkChoice {
+  /** For merge transition. Initialize transition store when merge fork condition is met */
+  initializeTransitionStore(transitionStore: ITransitionStore): void;
   /**
    * Returns the block root of an ancestor of `block_root` at the given `slot`. (Note: `slot` refers
    * to the block that is *returned*, not the one that is supplied.)

packages/fork-choice/src/index.ts

@@ -4,6 +4,7 @@ export {IProtoBlock, IProtoNode} from "./protoArray/interface";
 export {ForkChoice} from "./forkChoice/forkChoice";
 export {IForkChoice} from "./forkChoice/interface";
 export {ForkChoiceStore, IForkChoiceStore, CheckpointWithHex} from "./forkChoice/store";
+export {ITransitionStore} from "./forkChoice/transitionStore";
 export {InvalidAttestation, InvalidAttestationCode, InvalidBlock, InvalidBlockCode} from "./forkChoice/errors";
 
 export {IForkChoiceMetrics} from "./metrics";

packages/lodestar/src/chain/chain.ts

@@ -6,7 +6,7 @@ import fs from "fs";
 import {ForkName} from "@chainsafe/lodestar-params";
 import {CachedBeaconState, computeStartSlotAtEpoch} from "@chainsafe/lodestar-beacon-state-transition";
 import {IBeaconConfig} from "@chainsafe/lodestar-config";
-import {IForkChoice} from "@chainsafe/lodestar-fork-choice";
+import {IForkChoice, ITransitionStore} from "@chainsafe/lodestar-fork-choice";
 import {allForks, ForkDigest, Number64, Root, phase0, Slot} from "@chainsafe/lodestar-types";
 import {ILogger} from "@chainsafe/lodestar-utils";
 import {fromHexString, TreeBacked} from "@chainsafe/ssz";
@@ -44,14 +44,6 @@ import {ForkDigestContext, IForkDigestContext} from "../util/forkDigestContext";
 import {LightClientIniter} from "./lightClient";
 import {Archiver} from "./archiver";
 
-export interface IBeaconChainModules {
-  config: IBeaconConfig;
-  db: IBeaconDb;
-  logger: ILogger;
-  metrics: IMetrics | null;
-  anchorState: TreeBacked<allForks.BeaconState>;
-}
-
 export class BeaconChain implements IBeaconChain {
   readonly genesisTime: Number64;
   readonly genesisValidatorsRoot: Root;
@@ -96,7 +88,24 @@ export class BeaconChain implements IBeaconChain {
   private readonly archiver: Archiver;
   private abortController = new AbortController();
 
-  constructor(opts: IChainOptions, {config, db, logger, metrics, anchorState}: IBeaconChainModules) {
+  constructor(
+    opts: IChainOptions,
+    {
+      config,
+      db,
+      logger,
+      metrics,
+      anchorState,
+      transitionStore,
+    }: {
+      config: IBeaconConfig;
+      db: IBeaconDb;
+      logger: ILogger;
+      metrics: IMetrics | null;
+      anchorState: TreeBacked<allForks.BeaconState>;
+      transitionStore: ITransitionStore | null;
+    }
+  ) {
     this.opts = opts;
     this.config = config;
     this.db = db;
@@ -117,7 +126,7 @@ export class BeaconChain implements IBeaconChain {
     const stateCache = new StateContextCache({metrics});
     const checkpointStateCache = new CheckpointStateCache({metrics});
     const cachedState = restoreStateCaches(config, stateCache, checkpointStateCache, anchorState);
-    const forkChoice = initializeForkChoice(config, emitter, clock.currentSlot, cachedState, metrics);
+    const forkChoice = initializeForkChoice(config, transitionStore, emitter, clock.currentSlot, cachedState, metrics);
     const regen = new QueuedStateRegenerator({
       config,
       emitter,

packages/lodestar/src/chain/forkChoice/index.ts

@@ -5,19 +5,25 @@
 import {toHexString} from "@chainsafe/ssz";
 import {allForks, Slot} from "@chainsafe/lodestar-types";
 import {IChainForkConfig} from "@chainsafe/lodestar-config";
-import {ForkChoice, ProtoArray, ForkChoiceStore} from "@chainsafe/lodestar-fork-choice";
+import {ForkChoice, ProtoArray, ForkChoiceStore, ITransitionStore} from "@chainsafe/lodestar-fork-choice";
 
 import {computeAnchorCheckpoint} from "../initState";
 import {ChainEventEmitter} from "../emitter";
 import {getEffectiveBalances, CachedBeaconState} from "@chainsafe/lodestar-beacon-state-transition";
 import {IMetrics} from "../../metrics";
 import {ChainEvent} from "../emitter";
+import {IBeaconDb} from "../../db";
+
+export type ForkChoiceOpts = {
+  terminalTotalDifficulty?: bigint;
+};
 
 /**
  * Fork Choice extended with a ChainEventEmitter
  */
 export function initializeForkChoice(
   config: IChainForkConfig,
+  transitionStore: ITransitionStore | null,
   emitter: ChainEventEmitter,
   currentSlot: Slot,
   state: CachedBeaconState<allForks.BeaconState>,
@@ -37,9 +43,6 @@ export function initializeForkChoice(
   // TODO - PERFORMANCE WARNING - NAIVE CODE
   const justifiedBalances = getEffectiveBalances(state);
 
-  // TODO: Create and persist transition store
-  const transitionStore = null;
-
   return new ForkChoice(
     config,
 
@@ -65,3 +68,19 @@ export function initializeForkChoice(
     metrics
   );
 }
+
+/**
+ * Initialize TransitionStore with locally persisted value, overriding it with user provided option.
+ */
+export async function initializeTransitionStore(opts: ForkChoiceOpts, db: IBeaconDb): Promise<ITransitionStore | null> {
+  if (opts.terminalTotalDifficulty !== undefined) {
+    return {terminalTotalDifficulty: opts.terminalTotalDifficulty};
+  }
+
+  const terminalTotalDifficulty = await db.totalTerminalDifficulty.get();
+  if (terminalTotalDifficulty !== null) {
+    return {terminalTotalDifficulty: opts.terminalTotalDifficulty ?? terminalTotalDifficulty};
+  }
+
+  return null;
+}

packages/lodestar/src/chain/options.ts

@@ -1,11 +1,13 @@
 import {BlockProcessOpts} from "./blocks/process";
+import {ForkChoiceOpts} from "./forkChoice";
 
 // eslint-disable-next-line @typescript-eslint/ban-types
-export type IChainOptions = BlockProcessOpts & {
-  useSingleThreadVerifier?: boolean;
-  persistInvalidSszObjects?: boolean;
-  persistInvalidSszObjectsDir: string;
-};
+export type IChainOptions = BlockProcessOpts &
+  ForkChoiceOpts & {
+    useSingleThreadVerifier?: boolean;
+    persistInvalidSszObjects?: boolean;
+    persistInvalidSszObjectsDir: string;
+  };
 
 export const defaultChainOptions: IChainOptions = {
   useSingleThreadVerifier: false,

packages/lodestar/src/db/beacon.ts

@@ -25,6 +25,7 @@ import {
   PreGenesisStateLastProcessedBlock,
   LatestFinalizedUpdate,
   LatestNonFinalizedUpdate,
+  TotalTerminalDifficulty,
 } from "./single";
 import {PendingBlockRepository} from "./repositories/pendingBlock";
 
@@ -53,6 +54,7 @@ export class BeaconDb extends DatabaseService implements IBeaconDb {
   lightclientFinalizedCheckpoint: LightclientFinalizedCheckpoint;
   lightClientInitProof: LightClientInitProofRepository;
   lightClientSyncCommitteeProof: LightClientSyncCommitteeProofRepository;
+  totalTerminalDifficulty: TotalTerminalDifficulty;
 
   constructor(opts: IDatabaseApiOptions) {
     super(opts);
@@ -81,6 +83,8 @@ export class BeaconDb extends DatabaseService implements IBeaconDb {
       this.db,
       this.metrics
     );
+    // merge
+    this.totalTerminalDifficulty = new TotalTerminalDifficulty(this.config, this.db, this.metrics);
   }
 
   /**

packages/lodestar/src/db/interface.ts

@@ -25,6 +25,7 @@ import {
   PreGenesisStateLastProcessedBlock,
   LatestFinalizedUpdate,
   LatestNonFinalizedUpdate,
+  TotalTerminalDifficulty,
 } from "./single";
 import {PendingBlockRepository} from "./repositories/pendingBlock";
 
@@ -69,6 +70,7 @@ export interface IBeaconDb {
   lightclientFinalizedCheckpoint: LightclientFinalizedCheckpoint;
   lightClientInitProof: LightClientInitProofRepository;
   lightClientSyncCommitteeProof: LightClientSyncCommitteeProofRepository;
+  totalTerminalDifficulty: TotalTerminalDifficulty;
 
   processBlockOperations(signedBlock: allForks.SignedBeaconBlock): Promise<void>;
 

packages/lodestar/src/db/single/index.ts

@@ -2,3 +2,4 @@ export * from "./latestFinalizedUpdate";
 export * from "./latestNonFinalizedUpdate";
 export * from "./preGenesisState";
 export * from "./preGenesisStateLastProcessedBlock";
+export {TotalTerminalDifficulty} from "./totalTerminalDifficulty";

packages/lodestar/src/db/single/totalTerminalDifficulty.ts

@@ -0,0 +1,33 @@
+import {Type} from "@chainsafe/ssz";
+import {IChainForkConfig} from "@chainsafe/lodestar-config";
+import {ssz} from "@chainsafe/lodestar-types";
+import {IDatabaseController, Bucket, IDbMetrics} from "@chainsafe/lodestar-db";
+
+export class TotalTerminalDifficulty {
+  private readonly bucket = Bucket.merge_totalTerminalDifficulty;
+  private readonly key = Buffer.from(new Uint8Array([this.bucket]));
+  private readonly type: Type<bigint>;
+  private readonly db: IDatabaseController<Buffer, Buffer>;
+  private readonly metrics?: IDbMetrics;
+
+  constructor(config: IChainForkConfig, db: IDatabaseController<Buffer, Buffer>, metrics?: IDbMetrics) {
+    this.db = db;
+    this.type = ssz.Uint256;
+    this.metrics = metrics;
+  }
+
+  async put(value: bigint): Promise<void> {
+    this.metrics?.dbWrites.labels({bucket: "merge_totalTerminalDifficulty"}).inc();
+    await this.db.put(this.key, this.type.serialize(value) as Buffer);
+  }
+
+  async get(): Promise<bigint | null> {
+    this.metrics?.dbReads.labels({bucket: "merge_totalTerminalDifficulty"}).inc();
+    const value = await this.db.get(this.key);
+    return value ? this.type.deserialize(value) : null;
+  }
+
+  async delete(): Promise<void> {
+    await this.db.delete(this.key);
+  }
+}

packages/lodestar/src/node/nodejs.ts

@@ -14,7 +14,7 @@ import {Api} from "@chainsafe/lodestar-api";
 import {IBeaconDb} from "../db";
 import {INetwork, Network, getReqRespHandlers} from "../network";
 import {BeaconSync, IBeaconSync} from "../sync";
-import {BeaconChain, IBeaconChain, initBeaconMetrics} from "../chain";
+import {BeaconChain, IBeaconChain, initBeaconMetrics, initializeTransitionStore} from "../chain";
 import {createMetrics, IMetrics, HttpMetricsServer} from "../metrics";
 import {getApi, RestApi} from "../api";
 import {IBeaconNodeOptions} from "./options";
@@ -131,6 +131,7 @@ export class BeaconNode {
       logger: logger.child(opts.logger.chain),
       metrics,
       anchorState,
+      transitionStore: await initializeTransitionStore(opts.chain, db),
     });
 
     // Load persisted data from disk to in-memory caches

packages/lodestar/test/unit/chain/forkChoice/forkChoice.test.ts

@@ -46,7 +46,8 @@ describe("LodestarForkChoice", function () {
   beforeEach(() => {
     const emitter = new ChainEventEmitter();
     const currentSlot = 40;
-    forkChoice = initializeForkChoice(config, emitter, currentSlot, state);
+    const transitionStore = null;
+    forkChoice = initializeForkChoice(config, transitionStore, emitter, currentSlot, state);
   });
 
   describe("forkchoice", function () {

packages/lodestar/test/utils/mocks/chain/chain.ts

@@ -204,6 +204,7 @@ function mockForkChoice(): IForkChoice {
   const checkpoint: CheckpointWithHex = {epoch: 0, root, rootHex};
 
   return {
+    initializeTransitionStore: () => {},
     getAncestor: () => rootHex,
     getHeadRoot: () => rootHex,
     getHead: () => block,

packages/spec-test-runner/test/spec/allForks/forkChoice.ts

@@ -46,7 +46,8 @@ export function forkChoiceTest(fork: ForkName): void {
         let state = createCachedBeaconState(config, tbState);
 
         const emitter = new ChainEventEmitter();
-        const forkchoice = initializeForkChoice(config, emitter, currentSlot, state);
+        const transitionStore = null;
+        const forkchoice = initializeForkChoice(config, transitionStore, emitter, currentSlot, state);
 
         const checkpointStateCache = new CheckpointStateCache({});
         const stateCache = new Map<string, CachedBeaconState<allForks.BeaconState>>();