refactor(gov)!: use collections for deposit state management

x/gov/genesis.go

@@ -34,7 +34,10 @@ func InitGenesis(ctx sdk.Context, ak types.AccountKeeper, bk types.BankKeeper, k
 
 	var totalDeposits sdk.Coins
 	for _, deposit := range data.Deposits {
-		k.SetDeposit(ctx, *deposit)
+		err := k.SetDeposit(ctx, *deposit)
+		if err != nil {
+			panic(err)
+		}
 		totalDeposits = totalDeposits.Add(deposit.Amount...)
 	}
 

x/gov/keeper/deposit.go

@@ -2,6 +2,8 @@ package keeper
 
 import (
 	"context"
+	"cosmossdk.io/collections"
+	stderr "errors"
 	"fmt"
 
 	"cosmossdk.io/errors"
@@ -15,85 +17,38 @@ import (
 	v1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1"
 )
 
-// GetDeposit gets the deposit of a specific depositor on a specific proposal
-func (keeper Keeper) GetDeposit(ctx context.Context, proposalID uint64, depositorAddr sdk.AccAddress) (deposit v1.Deposit, err error) {
-	store := keeper.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.DepositKey(proposalID, depositorAddr))
-	if err != nil {
-		return deposit, err
-	}
-
-	if bz == nil {
-		return deposit, types.ErrDepositNotFound
-	}
-
-	err = keeper.cdc.Unmarshal(bz, &deposit)
-	if err != nil {
-		return deposit, err
-	}
-
-	return deposit, nil
-}
-
 // SetDeposit sets a Deposit to the gov store
 func (keeper Keeper) SetDeposit(ctx context.Context, deposit v1.Deposit) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	bz, err := keeper.cdc.Marshal(&deposit)
-	if err != nil {
-		return err
-	}
-
 	depositor, err := keeper.authKeeper.StringToBytes(deposit.Depositor)
 	if err != nil {
 		return err
 	}
-
-	return store.Set(types.DepositKey(deposit.ProposalId, depositor), bz)
-}
-
-// GetAllDeposits returns all the deposits from the store
-func (keeper Keeper) GetAllDeposits(ctx context.Context) (deposits v1.Deposits, err error) {
-	err = keeper.IterateAllDeposits(ctx, func(deposit v1.Deposit) error {
-		deposits = append(deposits, &deposit)
-		return nil
-	})
-
-	return
+	return keeper.Deposits.Set(ctx, collections.Join(deposit.ProposalId, sdk.AccAddress(depositor)), deposit)
 }
 
 // GetDeposits returns all the deposits of a proposal
 func (keeper Keeper) GetDeposits(ctx context.Context, proposalID uint64) (deposits v1.Deposits, err error) {
-	err = keeper.IterateDeposits(ctx, proposalID, func(deposit v1.Deposit) error {
+	err = keeper.IterateDeposits(ctx, proposalID, func(_ collections.Pair[uint64, sdk.AccAddress], deposit v1.Deposit) bool {
 		deposits = append(deposits, &deposit)
-		return nil
+		return false
 	})
 
 	return
 }
 
 // DeleteAndBurnDeposits deletes and burns all the deposits on a specific proposal.
 func (keeper Keeper) DeleteAndBurnDeposits(ctx context.Context, proposalID uint64) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-
-	err := keeper.IterateDeposits(ctx, proposalID, func(deposit v1.Deposit) error {
-		err := keeper.bankKeeper.BurnCoins(ctx, types.ModuleName, deposit.Amount)
-		if err != nil {
-			return err
-		}
-
-		depositor, err := keeper.authKeeper.StringToBytes(deposit.Depositor)
-		if err != nil {
-			return err
-		}
-
-		err = store.Delete(types.DepositKey(proposalID, depositor))
-		if err != nil {
-			return err
-		}
-		return nil
+	coinsToBurn := sdk.NewCoins()
+	err := keeper.IterateDeposits(ctx, proposalID, func(key collections.Pair[uint64, sdk.AccAddress], deposit v1.Deposit) bool {
+		coinsToBurn = coinsToBurn.Add(deposit.Amount...)
+		_ = keeper.Deposits.Remove(ctx, key) // can't error, otherwise the iterator wouldn't report it
+		return false
 	})
+	if err != nil {
+		return err
+	}
 
-	return err
+	return keeper.bankKeeper.BurnCoins(ctx, types.ModuleName, coinsToBurn)
 }
 
 // IterateAllDeposits iterates over all the stored deposits and performs a callback function.
@@ -123,28 +78,12 @@ func (keeper Keeper) IterateAllDeposits(ctx context.Context, cb func(deposit v1.
 }
 
 // IterateDeposits iterates over all the proposals deposits and performs a callback function
-func (keeper Keeper) IterateDeposits(ctx context.Context, proposalID uint64, cb func(deposit v1.Deposit) error) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	iterator := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), types.DepositsKey(proposalID))
-
-	defer iterator.Close()
-
-	for ; iterator.Valid(); iterator.Next() {
-		var deposit v1.Deposit
-		err := keeper.cdc.Unmarshal(iterator.Value(), &deposit)
-		if err != nil {
-			return err
-		}
-
-		err = cb(deposit)
-		// exit early without error if cb returns ErrStopIterating
-		if errors.IsOf(err, errors.ErrStopIterating) {
-			return nil
-		} else if err != nil {
-			return err
-		}
+func (keeper Keeper) IterateDeposits(ctx context.Context, proposalID uint64, cb func(key collections.Pair[uint64, sdk.AccAddress], value v1.Deposit) bool) error {
+	pair := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](proposalID)
+	err := keeper.Deposits.Walk(ctx, pair, cb)
+	if err != nil && !stderr.Is(err, collections.ErrInvalidIterator) {
+		return err
 	}
-
 	return nil
 }
 
@@ -196,12 +135,12 @@ func (keeper Keeper) AddDeposit(ctx context.Context, proposalID uint64, deposito
 	}
 
 	// Add or update deposit object
-	deposit, err := keeper.GetDeposit(ctx, proposalID, depositorAddr)
+	deposit, err := keeper.Deposits.Get(ctx, collections.Join(proposalID, depositorAddr))
 	switch {
 	case err == nil:
 		// deposit exists
 		deposit.Amount = sdk.NewCoins(deposit.Amount...).Add(depositAmount...)
-	case errors.IsOf(err, types.ErrDepositNotFound):
+	case errors.IsOf(err, collections.ErrNotFound):
 		// deposit doesn't exist
 		deposit = v1.NewDeposit(proposalID, depositorAddr, depositAmount)
 	default:
@@ -233,7 +172,6 @@ func (keeper Keeper) AddDeposit(ctx context.Context, proposalID uint64, deposito
 // send to a destAddress if defined or burn otherwise.
 // Remaining funds are send back to the depositor.
 func (keeper Keeper) ChargeDeposit(ctx context.Context, proposalID uint64, destAddress, proposalCancelRate string) error {
-	store := keeper.storeService.OpenKVStore(ctx)
 	rate := sdkmath.LegacyMustNewDecFromStr(proposalCancelRate)
 	var cancellationCharges sdk.Coins
 
@@ -275,11 +213,7 @@ func (keeper Keeper) ChargeDeposit(ctx context.Context, proposalID uint64, destA
 				return err
 			}
 		}
-
-		err = store.Delete(types.DepositKey(deposit.ProposalId, depositerAddress))
-		if err != nil {
-			return err
-		}
+		return keeper.Deposits.Remove(ctx, collections.Join(deposit.ProposalId, sdk.AccAddress(depositerAddress)))
 	}
 
 	// burn the cancellation fee or sent the cancellation charges to destination address.
@@ -317,25 +251,19 @@ func (keeper Keeper) ChargeDeposit(ctx context.Context, proposalID uint64, destA
 
 // RefundAndDeleteDeposits refunds and deletes all the deposits on a specific proposal.
 func (keeper Keeper) RefundAndDeleteDeposits(ctx context.Context, proposalID uint64) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-
-	err := keeper.IterateDeposits(ctx, proposalID, func(deposit v1.Deposit) error {
-		depositor, err := keeper.authKeeper.StringToBytes(deposit.Depositor)
-		if err != nil {
-			return err
-		}
-
+	var err error
+	iterErr := keeper.IterateDeposits(ctx, proposalID, func(key collections.Pair[uint64, sdk.AccAddress], deposit v1.Deposit) bool {
+		depositor := key.K2()
 		err = keeper.bankKeeper.SendCoinsFromModuleToAccount(ctx, types.ModuleName, depositor, deposit.Amount)
 		if err != nil {
-			return err
-		}
-
-		err = store.Delete(types.DepositKey(proposalID, depositor))
-		if err != nil {
-			return err
+			return true
 		}
-		return nil
+		_ = keeper.Deposits.Remove(ctx, key) // cannot error
+		return false
 	})
+	if iterErr != nil {
+		return iterErr
+	}
 
 	return err
 }

x/gov/keeper/deposit_test.go

@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"cosmossdk.io/collections"
 	"fmt"
 	"testing"
 
@@ -12,7 +13,6 @@ import (
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	authtypes "github.com/cosmos/cosmos-sdk/x/auth/types"
 	disttypes "github.com/cosmos/cosmos-sdk/x/distribution/types"
-	"github.com/cosmos/cosmos-sdk/x/gov/types"
 	v1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1"
 )
 
@@ -68,8 +68,8 @@ func TestDeposits(t *testing.T) {
 			require.True(t, sdk.NewCoins(proposal.TotalDeposit...).Equal(sdk.NewCoins()))
 
 			// Check no deposits at beginning
-			_, err = govKeeper.GetDeposit(ctx, proposalID, TestAddrs[1])
-			require.ErrorIs(t, err, types.ErrDepositNotFound)
+			_, err = govKeeper.Deposits.Get(ctx, collections.Join(proposalID, TestAddrs[1]))
+			require.ErrorIs(t, err, collections.ErrNotFound)
 			proposal, err = govKeeper.GetProposal(ctx, proposalID)
 			require.Nil(t, err)
 			require.Nil(t, proposal.VotingStartTime)
@@ -78,7 +78,7 @@ func TestDeposits(t *testing.T) {
 			votingStarted, err := govKeeper.AddDeposit(ctx, proposalID, TestAddrs[0], fourStake)
 			require.NoError(t, err)
 			require.False(t, votingStarted)
-			deposit, err := govKeeper.GetDeposit(ctx, proposalID, TestAddrs[0])
+			deposit, err := govKeeper.Deposits.Get(ctx, collections.Join(proposalID, TestAddrs[0]))
 			require.Nil(t, err)
 			require.Equal(t, fourStake, sdk.NewCoins(deposit.Amount...))
 			require.Equal(t, TestAddrs[0].String(), deposit.Depositor)
@@ -91,7 +91,7 @@ func TestDeposits(t *testing.T) {
 			votingStarted, err = govKeeper.AddDeposit(ctx, proposalID, TestAddrs[0], fiveStake)
 			require.NoError(t, err)
 			require.False(t, votingStarted)
-			deposit, err = govKeeper.GetDeposit(ctx, proposalID, TestAddrs[0])
+			deposit, err = govKeeper.Deposits.Get(ctx, collections.Join(proposalID, TestAddrs[0]))
 			require.Nil(t, err)
 			require.Equal(t, fourStake.Add(fiveStake...), sdk.NewCoins(deposit.Amount...))
 			require.Equal(t, TestAddrs[0].String(), deposit.Depositor)
@@ -104,7 +104,7 @@ func TestDeposits(t *testing.T) {
 			votingStarted, err = govKeeper.AddDeposit(ctx, proposalID, TestAddrs[1], fourStake)
 			require.NoError(t, err)
 			require.True(t, votingStarted)
-			deposit, err = govKeeper.GetDeposit(ctx, proposalID, TestAddrs[1])
+			deposit, err = govKeeper.Deposits.Get(ctx, collections.Join(proposalID, TestAddrs[1]))
 			require.Nil(t, err)
 			require.Equal(t, TestAddrs[1].String(), deposit.Depositor)
 			require.Equal(t, fourStake, sdk.NewCoins(deposit.Amount...))
@@ -120,7 +120,12 @@ func TestDeposits(t *testing.T) {
 
 			// Test deposit iterator
 			// NOTE order of deposits is determined by the addresses
-			deposits, _ := govKeeper.GetAllDeposits(ctx)
+			var deposits v1.Deposits
+			err = govKeeper.Deposits.Walk(ctx, nil, func(_ collections.Pair[uint64, sdk.AccAddress], deposit v1.Deposit) bool {
+				deposits = append(deposits, &deposit)
+				return false
+			})
+			require.NoError(t, err)
 			require.Len(t, deposits, 2)
 			propDeposits, _ := govKeeper.GetDeposits(ctx, proposalID)
 			require.Equal(t, deposits, propDeposits)
@@ -130,12 +135,12 @@ func TestDeposits(t *testing.T) {
 			require.Equal(t, fourStake, sdk.NewCoins(deposits[1].Amount...))
 
 			// Test Refund Deposits
-			deposit, err = govKeeper.GetDeposit(ctx, proposalID, TestAddrs[1])
+			deposit, err = govKeeper.Deposits.Get(ctx, collections.Join(proposalID, TestAddrs[1]))
 			require.Nil(t, err)
 			require.Equal(t, fourStake, sdk.NewCoins(deposit.Amount...))
 			govKeeper.RefundAndDeleteDeposits(ctx, proposalID)
-			deposit, err = govKeeper.GetDeposit(ctx, proposalID, TestAddrs[1])
-			require.ErrorIs(t, err, types.ErrDepositNotFound)
+			deposit, err = govKeeper.Deposits.Get(ctx, collections.Join(proposalID, TestAddrs[1]))
+			require.ErrorIs(t, err, collections.ErrNotFound)
 			require.Equal(t, addr0Initial, bankKeeper.GetAllBalances(ctx, TestAddrs[0]))
 			require.Equal(t, addr1Initial, bankKeeper.GetAllBalances(ctx, TestAddrs[1]))
 

x/gov/keeper/grpc_query.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"context"
+	"cosmossdk.io/collections"
 
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
@@ -91,9 +92,9 @@ func (q queryServer) Proposals(ctx context.Context, req *v1.QueryProposalsReques
 				if err != nil {
 					return nil, err
 				}
-				_, err = q.k.GetDeposit(ctx, p.Id, depositor)
+				has, err := q.k.Deposits.Has(ctx, collections.Join(p.Id, sdk.AccAddress(depositor)))
 				// if no error, deposit found, matchDepositor = true
-				matchDepositor = err == nil
+				matchDepositor = err == nil && has
 			}
 
 			if matchVoter && matchDepositor && matchStatus {
@@ -225,7 +226,7 @@ func (q queryServer) Deposit(ctx context.Context, req *v1.QueryDepositRequest) (
 	if err != nil {
 		return nil, err
 	}
-	deposit, err := q.k.GetDeposit(ctx, req.ProposalId, depositor)
+	deposit, err := q.k.Deposits.Get(ctx, collections.Join(req.ProposalId, sdk.AccAddress(depositor)))
 	if err != nil {
 		if errors.IsOf(err, types.ErrDepositNotFound) {
 			return nil, status.Errorf(codes.InvalidArgument,
@@ -248,19 +249,10 @@ func (q queryServer) Deposits(ctx context.Context, req *v1.QueryDepositsRequest)
 	}
 
 	var deposits []*v1.Deposit
-
-	store := q.k.storeService.OpenKVStore(ctx)
-	depositStore := prefix.NewStore(runtime.KVStoreAdapter(store), types.DepositsKey(req.ProposalId))
-
-	pageRes, err := query.Paginate(depositStore, req.Pagination, func(key, value []byte) error {
-		var deposit v1.Deposit
-		if err := q.k.cdc.Unmarshal(value, &deposit); err != nil {
-			return err
-		}
-
+	_, pageRes, err := query.CollectionFilteredPaginate(ctx, q.k.Deposits, req.Pagination, func(_ collections.Pair[uint64, sdk.AccAddress], deposit v1.Deposit) (bool, error) {
 		deposits = append(deposits, &deposit)
-		return nil
-	})
+		return false, nil // we don't include results as they're being appended to the slice above.
+	}, query.WithCollectionPaginationPairPrefix[uint64, sdk.AccAddress](req.ProposalId))
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}

x/gov/keeper/keeper.go

@@ -53,6 +53,7 @@ type Keeper struct {
 	Schema       collections.Schema
 	Constitution collections.Item[string]
 	Params       collections.Item[v1.Params]
+	Deposits     collections.Map[collections.Pair[uint64, sdk.AccAddress], v1.Deposit]
 }
 
 // GetAuthority returns the x/gov module's authority.
@@ -99,6 +100,7 @@ func NewKeeper(
 		authority:    authority,
 		Constitution: collections.NewItem(sb, types.ConstitutionKey, "constitution", collections.StringValue),
 		Params:       collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[v1.Params](cdc)),
+		Deposits:     collections.NewMap(sb, types.DepositsKeyPrefix, "deposits", collections.PairKeyCodec(collections.Uint64Key, sdk.AddressKeyAsIndexKey(sdk.AccAddressKey)), codec.CollValue[v1.Deposit](cdc)),
 	}
 	schema, err := sb.Build()
 	if err != nil {

x/gov/keeper/proposal.go

@@ -2,6 +2,7 @@ package keeper
 
 import (
 	"context"
+	"cosmossdk.io/collections"
 	"errors"
 	"fmt"
 	"time"
@@ -323,7 +324,7 @@ func (keeper Keeper) GetProposalsFiltered(ctx context.Context, params v1.QueryPr
 
 		// match depositor (if supplied)
 		if len(params.Depositor) > 0 {
-			_, err = keeper.GetDeposit(ctx, p.Id, params.Depositor)
+			_, err = keeper.Deposits.Get(ctx, collections.Join(p.Id, params.Depositor))
 			// if no error, deposit found, matchDepositor = true
 			matchDepositor = err == nil
 		}