From 38b92c01714548c7d0d1b341f267abceda0c2330 Mon Sep 17 00:00:00 2001
From: Potuz <potuz@prysmaticlabs.com>
Date: Wed, 11 Sep 2024 14:58:07 -0300
Subject: [PATCH] Avoid allocating in loops (#14439)

* Avoid allocating in loops

* change process slashings

* add changelog

* Kasey's review

* Only marked trie dirty if changed
---
 CHANGELOG.md                                  |  1 +
 beacon-chain/core/epoch/epoch_processing.go   | 21 +++++++++++++------
 .../core/epoch/precompute/slashing.go         | 17 +++++++--------
 3 files changed, 24 insertions(+), 15 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b5858dcd92..12248dd99f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -50,6 +50,7 @@ The format is based on Keep a Changelog, and this project adheres to Semantic Ve
 - Core: Fix process effective balance update to safe copy validator for Electra.
 - `== nil` checks before calling `IsNil()` on interfaces to prevent panics.
 - Core: Fixed slash processing causing extra hashing
+- Core: Fixed extra allocations when processing slashings
 
 ### Security
 
diff --git a/beacon-chain/core/epoch/epoch_processing.go b/beacon-chain/core/epoch/epoch_processing.go
index 499b3baada..e1adcc3902 100644
--- a/beacon-chain/core/epoch/epoch_processing.go
+++ b/beacon-chain/core/epoch/epoch_processing.go
@@ -177,18 +177,27 @@ func ProcessSlashings(st state.BeaconState, slashingMultiplier uint64) (state.Be
 	// below equally.
 	increment := params.BeaconConfig().EffectiveBalanceIncrement
 	minSlashing := math.Min(totalSlashing*slashingMultiplier, totalBalance)
-	err = st.ApplyToEveryValidator(func(idx int, val state.ReadOnlyValidator) (newVal *ethpb.Validator, err error) {
+	bals := st.Balances()
+	changed := false
+	err = st.ReadFromEveryValidator(func(idx int, val state.ReadOnlyValidator) error {
 		correctEpoch := (currentEpoch + exitLength/2) == val.WithdrawableEpoch()
 		if val.Slashed() && correctEpoch {
 			penaltyNumerator := val.EffectiveBalance() / increment * minSlashing
 			penalty := penaltyNumerator / totalBalance * increment
-			if err = helpers.DecreaseBalance(st, primitives.ValidatorIndex(idx), penalty); err != nil {
-				return
-			}
+			bals[idx] = helpers.DecreaseBalanceWithVal(bals[idx], penalty)
+			changed = true
 		}
-		return
+		return nil
 	})
-	return st, err
+	if err != nil {
+		return nil, err
+	}
+	if changed {
+		if err := st.SetBalances(bals); err != nil {
+			return nil, err
+		}
+	}
+	return st, nil
 }
 
 // ProcessEth1DataReset processes updates to ETH1 data votes during epoch processing.
diff --git a/beacon-chain/core/epoch/precompute/slashing.go b/beacon-chain/core/epoch/precompute/slashing.go
index 69bf6d7147..02458f4a90 100644
--- a/beacon-chain/core/epoch/precompute/slashing.go
+++ b/beacon-chain/core/epoch/precompute/slashing.go
@@ -5,9 +5,7 @@ import (
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/core/time"
 	"github.com/prysmaticlabs/prysm/v5/beacon-chain/state"
 	"github.com/prysmaticlabs/prysm/v5/config/params"
-	"github.com/prysmaticlabs/prysm/v5/consensus-types/primitives"
 	"github.com/prysmaticlabs/prysm/v5/math"
-	ethpb "github.com/prysmaticlabs/prysm/v5/proto/prysm/v1alpha1"
 )
 
 // ProcessSlashingsPrecompute processes the slashed validators during epoch processing.
@@ -44,17 +42,18 @@ func ProcessSlashingsPrecompute(s state.BeaconState, pBal *Balance) error {
 	}
 
 	increment := params.BeaconConfig().EffectiveBalanceIncrement
-	validatorFunc := func(idx int, val state.ReadOnlyValidator) (newVal *ethpb.Validator, err error) {
+	bals := s.Balances()
+	validatorFunc := func(idx int, val state.ReadOnlyValidator) error {
 		correctEpoch := epochToWithdraw == val.WithdrawableEpoch()
 		if val.Slashed() && correctEpoch {
 			penaltyNumerator := val.EffectiveBalance() / increment * minSlashing
 			penalty := penaltyNumerator / pBal.ActiveCurrentEpoch * increment
-			if err = helpers.DecreaseBalance(s, primitives.ValidatorIndex(idx), penalty); err != nil {
-				return
-			}
+			bals[idx] = helpers.DecreaseBalanceWithVal(bals[idx], penalty)
 		}
-		return
+		return nil
 	}
-
-	return s.ApplyToEveryValidator(validatorFunc)
+	if err := s.ReadFromEveryValidator(validatorFunc); err != nil {
+		return err
+	}
+	return s.SetBalances(bals)
 }