Delete contract-related states unpon unregister

x/dex/keeper/common.go

@@ -0,0 +1,12 @@
+package keeper
+
+import sdk "github.com/cosmos/cosmos-sdk/types"
+
+func (k Keeper) removeAllForPrefix(ctx sdk.Context, prefix []byte) {
+	store := ctx.KVStore(k.storeKey)
+	iter := sdk.KVStorePrefixIterator(store, prefix)
+	defer iter.Close()
+	for ; iter.Valid(); iter.Next() {
+		store.Delete(iter.Key())
+	}
+}

x/dex/keeper/long_book.go

@@ -104,6 +104,10 @@ func (k Keeper) GetLongAllocationForOrderID(ctx sdk.Context, contractAddr string
 	return nil, false
 }
 
+func (k Keeper) RemoveAllLongBooksForContract(ctx sdk.Context, contractAddr string) {
+	k.removeAllForPrefix(ctx, types.OrderBookContractPrefix(true, contractAddr))
+}
+
 func GetKeyForLongBook(longBook types.LongBook) []byte {
 	return GetKeyForPrice(longBook.Entry.Price)
 }

x/dex/keeper/match_result.go

@@ -35,3 +35,11 @@ func (k Keeper) GetMatchResultState(ctx sdk.Context, contractAddr string) (*type
 	}
 	return &result, true
 }
+
+func (k Keeper) DeleteMatchResultState(ctx sdk.Context, contractAddr string) {
+	store := prefix.NewStore(
+		ctx.KVStore(k.storeKey),
+		types.MatchResultPrefix(contractAddr),
+	)
+	store.Delete([]byte(MatchResultKey))
+}

x/dex/keeper/msgserver/msg_server_unregister_contract.go

@@ -31,5 +31,12 @@ func (k msgServer) UnregisterContract(goCtx context.Context, msg *types.MsgUnreg
 		return nil, err
 	}
 	k.DeleteContract(ctx, msg.ContractAddr)
+	k.RemoveAllLongBooksForContract(ctx, msg.ContractAddr)
+	k.RemoveAllShortBooksForContract(ctx, msg.ContractAddr)
+	k.RemoveAllPricesForContract(ctx, msg.ContractAddr)
+	k.DeleteMatchResultState(ctx, msg.ContractAddr)
+	k.DeleteNextOrderID(ctx, msg.ContractAddr)
+	k.DeleteAllRegisteredPairsForContract(ctx, msg.ContractAddr)
+	k.RemoveAllTriggeredOrders(ctx, msg.ContractAddr)
 	return &types.MsgUnregisterContractResponse{}, nil
 }

x/dex/keeper/msgserver/msg_server_unregister_contract_test.go

