perf: rust backed PTC sampling

packages/state-transition/package.json

@@ -66,7 +66,7 @@
     "@chainsafe/persistent-ts": "^1.0.0",
     "@chainsafe/pubkey-index-map": "^3.0.0",
     "@chainsafe/ssz": "^1.4.0",
-    "@chainsafe/swap-or-not-shuffle": "^1.2.1",
+    "@chainsafe/swap-or-not-shuffle": "file:../../chainsafe-swap-or-not-shuffle-1.2.1-local.tgz",
     "@lodestar/config": "workspace:^",
     "@lodestar/params": "workspace:^",
     "@lodestar/types": "workspace:^",

packages/state-transition/src/epoch/processPtcWindow.ts

@@ -21,7 +21,7 @@ export function processPtcWindow(state: CachedBeaconStateGloas, cache: EpochTran
   const newNextPayloadTimelinessCommittees = computePayloadTimelinessCommitteesForEpoch(
     state,
     nextEpoch,
-    nextEpochShuffling.committees,
+    nextEpochShuffling,
     state.epochCtx.effectiveBalanceIncrements
   );
 

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

@@ -191,7 +191,7 @@ export function initializePtcWindow(state: CachedBeaconStateFulu): Uint32Array[]
       ...computePayloadTimelinessCommitteesForEpoch(
         state,
         epoch,
-        shuffling.committees,
+        shuffling,
         state.epochCtx.effectiveBalanceIncrements
       )
     );

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

@@ -1,5 +1,7 @@
 import {digest} from "@chainsafe/as-sha256";
 import {
+  computePtcIndices,
+  computePtcIndicesForEpoch,
   computeProposerIndex as nativeComputeProposerIndex,
   computeSyncCommitteeIndices as nativeComputeSyncCommitteeIndices,
 } from "@chainsafe/swap-or-not-shuffle";
@@ -23,6 +25,7 @@ import {assert, bytesToBigInt, bytesToInt, intToBytes} from "@lodestar/utils";
 import {EffectiveBalanceIncrements} from "../cache/effectiveBalanceIncrements.js";
 import {BeaconStateAllForks, CachedBeaconStateAllForks} from "../types.js";
 import {computeEpochAtSlot, computeStartSlotAtEpoch} from "./epoch.js";
+import {EpochShuffling} from "./epochShuffling.js";
 
 /**
  * Compute proposer indices for an epoch
@@ -269,9 +272,12 @@ export function getNextSyncCommitteeIndices(
 }
 
 /**
- * Compute PTC for all slots in an epoch eagerly.
+ * Naive JS version of `computePayloadTimelinessCommitteesForEpoch`.
+ * Used to verify the optimized Rust-backed version. Not for production use.
+ *
+ * SLOW CODE - 🐢
  */
