Add migration function for stateful order expirations

protocol/x/clob/keeper/stateful_order_state.go

@@ -230,6 +230,7 @@ func (k Keeper) RemoveExpiredStatefulOrders(ctx sdk.Context, blockTime time.Time
 			[]byte(fmt.Sprintf(types.StatefulOrdersExpirationsKeyPrefix, sdk.FormatTimeString(blockTime))),
 		),
 	)
+	defer it.Close()
 	for ; it.Valid(); it.Next() {
 		var orderId types.OrderId
 		k.cdc.MustUnmarshal(it.Value(), &orderId)
@@ -454,3 +455,49 @@ func (k Keeper) SetStatefulOrderCount(
 		)
 	}
 }
+
+// MigrateOrderExpirationState migrates order expiration state from slices based on time to individual keys.
+// Deprecated: Only intended for use in the v5.2 upgrade handler.
+func (k Keeper) MigrateOrderExpirationState(ctx sdk.Context) {
+	store := ctx.KVStore(k.storeKey)
+	prefixStore := prefix.NewStore(store, []byte(types.LegacyStatefulOrdersTimeSlicePrefix))
+	it := prefixStore.Iterator(nil, nil)
+	defer it.Close()
+	for ; it.Valid(); it.Next() {
+		time, err := sdk.ParseTimeBytes(it.Key())
+		if err != nil {
+			panic(fmt.Sprintf("migration failed due to malformed time: %s", it.Key()))
+		}
+		var orders types.StatefulOrderTimeSliceValue
+		k.cdc.MustUnmarshal(it.Value(), &orders)
+		for _, orderId := range orders.OrderIds {
+			k.AddStatefulOrderIdExpiration(ctx, time, orderId)
+		}
+		prefixStore.Delete(it.Key())
+	}
+}
+
+// LegacySetStatefulOrdersTimeSliceInState sets a sorted list of order IDs in state at a `goodTilBlockTime`.
+// This function automatically sorts the order IDs before writing them to state.
+// Deprecated: Only intended for testing MigrateOrderExpirationState.
+func (k Keeper) LegacySetStatefulOrdersTimeSliceInState(
+	ctx sdk.Context,
+	goodTilBlockTime time.Time,
+	orderIds []types.OrderId,
+) {
+	// Sort the order IDs.
+	types.MustSortAndHaveNoDuplicates(orderIds)
+
+	statefulOrderPlacement := types.StatefulOrderTimeSliceValue{
+		OrderIds: orderIds,
+	}
+	b := k.cdc.MustMarshal(&statefulOrderPlacement)
+	store := prefix.NewStore(
+		ctx.KVStore(k.storeKey),
+		[]byte(types.LegacyStatefulOrdersTimeSlicePrefix),
+	)
+	store.Set(
+		sdk.FormatTimeBytes(goodTilBlockTime),
+		b,
+	)
+}

protocol/x/clob/keeper/stateful_order_state_test.go

@@ -1072,3 +1072,59 @@ func TestGetAllStatefulOrders(t *testing.T) {
 		})
 	}
 }
+
+func TestMigrateOrderExpirationState(t *testing.T) {
+	tests := map[string]struct {
+		timeSlicesToOrderIds map[time.Time][]types.OrderId
+	}{
+		"Multiple time slices": {
+			timeSlicesToOrderIds: map[time.Time][]types.OrderId{
+				constants.Time_21st_Feb_2021: {
+					constants.ConditionalOrder_Alice_Num0_Id0_Clob0_Buy5_Price10_GTBT15_StopLoss20.OrderId,
+					constants.ConditionalOrder_Alice_Num1_Id0_Clob0_Sell5_Price10_GTB15.OrderId,
+				},
+				constants.Time_21st_Feb_2021.Add(1): {
+					constants.ConditionalOrder_Alice_Num1_Id1_Clob0_Sell50_Price5_GTB30.OrderId,
+				},
+				constants.Time_21st_Feb_2021.Add(77): {
+					constants.LongTermOrder_Alice_Num1_Id1_Clob0_Sell25_Price30_GTBT10.OrderId,
+					constants.LongTermOrder_Alice_Num0_Id0_Clob0_Buy5_Price10_GTBT20.OrderId,
+				},
+			},
+		},
+		"No time slices": {
+			timeSlicesToOrderIds: map[time.Time][]types.OrderId{},
+		},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			memClob := memclob.NewMemClobPriceTimePriority(false)
+			ks := keepertest.NewClobKeepersTestContext(t, memClob, &mocks.BankKeeper{}, &mocks.IndexerEventManager{})
+
+			for timestamp, orderIds := range tc.timeSlicesToOrderIds {
+				ks.ClobKeeper.LegacySetStatefulOrdersTimeSliceInState(ks.Ctx, timestamp, orderIds)
+			}
+
+			ks.ClobKeeper.MigrateOrderExpirationState(ks.Ctx)
+
+			oldStore := prefix.NewStore(
+				ks.Ctx.KVStore(ks.StoreKey),
+				[]byte(types.LegacyStatefulOrdersTimeSlicePrefix),
+			)
+			it := oldStore.Iterator(nil, nil)
+			require.False(t, it.Valid())
+
+			for goodTilTime, expectedOrderIds := range tc.timeSlicesToOrderIds {
+				orderIds := ks.ClobKeeper.GetStatefulOrderIdExpirations(ks.Ctx, goodTilTime)
+				require.ElementsMatch(
+					t,
+					expectedOrderIds,
+					orderIds,
+					"Mismatch of order IDs for timestamp",
+					goodTilTime.String(),
+				)
+			}
+		})
+	}
+}