feat: use rust shuffle

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

@@ -4,11 +4,11 @@ import {
   IShufflingCache,
   ShufflingBuildProps,
   computeEpochShuffling,
+  computeEpochShufflingAsync,
 } from "@lodestar/state-transition";
 import {Epoch, RootHex} from "@lodestar/types";
 import {LodestarError, Logger, MapDef, pruneSetToMax} from "@lodestar/utils";
 import {Metrics} from "../metrics/metrics.js";
-import {callInNextEventLoop} from "../util/eventLoop.js";
 
 /**
  * Same value to CheckpointBalancesCache, with the assumption that we don't have to use it for old epochs. In the worse case:
@@ -178,14 +178,19 @@ export class ShufflingCache implements IShufflingCache {
     this.insertPromise(epoch, decisionRoot);
     /**
      * TODO: (@matthewkeil) This will get replaced by a proper build queue and a worker to do calculations
-     * on a NICE thread with a rust implementation
+     * on a NICE thread
      */
-    callInNextEventLoop(() => {
-      const timer = this.metrics?.shufflingCache.shufflingCalculationTime.startTimer({source: "build"});
-      const shuffling = computeEpochShuffling(state, activeIndices, epoch);
-      timer?.();
-      this.set(shuffling, decisionRoot);
-    });
+    const timer = this.metrics?.shufflingCache.shufflingCalculationTime.startTimer({source: "build"});
+    computeEpochShufflingAsync(state, activeIndices, epoch)
+      .then((shuffling) => {
+        this.set(shuffling, decisionRoot);
+      })
+      .catch((err) =>
+        this.logger?.error(`error building shuffling for epoch ${epoch} at decisionRoot ${decisionRoot}`, {}, err)
+      )
+      .finally(() => {
+        timer?.();
+      });
   }
 
   /**

packages/beacon-node/test/spec/presets/shuffling.test.ts

@@ -1,24 +1,27 @@
 import path from "node:path";
-import {unshuffleList} from "@lodestar/state-transition";
+import {unshuffleList} from "@chainsafe/swap-or-not-shuffle";
 import {InputType} from "@lodestar/spec-test-util";
 import {bnToNum, fromHex} from "@lodestar/utils";
-import {ACTIVE_PRESET} from "@lodestar/params";
+import {ACTIVE_PRESET, SHUFFLE_ROUND_COUNT} from "@lodestar/params";
 import {RunnerType, TestRunnerFn} from "../utils/types.js";
 import {ethereumConsensusSpecsTests} from "../specTestVersioning.js";
 import {specTestIterator} from "../utils/specTestIterator.js";
 
-const shuffling: TestRunnerFn<ShufflingTestCase, number[]> = () => {
+const shuffling: TestRunnerFn<ShufflingTestCase, string> = () => {
   return {
     testFunction: (testcase) => {
       const seed = fromHex(testcase.mapping.seed);
-      const output = Array.from({length: bnToNum(testcase.mapping.count)}, (_, i) => i);
-      unshuffleList(output, seed);
-      return output;
+      const output = unshuffleList(
+        Uint32Array.from(Array.from({length: bnToNum(testcase.mapping.count)}, (_, i) => i)),
+        seed,
+        SHUFFLE_ROUND_COUNT
+      );
+      return Buffer.from(output).toString("hex");
     },
     options: {
       inputTypes: {mapping: InputType.YAML},
       timeout: 10000,
-      getExpected: (testCase) => testCase.mapping.mapping.map((value) => bnToNum(value)),
+      getExpected: (testCase) => Buffer.from(testCase.mapping.mapping.map((value) => bnToNum(value))).toString("hex"),
       // Do not manually skip tests here, do it in packages/beacon-node/test/spec/presets/index.test.ts
     },
   };

packages/state-transition/package.json

@@ -63,6 +63,7 @@
     "@chainsafe/persistent-merkle-tree": "^0.8.0",
     "@chainsafe/persistent-ts": "^0.19.1",
     "@chainsafe/ssz": "^0.17.1",
+    "@chainsafe/swap-or-not-shuffle": "^0.0.1",
     "@lodestar/config": "^1.22.0",
     "@lodestar/params": "^1.22.0",
     "@lodestar/types": "^1.22.0",

packages/state-transition/src/util/epochShuffling.ts

@@ -1,3 +1,4 @@
+import {asyncUnshuffleList, unshuffleList} from "@chainsafe/swap-or-not-shuffle";
 import {Epoch, RootHex, ssz, ValidatorIndex} from "@lodestar/types";
 import {GaugeExtra, intDiv, Logger, NoLabels, toRootHex} from "@lodestar/utils";
 import {
@@ -6,11 +7,11 @@ import {
   MAX_COMMITTEES_PER_SLOT,
   SLOTS_PER_EPOCH,
   TARGET_COMMITTEE_SIZE,
+  SHUFFLE_ROUND_COUNT,
 } from "@lodestar/params";
 import {BeaconConfig} from "@lodestar/config";
 import {BeaconStateAllForks} from "../types.js";
 import {getSeed} from "./seed.js";
-import {unshuffleList} from "./shuffle.js";
 import {computeStartSlotAtEpoch} from "./epoch.js";
 import {getBlockRootAtSlot} from "./blockRoot.js";
 import {computeAnchorCheckpoint} from "./computeAnchorCheckpoint.js";
@@ -102,42 +103,67 @@ export function computeCommitteeCount(activeValidatorCount: number): number {
   return Math.max(1, Math.min(MAX_COMMITTEES_PER_SLOT, committeesPerSlot));
 }
 
-export function computeEpochShuffling(
-  state: BeaconStateAllForks,
-  activeIndices: Uint32Array,
-  epoch: Epoch
-): EpochShuffling {
-  const activeValidatorCount = activeIndices.length;
-
-  const shuffling = activeIndices.slice();
-  const seed = getSeed(state, epoch, DOMAIN_BEACON_ATTESTER);
-  unshuffleList(shuffling, seed);
-
+function buildCommitteesFromShuffling(
+  activeValidatorCount: number,
+  shuffling: Uint32Array
+): Pick<EpochShuffling, "committees" | "committeesPerSlot"> {
   const committeesPerSlot = computeCommitteeCount(activeValidatorCount);
-
   const committeeCount = committeesPerSlot * SLOTS_PER_EPOCH;
 
-  const committees: Uint32Array[][] = [];
+  const committees = new Array<Uint32Array[]>(SLOTS_PER_EPOCH);
   for (let slot = 0; slot < SLOTS_PER_EPOCH; slot++) {
-    const slotCommittees: Uint32Array[] = [];
+    const slotCommittees = new Array<Uint32Array>(committeesPerSlot);
+
     for (let committeeIndex = 0; committeeIndex < committeesPerSlot; committeeIndex++) {
       const index = slot * committeesPerSlot + committeeIndex;
       const startOffset = Math.floor((activeValidatorCount * index) / committeeCount);
       const endOffset = Math.floor((activeValidatorCount * (index + 1)) / committeeCount);
       if (!(startOffset <= endOffset)) {
         throw new Error(`Invalid offsets: start ${startOffset} must be less than or equal end ${endOffset}`);
       }
-      slotCommittees.push(shuffling.subarray(startOffset, endOffset));
+      slotCommittees[committeeIndex] = shuffling.subarray(startOffset, endOffset);
     }
-    committees.push(slotCommittees);
+
+    committees[slot] = slotCommittees;
   }
 
+  return {
+    committees,
+    committeesPerSlot,
+  };
+}
+
+export function computeEpochShuffling(
+  state: BeaconStateAllForks,
+  activeIndices: Uint32Array,
+  epoch: Epoch
+): EpochShuffling {
+  const activeValidatorCount = activeIndices.length;
+  const seed = getSeed(state, epoch, DOMAIN_BEACON_ATTESTER);
+  const shuffling = unshuffleList(activeIndices, seed, SHUFFLE_ROUND_COUNT);
+
   return {
     epoch,
     activeIndices,
     shuffling,
-    committees,
-    committeesPerSlot,
+    ...buildCommitteesFromShuffling(activeValidatorCount, shuffling),
+  };
+}
+
+export async function computeEpochShufflingAsync(
+  state: BeaconStateAllForks,
+  activeIndices: Uint32Array,
+  epoch: Epoch
+): Promise<EpochShuffling> {
+  const activeValidatorCount = activeIndices.length;
+  const seed = getSeed(state, epoch, DOMAIN_BEACON_ATTESTER);
+  const shuffling = await asyncUnshuffleList(activeIndices, seed, SHUFFLE_ROUND_COUNT);
+
+  return {
+    epoch,
+    activeIndices,
+    shuffling,
+    ...buildCommitteesFromShuffling(activeValidatorCount, shuffling),
   };
 }
 

packages/state-transition/src/util/index.ts

@@ -17,7 +17,6 @@ export * from "./genesis.js";
 export * from "./interop.js";
 export * from "./rootCache.js";
 export * from "./seed.js";
-export * from "./shuffle.js";
 export * from "./shufflingDecisionRoot.js";
 export * from "./signatureSets.js";
 export * from "./signingRoot.js";