add governance support for superfluid

x/superfluid/keeper/stake.go

@@ -350,3 +350,64 @@ func (k Keeper) forceUndelegateAndBurnOsmoTokens(ctx sdk.Context,
 // Eugen’s point: Only rewards message needs to be updated. Rest of messages are fine
 // Queries need to be updated
 // We can do this at the very end though, since it just relates to queries.
+
+// IterateBondedValidatorsByPower implements govtypes.StakingKeeper
+func (k Keeper) IterateBondedValidatorsByPower(ctx sdk.Context, fn func(int64, stakingtypes.ValidatorI) bool) {
+	k.sk.IterateBondedValidatorsByPower(ctx, fn)
+}
+
+// TotalBondedTokens implements govtypes.StakingKeeper
+func (k Keeper) TotalBondedTokens(ctx sdk.Context) sdk.Int {
+	return k.sk.TotalBondedTokens(ctx)
+}
+
+// IterateDelegations implements govtypes.StakingKeeper
+func (k Keeper) IterateDelegations(ctx sdk.Context, delegator sdk.AccAddress, fn func(int64, stakingtypes.DelegationI) bool) {
+	synthlocks := k.lk.GetAllSyntheticLockupsByAddr(ctx, delegator)
+	for i, lock := range synthlocks {
+		// get locked coin from the lock ID
+		interim, ok := k.GetIntermediaryAccountFromLockId(ctx, lock.UnderlyingLockId)
+		if !ok {
+			panic(fmt.Sprintf("intermediary account retrieval failed with underlying lock(Lock: %+v)", lock))
+		}
+		lock, err := k.lk.GetLockByID(ctx, lock.UnderlyingLockId)
+		if err != nil {
+			panic(fmt.Sprintf("lockup retrieval failed with underlying lock(Lock: %+v; Error: %s)", lock, err))
+		}
+		coin, err := lock.SingleCoin()
+		if err != nil {
+			panic(fmt.Sprintf("no single coin in the lock(Lock: %+v; Error: %s)", lock, err))
+		}
+
+		// get osmo-equivalent token amount
+		amount := k.GetSuperfluidOSMOTokens(ctx, interim.Denom, coin.Amount)
+
+		// get validator shares equivalent to the token amount
+		valAddr, err := sdk.ValAddressFromBech32(interim.ValAddr)
+		if err != nil {
+			panic(fmt.Sprintf("validator address decoding failed(Lock: %+v; Address: %s; Error: %s)", lock, interim.ValAddr, err))
+		}
+		validator, found := k.sk.GetValidator(ctx, valAddr)
+		if !found {
+			panic(fmt.Sprintf("validator not exists(Lock: %+v; Address: %s)", lock, valAddr))
+		}
+		shares, err := validator.SharesFromTokens(amount)
+		if err != nil {
+			// tokens are not valid. continue.
+			continue
+		}
+
+		// construct delegation and call callback
+		delegation := stakingtypes.Delegation{
+			DelegatorAddress: delegator.String(),
+			ValidatorAddress: interim.ValAddr,
+			Shares:           shares,
+		}
+		fn(int64(i), delegation)
+	}
+
+	// call the callback with the non-superfluid delegations
+	k.sk.IterateDelegations(ctx, delegator, func(i int64, delegation stakingtypes.DelegationI) (stop bool) {
+		return fn(int64(len(synthlocks))+i, delegation) // add len(synthlocks) to index to avoid overlapping
+	})
+}

x/superfluid/keeper/stake_test.go

@@ -1,13 +1,15 @@
 package keeper_test
 
 import (
+	"fmt"
 	"time"
 
+	abci "github.com/tendermint/tendermint/abci/types"
+
 	lockuptypes "github.com/osmosis-labs/osmosis/v7/x/lockup/types"
 	minttypes "github.com/osmosis-labs/osmosis/v7/x/mint/types"
 	"github.com/osmosis-labs/osmosis/v7/x/superfluid/keeper"
 	"github.com/osmosis-labs/osmosis/v7/x/superfluid/types"
-	abci "github.com/tendermint/tendermint/abci/types"
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -895,3 +897,105 @@ func (suite *KeeperTestSuite) TestRefreshIntermediaryDelegationAmounts() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestSuperfluidDelegationGovernanceVoting() {
+	testCases := []struct {
+		name             string
+		validatorStats   []stakingtypes.BondStatus
+		superDelegations [][]superfluidDelegation
+	}{
+		{
+			"with single validator and single delegation",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}}},
+		},
+		{
+			"with single validator and additional delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}, {0, 0, 0, 1000000}}},
+		},
+		{
+			"with multiple validator and multiple superfluid delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}}, {{1, 1, 0, 1000000}}},
+		},
+		{
+			"with single validator and multiple denom superfluid delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}, {0, 0, 1, 1000000}}},
+		},
+		{
+			"with multiple validators and multiple denom superfluid delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}, {0, 1, 1, 1000000}}},
+		},
+		{
+			"many delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{
+				{{0, 0, 0, 1000000}, {0, 1, 1, 1000000}},
+				{{1, 0, 0, 1000000}, {1, 0, 1, 1000000}},
+				{{2, 1, 1, 1000000}, {2, 1, 0, 1000000}},
+				{{3, 0, 0, 1000000}, {3, 1, 1, 1000000}},
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		suite.Run(tc.name, func() {
+			suite.SetupTest()
+
+			denoms, _ := suite.SetupGammPoolsAndSuperfluidAssets([]sdk.Dec{sdk.NewDec(20), sdk.NewDec(20)})
+
+			// Generate delegator addresses
+			delAddrs := CreateRandomAccounts(len(tc.superDelegations))
+
+			// setup validators
+			valAddrs := suite.SetupValidators(tc.validatorStats)
+
+			// setup superfluid delegations
+			for _, sfdel := range tc.superDelegations {
+				fmt.Println(sfdel)
+				intermediaryAccs, _ := suite.SetupSuperfluidDelegations(delAddrs, valAddrs, sfdel, denoms)
+				suite.checkIntermediaryAccountDelegations(intermediaryAccs)
+			}
+
+			// all expected delegated amounts to a validator from a delegator
+			delegatedAmount := func(delidx, validx int) sdk.Int {
+				res := sdk.ZeroInt()
+				for _, del := range tc.superDelegations[delidx] {
+					if del.valIndex == int64(validx) {
+						res = res.AddRaw(del.lpAmount)
+					}
+				}
+				return res
+			}
+			for delidx := range tc.superDelegations {
+				// store all actual delegations to a validator
+				sharePerValidatorMap := make(map[string]sdk.Dec)
+				for validx := range tc.validatorStats {
+					sharePerValidatorMap[valAddrs[validx].String()] = sdk.ZeroDec()
+				}
+				addToSharePerValidatorMap := func(val sdk.ValAddress, share sdk.Dec) {
+					if existing, ok := sharePerValidatorMap[val.String()]; ok {
+						share.AddMut(existing)
+					}
+					sharePerValidatorMap[val.String()] = share
+				}
+
+				// iterate delegations and add eligible shares to the sharePerValidatorMap
+				suite.App.SuperfluidKeeper.IterateDelegations(suite.Ctx, delAddrs[delidx], func(_ int64, del stakingtypes.DelegationI) bool {
+					addToSharePerValidatorMap(del.GetValidatorAddr(), del.GetShares())
+					return false
+				})
+
+				// check if the expected delegated amount equals to actual
+				for validx := range tc.validatorStats {
+					suite.Equal(delegatedAmount(delidx, validx).Int64()*10, sharePerValidatorMap[valAddrs[validx].String()].RoundInt().Int64())
+				}
+			}
+		})
+	}
+}

x/superfluid/types/expected_keepers.go

@@ -69,6 +69,10 @@ type StakingKeeper interface {
 	GetUnbondingDelegation(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) (ubd stakingtypes.UnbondingDelegation, found bool)
 	UnbondingTime(ctx sdk.Context) time.Duration
 	GetParams(ctx sdk.Context) stakingtypes.Params
+
+	IterateBondedValidatorsByPower(ctx sdk.Context, fn func(int64, stakingtypes.ValidatorI) bool)
+	TotalBondedTokens(ctx sdk.Context) sdk.Int
+	IterateDelegations(ctx sdk.Context, delegator sdk.AccAddress, fn func(int64, stakingtypes.DelegationI) bool)
 }
 
 // DistrKeeper expected distribution keeper.