[ValSet-Pref] Add Msgs for Delegate and Undelegate Tokens (#3260)

* refactored twap api.go for geometric TWAP

* working delegate tokens function

* added multiplier

* nit

* nit

* removed multiplier

* adam comments

* matts feedback

* added undelegate logic

* adams feedback

* matts comments

x/valset-pref/README.md

@@ -29,13 +29,13 @@ How does this module work?
 
 Staking Calculation 
 
-- The user provides an amount to delegate and our `MsgStakeToValidatorSet` divides the amount based on validator weight distribution.
+- The user provides an amount to delegate and our `MsgDelegateToValidatorSet` divides the amount based on validator weight distribution.
   For example: Stake 100osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
   our delegate logic will attempt to delegate (100 * 0.5) 50osmo for ValA , (100 * 0.3) 30osmo from ValB and (100 * 0.2) 20osmo from ValC.
 
 UnStaking Calculation 
 
-- The user provides an amount to undelegate and our `MsgUnStakeFromValidatorSet` divides the amount based on validator weight distribution.
+- The user provides an amount to undelegate and our `MsgUnDelegateToValidatorSet` divides the amount based on validator weight distribution.
 - Here, the user can either undelegate the entire amount or partial amount 
   - Entire amount unstaking: UnStake 100osmo from validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2},
     our undelegate logic will attempt to undelegate 50osmo from ValA , 30osmo from ValB, 20osmo from ValC
@@ -88,10 +88,10 @@ restaking to a new set is going to happen behind the scenes.
 - Update the `KVStore` value for the specific owner address key.
 - Run the undelegate logic and restake the tokens with updated weights. 
 
-### StakeToValidatorSet
+### DelegateToValidatorSet
 
 Gets the existing validator-set of the delegator and delegates the given amount. The given amount 
