fix: Fix superfluid delegation edge case on jailed validator

CHANGELOG.md

@@ -62,7 +62,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### State Machine Breaking
 
 * [#8732](https://github.com/osmosis-labs/osmosis/pull/8732) fix: iterate delegations continue instead of erroring
-
+* [#8778](https://github.com/osmosis-labs/osmosis/pull/8778) fix: Fix superfluid delegation edge case on jailed validator
 ## v26.0.1
 
 ### State Machine Breaking

x/superfluid/keeper/export_test.go

@@ -85,3 +85,9 @@ func (k Keeper) ConvertUnlockedToStake(ctx sdk.Context, sender sdk.AccAddress, v
 func (k Keeper) DelegateBaseOnValsetPref(ctx sdk.Context, sender sdk.AccAddress, valAddr, originalSuperfluidValAddr string, totalAmtToStake osmomath.Int) error {
 	return k.delegateBaseOnValsetPref(ctx, sender, valAddr, originalSuperfluidValAddr, totalAmtToStake)
 }
+
+func (k Keeper) ForceUndelegateAndBurnOsmoTokens(ctx sdk.Context,
+	osmoAmount osmomath.Int, intermediaryAcc types.SuperfluidIntermediaryAccount,
+) error {
+	return k.forceUndelegateAndBurnOsmoTokens(ctx, osmoAmount, intermediaryAcc)
+}

x/superfluid/keeper/stake.go

@@ -8,6 +8,7 @@ import (
 
 	addresscodec "cosmossdk.io/core/address"
 	errorsmod "cosmossdk.io/errors"
+	math "cosmossdk.io/math"
 
 	"github.com/osmosis-labs/osmosis/osmomath"
 	"github.com/osmosis-labs/osmosis/osmoutils"
@@ -503,17 +504,28 @@ func (k Keeper) forceUndelegateAndBurnOsmoTokens(ctx sdk.Context,
 	if err != nil {
 		return err
 	}
-	// TODO: Better understand and decide between ValidateUnbondAmount and SharesFromTokens
-	// briefly looked into it, did not understand what's correct.
-	// TODO: ensure that intermediate account has at least osmoAmount staked.
+
 	shares, err := k.sk.ValidateUnbondAmount(
 		ctx, intermediaryAcc.GetAccAddress(), valAddr, osmoAmount,
 	)
-	if err == stakingtypes.ErrNoDelegation {
-		return nil
-	} else if err != nil {
-		return err
+
+	if err != nil {
+		// if ValidateUnbondAmount has failed it indicates that the amount we're trying to unbond is
+		// greater then what validator has this can be due to different factors (e.g jail)
+		// in this case we brun min(amount_tpo_burn, total_delegation_share)
+		validator, err := k.sk.GetValidator(ctx, valAddr)
+		if err != nil {
+			return err
+		}
+		del, err := k.sk.GetDelegation(ctx, intermediaryAcc.GetAccAddress(), valAddr)
+		if err != nil {
+			return err
+		}
+
+		valShares, err := validator.SharesFromTokens(osmoAmount)
+		shares = math.LegacyMinDec(del.Shares, valShares)
 	}
+
 	err = osmoutils.ApplyFuncIfNoError(ctx, func(cacheCtx sdk.Context) error {
 		undelegatedCoins, err := k.sk.InstantUndelegate(cacheCtx, intermediaryAcc.GetAccAddress(), valAddr, shares)
 		if err != nil {

x/superfluid/keeper/stake_test.go

@@ -1,6 +1,7 @@
 package keeper_test
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/osmosis-labs/osmosis/osmomath"
@@ -14,6 +15,8 @@ import (
 	"github.com/osmosis-labs/osmosis/v26/x/superfluid/types"
 
 	errorsmod "cosmossdk.io/errors"
+	"cosmossdk.io/math"
+	evidencetypes "cosmossdk.io/x/evidence/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/x/bank/testutil"
 	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
@@ -1049,6 +1052,94 @@ func (s *KeeperTestSuite) TestRefreshIntermediaryDelegationAmounts() {
 	}
 }
 
+func (s *KeeperTestSuite) TestForceUndelegateAndBurnOsmoTokens() {
+	testCases := []struct {
+		name                string
+		validatorStats      []stakingtypes.BondStatus
+		superDelegations    []superfluidDelegation
+		jailed              bool
+		jailValWithSmallAmt bool
+
+		expectedShareDiff math.LegacyDec
+	}{
+		{
+			"with single validator and single superfluid delegation and single undelegation",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[]superfluidDelegation{{0, 0, 0, 1000000}},
+			false,
+			false,
+			math.LegacyMustNewDecFromStr("10"),
+		},
+		{
+			"jailed validator where superfluid delegation was major delegation",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[]superfluidDelegation{{0, 0, 0, 1}},
+			true,
+			true,
+			math.LegacyDec{},
+		},
+		{
+			"jailed validator where superfluid delegation was major delegation",
+			[]stakingtypes.BondStatus{stakingtypes.Bonded},
+			[]superfluidDelegation{{0, 0, 0, 1000000}},
+			true,
+			false,
+			math.LegacyMustNewDecFromStr("10.526315789473684210"),
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+		s.Run(tc.name, func() {
+			s.SetupTest()
+
+			// setup validators
+			valAddrs := s.SetupValidators(tc.validatorStats)
+
+			denoms, _ := s.SetupGammPoolsAndSuperfluidAssets([]osmomath.Dec{osmomath.NewDec(20), osmomath.NewDec(20)})
+
+			// setup superfluid delegations
+			_, intermediaryAccs, _ := s.setupSuperfluidDelegations(valAddrs, tc.superDelegations, denoms)
+
+			delegationBeforeUndelegate, err := s.App.StakingKeeper.GetDelegation(s.Ctx, intermediaryAccs[0].GetAccAddress(), valAddrs[0])
+			s.Require().NoError(err)
+
+			// jail the validator as part of set up
+			if tc.jailed {
+				validator, err := s.App.StakingKeeper.GetValidator(s.Ctx, valAddrs[0])
+				s.Require().NoError(err)
+				s.Ctx = s.Ctx.WithBlockHeight(100)
+				consAddr, err := validator.GetConsAddr()
+				s.Require().NoError(err)
+				// slash by slash factor
+				power := sdk.TokensToConsensusPower(validator.Tokens, sdk.DefaultPowerReduction)
+
+				// Note: this calls BeforeValidatorSlashed hook
+				s.handleEquivocationEvidence(s.Ctx, &evidencetypes.Equivocation{
+					Height:           80,
+					Time:             time.Time{},
+					Power:            power,
+					ConsensusAddress: sdk.ConsAddress(consAddr).String(),
+				})
+				val, err := s.App.StakingKeeper.GetValidatorByConsAddr(s.Ctx, consAddr)
+				s.Require().NoError(err)
+				s.Require().Equal(val.Jailed, true)
+			}
+
+			err = s.App.SuperfluidKeeper.ForceUndelegateAndBurnOsmoTokens(s.Ctx, math.NewInt(10), intermediaryAccs[0])
+			s.Require().NoError(err)
+
+			if !tc.jailValWithSmallAmt {
+				delegationAfterUndelegate, err := s.App.StakingKeeper.GetDelegation(s.Ctx, intermediaryAccs[0].GetAccAddress(), valAddrs[0])
+				s.Require().NoError(err)
+
+				shareDiff := delegationBeforeUndelegate.Shares.Sub(delegationAfterUndelegate.Shares)
+				fmt.Println("share diff: ", shareDiff.String())
+				s.Require().True(shareDiff.Equal(tc.expectedShareDiff))
+			}
+		})
+	}
+}
 func (s *KeeperTestSuite) TestUnbondConvertAndStake() {
 	defaultJoinTime := s.Ctx.BlockTime()
 	type tc struct {