refactor(x/distribution)!: migrate to collections (#17115)

Co-authored-by: Julien Robert <[email protected]>

CHANGELOG.md

@@ -56,6 +56,9 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* (x/distribution) [#17115](https://github.com/cosmos/cosmos-sdk/pull/17115) Use collections for `PreviousProposer` and `ValidatorSlashEvents`:
+    * remove from `Keeper`: `GetPreviousProposerConsAddr`, `SetPreviousProposerConsAddr`, `GetValidatorHistoricalReferenceCount`, `GetValidatorSlashEvent`, `SetValidatorSlashEvent`.
+* (x/slashing) [17063](https://github.com/cosmos/cosmos-sdk/pull/17063) Use collections for `HistoricalInfo`:
 * (x/feegrant) [16535](https://github.com/cosmos/cosmos-sdk/pull/16535) Use collections for `FeeAllowance`, `FeeAllowanceQueue`.
 * (x/staking) [17063](https://github.com/cosmos/cosmos-sdk/pull/17063) Use collections for `HistoricalInfo`:
     * remove `Keeper`: `GetHistoricalInfo`, `SetHistoricalInfo`,
@@ -94,6 +97,10 @@ Ref: https://keepachangelog.com/en/1.0.0/
 * (server) [#17177](https://github.com/cosmos/cosmos-sdk/pull/17177) Remove `iavl-lazy-loading` configuration.
 * (rosetta) [#16276](https://github.com/cosmos/cosmos-sdk/issues/16276) Rosetta migration to standalone repo.
 
+### State Machine Breaking
+
+* (x/distribution) [#17115](https://github.com/cosmos/cosmos-sdk/pull/17115) Migrate `PreviousProposer` to collections.
+
 ## [v0.50.0-beta.0](https://github.com/cosmos/cosmos-sdk/releases/tag/v0.50.0-beta.0) - 2023-07-19
 
 ### Features

go.mod

@@ -4,7 +4,7 @@ module github.com/cosmos/cosmos-sdk
 
 require (
 	cosmossdk.io/api v0.7.0
-	cosmossdk.io/collections v0.3.0
+	cosmossdk.io/collections v0.3.1-0.20230727092431-f0f777fa3cb7
 	cosmossdk.io/core v0.9.0
 	cosmossdk.io/depinject v1.0.0-alpha.3
 	cosmossdk.io/errors v1.0.0

go.sum

@@ -37,8 +37,8 @@ cloud.google.com/go/storage v1.10.0/go.mod h1:FLPqc6j+Ki4BU591ie1oL6qBQGu2Bl/tZ9
 cloud.google.com/go/storage v1.14.0/go.mod h1:GrKmX003DSIwi9o29oFT7YDnHYwZoctc3fOKtUw0Xmo=
 cosmossdk.io/api v0.7.0 h1:QsEMIWuv9xWDbF2HZnW4Lpu1/SejCztPu0LQx7t6MN4=
 cosmossdk.io/api v0.7.0/go.mod h1:kJFAEMLN57y0viszHDPLMmieF0471o5QAwwApa+270M=
-cosmossdk.io/collections v0.3.0 h1:v0eEqLBxebAV+t+Ahwf9tSJOu95HVLINwROXx2TTZ08=
-cosmossdk.io/collections v0.3.0/go.mod h1:CHE1+niUElL9ikCpevRZcp0yqQ4TU0TrEEGirN0mvIg=
+cosmossdk.io/collections v0.3.1-0.20230727092431-f0f777fa3cb7 h1:4XhcAIVBXPwCFTT9abOzZZaEadbRaVws8A6UTr2KGIE=
+cosmossdk.io/collections v0.3.1-0.20230727092431-f0f777fa3cb7/go.mod h1:+KJND4gZHilxE3meopEl/Uet6IAw3PyiSbgeOrFzAZE=
 cosmossdk.io/core v0.9.0 h1:30ScAOHDIUOCg1DKAwqkho9wuQJnu7GUrMcg0XLioic=
 cosmossdk.io/core v0.9.0/go.mod h1:NFgl5r41Q36+RixTvyrfsS6qQ65agCbZ1FTpnN7/G1Y=
 cosmossdk.io/depinject v1.0.0-alpha.3 h1:6evFIgj//Y3w09bqOUOzEpFj5tsxBqdc5CfkO7z+zfw=

simapp/export.go

@@ -105,7 +105,10 @@ func (app *SimApp) prepForZeroHeightGenesis(ctx sdk.Context, jailAllowedAddrs []
 	}
 
 	// clear validator slash events
-	app.DistrKeeper.DeleteAllValidatorSlashEvents(ctx)
+	err = app.DistrKeeper.ValidatorSlashEvents.Clear(ctx, nil)
+	if err != nil {
+		panic(err)
+	}
 
 	// clear validator historical rewards
 	err = app.DistrKeeper.ValidatorHistoricalRewards.Clear(ctx, nil)

tests/integration/distribution/keeper/grpc_query_test.go

@@ -236,7 +236,12 @@ func TestGRPCValidatorSlashes(t *testing.T) {
 	}
 
 	for i, slash := range slashes {
-		assert.NilError(t, f.distrKeeper.SetValidatorSlashEvent(f.sdkCtx, f.valAddr, uint64(i+2), 0, slash))
+		err := f.distrKeeper.ValidatorSlashEvents.Set(
+			f.sdkCtx,
+			collections.Join3(f.valAddr, uint64(i+2), uint64(0)),
+			slash,
+		)
+		assert.NilError(t, err)
 	}
 
 	var (

tests/integration/distribution/keeper/msg_server_test.go

@@ -275,8 +275,8 @@ func TestMsgWithdrawDelegatorReward(t *testing.T) {
 	}
 	height := f.app.LastBlockHeight()
 
-	_, err = f.distrKeeper.GetPreviousProposerConsAddr(f.sdkCtx)
-	assert.Error(t, err, "previous proposer not set")
+	_, err = f.distrKeeper.PreviousProposer.Get(f.sdkCtx)
+	assert.ErrorIs(t, err, collections.ErrNotFound)
 
 	for _, tc := range testCases {
 		tc := tc
@@ -315,7 +315,7 @@ func TestMsgWithdrawDelegatorReward(t *testing.T) {
 				assert.DeepEqual(t, rewards, initOutstandingRewards.Sub(curOutstandingRewards.Rewards))
 			}
 
-			prevProposerConsAddr, err := f.distrKeeper.GetPreviousProposerConsAddr(f.sdkCtx)
+			prevProposerConsAddr, err := f.distrKeeper.PreviousProposer.Get(f.sdkCtx)
 			assert.NilError(t, err)
 			assert.Assert(t, prevProposerConsAddr.Empty() == false)
 			assert.DeepEqual(t, prevProposerConsAddr, valConsAddr)

types/query/collections_pagination.go

@@ -19,6 +19,15 @@ func WithCollectionPaginationPairPrefix[K1, K2 any](prefix K1) func(o *Collectio
 	}
 }
 
+// WithCollectionPaginationTriplePrefix applies a prefix to a collection, whose key is a collection.Triple,
+// being paginated that needs prefixing.
+func WithCollectionPaginationTriplePrefix[K1, K2, K3 any](prefix K1) func(o *CollectionsPaginateOptions[collections.Triple[K1, K2, K3]]) {
+	return func(o *CollectionsPaginateOptions[collections.Triple[K1, K2, K3]]) {
+		prefix := collections.TriplePrefix[K1, K2, K3](prefix)
+		o.Prefix = &prefix
+	}
+}
+
 // CollectionsPaginateOptions provides extra options for pagination in collections.
 type CollectionsPaginateOptions[K any] struct {
 	// Prefix allows to optionally set a prefix for the pagination.

x/distribution/abci.go

@@ -30,5 +30,5 @@ func BeginBlocker(ctx sdk.Context, k keeper.Keeper) error {
 
 	// record the proposer for when we payout on the next block
 	consAddr := sdk.ConsAddress(ctx.BlockHeader().ProposerAddress)
-	return k.SetPreviousProposerConsAddr(ctx, consAddr)
+	return k.PreviousProposer.Set(ctx, consAddr)
 }

x/distribution/keeper/delegation.go

@@ -120,7 +120,7 @@ func (k Keeper) CalculateDelegationRewards(ctx context.Context, val stakingtypes
 	// for them for the stake sanity check below.
 	endingHeight := uint64(sdkCtx.BlockHeight())
 	if endingHeight > startingHeight {
-		k.IterateValidatorSlashEventsBetween(ctx, valAddr, startingHeight, endingHeight,
+		err = k.IterateValidatorSlashEventsBetween(ctx, valAddr, startingHeight, endingHeight,
 			func(height uint64, event types.ValidatorSlashEvent) (stop bool) {
 				endingPeriod := event.ValidatorPeriod
 				if endingPeriod > startingPeriod {
@@ -138,6 +138,9 @@ func (k Keeper) CalculateDelegationRewards(ctx context.Context, val stakingtypes
 				return false
 			},
 		)
+		if err != nil {
+			return sdk.DecCoins{}, err
+		}
 	}
 
 	// A total stake sanity check; Recalculated final stake should be less than or

x/distribution/keeper/delegation_test.go

@@ -1,12 +1,14 @@
 package keeper_test
 
 import (
+	"errors"
 	"testing"
 
 	cmtproto "github.com/cometbft/cometbft/proto/tendermint/types"
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/math"
 	storetypes "cosmossdk.io/store/types"
 
@@ -72,13 +74,13 @@ func TestCalculateRewardsBasic(t *testing.T) {
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
 
 	// historical count should be 2 (once for validator init, once for delegation init)
-	require.Equal(t, uint64(2), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 2, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// end period
 	endingPeriod, _ := distrKeeper.IncrementValidatorPeriod(ctx, val)
 
 	// historical count should be 2 still
-	require.Equal(t, uint64(2), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 2, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// calculate delegation rewards
 	rewards, err := distrKeeper.CalculateDelegationRewards(ctx, val, del, endingPeriod)
@@ -108,6 +110,22 @@ func TestCalculateRewardsBasic(t *testing.T) {
 	require.Equal(t, sdk.DecCoins{{Denom: sdk.DefaultBondDenom, Amount: math.LegacyNewDec(initial / 2)}}, valCommission.Commission)
 }
 
+func getValHistoricalReferenceCount(k keeper.Keeper, ctx sdk.Context) int {
+	count := 0
+	err := k.ValidatorHistoricalRewards.Walk(
+		ctx, nil, func(key collections.Pair[sdk.ValAddress, uint64], rewards disttypes.ValidatorHistoricalRewards) (stop bool, err error) {
+			count += int(rewards.ReferenceCount)
+			return false, nil
+		},
+	)
+
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		panic(err)
+	}
+
+	return count
+}
+
 func TestCalculateRewardsAfterSlash(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	key := storetypes.NewKVStoreKey(disttypes.StoreKey)
@@ -486,7 +504,7 @@ func TestWithdrawDelegationRewardsBasic(t *testing.T) {
 	require.NoError(t, distrKeeper.AllocateTokensToValidator(ctx, val, tokens))
 
 	// historical count should be 2 (initial + latest for delegation)
-	require.Equal(t, uint64(2), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 2, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// withdraw rewards (the bank keeper should be called with the right amount of tokens to transfer)
 	expRewards := sdk.Coins{sdk.NewCoin(sdk.DefaultBondDenom, initial.QuoRaw(2))}
@@ -495,7 +513,7 @@ func TestWithdrawDelegationRewardsBasic(t *testing.T) {
 	require.Nil(t, err)
 
 	// historical count should still be 2 (added one record, cleared one)
-	require.Equal(t, uint64(2), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 2, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// withdraw commission (the bank keeper should be called with the right amount of tokens to transfer)
 	expCommission := sdk.Coins{sdk.NewCoin(sdk.DefaultBondDenom, initial.QuoRaw(2))}
@@ -808,7 +826,7 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	require.NoError(t, distrKeeper.AllocateTokensToValidator(ctx, val, tokens))
 
 	// historical count should be 2 (validator init, delegation init)
-	require.Equal(t, uint64(2), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 2, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// second delegation
 	_, del2, err := distrtestutil.Delegate(
@@ -830,7 +848,7 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	require.NoError(t, err)
 
 	// historical count should be 3 (second delegation init)
-	require.Equal(t, uint64(3), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 3, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// next block
 	ctx = ctx.WithBlockHeight(ctx.BlockHeight() + 1)
@@ -851,7 +869,7 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 	require.NoError(t, err)
 
 	// historical count should be 3 (validator init + two delegations)
-	require.Equal(t, uint64(3), distrKeeper.GetValidatorHistoricalReferenceCount(ctx))
+	require.Equal(t, 3, getValHistoricalReferenceCount(distrKeeper, ctx))
 
 	// validator withdraws commission
 	expCommission := sdk.Coins{sdk.NewCoin(sdk.DefaultBondDenom, math.NewInt(initial))}

x/distribution/keeper/genesis.go

@@ -47,7 +47,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data types.GenesisState) {
 		}
 	}
 
-	if err = k.SetPreviousProposerConsAddr(ctx, previousProposer); err != nil {
+	if err = k.PreviousProposer.Set(ctx, previousProposer); err != nil {
 		panic(err)
 	}
 
@@ -112,7 +112,17 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data types.GenesisState) {
 		if err != nil {
 			panic(err)
 		}
-		err = k.SetValidatorSlashEvent(ctx, valAddr, evt.Height, evt.Period, evt.ValidatorSlashEvent)
+
+		err = k.ValidatorSlashEvents.Set(
+			ctx,
+			collections.Join3(
+				sdk.ValAddress(valAddr),
+				evt.Height,
+				evt.Period,
+			),
+			evt.ValidatorSlashEvent,
+		)
+
 		if err != nil {
 			panic(err)
 		}
@@ -160,7 +170,7 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 		panic(err)
 	}
 
-	pp, err := k.GetPreviousProposerConsAddr(ctx)
+	pp, err := k.PreviousProposer.Get(ctx)
 	if err != nil {
 		panic(err)
 	}
@@ -234,17 +244,23 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 	}
 
 	slashes := make([]types.ValidatorSlashEventRecord, 0)
-	k.IterateValidatorSlashEvents(ctx,
-		func(val sdk.ValAddress, height uint64, event types.ValidatorSlashEvent) (stop bool) {
+	err = k.ValidatorSlashEvents.Walk(
+		ctx,
+		nil,
+		func(k collections.Triple[sdk.ValAddress, uint64, uint64], event types.ValidatorSlashEvent) (stop bool, err error) {
 			slashes = append(slashes, types.ValidatorSlashEventRecord{
-				ValidatorAddress:    val.String(),
-				Height:              height,
+				ValidatorAddress:    k.K1().String(),
+				Height:              k.K2(),
 				Period:              event.ValidatorPeriod,
 				ValidatorSlashEvent: event,
 			})
-			return false
+			return false, nil
 		},
 	)
 
+	if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+		panic(err)
+	}
+
 	return types.NewGenesisState(params, feePool, dwi, pp, outstanding, acc, his, cur, dels, slashes)
 }

x/distribution/keeper/grpc_query.go

@@ -8,9 +8,7 @@ import (
 
 	"cosmossdk.io/collections"
 	"cosmossdk.io/errors"
-	"cosmossdk.io/store/prefix"
 
-	"github.com/cosmos/cosmos-sdk/runtime"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/query"
 	"github.com/cosmos/cosmos-sdk/x/distribution/types"
@@ -178,28 +176,21 @@ func (k Querier) ValidatorSlashes(ctx context.Context, req *types.QueryValidator
 		return nil, status.Errorf(codes.InvalidArgument, "invalid validator address")
 	}
 
-	store := runtime.KVStoreAdapter(k.storeService.OpenKVStore(ctx))
-	slashesStore := prefix.NewStore(store, types.GetValidatorSlashEventPrefix(valAddr))
-
-	events, pageRes, err := query.GenericFilteredPaginate(k.cdc, slashesStore, req.Pagination, func(key []byte, result *types.ValidatorSlashEvent) (*types.ValidatorSlashEvent, error) {
-		if result.ValidatorPeriod < req.StartingHeight || result.ValidatorPeriod > req.EndingHeight {
-			return nil, nil
+	events, pageRes, err := query.CollectionFilteredPaginate(ctx, k.ValidatorSlashEvents, req.Pagination, func(key collections.Triple[sdk.ValAddress, uint64, uint64], ev types.ValidatorSlashEvent) (include bool, err error) {
+		if ev.ValidatorPeriod < req.StartingHeight || ev.ValidatorPeriod > req.EndingHeight {
+			return false, nil
 		}
-
-		return result, nil
-	}, func() *types.ValidatorSlashEvent {
-		return &types.ValidatorSlashEvent{}
-	})
+		return true, nil
+	}, func(_ collections.Triple[sdk.ValAddress, uint64, uint64], value types.ValidatorSlashEvent) (types.ValidatorSlashEvent, error) {
+		return value, nil
+	},
+		query.WithCollectionPaginationTriplePrefix[sdk.ValAddress, uint64, uint64](valAddr),
+	)
 	if err != nil {
 		return nil, err
 	}
 
-	slashes := []types.ValidatorSlashEvent{}
-	for _, event := range events {
-		slashes = append(slashes, *event)
-	}
-
-	return &types.QueryValidatorSlashesResponse{Slashes: slashes, Pagination: pageRes}, nil
+	return &types.QueryValidatorSlashesResponse{Slashes: events, Pagination: pageRes}, nil
 }
 
 // DelegationRewards the total rewards accrued by a delegation

x/distribution/keeper/hooks.go

@@ -109,7 +109,10 @@ func (h Hooks) AfterValidatorRemoved(ctx context.Context, _ sdk.ConsAddress, val
 	}
 
 	// clear slashes
-	h.k.DeleteValidatorSlashEvents(ctx, valAddr)
+	err = h.k.ValidatorSlashEvents.Clear(ctx, collections.NewPrefixedTripleRange[sdk.ValAddress, uint64, uint64](valAddr))
+	if err != nil {
+		return err
+	}
 
 	// clear historical rewards
 	err = h.k.ValidatorHistoricalRewards.Clear(ctx, collections.NewPrefixedPairRange[sdk.ValAddress, uint64](valAddr))

x/distribution/keeper/invariants.go

@@ -143,16 +143,34 @@ func ReferenceCountInvariant(k Keeper) sdk.Invariant {
 		}
 
 		slashCount := uint64(0)
-		k.IterateValidatorSlashEvents(ctx,
-			func(_ sdk.ValAddress, _ uint64, _ types.ValidatorSlashEvent) (stop bool) {
+		err = k.ValidatorSlashEvents.Walk(
+			ctx,
+			nil,
+			func(k collections.Triple[sdk.ValAddress, uint64, uint64], event types.ValidatorSlashEvent) (stop bool, err error) {
 				slashCount++
-				return false
-			})
+				return false, nil
+			},
+		)
+
+		if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+			panic(err)
+		}
 
 		// one record per validator (last tracked period), one record per
 		// delegation (previous period), one record per slash (previous period)
 		expected := valCount + uint64(len(dels)) + slashCount
-		count := k.GetValidatorHistoricalReferenceCount(ctx)
+		count := uint64(0)
+		err = k.ValidatorHistoricalRewards.Walk(
+			ctx, nil, func(key collections.Pair[sdk.ValAddress, uint64], rewards types.ValidatorHistoricalRewards) (stop bool, err error) {
+				count += uint64(rewards.ReferenceCount)
+				return false, nil
+			},
+		)
+
+		if err != nil && !errors.Is(err, collections.ErrInvalidIterator) {
+			panic(err)
+		}
+
 		broken := count != expected
 
 		return sdk.FormatInvariant(types.ModuleName, "reference count",