refactor(feegrant): collection migration

x/feegrant/keeper/grpc_query.go

@@ -1,18 +1,16 @@
 package keeper
 
 import (
-	"bytes"
 	"context"
 
 	"github.com/cosmos/gogoproto/proto"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
-	"cosmossdk.io/store/prefix"
+	"cosmossdk.io/collections"
 	"cosmossdk.io/x/feegrant"
 
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
-	"github.com/cosmos/cosmos-sdk/runtime"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/query"
 )
@@ -76,19 +74,11 @@ func (q Keeper) Allowances(c context.Context, req *feegrant.QueryAllowancesReque
 
 	var grants []*feegrant.Grant
 
-	store := q.storeService.OpenKVStore(ctx)
-	grantsStore := prefix.NewStore(runtime.KVStoreAdapter(store), feegrant.FeeAllowancePrefixByGrantee(granteeAddr))
-
-	pageRes, err := query.Paginate(grantsStore, req.Pagination, func(key, value []byte) error {
-		var grant feegrant.Grant
-
-		if err := q.cdc.Unmarshal(value, &grant); err != nil {
-			return err
-		}
-
+	_, pageRes, err := query.CollectionFilteredPaginate(ctx, q.FeeAllowance, req.Pagination, func(key collections.Pair[sdk.AccAddress, sdk.AccAddress], grant feegrant.Grant) (include bool, err error) {
 		grants = append(grants, &grant)
-		return nil
-	})
+		return false, nil
+	}, query.WithCollectionPaginationPairPrefix[sdk.AccAddress, sdk.AccAddress](granteeAddr))
+
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}
@@ -109,19 +99,31 @@ func (q Keeper) AllowancesByGranter(c context.Context, req *feegrant.QueryAllowa
 
 	ctx := sdk.UnwrapSDKContext(c)
 
-	store := q.storeService.OpenKVStore(ctx)
-	prefixStore := prefix.NewStore(runtime.KVStoreAdapter(store), feegrant.FeeAllowanceKeyPrefix)
-	grants, pageRes, err := query.GenericFilteredPaginate(q.cdc, prefixStore, req.Pagination, func(key []byte, grant *feegrant.Grant) (*feegrant.Grant, error) {
-		// ParseAddressesFromFeeAllowanceKey expects the full key including the prefix.
-		granter, _ := feegrant.ParseAddressesFromFeeAllowanceKey(append(feegrant.FeeAllowanceKeyPrefix, key...))
-		if !bytes.Equal(granter, granterAddr) {
-			return nil, nil
-		}
-
-		return grant, nil
-	}, func() *feegrant.Grant {
-		return &feegrant.Grant{}
-	})
+	var grants []*feegrant.Grant
+	_, pageRes, err := query.CollectionFilteredPaginate(ctx, q.FeeAllowance, req.Pagination,
+		func(key collections.Pair[sdk.AccAddress, sdk.AccAddress], grant feegrant.Grant) (include bool, err error) {
+			if !sdk.AccAddress(granterAddr).Equals(key.K2()) {
+				return false, nil
+			}
+
+			grants = append(grants, &grant)
+			return false, nil
+		},
+	)
+
+	// store := q.storeService.OpenKVStore(ctx)
+	// prefixStore := prefix.NewStore(runtime.KVStoreAdapter(store), feegrant.FeeAllowanceKeyPrefix)
+	// grants, pageRes, err := query.GenericFilteredPaginate(q.cdc, prefixStore, req.Pagination, func(key []byte, grant *feegrant.Grant) (*feegrant.Grant, error) {
+	// 	// ParseAddressesFromFeeAllowanceKey expects the full key including the prefix.
+	// 	granter, _ := feegrant.ParseAddressesFromFeeAllowanceKey(append(feegrant.FeeAllowanceKeyPrefix, key...))
+	// 	if !bytes.Equal(granter, granterAddr) {
+	// 		return nil, nil
+	// 	}
+
+	// 	return grant, nil
+	// }, func() *feegrant.Grant {
+	// 	return &feegrant.Grant{}
+	// })
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}

x/feegrant/keeper/keeper.go

