Add processRewardsAndPenaltiesAllForks

packages/beacon-state-transition/src/allForks/epoch/processRewardsAndPenalties.ts

@@ -0,0 +1,36 @@
+import {CachedBeaconStateAltair, CachedBeaconStatePhase0, CachedBeaconStateAllForks, IEpochProcess} from "../../types";
+import {ForkName, GENESIS_EPOCH} from "@chainsafe/lodestar-params";
+import {getAttestationDeltas as getAttestationDeltasPhase0} from "../../phase0/epoch/getAttestationDeltas";
+import {getRewardsAndPenalties as getRewardsPenaltiesAltair} from "../../altair/epoch/getRewardsAndPenalties";
+
+/**
+ * Iterate over all validator and compute rewards and penalties to apply to balances.
+ *
+ * PERF: Cost = 'proportional' to $VALIDATOR_COUNT. Extra work is done per validator the more status flags
+ * are true, worst case: FLAG_UNSLASHED + FLAG_ELIGIBLE_ATTESTER + FLAG_PREV_*
+ */
+export function processRewardsAndPenaltiesAllForks<T extends CachedBeaconStateAllForks>(
+  fork: ForkName,
+  state: T,
+  epochProcess: IEpochProcess
+): void {
+  // No rewards are applied at the end of `GENESIS_EPOCH` because rewards are for work done in the previous epoch
+  if (epochProcess.currentEpoch === GENESIS_EPOCH) {
+    return;
+  }
+
+  const [rewards, penalties] =
+    fork === ForkName.phase0
+      ? getAttestationDeltasPhase0(state as CachedBeaconStatePhase0, epochProcess)
+      : getRewardsPenaltiesAltair(state as CachedBeaconStateAltair, epochProcess);
+
+  const deltas: number[] = [];
+  for (let i = 0, len = rewards.length; i < len; i++) {
+    deltas.push(rewards[i] - penalties[i]);
+  }
+
+  // important: do not change state one balance at a time
+  // set them all at once, constructing the tree in one go
+  // cache the balances array, too
+  epochProcess.balances = state.balanceList.updateAll(deltas);
+}

packages/beacon-state-transition/src/altair/epoch/balance.ts → packages/beacon-state-transition/src/altair/epoch/getRewardsAndPenalties.ts

