refactor(gov)!: use collections for Vote state management. (#16164)

Co-authored-by: unknown unknown <unknown@unknown>

CHANGELOG.md

@@ -196,14 +196,19 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (types/math) [#16040](https://github.com/cosmos/cosmos-sdk/pull/16040) Remove unused aliases in math.go
 * (x/gov) [#16106](https://github.com/cosmos/cosmos-sdk/pull/16106) Remove gRPC query methods from Keeper
 * (x/gov) [#16118](https://github.com/cosmos/cosmos-sdk/pull/16118/) Use collections for constituion and params state management.
-* (x/gov) [](https://github.com/cosmos/cosmos-sdk/pull/16127) Use collections for deposit management:
+* (x/gov) [#16127](https://github.com/cosmos/cosmos-sdk/pull/16127) Use collections for deposit state management:
   - The following methods are removed from the gov keeper: `GetDeposit`, `GetAllDeposits`, `IterateAllDeposits`.
   - The following functions are removed from the gov types: `DepositKey`, `DepositsKey`.
+* (x/gov) [#16164](https://github.com/cosmos/cosmos-sdk/pull/16164) Use collections for vote state management:
+  - Removed: types `VoteKey`, `VoteKeys`
+  - Removed: keeper `IterateVotes`, `IterateAllVotes`, `GetVotes`, `GetVote`, `SetVote`
+
 * (sims) [#16155](https://github.com/cosmos/cosmos-sdk/pull/16155) 
   * `simulation.NewOperationMsg` now marshals the operation msg as proto bytes instead of legacy amino JSON bytes.
   * `simulation.NewOperationMsg` is now 2-arity instead of 3-arity with the obsolete argument `codec.ProtoCodec` removed.
   * The field `OperationMsg.Msg` is now of type `[]byte` instead of `json.RawMessage`.
 
+
 ### Client Breaking Changes
 
 * (x/staking) [#15701](https://github.com/cosmos/cosmos-sdk/pull/15701) `HistoricalInfoKey` now has a binary format.

x/gov/genesis.go

@@ -1,8 +1,11 @@
 package gov
 
 import (
+	"errors"
 	"fmt"
 
+	"cosmossdk.io/collections"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/gov/keeper"
 	"github.com/cosmos/cosmos-sdk/x/gov/types"
@@ -42,7 +45,14 @@ func InitGenesis(ctx sdk.Context, ak types.AccountKeeper, bk types.BankKeeper, k
 	}
 
 	for _, vote := range data.Votes {
-		k.SetVote(ctx, *vote)
+		addr, err := ak.StringToBytes(vote.Voter)
+		if err != nil {
+			panic(err)
+		}
+		err = k.Votes.Set(ctx, collections.Join(vote.ProposalId, sdk.AccAddress(addr)), *vote)
+		if err != nil {
+			panic(err)
+		}
 	}
 
 	for _, proposal := range data.Proposals {
@@ -90,21 +100,22 @@ func ExportGenesis(ctx sdk.Context, k *keeper.Keeper) (*v1.GenesisState, error)
 	}
 
 	var proposalsDeposits v1.Deposits
-	var proposalsVotes v1.Votes
-	for _, proposal := range proposals {
-		deposits, err := k.GetDeposits(ctx, proposal.Id)
-		if err != nil {
-			return nil, err
-		}
-
-		proposalsDeposits = append(proposalsDeposits, deposits...)
-
-		votes, err := k.GetVotes(ctx, proposal.Id)
-		if err != nil {
-			return nil, err
-		}
+	err = k.Deposits.Walk(ctx, nil, func(_ collections.Pair[uint64, sdk.AccAddress], value v1.Deposit) (stop bool, err error) {
+		proposalsDeposits = append(proposalsDeposits, &value)
+		return false, nil
+	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		panic(err)
+	}
 
-		proposalsVotes = append(proposalsVotes, votes...)
+	// export proposals votes
+	var proposalsVotes v1.Votes
+	err = k.Votes.Walk(ctx, nil, func(_ collections.Pair[uint64, sdk.AccAddress], value v1.Vote) (stop bool, err error) {
+		proposalsVotes = append(proposalsVotes, &value)
+		return false, nil
+	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		panic(err)
 	}
 
 	return &v1.GenesisState{

x/gov/keeper/grpc_query.go

@@ -82,9 +82,9 @@ func (q queryServer) Proposals(ctx context.Context, req *v1.QueryProposalsReques
 					return nil, err
 				}
 
-				_, err = q.k.GetVote(ctx, p.Id, voter)
+				has, err := q.k.Votes.Has(ctx, collections.Join(p.Id, sdk.AccAddress(voter)))
 				// if no error, vote found, matchVoter = true
-				matchVoter = err == nil
+				matchVoter = err == nil && has
 			}
 
 			// match depositor (if supplied)
@@ -131,9 +131,9 @@ func (q queryServer) Vote(ctx context.Context, req *v1.QueryVoteRequest) (*v1.Qu
 	if err != nil {
 		return nil, err
 	}
-	vote, err := q.k.GetVote(ctx, req.ProposalId, voter)
+	vote, err := q.k.Votes.Get(ctx, collections.Join(req.ProposalId, sdk.AccAddress(voter)))
 	if err != nil {
-		if errors.IsOf(err, types.ErrVoteNotFound) {
+		if errors.IsOf(err, collections.ErrNotFound) {
 			return nil, status.Errorf(codes.InvalidArgument,
 				"voter: %v not found for proposal: %v", req.Voter, req.ProposalId)
 		}
@@ -154,18 +154,10 @@ func (q queryServer) Votes(ctx context.Context, req *v1.QueryVotesRequest) (*v1.
 	}
 
 	var votes v1.Votes
-	store := q.k.storeService.OpenKVStore(ctx)
-	votesStore := prefix.NewStore(runtime.KVStoreAdapter(store), types.VotesKey(req.ProposalId))
-
-	pageRes, err := query.Paginate(votesStore, req.Pagination, func(key, value []byte) error {
-		var vote v1.Vote
-		if err := q.k.cdc.Unmarshal(value, &vote); err != nil {
-			return err
-		}
-
-		votes = append(votes, &vote)
-		return nil
-	})
+	_, pageRes, err := query.CollectionFilteredPaginate(ctx, q.k.Votes, req.Pagination, func(_ collections.Pair[uint64, sdk.AccAddress], value v1.Vote) (include bool, err error) {
+		votes = append(votes, &value)
+		return false, nil // not including results because they're being appended.
+	}, query.WithCollectionPaginationPairPrefix[uint64, sdk.AccAddress](req.ProposalId))
 	if err != nil {
 		return nil, status.Error(codes.Internal, err.Error())
 	}

x/gov/keeper/keeper.go

@@ -54,6 +54,7 @@ type Keeper struct {
 	Constitution collections.Item[string]
 	Params       collections.Item[v1.Params]
 	Deposits     collections.Map[collections.Pair[uint64, sdk.AccAddress], v1.Deposit]
+	Votes        collections.Map[collections.Pair[uint64, sdk.AccAddress], v1.Vote]
 }
 
 // GetAuthority returns the x/gov module's authority.
@@ -101,6 +102,7 @@ func NewKeeper(
 		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)), //nolint: staticcheck // Needed to retain state compatibility
+		Votes:        collections.NewMap(sb, types.VotesKeyPrefix, "votes", collections.PairKeyCodec(collections.Uint64Key, sdk.AddressKeyAsIndexKey(sdk.AccAddressKey)), codec.CollValue[v1.Vote](cdc)),          //nolint: staticcheck // Needed to retain state compatibility
 	}
 	schema, err := sb.Build()
 	if err != nil {

x/gov/keeper/migrations.go

@@ -38,7 +38,7 @@ func (m Migrator) Migrate3to4(ctx sdk.Context) error {
 	return v4.MigrateStore(ctx, m.keeper.storeService, m.legacySubspace, m.keeper.cdc)
 }
 
-// Migrate3to4 migrates from version 4 to 5.
+// Migrate4to5 migrates from version 4 to 5.
 func (m Migrator) Migrate4to5(ctx sdk.Context) error {
 	return v5.MigrateStore(ctx, m.keeper.storeService, m.keeper.cdc)
 }

x/gov/keeper/proposal.go

@@ -318,9 +318,9 @@ func (keeper Keeper) GetProposalsFiltered(ctx context.Context, params v1.QueryPr
 
 		// match voter address (if supplied)
 		if len(params.Voter) > 0 {
-			_, err = keeper.GetVote(ctx, p.Id, params.Voter)
+			has, err := keeper.Votes.Has(ctx, collections.Join(p.Id, params.Voter))
 			// if no error, vote found, matchVoter = true
-			matchVoter = err == nil
+			matchVoter = err == nil && has
 		}
 
 		// match depositor (if supplied)

x/gov/keeper/proposal_test.go

@@ -7,6 +7,8 @@ import (
 	"testing"
 	"time"
 
+	"cosmossdk.io/collections"
+
 	"github.com/stretchr/testify/require"
 
 	"github.com/cosmos/cosmos-sdk/testutil/testdata"
@@ -200,7 +202,7 @@ func (suite *KeeperTestSuite) TestGetProposalsFiltered() {
 				d := v1.NewDeposit(proposalID, addr1, nil)
 				v := v1.NewVote(proposalID, addr1, v1.NewNonSplitVoteOption(v1.OptionYes), "")
 				suite.govKeeper.SetDeposit(suite.ctx, d)
-				suite.govKeeper.SetVote(suite.ctx, v)
+				require.NoError(suite.T(), suite.govKeeper.Votes.Set(suite.ctx, collections.Join(proposalID, addr1), v))
 			}
 
 			suite.govKeeper.SetProposal(suite.ctx, p)

x/gov/keeper/tally.go

@@ -2,6 +2,9 @@ package keeper
 
 import (
 	"context"
+	"errors"
+
+	"cosmossdk.io/collections"
 
 	"cosmossdk.io/math"
 
@@ -37,12 +40,12 @@ func (keeper Keeper) Tally(ctx context.Context, proposal v1.Proposal) (passes, b
 
 		return false
 	})
-
-	err = keeper.IterateVotes(ctx, proposal.Id, func(vote v1.Vote) error {
+	rng := collections.NewPrefixedPairRange[uint64, sdk.AccAddress](proposal.Id)
+	err = keeper.Votes.Walk(ctx, rng, func(key collections.Pair[uint64, sdk.AccAddress], vote v1.Vote) (bool, error) {
 		// if validator, just record it in the map
 		voter, err := keeper.authKeeper.StringToBytes(vote.Voter)
 		if err != nil {
-			return err
+			return false, err
 		}
 
 		valAddrStr := sdk.ValAddress(voter).String()
@@ -75,10 +78,10 @@ func (keeper Keeper) Tally(ctx context.Context, proposal v1.Proposal) (passes, b
 			return false
 		})
 
-		return keeper.deleteVote(ctx, vote.ProposalId, voter)
+		return false, keeper.Votes.Remove(ctx, collections.Join(vote.ProposalId, sdk.AccAddress(voter)))
 	})
 
-	if err != nil {
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
 		return false, false, tallyResults, err
 	}
 

x/gov/keeper/vote.go

@@ -4,10 +4,9 @@ import (
 	"context"
 	"fmt"
 
-	"cosmossdk.io/errors"
-	storetypes "cosmossdk.io/store/types"
+	"cosmossdk.io/collections"
 
-	"github.com/cosmos/cosmos-sdk/runtime"
+	"cosmossdk.io/errors"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/gov/types"
 	v1 "github.com/cosmos/cosmos-sdk/x/gov/types/v1"
@@ -38,7 +37,7 @@ func (keeper Keeper) AddVote(ctx context.Context, proposalID uint64, voterAddr s
 	}
 
 	vote := v1.NewVote(proposalID, voterAddr, options, metadata)
-	err = keeper.SetVote(ctx, vote)
+	err = keeper.Votes.Set(ctx, collections.Join(proposalID, voterAddr), vote)
 	if err != nil {
 		return err
 	}
@@ -58,118 +57,8 @@ func (keeper Keeper) AddVote(ctx context.Context, proposalID uint64, voterAddr s
 	return nil
 }
 
-// GetAllVotes returns all the votes from the store
-func (keeper Keeper) GetAllVotes(ctx context.Context) (votes v1.Votes, err error) {
-	err = keeper.IterateAllVotes(ctx, func(vote v1.Vote) error {
-		votes = append(votes, &vote)
-		return nil
-	})
-	return
-}
-
-// GetVotes returns all the votes from a proposal
-func (keeper Keeper) GetVotes(ctx context.Context, proposalID uint64) (votes v1.Votes, err error) {
-	err = keeper.IterateVotes(ctx, proposalID, func(vote v1.Vote) error {
-		votes = append(votes, &vote)
-		return nil
-	})
-	return
-}
-
-// GetVote gets the vote from an address on a specific proposal
-func (keeper Keeper) GetVote(ctx context.Context, proposalID uint64, voterAddr sdk.AccAddress) (vote v1.Vote, err error) {
-	store := keeper.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.VoteKey(proposalID, voterAddr))
-	if err != nil {
-		return vote, err
-	}
-
-	if bz == nil {
-		return vote, types.ErrVoteNotFound
-	}
-
-	err = keeper.cdc.Unmarshal(bz, &vote)
-	if err != nil {
-		return vote, err
-	}
-
-	return vote, nil
-}
-
-// SetVote sets a Vote to the gov store
-func (keeper Keeper) SetVote(ctx context.Context, vote v1.Vote) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	bz, err := keeper.cdc.Marshal(&vote)
-	if err != nil {
-		return err
-	}
-
-	addr, err := keeper.authKeeper.StringToBytes(vote.Voter)
-	if err != nil {
-		return err
-	}
-
-	return store.Set(types.VoteKey(vote.ProposalId, addr), bz)
-}
-
-// IterateAllVotes iterates over all the stored votes and performs a callback function
-func (keeper Keeper) IterateAllVotes(ctx context.Context, cb func(vote v1.Vote) error) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	iterator := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), types.VotesKeyPrefix)
-
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		var vote v1.Vote
-		err := keeper.cdc.Unmarshal(iterator.Value(), &vote)
-		if err != nil {
-			return err
-		}
-
-		err = cb(vote)
-		// exit early without error if cb returns ErrStopIterating
-		if errors.IsOf(err, errors.ErrStopIterating) {
-			return nil
-		} else if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
-// IterateVotes iterates over all the proposals votes and performs a callback function
-func (keeper Keeper) IterateVotes(ctx context.Context, proposalID uint64, cb func(vote v1.Vote) error) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	iterator := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), types.VotesKey(proposalID))
-
-	defer iterator.Close()
-	for ; iterator.Valid(); iterator.Next() {
-		var vote v1.Vote
-		err := keeper.cdc.Unmarshal(iterator.Value(), &vote)
-		if err != nil {
-			return err
-		}
-
-		err = cb(vote)
-		// exit early without error if cb returns ErrStopIterating
-		if errors.IsOf(err, errors.ErrStopIterating) {
-			return nil
-		} else if err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
 // deleteVotes deletes the all votes from a given proposalID.
 func (keeper Keeper) deleteVotes(ctx context.Context, proposalID uint64) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	return store.Delete(types.VotesKey(proposalID))
-}
-
-// deleteVote deletes a vote from a given proposalID and voter from the store
-func (keeper Keeper) deleteVote(ctx context.Context, proposalID uint64, voterAddr sdk.AccAddress) error {
-	store := keeper.storeService.OpenKVStore(ctx)
-	return store.Delete(types.VoteKey(proposalID, voterAddr))
+	// TODO(tip): fix https://github.com/cosmos/cosmos-sdk/issues/16162
+	return nil
 }