feat(x/protocolpool)!: allow any coins in continuous funds

simapp/app.go

@@ -377,7 +377,7 @@ func NewSimApp(
 		panic(err)
 	}
 
-	app.PoolKeeper = poolkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), logger.With(log.ModuleKey, "x/protocolpool")), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, govModuleAddr)
+	app.PoolKeeper = poolkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), logger.With(log.ModuleKey, "x/protocolpool")), app.AuthKeeper, app.BankKeeper, govModuleAddr)
 
 	app.DistrKeeper = distrkeeper.NewKeeper(appCodec, runtime.NewEnvironment(runtime.NewKVStoreService(keys[distrtypes.StoreKey]), logger.With(log.ModuleKey, "x/distribution")), app.AuthKeeper, app.BankKeeper, app.StakingKeeper, cometService, authtypes.FeeCollectorName, govModuleAddr)
 

tests/integration/distribution/keeper/msg_server_test.go

@@ -138,7 +138,7 @@ func initFixture(t *testing.T) *fixture {
 	stakingKeeper := stakingkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[stakingtypes.StoreKey]), log.NewNopLogger(), runtime.EnvWithQueryRouterService(grpcRouter), runtime.EnvWithMsgRouterService(msgRouter)), accountKeeper, bankKeeper, consensusParamsKeeper, authority.String(), addresscodec.NewBech32Codec(sdk.Bech32PrefixValAddr), addresscodec.NewBech32Codec(sdk.Bech32PrefixConsAddr), cometService)
 	require.NoError(t, stakingKeeper.Params.Set(newCtx, stakingtypes.DefaultParams()))
 
-	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, stakingKeeper, authority.String())
+	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, authority.String())
 
 	distrKeeper := distrkeeper.NewKeeper(
 		cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[distrtypes.StoreKey]), logger), accountKeeper, bankKeeper, stakingKeeper, cometService, distrtypes.ModuleName, authority.String(),

tests/integration/gov/keeper/keeper_test.go

@@ -119,7 +119,7 @@ func initFixture(tb testing.TB) *fixture {
 
 	stakingKeeper := stakingkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[stakingtypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, consensusParamsKeeper, authority.String(), addresscodec.NewBech32Codec(sdk.Bech32PrefixValAddr), addresscodec.NewBech32Codec(sdk.Bech32PrefixConsAddr), runtime.NewContextAwareCometInfoService())
 
-	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, stakingKeeper, authority.String())
+	poolKeeper := poolkeeper.NewKeeper(cdc, runtime.NewEnvironment(runtime.NewKVStoreService(keys[pooltypes.StoreKey]), log.NewNopLogger()), accountKeeper, bankKeeper, authority.String())
 
 	// set default staking params
 	err := stakingKeeper.Params.Set(newCtx, stakingtypes.DefaultParams())

x/protocolpool/depinject.go

@@ -37,7 +37,6 @@ type ModuleInputs struct {
 
 	AccountKeeper types.AccountKeeper
 	BankKeeper    types.BankKeeper
-	StakingKeeper types.StakingKeeper
 }
 
 type ModuleOutputs struct {
@@ -59,7 +58,7 @@ func ProvideModule(in ModuleInputs) ModuleOutputs {
 		panic(err)
 	}
 
-	k := keeper.NewKeeper(in.Codec, in.Environment, in.AccountKeeper, in.BankKeeper, in.StakingKeeper, authorityAddr)
+	k := keeper.NewKeeper(in.Codec, in.Environment, in.AccountKeeper, in.BankKeeper, authorityAddr)
 	m := NewAppModule(in.Codec, k, in.AccountKeeper, in.BankKeeper)
 
 	return ModuleOutputs{

x/protocolpool/keeper/genesis.go

@@ -2,11 +2,9 @@ package keeper
 
 import (
 	"context"
-	"errors"
 	"fmt"
 	"time"
 
-	"cosmossdk.io/math"
 	"cosmossdk.io/x/protocolpool/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -56,19 +54,19 @@ func (k Keeper) InitGenesis(ctx context.Context, data *types.GenesisState) error
 		return fmt.Errorf("failed to set last balance: %w", err)
 	}
 
-	totalToBeDistributed := math.ZeroInt()
+	totalToBeDistributed := sdk.NewCoins()
 	for _, distribution := range data.Distributions {
-		totalToBeDistributed = totalToBeDistributed.Add(distribution.Amount)
+		totalToBeDistributed = totalToBeDistributed.Add(distribution.Amount.Amount...)
 		if err := k.Distributions.Set(ctx, *distribution.Time, distribution.Amount); err != nil {
 			return fmt.Errorf("failed to set distribution: %w", err)
 		}
 	}
 
 	// sanity check to avoid trying to distribute more than what is available
-	if data.LastBalance.LT(totalToBeDistributed) {
-		return errors.New("total to be distributed is greater than the last balance")
-	}
 
+	if totalToBeDistributed.IsAnyGT(data.LastBalance.Amount) || !totalToBeDistributed.DenomsSubsetOf(data.LastBalance.Amount) {
+		return fmt.Errorf("total to be distributed is greater than the last balance: %s > %s", totalToBeDistributed, data.LastBalance.Amount)
+	}
 	return nil
 }
 
@@ -112,12 +110,14 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*types.GenesisState, error)
 
 	genState := types.NewGenesisState(cf, budget)
 