-export function computePayloadTimelinessCommitteesForEpoch(
+export function naiveComputePayloadTimelinessCommitteesForEpoch(
   state: BeaconStateAllForks,
   epoch: number,
   committees: Uint32Array[][],
@@ -293,15 +299,60 @@ export function computePayloadTimelinessCommitteesForEpoch(
     slotSeedView.setUint32(epochSeed.length + 4, 0, true);
     const slotSeed = digest(slotSeedInput);
 
-    result[i] = computePayloadTimelinessCommitteeForSlot(slotSeed, committees[i], effectiveBalanceIncrements);
+    result[i] = naiveComputePayloadTimelinessCommitteeForSlot(slotSeed, committees[i], effectiveBalanceIncrements);
   }
   return result;
 }
 
 /**
- * Compute PTC for a single slot.
+ * Compute PTC for all slots in an epoch eagerly.
  */
-export function computePayloadTimelinessCommitteeForSlot(
+export function computePayloadTimelinessCommitteesForEpoch(
+  state: BeaconStateAllForks,
+  epoch: number,
+  epochShuffling: EpochShuffling,
+  effectiveBalanceIncrements: EffectiveBalanceIncrements
+): Uint32Array[] {
+  const epochSeed = getSeed(state, epoch, DOMAIN_PTC_ATTESTER);
+  const startSlot = epoch * SLOTS_PER_EPOCH;
+  const {committees, shuffling} = epochShuffling;
+
+  const slotOffsets = new Uint32Array(SLOTS_PER_EPOCH + 1);
+  for (let i = 0; i < SLOTS_PER_EPOCH; i++) {
+    let slotLen = 0;
+    for (const c of committees[i]) slotLen += c.length;
+    slotOffsets[i + 1] = slotOffsets[i] + slotLen;
+  }
+  if (shuffling.length === 0) {
+    throw Error("Validator indices must not be empty");
+  }
+
+  const flat = computePtcIndicesForEpoch(
+    epochSeed,
+    startSlot,
+    SLOTS_PER_EPOCH,
+    shuffling,
+    slotOffsets,
+    effectiveBalanceIncrements,
+    PTC_SIZE,
+    MAX_EFFECTIVE_BALANCE_ELECTRA,
+    EFFECTIVE_BALANCE_INCREMENT
+  );
+
+  const result = new Array<Uint32Array>(SLOTS_PER_EPOCH);
+  for (let i = 0; i < SLOTS_PER_EPOCH; i++) {
+    result[i] = flat.subarray(i * PTC_SIZE, (i + 1) * PTC_SIZE);
+  }
+  return result;
+}
+
+/**
+ * Naive JS version of `computePayloadTimelinessCommitteeForSlot`.
+ * Used to verify the optimized Rust-backed version. Not for production use.
+ *
+ * SLOW CODE - 🐢
+ */
+export function naiveComputePayloadTimelinessCommitteeForSlot(
   slotSeed: Uint8Array,
   slotCommittees: Uint32Array[],
   effectiveBalanceIncrements: EffectiveBalanceIncrements
@@ -317,6 +368,31 @@ export function computePayloadTimelinessCommitteeForSlot(
   return computePayloadTimelinessCommitteeIndices(effectiveBalanceIncrements, allIndices, slotSeed);
 }
 
+/**
+ * Compute PTC for a single slot.
+ */
+export function computePayloadTimelinessCommitteeForSlot(
+  slotSeed: Uint8Array,
+  slotCommittees: Uint32Array[],
+  effectiveBalanceIncrements: EffectiveBalanceIncrements
+): Uint32Array {
+  const totalLen = slotCommittees.reduce((sum, c) => sum + c.length, 0);
+  const allIndices = new Uint32Array(totalLen);
+  let offset = 0;
+  for (const c of slotCommittees) {
+    allIndices.set(c, offset);
+    offset += c.length;
+  }
+  return computePtcIndices(
+    slotSeed,
+    allIndices,
+    effectiveBalanceIncrements,
+    PTC_SIZE,
+    MAX_EFFECTIVE_BALANCE_ELECTRA,
+    EFFECTIVE_BALANCE_INCREMENT
+  );
+}
+
 /**
  * Optimized version of PTC indices computation.
  * Avoids BigInt conversions and uses DataView for efficient byte reading.

packages/state-transition/test/perf/util/seed.test.ts

@@ -1,11 +1,15 @@
 import {bench, describe} from "@chainsafe/benchmark";
 import {ForkSeq} from "@lodestar/params";
 import {fromHex} from "@lodestar/utils";
-import {generatePerfTestCachedStateAltair} from "../../../src/testUtils/util.js";
+import {generatePerfTestCachedStateAltair, generatePerfTestCachedStateElectra} from "../../../src/testUtils/util.js";
 import {
+  computePayloadTimelinessCommitteeForSlot,
+  computePayloadTimelinessCommitteesForEpoch,
   computeProposerIndex,
   computeShuffledIndex,
   getNextSyncCommitteeIndices,
+  naiveComputePayloadTimelinessCommitteeForSlot,
+  naiveComputePayloadTimelinessCommitteesForEpoch,
   naiveComputeProposerIndex,
   naiveGetNextSyncCommitteeIndices,
 } from "../../../src/util/seed.js";
@@ -89,3 +93,89 @@ describe("computeShuffledIndex", () => {
     });
   }
 });
+
+describe("computePayloadTimelinessCommitteeForSlot - pure TS vs Rust magic (250k-1M validators)", () => {
+  for (const vc of [250_000, 1_000_000]) {
+    const seed = new Uint8Array(32).fill(1);
+    const indices = new Uint32Array(Array.from({length: vc}, (_, i) => i));
+    const effectiveBalanceIncrements = new Uint16Array(vc).fill(32);
+    const slotCommittees = [indices]; // single committee spanning all validators
+
+    bench({
+      id: `naive TS - naiveComputePayloadTimelinessCommitteeForSlot - ${vc} validators`,
+      fn: () => {
+        naiveComputePayloadTimelinessCommitteeForSlot(seed, slotCommittees, effectiveBalanceIncrements);
+      },
+    });
+
+    bench({
+      id: `(uses native Rust) -computePayloadTimelinessCommitteeForSlot - ${vc} validators`,
+      fn: () => {
+        computePayloadTimelinessCommitteeForSlot(seed, slotCommittees, effectiveBalanceIncrements);
+      },
+    });
+  }
+});
+
+describe("computePayloadTimelinessCommitteesForEpoch - pure TS vs  Rust magic (250k -1M validators)", () => {
+  for (const vc of [250_000, 1_000_000]) {
+    const cachedState = generatePerfTestCachedStateElectra({goBackOneSlot: false, vc});
+    const {epochCtx} = cachedState;
+    const epoch = epochCtx.epoch;
+    const {effectiveBalanceIncrements} = epochCtx;
+
+    // eslint-disable-next-line no-console
+    console.log(`[vc=${vc}] effectiveBalanceIncrements[0]=${effectiveBalanceIncrements[0]}`);
+
+    const naiveResult = naiveComputePayloadTimelinessCommitteesForEpoch(
+      cachedState,
+      epoch,
+      epochCtx.currentShuffling.committees,
+      effectiveBalanceIncrements
+    );
+    const rustResult = computePayloadTimelinessCommitteesForEpoch(
+      cachedState,
+      epoch,
+      epochCtx.currentShuffling,
+      effectiveBalanceIncrements
+    );
+    for (let i = 0; i < naiveResult.length; i++) {
+      const naive = naiveResult[i];
+      const rust = rustResult[i];
+      if (naive.length !== rust.length) {
+        throw new Error(`PTC length mismatch at slot ${i} (vc=${vc}): naive=${naive.length} rust=${rust.length}`);
+      }
+      for (let j = 0; j < naive.length; j++) {
+        if (naive[j] !== rust[j]) {
+          throw new Error(
+            `PTC index mismatch at slot ${i} position ${j} (vc=${vc}): naive=${naive[j]} rust=${rust[j]}`
+          );
+        }
+      }
+    }
+
+    bench({
+      id: `naive TS - naiveComputePayloadTimelinessCommitteesForEpoch - ${vc} validators`,
+      fn: () => {
+        naiveComputePayloadTimelinessCommitteesForEpoch(
+          cachedState,
+          epoch,
+          epochCtx.currentShuffling.committees,
+          effectiveBalanceIncrements
+        );
+      },
+    });
+
+    bench({
+      id: `(uses native Rust) -computePayloadTimelinessCommitteesForEpoch - ${vc} validators`,
+      fn: () => {
+        computePayloadTimelinessCommitteesForEpoch(
+          cachedState,
+          epoch,
+          epochCtx.currentShuffling,
+          effectiveBalanceIncrements
+        );
+      },
+    });
+  }
+});