refactor: add PokeTokenWeights to Pool interface

x/gamm/keeper/pool.go

@@ -34,7 +34,7 @@ func (k Keeper) GetPool(ctx sdk.Context, poolId uint64) (types.PoolI, error) {
 		return nil, err
 	}
 
-	// pool.PokeTokenWeights(ctx.BlockTime())
+	pool.PokeTokenWeights(ctx.BlockTime())
 
 	return pool, nil
 }

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

@@ -380,52 +380,62 @@ func (pa *Pool) updateAllWeights(newWeights []PoolAsset) {
 // PokeTokenWeights checks to see if the pool's token weights need to be updated,
 // and if so, does so.
 func (pa *Pool) PokeTokenWeights(blockTime time.Time) {
-	// Pool weights aren't changing, do nothing.
-	poolWeightsChanging := (pa.PoolParams.SmoothWeightChangeParams != nil)
+	// check if pool weights didn't change
+	poolWeightsChanging := pa.PoolParams.SmoothWeightChangeParams != nil
 	if !poolWeightsChanging {
 		return
 	}
-	// Pool weights are changing.
+
 	// TODO: Add intra-block cache check that we haven't already poked
 	// the block yet.
 	params := *pa.PoolParams.SmoothWeightChangeParams
 
-	// the weights w(t) for the pool at time `t` is the following:
-	//   t <= start_time: w(t) = initial_pool_weights
-	//   start_time < t <= start_time + duration:
+	// The weights w(t) for the pool at time `t` is defined in one of three
+	// possible ways:
+	//
+	// 1. t <= start_time: w(t) = initial_pool_weights
+	//
+	// 2. start_time < t <= start_time + duration:
 	//     w(t) = initial_pool_weights + (t - start_time) *
 	//       (target_pool_weights - initial_pool_weights) / (duration)
-	//   t > start_time + duration: w(t) = target_pool_weights
-
-	// t <= StartTime
-	if blockTime.Before(params.StartTime) || params.StartTime.Equal(blockTime) {
-		// Do nothing
+	//
+	// 3. t > start_time + duration: w(t) = target_pool_weights
+	switch {
+	case blockTime.Before(params.StartTime) || params.StartTime.Equal(blockTime):
+		// case 1: t <= start_time
 		return
-	} else if blockTime.After(params.StartTime.Add(params.Duration)) {
-		// t > start_time + duration
+
+	case blockTime.After(params.StartTime.Add(params.Duration)):
+		// case 2: start_time < t <= start_time + duration:
+
 		// Update weights to be the target weights.
-		// TODO: When we add support for adding new assets via this method,
-		// 		 Ensure the new asset has some token sent with it.
+		//
+		// TODO: When we add support for adding new assets via this method, ensure
+		// the new asset has some token sent with it.
 		pa.updateAllWeights(params.TargetPoolWeights)
-		// We've finished updating weights, so delete this parameter
+
+		// we've finished updating the weights, so reset the following fields
 		pa.PoolParams.SmoothWeightChangeParams = nil
 		return
-	} else {
-		//	w(t) = initial_pool_weights + (t - start_time) *
-		//       (target_pool_weights - initial_pool_weights) / (duration)
-		// We first compute percent duration elapsed = (t - start_time) / duration, via Unix time.
+
+	default:
+		// case 3: t > start_time + duration: w(t) = target_pool_weights
+
 		shiftedBlockTime := blockTime.Sub(params.StartTime).Milliseconds()
 		percentDurationElapsed := sdk.NewDec(shiftedBlockTime).QuoInt64(params.Duration.Milliseconds())
-		// If the duration elapsed is equal to the total time,
-		// or a rounding error makes it seem like it is, just set to target weight
+
+		// If the duration elapsed is equal to the total time, or a rounding error
+		// makes it seem like it is, just set to target weight.
 		if percentDurationElapsed.GTE(sdk.OneDec()) {
 			pa.updateAllWeights(params.TargetPoolWeights)
 			return
 		}
+
+		// below will be auto-truncated according to internal weight precision routine
 		totalWeightsDiff := subPoolAssetWeights(params.TargetPoolWeights, params.InitialPoolWeights)
-		// Below will be auto-truncated according to internal weight precision routine.
 		scaledDiff := poolAssetsMulDec(totalWeightsDiff, percentDurationElapsed)
 		updatedWeights := addPoolAssetWeights(params.InitialPoolWeights, scaledDiff)
+
 		pa.updateAllWeights(updatedWeights)
 	}
 }

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

@@ -2,7 +2,8 @@ package stableswap
 
 import (
 	"encoding/json"
-	fmt "fmt"
+	"fmt"
+	"time"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 
@@ -84,3 +85,7 @@ func (pa *Pool) ExitPool(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (e
 func (pa Pool) CalcExitPoolShares(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error) {
 	return sdk.Coins{}, types.ErrNotImplemented
 }
+
+func (pa *Pool) PokeTokenWeights(blockTime time.Time) {
+	panic("not implemented yet")
+}

x/gamm/types/pool.go

@@ -1,8 +1,11 @@
 package types
 
 import (
+	"time"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	proto "github.com/gogo/protobuf/proto"
+
 	"github.com/osmosis-labs/osmosis/v7/v043_temp/address"
 )
 
@@ -65,6 +68,10 @@ type PoolI interface {
 	// CalcExitPoolShares returns how many coins ExitPool would return on these arguments.
 	// This does not mutate the pool, or state.
 	CalcExitPoolShares(ctx sdk.Context, numShares sdk.Int, exitFee sdk.Dec) (exitedCoins sdk.Coins, err error)
+
+	// PokeTokenWeights determines if a pool's weights need to be updated and
+	// updates them if so.
+	PokeTokenWeights(blockTime time.Time)
 }
 
 func NewPoolAddress(poolId uint64) sdk.AccAddress {