add governance support for superfluid

x/superfluid/keeper/stake.go

@@ -350,3 +350,77 @@ 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
+// Iterates through staking keeper's delegations, and then all of the superfluid delegations.
+// Iterates through staking keeper's delegations, and then all of the superfluid delegations.
+func (k Keeper) IterateDelegations(ctx sdk.Context, delegator sdk.AccAddress, fn func(int64, stakingtypes.DelegationI) bool) {
+	// call the callback with the non-superfluid delegations
+	var index int64
+	k.sk.IterateDelegations(ctx, delegator, func(i int64, delegation stakingtypes.DelegationI) (stop bool) {
+		index = i
+		return fn(i, delegation)
+	})
+
+	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 {
+			ctx.Logger().Error("intermediary account retrieval failed with underlying lock", "Lock", lock)
+			continue
+		}
+
+		lock, err := k.lk.GetLockByID(ctx, lock.UnderlyingLockId)
+		if err != nil {
+			ctx.Logger().Error("lockup retrieval failed with underlying lock", "Lock", lock, "Error", err)
+			continue
+		}
+
+		coin, err := lock.SingleCoin()
+		if err != nil {
+			ctx.Logger().Error("lock fails to meet expected invariant, it contains multiple coins", "Lock", lock, "Error", err)
+			continue
+		}
+
+		// 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 {
+			ctx.Logger().Error("failed to decode validator address", "Intermediary", interim.ValAddr, "LockID", lock.ID, "Error", err)
+			continue
+		}
+
+		validator, found := k.sk.GetValidator(ctx, valAddr)
+		if !found {
+			ctx.Logger().Error("validator does not exist for lock", "Validator", valAddr, "LockID", lock.ID)
+			continue
+		}
+
+		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(index+int64(i), delegation)
+	}
+}

x/superfluid/keeper/stake_test.go

@@ -1,18 +1,26 @@
 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"
 )
 
+type normalDelegation struct {
+	delIndex   int64
+	valIndex   int64
+	coinAmount int64
+}
+
 type superfluidDelegation struct {
 	delIndex int64
 	valIndex int64
@@ -30,6 +38,15 @@ type osmoEquivalentMultiplier struct {
 	price   sdk.Dec
 }
 
+func (suite *KeeperTestSuite) SetupNormalDelegation(delAddrs []sdk.AccAddress, valAddrs []sdk.ValAddress, del normalDelegation) error {
+	val, found := suite.App.StakingKeeper.GetValidator(suite.Ctx, valAddrs[del.valIndex])
+	if !found {
+		return fmt.Errorf("validator not found")
+	}
+	_, err := suite.App.StakingKeeper.Delegate(suite.Ctx, delAddrs[del.delIndex], sdk.NewIntFromUint64(uint64(del.coinAmount)), stakingtypes.Bonded, val, false)
+	return err
+}
+
 func (suite *KeeperTestSuite) SetupSuperfluidDelegations(delAddrs []sdk.AccAddress, valAddrs []sdk.ValAddress, superDelegations []superfluidDelegation, denoms []string) ([]types.SuperfluidIntermediaryAccount, []lockuptypes.PeriodLock) {
 	flagIntermediaryAcc := make(map[string]bool)
 	intermediaryAccs := []types.SuperfluidIntermediaryAccount{}
@@ -895,3 +912,131 @@ func (suite *KeeperTestSuite) TestRefreshIntermediaryDelegationAmounts() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestSuperfluidDelegationGovernanceVoting() {
+	testCases := []struct {
+		name              string
+		validatorStats    []stakingtypes.BondStatus
+		superDelegations  [][]superfluidDelegation
+		normalDelegations []normalDelegation
+	}{
+		{
+			"with single validator and single delegation",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}}},
+			nil,
+		},
+		{
+			"with single validator and additional delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}, {0, 0, 0, 1000000}}},
+			nil,
+		},
+		{
+			"with multiple validator and multiple superfluid delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}}, {{1, 1, 0, 1000000}}},
+			nil,
+		},
+		{
+			"with single validator and multiple denom superfluid delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}, {0, 0, 1, 1000000}}},
+			nil,
+		},
+		{
+			"with multiple validators and multiple denom superfluid delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded, stakingtypes.Bonded},
+			[][]superfluidDelegation{{{0, 0, 0, 1000000}, {0, 1, 1, 1000000}}},
+			nil,
+		},
+		{
+			"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}},
+			},
+			nil,
+		},
+		{
+			"with normal delegations",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[][]superfluidDelegation{
+				{{0, 0, 0, 1000000}, {0, 0, 1, 1000000}},
+			},
+			[]normalDelegation{
+				{0, 0, 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 {
+				intermediaryAccs, _ := suite.SetupSuperfluidDelegations(delAddrs, valAddrs, sfdel, denoms)
+				suite.checkIntermediaryAccountDelegations(intermediaryAccs)
+			}
+
+			// setup normal delegations
+			for _, del := range tc.normalDelegations {
+				err := suite.SetupNormalDelegation(delAddrs, valAddrs, del)
+				suite.NoError(err)
+			}
+
+			// 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)
+					}
+				}
+				if len(tc.normalDelegations) != 0 {
+					del := tc.normalDelegations[delidx]
+					res = res.AddRaw(del.coinAmount / 10) // LP price is 10 osmo in this test
+				}
+				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.