From 5046a6d19a5e1127ef106b0318740cd1bccb429e Mon Sep 17 00:00:00 2001
From: Adam Tucker <adamleetucker@outlook.com>
Date: Sun, 29 Oct 2023 21:18:01 -0600
Subject: [PATCH] InstantUndelegate + basic test

---
 .../staking/keeper/delegation_test.go         | 38 +++++++++++++++++++
 x/staking/keeper/delegation.go                | 34 +++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/tests/integration/staking/keeper/delegation_test.go b/tests/integration/staking/keeper/delegation_test.go
index a68c68ee13e9..35f56f7f90bf 100644
--- a/tests/integration/staking/keeper/delegation_test.go
+++ b/tests/integration/staking/keeper/delegation_test.go
@@ -113,3 +113,41 @@ func TestUnbondingDelegationsMaxEntries(t *testing.T) {
 	assert.Assert(math.IntEq(t, newBonded, oldBonded.SubRaw(1)))
 	assert.Assert(math.IntEq(t, newNotBonded, oldNotBonded.AddRaw(1)))
 }
+
+func TestInstantUndelegate(t *testing.T) {
+	_, app, ctx := createTestInput(t)
+
+	delAddrs := simapp.AddTestAddrsIncremental(app, ctx, 1, sdk.NewInt(10000))
+	valAddrs := simtestutil.ConvertAddrsToValAddrs(delAddrs)
+
+	startTokens := app.StakingKeeper.TokensFromConsensusPower(ctx, 10)
+	notBondedPool := app.StakingKeeper.GetNotBondedPool(ctx)
+
+	require.NoError(t, banktestutil.FundModuleAccount(app.BankKeeper, ctx, notBondedPool.GetName(), sdk.NewCoins(sdk.NewCoin(app.StakingKeeper.BondDenom(ctx), startTokens))))
+	app.AccountKeeper.SetModuleAccount(ctx, notBondedPool)
+
+	// create a validator and a delegator to that validator
+	// note this validator starts not-bonded
+	validator := testutil.NewValidator(t, valAddrs[0], PKs[0])
+
+	validator, issuedShares := validator.AddTokensFromDel(startTokens)
+	require.Equal(t, startTokens, issuedShares.RoundInt())
+
+	validator = keeper.TestingUpdateValidator(app.StakingKeeper, ctx, validator, true)
+
+	delegation := types.NewDelegation(delAddrs[0], valAddrs[0], issuedShares)
+	app.StakingKeeper.SetDelegation(ctx, delegation)
+
+	bondTokens := app.StakingKeeper.TokensFromConsensusPower(ctx, 6)
+
+	oldBal := app.BankKeeper.GetBalance(ctx, delAddrs[0], app.StakingKeeper.BondDenom(ctx))
+
+	res, err := app.StakingKeeper.InstantUndelegate(ctx, delAddrs[0], valAddrs[0], sdk.NewDecFromInt(bondTokens))
+	require.NoError(t, err)
+
+	require.Equal(t, res, sdk.NewCoins(sdk.NewCoin(app.StakingKeeper.BondDenom(ctx), bondTokens)))
+
+	newBal := app.BankKeeper.GetBalance(ctx, delAddrs[0], app.StakingKeeper.BondDenom(ctx))
+
+	require.Equal(t, oldBal.Add(sdk.NewCoin(app.StakingKeeper.BondDenom(ctx), bondTokens)), newBal)
+}
diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go
index 606cb1e51c89..ec6fa815217d 100644
--- a/x/staking/keeper/delegation.go
+++ b/x/staking/keeper/delegation.go
@@ -1151,6 +1151,40 @@ func (k Keeper) Undelegate(
 	return completionTime, returnAmount, nil
 }
 
+// InstantUndelegate allows another module account to undelegate while bypassing unbonding time.
+// This function is a combination of Undelegate and CompleteUnbonding,
+// but skips the creation and deletion of UnbondingDelegationEntry
+func (k Keeper) InstantUndelegate(
+	ctx sdk.Context, delAddr sdk.AccAddress, valAddr sdk.ValAddress, sharesAmount sdk.Dec,
+) (sdk.Coins, error) {
+
+	validator, found := k.GetValidator(ctx, valAddr)
+	if !found {
+		return nil, types.ErrNoDelegatorForAddress
+	}
+
+	returnAmount, err := k.Unbond(ctx, delAddr, valAddr, sharesAmount)
+	if err != nil {
+		return nil, err
+	}
+
+	bondDenom := k.GetParams(ctx).BondDenom
+
+	amt := sdk.NewCoin(bondDenom, returnAmount)
+	res := sdk.NewCoins(amt)
+
+	moduleName := types.NotBondedPoolName
+	if validator.IsBonded() {
+		moduleName = types.BondedPoolName
+	}
+	err = k.bankKeeper.UndelegateCoinsFromModuleToAccount(ctx, moduleName, delAddr, res)
+	if err != nil {
+		return nil, err
+	}
+	return res, nil
+
+}
+
 // CompleteUnbonding completes the unbonding of all mature entries in the
 // retrieved unbonding delegation object and returns the total unbonding balance
 // or an error upon failure.