-will be divided based on the weights distributed to the validators. The weights will be unchanged! 
+will be divided based on the weights distributed to the validators. The weights will be unchanged 
 
 ```go
     string delegator = 1 [ (gogoproto.moretags) = "yaml:\"delegator\"" ];

x/valset-pref/keeper_test.go

@@ -29,21 +29,26 @@ func (suite *KeeperTestSuite) SetupMultipleValidators(numValidator int) []string
 }
 
 func (suite *KeeperTestSuite) PrepareDelegateToValidatorSet() []types.ValidatorPreference {
-	valAddrs := suite.SetupMultipleValidators(3)
+	valAddrs := suite.SetupMultipleValidators(4)
 	valPreferences := []types.ValidatorPreference{
 		{
 			ValOperAddress: valAddrs[0],
-			Weight:         sdk.NewDecWithPrec(5, 1),
+			Weight:         sdk.NewDecWithPrec(2, 1), // 0.2
 		},
 		{
 			ValOperAddress: valAddrs[1],
-			Weight:         sdk.NewDecWithPrec(3, 1),
+			Weight:         sdk.NewDecWithPrec(332, 3), // 0.332
 		},
 		{
 			ValOperAddress: valAddrs[2],
-			Weight:         sdk.NewDecWithPrec(2, 1),
+			Weight:         sdk.NewDecWithPrec(12, 2), // 0.12
+		},
+		{
+			ValOperAddress: valAddrs[3],
+			Weight:         sdk.NewDecWithPrec(348, 3), // 0.348
 		},
 	}
+
 	return valPreferences
 }
 

x/valset-pref/msg_server.go

@@ -24,7 +24,7 @@ var _ types.MsgServer = msgServer{}
 func (server msgServer) SetValidatorSetPreference(goCtx context.Context, msg *types.MsgSetValidatorSetPreference) (*types.MsgSetValidatorSetPreferenceResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	err := server.keeper.SetupValidatorSetPreference(ctx, msg.Delegator, msg.Preferences)
+	err := server.keeper.SetValidatorSetPreference(ctx, msg.Delegator, msg.Preferences)
 	if err != nil {
 		return nil, err
 	}
@@ -38,10 +38,24 @@ func (server msgServer) SetValidatorSetPreference(goCtx context.Context, msg *ty
 }
 
 func (server msgServer) DelegateToValidatorSet(goCtx context.Context, msg *types.MsgDelegateToValidatorSet) (*types.MsgDelegateToValidatorSetResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	err := server.keeper.DelegateToValidatorSet(ctx, msg.Delegator, msg.Coin)
+	if err != nil {
+		return nil, err
+	}
+
 	return &types.MsgDelegateToValidatorSetResponse{}, nil
 }
 
 func (server msgServer) UndelegateFromValidatorSet(goCtx context.Context, msg *types.MsgUndelegateFromValidatorSet) (*types.MsgUndelegateFromValidatorSetResponse, error) {
+	ctx := sdk.UnwrapSDKContext(goCtx)
+
+	err := server.keeper.UndelegateFromValidatorSet(ctx, msg.Delegator, msg.Coin)
+	if err != nil {
+		return nil, err
+	}
+
 	return &types.MsgUndelegateFromValidatorSetResponse{}, nil
 }
 

x/valset-pref/msg_server_test.go

@@ -124,3 +124,171 @@ func (suite *KeeperTestSuite) TestSetValidatorSetPreference() {
 		})
 	}
 }
+
+func (suite *KeeperTestSuite) TestDelegateToValidatorSet() {
+	suite.SetupTest()
+
+	// prepare validators to delegate to
+	preferences := suite.PrepareDelegateToValidatorSet()
+
+	amountToFund := sdk.Coins{sdk.NewInt64Coin(sdk.DefaultBondDenom, 100_000_000)} // 100 osmo
+
+	tests := []struct {
+		name           string
+		delegator      sdk.AccAddress
+		coin           sdk.Coin  // amount to delegate
+		expectedShares []sdk.Dec // expected shares after delegation
+		expectPass     bool
+		valSetExists   bool
+	}{
+		{
+			name:           "Delegate to valid validators",
+			delegator:      sdk.AccAddress([]byte("addr1---------------")),
+			coin:           sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(10_000_000)),
+			expectedShares: []sdk.Dec{sdk.NewDec(2_000_000), sdk.NewDec(3_320_000), sdk.NewDec(1_200_000), sdk.NewDec(3_480_000)},
+			expectPass:     true,
+		},
+		{
+			name:           "Delegate more tokens to existing validator-set",
+			delegator:      sdk.AccAddress([]byte("addr1---------------")),
+			coin:           sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(10_000_000)),
+			expectedShares: []sdk.Dec{sdk.NewDec(4_000_000), sdk.NewDec(6_640_000), sdk.NewDec(2_400_000), sdk.NewDec(6_960_000)},
+			expectPass:     true,
+			valSetExists:   true,
+		},
+		{
+			name:       "User does not have enough tokens to stake",
+			delegator:  sdk.AccAddress([]byte("addr3---------------")),
+			coin:       sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(200_000_000)),
+			expectPass: false,
+		},
+	}
+
+	for _, test := range tests {
+		suite.Run(test.name, func() {
+			// setup message server
+			msgServer := valPref.NewMsgServerImpl(suite.App.ValidatorSetPreferenceKeeper)
+			c := sdk.WrapSDKContext(suite.Ctx)
+
+			// if validatorSetExist no need to refund and setValSet again
+			if !test.valSetExists {
+				suite.FundAcc(test.delegator, amountToFund)
+
+				_, err := msgServer.SetValidatorSetPreference(c, types.NewMsgSetValidatorSetPreference(test.delegator, preferences))
+				suite.Require().NoError(err)
+			}
+
+			// call the create validator set preference
+			_, err := msgServer.DelegateToValidatorSet(c, types.NewMsgDelegateToValidatorSet(test.delegator, test.coin))
+			if test.expectPass {
+				suite.Require().NoError(err)
+
+				// check if the user balance decreased
+				balance := suite.App.BankKeeper.GetBalance(suite.Ctx, test.delegator, sdk.DefaultBondDenom)
+				expectedBalance := amountToFund[0].Amount.Sub(test.coin.Amount)
+				if test.valSetExists {
+					expectedBalance = balance.Amount
+				}
+
+				suite.Require().Equal(expectedBalance, balance.Amount)
+
+				// check if the expectedShares matches after delegation
+				for i, val := range preferences {
+					valAddr, err := sdk.ValAddressFromBech32(val.ValOperAddress)
+					suite.Require().NoError(err)
+
+					// guarantees that the delegator exists because we check it in DelegateToValidatorSet
+					del, _ := suite.App.StakingKeeper.GetDelegation(suite.Ctx, test.delegator, valAddr)
+					suite.Require().Equal(del.Shares, test.expectedShares[i])
+				}
+
+			} else {
+				suite.Require().Error(err)
+			}
+		})
+	}
+}
+
+func (suite *KeeperTestSuite) TestUnDelegateFromValidatorSet() {
+	tests := []struct {
+		name           string
+		delegator      sdk.AccAddress
+		coinToStake    sdk.Coin
+		coinToUnStake  sdk.Coin
+		expectedShares []sdk.Dec // expected shares after undelegation
+		expectPass     bool
+	}{
+		{
+			name:           "Unstake half from the ValSet",
+			delegator:      sdk.AccAddress([]byte("addr1---------------")),
+			coinToStake:    sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)), // delegate 20osmo
+			coinToUnStake:  sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(10_000_000)), // undelegate 10osmo
+			expectedShares: []sdk.Dec{sdk.NewDec(2_000_000), sdk.NewDec(3_320_000), sdk.NewDec(1_200_000), sdk.NewDec(3_480_000)},
+			expectPass:     true,
+		},
+		{
+			name:           "Unstake x amount from ValSet",
+			delegator:      sdk.AccAddress([]byte("addr2---------------")),
+			coinToStake:    sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),                                           // delegate 20osmo
+			coinToUnStake:  sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(15_000_000)),                                           // undelegate 15osmo
+			expectedShares: []sdk.Dec{sdk.NewDec(1_000_000), sdk.NewDec(1_660_000), sdk.NewDec(600_000), sdk.NewDec(1_740_000)}, // validatorDelegatedShares - (weight * coinToUnstake)
+			expectPass:     true,
+		},
+		{
+			name:          "Unstake everything",
+			delegator:     sdk.AccAddress([]byte("addr3---------------")),
+			coinToStake:   sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			coinToUnStake: sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			expectPass:    true,
+		},
+		{
+			name:          "Unstake more amount than the staked amount",
+			delegator:     sdk.AccAddress([]byte("addr4---------------")),
+			coinToStake:   sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(20_000_000)),
+			coinToUnStake: sdk.NewCoin(sdk.DefaultBondDenom, sdk.NewInt(40_000_000)),
+			expectPass:    false,
+		},
+	}
+
+	for _, test := range tests {
+		suite.Run(test.name, func() {
+			suite.SetupTest()
+
+			suite.FundAcc(test.delegator, sdk.Coins{sdk.NewInt64Coin(sdk.DefaultBondDenom, 100_000_000)}) // 100 osmo
+
+			// setup message server
+			msgServer := valPref.NewMsgServerImpl(suite.App.ValidatorSetPreferenceKeeper)
+			c := sdk.WrapSDKContext(suite.Ctx)
+
+			// call the create validator set preference
+			preferences := suite.PrepareDelegateToValidatorSet()
+
+			_, err := msgServer.SetValidatorSetPreference(c, types.NewMsgSetValidatorSetPreference(test.delegator, preferences))
+			suite.Require().NoError(err)
+
+			// call the create validator set preference
+			_, err = msgServer.DelegateToValidatorSet(c, types.NewMsgDelegateToValidatorSet(test.delegator, test.coinToStake))
+			suite.Require().NoError(err)
+
+			_, err = msgServer.UndelegateFromValidatorSet(c, types.NewMsgUndelegateFromValidatorSet(test.delegator, test.coinToUnStake))
+			if test.expectPass {
+				suite.Require().NoError(err)
+
+				// check if the expectedShares matches after undelegation
+				for i, val := range preferences {
+					valAddr, err := sdk.ValAddressFromBech32(val.ValOperAddress)
+					suite.Require().NoError(err)
+
+					// guarantees that the delegator exists because we check it in UnDelegateToValidatorSet
+					del, found := suite.App.StakingKeeper.GetDelegation(suite.Ctx, test.delegator, valAddr)
+					if found {
+						suite.Require().Equal(del.GetShares(), test.expectedShares[i])
+					}
+				}
+
+			} else {
+				suite.Require().Error(err)
+			}
+		})
+	}
+}

x/valset-pref/types/expected_interfaces.go

@@ -15,3 +15,7 @@ type StakingInterface interface {
 	GetDelegation(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress) (delegation stakingtypes.Delegation, found bool)
 	Undelegate(ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress, sharesAmount sdk.Dec) (time.Time, error)
 }
+
+type BankKeeper interface {
+	GetBalance(ctx sdk.Context, addr sdk.AccAddress, denom string) sdk.Coin
+}

x/valset-pref/types/msgs.go

@@ -74,7 +74,7 @@ const (
 var _ sdk.Msg = &MsgDelegateToValidatorSet{}
 
 // NewMsgMsgStakeToValidatorSet creates a msg to stake to a validator set.
-func NewMsgMsgStakeToValidatorSet(delegator sdk.AccAddress, coin sdk.Coin) *MsgDelegateToValidatorSet {
+func NewMsgDelegateToValidatorSet(delegator sdk.AccAddress, coin sdk.Coin) *MsgDelegateToValidatorSet {
 	return &MsgDelegateToValidatorSet{
 		Delegator: delegator.String(),
 		Coin:      coin,

x/valset-pref/validator_set.go

@@ -9,7 +9,7 @@ import (
 	"github.com/osmosis-labs/osmosis/v13/x/valset-pref/types"
 )
 
-func (k Keeper) SetupValidatorSetPreference(ctx sdk.Context, delegator string, preferences []types.ValidatorPreference) error {
+func (k Keeper) SetValidatorSetPreference(ctx sdk.Context, delegator string, preferences []types.ValidatorPreference) error {
 	// check if a user already has a validator-set created
 	existingValidators, found := k.GetValidatorSetPreference(ctx, delegator)
 	if found {
@@ -29,6 +29,92 @@ func (k Keeper) SetupValidatorSetPreference(ctx sdk.Context, delegator string, p
 	return nil
 }
 
+// DelegateToValidatorSet delegates to a delegators existing validator-set.
+// For ex: delegate 10osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
+// our delegate logic would attempt to delegate 5osmo to A , 2osmo to B, 3osmo to C
+func (k Keeper) DelegateToValidatorSet(ctx sdk.Context, delegatorAddr string, coin sdk.Coin) error {
+	// get the existing validator set preference from store
+	existingSet, found := k.GetValidatorSetPreference(ctx, delegatorAddr)
+	if !found {
+		return fmt.Errorf("user %s doesn't have validator set", delegatorAddr)
+	}
+
+	delegator, err := sdk.AccAddressFromBech32(delegatorAddr)
+	if err != nil {
+		return err
+	}
+
+	// loop through the validatorSetPreference and delegate the proportion of the tokens based on weights
+	for _, val := range existingSet.Preferences {
+		_, validator, err := k.getValAddrAndVal(ctx, val.ValOperAddress)
+		if err != nil {
+			return err
+		}
+
+		// tokenAmt takes the amount to delegate, calculated by {val_distribution_weight * tokenAmt}
+		tokenAmt := val.Weight.Mul(coin.Amount.ToDec()).TruncateInt()
+
+		// TODO: What happens here if validator unbonding
+		// Delegate the unbonded tokens
+		_, err = k.stakingKeeper.Delegate(ctx, delegator, tokenAmt, stakingtypes.Unbonded, validator, true)
+		if err != nil {
+			return err
+		}
+	}
+
+	return nil
+}
+
+// UndelegateFromValidatorSet undelegates {coin} amount from the validator set.
+// For ex: userA has staked 10tokens with weight {Val->0.5, ValB->0.3, ValC->0.2}
+// undelegate 6osmo with validator-set {ValA -> 0.5, ValB -> 0.3, ValC -> 0.2}
+// our undelegate logic would attempt to undelegate 3osmo from A , 1.8osmo from B, 1.2osmo from C
+func (k Keeper) UndelegateFromValidatorSet(ctx sdk.Context, delegatorAddr string, coin sdk.Coin) error {
+	// get the existing validator set preference
+	existingSet, found := k.GetValidatorSetPreference(ctx, delegatorAddr)
+	if !found {
+		return fmt.Errorf("user %s doesn't have validator set", delegatorAddr)
+	}
+
+	delegator, err := sdk.AccAddressFromBech32(delegatorAddr)
+	if err != nil {
+		return err
+	}
+
+	// the total amount the user wants to undelegate
+	tokenAmt := sdk.NewDec(coin.Amount.Int64())
+
+	totalAmountFromWeights := sdk.NewDec(0)
+	for _, val := range existingSet.Preferences {
+		totalAmountFromWeights = totalAmountFromWeights.Add(val.Weight.Mul(tokenAmt))
+	}
+
+	if !totalAmountFromWeights.Equal(tokenAmt) {
+		return fmt.Errorf("The undelegate total do not add up with the amount calculated from weights expected %s got %s", tokenAmt, totalAmountFromWeights)
+	}
+
+	for _, val := range existingSet.Preferences {
+		// Calculate the amount to undelegate based on the existing weights
+		amountToUnDelegate := val.Weight.Mul(tokenAmt)
+
+		valAddr, validator, err := k.getValAddrAndVal(ctx, val.ValOperAddress)
+		if err != nil {
+			return err
+		}
+
+		sharesAmt, err := validator.SharesFromTokens(amountToUnDelegate.TruncateInt())
+		if err != nil {
+			return err
+		}
+
+		_, err = k.stakingKeeper.Undelegate(ctx, delegator, valAddr, sharesAmt) // this has to be shares amount
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
 // GetValAddrAndVal checks if the validator address is valid and the validator provided exists on chain.
 func (k Keeper) getValAddrAndVal(ctx sdk.Context, valOperAddress string) (sdk.ValAddress, stakingtypes.Validator, error) {
 	valAddr, err := sdk.ValAddressFromBech32(valOperAddress)