feature: Add accumulating rewards to NAS

proto/pstake/liquidstake/v1beta1/liquidstake.proto

@@ -212,4 +212,10 @@ message NetAmountState {
     (gogoproto.customtype) = "github.com/cosmos/cosmos-sdk/types.Int",
     (gogoproto.nullable) = false
   ];
+
+  // accumulating_rewards define the accumulating rewards of proxy account
+  string accumulating_rewards = 9 [
+    (gogoproto.customtype) = "cosmossdk.io/math.Int",
+    (gogoproto.nullable) = false
+  ];
 }

x/liquidstake/keeper/keeper.go

@@ -3,6 +3,7 @@ package keeper
 import (
 	"fmt"
 
+	sdkmath "cosmossdk.io/math"
 	"github.com/cometbft/cometbft/libs/log"
 	"github.com/cosmos/cosmos-sdk/baseapp"
 	"github.com/cosmos/cosmos-sdk/codec"
@@ -91,6 +92,32 @@ func (k Keeper) GetParams(ctx sdk.Context) (params types.Params) {
 	return params
 }
 
+// GetAccumulatingRewards returns the accumulating rewards in the proxy account.
+func (k Keeper) GetAccumulatingRewards(ctx sdk.Context) sdkmath.Int {
+	store := ctx.KVStore(k.storeKey)
+
+	bz := store.Get(types.AccumulatingRewardsKey)
+	if bz == nil {
+		return sdkmath.ZeroInt()
+	}
+
+	var acRewards sdkmath.Int
+	if err := acRewards.Unmarshal(bz); err != nil {
+		panic(err)
+	}
+	return acRewards
+}
+
+// SetAccumulatingRewards sets the accumulating rewards in the proxy account.
+func (k Keeper) SetAccumulatingRewards(ctx sdk.Context, acRewards sdkmath.Int) {
+	store := ctx.KVStore(k.storeKey)
+	bz, err := acRewards.Marshal()
+	if err != nil {
+		panic(err)
+	}
+	store.Set(types.AccumulatingRewardsKey, bz)
+}
+
 // GetCodec return codec.Codec object used by the keeper
 func (k Keeper) GetCodec() codec.BinaryCodec { return k.cdc }
 

x/liquidstake/keeper/liquidstake.go

@@ -43,6 +43,7 @@ func (k Keeper) GetNetAmountState(ctx sdk.Context) (nas types.NetAmountState) {
 		TotalRemainingRewards: totalRemainingRewards,
 		TotalUnbondingBalance: totalUnbondingBalance,
 		ProxyAccBalance:       k.GetProxyAccBalance(ctx, types.LiquidStakeProxyAcc).Amount,
+		AccumulatingRewards:   k.GetAccumulatingRewards(ctx),
 	}
 
 	nas.NetAmount = nas.CalcNetAmount()
@@ -715,6 +716,9 @@ func (k Keeper) LiquidDelegate(ctx sdk.Context, proxyAcc sdk.AccAddress, activeV
 		return types.ErrInvalidActiveLiquidValidators
 	}
 	weightedAmt[0] = weightedAmt[0].Add(crumb)
+
+	acRewards := k.GetAccumulatingRewards(ctx)
+
 	for i, val := range activeVals {
 		if !weightedAmt[i].IsPositive() {
 			continue
@@ -724,6 +728,12 @@ func (k Keeper) LiquidDelegate(ctx sdk.Context, proxyAcc sdk.AccAddress, activeV
 		if err != nil {
 			return errorsmod.Wrapf(err, "failed to delegate to validator %s", val.GetOperator())
 		}
+
+		if weightedAmt[i].GTE(acRewards) {
+			acRewards = sdk.ZeroInt()
+		} else {
+			acRewards = acRewards.Sub(weightedAmt[i])
+		}
 	}
 	return nil
 }
@@ -792,6 +802,15 @@ func (k Keeper) LiquidUnstake(
 				return time.Time{}, sdk.ZeroInt(), []stakingtypes.UnbondingDelegation{}, sdk.ZeroInt(), err
 			}
 
+			// update accumulating rewards
+			acRewards := k.GetAccumulatingRewards(ctx)
+			if unbondingAmountInt.GTE(acRewards) {
+				acRewards = sdk.ZeroInt()
+			} else {
+				acRewards = acRewards.Sub(unbondingAmountInt)
+			}
+			k.SetAccumulatingRewards(ctx, acRewards)
+
 			return time.Time{}, sdk.ZeroInt(), []stakingtypes.UnbondingDelegation{}, unbondingAmountInt, nil
 		}
 
@@ -987,6 +1006,9 @@ func (k Keeper) CheckDelegationStates(ctx sdk.Context, proxyAcc sdk.AccAddress)
 }
 
 func (k Keeper) WithdrawLiquidRewards(ctx sdk.Context, proxyAcc sdk.AccAddress) {
+	acRewards := k.GetAccumulatingRewards(ctx)
+	bondDenom := k.stakingKeeper.BondDenom(ctx)
+
 	// iterate over all the delegations (even those out of the active set) and withdraw rewards
 	k.stakingKeeper.IterateDelegations(
 		ctx, proxyAcc,
@@ -1014,13 +1036,27 @@ func (k Keeper) WithdrawLiquidRewards(ctx sdk.Context, proxyAcc sdk.AccAddress)
 				)
 				// no need to return here, will be picked up in the next epoch
 			} else {
+				var msgWithdrawResponse distributiontypes.MsgWithdrawDelegatorRewardResponse
+				if err := k.cdc.Unmarshal(res.MsgResponses[0].Value, &msgWithdrawResponse); err != nil {
+					k.Logger(ctx).Error(
+						"failed to unmarshal withdraw rewards response",
+						types.ErrorKeyVal,
+						err,
+					)
+				}
+
+				amount := msgWithdrawResponse.Amount.AmountOf(bondDenom)
+				acRewards = acRewards.Add(amount)
+
 				// emit the events
 				ctx.EventManager().EmitEvents(res.GetEvents())
 			}
 
 			return false
 		},
 	)