-	genState.LastBalance, err = k.LastBalance.Get(ctx)
+	lastBalance, err := k.LastBalance.Get(ctx)
 	if err != nil {
 		return nil, err
 	}
 
-	err = k.Distributions.Walk(ctx, nil, func(key time.Time, value math.Int) (stop bool, err error) {
+	genState.LastBalance = lastBalance
+
+	err = k.Distributions.Walk(ctx, nil, func(key time.Time, value types.DistributionAmount) (stop bool, err error) {
 		genState.Distributions = append(genState.Distributions, &types.Distribution{
 			Time:   &key,
 			Amount: value,

x/protocolpool/keeper/genesis_test.go

@@ -32,15 +32,15 @@ func (suite *KeeperTestSuite) TestInitGenesis() {
 	)
 
 	gs.Distributions = append(gs.Distributions, &types.Distribution{
-		Amount: math.OneInt(),
+		Amount: types.DistributionAmount{Amount: sdk.NewCoins(sdk.NewCoin("stake", math.NewInt(100)))},
 		Time:   &time.Time{},
 	})
 
 	err := suite.poolKeeper.InitGenesis(suite.ctx, gs)
 	suite.Require().ErrorContains(err, "total to be distributed is greater than the last balance")
 
 	// Set last balance
-	gs.LastBalance = math.NewInt(1)
+	gs.LastBalance = types.DistributionAmount{Amount: sdk.NewCoins(sdk.NewCoin("stake", math.NewInt(101)))}
 	err = suite.poolKeeper.InitGenesis(suite.ctx, gs)
 	suite.Require().NoError(err)
 
@@ -49,5 +49,5 @@ func (suite *KeeperTestSuite) TestInitGenesis() {
 	suite.Require().NoError(err)
 	suite.Require().Equal(gs.ContinuousFund, exportedGenState.ContinuousFund)
 	suite.Require().Equal(gs.Budget, exportedGenState.Budget)
-	suite.Require().Equal(math.OneInt(), exportedGenState.LastBalance)
+	suite.Require().Equal(math.NewInt(101), exportedGenState.LastBalance.Amount.AmountOf("stake"))
 }

x/protocolpool/keeper/keeper.go

@@ -21,9 +21,8 @@
 type Keeper struct {
 	appmodule.Environment
 
-	authKeeper    types.AccountKeeper
-	bankKeeper    types.BankKeeper
-	stakingKeeper types.StakingKeeper
+	authKeeper types.AccountKeeper
+	bankKeeper types.BankKeeper
 
 	cdc codec.BinaryCodec
 
@@ -34,16 +33,16 @@
 	BudgetProposal collections.Map[sdk.AccAddress, types.Budget]
 	ContinuousFund collections.Map[sdk.AccAddress, types.ContinuousFund]
 	// RecipientFundDistribution key: RecipientAddr | value: Claimable amount
-	RecipientFundDistribution collections.Map[sdk.AccAddress, math.Int]
-	Distributions             collections.Map[time.Time, math.Int] // key: time.Time | value: amount
-	LastBalance               collections.Item[math.Int]
+	RecipientFundDistribution collections.Map[sdk.AccAddress, types.DistributionAmount]
+	Distributions             collections.Map[time.Time, types.DistributionAmount] // key: time.Time, denom | value: amounts
+	LastBalance               collections.Item[types.DistributionAmount]
 }
 
 const (
 	errModuleAccountNotSet = "%s module account has not been set"
 )
 
-func NewKeeper(cdc codec.BinaryCodec, env appmodule.Environment, ak types.AccountKeeper, bk types.BankKeeper, sk types.StakingKeeper, authority string,
+func NewKeeper(cdc codec.BinaryCodec, env appmodule.Environment, ak types.AccountKeeper, bk types.BankKeeper, authority string,
 ) Keeper {
 	// ensure pool module account is set
 	if addr := ak.GetModuleAddress(types.ModuleName); addr == nil {
@@ -64,14 +63,13 @@
 		Environment:               env,
 		authKeeper:                ak,
 		bankKeeper:                bk,
-		stakingKeeper:             sk,
 		cdc:                       cdc,
 		authority:                 authority,
 		BudgetProposal:            collections.NewMap(sb, types.BudgetKey, "budget", sdk.AccAddressKey, codec.CollValue[types.Budget](cdc)),
 		ContinuousFund:            collections.NewMap(sb, types.ContinuousFundKey, "continuous_fund", sdk.AccAddressKey, codec.CollValue[types.ContinuousFund](cdc)),
-		RecipientFundDistribution: collections.NewMap(sb, types.RecipientFundDistributionKey, "recipient_fund_distribution", sdk.AccAddressKey, sdk.IntValue),
-		Distributions:             collections.NewMap(sb, types.DistributionsKey, "distributions", sdk.TimeKey, sdk.IntValue),
-		LastBalance:               collections.NewItem(sb, types.LastBalanceKey, "last_balance", sdk.IntValue),
+		RecipientFundDistribution: collections.NewMap(sb, types.RecipientFundDistributionKey, "recipient_fund_distribution", sdk.AccAddressKey, codec.CollValue[types.DistributionAmount](cdc)),
+		Distributions:             collections.NewMap(sb, types.DistributionsKey, "distributions", sdk.TimeKey, codec.CollValue[types.DistributionAmount](cdc)),
+		LastBalance:               collections.NewItem(sb, types.LastBalanceKey, "last_balance", codec.CollValue[types.DistributionAmount](cdc)),
 	}
 
 	schema, err := sb.Build()
@@ -114,34 +112,27 @@
 	return k.bankKeeper.GetAllBalances(ctx, moduleAccount.GetAddress()), nil
 }
 
-func (k Keeper) withdrawRecipientFunds(ctx context.Context, recipient []byte) (sdk.Coin, error) {
+func (k Keeper) withdrawRecipientFunds(ctx context.Context, recipient []byte) (sdk.Coins, error) {
 	// get allocated continuous fund
 	fundsAllocated, err := k.RecipientFundDistribution.Get(ctx, recipient)
 	if err != nil {
 		if errors.Is(err, collections.ErrNotFound) {
-			return sdk.Coin{}, types.ErrNoRecipientFound
+			return nil, types.ErrNoRecipientFound
 		}
-		return sdk.Coin{}, err
-	}
-
-	denom, err := k.stakingKeeper.BondDenom(ctx)
-	if err != nil {
-		return sdk.Coin{}, err
+		return nil, err
 	}
-
 	// Distribute funds to the recipient from pool module account
-	withdrawnAmount := sdk.NewCoin(denom, fundsAllocated)
-	err = k.DistributeFromStreamFunds(ctx, sdk.NewCoins(withdrawnAmount), recipient)
+	err = k.DistributeFromStreamFunds(ctx, fundsAllocated.Amount, recipient)
 	if err != nil {
-		return sdk.Coin{}, fmt.Errorf("error while distributing funds: %w", err)
+		return nil, fmt.Errorf("error while distributing funds: %w", err)
 	}
 
 	// reset fund distribution
-	err = k.RecipientFundDistribution.Set(ctx, recipient, math.ZeroInt())
+	err = k.RecipientFundDistribution.Set(ctx, recipient, types.DistributionAmount{Amount: sdk.NewCoins()})
 	if err != nil {
-		return sdk.Coin{}, err
+		return nil, err
 	}
-	return withdrawnAmount, nil
+	return fundsAllocated.Amount, nil
 }
 
 // SetToDistribute sets the amount to be distributed among recipients.
@@ -152,30 +143,29 @@
 		return errorsmod.Wrapf(sdkerrors.ErrUnknownAddress, "module account %s does not exist", types.ProtocolPoolDistrAccount)
 	}
 
-	denom, err := k.stakingKeeper.BondDenom(ctx)
-	if err != nil {
-		return err
-	}
-
 	currentBalance := k.bankKeeper.GetAllBalances(ctx, moduleAccount.GetAddress())
-	distributionBalance := currentBalance.AmountOf(denom)
 
 	// if the balance is zero, return early
-	if distributionBalance.IsZero() {
+	if currentBalance.IsZero() {
 		return nil
 	}
 
+	// if the balance does not have any of the allowed denoms, return early // TODO
+
 	lastBalance, err := k.LastBalance.Get(ctx)
 	if err != nil {
 		if errors.Is(err, collections.ErrNotFound) {
-			lastBalance = math.ZeroInt()
+			lastBalance = types.DistributionAmount{Amount: sdk.NewCoins()}
 		} else {
 			return err
 		}
 	}
 
 	// Calculate the amount to be distributed
-	amountToDistribute := distributionBalance.Sub(lastBalance)
+	amountToDistribute, anyNegative := currentBalance.SafeSub(lastBalance.Amount...)
+	if anyNegative {
+		return errors.New("error while calculating the amount to distribute, result can't be negative")
+	}
 
 	// Check if there are any recipients to distribute to, if not, send straight to the community pool and avoid
 	// setting the distributions
@@ -190,24 +180,23 @@
 
 	// if there are no continuous funds, send all the funds to the community pool and reset the last balance
 	if !hasContinuousFunds {
-		poolCoins := sdk.NewCoins(sdk.NewCoin(denom, amountToDistribute))
-		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.ModuleName, poolCoins); err != nil {
+		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.ModuleName, amountToDistribute); err != nil {
 			return err
 		}
 
-		if !lastBalance.IsZero() { // only reset if the last balance is not zero (so we leave it at zero/nil)
-			return k.LastBalance.Set(ctx, math.ZeroInt())
+		if !lastBalance.Amount.IsZero() { // only reset if the last balance is not zero (so we leave it at zero)
+			return k.LastBalance.Set(ctx, types.DistributionAmount{Amount: sdk.NewCoins()})
 		}
 
 		return nil
 	}
 
-	if err = k.Distributions.Set(ctx, k.HeaderService.HeaderInfo(ctx).Time, amountToDistribute); err != nil {
+	if err = k.Distributions.Set(ctx, k.HeaderService.HeaderInfo(ctx).Time, types.DistributionAmount{Amount: amountToDistribute}); err != nil {
 		return fmt.Errorf("error while setting Distributions: %w", err)
 	}
 
 	// Update the last balance
-	return k.LastBalance.Set(ctx, distributionBalance)
+	return k.LastBalance.Set(ctx, types.DistributionAmount{Amount: currentBalance})
 }
 
 func (k Keeper) IterateAndUpdateFundsDistribution(ctx context.Context) error {
@@ -229,12 +218,15 @@
 	}
 
 	// next we iterate over the distributions, calculate each recipient's share and the remaining pool funds
-	toDistribute := map[string]math.Int{}
-	poolFunds := math.ZeroInt()
-	fullAmountToDistribute := math.ZeroInt()
+	toDistribute := map[string]sdk.Coins{}
+	poolFunds := sdk.NewCoins()
+	amountToDistribute := sdk.NewCoins() // amount assigned to distributions
+	allDistributions := sdk.NewCoins()   // total amount distributed to the pool, to then calculate the remaining pool funds
 
-	if err = k.Distributions.Walk(ctx, nil, func(key time.Time, amount math.Int) (stop bool, err error) {
+	if err = k.Distributions.Walk(ctx, nil, func(key time.Time, amount types.DistributionAmount) (stop bool, err error) {
 		percentageToDistribute := math.LegacyZeroDec()
+		allDistributions = allDistributions.Add(amount.Amount...)
+
 		for _, f := range funds {
 			if f.Expiry != nil && f.Expiry.Before(key) {
 				continue
@@ -244,21 +236,21 @@
 
 			_, ok := toDistribute[f.Recipient]
 			if !ok {
-				toDistribute[f.Recipient] = math.ZeroInt()
+				toDistribute[f.Recipient] = sdk.NewCoins()
+			}
+
+			for _, denom := range amount.Amount.Denoms() {
+				am := sdk.NewCoin(denom, f.Percentage.MulInt(amount.Amount.AmountOf(denom)).TruncateInt())
+				toDistribute[f.Recipient] = toDistribute[f.Recipient].Add(am)
+				amountToDistribute = amountToDistribute.Add(am)
 			}
-			amountToDistribute := f.Percentage.MulInt(amount).TruncateInt()
-			toDistribute[f.Recipient] = toDistribute[f.Recipient].Add(amountToDistribute)
-			fullAmountToDistribute = fullAmountToDistribute.Add(amountToDistribute)
 		}
 
 		// sanity check for max percentage
 		if percentageToDistribute.GT(math.LegacyOneDec()) {
 			return true, errors.New("total funds percentage cannot exceed 100")
 		}
 
-		remaining := math.LegacyOneDec().Sub(percentageToDistribute).MulInt(amount).RoundInt()
-		poolFunds = poolFunds.Add(remaining)
-
 		return false, nil
 	}); err != nil {
 		return err
@@ -269,26 +261,20 @@
 		return err
 	}
 
-	if err = k.LastBalance.Set(ctx, math.ZeroInt()); err != nil {
+	if err = k.LastBalance.Set(ctx, types.DistributionAmount{Amount: sdk.NewCoins()}); err != nil {
 		return err
 	}
 
 	// send the funds to the stream account to be distributed later, and the remaining to the community pool
-	bondDenom, err := k.stakingKeeper.BondDenom(ctx)
-	if err != nil {
-		return err
-	}
-
-	streamAmt := sdk.NewCoins(sdk.NewCoin(bondDenom, fullAmountToDistribute))
-	if !streamAmt.IsZero() {
-		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.StreamAccount, streamAmt); err != nil {
+	if !amountToDistribute.IsZero() {
+		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.StreamAccount, amountToDistribute); err != nil {
 			return err
 		}
 	}
 
+	poolFunds = allDistributions.Sub(amountToDistribute...)
 	if !poolFunds.IsZero() {
-		poolCoins := sdk.NewCoins(sdk.NewCoin(bondDenom, poolFunds))
-		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.ModuleName, poolCoins); err != nil {
+		if err := k.bankKeeper.SendCoinsFromModuleToModule(ctx, types.ProtocolPoolDistrAccount, types.ModuleName, poolFunds); err != nil {
 			return err
 		}
 	}
@@ -310,14 +296,14 @@
 		toClaim, err := k.RecipientFundDistribution.Get(ctx, bzAddr)
 		if err != nil {
 			if errors.Is(err, collections.ErrNotFound) {
-				toClaim = math.ZeroInt()
+				toClaim = types.DistributionAmount{Amount: sdk.NewCoins()}
 			} else {
 				return err
 			}
 		}
 
-		amount := toClaim.Add(toDistribute[recipient])
-		if err = k.RecipientFundDistribution.Set(ctx, bzAddr, amount); err != nil {
+		toClaim.Amount = toClaim.Amount.Add(toDistribute[recipient]...)
+		if err = k.RecipientFundDistribution.Set(ctx, bzAddr, toClaim); err != nil {
 			return err
 		}
 	}