ProcessRewardsAndPenaltiesPrecompute taking in rewards funcs (#8737)

Co-authored-by: prylabs-bulldozer[bot] <58059840+prylabs-bulldozer[bot]@users.noreply.github.com>
prysmaticlabs · Apr 13, 2021 · 8b494fb · 8b494fb
1 parent 6aa1297
commit 8b494fb
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 10 deletions.
diff --git a/beacon-chain/core/epoch/precompute/reward_penalty.go b/beacon-chain/core/epoch/precompute/reward_penalty.go
@@ -9,12 +9,17 @@ import (
 	"github.com/prysmaticlabs/prysm/shared/params"
 )
 
+type attesterRewardsFunc func(iface.ReadOnlyBeaconState, *Balance, []*Validator) ([]uint64, []uint64, error)
+type proposerRewardsFunc func(iface.ReadOnlyBeaconState, *Balance, []*Validator) ([]uint64, error)
+
 // ProcessRewardsAndPenaltiesPrecompute processes the rewards and penalties of individual validator.
 // This is an optimized version by passing in precomputed validator attesting records and and total epoch balances.
 func ProcessRewardsAndPenaltiesPrecompute(
 	state iface.BeaconState,
 	pBal *Balance,
 	vp []*Validator,
+	attRewardsFunc attesterRewardsFunc,
+	proRewardsFunc proposerRewardsFunc,
 ) (iface.BeaconState, error) {
 	// Can't process rewards and penalties in genesis epoch.
 	if helpers.CurrentEpoch(state) == 0 {
@@ -27,13 +32,13 @@ func ProcessRewardsAndPenaltiesPrecompute(
 		return state, errors.New("precomputed registries not the same length as state registries")
 	}
 
-	attsRewards, attsPenalties, err := AttestationsDelta(state, pBal, vp)
+	attsRewards, attsPenalties, err := attRewardsFunc(state, pBal, vp)
 	if err != nil {
-		return nil, errors.Wrap(err, "could not get attestation delta")
+		return nil, errors.Wrap(err, "could not get attester attestation delta")
 	}
-	proposerRewards, err := ProposersDelta(state, pBal, vp)
+	proposerRewards, err := proRewardsFunc(state, pBal, vp)
 	if err != nil {
-		return nil, errors.Wrap(err, "could not get attestation delta")
+		return nil, errors.Wrap(err, "could not get proposer attestation delta")
 	}
 	validatorBals := state.Balances()
 	for i := 0; i < numOfVals; i++ {

diff --git a/beacon-chain/core/epoch/precompute/reward_penalty_test.go b/beacon-chain/core/epoch/precompute/reward_penalty_test.go
@@ -42,7 +42,7 @@ func TestProcessRewardsAndPenaltiesPrecompute(t *testing.T) {
 	vp, bp, err = ProcessAttestations(context.Background(), beaconState, vp, bp)
 	require.NoError(t, err)
 
-	processedState, err := ProcessRewardsAndPenaltiesPrecompute(beaconState, bp, vp)
+	processedState, err := ProcessRewardsAndPenaltiesPrecompute(beaconState, bp, vp, AttestationsDelta, ProposersDelta)
 	require.NoError(t, err)
 	beaconState, ok := processedState.(*stateV0.BeaconState)
 	require.Equal(t, true, ok)

diff --git a/beacon-chain/core/epoch/spectest/rewards_and_penalties_test.go b/beacon-chain/core/epoch/spectest/rewards_and_penalties_test.go
@@ -34,7 +34,7 @@ func processRewardsAndPenaltiesPrecomputeWrapper(t *testing.T, st iface.BeaconSt
 	vp, bp, err = precompute.ProcessAttestations(ctx, st, vp, bp)
 	require.NoError(t, err)
 
-	st, err = precompute.ProcessRewardsAndPenaltiesPrecompute(st, bp, vp)
+	st, err = precompute.ProcessRewardsAndPenaltiesPrecompute(st, bp, vp, precompute.AttestationsDelta, precompute.ProposersDelta)
 	require.NoError(t, err, "Could not process reward")
 
 	return st, nil

diff --git a/beacon-chain/core/state/transition.go b/beacon-chain/core/state/transition.go
@@ -407,7 +407,7 @@ func ProcessEpochPrecompute(ctx context.Context, state iface.BeaconState) (iface
 		return nil, errors.Wrap(err, "could not process justification")
 	}
 
-	state, err = precompute.ProcessRewardsAndPenaltiesPrecompute(state, bp, vp)
+	state, err = precompute.ProcessRewardsAndPenaltiesPrecompute(state, bp, vp, precompute.AttestationsDelta, precompute.ProposersDelta)
 	if err != nil {
 		return nil, errors.Wrap(err, "could not process rewards and penalties")
 	}

diff --git a/beacon-chain/rpc/beacon/validators.go b/beacon-chain/rpc/beacon/validators.go
@@ -668,7 +668,7 @@ func (bs *Server) GetValidatorPerformance(
 	if err != nil {
 		return nil, err
 	}
-	headState, err = precompute.ProcessRewardsAndPenaltiesPrecompute(headState, bp, vp)
+	headState, err = precompute.ProcessRewardsAndPenaltiesPrecompute(headState, bp, vp, precompute.AttestationsDelta, precompute.ProposersDelta)
 	if err != nil {
 		return nil, err
 	}

diff --git a/beacon-chain/rpc/beacon/validators_test.go b/beacon-chain/rpc/beacon/validators_test.go
@@ -1783,7 +1783,7 @@ func TestGetValidatorPerformance_Indices(t *testing.T) {
 	require.NoError(t, err)
 	vp, bp, err = precompute.ProcessAttestations(ctx, c, vp, bp)
 	require.NoError(t, err)
-	_, err = precompute.ProcessRewardsAndPenaltiesPrecompute(c, bp, vp)
+	_, err = precompute.ProcessRewardsAndPenaltiesPrecompute(c, bp, vp, precompute.AttestationsDelta, precompute.ProposersDelta)
 	require.NoError(t, err)
 	farFuture := params.BeaconConfig().FarFutureSlot
 	want := &ethpb.ValidatorPerformanceResponse{
@@ -1856,7 +1856,7 @@ func TestGetValidatorPerformance_IndicesPubkeys(t *testing.T) {
 	require.NoError(t, err)
 	vp, bp, err = precompute.ProcessAttestations(ctx, c, vp, bp)
 	require.NoError(t, err)
-	_, err = precompute.ProcessRewardsAndPenaltiesPrecompute(c, bp, vp)
+	_, err = precompute.ProcessRewardsAndPenaltiesPrecompute(c, bp, vp, precompute.AttestationsDelta, precompute.ProposersDelta)
 	require.NoError(t, err)
 	farFuture := params.BeaconConfig().FarFutureSlot
 	want := &ethpb.ValidatorPerformanceResponse{