+
+	k.SetAccumulatingRewards(ctx, acRewards)
 }
 
 // GetLiquidValidator get a single liquid validator

x/liquidstake/keeper/rebalancing.go

@@ -260,6 +260,15 @@ func (k Keeper) AutocompoundStakingRewards(ctx sdk.Context, whitelistedValsMap t
 		return
 	}
 
+	// subtract fee from accumulated rewards
+	acRewards := k.GetAccumulatingRewards(ctx)
+	if autocompoundFee.Amount.GTE(acRewards) {
+		acRewards = sdk.ZeroInt()
+	} else {
+		acRewards = acRewards.Sub(autocompoundFee.Amount)
+	}
+	k.SetAccumulatingRewards(ctx, acRewards)
+
 	ctx.EventManager().EmitEvents(sdk.Events{
 		sdk.NewEvent(
 			types.EventTypeAutocompound,

x/liquidstake/types/keys.go

@@ -28,6 +28,9 @@ var (
 
 	// LiquidValidatorsKey defines prefix for each key to a liquid validator
 	LiquidValidatorsKey = []byte{0x02}
+
+	// AccumulatingRewardsKey defines prefix for each key to accumulating rewards in proxy account
+	AccumulatingRewardsKey = []byte{0x03}
 )
 
 // GetLiquidValidatorKey creates the key for the liquid validator with address

x/liquidstake/types/liquidstake.go

@@ -182,7 +182,7 @@ func DeductFeeRate(input, feeRate math.LegacyDec) (feeDeductedOutput math.Legacy
 }
 
 func (nas NetAmountState) CalcNetAmount() math.LegacyDec {
-	return math.LegacyNewDecFromInt(nas.TotalLiquidTokens.Add(nas.TotalUnbondingBalance))
+	return math.LegacyNewDecFromInt(nas.TotalLiquidTokens.Add(nas.TotalUnbondingBalance).Add(nas.AccumulatingRewards))
 }
 
 func (nas NetAmountState) CalcMintRate() math.LegacyDec {