From 69f03cd17f2264c8be576f1b71abb02322146bf6 Mon Sep 17 00:00:00 2001
From: Julien Robert <julien@rbrt.fr>
Date: Fri, 1 Mar 2024 18:24:16 +0100
Subject: [PATCH] feat(x/gov): add custom fuction to calculate vote results and
 vp (#19592)

---
 docs/architecture/adr-069-gov-improvements.md | 18 ++++----
 simapp/app.go                                 |  2 +-
 tests/integration/gov/keeper/keeper_test.go   |  2 +-
 x/gov/CHANGELOG.md                            |  2 +
 x/gov/depinject.go                            |  2 +-
 x/gov/keeper/common_test.go                   |  2 +-
 x/gov/keeper/config.go                        | 42 +++++++++++++++++++
 x/gov/keeper/deposit.go                       |  6 +--
 x/gov/keeper/export_test.go                   |  2 +-
 x/gov/keeper/keeper.go                        |  7 ++--
 x/gov/keeper/msg_server.go                    |  4 +-
 x/gov/keeper/proposal.go                      |  2 +-
 x/gov/keeper/tally.go                         |  9 +++-
 x/gov/types/config.go                         | 20 ---------
 14 files changed, 73 insertions(+), 47 deletions(-)
 create mode 100644 x/gov/keeper/config.go
 delete mode 100644 x/gov/types/config.go

diff --git a/docs/architecture/adr-069-gov-improvements.md b/docs/architecture/adr-069-gov-improvements.md
index 6d5c80fad288..af5b12645205 100644
--- a/docs/architecture/adr-069-gov-improvements.md
+++ b/docs/architecture/adr-069-gov-improvements.md
@@ -164,19 +164,15 @@ Due to the vote option change, each proposal can have the same tallying method.
 
 However, chains may want to change the tallying function (weighted vote per voting power) of `x/gov` for a different algorithm (using a quadratic function on the voter stake, for instance).
 
-The custom tallying function can be passed to the `x/gov` keeper with the following interface:
+The custom tallying function can be passed to the `x/gov` keeper config:
 
 ```go
-type Tally interface{
-    // to be decided
-
-    // Calculate calculates the tally result
-    Calculate(proposal v1.Proposal, govKeeper GovKeeper, stakingKeeper StakingKeeper) govv1.TallyResult
-    // IsAccepted returns true if the proposal passes/is accepted
-    IsAccepted() bool
-    // BurnDeposit returns true if the proposal deposit should be burned
-    BurnDeposit() bool
-}
+type CalculateVoteResultsAndVotingPowerFn func(
+	ctx context.Context,
+	keeper Keeper,
+	proposalID uint64,
+	validators map[string]v1.ValidatorGovInfo,
+) (totalVoterPower math.LegacyDec, results map[v1.VoteOption]math.LegacyDec, err error)
 ```
 
 ## Consequences
diff --git a/simapp/app.go b/simapp/app.go
index 510af074c4aa..2c1e47b95bf5 100644
--- a/simapp/app.go
+++ b/simapp/app.go
@@ -375,7 +375,7 @@ func NewSimApp(
 	// by granting the governance module the right to execute the message.
 	// See: https://docs.cosmos.network/main/modules/gov#proposal-messages
 	govRouter := govv1beta1.NewRouter()
-	govConfig := govtypes.DefaultConfig()
+	govConfig := govkeeper.DefaultConfig()
 	/*
 		Example of setting gov params:
 		govConfig.MaxMetadataLen = 10000
diff --git a/tests/integration/gov/keeper/keeper_test.go b/tests/integration/gov/keeper/keeper_test.go
index 39756b834f34..68561f4a7336 100644
--- a/tests/integration/gov/keeper/keeper_test.go
+++ b/tests/integration/gov/keeper/keeper_test.go
@@ -112,7 +112,7 @@ func initFixture(tb testing.TB) *fixture {
 		stakingKeeper,
 		poolKeeper,
 		router,
-		types.DefaultConfig(),
+		keeper.DefaultConfig(),
 		authority.String(),
 	)
 	assert.NilError(tb, govKeeper.ProposalID.Set(newCtx, 1))
diff --git a/x/gov/CHANGELOG.md b/x/gov/CHANGELOG.md
index cc844666388d..0524fd23ff8b 100644
--- a/x/gov/CHANGELOG.md
+++ b/x/gov/CHANGELOG.md
@@ -27,6 +27,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### Features
 
+* [#19592](https://github.com/cosmos/cosmos-sdk/pull/19592) Add custom tally function.
 * [#19304](https://github.com/cosmos/cosmos-sdk/pull/19304) Add `MsgSudoExec` for allowing executing any message as a sudo.
 * [#19101](https://github.com/cosmos/cosmos-sdk/pull/19101) Add message based params configuration.
 * [#18532](https://github.com/cosmos/cosmos-sdk/pull/18532) Add SPAM vote to proposals.
@@ -59,6 +60,7 @@ Ref: https://keepachangelog.com/en/1.0.0/
 
 ### API Breaking Changes
 
+* [#19592](https://github.com/cosmos/cosmos-sdk/pull/19592) `types.Config` and `types.DefaultConfig` have been moved to the keeper package in order to support the custom tallying function.
 * [#19349](https://github.com/cosmos/cosmos-sdk/pull/19349) Simplify state management in `x/gov`. Note `k.VotingPeriodProposals` and `k.SetProposal` are no longer needed and have been removed.
 * [#18532](https://github.com/cosmos/cosmos-sdk/pull/18532) All functions that were taking an expedited bool parameter now take a `ProposalType` parameter instead.
 * [#17496](https://github.com/cosmos/cosmos-sdk/pull/17496) in `x/gov/types/v1beta1/vote.go` `NewVote` was removed, constructing the struct is required for this type.
diff --git a/x/gov/depinject.go b/x/gov/depinject.go
index 1b99aeccc770..3776ddf46025 100644
--- a/x/gov/depinject.go
+++ b/x/gov/depinject.go
@@ -60,7 +60,7 @@ type ModuleOutputs struct {
 }
 
 func ProvideModule(in ModuleInputs) ModuleOutputs {
-	defaultConfig := govtypes.DefaultConfig()
+	defaultConfig := keeper.DefaultConfig()
 	if in.Config.MaxTitleLen != 0 {
 		defaultConfig.MaxTitleLen = in.Config.MaxTitleLen
 	}
diff --git a/x/gov/keeper/common_test.go b/x/gov/keeper/common_test.go
index be8d44ea20c9..fb78f6fca0ca 100644
--- a/x/gov/keeper/common_test.go
+++ b/x/gov/keeper/common_test.go
@@ -130,7 +130,7 @@ func setupGovKeeper(t *testing.T, expectations ...func(sdk.Context, mocks)) (
 
 	// Gov keeper initializations
 
-	govKeeper := keeper.NewKeeper(encCfg.Codec, storeService, m.acctKeeper, m.bankKeeper, m.stakingKeeper, m.poolKeeper, baseApp.MsgServiceRouter(), types.DefaultConfig(), govAcct.String())
+	govKeeper := keeper.NewKeeper(encCfg.Codec, storeService, m.acctKeeper, m.bankKeeper, m.stakingKeeper, m.poolKeeper, baseApp.MsgServiceRouter(), keeper.DefaultConfig(), govAcct.String())
 	require.NoError(t, govKeeper.ProposalID.Set(ctx, 1))
 	govRouter := v1beta1.NewRouter() // Also register legacy gov handlers to test them too.
 	govRouter.AddRoute(types.RouterKey, v1beta1.ProposalHandler)
diff --git a/x/gov/keeper/config.go b/x/gov/keeper/config.go
new file mode 100644
index 000000000000..5e8f3a7d0c6a
--- /dev/null
+++ b/x/gov/keeper/config.go
@@ -0,0 +1,42 @@
+package keeper
+
+import (
+	"context"
+
+	"cosmossdk.io/math"
+	v1 "cosmossdk.io/x/gov/types/v1"
+)
+
+// CalculateVoteResultsAndVotingPowerFn is a function signature for calculating vote results and voting power
+// It can be overridden to customize the voting power calculation for proposals
+// It gets the proposal tallied and the validators governance infos (bonded tokens, voting power, etc.)
+// It must return the total voting power and the results of the vote
+type CalculateVoteResultsAndVotingPowerFn func(
+	ctx context.Context,
+	keeper Keeper,
+	proposalID uint64,
+	validators map[string]v1.ValidatorGovInfo,
+) (totalVoterPower math.LegacyDec, results map[v1.VoteOption]math.LegacyDec, err error)
+
+// Config is a config struct used for initializing the gov module to avoid using globals.
+type Config struct {
+	// MaxTitleLen defines the amount of characters that can be used for proposal title
+	MaxTitleLen uint64
+	// MaxMetadataLen defines the amount of characters that can be used for proposal metadata
+	MaxMetadataLen uint64
+	// MaxSummaryLen defines the amount of characters that can be used for proposal summary
+	MaxSummaryLen uint64
+	// CalculateVoteResultsAndVotingPowerFn is a function signature for calculating vote results and voting power
+	// Keeping it nil will use the default implementation
+	CalculateVoteResultsAndVotingPowerFn CalculateVoteResultsAndVotingPowerFn
+}
+
+// DefaultConfig returns the default config for gov.
+func DefaultConfig() Config {
+	return Config{
+		MaxTitleLen:                          255,
+		MaxMetadataLen:                       255,
+		MaxSummaryLen:                        10200,
+		CalculateVoteResultsAndVotingPowerFn: nil,
+	}
+}
diff --git a/x/gov/keeper/deposit.go b/x/gov/keeper/deposit.go
index 34512ae253bf..ece9eac79059 100644
--- a/x/gov/keeper/deposit.go
+++ b/x/gov/keeper/deposit.go
@@ -98,7 +98,7 @@ func (k Keeper) AddDeposit(ctx context.Context, proposalID uint64, depositorAddr
 	}
 
 	// the deposit must only contain valid denoms (listed in the min deposit param)
-	if err := k.validateDepositDenom(ctx, params, depositAmount); err != nil {
+	if err := k.validateDepositDenom(params, depositAmount); err != nil {
 		return false, err
 	}
 
@@ -280,7 +280,7 @@ func (k Keeper) ChargeDeposit(ctx context.Context, proposalID uint64, destAddres
 // validateInitialDeposit validates if initial deposit is greater than or equal to the minimum
 // required at the time of proposal submission. This threshold amount is determined by
 // the deposit parameters. Returns nil on success, error otherwise.
-func (k Keeper) validateInitialDeposit(ctx context.Context, params v1.Params, initialDeposit sdk.Coins, proposalType v1.ProposalType) error {
+func (k Keeper) validateInitialDeposit(params v1.Params, initialDeposit sdk.Coins, proposalType v1.ProposalType) error {
 	if !initialDeposit.IsValid() || initialDeposit.IsAnyNegative() {
 		return errors.Wrap(sdkerrors.ErrInvalidCoins, initialDeposit.String())
 	}
@@ -311,7 +311,7 @@ func (k Keeper) validateInitialDeposit(ctx context.Context, params v1.Params, in
 }
 
 // validateDepositDenom validates if the deposit denom is accepted by the governance module.
-func (k Keeper) validateDepositDenom(ctx context.Context, params v1.Params, depositAmount sdk.Coins) error {
+func (k Keeper) validateDepositDenom(params v1.Params, depositAmount sdk.Coins) error {
 	denoms := []string{}
 	acceptedDenoms := make(map[string]bool, len(params.MinDeposit))
 	for _, coin := range params.MinDeposit {
diff --git a/x/gov/keeper/export_test.go b/x/gov/keeper/export_test.go
index 1421e96781d1..8294a34d95b7 100644
--- a/x/gov/keeper/export_test.go
+++ b/x/gov/keeper/export_test.go
@@ -14,5 +14,5 @@ func (k Keeper) ValidateInitialDeposit(ctx sdk.Context, initialDeposit sdk.Coins
 		return err
 	}
 
-	return k.validateInitialDeposit(ctx, params, initialDeposit, proposalType)
+	return k.validateInitialDeposit(params, initialDeposit, proposalType)
 }
diff --git a/x/gov/keeper/keeper.go b/x/gov/keeper/keeper.go
index 6a82442a401c..6ffa959c71c1 100644
--- a/x/gov/keeper/keeper.go
+++ b/x/gov/keeper/keeper.go
@@ -42,7 +42,8 @@ type Keeper struct {
 	// Msg server router
 	router baseapp.MessageRouter
 
-	config types.Config
+	// Config represent extra module configuration
+	config Config
 
 	// the address capable of executing a MsgUpdateParams message. Typically, this
 	// should be the x/gov module account.
@@ -88,7 +89,7 @@ func (k Keeper) GetAuthority() string {
 func NewKeeper(
 	cdc codec.Codec, storeService corestoretypes.KVStoreService, authKeeper types.AccountKeeper,
 	bankKeeper types.BankKeeper, sk types.StakingKeeper, pk types.PoolKeeper,
-	router baseapp.MessageRouter, config types.Config, authority string,
+	router baseapp.MessageRouter, config Config, authority string,
 ) *Keeper {
 	// ensure governance module account is set
 	if addr := authKeeper.GetModuleAddress(types.ModuleName); addr == nil {
@@ -99,7 +100,7 @@ func NewKeeper(
 		panic(fmt.Sprintf("invalid authority address: %s", authority))
 	}
 
-	defaultConfig := types.DefaultConfig()
+	defaultConfig := DefaultConfig()
 	// If MaxMetadataLen not set by app developer, set to default value.
 	if config.MaxTitleLen == 0 {
 		config.MaxTitleLen = defaultConfig.MaxTitleLen
diff --git a/x/gov/keeper/msg_server.go b/x/gov/keeper/msg_server.go
index 9fd9b8b6c86f..e6cc63a2b8f6 100644
--- a/x/gov/keeper/msg_server.go
+++ b/x/gov/keeper/msg_server.go
@@ -76,11 +76,11 @@ func (k msgServer) SubmitProposal(goCtx context.Context, msg *v1.MsgSubmitPropos
 	if msg.Expedited { // checking for backward compatibility
 		msg.ProposalType = v1.ProposalType_PROPOSAL_TYPE_EXPEDITED
 	}
-	if err := k.validateInitialDeposit(ctx, params, msg.GetInitialDeposit(), msg.ProposalType); err != nil {
+	if err := k.validateInitialDeposit(params, msg.GetInitialDeposit(), msg.ProposalType); err != nil {
 		return nil, err
 	}
 
-	if err := k.validateDepositDenom(ctx, params, msg.GetInitialDeposit()); err != nil {
+	if err := k.validateDepositDenom(params, msg.GetInitialDeposit()); err != nil {
 		return nil, err
 	}
 
diff --git a/x/gov/keeper/proposal.go b/x/gov/keeper/proposal.go
index 73d4566fa5a0..8e7801f8f520 100644
--- a/x/gov/keeper/proposal.go
+++ b/x/gov/keeper/proposal.go
@@ -264,7 +264,7 @@ func (k Keeper) ActivateVotingPeriod(ctx context.Context, proposal v1.Proposal)
 			customMessageParams, err := k.MessageBasedParams.Get(ctx, sdk.MsgTypeURL(proposal.Messages[0]))
 			if err == nil {
 				votingPeriod = customMessageParams.VotingPeriod
-			} else if err != nil && !errors.Is(err, collections.ErrNotFound) {
+			} else if !errors.Is(err, collections.ErrNotFound) {
 				return err
 			}
 		}
diff --git a/x/gov/keeper/tally.go b/x/gov/keeper/tally.go
index bf45e796e9c1..8f87ec3570d2 100644
--- a/x/gov/keeper/tally.go
+++ b/x/gov/keeper/tally.go
@@ -18,7 +18,11 @@ func (k Keeper) Tally(ctx context.Context, proposal v1.Proposal) (passes, burnDe
 		return false, false, v1.TallyResult{}, err
 	}
 
-	totalVoterPower, results, err := k.calculateVoteResultsAndVotingPower(ctx, proposal.Id, validators)
+	if k.config.CalculateVoteResultsAndVotingPowerFn == nil {
+		k.config.CalculateVoteResultsAndVotingPowerFn = defaultCalculateVoteResultsAndVotingPower
+	}
+
+	totalVoterPower, results, err := k.config.CalculateVoteResultsAndVotingPowerFn(ctx, k, proposal.Id, validators)
 	if err != nil {
 		return false, false, v1.TallyResult{}, err
 	}
@@ -232,8 +236,9 @@ func (k Keeper) getCurrentValidators(ctx context.Context) (map[string]v1.Validat
 
 // calculateVoteResultsAndVotingPower iterate over all votes, tally up the voting power of each validator
 // and returns the votes results from voters
-func (k Keeper) calculateVoteResultsAndVotingPower(
+func defaultCalculateVoteResultsAndVotingPower(
 	ctx context.Context,
+	k Keeper,
 	proposalID uint64,
 	validators map[string]v1.ValidatorGovInfo,
 ) (math.LegacyDec, map[v1.VoteOption]math.LegacyDec, error) {
diff --git a/x/gov/types/config.go b/x/gov/types/config.go
deleted file mode 100644
index 919d54e5b9ad..000000000000
--- a/x/gov/types/config.go
+++ /dev/null
@@ -1,20 +0,0 @@
-package types
-
-// Config is a config struct used for initializing the gov module to avoid using globals.
-type Config struct {
-	// MaxTitleLen defines the amount of characters that can be used for proposal title
-	MaxTitleLen uint64
-	// MaxMetadataLen defines the amount of characters that can be used for proposal metadata.
-	MaxMetadataLen uint64
-	// MaxSummaryLen defines the amount of characters that can be used for proposal summary
-	MaxSummaryLen uint64
-}
-
-// DefaultConfig returns the default config for gov.
-func DefaultConfig() Config {
-	return Config{
-		MaxTitleLen:    255,
-		MaxMetadataLen: 255,
-		MaxSummaryLen:  10200,
-	}
-}