From 1b4ae82744ea66fd8e6df1d8a8c2a94a5b2d1f0d Mon Sep 17 00:00:00 2001
From: Unity Chaos <86889904+UnityChaos@users.noreply.github.com>
Date: Sun, 20 Feb 2022 00:22:40 -0800
Subject: [PATCH] InstantUndelegate (#117)

* InstantUndelegate + basic test

* Update x/staking/keeper/delegation.go

Co-authored-by: Dev Ojha <ValarDragon@users.noreply.github.com>

* Add comments on InstantUndelegate function

* Update x/staking/keeper/delegation_test.go

Co-authored-by: Dev Ojha <ValarDragon@users.noreply.github.com>
---
 x/staking/keeper/delegation.go      | 34 +++++++++++++++++++++++++
 x/staking/keeper/delegation_test.go | 39 +++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+)

diff --git a/x/staking/keeper/delegation.go b/x/staking/keeper/delegation.go
index 6ed92e37aab9..5fc2f4e2e0ed 100644
--- a/x/staking/keeper/delegation.go
+++ b/x/staking/keeper/delegation.go
@@ -751,6 +751,40 @@ func (k Keeper) Undelegate(
 	return completionTime, 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.
diff --git a/x/staking/keeper/delegation_test.go b/x/staking/keeper/delegation_test.go
index cac61af1fbe2..f92ac9b95f6b 100644
--- a/x/staking/keeper/delegation_test.go
+++ b/x/staking/keeper/delegation_test.go
@@ -177,6 +177,45 @@ func TestUnbondingDelegation(t *testing.T) {
 	require.Equal(t, 0, len(resUnbonds))
 }
 
+// TODO: Document what this test is doing, or add more in-line comments.
+func TestInstantUndelegate(t *testing.T) {
+	_, app, ctx := createTestInput()
+
+	delAddrs := simapp.AddTestAddrsIncremental(app, ctx, 1, sdk.NewInt(10000))
+	valAddrs := simapp.ConvertAddrsToValAddrs(delAddrs)
+
+	startTokens := app.StakingKeeper.TokensFromConsensusPower(ctx, 10)
+	notBondedPool := app.StakingKeeper.GetNotBondedPool(ctx)
+
+	require.NoError(t, simapp.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 := teststaking.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], bondTokens.ToDec())
+	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)
+}
+
 func TestUnbondDelegation(t *testing.T) {
 	_, app, ctx := createTestInput()