Partial unlocking implementation (#893)

* add partial unlock logic

* add partial unlock test

* fix lint

* Apply suggestions from code review

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

* add test for case after unlocking period

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

Author: mconcat

proto/osmosis/lockup/tx.proto

@@ -42,5 +42,10 @@ message MsgBeginUnlockingAllResponse { repeated PeriodLock unlocks = 1; }
 message MsgBeginUnlocking {
   string owner = 1 [ (gogoproto.moretags) = "yaml:\"owner\"" ];
   uint64 ID = 2;
+  // Amount of unlocking coins. Unlock all if not set.
+  repeated cosmos.base.v1beta1.Coin coins = 3 [
+    (gogoproto.nullable) = false,
+    (gogoproto.castrepeated) = "github.com/cosmos/cosmos-sdk/types.Coins"
+  ];
 }
 message MsgBeginUnlockingResponse { bool success = 1; }

x/lockup/client/cli/tx.go

@@ -128,6 +128,7 @@ func NewBeginUnlockByIDCmd() *cobra.Command {
 			msg := types.NewMsgBeginUnlocking(
 				clientCtx.GetFromAddress(),
 				uint64(id),
+				nil,
 			)
 
 			return tx.GenerateOrBroadcastTxWithFactory(clientCtx, txf, msg)

x/lockup/keeper/iterator.go

@@ -201,7 +201,7 @@ func (k Keeper) beginUnlockFromIterator(ctx sdk.Context, iterator db.Iterator) (
 
 	locks := k.getLocksFromIterator(ctx, iterator)
 	for _, lock := range locks {
-		err := k.BeginUnlock(ctx, lock)
+		err := k.BeginUnlock(ctx, lock, nil)
 		if err != nil {
 			return locks, err
 		}

x/lockup/keeper/lock.go

@@ -429,22 +429,60 @@ func (k Keeper) Lock(ctx sdk.Context, lock types.PeriodLock) error {
 	return nil
 }
 
+// splitLock splits a lock with the given amount, and stores split new lock to the state
+func (k Keeper) splitLock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Coins) (types.PeriodLock, error) {
+	if lock.IsUnlocking() {
+		return types.PeriodLock{}, fmt.Errorf("cannot split unlocking lock")
+	}
+	lock.Coins = lock.Coins.Sub(coins)
+	err := k.setLock(ctx, lock)
+	if err != nil {
+		return types.PeriodLock{}, err
+	}
+
+	splitLockID := k.GetLastLockID(ctx) + 1
+	k.SetLastLockID(ctx, splitLockID)
+
+	splitLock := types.NewPeriodLock(splitLockID, lock.OwnerAddress(), lock.Duration, lock.EndTime, coins)
+	err = k.setLock(ctx, splitLock)
+	return splitLock, err
+}
+
 // BeginUnlock is a utility to start unlocking coins from NotUnlocking queue
-func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock) error {
-	// remove lock refs from not unlocking queue
+func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock, coins sdk.Coins) error {
+	// sanity check
+	if !coins.IsAllLTE(lock.Coins) {
+		return fmt.Errorf("requested amount to unlock exceedes locked tokens")
+	}
+
+	// If the amount were unlocking is empty, or the entire coins amount, unlock the entire lock.
+	// Otherwise, split the lock into two locks, and fully unlock the newly created lock.
+	// (By virtue, the newly created lock we split into should have the unlock amount)
+	if len(coins) != 0 && !coins.IsEqual(lock.Coins) {
+		// prohibit partial unlock if other locks are referring
+		if k.HasAnySyntheticLockups(ctx, lock.ID) {
+			return fmt.Errorf("cannot partial unlock a lock with synthetic lockup")
+		}
+
+		splitLock, err := k.splitLock(ctx, lock, coins)
+		if err != nil {
+			return err
+		}
+		lock = splitLock
+	}
+
+	// remove lock refs from not unlocking queue if exists
 	err := k.deleteLockRefs(ctx, types.KeyPrefixNotUnlocking, lock)
 	if err != nil {
 		return err
 	}
 
 	// store lock with end time set
 	lock.EndTime = ctx.BlockTime().Add(lock.Duration)
-	store := ctx.KVStore(k.storeKey)
-	bz, err := proto.Marshal(&lock)
+	err = k.setLock(ctx, lock)
 	if err != nil {
 		return err
 	}
-	store.Set(lockStoreKey(lock.ID), bz)
 
 	// add lock refs into unlocking queue
 	err = k.addLockRefs(ctx, types.KeyPrefixUnlocking, lock)
@@ -456,11 +494,7 @@ func (k Keeper) BeginUnlock(ctx sdk.Context, lock types.PeriodLock) error {
 		return nil
 	}
 
-	lockOwner, err := sdk.AccAddressFromBech32(lock.Owner)
-	if err != nil {
-		panic(err)
-	}
-	k.hooks.OnStartUnlock(ctx, lockOwner, lock.ID, lock.Coins, lock.Duration, lock.EndTime)
+	k.hooks.OnStartUnlock(ctx, lock.OwnerAddress(), lock.ID, lock.Coins, lock.Duration, lock.EndTime)
 
 	return nil
 }
@@ -484,7 +518,7 @@ func (k Keeper) Unlock(ctx sdk.Context, lock types.PeriodLock) error {
 // TODO: Revisit for Superfluid Staking
 func (k Keeper) ForceUnlock(ctx sdk.Context, lock types.PeriodLock) error {
 	if !lock.IsUnlocking() {
-		err := k.BeginUnlock(ctx, lock)
+		err := k.BeginUnlock(ctx, lock, nil)
 		if err != nil {
 			return err
 		}

x/lockup/keeper/lock_test.go

@@ -65,7 +65,7 @@ func (suite *KeeperTestSuite) TestBeginUnlockPeriodLock() {
 	suite.Require().Equal(locks[0].IsUnlocking(), false)
 
 	// begin unlock
-	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, locks[0])
+	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, locks[0], nil)
 	suite.Require().NoError(err)
 
 	// check locks
@@ -127,7 +127,7 @@ func (suite *KeeperTestSuite) TestUnlockPeriodLockByID() {
 	// begin unlock
 	lock, err = lockKeeper.GetLockByID(suite.ctx, 1)
 	suite.Require().NoError(err)
-	err = lockKeeper.BeginUnlock(suite.ctx, *lock)
+	err = lockKeeper.BeginUnlock(suite.ctx, *lock, nil)
 	suite.Require().NoError(err)
 
 	// unlock 1s after begin unlock
@@ -190,14 +190,82 @@ func (suite *KeeperTestSuite) TestUnlock() {
 	suite.Require().NoError(err)
 
 	// begin unlock with lock object
-	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock)
+	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, nil)
 	suite.Require().NoError(err)
 
 	// unlock with lock object
 	err = suite.app.LockupKeeper.Unlock(suite.ctx.WithBlockTime(now.Add(time.Second)), lock)
 	suite.Require().NoError(err)
 }
 
+func (suite *KeeperTestSuite) TestPartialUnlock() {
+	suite.SetupTest()
+	now := suite.ctx.BlockTime()
+
+	addr1 := sdk.AccAddress([]byte("addr1---------------"))
+	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
+
+	// lock with balance
+	err := simapp.FundAccount(suite.app.BankKeeper, suite.ctx, addr1, coins)
+	suite.Require().NoError(err)
+	lock, err := suite.app.LockupKeeper.LockTokens(suite.ctx, addr1, coins, time.Second)
+	suite.Require().NoError(err)
+
+	// check unlocking coins
+	unlockings := suite.app.LockupKeeper.GetAccountUnlockingCoins(suite.ctx, addr1)
+	suite.Require().Equal(len(unlockings), 0)
+
+	// check locked coins
+	locked := suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx, addr1)
+	suite.Require().Equal(len(locked), 1)
+	suite.Require().Equal(locked[0].Amount.Int64(), int64(10))
+
+	// test exceeding coins
+	exceedingCoins := sdk.Coins{sdk.NewInt64Coin("stake", 15)}
+	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, exceedingCoins)
+	suite.Require().Error(err)
+
+	// test invalid coins
+	invalidCoins := sdk.Coins{sdk.NewInt64Coin("unknown", 1)}
+	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, invalidCoins)
+	suite.Require().Error(err)
+
+	// begin unlock partial amount
+	partialCoins := sdk.Coins{sdk.NewInt64Coin("stake", 1)}
+	err = suite.app.LockupKeeper.BeginUnlock(suite.ctx, lock, partialCoins)
+	suite.Require().NoError(err)
+
+	// check unlocking coins
+	unlockings = suite.app.LockupKeeper.GetAccountUnlockingCoins(suite.ctx, addr1)
+	suite.Require().Equal(len(unlockings), 1)
+	suite.Require().Equal(unlockings[0].Amount.Int64(), int64(1))
+
+	// check locked coins
+	locked = suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx, addr1)
+	suite.Require().Equal(len(locked), 1)
+	suite.Require().Equal(locked[0].Amount.Int64(), int64(10))
+
+	// check locked coins after the unlocking period
+	locked = suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx.WithBlockTime(now.Add(time.Second)), addr1)
+	suite.Require().Equal(len(locked), 1)
+	suite.Require().Equal(locked[0].Amount.Int64(), int64(9))
+
+	// Finish unlocking partial unlock
+	partialUnlock := suite.app.LockupKeeper.GetAccountPeriodLocks(suite.ctx, addr1)[1]
+	err = suite.app.LockupKeeper.Unlock(suite.ctx.WithBlockTime(now.Add(time.Second)), partialUnlock)
+	suite.Require().NoError(err)
+
+	// check unlocking coins
+	unlockings = suite.app.LockupKeeper.GetAccountUnlockingCoins(suite.ctx, addr1)
+	suite.Require().Equal(len(unlockings), 0)
+
+	// check locked coins
+	locked = suite.app.LockupKeeper.GetAccountLockedCoins(suite.ctx, addr1)
+	suite.Require().Equal(len(locked), 1)
+	suite.Require().Equal(locked[0].Amount.Int64(), int64(9))
+
+}
+
 func (suite *KeeperTestSuite) TestModuleLockedCoins() {
 	suite.SetupTest()
 

x/lockup/keeper/msg_server.go

@@ -78,7 +78,8 @@ func (server msgServer) BeginUnlocking(goCtx context.Context, msg *types.MsgBegi
 	if err != nil {
 		return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
 	}
-	err = server.keeper.BeginUnlock(ctx, *lock)
+
+	err = server.keeper.BeginUnlock(ctx, *lock, msg.Coins)
 	if err != nil {
 		return nil, sdkerrors.Wrap(sdkerrors.ErrInvalidRequest, err.Error())
 	}

x/lockup/keeper/synthetic_lock.go

@@ -46,6 +46,7 @@ func (k Keeper) GetSyntheticLockup(ctx sdk.Context, lockID uint64, suffix string
 func (k Keeper) GetAllSyntheticLockupsByLockup(ctx sdk.Context, lockID uint64) []types.SyntheticLock {
 	store := ctx.KVStore(k.storeKey)
 	iterator := sdk.KVStorePrefixIterator(store, combineKeys(types.KeyPrefixSyntheticLockup, sdk.Uint64ToBigEndian(lockID)))
+	defer iterator.Close()
 
 	synthLocks := []types.SyntheticLock{}
 	for ; iterator.Valid(); iterator.Next() {
@@ -59,9 +60,17 @@ func (k Keeper) GetAllSyntheticLockupsByLockup(ctx sdk.Context, lockID uint64) [
 	return synthLocks
 }
 
+func (k Keeper) HasAnySyntheticLockups(ctx sdk.Context, lockID uint64) bool {
+	store := ctx.KVStore(k.storeKey)
+	iterator := sdk.KVStorePrefixIterator(store, combineKeys(types.KeyPrefixSyntheticLockup, sdk.Uint64ToBigEndian(lockID)))
+	defer iterator.Close()
+	return iterator.Valid()
+}
+
 func (k Keeper) GetAllSyntheticLockups(ctx sdk.Context) []types.SyntheticLock {
 	store := ctx.KVStore(k.storeKey)
 	iterator := sdk.KVStorePrefixIterator(store, types.KeyPrefixSyntheticLockup)
+	defer iterator.Close()
 
 	synthLocks := []types.SyntheticLock{}
 	for ; iterator.Valid(); iterator.Next() {

x/lockup/types/lock.go

@@ -29,6 +29,15 @@ func (p SyntheticLock) IsUnlocking() bool {
 	return !p.EndTime.Equal(time.Time{})
 }
 
+// OwnerAddress returns locks owner address
+func (p PeriodLock) OwnerAddress() sdk.AccAddress {
+	addr, err := sdk.AccAddressFromBech32(p.Owner)
+	if err != nil {
+		panic(err)
+	}
+	return addr
+}
+
 func SumLocksByDenom(locks []PeriodLock, denom string) sdk.Int {
 	sum := sdk.NewInt(0)
 	// validate the denom once, so we can avoid the expensive validate check in the hot loop.

x/lockup/types/msgs.go

@@ -66,10 +66,11 @@ func (m MsgBeginUnlockingAll) GetSigners() []sdk.AccAddress {
 var _ sdk.Msg = &MsgBeginUnlocking{}
 
 // NewMsgBeginUnlocking creates a message to begin unlocking the tokens of a specific lock
-func NewMsgBeginUnlocking(owner sdk.AccAddress, id uint64) *MsgBeginUnlocking {
+func NewMsgBeginUnlocking(owner sdk.AccAddress, id uint64, coins sdk.Coins) *MsgBeginUnlocking {
 	return &MsgBeginUnlocking{
 		Owner: owner.String(),
 		ID:    id,
+		Coins: coins,
 	}
 }