@@ -67,6 +67,25 @@ func TestUnregisterContractSetSiblings(t *testing.T) {
 	require.Equal(t, int64(8900000), balance.Amount.Int64())
 
 	handler := dex.NewHandler(keeper)
+	tickSize := sdk.OneDec()
+	_, err = handler(ctx, &types.MsgRegisterPairs{
+		Creator: testAccount.String(),
+		Batchcontractpair: []types.BatchContractPair{
+			{
+				ContractAddr: contractAddr.String(),
+				Pairs: []*types.Pair{
+					{
+						PriceDenom:       "usei",
+						AssetDenom:       "uatom",
+						PriceTicksize:    &tickSize,
+						QuantityTicksize: &tickSize,
+					},
+				},
+			},
+		},
+	})
+	require.NoError(t, err)
+
 	_, err = handler(ctx, &types.MsgUnregisterContract{
 		Creator:      testAccount.String(),
 		ContractAddr: contractAddr.String(),
@@ -76,4 +95,6 @@ func TestUnregisterContractSetSiblings(t *testing.T) {
 	require.Error(t, err)
 	balance = keeper.BankKeeper.GetBalance(ctx, testAccount, "usei")
 	require.Equal(t, int64(9900000), balance.Amount.Int64())
+	pairs := keeper.GetAllRegisteredPairs(ctx, contractAddr.String())
+	require.Empty(t, pairs)
 }

x/dex/keeper/order_placement.go

@@ -25,3 +25,9 @@ func (k Keeper) SetNextOrderID(ctx sdk.Context, contractAddr string, nextID uint
 	binary.BigEndian.PutUint64(bz, nextID)
 	store.Set(byteKey, bz)
 }
+
+func (k Keeper) DeleteNextOrderID(ctx sdk.Context, contractAddr string) {
+	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.NextOrderIDPrefix(contractAddr))
+	byteKey := types.KeyPrefix(types.NextOrderIDKey)
+	store.Delete(byteKey)
+}

x/dex/keeper/pair.go

@@ -50,3 +50,7 @@ func (k Keeper) GetAllRegisteredPairs(ctx sdk.Context, contractAddr string) []ty
 
 	return list
 }
+
+func (k Keeper) DeleteAllRegisteredPairsForContract(ctx sdk.Context, contractAddr string) {
+	k.removeAllForPrefix(ctx, types.RegisteredPairPrefix(contractAddr))
+}

x/dex/keeper/price.go

@@ -59,6 +59,10 @@ func (k Keeper) GetAllPrices(ctx sdk.Context, contractAddr string, pair types.Pa
 	return
 }
 
+func (k Keeper) RemoveAllPricesForContract(ctx sdk.Context, contractAddr string) {
+	k.removeAllForPrefix(ctx, types.PriceContractPrefix(contractAddr))
+}
+
 func GetKeyForTs(ts uint64) []byte {
 	tsKey := make([]byte, 8)
 	binary.BigEndian.PutUint64(tsKey, ts)

x/dex/keeper/short_book.go

@@ -89,6 +89,10 @@ func (k Keeper) GetShortAllocationForOrderID(ctx sdk.Context, contractAddr strin
 	return nil, false
 }
 
+func (k Keeper) RemoveAllShortBooksForContract(ctx sdk.Context, contractAddr string) {
+	k.removeAllForPrefix(ctx, types.OrderBookContractPrefix(false, contractAddr))
+}
+
 func GetKeyForShortBook(shortBook types.ShortBook) []byte {
 	return GetKeyForPrice(shortBook.Entry.Price)
 }

x/dex/keeper/trigger_book.go

@@ -60,6 +60,10 @@ func (k Keeper) GetAllTriggeredOrders(ctx sdk.Context, contractAddr string) (lis
 	return
 }
 
+func (k Keeper) RemoveAllTriggeredOrders(ctx sdk.Context, contractAddr string) {
+	k.removeAllForPrefix(ctx, types.ContractKeyPrefix(types.TriggerBookKey, contractAddr))
+}
+
 func GetKeyForOrderID(orderID uint64) []byte {
 	key := make([]byte, 8)
 	binary.LittleEndian.PutUint64(key, orderID)

x/dex/types/keys.go

@@ -45,16 +45,20 @@ func PairPrefix(priceDenom string, assetDenom string) []byte {
 }
 
 func OrderBookPrefix(long bool, contractAddr string, priceDenom string, assetDenom string) []byte {
+	return append(
+		OrderBookContractPrefix(long, contractAddr),
+		PairPrefix(priceDenom, assetDenom)...,
+	)
+}
+
+func OrderBookContractPrefix(long bool, contractAddr string) []byte {
 	var prefix []byte
 	if long {
 		prefix = KeyPrefix(LongBookKey)
 	} else {
 		prefix = KeyPrefix(ShortBookKey)
 	}
-	return append(
-		append(prefix, AddressKeyPrefix(contractAddr)...),
-		PairPrefix(priceDenom, assetDenom)...,
-	)
+	return append(prefix, AddressKeyPrefix(contractAddr)...)
 }
 
 func TriggerOrderBookPrefix(contractAddr string, priceDenom string, assetDenom string) []byte {
@@ -73,11 +77,15 @@ func TwapPrefix(contractAddr string) []byte {
 // `Price` constant + contract + price denom + asset denom
 func PricePrefix(contractAddr string, priceDenom string, assetDenom string) []byte {
 	return append(
-		append(KeyPrefix(PriceKey), AddressKeyPrefix(contractAddr)...),
+		PriceContractPrefix(contractAddr),
 		PairPrefix(priceDenom, assetDenom)...,
 	)
 }
 
+func PriceContractPrefix(contractAddr string) []byte {
+	return append(KeyPrefix(PriceKey), AddressKeyPrefix(contractAddr)...)
+}
+
 func SettlementEntryPrefix(contractAddr string, priceDenom string, assetDenom string) []byte {
 	return append(
 		append(KeyPrefix(SettlementEntryKey), AddressKeyPrefix(contractAddr)...),