Revert "Simplify Effective Balance Calculation (#8743)" (#8844)

* Revert "Simplify Effective Balance Calculation (#8743)"

This reverts commit 45d2df1.

* add flag

beacon-chain/core/epoch/BUILD.bazel

@@ -16,6 +16,7 @@ go_library(
         "//beacon-chain/state/stateV0:go_default_library",
         "//proto/beacon/p2p/v1:go_default_library",
         "//shared/attestationutil:go_default_library",
+        "//shared/featureconfig:go_default_library",
         "//shared/mathutil:go_default_library",
         "//shared/params:go_default_library",
         "@com_github_pkg_errors//:go_default_library",

beacon-chain/core/epoch/epoch_processing.go

@@ -17,6 +17,7 @@ import (
 	"github.com/prysmaticlabs/prysm/beacon-chain/state/stateV0"
 	pb "github.com/prysmaticlabs/prysm/proto/beacon/p2p/v1"
 	"github.com/prysmaticlabs/prysm/shared/attestationutil"
+	"github.com/prysmaticlabs/prysm/shared/featureconfig"
 	"github.com/prysmaticlabs/prysm/shared/mathutil"
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
@@ -257,18 +258,40 @@ func ProcessEffectiveBalanceUpdates(state iface.BeaconState) (iface.BeaconState,
 		balance := bals[idx]
 
 		if balance+downwardThreshold < val.EffectiveBalance || val.EffectiveBalance+upwardThreshold < balance {
-			effectiveBal := maxEffBalance
-			if effectiveBal > balance-balance%effBalanceInc {
-				effectiveBal = balance - balance%effBalanceInc
+			newVal := stateV0.CopyValidator(val)
+			newVal.EffectiveBalance = maxEffBalance
+			if newVal.EffectiveBalance > balance-balance%effBalanceInc {
+				newVal.EffectiveBalance = balance - balance%effBalanceInc
 			}
-			if effectiveBal != val.EffectiveBalance {
-				newVal := stateV0.CopyValidator(val)
-				newVal.EffectiveBalance = effectiveBal
-				return true, newVal, nil
+			return true, newVal, nil
+		}
+		return false, val, nil
+	}
+
+	if featureconfig.Get().EnableOptimizedBalanceUpdate {
+		validatorFunc = func(idx int, val *ethpb.Validator) (bool, *ethpb.Validator, error) {
+			if val == nil {
+				return false, nil, fmt.Errorf("validator %d is nil in state", idx)
+			}
+			if idx >= len(bals) {
+				return false, nil, fmt.Errorf("validator index exceeds validator length in state %d >= %d", idx, len(state.Balances()))
+			}
+			balance := bals[idx]
+
+			if balance+downwardThreshold < val.EffectiveBalance || val.EffectiveBalance+upwardThreshold < balance {
+				effectiveBal := maxEffBalance
+				if effectiveBal > balance-balance%effBalanceInc {
+					effectiveBal = balance - balance%effBalanceInc
+				}
+				if effectiveBal != val.EffectiveBalance {
+					newVal := stateV0.CopyValidator(val)
+					newVal.EffectiveBalance = effectiveBal
+					return true, newVal, nil
+				}
+				return false, val, nil
 			}
 			return false, val, nil
 		}
-		return false, val, nil
 	}
 
 	if err := state.ApplyToEveryValidator(validatorFunc); err != nil {

shared/featureconfig/config.go

@@ -48,6 +48,7 @@ type Flags struct {
 	DisableAttestingHistoryDBCache     bool // DisableAttestingHistoryDBCache for the validator client increases disk reads/writes.
 	UpdateHeadTimely                   bool // UpdateHeadTimely updates head right after state transition.
 	ProposerAttsSelectionUsingMaxCover bool // ProposerAttsSelectionUsingMaxCover enables max-cover algorithm when selecting attestations for proposing.
+	EnableOptimizedBalanceUpdate       bool // EnableOptimizedBalanceUpdate uses an updated method of performing balance updates.
 
 	// Logging related toggles.
 	DisableGRPCConnectionLogs bool // Disables logging when a new grpc client has connected.
@@ -192,6 +193,10 @@ func ConfigureBeaconChain(ctx *cli.Context) {
 		log.WithField(disableProposerAttsSelectionUsingMaxCover.Name, disableProposerAttsSelectionUsingMaxCover.Usage).Warn(enabledFeatureFlag)
 		cfg.ProposerAttsSelectionUsingMaxCover = false
 	}
+	if ctx.Bool(enableOptimizedBalanceUpdate.Name) {
+		log.WithField(enableOptimizedBalanceUpdate.Name, enableOptimizedBalanceUpdate.Usage).Warn(enabledFeatureFlag)
+		cfg.EnableOptimizedBalanceUpdate = true
+	}
 	Init(cfg)
 }
 

shared/featureconfig/flags.go

@@ -118,13 +118,18 @@ var (
 		Name:  "enable-slashing-protection-pruning",
 		Usage: "Enables the pruning of the validator client's slashing protectin database",
 	}
+	enableOptimizedBalanceUpdate = &cli.BoolFlag{
+		Name:  "enable-optimized-balance-update",
+		Usage: "Enables the optimized method of updating validator balances.",
+	}
 )
 
 // devModeFlags holds list of flags that are set when development mode is on.
 var devModeFlags = []cli.Flag{
 	enableLargerGossipHistory,
 	forceOptMaxCoverAggregationStategy,
 	updateHeadTimely,
+	enableOptimizedBalanceUpdate,
 }
 
 // ValidatorFlags contains a list of all the feature flags that apply to the validator client.
@@ -173,6 +178,7 @@ var BeaconChainFlags = append(deprecatedFlags, []cli.Flag{
 	forceOptMaxCoverAggregationStategy,
 	updateHeadTimely,
 	disableProposerAttsSelectionUsingMaxCover,
+	enableOptimizedBalanceUpdate,
 }...)
 
 // E2EBeaconChainFlags contains a list of the beacon chain feature flags to be tested in E2E.