@@ -2,9 +2,11 @@ package keeper
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
@@ -15,26 +17,46 @@ import (
 	"github.com/cosmos/cosmos-sdk/codec"
 	"github.com/cosmos/cosmos-sdk/runtime"
 	sdk "github.com/cosmos/cosmos-sdk/types"
+	"github.com/cosmos/cosmos-sdk/types/address"
 	sdkerrors "github.com/cosmos/cosmos-sdk/types/errors"
 	"github.com/cosmos/cosmos-sdk/x/auth/ante"
 )
 
 // Keeper manages state of all fee grants, as well as calculating approval.
 // It must have a codec with all available allowances registered.
 type Keeper struct {
-	cdc          codec.BinaryCodec
-	storeService store.KVStoreService
-	authKeeper   feegrant.AccountKeeper
+	cdc               codec.BinaryCodec
+	storeService      store.KVStoreService
+	authKeeper        feegrant.AccountKeeper
+	Schema            collections.Schema
+	FeeAllowance      collections.Map[collections.Pair[sdk.AccAddress, sdk.AccAddress], feegrant.Grant]
+	FeeAllowanceQueue collections.Map[collections.Pair[time.Time, []byte], bool]
 }
 
 var _ ante.FeegrantKeeper = &Keeper{}
 
 // NewKeeper creates a fee grant Keeper
 func NewKeeper(cdc codec.BinaryCodec, storeService store.KVStoreService, ak feegrant.AccountKeeper) Keeper {
+	sb := collections.NewSchemaBuilder(storeService)
+
 	return Keeper{
 		cdc:          cdc,
 		storeService: storeService,
 		authKeeper:   ak,
+		FeeAllowance: collections.NewMap(
+			sb,
+			feegrant.FeeAllowanceKeyPrefix,
+			"allowances",
+			collections.PairKeyCodec(sdk.AccAddressKey, sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+			codec.CollValue[feegrant.Grant](cdc),
+		),
+		FeeAllowanceQueue: collections.NewMap(
+			sb,
+			feegrant.FeeAllowanceQueueKeyPrefix,
+			"allowancesQueue",
+			collections.PairKeyCodec(sdk.TimeKey, collections.BytesKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+			collections.BoolValue,
+		),
 	}
 }
 
@@ -57,8 +79,7 @@ func (k Keeper) GrantAllowance(ctx context.Context, granter, grantee sdk.AccAddr
 		k.authKeeper.SetAccount(ctx, granteeAcc)
 	}
 
-	store := k.storeService.OpenKVStore(ctx)
-	key := feegrant.FeeAllowanceKey(granter, grantee)
+	ggp := granteeGranterPair(grantee, granter)
 
 	exp, err := feeAllowance.ExpiresAt()
 	if err != nil {
@@ -75,7 +96,7 @@ func (k Keeper) GrantAllowance(ctx context.Context, granter, grantee sdk.AccAddr
 	if exp != nil {
 		// `key` formed here with the prefix of `FeeAllowanceKeyPrefix` (which is `0x00`)
 		// remove the 1st byte and reuse the remaining key as it is
-		err = k.addToFeeAllowanceQueue(ctx, key[1:], exp)
+		err = k.addToFeeAllowanceQueue(ctx, ggp, exp)
 		if err != nil {
 			return err
 		}
@@ -86,13 +107,7 @@ func (k Keeper) GrantAllowance(ctx context.Context, granter, grantee sdk.AccAddr
 		return err
 	}
 
-	bz, err := k.cdc.Marshal(&grant)
-	if err != nil {
-		return err
-	}
-
-	err = store.Set(key, bz)
-	if err != nil {
+	if err := k.FeeAllowance.Set(ctx, collections.Join(grantee, granter), grant); err != nil {
 		return err
 	}
 
@@ -109,9 +124,6 @@ func (k Keeper) GrantAllowance(ctx context.Context, granter, grantee sdk.AccAddr
 
 // UpdateAllowance updates the existing grant.
 func (k Keeper) UpdateAllowance(ctx context.Context, granter, grantee sdk.AccAddress, feeAllowance feegrant.FeeAllowanceI) error {
-	store := k.storeService.OpenKVStore(ctx)
-	key := feegrant.FeeAllowanceKey(granter, grantee)
-
 	_, err := k.getGrant(ctx, granter, grantee)
 	if err != nil {
 		return err
@@ -122,13 +134,7 @@ func (k Keeper) UpdateAllowance(ctx context.Context, granter, grantee sdk.AccAdd
 		return err
 	}
 
-	bz, err := k.cdc.Marshal(&grant)
-	if err != nil {
-		return err
-	}
-
-	err = store.Set(key, bz)
-	if err != nil {
+	if err := k.FeeAllowance.Set(ctx, collections.Join(grantee, granter), grant); err != nil {
 		return err
 	}
 
@@ -150,10 +156,7 @@ func (k Keeper) revokeAllowance(ctx context.Context, granter, grantee sdk.AccAdd
 		return err
 	}
 
-	store := k.storeService.OpenKVStore(ctx)
-	key := feegrant.FeeAllowanceKey(granter, grantee)
-	err = store.Delete(key)
-	if err != nil {
+	if err := k.FeeAllowance.Remove(ctx, collections.Join(grantee, granter)); err != nil {
 		return err
 	}
 
@@ -171,7 +174,7 @@ func (k Keeper) revokeAllowance(ctx context.Context, granter, grantee sdk.AccAdd
 // If there is none, it returns nil, nil.
 // Returns an error on parsing issues
 func (k Keeper) GetAllowance(ctx context.Context, granter, grantee sdk.AccAddress) (feegrant.FeeAllowanceI, error) {
-	grant, err := k.getGrant(ctx, granter, grantee)
+	grant, err := k.FeeAllowance.Get(ctx, collections.Join(grantee, granter))
 	if err != nil {
 		return nil, err
 	}
@@ -181,22 +184,11 @@ func (k Keeper) GetAllowance(ctx context.Context, granter, grantee sdk.AccAddres
 
 // getGrant returns entire grant between both accounts
 func (k Keeper) getGrant(ctx context.Context, granter, grantee sdk.AccAddress) (*feegrant.Grant, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	key := feegrant.FeeAllowanceKey(granter, grantee)
-	bz, err := store.Get(key)
+	feegrant, err := k.FeeAllowance.Get(ctx, collections.Join(grantee, granter))
 	if err != nil {
 		return nil, err
 	}
 
-	if len(bz) == 0 {
-		return nil, sdkerrors.ErrNotFound.Wrap("fee-grant not found")
-	}
-
-	var feegrant feegrant.Grant
-	if err := k.cdc.Unmarshal(bz, &feegrant); err != nil {
-		return nil, err
-	}
-
 	return &feegrant, nil
 }
 
@@ -208,14 +200,12 @@ func (k Keeper) IterateAllFeeAllowances(ctx context.Context, cb func(grant feegr
 	iter := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), feegrant.FeeAllowanceKeyPrefix)
 	defer iter.Close()
 
-	stop := false
-	for ; iter.Valid() && !stop; iter.Next() {
-		bz := iter.Value()
-		var feeGrant feegrant.Grant
-		if err := k.cdc.Unmarshal(bz, &feeGrant); err != nil {
-			return err
-		}
-		stop = cb(feeGrant)
+	err := k.FeeAllowance.Walk(ctx, nil, func(key collections.Pair[sdk.AccAddress, sdk.AccAddress], grant feegrant.Grant) (stop bool, err error) {
+		stop = cb(grant)
+		return stop, err
+	})
+	if err != nil {
+		return err
 	}
 
 	return nil
@@ -307,31 +297,45 @@ func (k Keeper) ExportGenesis(ctx context.Context) (*feegrant.GenesisState, erro
 }
 
 func (k Keeper) addToFeeAllowanceQueue(ctx context.Context, grantKey []byte, exp *time.Time) error {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Set(feegrant.FeeAllowancePrefixQueue(exp, grantKey), []byte{})
+	return k.FeeAllowanceQueue.Set(ctx, collections.Join(*exp, grantKey), true)
 }
 
 // RemoveExpiredAllowances iterates grantsByExpiryQueue and deletes the expired grants.
 func (k Keeper) RemoveExpiredAllowances(ctx context.Context) error {
 	exp := sdk.UnwrapSDKContext(ctx).BlockTime()
-	store := k.storeService.OpenKVStore(ctx)
-	iterator, err := store.Iterator(feegrant.FeeAllowanceQueueKeyPrefix, storetypes.InclusiveEndBytes(feegrant.AllowanceByExpTimeKey(&exp)))
-	if err != nil {
-		return err
-	}
-	defer iterator.Close()
+	rng := collections.NewPrefixUntilPairRange[time.Time, []byte](exp)
 
-	for ; iterator.Valid(); iterator.Next() {
-		err = store.Delete(iterator.Key())
-		if err != nil {
-			return err
+	err := k.FeeAllowanceQueue.Walk(ctx, rng, func(key collections.Pair[time.Time, []byte], value bool) (bool, error) {
+		grantee, granter := parseGranteeGranterPair(key.K2())
+
+		if err := k.FeeAllowance.Remove(ctx, collections.Join(sdk.AccAddress(grantee), sdk.AccAddress(granter))); err != nil {
+			return true, err
 		}
 
-		granter, grantee := feegrant.ParseAddressesFromFeeAllowanceQueueKey(iterator.Key())
-		err = store.Delete(feegrant.FeeAllowanceKey(granter, grantee))
-		if err != nil {
-			return err
+		if err := k.FeeAllowanceQueue.Remove(ctx, key); err != nil {
+			return true, err
 		}
+
+		return false, nil
+	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		return err
 	}
+
 	return nil
 }
+
+func granteeGranterPair(grantee, granter sdk.AccAddress) []byte {
+	return append(address.MustLengthPrefix(grantee.Bytes()), granter.Bytes()...)
+}
+
+func parseGranteeGranterPair(key []byte) (sdk.AccAddress, sdk.AccAddress) {
+	granteeAddrLen, _ := sdk.ParseLengthPrefixedBytes(key, 0, 1)
+	key = key[1:]
+
+	grantee, _ := sdk.ParseLengthPrefixedBytes(key, 0, int(granteeAddrLen[0]))
+	key = key[int(granteeAddrLen[0]):]
+
+	granter := key
+	return grantee, granter
+}

x/feegrant/keeper/keeper_test.go

@@ -291,7 +291,7 @@ func (suite *KeeperTestSuite) TestUseGrantedFee() {
 	// verify: feegrant is revoked
 	_, err = suite.feegrantKeeper.GetAllowance(ctx, suite.addrs[0], suite.addrs[2])
 	suite.Error(err)
-	suite.Contains(err.Error(), "fee-grant not found")
+	suite.Contains(err.Error(), "not found")
 }
 
 func (suite *KeeperTestSuite) TestIterateGrants() {
@@ -322,6 +322,7 @@ func (suite *KeeperTestSuite) TestPruneGrants() {
 	eth := sdk.NewCoins(sdk.NewInt64Coin("eth", 123))
 	now := suite.ctx.BlockTime()
 	oneYearExpiry := now.AddDate(1, 0, 0)
+	now = now.Add(1000)
 
 	testCases := []struct {
 		name      string
@@ -382,7 +383,6 @@ func (suite *KeeperTestSuite) TestPruneGrants() {
 			grantee: suite.addrs[2],
 			allowance: &feegrant.BasicAllowance{
 				SpendLimit: eth,
-				Expiration: &oneYearExpiry,
 			},
 		},
 	}
@@ -395,14 +395,18 @@ func (suite *KeeperTestSuite) TestPruneGrants() {
 			}
 			err := suite.feegrantKeeper.GrantAllowance(suite.ctx, tc.granter, tc.grantee, tc.allowance)
 			suite.NoError(err)
-			suite.feegrantKeeper.RemoveExpiredAllowances(tc.ctx)
+			err = suite.feegrantKeeper.RemoveExpiredAllowances(tc.ctx)
+			suite.NoError(err)
+
 			grant, err := suite.feegrantKeeper.GetAllowance(tc.ctx, tc.granter, tc.grantee)
 			if tc.expErrMsg != "" {
 				suite.Error(err)
 				suite.Contains(err.Error(), tc.expErrMsg)
 			} else {
+				suite.NoError(err)
 				suite.NotNil(grant)
 			}
+
 			if tc.postRun != nil {
 				tc.postRun()
 			}