fix: cached balances in epoch transition (#7018)

* fix: update cached balances in epoch transition

* fix: more comments in processEffectiveBalanceUpdates()

packages/state-transition/src/block/processAttestationsAltair.ts

@@ -75,6 +75,7 @@ export function processAttestationsAltair(
     // For each participant, update their participation
     // In epoch processing, this participation info is used to calculate balance updates
     let totalBalanceIncrementsWithWeight = 0;
+    const validators = state.validators;
     for (const index of attestingIndices) {
       const flags = epochParticipation.get(index);
 
@@ -104,7 +105,7 @@ export function processAttestationsAltair(
       // TODO: describe issue. Compute progressive target balances
       // When processing each attestation, increase the cummulative target balance. Only applies post-altair
       if ((flagsNewSet & TIMELY_TARGET) === TIMELY_TARGET) {
-        const validator = state.validators.getReadonly(index);
+        const validator = validators.getReadonly(index);
         if (!validator.slashed) {
           if (inCurrentEpoch) {
             epochCtx.currentTargetUnslashedBalanceIncrements += effectiveBalanceIncrements[index];

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

@@ -133,15 +133,15 @@ export function processEpoch(
       const timer = metrics?.epochTransitionStepTime.startTimer({
         step: EpochTransitionStep.processPendingBalanceDeposits,
       });
-      processPendingBalanceDeposits(stateElectra);
+      processPendingBalanceDeposits(stateElectra, cache);
       timer?.();
     }
 
     {
       const timer = metrics?.epochTransitionStepTime.startTimer({
         step: EpochTransitionStep.processPendingConsolidations,
       });
-      processPendingConsolidations(stateElectra);
+      processPendingConsolidations(stateElectra, cache);
       timer?.();
     }
   }

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

@@ -39,12 +39,10 @@ export function processEffectiveBalanceUpdates(
 
   // update effective balances with hysteresis
 
-  // epochTransitionCache.balances is set in processRewardsAndPenalties(), so it's recycled here for performance.
-  // It defaults to `state.balances.getAll()` to make Typescript happy and for spec tests
-  const balances = state.balances.getAll();
-
-  // TODO: (@matthewkeil) This was causing additional failures but should not.  Check the EpochTransitionCache for why
-  // const balances = cache.balances ?? state.balances.getAll();
+  // epochTransitionCache.balances is initialized in processRewardsAndPenalties()
+  // and updated in processPendingBalanceDeposits() and processPendingConsolidations()
+  // so it's recycled here for performance.
+  const balances = cache.balances ?? state.balances.getAll();
 
   for (let i = 0, len = balances.length; i < len; i++) {
     const balance = balances[i];

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

@@ -1,5 +1,5 @@
 import {FAR_FUTURE_EPOCH} from "@lodestar/params";
-import {CachedBeaconStateElectra} from "../types.js";
+import {CachedBeaconStateElectra, EpochTransitionCache} from "../types.js";
 import {increaseBalance} from "../util/balance.js";
 import {getActivationExitChurnLimit} from "../util/validator.js";
 import {getCurrentEpoch} from "../util/epoch.js";
@@ -13,24 +13,29 @@ import {getCurrentEpoch} from "../util/epoch.js";
  *
  * TODO Electra: Update ssz library to support batch push to `pendingBalanceDeposits`
  */
-export function processPendingBalanceDeposits(state: CachedBeaconStateElectra): void {
+export function processPendingBalanceDeposits(state: CachedBeaconStateElectra, cache: EpochTransitionCache): void {
   const availableForProcessing = state.depositBalanceToConsume + BigInt(getActivationExitChurnLimit(state.epochCtx));
   const currentEpoch = getCurrentEpoch(state);
   let processedAmount = 0n;
   let nextDepositIndex = 0;
   const depositsToPostpone = [];
+  const validators = state.validators;
+  const cachedBalances = cache.balances;
 
   for (const deposit of state.pendingBalanceDeposits.getAllReadonly()) {
     const {amount, index: depositIndex} = deposit;
-    const validator = state.validators.get(depositIndex);
+    const validator = validators.getReadonly(depositIndex);
 
     // Validator is exiting, postpone the deposit until after withdrawable epoch
     if (validator.exitEpoch < FAR_FUTURE_EPOCH) {
       if (currentEpoch <= validator.withdrawableEpoch) {
         depositsToPostpone.push(deposit);
       } else {
         // Deposited balance will never become active. Increase balance but do not consume churn
-        increaseBalance(state, deposit.index, Number(amount));
+        increaseBalance(state, depositIndex, Number(amount));
+        if (cachedBalances) {
+          cachedBalances[depositIndex] += Number(amount);
+        }
       }
     } else {
       // Validator is not exiting, attempt to process deposit
@@ -39,7 +44,10 @@ export function processPendingBalanceDeposits(state: CachedBeaconStateElectra):
         break;
       } else {
         // Deposit fits in the churn, process it. Increase balance and consume churn.
-        increaseBalance(state, deposit.index, Number(amount));
+        increaseBalance(state, depositIndex, Number(amount));
+        if (cachedBalances) {
+          cachedBalances[depositIndex] += Number(amount);
+        }
         processedAmount = processedAmount + amount;
       }
     }

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

@@ -1,4 +1,4 @@
-import {CachedBeaconStateElectra} from "../types.js";
+import {CachedBeaconStateElectra, EpochTransitionCache} from "../types.js";
 import {decreaseBalance, increaseBalance} from "../util/balance.js";
 import {getActiveBalance} from "../util/validator.js";
 import {switchToCompoundingValidator} from "../util/electra.js";
@@ -16,12 +16,14 @@ import {switchToCompoundingValidator} from "../util/electra.js";
  * Dequeue all processed consolidations from `state.pendingConsolidation`
  *
  */
-export function processPendingConsolidations(state: CachedBeaconStateElectra): void {
+export function processPendingConsolidations(state: CachedBeaconStateElectra, cache: EpochTransitionCache): void {
   let nextPendingConsolidation = 0;
+  const validators = state.validators;
+  const cachedBalances = cache.balances;
 
   for (const pendingConsolidation of state.pendingConsolidations.getAllReadonly()) {
     const {sourceIndex, targetIndex} = pendingConsolidation;
-    const sourceValidator = state.validators.getReadonly(sourceIndex);
+    const sourceValidator = validators.getReadonly(sourceIndex);
 
     if (sourceValidator.slashed) {
       nextPendingConsolidation++;
@@ -37,6 +39,10 @@ export function processPendingConsolidations(state: CachedBeaconStateElectra): v
     const activeBalance = getActiveBalance(state, sourceIndex);
     decreaseBalance(state, sourceIndex, activeBalance);
     increaseBalance(state, targetIndex, activeBalance);
+    if (cachedBalances) {
+      cachedBalances[sourceIndex] -= activeBalance;
+      cachedBalances[targetIndex] += activeBalance;
+    }
 
     nextPendingConsolidation++;
   }

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

@@ -23,11 +23,10 @@ export function processSyncCommitteeUpdates(fork: ForkSeq, state: CachedBeaconSt
       activeValidatorIndices,
       effectiveBalanceIncrements
     );
+    const validators = state.validators;
 
     // Using the index2pubkey cache is slower because it needs the serialized pubkey.
-    const nextSyncCommitteePubkeys = nextSyncCommitteeIndices.map(
-      (index) => state.validators.getReadonly(index).pubkey
-    );
+    const nextSyncCommitteePubkeys = nextSyncCommitteeIndices.map((index) => validators.getReadonly(index).pubkey);
 
     // Rotate syncCommittee in state
     state.currentSyncCommittee = state.nextSyncCommittee;