Add tokens to lock keeper function

x/lockup/keeper/lock.go

@@ -281,6 +281,51 @@ func (k Keeper) UnlockPeriodLockByID(ctx sdk.Context, LockID uint64) (*types.Per
 	return lock, err
 }
 
+func (k Keeper) addTokensToLock(ctx sdk.Context, lock *types.PeriodLock, coins sdk.Coins) error {
+	lock.Coins = lock.Coins.Add(coins...)
+
+	err := k.setLock(ctx, *lock)
+	if err != nil {
+		return err
+	}
+
+	// modifications to accumulation store
+	for _, coin := range lock.Coins {
+		// remove previous store for the lock ID and add again with updated value
+		k.accumulationStore(ctx, coin.Denom).Remove(accumulationKey(lock.Duration, lock.ID))
+		k.accumulationStore(ctx, coin.Denom).Set(accumulationKey(lock.Duration, lock.ID), coin.Amount)
+	}
+
+	return nil
+}
+
+// AddTokensToLock locks more tokens into a lockup
+// This also saves the lock to the store.
+func (k Keeper) AddTokensToLockByID(ctx sdk.Context, owner sdk.AccAddress, lockID uint64, coins sdk.Coins) (*types.PeriodLock, error) {
+	lock, err := k.GetLockByID(ctx, lockID)
+	if err != nil {
+		return nil, err
+	}
+	if lock.Owner != owner.String() {
+		return nil, types.ErrNotLockOwner
+	}
+	if err := k.bk.SendCoinsFromAccountToModule(ctx, owner, types.ModuleName, coins); err != nil {
+		return nil, err
+	}
+
+	err = k.addTokensToLock(ctx, lock, coins)
+	if err != nil {
+		return nil, err
+	}
+
+	if k.hooks == nil {
+		return lock, nil
+	}
+
+	k.hooks.OnTokenLocked(ctx, owner, lock.ID, coins, lock.Duration, lock.EndTime)
+	return lock, nil
+}
+
 // LockTokens lock tokens from an account for specified duration
 func (k Keeper) LockTokens(ctx sdk.Context, owner sdk.AccAddress, coins sdk.Coins, duration time.Duration) (types.PeriodLock, error) {
 	ID := k.GetLastLockID(ctx) + 1
@@ -311,12 +356,10 @@ func (k Keeper) ClearAllLockRefKeys(ctx sdk.Context) {
 
 // ResetLock reset lock to lock's previous state on InitGenesis
 func (k Keeper) ResetLock(ctx sdk.Context, lock types.PeriodLock) error {
-	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)
 
 	// store refs by the status of unlock
 	if lock.IsUnlocking() {
@@ -331,6 +374,17 @@ func (k Keeper) ResetLock(ctx sdk.Context, lock types.PeriodLock) error {
 	return k.addLockRefs(ctx, types.KeyPrefixNotUnlocking, lock)
 }
 
+// setLock is a utility to store lock object into the store
+func (k Keeper) setLock(ctx sdk.Context, lock types.PeriodLock) error {
+	store := ctx.KVStore(k.storeKey)
+	bz, err := proto.Marshal(&lock)
+	if err != nil {
+		return err
+	}
+	store.Set(lockStoreKey(lock.ID), bz)
+	return nil
+}
+
 // Lock is a utility to lock coins into module account
 func (k Keeper) Lock(ctx sdk.Context, lock types.PeriodLock) error {
 	owner, err := sdk.AccAddressFromBech32(lock.Owner)

x/lockup/keeper/lock_test.go

@@ -258,6 +258,65 @@ func (suite *KeeperTestSuite) TestLockTokensAlot() {
 	// panic(1)
 }
 
+func (suite *KeeperTestSuite) TestAddTokensToLock() {
+	suite.SetupTest()
+
+	// lock coins
+	addr1 := sdk.AccAddress([]byte("addr1---------------"))
+	coins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
+	suite.LockTokens(addr1, coins, time.Second)
+
+	// check locks
+	locks, err := suite.app.LockupKeeper.GetPeriodLocks(suite.ctx)
+	suite.Require().NoError(err)
+	suite.Require().Len(locks, 1)
+	suite.Require().Equal(locks[0].Coins, coins)
+	// check accumulation store is correctly updated
+	accum := suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, types.QueryCondition{
+		LockQueryType: types.ByDuration,
+		Denom:         "stake",
+		Duration:      time.Second,
+	})
+	suite.Require().Equal(accum.String(), "10")
+
+	// add more tokens to lock
+	addCoins := sdk.Coins{sdk.NewInt64Coin("stake", 10)}
+	suite.app.BankKeeper.SetBalances(suite.ctx, addr1, addCoins)
+	_, err = suite.app.LockupKeeper.AddTokensToLockByID(suite.ctx, addr1, locks[0].ID, addCoins)
+	suite.Require().NoError(err)
+
+	// check locks after adding tokens to lock
+	locks, err = suite.app.LockupKeeper.GetPeriodLocks(suite.ctx)
+	suite.Require().NoError(err)
+	suite.Require().Len(locks, 1)
+	suite.Require().Equal(locks[0].Coins, coins.Add(addCoins...))
+
+	// check accumulation store is correctly updated
+	accum = suite.app.LockupKeeper.GetPeriodLocksAccumulation(suite.ctx, types.QueryCondition{
+		LockQueryType: types.ByDuration,
+		Denom:         "stake",
+		Duration:      time.Second,
+	})
+	suite.Require().Equal(accum.String(), "20")
+
+	// try to add tokens to unavailable lock
+	cacheCtx, _ := suite.ctx.CacheContext()
+	suite.app.BankKeeper.SetBalances(cacheCtx, addr1, addCoins)
+	_, err = suite.app.LockupKeeper.AddTokensToLockByID(cacheCtx, addr1, 1111, addCoins)
+	suite.Require().Error(err)
+
+	// try to add tokens with lack balance
+	cacheCtx, _ = suite.ctx.CacheContext()
+	_, err = suite.app.LockupKeeper.AddTokensToLockByID(cacheCtx, addr1, locks[0].ID, addCoins)
+	suite.Require().Error(err)
+
+	// try to add tokens to lock that is owned by others
+	addr2 := sdk.AccAddress([]byte("addr2---------------"))
+	suite.app.BankKeeper.SetBalances(cacheCtx, addr2, addCoins)
+	_, err = suite.app.LockupKeeper.AddTokensToLockByID(cacheCtx, addr2, locks[0].ID, addCoins)
+	suite.Require().Error(err)
+}
+
 func (suite *KeeperTestSuite) TestEndblockerWithdrawAllMaturedLockups() {
 	suite.SetupTest()
 

x/lockup/spec/05_keeper.md

@@ -46,6 +46,8 @@ type Keeper interface {
     UnlockPeriodLockByID(sdk.Context, LockID uint64) (*types.PeriodLock, error)
     // LockTokens lock tokens from an account for specified duration
     LockTokens(sdk.Context, owner sdk.AccAddress, coins sdk.Coins, duration time.Duration) (types.PeriodLock, error)
+    // AddTokensToLock locks more tokens into a lockup
+    AddTokensToLock(ctx sdk.Context, owner sdk.AccAddress, lockID uint64, coins sdk.Coins) (*types.PeriodLock, error)
     // Lock is a utility to lock coins into module account
     Lock(sdk.Context, lock types.PeriodLock) error
     // Unlock is a utility to unlock coins from module account