@@ -41,10 +41,7 @@ interface IRewardPenaltyItem {
  *   - unslashed:          100%
  *   - eligibleAttester:   98%
  */
-export function getRewardsPenaltiesDeltas(
-  state: CachedBeaconStateAltair,
-  process: IEpochProcess
-): [number[], number[]] {
+export function getRewardsAndPenalties(state: CachedBeaconStateAltair, process: IEpochProcess): [number[], number[]] {
   // TODO: Is there a cheaper way to measure length that going to `state.validators`?
   const validatorCount = state.validators.length;
   const activeIncrements = process.totalActiveStakeByIncrement;

packages/beacon-state-transition/src/altair/epoch/index.ts

@@ -15,7 +15,7 @@ import {processInactivityUpdates} from "./processInactivityUpdates";
 import {processSyncCommitteeUpdates} from "./processSyncCommitteeUpdates";
 
 // For spec tests
-export {getRewardsPenaltiesDeltas} from "./balance";
+export {getRewardsAndPenalties} from "./getRewardsAndPenalties";
 
 export {
   processInactivityUpdates,

packages/beacon-state-transition/src/altair/epoch/processRewardsAndPenalties.ts

@@ -1,6 +1,6 @@
-import {CachedBeaconStateAltair, IEpochProcess} from "../../types";
-import {GENESIS_EPOCH} from "@chainsafe/lodestar-params";
-import {getRewardsPenaltiesDeltas} from "./balance";
+import {CachedBeaconStateAltair, CachedBeaconStateAllForks, IEpochProcess} from "../../types";
+import {ForkName} from "@chainsafe/lodestar-params";
+import {processRewardsAndPenaltiesAllForks} from "../../allForks/epoch/processRewardsAndPenalties";
 
 /**
  * Iterate over all validator and compute rewards and penalties to apply to balances.
@@ -9,38 +9,5 @@ import {getRewardsPenaltiesDeltas} from "./balance";
  * are true, worst case: FLAG_UNSLASHED + FLAG_ELIGIBLE_ATTESTER + FLAG_PREV_*
  */
 export function processRewardsAndPenalties(state: CachedBeaconStateAltair, epochProcess: IEpochProcess): void {
-  if (state.currentShuffling.epoch == GENESIS_EPOCH) {
-    return;
-  }
-
-  const [rewards, penalties] = getRewardsPenaltiesDeltas(state, epochProcess);
-  const deltas = rewards.map((_, i) => Number(rewards[i] - penalties[i]));
-  // important: do not change state one balance at a time
-  // set them all at once, constructing the tree in one go
-  // cache the balances array, too
-  epochProcess.balances = state.balanceList.updateAll(deltas);
+  processRewardsAndPenaltiesAllForks(ForkName.altair, state as CachedBeaconStateAllForks, epochProcess);
 }
-
-// // naive version, leave here for debugging purposes
-// function processRewardsAndPenaltiesNAIVE() {
-//   const flagDeltas = Array.from({length: PARTICIPATION_FLAG_WEIGHTS.length}, (_, flag) =>
-//     getFlagIndexDeltas(state, process, flag)
-//   );
-
-//   const inactivityPenaltyDeltas = getInactivityPenaltyDeltas(state, process);
-//   flagDeltas.push(inactivityPenaltyDeltas);
-
-//   const newBalances = new BigUint64Array(balances.length);
-//   balances.forEach((balance, i) => {
-//     let newBalance = balance;
-//     for (const [rewards, penalties] of flagDeltas) {
-//       const b = newBalance + BigInt(rewards[i] - penalties[i]);
-//       if (b > 0) {
-//         newBalance = b;
-//       } else {
-//         newBalance = BigInt(0);
-//       }
-//     }
-//     newBalances[i] = newBalance;
-//   });
-// }

packages/beacon-state-transition/src/phase0/epoch/processRewardsAndPenalties.ts

@@ -1,7 +1,6 @@
-import {GENESIS_EPOCH} from "@chainsafe/lodestar-params";
-
-import {CachedBeaconStatePhase0, IEpochProcess} from "../../types";
-import {getAttestationDeltas} from "./getAttestationDeltas";
+import {ForkName} from "@chainsafe/lodestar-params";
+import {processRewardsAndPenaltiesAllForks} from "../../allForks/epoch/processRewardsAndPenalties";
+import {CachedBeaconStatePhase0, CachedBeaconStateAllForks, IEpochProcess} from "../../types";
 
 /**
  * Iterate over all validator and compute rewards and penalties to apply to balances.
@@ -10,14 +9,5 @@ import {getAttestationDeltas} from "./getAttestationDeltas";
  * are true, worst case: FLAG_UNSLASHED + FLAG_ELIGIBLE_ATTESTER + FLAG_PREV_*
  */
 export function processRewardsAndPenalties(state: CachedBeaconStatePhase0, epochProcess: IEpochProcess): void {
-  // No rewards are applied at the end of `GENESIS_EPOCH` because rewards are for work done in the previous epoch
-  if (epochProcess.currentEpoch === GENESIS_EPOCH) {
-    return;
-  }
-  const [rewards, penalties] = getAttestationDeltas(state, epochProcess);
-  const deltas = rewards.map((_, i) => rewards[i] - penalties[i]);
-  // important: do not change state one balance at a time
-  // set them all at once, constructing the tree in one go
-  // cache the balances array, too
-  epochProcess.balances = state.balanceList.updateAll(deltas);
+  processRewardsAndPenaltiesAllForks(ForkName.phase0, state as CachedBeaconStateAllForks, epochProcess);
 }

packages/lodestar/test/spec/allForks/rewards.ts

@@ -88,7 +88,7 @@ export function rewardsAltair(fork: ForkName): void {
         //   + set all inactivityScores to zero
         // - To get inactivity_penalty_deltas set TIMELY_HEAD_FLAG_INDEX | TIMELY_SOURCE_FLAG_INDEX to false
         //   + set PARTICIPATION_FLAG_WEIGHTS[TIMELY_TARGET_FLAG_INDEX] to zero
-        return altair.getRewardsPenaltiesDeltas(state, epochProcess);
+        return altair.getRewardsAndPenalties(state, epochProcess);
       },
       {
         inputTypes: inputTypeSszTreeBacked,