refactor(distribution)!: use collections for DelegatorWithdrawAddress…

… state management. (cosmos#16440)

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

CHANGELOG.md

@@ -255,6 +255,8 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (sims) [#16052](https://github.com/cosmos/cosmos-sdk/pull/16062) `GetOrGenerate` no longer requires a codec argument is now 4-arity instead of 5-arity.
 * (baseapp) [#16342](https://github.com/cosmos/cosmos-sdk/pull/16342) NewContext was renamed to NewContextLegacy. The replacement (NewContext) now does not take a header, instead you should set the header via `WithHeaderInfo` or `WithBlockHeight`. Note that `WithBlockHeight` will soon be depreacted and its recommneded to use `WithHeaderInfo`
 * (x/auth) [#16112](https://github.com/cosmos/cosmos-sdk/issues/16112) `helpers.AddGenesisAccount` has been moved to `x/genutil` to remove the cyclic dependency between `x/auth` and `x/genutil`.
+* (x/distribution) [#16440](https://github.com/cosmos/cosmos-sdk/pull/16440) use collections for `DelegatorWithdrawAddresState`:
+  * remove `Keeper`: `SetDelegatorWithdrawAddr`, `DeleteDelegatorWithdrawAddr`, `IterateDelegatorWithdrawAddrs`.
 
 ### Client Breaking Changes
 

types/collections.go

@@ -94,26 +94,26 @@ func (a genericAddressKey[T]) SizeNonTerminal(key T) int {
 	return collections.BytesKey.SizeNonTerminal(key)
 }
 
-// Deprecated: genericAddressIndexKey is a special key codec used to retain state backwards compatibility
+// Deprecated: lengthPrefixedAddressKey is a special key codec used to retain state backwards compatibility
 // when a generic address key (be: AccAddress, ValAddress, ConsAddress), is used as an index key.
-// More docs can be found in the AddressKeyAsIndexKey function.
-type genericAddressIndexKey[T addressUnion] struct {
+// More docs can be found in the LengthPrefixedAddressKey function.
+type lengthPrefixedAddressKey[T addressUnion] struct {
 	collcodec.KeyCodec[T]
 }
 
-func (g genericAddressIndexKey[T]) Encode(buffer []byte, key T) (int, error) {
+func (g lengthPrefixedAddressKey[T]) Encode(buffer []byte, key T) (int, error) {
 	return g.EncodeNonTerminal(buffer, key)
 }
 
-func (g genericAddressIndexKey[T]) Decode(buffer []byte) (int, T, error) {
+func (g lengthPrefixedAddressKey[T]) Decode(buffer []byte) (int, T, error) {
 	return g.DecodeNonTerminal(buffer)
 }
 
-func (g genericAddressIndexKey[T]) Size(key T) int { return g.SizeNonTerminal(key) }
+func (g lengthPrefixedAddressKey[T]) Size(key T) int { return g.SizeNonTerminal(key) }
 
-func (g genericAddressIndexKey[T]) KeyType() string { return "index_key/" + g.KeyCodec.KeyType() }
+func (g lengthPrefixedAddressKey[T]) KeyType() string { return "index_key/" + g.KeyCodec.KeyType() }
 
-// Deprecated: AddressKeyAsIndexKey implements an SDK backwards compatible indexing key encoder
+// Deprecated: LengthPrefixedAddressKey implements an SDK backwards compatible indexing key encoder
 // for addresses.
 // The status quo in the SDK is that address keys are length prefixed even when they're the
 // last part of a composite key. This should never be used unless to retain state compatibility.
@@ -122,8 +122,8 @@ func (g genericAddressIndexKey[T]) KeyType() string { return "index_key/" + g.Ke
 // byte to the string, then when you know when the string part finishes, it's logical that the
 // part which remains is the address key. In the SDK instead we prepend to the address key its
 // length too.
-func AddressKeyAsIndexKey[T addressUnion](keyCodec collcodec.KeyCodec[T]) collcodec.KeyCodec[T] {
-	return genericAddressIndexKey[T]{
+func LengthPrefixedAddressKey[T addressUnion](keyCodec collcodec.KeyCodec[T]) collcodec.KeyCodec[T] {
+	return lengthPrefixedAddressKey[T]{
 		keyCodec,
 	}
 }

types/collections_test.go

@@ -21,7 +21,7 @@ func TestCollectionsCorrectness(t *testing.T) {
 	})
 
 	t.Run("AddressIndexingKey", func(t *testing.T) {
-		colltest.TestKeyCodec(t, AddressKeyAsIndexKey(AccAddressKey), AccAddress{0x2, 0x5, 0x8})
+		colltest.TestKeyCodec(t, LengthPrefixedAddressKey(AccAddressKey), AccAddress{0x2, 0x5, 0x8})
 	})
 
 	t.Run("Time", func(t *testing.T) {

x/bank/keeper/view.go

@@ -45,7 +45,7 @@ func newBalancesIndexes(sb *collections.SchemaBuilder) BalancesIndexes {
 	return BalancesIndexes{
 		Denom: indexes.NewReversePair[math.Int](
 			sb, types.DenomAddressPrefix, "address_by_denom_index",
-			collections.PairKeyCodec(sdk.AddressKeyAsIndexKey(sdk.AccAddressKey), collections.StringKey), // nolint:staticcheck // Note: refer to the AddressKeyAsIndexKey docs to understand why we do this.
+			collections.PairKeyCodec(sdk.LengthPrefixedAddressKey(sdk.AccAddressKey), collections.StringKey), // nolint:staticcheck // Note: refer to the LengthPrefixedAddressKey docs to understand why we do this.
 		),
 	}
 }

x/distribution/keeper/genesis.go

@@ -1,8 +1,10 @@
 package keeper
 
 import (
+	"errors"
 	"fmt"
 
+	"cosmossdk.io/collections"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/distribution/types"
 )
@@ -29,7 +31,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data types.GenesisState) {
 		if err != nil {
 			panic(err)
 		}
-		err = k.SetDelegatorWithdrawAddr(ctx, delegatorAddress, withdrawAddress)
+		err = k.DelegatorsWithdrawAddress.Set(ctx, delegatorAddress, withdrawAddress)
 		if err != nil {
 			panic(err)
 		}
@@ -143,14 +145,17 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 		panic(err)
 	}
 
-	dwi := make([]types.DelegatorWithdrawInfo, 0)
-	k.IterateDelegatorWithdrawAddrs(ctx, func(del, addr sdk.AccAddress) (stop bool) {
+	var dwi []types.DelegatorWithdrawInfo
+	err = k.DelegatorsWithdrawAddress.Walk(ctx, nil, func(key, value sdk.AccAddress) (stop bool, err error) {
 		dwi = append(dwi, types.DelegatorWithdrawInfo{
-			DelegatorAddress: del.String(),
-			WithdrawAddress:  addr.String(),
+			DelegatorAddress: key.String(),
+			WithdrawAddress:  value.String(),
 		})
-		return false
+		return false, nil
 	})
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		panic(err)
+	}
 
 	pp, err := k.GetPreviousProposerConsAddr(ctx)
 	if err != nil {

x/distribution/keeper/keeper.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 
 	"cosmossdk.io/collections"
+	collcodec "cosmossdk.io/collections/codec"
 	"cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
@@ -26,9 +27,10 @@ type Keeper struct {
 	// should be the x/gov module account.
 	authority string
 
-	Schema  collections.Schema
-	Params  collections.Item[types.Params]
-	FeePool collections.Item[types.FeePool]
+	Schema                    collections.Schema
+	Params                    collections.Item[types.Params]
+	FeePool                   collections.Item[types.FeePool]
+	DelegatorsWithdrawAddress collections.Map[sdk.AccAddress, sdk.AccAddress]
 
 	feeCollectorName string // name of the FeeCollector ModuleAccount
 }
@@ -55,6 +57,13 @@ func NewKeeper(
 		authority:        authority,
 		Params:           collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
 		FeePool:          collections.NewItem(sb, types.FeePoolKey, "fee_pool", codec.CollValue[types.FeePool](cdc)),
+		DelegatorsWithdrawAddress: collections.NewMap(
+			sb,
+			types.DelegatorWithdrawAddrPrefix,
+			"delegators_withdraw_address",
+			sdk.LengthPrefixedAddressKey(sdk.AccAddressKey), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+			collcodec.KeyToValueCodec(sdk.AccAddressKey),
+		),
 	}
 
 	schema, err := sb.Build()
@@ -99,8 +108,7 @@ func (k Keeper) SetWithdrawAddr(ctx context.Context, delegatorAddr, withdrawAddr
 		),
 	)
 
-	k.SetDelegatorWithdrawAddr(ctx, delegatorAddr, withdrawAddr)
-	return nil
+	return k.DelegatorsWithdrawAddress.Set(ctx, delegatorAddr, withdrawAddr)
 }
 
 // withdraw rewards from a delegation

x/distribution/keeper/keeper_test.go

@@ -66,6 +66,9 @@ func TestSetWithdrawAddr(t *testing.T) {
 
 	err = distrKeeper.SetWithdrawAddr(ctx, delegatorAddr, withdrawAddr)
 	require.Nil(t, err)
+	addr, err := distrKeeper.GetDelegatorWithdrawAddr(ctx, delegatorAddr)
+	require.NoError(t, err)
+	require.Equal(t, addr, withdrawAddr)
 
 	require.Error(t, distrKeeper.SetWithdrawAddr(ctx, delegatorAddr, distrAcc.GetAddress()))
 }

x/distribution/keeper/store.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"errors"
 
+	"cosmossdk.io/collections"
 	gogotypes "github.com/cosmos/gogoproto/types"
 
 	storetypes "cosmossdk.io/store/types"
@@ -15,38 +16,11 @@ import (
 
 // get the delegator withdraw address, defaulting to the delegator address
 func (k Keeper) GetDelegatorWithdrawAddr(ctx context.Context, delAddr sdk.AccAddress) (sdk.AccAddress, error) {
-	store := k.storeService.OpenKVStore(ctx)
-	b, err := store.Get(types.GetDelegatorWithdrawAddrKey(delAddr))
-	if b == nil {
-		return delAddr, err
-	}
-	return sdk.AccAddress(b), nil
-}
-
-// set the delegator withdraw address
-func (k Keeper) SetDelegatorWithdrawAddr(ctx context.Context, delAddr, withdrawAddr sdk.AccAddress) error {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Set(types.GetDelegatorWithdrawAddrKey(delAddr), withdrawAddr.Bytes())
-}
-
-// delete a delegator withdraw addr
-func (k Keeper) DeleteDelegatorWithdrawAddr(ctx context.Context, delAddr, withdrawAddr sdk.AccAddress) error {
-	store := k.storeService.OpenKVStore(ctx)
-	return store.Delete(types.GetDelegatorWithdrawAddrKey(delAddr))
-}
-
-// iterate over delegator withdraw addrs
-func (k Keeper) IterateDelegatorWithdrawAddrs(ctx context.Context, handler func(del, addr sdk.AccAddress) (stop bool)) {
-	store := k.storeService.OpenKVStore(ctx)
-	iter := storetypes.KVStorePrefixIterator(runtime.KVStoreAdapter(store), types.DelegatorWithdrawAddrPrefix)
-	defer iter.Close()
-	for ; iter.Valid(); iter.Next() {
-		addr := sdk.AccAddress(iter.Value())
-		del := types.GetDelegatorWithdrawInfoAddress(iter.Key())
-		if handler(del, addr) {
-			break
-		}
+	addr, err := k.DelegatorsWithdrawAddress.Get(ctx, delAddr)
+	if err != nil && errors.Is(err, collections.ErrNotFound) {
+		return delAddr, nil
 	}
+	return addr, err
 }
 
 // GetPreviousProposerConsAddr returns the proposer consensus address for the

x/distribution/migrations/v2/store_test.go

@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"testing"
 
+	"github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/stretchr/testify/require"
 
 	storetypes "cosmossdk.io/store/types"
@@ -52,7 +53,7 @@ func TestStoreMigration(t *testing.T) {
 		{
 			"DelegatorWithdrawAddr",
 			v1.GetDelegatorWithdrawAddrKey(addr2),
-			types.GetDelegatorWithdrawAddrKey(addr2),
+			append(types.DelegatorWithdrawAddrPrefix, address.MustLengthPrefix(addr2.Bytes())...),
 		},
 		{
 			"DelegatorStartingInfo",

x/distribution/simulation/decoder_test.go

@@ -44,7 +44,6 @@ func TestDecodeDistributionStore(t *testing.T) {
 			{Key: types.FeePoolKey, Value: cdc.MustMarshal(&feePool)},
 			{Key: types.ProposerKey, Value: consAddr1.Bytes()},
 			{Key: types.GetValidatorOutstandingRewardsKey(valAddr1), Value: cdc.MustMarshal(&outstanding)},
-			{Key: types.GetDelegatorWithdrawAddrKey(delAddr1), Value: delAddr1.Bytes()},
 			{Key: types.GetDelegatorStartingInfoKey(valAddr1, delAddr1), Value: cdc.MustMarshal(&info)},
 			{Key: types.GetValidatorHistoricalRewardsKey(valAddr1, 100), Value: cdc.MustMarshal(&historicalRewards)},
 			{Key: types.GetValidatorCurrentRewardsKey(valAddr1), Value: cdc.MustMarshal(&currentRewards)},
@@ -61,7 +60,6 @@ func TestDecodeDistributionStore(t *testing.T) {
 		{"FeePool", fmt.Sprintf("%v\n%v", feePool, feePool)},
 		{"Proposer", fmt.Sprintf("%v\n%v", consAddr1, consAddr1)},
 		{"ValidatorOutstandingRewards", fmt.Sprintf("%v\n%v", outstanding, outstanding)},
-		{"DelegatorWithdrawAddr", fmt.Sprintf("%v\n%v", delAddr1, delAddr1)},
 		{"DelegatorStartingInfo", fmt.Sprintf("%v\n%v", info, info)},
 		{"ValidatorHistoricalRewards", fmt.Sprintf("%v\n%v", historicalRewards, historicalRewards)},
 		{"ValidatorCurrentRewards", fmt.Sprintf("%v\n%v", currentRewards, currentRewards)},

x/distribution/types/keys.go

@@ -48,12 +48,12 @@ var (
 	ProposerKey                       = []byte{0x01}             // key for the proposer operator address
 	ValidatorOutstandingRewardsPrefix = []byte{0x02}             // key for outstanding rewards
 
-	DelegatorWithdrawAddrPrefix          = []byte{0x03} // key for delegator withdraw address
-	DelegatorStartingInfoPrefix          = []byte{0x04} // key for delegator starting info
-	ValidatorHistoricalRewardsPrefix     = []byte{0x05} // key for historical validators rewards / stake
-	ValidatorCurrentRewardsPrefix        = []byte{0x06} // key for current validator rewards
-	ValidatorAccumulatedCommissionPrefix = []byte{0x07} // key for accumulated validator commission
-	ValidatorSlashEventPrefix            = []byte{0x08} // key for validator slash fraction
+	DelegatorWithdrawAddrPrefix          = collections.NewPrefix(3) // key for delegator withdraw address
+	DelegatorStartingInfoPrefix          = []byte{0x04}             // key for delegator starting info
+	ValidatorHistoricalRewardsPrefix     = []byte{0x05}             // key for historical validators rewards / stake
+	ValidatorCurrentRewardsPrefix        = []byte{0x06}             // key for current validator rewards
+	ValidatorAccumulatedCommissionPrefix = []byte{0x07}             // key for accumulated validator commission
+	ValidatorSlashEventPrefix            = []byte{0x08}             // key for validator slash fraction
 
 	ParamsKey = collections.NewPrefix(9) // key for distribution module params
 )
@@ -160,11 +160,6 @@ func GetValidatorOutstandingRewardsKey(valAddr sdk.ValAddress) []byte {
 	return append(ValidatorOutstandingRewardsPrefix, address.MustLengthPrefix(valAddr.Bytes())...)
 }
 
-// GetDelegatorWithdrawAddrKey creates the key for a delegator's withdraw addr.
-func GetDelegatorWithdrawAddrKey(delAddr sdk.AccAddress) []byte {
-	return append(DelegatorWithdrawAddrPrefix, address.MustLengthPrefix(delAddr.Bytes())...)
-}
-
 // GetDelegatorStartingInfoKey creates the key for a delegator's starting info.
 func GetDelegatorStartingInfoKey(v sdk.ValAddress, d sdk.AccAddress) []byte {
 	return append(append(DelegatorStartingInfoPrefix, address.MustLengthPrefix(v.Bytes())...), address.MustLengthPrefix(d.Bytes())...)

x/gov/keeper/keeper.go

@@ -103,8 +103,8 @@ 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)), // nolint: staticcheck // sdk.AddressKeyAsIndexKey is 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 // sdk.AddressKeyAsIndexKey is needed to retain state compatibility
+		Deposits:               collections.NewMap(sb, types.DepositsKeyPrefix, "deposits", collections.PairKeyCodec(collections.Uint64Key, sdk.LengthPrefixedAddressKey(sdk.AccAddressKey)), codec.CollValue[v1.Deposit](cdc)), // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
+		Votes:                  collections.NewMap(sb, types.VotesKeyPrefix, "votes", collections.PairKeyCodec(collections.Uint64Key, sdk.LengthPrefixedAddressKey(sdk.AccAddressKey)), codec.CollValue[v1.Vote](cdc)),          // nolint: staticcheck // sdk.LengthPrefixedAddressKey is needed to retain state compatibility
 		ProposalID:             collections.NewSequence(sb, types.ProposalIDKey, "proposal_id"),
 		Proposals:              collections.NewMap(sb, types.ProposalsKeyPrefix, "proposals", collections.Uint64Key, codec.CollValue[v1.Proposal](cdc)),
 		ActiveProposalsQueue:   collections.NewMap(sb, types.ActiveProposalQueuePrefix, "active_proposals_queue", collections.PairKeyCodec(sdk.TimeKey, collections.Uint64Key), collections.Uint64Value),     // sdk.TimeKey is needed to retain state compatibility