chore: moved simple sum sub to store file

x/incentive/keeper/btc_staking_gauge_test.go

@@ -24,14 +24,14 @@ func FuzzRewardBTCStaking(f *testing.F) {
 		// mock bank keeper
 		bankKeeper := types.NewMockBankKeeper(ctrl)
 
-		// create incentive keeper
-		keeper, ctx := testkeeper.IncentiveKeeper(t, bankKeeper, nil, nil)
+		// create incentive k
+		k, ctx := testkeeper.IncentiveKeeper(t, bankKeeper, nil, nil)
 		height := datagen.RandomInt(r, 1000)
 		ctx = datagen.WithCtxHeight(ctx, height)
 
 		// set a random gauge
 		gauge := datagen.GenRandomGauge(r)
-		keeper.SetBTCStakingGauge(ctx, height, gauge)
+		k.SetBTCStakingGauge(ctx, height, gauge)
 
 		// generate a random voting power distribution cache
 		dc, err := datagen.GenRandomVotingPowerDistCache(r, 100)
@@ -57,7 +57,7 @@ func FuzzRewardBTCStaking(f *testing.F) {
 			fpAddr := fp.GetAddress()
 
 			for _, btcDel := range fp.BtcDels {
-				err := keeper.BtcDelegationActivated(ctx, fpAddr, btcDel.GetAddress(), btcDel.TotalSat)
+				err := k.BtcDelegationActivated(ctx, fpAddr, btcDel.GetAddress(), btcDel.TotalSat)
 				require.NoError(t, err)
 
 				btcDelPortion := fp.GetBTCDelPortion(btcDel)
@@ -70,21 +70,21 @@ func FuzzRewardBTCStaking(f *testing.F) {
 		}
 
 		// distribute rewards in the gauge to finality providers/delegations
-		keeper.RewardBTCStaking(ctx, height, dc)
+		k.RewardBTCStaking(ctx, height, dc)
 
 		for _, fp := range dc.FinalityProviders {
 			fpAddr := fp.GetAddress()
 			for _, btcDel := range fp.BtcDels {
 				delAddr := btcDel.GetAddress()
-				delRwd, err := keeper.GetBTCDelegationRewardsTracker(ctx, fpAddr, delAddr)
+				delRwd, err := k.GetBTCDelegationRewardsTracker(ctx, fpAddr, delAddr)
 				require.NoError(t, err)
 				require.Equal(t, delRwd.TotalActiveSat.Uint64(), btcDel.TotalSat)
 
 				// makes sure the rewards added reach the delegation gauge
-				err = keeper.SendBtcDelegationRewardsToGauge(ctx, fpAddr, delAddr)
+				err = k.SendBtcDelegationRewardsToGauge(ctx, fpAddr, delAddr)
 				require.NoError(t, err)
 			}
-			fpCurrentRwd, err := keeper.GetFinalityProviderCurrentRewards(ctx, fpAddr)
+			fpCurrentRwd, err := k.GetFinalityProviderCurrentRewards(ctx, fpAddr)
 			require.NoError(t, err)
 			require.Equal(t, fpCurrentRwd.TotalActiveSat.Uint64(), fp.TotalBondedSat)
 		}
@@ -93,7 +93,7 @@ func FuzzRewardBTCStaking(f *testing.F) {
 		for addrStr, reward := range fpRewardMap {
 			addr, err := sdk.AccAddressFromBech32(addrStr)
 			require.NoError(t, err)
-			rg := keeper.GetRewardGauge(ctx, types.FinalityProviderType, addr)
+			rg := k.GetRewardGauge(ctx, types.FinalityProviderType, addr)
 			require.NotNil(t, rg)
 			require.Equal(t, reward, rg.Coins)
 		}
@@ -102,13 +102,14 @@ func FuzzRewardBTCStaking(f *testing.F) {
 		for addrStr, reward := range btcDelRewardMap {
 			addr, err := sdk.AccAddressFromBech32(addrStr)
 			require.NoError(t, err)
-			rg := keeper.GetRewardGauge(ctx, types.BTCDelegationType, addr)
+			rg := k.GetRewardGauge(ctx, types.BTCDelegationType, addr)
 			require.NotNil(t, rg)
 
 			// A little bit of rewards could be lost in the process due to precision points
 			// so 0.1% difference can be considered okay
 			allowedMarginError := CalculatePointOnePercent(reward)
-			require.Truef(t, reward.Sub(rg.Coins...).IsAllLT(allowedMarginError),
+			diff, _ := reward.SafeSub(rg.Coins...)
+			require.Truef(t, diff.IsAllLT(allowedMarginError),
 				"BTC delegation failed within the margin of error: %s\nRewards: %s\nGauge: %s",
 				allowedMarginError.String(), reward.String(), rg.Coins.String(),
 			)

x/incentive/keeper/reward_tracker.go

@@ -65,7 +65,7 @@ func (k Keeper) BtcDelegationActivated(ctx context.Context, fp, del sdk.AccAddre
 func (k Keeper) BtcDelegationUnbonded(ctx context.Context, fp, del sdk.AccAddress, sat uint64) error {
 	amtSat := sdkmath.NewIntFromUint64(sat)
 	return k.btcDelegationModifiedWithPreInitDel(ctx, fp, del, func(ctx context.Context, fp, del sdk.AccAddress) error {
-		return k.SubDelegationSat(ctx, fp, del, amtSat)
+		return k.subDelegationSat(ctx, fp, del, amtSat)
 	})
 }
 
@@ -270,16 +270,6 @@ func (k Keeper) addFinalityProviderStaked(ctx context.Context, fp sdk.AccAddress
 	return k.setFinalityProviderCurrentRewards(ctx, fp, fpCurrentRwd)
 }
 
-func (k Keeper) subFinalityProviderStaked(ctx context.Context, fp sdk.AccAddress, amt sdkmath.Int) error {
-	fpCurrentRwd, err := k.GetFinalityProviderCurrentRewards(ctx, fp)
-	if err != nil {
-		return err
-	}
-
-	fpCurrentRwd.SubTotalActiveSat(amt)
-	return k.setFinalityProviderCurrentRewards(ctx, fp, fpCurrentRwd)
-}
-
 func (k Keeper) AddFinalityProviderRewardsForDelegationsBTC(ctx context.Context, fp sdk.AccAddress, rwd sdk.Coins) error {
 	fpCurrentRwd, err := k.GetFinalityProviderCurrentRewards(ctx, fp)
 	if err != nil {
@@ -309,20 +299,3 @@ func (k Keeper) AddDelegationSat(ctx context.Context, fp, del sdk.AccAddress, am
 
 	return k.addFinalityProviderStaked(ctx, fp, amt)
 }
-
-// SubDelegationSat there is no need to check if the fp or delegation exists, because they should exist
-// otherwise it is probably a programming error calling to subtract the amount of active sat without
-// having any sat added in the first place.
-func (k Keeper) SubDelegationSat(ctx context.Context, fp, del sdk.AccAddress, amt sdkmath.Int) error {
-	btcDelRwdTracker, err := k.GetBTCDelegationRewardsTracker(ctx, fp, del)
-	if err != nil {
-		return err
-	}
-
-	btcDelRwdTracker.SubTotalActiveSat(amt)
-	if err := k.setBTCDelegationRewardsTracker(ctx, fp, del, btcDelRwdTracker); err != nil {
-		return err
-	}
-
-	return k.subFinalityProviderStaked(ctx, fp, amt)
-}

x/incentive/keeper/reward_tracker_store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/binary"
 
+	sdkmath "cosmossdk.io/math"
 	"cosmossdk.io/store/prefix"
 	"github.com/babylonlabs-io/babylon/x/incentive/types"
 	"github.com/cosmos/cosmos-sdk/runtime"
@@ -225,3 +226,34 @@ func (k Keeper) setFinalityProviderHistoricalRewards(ctx context.Context, fp sdk
 	k.storeFpHistoricalRewards(ctx, fp).Set(key, bz)
 	return nil
 }
+
+// subDelegationSat subtracts an amount of active stake from the BTCDelegationRewardsTracker
+// and the FinalityProviderCurrentRewards.
+// There is no need to check if the fp or delegation exists, because they should exist
+// otherwise it is probably a programming error calling to subtract the amount of active sat without
+// having any sat added in the first place that created the structures.
+func (k Keeper) subDelegationSat(ctx context.Context, fp, del sdk.AccAddress, amt sdkmath.Int) error {
+	btcDelRwdTracker, err := k.GetBTCDelegationRewardsTracker(ctx, fp, del)
+	if err != nil {
+		return err
+	}
+
+	btcDelRwdTracker.SubTotalActiveSat(amt)
+	if err := k.setBTCDelegationRewardsTracker(ctx, fp, del, btcDelRwdTracker); err != nil {
+		return err
+	}
+
+	return k.subFinalityProviderStaked(ctx, fp, amt)
+}
+
+// subFinalityProviderStaked subtracts an amount of active stake from the
+// FinalityProviderCurrentRewards, it errors out if the finality provider does not exist.
+func (k Keeper) subFinalityProviderStaked(ctx context.Context, fp sdk.AccAddress, amt sdkmath.Int) error {
+	fpCurrentRwd, err := k.GetFinalityProviderCurrentRewards(ctx, fp)
+	if err != nil {
+		return err
+	}
+
+	fpCurrentRwd.SubTotalActiveSat(amt)
+	return k.setFinalityProviderCurrentRewards(ctx, fp, fpCurrentRwd)
+}

x/incentive/keeper/reward_tracker_store_test.go

@@ -264,9 +264,121 @@ func FuzzCheckFinalityProviderHistoricalRewards(f *testing.F) {
 	})
 }
 
-func NewKeeperWithCtx(t *testing.T) (Keeper, sdk.Context) {
+func TestAddSubDelegationSat(t *testing.T) {
+	k, ctx := NewKeeperWithCtx(t)
+
+	fp1, del1 := datagen.GenRandomAddress(), datagen.GenRandomAddress()
+	fp2, del2 := datagen.GenRandomAddress(), datagen.GenRandomAddress()
+	amtFp1Del1, amtFp1Del2, amtFp2Del2, amtFp2Del1 := math.NewInt(2000), math.NewInt(4000), math.NewInt(500), math.NewInt(700)
+
+	_, err := k.GetBTCDelegationRewardsTracker(ctx, fp1, del1)
+	require.EqualError(t, err, types.ErrBTCDelegationRewardsTrackerNotFound.Error())
+
+	// adds 2000 for fp1, del1
+	// fp1       => 2000
+	// fp1, del1 => 2000
+	err = k.AddDelegationSat(ctx, fp1, del1, amtFp1Del1)
+	require.NoError(t, err)
+	checkFpTotalSat(t, ctx, k, fp1, amtFp1Del1)
+	checkFpDelTotalSat(t, ctx, k, fp1, del1, amtFp1Del1)
+
+	btcDelRwdFp1Del1, err := k.GetBTCDelegationRewardsTracker(ctx, fp1, del1)
+	require.NoError(t, err)
+	// if the normal flow with initilize BTC delegation would have been called,
+	// it would start as 1.
+	require.Equal(t, btcDelRwdFp1Del1.StartPeriodCumulativeReward, uint64(0))
+
+	// adds 4000 for fp1, del2
+	// fp1       => 6000
+	// fp1, del1 => 2000
+	// fp1, del2 => 4000
+	err = k.AddDelegationSat(ctx, fp1, del2, amtFp1Del2)
+	require.NoError(t, err)
+
+	checkFpTotalSat(t, ctx, k, fp1, amtFp1Del1.Add(amtFp1Del2))
+	checkFpDelTotalSat(t, ctx, k, fp1, del2, amtFp1Del2)
+	checkFpDelTotalSat(t, ctx, k, fp1, del1, amtFp1Del1)
+
+	// adds 500 for fp2, del2
+	// fp1       => 6000
+	// fp2       =>  500
+	// fp1, del1 => 2000
+	// fp1, del2 => 4000
+	// fp2, del2 =>  500
+	err = k.AddDelegationSat(ctx, fp2, del2, amtFp2Del2)
+	require.NoError(t, err)
+	checkFpTotalSat(t, ctx, k, fp1, amtFp1Del1.Add(amtFp1Del2))
+	checkFpTotalSat(t, ctx, k, fp2, amtFp2Del2)
+	checkFpDelTotalSat(t, ctx, k, fp1, del1, amtFp1Del1)
+	checkFpDelTotalSat(t, ctx, k, fp1, del2, amtFp1Del2)
+	checkFpDelTotalSat(t, ctx, k, fp2, del2, amtFp2Del2)
+
+	// adds 700 for fp2, del1
+	// fp1       => 6000
+	// fp2       => 1200
+	// fp1, del1 => 2000
+	// fp1, del2 => 4000
+	// fp2, del1 =>  700
+	// fp2, del2 =>  500
+	err = k.AddDelegationSat(ctx, fp2, del1, amtFp2Del1)
+	require.NoError(t, err)
+	checkFpTotalSat(t, ctx, k, fp1, amtFp1Del1.Add(amtFp1Del2))
+	checkFpTotalSat(t, ctx, k, fp2, amtFp2Del2.Add(amtFp2Del1))
+	checkFpDelTotalSat(t, ctx, k, fp1, del1, amtFp1Del1)
+	checkFpDelTotalSat(t, ctx, k, fp1, del2, amtFp1Del2)
+	checkFpDelTotalSat(t, ctx, k, fp2, del1, amtFp2Del1)
+	checkFpDelTotalSat(t, ctx, k, fp2, del2, amtFp2Del2)
+
+	lastAmtFp1Del2 := math.NewInt(2000)
+	// adds 2000 for fp1, del2
+	// fp1       => 8000
+	// fp2       => 1200
+	// fp1, del1 => 2000
+	// fp1, del2 => 6000
+	// fp2, del1 =>  700
+	// fp2, del2 =>  500
+	err = k.AddDelegationSat(ctx, fp1, del2, lastAmtFp1Del2)
+	require.NoError(t, err)
+	checkFpTotalSat(t, ctx, k, fp1, amtFp1Del1.Add(amtFp1Del2).Add(lastAmtFp1Del2))
+	checkFpTotalSat(t, ctx, k, fp2, amtFp2Del2.Add(amtFp2Del1))
+	checkFpDelTotalSat(t, ctx, k, fp1, del1, amtFp1Del1)
+	checkFpDelTotalSat(t, ctx, k, fp1, del2, amtFp1Del2.Add(lastAmtFp1Del2))
+	checkFpDelTotalSat(t, ctx, k, fp2, del1, amtFp2Del1)
+	checkFpDelTotalSat(t, ctx, k, fp2, del2, amtFp2Del2)
+
+	subAmtFp2Del2 := math.NewInt(350)
+	// subtract 350 for fp2, del2
+	// fp1       => 8000
+	// fp2       =>  850
+	// fp1, del1 => 2000
+	// fp1, del2 => 6000
+	// fp2, del1 =>  700
+	// fp2, del2 =>  150
+	err = k.subDelegationSat(ctx, fp2, del2, subAmtFp2Del2)
+	require.NoError(t, err)
+	checkFpTotalSat(t, ctx, k, fp1, amtFp1Del1.Add(amtFp1Del2).Add(lastAmtFp1Del2))
+	checkFpTotalSat(t, ctx, k, fp2, amtFp2Del2.Add(amtFp2Del1).Sub(subAmtFp2Del2))
+	checkFpDelTotalSat(t, ctx, k, fp1, del1, amtFp1Del1)
+	checkFpDelTotalSat(t, ctx, k, fp1, del2, amtFp1Del2.Add(lastAmtFp1Del2))
+	checkFpDelTotalSat(t, ctx, k, fp2, del1, amtFp2Del1)
+	checkFpDelTotalSat(t, ctx, k, fp2, del2, amtFp2Del2.Sub(subAmtFp2Del2))
+}
+
+func checkFpTotalSat(t *testing.T, ctx sdk.Context, k *Keeper, fp sdk.AccAddress, expectedSat math.Int) {
+	rwd, err := k.GetFinalityProviderCurrentRewards(ctx, fp)
+	require.NoError(t, err)
+	require.Equal(t, expectedSat.String(), rwd.TotalActiveSat.String())
+}
+
+func checkFpDelTotalSat(t *testing.T, ctx sdk.Context, k *Keeper, fp, del sdk.AccAddress, expectedSat math.Int) {
+	rwd, err := k.GetBTCDelegationRewardsTracker(ctx, fp, del)
+	require.NoError(t, err)
+	require.Equal(t, expectedSat.String(), rwd.TotalActiveSat.String())
+}
+
+func NewKeeperWithCtx(t *testing.T) (*Keeper, sdk.Context) {
 	encConf := appparams.DefaultEncodingConfig()
 	ctx, kvStore := store.NewStoreWithCtx(t, types.ModuleName)
 	k := NewKeeper(encConf.Codec, kvStore, nil, nil, nil, addr.AccGov.String(), addr.AccFeeCollector.String())
-	return k, ctx
+	return &k, ctx
 }