refactor: remove PokePool from the PoolI interface, define on extension instead

CHANGELOG.md

@@ -55,6 +55,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Misc Improvements
 
 * [#2804](https://github.com/osmosis-labs/osmosis/pull/2804) Improve error handling and messages when parsing pool assets.
+* [#3035](https://github.com/osmosis-labs/osmosis/pull/3035) Remove `PokePool` from `PoolI` interface. Define on a new WeightedPoolExtension` instead.
 
 ## v12.0.0
 

x/gamm/keeper/invariants.go

@@ -12,7 +12,7 @@ import (
 
 const poolBalanceInvariantName = "pool-account-balance-equals-expected"
 
-// RegisterInvariants registers all governance invariants.
+// RegisterInvariants registers all gamm invariants.
 func RegisterInvariants(ir sdk.InvariantRegistry, keeper Keeper, bk types.BankKeeper) {
 	ir.RegisterRoute(types.ModuleName, poolBalanceInvariantName, PoolAccountInvariant(keeper, bk))
 }

x/gamm/keeper/keeper_test.go

@@ -10,6 +10,7 @@ import (
 	"github.com/osmosis-labs/osmosis/v12/app/apptesting"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
 	balancertypes "github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/stableswap"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
 )
 
@@ -34,9 +35,7 @@ func (suite *KeeperTestSuite) prepareCustomBalancerPool(
 	poolAssets []balancertypes.PoolAsset,
 	poolParams balancer.PoolParams,
 ) uint64 {
-	for _, acc := range suite.TestAccs {
-		suite.FundAcc(acc, balances)
-	}
+	suite.fundAllAccountsWith(balances)
 
 	poolID, err := suite.App.GAMMKeeper.CreatePool(
 		suite.Ctx,
@@ -46,3 +45,26 @@ func (suite *KeeperTestSuite) prepareCustomBalancerPool(
 
 	return poolID
 }
+
+func (suite *KeeperTestSuite) prepareCustomStableswapPool(
+	balances sdk.Coins,
+	poolParams stableswap.PoolParams,
+	initialLiquidity sdk.Coins,
+	scalingFactors []uint64,
+) uint64 {
+	suite.fundAllAccountsWith(balances)
+
+	poolID, err := suite.App.GAMMKeeper.CreatePool(
+		suite.Ctx,
+		stableswap.NewMsgCreateStableswapPool(suite.TestAccs[0], poolParams, initialLiquidity, scalingFactors, ""),
+	)
+	suite.Require().NoError(err)
+
+	return poolID
+}
+
+func (suite *KeeperTestSuite) fundAllAccountsWith(balances sdk.Coins) {
+	for _, acc := range suite.TestAccs {
+		suite.FundAcc(acc, balances)
+	}
+}

x/gamm/keeper/pool.go

@@ -22,8 +22,8 @@ func (k Keeper) UnmarshalPool(bz []byte) (types.PoolI, error) {
 	return acc, k.cdc.UnmarshalInterface(bz, &acc)
 }
 
-// GetPoolAndPoke returns a PoolI based on it's identifier if one exists. Prior
-// to returning the pool, the weights of the pool are updated via PokePool.
+// GetPoolAndPoke returns a PoolI based on it's identifier if one exists. If poolId corresponds
+// to a pool with weights (e.g. balancer), the weights of the pool are updated via PokePool prior to returning.
 // TODO: Consider rename to GetPool due to downstream API confusion.
 func (k Keeper) GetPoolAndPoke(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
 	store := ctx.KVStore(k.storeKey)
@@ -39,7 +39,9 @@ func (k Keeper) GetPoolAndPoke(ctx sdk.Context, poolId uint64) (types.PoolI, err
 		return nil, err
 	}
 
-	pool.PokePool(ctx.BlockTime())
+	if pokePool, ok := pool.(types.WeightedPoolExtension); ok {
+		pokePool.PokePool(ctx.BlockTime())
+	}
 
 	return pool, nil
 }
@@ -74,7 +76,9 @@ func (k Keeper) GetPoolsAndPoke(ctx sdk.Context) (res []types.PoolI, err error)
 			return nil, err
 		}
 
-		pool.PokePool(ctx.BlockTime())
+		if pokePool, ok := pool.(types.WeightedPoolExtension); ok {
+			pokePool.PokePool(ctx.BlockTime())
+		}
 		res = append(res, pool)
 	}
 

x/gamm/keeper/pool_test.go

@@ -1,5 +1,16 @@
 package keeper_test
 
+import (
+	"time"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	balancertypes "github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/stableswap"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
+)
+
 // import (
 // 	"math/rand"
 // 	"time"
@@ -220,3 +231,87 @@ package keeper_test
 // 			"Expected equal %s: %d, %d", amt.Denom, amt.Amount.Int64(), sdk.NewInt(1000).Int64())
 // 	}
 // }
+
+// TestGetPoolAndPoke tests that the right pools is returned from GetPoolAndPoke.
+// For the pools implementing the weighted extension, asserts that PokePool is called.
+func (suite *KeeperTestSuite) TestGetPoolAndPoke() {
+	const (
+		startTime = 1000
+		blockTime = startTime + 100
+	)
+
+	// N.B.: We make a copy because SmoothWeightChangeParams get mutated.
+	// We would like to avoid mutating global pool assets that are used in other tests.
+	defaultPoolAssetsCopy := make([]balancertypes.PoolAsset, 2)
+	copy(defaultPoolAssetsCopy, defaultPoolAssets)
+
+	startPoolWeightAssets := []balancertypes.PoolAsset{
+		{
+			Weight: defaultPoolAssets[0].Weight.Quo(sdk.NewInt(2)),
+			Token:  defaultPoolAssets[0].Token,
+		},
+		{
+			Weight: defaultPoolAssets[1].Weight.Mul(sdk.NewInt(3)),
+			Token:  defaultPoolAssets[1].Token,
+		},
+	}
+
+	tests := map[string]struct {
+		isPokePool bool
+		poolId     uint64
+	}{
+		"weighted pool - change weights": {
+			isPokePool: true,
+			poolId: suite.prepareCustomBalancerPool(defaultAcctFunds, startPoolWeightAssets, balancer.PoolParams{
+				SwapFee: defaultSwapFee,
+				ExitFee: defaultExitFee,
+				SmoothWeightChangeParams: &balancer.SmoothWeightChangeParams{
+					StartTime:          time.Unix(startTime, 0), // start time is before block time so the weights should change
+					Duration:           time.Hour,
+					InitialPoolWeights: startPoolWeightAssets,
+					TargetPoolWeights:  defaultPoolAssetsCopy,
+				},
+			}),
+		},
+		"non weighted pool": {
+			poolId: suite.prepareCustomStableswapPool(
+				defaultAcctFunds,
+				stableswap.PoolParams{
+					SwapFee: defaultSwapFee,
+					ExitFee: defaultExitFee,
+				},
+				sdk.NewCoins(sdk.NewCoin(defaultAcctFunds[0].Denom, defaultAcctFunds[0].Amount.QuoRaw(2)), sdk.NewCoin(defaultAcctFunds[1].Denom, defaultAcctFunds[1].Amount.QuoRaw(2))),
+				[]uint64{1, 1},
+			)},
+	}
+
+	for name, tc := range tests {
+		suite.Run(name, func() {
+			k := suite.App.GAMMKeeper
+			ctx := suite.Ctx.WithBlockTime(time.Unix(blockTime, 0))
+
+			pool, err := k.GetPoolAndPoke(ctx, tc.poolId)
+
+			suite.Require().NoError(err)
+			suite.Require().Equal(tc.poolId, pool.GetId())
+
+			if tc.isPokePool {
+				pokePool, ok := pool.(types.WeightedPoolExtension)
+				suite.Require().True(ok)
+
+				poolAssetWeight0, err := pokePool.GetTokenWeight(startPoolWeightAssets[0].Token.Denom)
+				suite.Require().NoError(err)
+
+				poolAssetWeight1, err := pokePool.GetTokenWeight(startPoolWeightAssets[1].Token.Denom)
+				suite.Require().NoError(err)
+
+				suite.Require().NotEqual(startPoolWeightAssets[0].Weight, poolAssetWeight0)
+				suite.Require().NotEqual(startPoolWeightAssets[1].Weight, poolAssetWeight1)
+				return
+			}
+
+			_, ok := pool.(types.WeightedPoolExtension)
+			suite.Require().False(ok)
+		})
+	}
+}

x/gamm/module.go

@@ -29,6 +29,7 @@ import (
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/client/cli"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/keeper"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/balancer"
+	"github.com/osmosis-labs/osmosis/v12/x/gamm/pool-models/stableswap"
 	simulation "github.com/osmosis-labs/osmosis/v12/x/gamm/simulation"
 	"github.com/osmosis-labs/osmosis/v12/x/gamm/types"
 )
@@ -88,7 +89,7 @@ func (b AppModuleBasic) GetQueryCmd() *cobra.Command {
 func (AppModuleBasic) RegisterInterfaces(registry codectypes.InterfaceRegistry) {
 	types.RegisterInterfaces(registry)
 	balancer.RegisterInterfaces(registry)
-	// stableswap.RegisterInterfaces(registry)
+	stableswap.RegisterInterfaces(registry)
 }
 
 type AppModule struct {

x/gamm/pool-models/balancer/pool.go

@@ -29,6 +29,7 @@ const (
 var (
 	_ types.PoolI                  = &Pool{}
 	_ types.PoolAmountOutExtension = &Pool{}
+	_ types.WeightedPoolExtension  = &Pool{}
 )
 
 // NewPool returns a weighted CPMM pool with the provided parameters, and initial assets.

x/gamm/pool-models/stableswap/pool.go

@@ -4,7 +4,6 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
-	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
@@ -317,9 +316,6 @@ func (p Pool) CalcExitPoolCoinsFromShares(ctx sdk.Context, exitingShares sdk.Int
 	return cfmm_common.CalcExitPool(ctx, &p, exitingShares, exitFee)
 }
 
-// no-op for stableswap
-func (p *Pool) PokePool(blockTime time.Time) {}
-
 // SetStableSwapScalingFactors sets scaling factors for pool to the given amount
 // It should only be able to be successfully called by the pool's ScalingFactorGovernor
 // TODO: move commented test for this function from x/gamm/keeper/pool_service_test.go once a pool_test.go file has been created for stableswap

x/gamm/types/pool.go

@@ -79,10 +79,6 @@ type PoolI interface {
 	// CalcExitPoolCoinsFromShares returns how many coins ExitPool would return on these arguments.
 	// This does not mutate the pool, or state.
 	CalcExitPoolCoinsFromShares(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error)
-
-	// PokePool determines if a pool's weights need to be updated and updates
-	// them if so.
-	PokePool(blockTime time.Time)
 }
 
 // PoolAmountOutExtension is an extension of the PoolI
@@ -126,6 +122,19 @@ type PoolAmountOutExtension interface {
 	IncreaseLiquidity(sharesOut sdk.Int, coinsIn sdk.Coins)
 }
 
+// WeightedPoolExtension is an extension of the PoolI interface
+// That defines an additional API for handling the pool's weights.
+type WeightedPoolExtension interface {
+	PoolI
+
+	// PokePool determines if a pool's weights need to be updated and updates
+	// them if so.
+	PokePool(blockTime time.Time)
+
+	// GetTokenWeight returns the weight of the specified token in the pool.
+	GetTokenWeight(denom string) (sdk.Int, error)
+}
+
 func NewPoolAddress(poolId uint64) sdk.AccAddress {
 	key := append([]byte("pool"), sdk.Uint64ToBigEndian(poolId)...)
 	return address.Module(ModuleName, key)