feat: Recover cancel unbonding logic (#110)

* Recover cancel unbonding logic

* Fix lint and remove unused func

* Redelegate after cancel unbonding

* Refactor cancel unbonding msg function

* fix-lint

* adjust cancel unbonding amount

* rename

* Add more return variables for more meaningful

---------

Co-authored-by: Hieu Vu <gundamaster5@gmail.com>

x/multi-staking/keeper/keeper.go

@@ -127,3 +127,19 @@ func (k Keeper) AdjustUnbondAmount(ctx sdk.Context, delAcc sdk.AccAddress, valAc
 
 	return validator.TokensFromShares(shares).RoundInt(), nil
 }
+
+func (k Keeper) AdjustCancelUnbondingAmount(ctx sdk.Context, delAcc sdk.AccAddress, valAcc sdk.ValAddress, creationHeight int64, amount math.Int) (adjustedAmount math.Int, err error) {
+	undelegation, found := k.stakingKeeper.GetUnbondingDelegation(ctx, delAcc, valAcc)
+	if !found {
+		return math.Int{}, fmt.Errorf("undelegation not found")
+	}
+
+	totalUnbondingAmount := math.ZeroInt()
+	for _, entry := range undelegation.Entries {
+		if entry.CreationHeight == creationHeight {
+			totalUnbondingAmount = totalUnbondingAmount.Add(entry.Balance)
+		}
+	}
+
+	return math.MinInt(totalUnbondingAmount, amount), nil
+}

x/multi-staking/keeper/msg_server.go

@@ -197,7 +197,7 @@ func (k msgServer) Undelegate(goCtx context.Context, msg *stakingtypes.MsgUndele
 func (k msgServer) CancelUnbondingDelegation(goCtx context.Context, msg *stakingtypes.MsgCancelUnbondingDelegation) (*stakingtypes.MsgCancelUnbondingDelegationResponse, error) {
 	ctx := sdk.UnwrapSDKContext(goCtx)
 
-	multiStakerAddr, valAcc, err := types.AccAddrAndValAddrFromStrings(msg.DelegatorAddress, msg.ValidatorAddress)
+	delAcc, valAcc, err := types.AccAddrAndValAddrFromStrings(msg.DelegatorAddress, msg.ValidatorAddress)
 	if err != nil {
 		return nil, err
 	}
@@ -206,23 +206,33 @@ func (k msgServer) CancelUnbondingDelegation(goCtx context.Context, msg *staking
 		return nil, fmt.Errorf("not allow coin")
 	}
 
-	unbondEntry, found := k.keeper.GetUnbondingEntryAtCreationHeight(ctx, multiStakerAddr, valAcc, msg.CreationHeight)
-	if !found {
-		return nil, fmt.Errorf("unbondEntry not found")
+	unlockID := types.MultiStakingUnlockID(msg.DelegatorAddress, msg.ValidatorAddress)
+	cancelUnlockingCoin, err := k.keeper.DecreaseUnlockEntryAmount(ctx, unlockID, msg.Amount.Amount, msg.CreationHeight)
+	if err != nil {
+		return nil, err
 	}
-	cancelUnbondingCoin := sdk.NewCoin(k.keeper.stakingKeeper.BondDenom(ctx), unbondEntry.Balance)
 
-	sdkMsg := &stakingtypes.MsgCancelUnbondingDelegation{
-		DelegatorAddress: msg.DelegatorAddress,
-		ValidatorAddress: msg.ValidatorAddress,
-		Amount:           cancelUnbondingCoin, // replace with cancelUnbondingCoin
+	cancelUnbondingAmount := cancelUnlockingCoin.BondValue()
+	cancelUnbondingAmount, err = k.keeper.AdjustCancelUnbondingAmount(ctx, delAcc, valAcc, msg.CreationHeight, cancelUnbondingAmount)
+	if err != nil {
+		return nil, err
 	}
+	cancelUnbondingCoin := sdk.NewCoin(k.keeper.stakingKeeper.BondDenom(ctx), cancelUnbondingAmount)
 
-	unlockID := types.MultiStakingUnlockID(msg.DelegatorAddress, msg.ValidatorAddress)
-	err = k.keeper.DeleteUnlockEntryAtCreationHeight(ctx, unlockID, msg.CreationHeight)
+	lockID := types.MultiStakingLockID(msg.DelegatorAddress, msg.ValidatorAddress)
+	lock := k.keeper.GetOrCreateMultiStakingLock(ctx, lockID)
+	err = lock.AddCoinToMultiStakingLock(cancelUnlockingCoin)
 	if err != nil {
 		return nil, err
 	}
+	k.keeper.SetMultiStakingLock(ctx, lock)
+
+	sdkMsg := &stakingtypes.MsgCancelUnbondingDelegation{
+		DelegatorAddress: msg.DelegatorAddress,
+		ValidatorAddress: msg.ValidatorAddress,
+		Amount:           cancelUnbondingCoin, // replace with cancelUnbondingCoin
+		CreationHeight:   msg.CreationHeight,
+	}
 
 	return k.stakingMsgServer.CancelUnbondingDelegation(sdk.WrapSDKContext(ctx), sdkMsg)
 }

x/multi-staking/keeper/unlock.go

@@ -5,6 +5,8 @@ import (
 
 	"github.com/realio-tech/multi-staking-module/x/multi-staking/types"
 
+	"cosmossdk.io/math"
+
 	sdk "github.com/cosmos/cosmos-sdk/types"
 )
 
@@ -43,7 +45,7 @@ func (k Keeper) SetMultiStakingUnlockEntry(
 	if found {
 		unlock.AddEntry(ctx.BlockHeight(), multistakingCoin)
 	} else {
-		unlock = types.NewMultiStakingUnlock(ctx.BlockHeight(), multistakingCoin)
+		unlock = types.NewMultiStakingUnlock(unlockID, ctx.BlockHeight(), multistakingCoin)
 	}
 
 	k.SetMultiStakingUnlock(ctx, unlock)
@@ -69,3 +71,34 @@ func (k Keeper) DeleteUnlockEntryAtCreationHeight(
 	k.SetMultiStakingUnlock(ctx, unlock)
 	return nil
 }
+
+func (k Keeper) DecreaseUnlockEntryAmount(
+	ctx sdk.Context, unlockID types.UnlockID,
+	amount math.Int, creationHeight int64,
+) (types.MultiStakingCoin, error) {
+	unlockRecord, found := k.GetMultiStakingUnlock(ctx, unlockID)
+	if !found {
+		return types.MultiStakingCoin{}, fmt.Errorf("not found unlock recored")
+	}
+
+	unlockEntryIndex, found := unlockRecord.FindEntryIndexByHeight(creationHeight)
+	// entryIndex exists
+	if !found {
+		return types.MultiStakingCoin{}, fmt.Errorf("unbonding delegation entry is not found at block height %d", creationHeight)
+	}
+
+	unlockEntry := unlockRecord.Entries[unlockEntryIndex]
+	err := unlockRecord.RemoveCoinFromEntry(unlockEntryIndex, amount)
+	if err != nil {
+		return types.MultiStakingCoin{}, err
+	}
+
+	// set the unlocking record or remove it if there are no more entries
+	if len(unlockRecord.Entries) == 0 {
+		k.DeleteMultiStakingUnlock(ctx, unlockID)
+	} else {
+		k.SetMultiStakingUnlock(ctx, unlockRecord)
+	}
+
+	return types.NewMultiStakingCoin(unlockEntry.UnlockingCoin.Denom, amount, unlockEntry.GetBondWeight()), nil
+}

x/multi-staking/types/unlock.go

@@ -27,52 +27,72 @@ func (e UnlockEntry) String() string {
 //
 //nolint:interfacer
 func NewMultiStakingUnlock(
-	creationHeight int64, weightedCoin MultiStakingCoin,
+	unlockID UnlockID, creationHeight int64, weightedCoin MultiStakingCoin,
 ) MultiStakingUnlock {
 	return MultiStakingUnlock{
+		UnlockID: unlockID,
 		Entries: []UnlockEntry{
 			NewUnlockEntry(creationHeight, weightedCoin),
 		},
 	}
 }
 
+func (unlock *MultiStakingUnlock) FindEntryIndexByHeight(creationHeight int64) (int, bool) {
+	for index, unlockEntry := range unlock.Entries {
+		if unlockEntry.CreationHeight == creationHeight {
+			return index, true
+		}
+	}
+	return -1, false
+}
+
 // AddEntry - append entry to the unbonding delegation
 func (unlock *MultiStakingUnlock) AddEntry(creationHeight int64, weightedCoin MultiStakingCoin) {
 	// Check the entries exists with creation_height and complete_time
-	entryIndex := -1
-	for index, ubdEntry := range unlock.Entries {
-		if ubdEntry.CreationHeight == creationHeight {
-			entryIndex = index
-			break
-		}
-	}
+	entryIndex, found := unlock.FindEntryIndexByHeight(creationHeight)
 	// entryIndex exists
-	if entryIndex != -1 {
-		ubdEntry := unlock.Entries[entryIndex]
-		ubdEntry.UnlockingCoin = ubdEntry.UnlockingCoin.Add(weightedCoin)
+	if !found {
+		unlockEntry := unlock.Entries[entryIndex]
+		unlockEntry.UnlockingCoin = unlockEntry.UnlockingCoin.Add(weightedCoin)
 
 		// update the entry
-		unlock.Entries[entryIndex] = ubdEntry
+		unlock.Entries[entryIndex] = unlockEntry
 	} else {
 		// append the new unbond delegation entry
 		entry := NewUnlockEntry(creationHeight, weightedCoin)
 		unlock.Entries = append(unlock.Entries, entry)
 	}
 }
 
-// RemoveEntry - remove entry at index i to the unbonding delegation
+// RemoveCoinFromEntry - remove multi staking coin from unlocking entry
+func (unlock *MultiStakingUnlock) RemoveCoinFromEntry(entryIndex int, amount math.Int) error {
+	unlockEntry := unlock.Entries[entryIndex]
+	if unlockEntry.UnlockingCoin.Amount.LT(amount) {
+		return fmt.Errorf("cancel amount is greater than the unlocking entry amount")
+	}
+
+	updatedAmount := unlockEntry.UnlockingCoin.Amount.Sub(amount)
+	if updatedAmount.IsZero() {
+		unlock.RemoveEntry(entryIndex)
+	} else {
+		unlock.Entries[entryIndex].UnlockingCoin.Amount = updatedAmount
+	}
+
+	return nil
+}
+
+// RemoveEntry - remove entry at index i to the multi staking unlock
+func (unlock *MultiStakingUnlock) RemoveEntry(i int) {
+	unlock.Entries = append(unlock.Entries[:i], unlock.Entries[i+1:]...)
+}
+
+// RemoveEntryAtCreationHeight - remove entry at creation height to the multi staking unlock
 func (unlock *MultiStakingUnlock) RemoveEntryAtCreationHeight(creationHeight int64) {
 	// Check the entries exists with creation_height and complete_time
-	entryIndex := -1
-	for index, ubdEntry := range unlock.Entries {
-		if ubdEntry.CreationHeight == creationHeight {
-			entryIndex = index
-			break
-		}
-	}
+	entryIndex, found := unlock.FindEntryIndexByHeight(creationHeight)
 	// entryIndex exists
-	if entryIndex != -1 {
-		unlock.Entries = append(unlock.Entries[:entryIndex], unlock.Entries[entryIndex+1:]...)
+	if !found {
+		unlock.RemoveEntry(entryIndex)
 	}
 }