feat: fix tranche key ordering

x/dex/genesis.go

@@ -21,7 +21,7 @@ func InitGenesis(ctx sdk.Context, k keeper.Keeper, genState types.GenesisState)
 			k.SetLimitOrderTranche(ctx, tranche)
 			if tranche.HasExpiration() {
 				// re-create expiration record
-				loExpiration := keeper.NewLimitOrderExpiration(tranche)
+				loExpiration := types.NewLimitOrderExpiration(tranche)
 				k.SetLimitOrderExpiration(ctx, loExpiration)
 			}
 		}

x/dex/keeper/inactive_limit_order_tranche.go

@@ -58,6 +58,12 @@ func (k Keeper) GetAllInactiveLimitOrderTranche(ctx sdk.Context) (list []*types.
 	return list
 }
 
+// GetInactiveLimitOrderTrancheIterator returns a store iterator over all inactive limit order tranches.
+func (k Keeper) GetInactiveLimitOrderTrancheIterator(ctx sdk.Context) storetypes.Iterator {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.KeyPrefix(types.InactiveLimitOrderTrancheKeyPrefix))
+	return storetypes.KVStorePrefixIterator(store, []byte{})
+}
+
 // UpdateInactiveTranche handles the logic for all updates to InactiveLimitOrderTranches
 // It will delete an InactiveTranche if there is no remaining MakerReserves or TakerReserves
 func (k Keeper) UpdateInactiveTranche(sdkCtx sdk.Context, tranche *types.LimitOrderTranche) {

x/dex/keeper/limit_order_expiration.go

@@ -10,19 +10,6 @@ import (
 	"github.com/neutron-org/neutron/v11/x/dex/types"
 )
 
-// Creates a new LimitOrderExpiration struct based on a LimitOrderTranche
-func NewLimitOrderExpiration(tranche *types.LimitOrderTranche) *types.LimitOrderExpiration {
-	trancheExpiry := tranche.ExpirationTime
-	if trancheExpiry == nil {
-		panic("Cannot create LimitOrderExpiration from tranche with nil ExpirationTime")
-	}
-
-	return &types.LimitOrderExpiration{
-		TrancheRef:     tranche.Key.KeyMarshal(),
-		ExpirationTime: *tranche.ExpirationTime,
-	}
-}
-
 // SetLimitOrderExpiration set a specific goodTilRecord in the store from its index
 func (k Keeper) SetLimitOrderExpiration(
 	ctx sdk.Context,

x/dex/keeper/limit_order_tranche.go

@@ -2,7 +2,6 @@ package keeper
 
 import (
 	"encoding/binary"
-	"fmt"
 	"time"
 
 	"cosmossdk.io/math"
@@ -208,10 +207,11 @@ func (k Keeper) GetAllLimitOrderTrancheAtIndex(
 	return trancheList
 }
 
-func (k Keeper) NewTrancheKey(ctx sdk.Context) string {
+// NextTrancheKey increments the tranche count and returns the next tranche key.
+func (k Keeper) NextTrancheKey(ctx sdk.Context) string {
 	trancheCount := k.GetTrancheCount(ctx)
 	k.IncrementTrancheCount(ctx)
-	return fmt.Sprintf("tk-%d", trancheCount)
+	return types.NewTrancheKey(trancheCount)
 }
 
 func (k Keeper) GetOrInitPlaceTranche(ctx sdk.Context,
@@ -232,15 +232,15 @@ func (k Keeper) GetOrInitPlaceTranche(ctx sdk.Context,
 		limitOrderTrancheKey := &types.LimitOrderTrancheKey{
 			TradePairId:           tradePairID,
 			TickIndexTakerToMaker: tickIndexTakerToMaker,
-			TrancheKey:            k.NewTrancheKey(ctx),
+			TrancheKey:            k.NextTrancheKey(ctx),
 		}
 		placeTranche, err = NewLimitOrderTranche(limitOrderTrancheKey, &JITGoodTilTime)
 		ctx.EventManager().EmitEvents(types.GetEventsIncTotalOrders(tradePairID))
 	case types.LimitOrderType_GOOD_TIL_TIME:
 		limitOrderTrancheKey := &types.LimitOrderTrancheKey{
 			TradePairId:           tradePairID,
 			TickIndexTakerToMaker: tickIndexTakerToMaker,
-			TrancheKey:            k.NewTrancheKey(ctx),
+			TrancheKey:            k.NextTrancheKey(ctx),
 		}
 		placeTranche, err = NewLimitOrderTranche(limitOrderTrancheKey, goodTil)
 		ctx.EventManager().EmitEvents(types.GetEventsIncExpiringOrders(tradePairID))
@@ -250,7 +250,7 @@ func (k Keeper) GetOrInitPlaceTranche(ctx sdk.Context,
 			limitOrderTrancheKey := &types.LimitOrderTrancheKey{
 				TradePairId:           tradePairID,
 				TickIndexTakerToMaker: tickIndexTakerToMaker,
-				TrancheKey:            k.NewTrancheKey(ctx),
+				TrancheKey:            k.NextTrancheKey(ctx),
 			}
 			placeTranche, err = NewLimitOrderTranche(limitOrderTrancheKey, nil)
 			ctx.EventManager().EmitEvents(types.GetEventsIncTotalOrders(tradePairID))

x/dex/keeper/limit_order_tranche_test.go

@@ -1,7 +1,6 @@
 package keeper_test
 
 import (
-	"fmt"
 	"testing"
 
 	"cosmossdk.io/math"
@@ -24,7 +23,7 @@ func createNLimitOrderTranches(
 		items[i] = types.MustNewLimitOrderTranche(
 			"TokenA",
 			"TokenB",
-			keeper.NewTrancheKey(ctx),
+			keeper.NextTrancheKey(ctx),
 			int64(i),
 			math.ZeroInt(),
 			math.ZeroInt(),
@@ -47,7 +46,7 @@ func TestGetLimitOrderTranche(t *testing.T) {
 			nullify.Fill(item),
 			nullify.Fill(rst),
 		)
-		require.Equal(t, fmt.Sprintf("tk-%d", n), item.Key.TrancheKey)
+		require.Equal(t, types.NewTrancheKey(uint64(n)), item.Key.TrancheKey)
 	}
 }
 

x/dex/keeper/limit_order_tranche_user_test.go

@@ -1,7 +1,6 @@
 package keeper_test
 
 import (
-	"fmt"
 	"strconv"
 	"testing"
 
@@ -20,7 +19,7 @@ func createNLimitOrderTrancheUser(keeper *keeper.Keeper, ctx sdk.Context, n int)
 	items := make([]*types.LimitOrderTrancheUser, n)
 	for i := range items {
 		val := &types.LimitOrderTrancheUser{
-			TrancheKey:            fmt.Sprintf("tk-%d", i),
+			TrancheKey:            types.NewTrancheKey(uint64(i)),
 			Address:               strconv.Itoa(i),
 			TradePairId:           &types.TradePairID{MakerDenom: "TokenA", TakerDenom: "TokenB"},
 			TickIndexTakerToMaker: int64(i),

x/dex/keeper/migrations.go

@@ -9,6 +9,7 @@ import (
 	v6 "github.com/neutron-org/neutron/v11/x/dex/migrations/v6"
 	v7 "github.com/neutron-org/neutron/v11/x/dex/migrations/v7"
 	v8 "github.com/neutron-org/neutron/v11/x/dex/migrations/v8"
+	v9 "github.com/neutron-org/neutron/v11/x/dex/migrations/v9"
 )
 
 // Migrator is a struct for handling in-place store migrations.
@@ -50,3 +51,8 @@ func (m Migrator) Migrate6to7(ctx sdk.Context) error {
 func (m Migrator) Migrate7to8(ctx sdk.Context) error {
 	return v8.MigrateStore(ctx, m.keeper.cdc, m.keeper.storeKey)
 }
+
+// Migrate8to9 migrates from version 8 to 9.
+func (m Migrator) Migrate8to9(ctx sdk.Context) error {
+	return v9.MigrateStore(ctx, m.keeper.cdc, &m.keeper)
+}

x/dex/keeper/place_limit_order.go

@@ -178,7 +178,7 @@ func (k Keeper) ExecutePlaceLimitOrder(
 		trancheUser.SharesOwned = trancheUser.SharesOwned.Add(amountToPlace)
 
 		if orderType.HasExpiration() {
-			goodTilRecord := NewLimitOrderExpiration(placeTranche)
+			goodTilRecord := types.NewLimitOrderExpiration(placeTranche)
 			k.SetLimitOrderExpiration(ctx, goodTilRecord)
 			ctx.GasMeter().ConsumeGas(types.ExpiringLimitOrderGas, "Expiring LimitOrder Fee")
 		}

x/dex/migrations/v9/store.go

@@ -0,0 +1,229 @@
+package v9
+
+import (
+	"fmt"
+	"strconv"
+	"strings"
+	"time"
+
+	storetypes "cosmossdk.io/store/types"
+	"github.com/cosmos/cosmos-sdk/codec"
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	dextypes "github.com/neutron-org/neutron/v11/x/dex/types"
+)
+
+// dexKeeper defines an interface with dex keeper methods required for the migration. It is defined
+// to avoid import loop (x/dex/migrations <-> x/dex/keeper).
+type dexKeeper interface {
+	GetAllLimitOrderExpiration(ctx sdk.Context) (list []*dextypes.LimitOrderExpiration)
+	GetLimitOrderTrancheByKey(ctx sdk.Context, key []byte) (tranche *dextypes.LimitOrderTranche, found bool)
+	RemoveLimitOrderExpiration(ctx sdk.Context, goodTilDate time.Time, trancheRef []byte)
+	SetLimitOrderExpiration(ctx sdk.Context, goodTilRecord *dextypes.LimitOrderExpiration)
+	GetAllTickLiquidity(ctx sdk.Context) (list []*dextypes.TickLiquidity)
+	RemoveLimitOrderTranche(ctx sdk.Context, trancheKey *dextypes.LimitOrderTrancheKey)
+	SetLimitOrderTranche(ctx sdk.Context, tranche *dextypes.LimitOrderTranche)
+	GetInactiveLimitOrderTrancheIterator(ctx sdk.Context) storetypes.Iterator
+	RemoveInactiveLimitOrderTranche(ctx sdk.Context, limitOrderTrancheKey *dextypes.LimitOrderTrancheKey)
+	SetInactiveLimitOrderTranche(ctx sdk.Context, limitOrderTranche *dextypes.LimitOrderTranche)
+	GetAllLimitOrderTrancheUser(ctx sdk.Context) (list []*dextypes.LimitOrderTrancheUser)
+	RemoveLimitOrderTrancheUser(ctx sdk.Context, trancheUser *dextypes.LimitOrderTrancheUser)
+	SetLimitOrderTrancheUser(ctx sdk.Context, limitOrderTrancheUser *dextypes.LimitOrderTrancheUser)
+}
+
+// MigrateStore performs in-place store migrations. It reconstructs the tranche keys for limit order
+// expirations, tranches, inactive tranches, and tranche user lists.
+func MigrateStore(ctx sdk.Context, cdc codec.BinaryCodec, dexKeeper dexKeeper) error {
+	ctx.Logger().Info("Starting dex store migration...")
+
+	ctx.Logger().Info("Reconstructing tranche keys...")
+	if err := ReconstructTrancheKeys(ctx, cdc, dexKeeper); err != nil {
+		return err
+	}
+
+	ctx.Logger().Info("Dex store migration completed")
+	return nil
+}
+
+func ReconstructTrancheKeys(ctx sdk.Context, cdc codec.BinaryCodec, k dexKeeper) error {
+	ctx.Logger().Info("Reconstructing LO expirations...")
+	if err := reconstructLoExpirations(ctx, k); err != nil {
+		return fmt.Errorf("failed to reconstruct LO expirations: %w", err)
+	}
+	ctx.Logger().Info("Done")
+
+	ctx.Logger().Info("Reconstructing LO tranches...")
+	if err := reconstructLoTranches(ctx, k); err != nil {
+		return fmt.Errorf("failed to reconstruct LO tranches: %w", err)
+	}
+	ctx.Logger().Info("Done")
+
+	ctx.Logger().Info("Reconstructing inactive LO tranches...")
+	if err := reconstructInactiveLoTranches(ctx, cdc, k); err != nil {
+		return fmt.Errorf("failed to reconstruct inactive LO tranches: %w", err)
+	}
+	ctx.Logger().Info("Done")
+
+	ctx.Logger().Info("Reconstructing LO tranche user lists...")
+	if err := reconstructLoTrancheUserLists(ctx, k); err != nil {
+		return fmt.Errorf("failed to reconstruct LO tranche user lists: %w", err)
+	}
+	ctx.Logger().Info("Done")
+
+	return nil
+}
+
+func reconstructLoExpirations(ctx sdk.Context, k dexKeeper) error {
+	allExpirations := k.GetAllLimitOrderExpiration(ctx) // total count varies but is expected to be small or even 0
+
+	expirationsToRemove := make([]dextypes.LimitOrderExpiration, 0)
+	expirationsToUpdate := make([]dextypes.LimitOrderExpiration, 0)
+	for _, expiration := range allExpirations {
+		tranche, found := k.GetLimitOrderTrancheByKey(ctx, expiration.TrancheRef)
+		if !found {
+			return fmt.Errorf("limit order tranche not found for expiration.TrancheRef %s", expiration.TrancheRef)
+		}
+
+		if !strings.HasPrefix(tranche.Key.TrancheKey, "tk-") {
+			continue
+		}
+
+		expirationsToRemove = append(expirationsToRemove, *expiration)
+
+		trancheIdxStr := strings.TrimPrefix(tranche.Key.TrancheKey, "tk-")
+		trancheIdx, err := strconv.ParseUint(trancheIdxStr, 10, 64)
+		if err != nil {
+			return fmt.Errorf("failed to parse tranche idx %s: %w", trancheIdxStr, err)
+		}
+		tranche.Key.TrancheKey = dextypes.NewTrancheKey(trancheIdx)
+		expirationsToUpdate = append(expirationsToUpdate, *dextypes.NewLimitOrderExpiration(tranche))
+	}
+
+	if len(expirationsToRemove) != len(expirationsToUpdate) {
+		return fmt.Errorf("mismatch in LO expirations to remove and update counts: %d != %d", len(expirationsToRemove), len(expirationsToUpdate))
+	}
+
+	for _, expiration := range expirationsToRemove {
+		k.RemoveLimitOrderExpiration(ctx, expiration.ExpirationTime, expiration.TrancheRef)
+	}
+	for _, expiration := range expirationsToUpdate {
+		k.SetLimitOrderExpiration(ctx, &expiration)
+	}
+	ctx.Logger().Info("LO expiration keys reconstructed", "count", len(expirationsToUpdate))
+
+	return nil
+}
+
+func reconstructLoTranches(ctx sdk.Context, k dexKeeper) error {
+	tickLiquidities := k.GetAllTickLiquidity(ctx) // there are only 600-ish entries, so getting all is fine
+
+	loTrancheKeysToRemove := make([]dextypes.LimitOrderTrancheKey, 0)
+	loTranchesToUpdate := make([]dextypes.LimitOrderTranche, 0)
+	for _, tickLiquidity := range tickLiquidities {
+		if loTranche := tickLiquidity.GetLimitOrderTranche(); loTranche != nil {
+			if !strings.HasPrefix(loTranche.Key.TrancheKey, "tk-") {
+				continue
+			}
+
+			loTrancheKeysToRemove = append(loTrancheKeysToRemove, *loTranche.Key)
+
+			trancheIdxStr := strings.TrimPrefix(loTranche.Key.TrancheKey, "tk-")
+			trancheIdx, err := strconv.ParseUint(trancheIdxStr, 10, 64)
+			if err != nil {
+				return fmt.Errorf("failed to parse tranche idx %s: %w", trancheIdxStr, err)
+			}
+			loTranche.Key.TrancheKey = dextypes.NewTrancheKey(trancheIdx)
+			loTranchesToUpdate = append(loTranchesToUpdate, *loTranche)
+		}
+	}
+
+	if len(loTrancheKeysToRemove) != len(loTranchesToUpdate) {
+		return fmt.Errorf("mismatch in LO tranches to remove and update counts: %d != %d", len(loTrancheKeysToRemove), len(loTranchesToUpdate))
+	}
+
+	for _, loTrancheKey := range loTrancheKeysToRemove {
+		k.RemoveLimitOrderTranche(ctx, &loTrancheKey)
+	}
+	for _, loTranche := range loTranchesToUpdate {
+		k.SetLimitOrderTranche(ctx, &loTranche)
+	}
+	ctx.Logger().Info("LO tranche keys reconstructed", "count", len(loTranchesToUpdate))
+
+	return nil
+}
+
+func reconstructInactiveLoTranches(ctx sdk.Context, cdc codec.BinaryCodec, k dexKeeper) error {
+	iter := k.GetInactiveLimitOrderTrancheIterator(ctx) // there are more than 400k entries -> iterating
+
+	inactiveKeysToRemove := make([]dextypes.LimitOrderTrancheKey, 0)
+	inactiveTranchesToUpdate := make([]dextypes.LimitOrderTranche, 0)
+	for ; iter.Valid(); iter.Next() {
+		var tranche dextypes.LimitOrderTranche
+		cdc.MustUnmarshal(iter.Value(), &tranche)
+
+		if !strings.HasPrefix(tranche.Key.TrancheKey, "tk-") {
+			continue
+		}
+
+		inactiveKeysToRemove = append(inactiveKeysToRemove, *tranche.Key)
+
+		trancheIdxStr := strings.TrimPrefix(tranche.Key.TrancheKey, "tk-")
+		trancheIdx, err := strconv.ParseUint(trancheIdxStr, 10, 64)
+		if err != nil {
+			iter.Close() //nolint:errcheck,gosec
+			return fmt.Errorf("failed to parse tranche idx %s: %w", trancheIdxStr, err)
+		}
+		tranche.Key.TrancheKey = dextypes.NewTrancheKey(trancheIdx)
+		inactiveTranchesToUpdate = append(inactiveTranchesToUpdate, tranche)
+	}
+	iter.Close() //nolint:errcheck,gosec
+
+	if len(inactiveKeysToRemove) != len(inactiveTranchesToUpdate) {
+		return fmt.Errorf("mismatch in inactive LO tranches to remove and update counts: %d != %d", len(inactiveKeysToRemove), len(inactiveTranchesToUpdate))
+	}
+
+	for _, key := range inactiveKeysToRemove {
+		k.RemoveInactiveLimitOrderTranche(ctx, &key)
+	}
+	for _, tranche := range inactiveTranchesToUpdate {
+		k.SetInactiveLimitOrderTranche(ctx, &tranche)
+	}
+	ctx.Logger().Info("inactive LO tranche keys reconstructed", "count", len(inactiveTranchesToUpdate))
+
+	return nil
+}
+
+func reconstructLoTrancheUserLists(ctx sdk.Context, k dexKeeper) error {
+	allUsers := k.GetAllLimitOrderTrancheUser(ctx) // there are only 300-ish entries, so getting all is fine
+
+	usersToRemove := make([]dextypes.LimitOrderTrancheUser, 0)
+	usersToUpdate := make([]dextypes.LimitOrderTrancheUser, 0)
+	for _, user := range allUsers {
+		if !strings.HasPrefix(user.TrancheKey, "tk-") {
+			continue
+		}
+
+		usersToRemove = append(usersToRemove, *user)
+
+		trancheIdxStr := strings.TrimPrefix(user.TrancheKey, "tk-")
+		trancheIdx, err := strconv.ParseUint(trancheIdxStr, 10, 64)
+		if err != nil {
+			return fmt.Errorf("failed to parse tranche idx %s: %w", trancheIdxStr, err)
+		}
+		user.TrancheKey = dextypes.NewTrancheKey(trancheIdx)
+		usersToUpdate = append(usersToUpdate, *user)
+	}
+
+	if len(usersToRemove) != len(usersToUpdate) {
+		return fmt.Errorf("mismatch in LO tranche user keys to remove and update counts: %d != %d", len(usersToRemove), len(usersToUpdate))
+	}
+
+	for _, user := range usersToRemove {
+		k.RemoveLimitOrderTrancheUser(ctx, &user)
+	}
+	for _, user := range usersToUpdate {
+		k.SetLimitOrderTrancheUser(ctx, &user)
+	}
+	ctx.Logger().Info("LO tranche user keys reconstructed", "count", len(usersToUpdate))
+
+	return nil
+}