Use cached context to get tokens in undelegations and redelegations.

x/ccv/provider/keeper/punish_validator.go

@@ -1,8 +1,6 @@
 package keeper
 
 import (
-	"time"
-
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	evidencetypes "github.com/cosmos/cosmos-sdk/x/evidence/types"
 
@@ -44,33 +42,25 @@ func (k Keeper) JailAndTombstoneValidator(ctx sdk.Context, providerAddr types.Pr
 	k.slashingKeeper.Tombstone(ctx, providerAddr.ToSdkConsAddr())
 }
 
-// ComputePowerToSlash computes the power to be slashed based on the tokens in non-matured (based on the
-// provider `now` time) `undelegations` and `redelegations`, as well as the current `power` of the validator
-func (k Keeper) ComputePowerToSlash(now time.Time, undelegations []stakingtypes.UnbondingDelegation,
+// ComputePowerToSlash computes the power to be slashed based on the tokens in non-matured `undelegations` and
+// `redelegations`, as well as the current `power` of the validator
+func (k Keeper) ComputePowerToSlash(ctx sdk.Context, validator stakingtypes.Validator, undelegations []stakingtypes.UnbondingDelegation,
 	redelegations []stakingtypes.Redelegation, power int64, powerReduction sdk.Int,
 ) int64 {
 	// compute the total numbers of tokens currently being undelegated
 	undelegationsInTokens := sdk.NewInt(0)
+
+	cachedCtx, _ := ctx.CacheContext()
 	for _, u := range undelegations {
-		for _, entry := range u.Entries {
-			if entry.IsMature(now) && !entry.OnHold() {
-				//  undelegation no longer eligible for slashing, skip it
-				continue
-			}
-			undelegationsInTokens = undelegationsInTokens.Add(entry.InitialBalance)
-		}
+		amountSlashed := k.stakingKeeper.SlashUnbondingDelegation(cachedCtx, u, 0, sdk.NewDec(1))
+		undelegationsInTokens = undelegationsInTokens.Add(amountSlashed)
 	}
 
 	// compute the total numbers of tokens currently being redelegated
 	redelegationsInTokens := sdk.NewInt(0)
 	for _, r := range redelegations {
-		for _, entry := range r.Entries {
-			if entry.IsMature(now) && !entry.OnHold() {
-				//  redelegation no longer eligible for slashing, skip it
-				continue
-			}
-			redelegationsInTokens = redelegationsInTokens.Add(entry.InitialBalance)
-		}
+		amountSlashed := k.stakingKeeper.SlashRedelegation(cachedCtx, validator, r, 0, sdk.NewDec(1))
+		redelegationsInTokens = redelegationsInTokens.Add(amountSlashed)
 	}
 
 	// The power we pass to staking's keeper `Slash` method is the current power of the validator together with the total
@@ -105,7 +95,7 @@ func (k Keeper) SlashValidator(ctx sdk.Context, providerAddr types.ProviderConsA
 	redelegations := k.stakingKeeper.GetRedelegationsFromSrcValidator(ctx, validator.GetOperator())
 	lastPower := k.stakingKeeper.GetLastValidatorPower(ctx, validator.GetOperator())
 	powerReduction := k.stakingKeeper.PowerReduction(ctx)
-	totalPower := k.ComputePowerToSlash(ctx.BlockHeader().Time, undelegations, redelegations, lastPower, powerReduction)
+	totalPower := k.ComputePowerToSlash(ctx, validator, undelegations, redelegations, lastPower, powerReduction)
 	slashFraction := k.slashingKeeper.SlashFractionDoubleSign(ctx)
 
 	k.stakingKeeper.Slash(ctx, providerAddr.ToSdkConsAddr(), 0, totalPower, slashFraction, stakingtypes.DoubleSign)

x/ccv/provider/keeper/punish_validator_test.go

@@ -176,7 +176,7 @@ func createRedelegation(initialBalances []int64, completionTimes []time.Time) st
 // TestComputePowerToSlash tests that `ComputePowerToSlash` computes the correct power to be slashed based on
 // the tokens in non-mature undelegation and redelegation entries, as well as the current power of the validator
 func TestComputePowerToSlash(t *testing.T) {
-	providerKeeper, ctx, ctrl, _ := testkeeper.GetProviderKeeperAndCtx(t, testkeeper.NewInMemKeeperParams(t))
+	providerKeeper, ctx, ctrl, mocks := testkeeper.GetProviderKeeperAndCtx(t, testkeeper.NewInMemKeeperParams(t))
 	defer ctrl.Finish()
 
 	// undelegation or redelegation entries with completion time `now` have matured
@@ -297,8 +297,39 @@ func TestComputePowerToSlash(t *testing.T) {
 		},
 	}
 
+	pubKey, _ := cryptocodec.FromTmPubKeyInterface(tmtypes.NewMockPV().PrivKey.PubKey())
+	validator, _ := stakingtypes.NewValidator(pubKey.Address().Bytes(), pubKey, stakingtypes.Description{})
+
 	for _, tc := range testCases {
-		actualPower := providerKeeper.ComputePowerToSlash(now,
+		gomock.InOrder(mocks.MockStakingKeeper.EXPECT().
+			SlashUnbondingDelegation(gomock.Any(), gomock.Any(), int64(0), sdk.NewDec(1)).
+			DoAndReturn(
+				func(_ sdk.Context, undelegation stakingtypes.UnbondingDelegation, _ int64, _ sdk.Dec) sdk.Int {
+					sum := sdk.NewInt(0)
+					for _, r := range undelegation.Entries {
+						if r.IsMature(ctx.BlockTime()) {
+							continue
+						}
+						sum = sum.Add(sdk.NewInt(r.InitialBalance.Int64()))
+					}
+					return sum
+				}).AnyTimes(),
+			mocks.MockStakingKeeper.EXPECT().
+				SlashRedelegation(gomock.Any(), gomock.Any(), gomock.Any(), int64(0), sdk.NewDec(1)).
+				DoAndReturn(
+					func(ctx sdk.Context, _ stakingtypes.Validator, redelegation stakingtypes.Redelegation, _ int64, _ sdk.Dec) sdk.Int {
+						sum := sdk.NewInt(0)
+						for _, r := range redelegation.Entries {
+							if r.IsMature(ctx.BlockTime()) {
+								continue
+							}
+							sum = sum.Add(sdk.NewInt(r.InitialBalance.Int64()))
+						}
+						return sum
+					}).AnyTimes(),
+		)
+
+		actualPower := providerKeeper.ComputePowerToSlash(ctx, validator,
 			tc.undelegations, tc.redelegations, tc.power, tc.powerReduction)
 
 		if tc.expectedPower != actualPower {
@@ -386,6 +417,32 @@ func TestSlashValidator(t *testing.T) {
 		mocks.MockStakingKeeper.EXPECT().
 			PowerReduction(ctx).
 			Return(powerReduction),
+		mocks.MockStakingKeeper.EXPECT().
+			SlashUnbondingDelegation(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(
+				func(_ sdk.Context, undelegation stakingtypes.UnbondingDelegation, _ int64, _ sdk.Dec) sdk.Int {
+					sum := sdk.NewInt(0)
+					for _, r := range undelegation.Entries {
+						if r.IsMature(ctx.BlockTime()) {
+							continue
+						}
+						sum = sum.Add(sdk.NewInt(r.InitialBalance.Int64()))
+					}
+					return sum
+				}).AnyTimes(),
+		mocks.MockStakingKeeper.EXPECT().
+			SlashRedelegation(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).
+			DoAndReturn(
+				func(_ sdk.Context, _ stakingtypes.Validator, redelegation stakingtypes.Redelegation, _ int64, _ sdk.Dec) sdk.Int {
+					sum := sdk.NewInt(0)
+					for _, r := range redelegation.Entries {
+						if r.IsMature(ctx.BlockTime()) {
+							continue
+						}
+						sum = sum.Add(sdk.NewInt(r.InitialBalance.Int64()))
+					}
+					return sum
+				}).AnyTimes(),
 		mocks.MockSlashingKeeper.EXPECT().
 			SlashFractionDoubleSign(ctx).
 			Return(slashFraction),

x/ccv/types/expected_keepers.go

@@ -30,6 +30,8 @@ type StakingKeeper interface {
 	// slash the validator and delegators of the validator, specifying offence height, offence power, and slash fraction
 	Jail(sdk.Context, sdk.ConsAddress) // jail a validator
 	Slash(sdk.Context, sdk.ConsAddress, int64, int64, sdk.Dec, stakingtypes.InfractionType)
+	SlashUnbondingDelegation(sdk.Context, stakingtypes.UnbondingDelegation, int64, sdk.Dec) sdk.Int
+	SlashRedelegation(sdk.Context, stakingtypes.Validator, stakingtypes.Redelegation, int64, sdk.Dec) sdk.Int
 	Unjail(ctx sdk.Context, addr sdk.ConsAddress)
 	GetValidator(ctx sdk.Context, addr sdk.ValAddress) (validator stakingtypes.Validator, found bool)
 	IterateLastValidatorPowers(ctx sdk.Context, cb func(addr sdk.ValAddress, power int64) (stop bool))