From 2843261161638dcc09a95602bffecfec37a82eef Mon Sep 17 00:00:00 2001
From: samricotta <37125168+samricotta@users.noreply.github.com>
Date: Mon, 22 May 2023 21:51:24 +0200
Subject: [PATCH] refactor(distribution)!: use collections for params state
 (#16211)

---
 CHANGELOG.md                                  |  3 +-
 .../distribution/keeper/grpc_query_test.go    |  6 ++--
 .../distribution/keeper/msg_server_test.go    | 34 +++++++++----------
 x/distribution/keeper/allocation_test.go      |  4 +--
 x/distribution/keeper/delegation_test.go      | 18 +++++-----
 x/distribution/keeper/genesis.go              |  4 +--
 x/distribution/keeper/grpc_query.go           |  2 +-
 x/distribution/keeper/keeper.go               | 15 +++++++-
 x/distribution/keeper/keeper_test.go          |  4 +--
 x/distribution/keeper/msg_server.go           |  2 +-
 x/distribution/keeper/params.go               | 29 ++--------------
 x/distribution/types/keys.go                  |  4 ++-
 12 files changed, 58 insertions(+), 67 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ddccd6a9b0..7463de3e6f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -214,11 +214,12 @@ Ref: https://keepachangelog.com/en/1.0.0/
     * Remove: errors unused errors
 
 * (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`.
 * (cli) [#16209](https://github.com/cosmos/cosmos-sdk/pull/16209) Add API `StartCmdWithOptions` to create customized start command.
-
+* (x/distribution) [#16211](https://github.com/cosmos/cosmos-sdk/pull/16211) Use collections for params state management.
 
 ### Client Breaking Changes
 
diff --git a/tests/integration/distribution/keeper/grpc_query_test.go b/tests/integration/distribution/keeper/grpc_query_test.go
index 1c2059cf0c..1e929f031c 100644
--- a/tests/integration/distribution/keeper/grpc_query_test.go
+++ b/tests/integration/distribution/keeper/grpc_query_test.go
@@ -18,7 +18,7 @@ func TestGRPCParams(t *testing.T) {
 	t.Parallel()
 	f := initFixture(t)
 
-	f.distrKeeper.SetParams(f.sdkCtx, types.DefaultParams())
+	f.distrKeeper.Params.Set(f.sdkCtx, types.DefaultParams())
 
 	qr := f.app.QueryHelper()
 	queryClient := types.NewQueryClient(qr)
@@ -51,7 +51,7 @@ func TestGRPCParams(t *testing.T) {
 					WithdrawAddrEnabled: true,
 				}
 
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 				expParams = params
 			},
 			msg: &types.QueryParamsRequest{},
@@ -369,7 +369,7 @@ func TestGRPCDelegatorWithdrawAddress(t *testing.T) {
 	t.Parallel()
 	f := initFixture(t)
 
-	f.distrKeeper.SetParams(f.sdkCtx, types.DefaultParams())
+	f.distrKeeper.Params.Set(f.sdkCtx, types.DefaultParams())
 
 	qr := f.app.QueryHelper()
 	queryClient := types.NewQueryClient(qr)
diff --git a/tests/integration/distribution/keeper/msg_server_test.go b/tests/integration/distribution/keeper/msg_server_test.go
index 98e2bc0990..c885c18013 100644
--- a/tests/integration/distribution/keeper/msg_server_test.go
+++ b/tests/integration/distribution/keeper/msg_server_test.go
@@ -149,7 +149,7 @@ func TestMsgWithdrawDelegatorReward(t *testing.T) {
 	f.distrKeeper.SetFeePool(f.sdkCtx, distrtypes.FeePool{
 		CommunityPool: sdk.NewDecCoins(sdk.DecCoin{Denom: "stake", Amount: math.LegacyNewDec(10000)}),
 	})
-	f.distrKeeper.SetParams(f.sdkCtx, distrtypes.DefaultParams())
+	f.distrKeeper.Params.Set(f.sdkCtx, distrtypes.DefaultParams())
 	initFeePool, err := f.distrKeeper.GetFeePool(f.sdkCtx)
 	assert.NilError(t, err)
 
@@ -316,7 +316,7 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 	t.Parallel()
 	f := initFixture(t)
 
-	f.distrKeeper.SetParams(f.sdkCtx, distrtypes.DefaultParams())
+	f.distrKeeper.Params.Set(f.sdkCtx, distrtypes.DefaultParams())
 
 	delAddr := sdk.AccAddress(PKS[0].Address())
 	withdrawAddr := sdk.AccAddress(PKS[1].Address())
@@ -331,9 +331,9 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 		{
 			name: "empty delegator address",
 			preRun: func() {
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				params.WithdrawAddrEnabled = true
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 			},
 			msg: &distrtypes.MsgSetWithdrawAddress{
 				DelegatorAddress: emptyDelAddr.String(),
@@ -345,9 +345,9 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 		{
 			name: "empty withdraw address",
 			preRun: func() {
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				params.WithdrawAddrEnabled = true
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 			},
 			msg: &distrtypes.MsgSetWithdrawAddress{
 				DelegatorAddress: delAddr.String(),
@@ -359,9 +359,9 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 		{
 			name: "both empty addresses",
 			preRun: func() {
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				params.WithdrawAddrEnabled = true
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 			},
 			msg: &distrtypes.MsgSetWithdrawAddress{
 				DelegatorAddress: emptyDelAddr.String(),
@@ -373,9 +373,9 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 		{
 			name: "withdraw address disabled",
 			preRun: func() {
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				params.WithdrawAddrEnabled = false
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 			},
 			msg: &distrtypes.MsgSetWithdrawAddress{
 				DelegatorAddress: delAddr.String(),
@@ -387,9 +387,9 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 		{
 			name: "valid msg with same delegator and withdraw address",
 			preRun: func() {
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				params.WithdrawAddrEnabled = true
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 			},
 			msg: &distrtypes.MsgSetWithdrawAddress{
 				DelegatorAddress: delAddr.String(),
@@ -400,9 +400,9 @@ func TestMsgSetWithdrawAddress(t *testing.T) {
 		{
 			name: "valid msg",
 			preRun: func() {
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				params.WithdrawAddrEnabled = true
-				assert.NilError(t, f.distrKeeper.SetParams(f.sdkCtx, params))
+				assert.NilError(t, f.distrKeeper.Params.Set(f.sdkCtx, params))
 			},
 			msg: &distrtypes.MsgSetWithdrawAddress{
 				DelegatorAddress: delAddr.String(),
@@ -754,7 +754,7 @@ func TestMsgUpdateParams(t *testing.T) {
 				assert.NilError(t, err)
 
 				// query the params and verify it has been updated
-				params, _ := f.distrKeeper.GetParams(f.sdkCtx)
+				params, _ := f.distrKeeper.Params.Get(f.sdkCtx)
 				assert.DeepEqual(t, distrtypes.DefaultParams(), params)
 			}
 		})
@@ -765,7 +765,7 @@ func TestMsgCommunityPoolSpend(t *testing.T) {
 	t.Parallel()
 	f := initFixture(t)
 
-	f.distrKeeper.SetParams(f.sdkCtx, distrtypes.DefaultParams())
+	f.distrKeeper.Params.Set(f.sdkCtx, distrtypes.DefaultParams())
 	f.distrKeeper.SetFeePool(f.sdkCtx, distrtypes.FeePool{
 		CommunityPool: sdk.NewDecCoins(sdk.DecCoin{Denom: "stake", Amount: math.LegacyNewDec(10000)}),
 	})
@@ -845,7 +845,7 @@ func TestMsgDepositValidatorRewardsPool(t *testing.T) {
 	t.Parallel()
 	f := initFixture(t)
 
-	f.distrKeeper.SetParams(f.sdkCtx, distrtypes.DefaultParams())
+	f.distrKeeper.Params.Set(f.sdkCtx, distrtypes.DefaultParams())
 	f.distrKeeper.SetFeePool(f.sdkCtx, distrtypes.FeePool{
 		CommunityPool: sdk.NewDecCoins(sdk.DecCoin{Denom: "stake", Amount: math.LegacyNewDec(100)}),
 	})
diff --git a/x/distribution/keeper/allocation_test.go b/x/distribution/keeper/allocation_test.go
index c4277ca871..a9546e69d7 100644
--- a/x/distribution/keeper/allocation_test.go
+++ b/x/distribution/keeper/allocation_test.go
@@ -102,7 +102,7 @@ func TestAllocateTokensToManyValidators(t *testing.T) {
 	)
 
 	// reset fee pool & set params
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
 
 	// create validator with 50% commission
@@ -237,7 +237,7 @@ func TestAllocateTokensTruncation(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 10% commission
 	valAddr0 := sdk.ValAddress(valConsAddr0)
diff --git a/x/distribution/keeper/delegation_test.go b/x/distribution/keeper/delegation_test.go
index e15828d838..f4ef44b546 100644
--- a/x/distribution/keeper/delegation_test.go
+++ b/x/distribution/keeper/delegation_test.go
@@ -48,7 +48,7 @@ func TestCalculateRewardsBasic(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -131,7 +131,7 @@ func TestCalculateRewardsAfterSlash(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -231,7 +231,7 @@ func TestCalculateRewardsAfterManySlashes(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -351,7 +351,7 @@ func TestCalculateRewardsMultiDelegator(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -446,7 +446,7 @@ func TestWithdrawDelegationRewardsBasic(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -519,7 +519,7 @@ func TestCalculateRewardsAfterManySlashesInSameBlock(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -631,7 +631,7 @@ func TestCalculateRewardsMultiDelegatorMultiSlash(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	valPower := int64(100)
 
@@ -763,7 +763,7 @@ func TestCalculateRewardsMultiDelegatorMultWithdraw(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
@@ -959,7 +959,7 @@ func Test100PercentCommissionReward(t *testing.T) {
 
 	// reset fee pool
 	distrKeeper.SetFeePool(ctx, disttypes.InitialFeePool())
-	distrKeeper.SetParams(ctx, disttypes.DefaultParams())
+	distrKeeper.Params.Set(ctx, disttypes.DefaultParams())
 
 	// create validator with 50% commission
 	valAddr := sdk.ValAddress(valConsAddr0)
diff --git a/x/distribution/keeper/genesis.go b/x/distribution/keeper/genesis.go
index 188b042cb7..ed67b6229e 100644
--- a/x/distribution/keeper/genesis.go
+++ b/x/distribution/keeper/genesis.go
@@ -13,7 +13,7 @@ func (k Keeper) InitGenesis(ctx sdk.Context, data types.GenesisState) {
 
 	k.SetFeePool(ctx, data.FeePool)
 
-	if err := k.SetParams(ctx, data.Params); err != nil {
+	if err := k.Params.Set(ctx, data.Params); err != nil {
 		panic(err)
 	}
 
@@ -114,7 +114,7 @@ func (k Keeper) ExportGenesis(ctx sdk.Context) *types.GenesisState {
 		panic(err)
 	}
 
-	params, err := k.GetParams(ctx)
+	params, err := k.Params.Get(ctx)
 	if err != nil {
 		panic(err)
 	}
diff --git a/x/distribution/keeper/grpc_query.go b/x/distribution/keeper/grpc_query.go
index 7d0916ca55..63a4c2e26a 100644
--- a/x/distribution/keeper/grpc_query.go
+++ b/x/distribution/keeper/grpc_query.go
@@ -28,7 +28,7 @@ func NewQuerier(keeper Keeper) Querier {
 
 // Params queries params of distribution module
 func (k Querier) Params(c context.Context, req *types.QueryParamsRequest) (*types.QueryParamsResponse, error) {
-	params, err := k.GetParams(c)
+	params, err := k.Keeper.Params.Get(c)
 	if err != nil {
 		return nil, err
 	}
diff --git a/x/distribution/keeper/keeper.go b/x/distribution/keeper/keeper.go
index 9bdbe8b9c7..e25fa99f46 100644
--- a/x/distribution/keeper/keeper.go
+++ b/x/distribution/keeper/keeper.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 
+	"cosmossdk.io/collections"
 	"cosmossdk.io/core/store"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/log"
@@ -25,6 +26,9 @@ type Keeper struct {
 	// should be the x/gov module account.
 	authority string
 
+	Schema collections.Schema
+	Params collections.Item[types.Params]
+
 	feeCollectorName string // name of the FeeCollector ModuleAccount
 }
 
@@ -39,7 +43,8 @@ func NewKeeper(
 		panic(fmt.Sprintf("%s module account has not been set", types.ModuleName))
 	}
 
-	return Keeper{
+	sb := collections.NewSchemaBuilder(storeService)
+	k := Keeper{
 		storeService:     storeService,
 		cdc:              cdc,
 		authKeeper:       ak,
@@ -47,7 +52,15 @@ func NewKeeper(
 		stakingKeeper:    sk,
 		feeCollectorName: feeCollectorName,
 		authority:        authority,
+		Params:           collections.NewItem(sb, types.ParamsKey, "params", codec.CollValue[types.Params](cdc)),
 	}
+
+	schema, err := sb.Build()
+	if err != nil {
+		panic(err)
+	}
+	k.Schema = schema
+	return k
 }
 
 // GetAuthority returns the x/distribution module's authority.
diff --git a/x/distribution/keeper/keeper_test.go b/x/distribution/keeper/keeper_test.go
index 8049680131..84f37d92e9 100644
--- a/x/distribution/keeper/keeper_test.go
+++ b/x/distribution/keeper/keeper_test.go
@@ -56,13 +56,13 @@ func TestSetWithdrawAddr(t *testing.T) {
 
 	params := types.DefaultParams()
 	params.WithdrawAddrEnabled = false
-	require.NoError(t, distrKeeper.SetParams(ctx, params))
+	require.NoError(t, distrKeeper.Params.Set(ctx, params))
 
 	err := distrKeeper.SetWithdrawAddr(ctx, delegatorAddr, withdrawAddr)
 	require.NotNil(t, err)
 
 	params.WithdrawAddrEnabled = true
-	require.NoError(t, distrKeeper.SetParams(ctx, params))
+	require.NoError(t, distrKeeper.Params.Set(ctx, params))
 
 	err = distrKeeper.SetWithdrawAddr(ctx, delegatorAddr, withdrawAddr)
 	require.Nil(t, err)
diff --git a/x/distribution/keeper/msg_server.go b/x/distribution/keeper/msg_server.go
index 3c93cf2b27..e2b5e737d1 100644
--- a/x/distribution/keeper/msg_server.go
+++ b/x/distribution/keeper/msg_server.go
@@ -132,7 +132,7 @@ func (k msgServer) UpdateParams(ctx context.Context, msg *types.MsgUpdateParams)
 		return nil, err
 	}
 
-	if err := k.SetParams(ctx, msg.Params); err != nil {
+	if err := k.Params.Set(ctx, msg.Params); err != nil {
 		return nil, err
 	}
 
diff --git a/x/distribution/keeper/params.go b/x/distribution/keeper/params.go
index f81df2b866..60b352b267 100644
--- a/x/distribution/keeper/params.go
+++ b/x/distribution/keeper/params.go
@@ -4,36 +4,11 @@ import (
 	"context"
 
 	"cosmossdk.io/math"
-
-	"github.com/cosmos/cosmos-sdk/x/distribution/types"
 )
 
-// GetParams returns the total set of distribution parameters.
-func (k Keeper) GetParams(ctx context.Context) (params types.Params, err error) {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := store.Get(types.ParamsKey)
-	if bz == nil || err != nil {
-		return params, err
-	}
-
-	err = k.cdc.Unmarshal(bz, &params)
-	return params, err
-}
-
-// SetParams sets the distribution parameters.
-// CONTRACT: This method performs no validation of the parameters.
-func (k Keeper) SetParams(ctx context.Context, params types.Params) error {
-	store := k.storeService.OpenKVStore(ctx)
-	bz, err := k.cdc.Marshal(&params)
-	if err != nil {
-		return err
-	}
-	return store.Set(types.ParamsKey, bz)
-}
-
 // GetCommunityTax returns the current distribution community tax.
 func (k Keeper) GetCommunityTax(ctx context.Context) (math.LegacyDec, error) {
-	params, err := k.GetParams(ctx)
+	params, err := k.Params.Get(ctx)
 	if err != nil {
 		return math.LegacyDec{}, err
 	}
@@ -44,7 +19,7 @@ func (k Keeper) GetCommunityTax(ctx context.Context) (math.LegacyDec, error) {
 // GetWithdrawAddrEnabled returns the current distribution withdraw address
 // enabled parameter.
 func (k Keeper) GetWithdrawAddrEnabled(ctx context.Context) (enabled bool, err error) {
-	params, err := k.GetParams(ctx)
+	params, err := k.Params.Get(ctx)
 	if err != nil {
 		return false, err
 	}
diff --git a/x/distribution/types/keys.go b/x/distribution/types/keys.go
index 07fb3d6e28..38858e78ca 100644
--- a/x/distribution/types/keys.go
+++ b/x/distribution/types/keys.go
@@ -3,6 +3,8 @@ package types
 import (
 	"encoding/binary"
 
+	"cosmossdk.io/collections"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/address"
 	"github.com/cosmos/cosmos-sdk/types/kv"
@@ -53,7 +55,7 @@ var (
 	ValidatorAccumulatedCommissionPrefix = []byte{0x07} // key for accumulated validator commission
 	ValidatorSlashEventPrefix            = []byte{0x08} // key for validator slash fraction
 
-	ParamsKey = []byte{0x09} // key for distribution module params
+	ParamsKey = collections.NewPrefix(9) // key for distribution module params
 )
 
 // GetValidatorOutstandingRewardsAddress creates an address from a validator's outstanding rewards key.