fix: correct createFromState() with cached current shuffling

ChainSafe · Dec 29, 2023 · 2b9d51b · 2b9d51b
1 parent 4c6c33e
commit 2b9d51b
Show file tree

Hide file tree

Showing 3 changed files with 64 additions and 7 deletions.
diff --git a/packages/state-transition/src/cache/epochCache.ts b/packages/state-transition/src/cache/epochCache.ts
@@ -309,8 +309,10 @@ export class EpochCache {
       if (cachedPreviousShuffling == null && isActiveValidator(validator, previousEpoch)) {
         previousActiveIndices.push(i);
       }
-      if (cachedCurrentShuffling == null && isActiveValidator(validator, currentEpoch)) {
-        currentActiveIndices.push(i);
+      if (isActiveValidator(validator, currentEpoch)) {
+        if (cachedCurrentShuffling == null) {
+          currentActiveIndices.push(i);
+        }
         // We track totalActiveBalanceIncrements as ETH to fit total network balance in a JS number (53 bits)
         totalActiveBalanceIncrements += effectiveBalanceIncrements[i];
       }

diff --git a/packages/state-transition/test/unit/cachedBeaconState.test.ts b/packages/state-transition/test/unit/cachedBeaconState.test.ts
@@ -1,13 +1,14 @@
 import {describe, it, expect} from "vitest";
-import {ssz} from "@lodestar/types";
+import {Epoch, ssz, RootHex} from "@lodestar/types";
 import {toHexString} from "@lodestar/utils";
-import {config} from "@lodestar/config/default";
+import {config as defaultConfig} from "@lodestar/config/default";
 import {createBeaconConfig} from "@lodestar/config";
 import {createCachedBeaconStateTest} from "../utils/state.js";
 import {PubkeyIndexMap} from "../../src/cache/pubkeyCache.js";
 import {createCachedBeaconState, loadUnfinalizedCachedBeaconState} from "../../src/cache/stateCache.js";
 import {interopPubkeysCached} from "../utils/interop.js";
 import {modifyStateSameValidator, newStateWithValidators} from "../utils/capella.js";
+import {EpochShuffling, getShufflingDecisionBlock} from "../../src/util/epochShuffling.js";
 
 describe("CachedBeaconState", () => {
   it("Clone and mutate", () => {
@@ -57,10 +58,11 @@ describe("CachedBeaconState", () => {
     const pubkeys = interopPubkeysCached(2 * numValidator);
 
     const stateView = newStateWithValidators(numValidator);
+    const config = createBeaconConfig(defaultConfig, stateView.genesisValidatorsRoot);
     const seedState = createCachedBeaconState(
       stateView,
       {
-        config: createBeaconConfig(config, stateView.genesisValidatorsRoot),
+        config,
         pubkey2index: new PubkeyIndexMap(),
         index2pubkey: [],
       },
@@ -131,6 +133,43 @@ describe("CachedBeaconState", () => {
         const newStateBytes = newCachedState.serialize();
         expect(newStateBytes).toEqual(stateBytes);
         expect(newCachedState.hashTreeRoot()).toEqual(state.hashTreeRoot());
+        const shufflingGetter = (shufflingEpoch: Epoch, dependentRoot: RootHex) => {
+          if (
+            shufflingEpoch === seedState.epochCtx.epoch - 1 &&
+            dependentRoot === getShufflingDecisionBlock(seedState, shufflingEpoch)
+          ) {
+            return seedState.epochCtx.previousShuffling;
+          }
+
+          if (
+            shufflingEpoch === seedState.epochCtx.epoch &&
+            dependentRoot === getShufflingDecisionBlock(seedState, shufflingEpoch)
+          ) {
+            return seedState.epochCtx.currentShuffling;
+          }
+
+          if (
+            shufflingEpoch === seedState.epochCtx.epoch + 1 &&
+            dependentRoot === getShufflingDecisionBlock(seedState, shufflingEpoch)
+          ) {
+            return seedState.epochCtx.nextShuffling;
+          }
+
+          return null;
+        };
+        const cachedState = createCachedBeaconState(
+          state,
+          {
+            config,
+            pubkey2index: new PubkeyIndexMap(),
+            index2pubkey: [],
+          },
+          {skipSyncCommitteeCache: true, shufflingGetter}
+        );
+        // validatorCountDelta < 0 is unrealistic and shuffling computation results in a different result
+        if (validatorCountDelta >= 0) {
+          expect(newCachedState.epochCtx).toEqual(cachedState.epochCtx);
+        }
 
         // confirm loadUnfinalizedCachedBeaconState() result
         for (let i = 0; i < newCachedState.validators.length; i++) {

diff --git a/packages/state-transition/test/utils/capella.ts b/packages/state-transition/test/utils/capella.ts
@@ -1,7 +1,12 @@
 import crypto from "node:crypto";
 import {ssz} from "@lodestar/types";
 import {config} from "@lodestar/config/default";
-import {BLS_WITHDRAWAL_PREFIX, ETH1_ADDRESS_WITHDRAWAL_PREFIX, SLOTS_PER_EPOCH} from "@lodestar/params";
+import {
+  BLS_WITHDRAWAL_PREFIX,
+  ETH1_ADDRESS_WITHDRAWAL_PREFIX,
+  SLOTS_PER_EPOCH,
+  SLOTS_PER_HISTORICAL_ROOT,
+} from "@lodestar/params";
 import {BeaconStateCapella, CachedBeaconStateCapella} from "../../src/index.js";
 import {createCachedBeaconStateTest} from "./state.js";
 import {mulberry32} from "./rand.js";
@@ -67,10 +72,17 @@ export function newStateWithValidators(numValidator: number): BeaconStateCapella
   const capellaStateType = ssz.capella.BeaconState;
   const stateView = capellaStateType.defaultViewDU();
   stateView.slot = config.CAPELLA_FORK_EPOCH * SLOTS_PER_EPOCH + 100;
+  for (let i = 0; i < SLOTS_PER_HISTORICAL_ROOT; i++) {
+    stateView.blockRoots.set(i, crypto.randomBytes(32));
+  }
 
   for (let i = 0; i < numValidator; i++) {
     const validator = ssz.phase0.Validator.defaultViewDU();
     validator.pubkey = pubkeys[i];
+    // make all validators active
+    validator.activationEpoch = 0;
+    validator.exitEpoch = Infinity;
+    validator.effectiveBalance = 32e9;
     stateView.validators.push(validator);
     stateView.balances.push(32);
     stateView.inactivityScores.push(0);
@@ -85,15 +97,19 @@ export function newStateWithValidators(numValidator: number): BeaconStateCapella
  * Modify a state without changing number of validators
  */
 export function modifyStateSameValidator(seedState: BeaconStateCapella): BeaconStateCapella {
+  const slotDiff = 10;
   const state = seedState.clone();
-  state.slot = seedState.slot + 10;
+  state.slot = seedState.slot + slotDiff;
   state.latestBlockHeader = ssz.phase0.BeaconBlockHeader.toViewDU({
     slot: state.slot,
     proposerIndex: 0,
     parentRoot: state.hashTreeRoot(),
     stateRoot: state.hashTreeRoot(),
     bodyRoot: ssz.phase0.BeaconBlockBody.hashTreeRoot(ssz.phase0.BeaconBlockBody.defaultValue()),
   });
+  for (let i = 1; i <= slotDiff; i++) {
+    state.blockRoots.set((seedState.slot + i) % SLOTS_PER_HISTORICAL_ROOT, crypto.randomBytes(32));
+  }
   state.blockRoots.set(0, crypto.randomBytes(32));
   state.stateRoots.set(0, crypto.randomBytes(32));
   state.historicalRoots.push(crypto.randomBytes(32));