refactor(x/accounts/defaults/lockup): Clean up some logic

x/accounts/defaults/lockup/continuous_locking_account.go

@@ -41,7 +41,7 @@ func (cva ContinuousLockingAccount) Init(ctx context.Context, msg *lockuptypes.M
 		return nil, sdkerrors.ErrInvalidRequest.Wrapf("invalid end time %s", msg.EndTime.String())
 	}
 
-	if msg.EndTime.Before(msg.StartTime) {
+	if msg.EndTime.Before(msg.StartTime) || msg.EndTime.Equal(msg.StartTime) {
 		return nil, sdkerrors.ErrInvalidRequest.Wrap("invalid start and end time (must be start before end)")
 	}
 
@@ -89,36 +89,20 @@ func (cva ContinuousLockingAccount) GetLockCoinsInfo(ctx context.Context, blockT
 	unlockedCoins = sdk.Coins{}
 	lockedCoins = sdk.Coins{}
 
-	// We must handle the case where the start time for a lockup account has
-	// been set into the future or when the start of the chain is not exactly
-	// known.
-	startTime, err := cva.StartTime.Get(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-	endTime, err := cva.EndTime.Get(ctx)
-	if err != nil {
-		return nil, nil, err
-	}
-	var originalVesting sdk.Coins
+	var originalLocking sdk.Coins
 	err = cva.IterateCoinEntries(ctx, cva.OriginalLocking, func(key string, value math.Int) (stop bool, err error) {
-		originalVesting = append(originalVesting, sdk.NewCoin(key, value))
-		vestedCoin, vestingCoin, err := cva.GetLockCoinInfoWithDenom(ctx, blockTime, key)
+		originalLocking = append(originalLocking, sdk.NewCoin(key, value))
+		unlockedCoin, lockedCoin, err := cva.GetLockCoinInfoWithDenom(ctx, blockTime, key)
 		if err != nil {
 			return true, err
 		}
-		unlockedCoins = append(unlockedCoins, *vestedCoin)
-		lockedCoins = append(lockedCoins, *vestingCoin)
+		unlockedCoins = append(unlockedCoins, *unlockedCoin)
+		lockedCoins = append(lockedCoins, *lockedCoin)
 		return false, nil
 	})
 	if err != nil {
 		return nil, nil, err
 	}
-	if startTime.After(blockTime) {
-		return unlockedCoins, originalVesting, nil
-	} else if endTime.Before(blockTime) {
-		return originalVesting, lockedCoins, nil
-	}
 
 	return unlockedCoins, lockedCoins, nil
 }

x/accounts/defaults/lockup/go.mod

@@ -11,10 +11,12 @@ require (
 	github.com/cosmos/gogoproto v1.4.12
 )
 
+require cosmossdk.io/x/staking v0.0.0-00010101000000-000000000000
+
 require (
 	buf.build/gen/go/cosmos/gogo-proto/protocolbuffers/go v1.32.0-20240130113600-88ef6483f90f.1 // indirect
 	buf.build/gen/go/tendermint/tendermint/protocolbuffers/go v1.32.0-20231117195010-33ed361a9051.1 // indirect
-	cosmossdk.io/api v0.7.3
+	cosmossdk.io/api v0.7.3 // indirect
 	cosmossdk.io/errors v1.0.1
 	cosmossdk.io/log v1.3.1 // indirect
 	cosmossdk.io/math v1.3.0

x/accounts/defaults/lockup/lockup.go

@@ -11,12 +11,14 @@ import (
 	"cosmossdk.io/collections"
 	collcodec "cosmossdk.io/collections/codec"
 	"cosmossdk.io/core/address"
+	"cosmossdk.io/core/branch"
 	"cosmossdk.io/core/header"
 	errorsmod "cosmossdk.io/errors"
 	"cosmossdk.io/math"
 	"cosmossdk.io/x/accounts/accountstd"
 	lockuptypes "cosmossdk.io/x/accounts/defaults/lockup/types"
 	banktypes "cosmossdk.io/x/bank/types"
+	stakingtypes "cosmossdk.io/x/staking/types"
 
 	codectypes "github.com/cosmos/cosmos-sdk/codec/types"
 	sdk "github.com/cosmos/cosmos-sdk/types"
@@ -53,6 +55,7 @@ func newBaseLockup(d accountstd.Dependencies) *BaseLockup {
 		WithdrawedCoins:  collections.NewMap(d.SchemaBuilder, WithdrawedCoinsPrefix, "withdrawed_coins", collections.StringKey, sdk.IntValue),
 		addressCodec:     d.AddressCodec,
 		headerService:    d.Environment.HeaderService,
+		branchService:    d.Environment.BranchService,
 		EndTime:          collections.NewItem(d.SchemaBuilder, EndTimePrefix, "end_time", collcodec.KeyToValueCodec[time.Time](sdk.TimeKey)),
 	}
 
@@ -66,6 +69,7 @@ type BaseLockup struct {
 	DelegatedFree    collections.Map[string, math.Int]
 	DelegatedLocking collections.Map[string, math.Int]
 	WithdrawedCoins  collections.Map[string, math.Int]
+	branchService    branch.Service
 	addressCodec     address.Codec
 	headerService    header.Service
 	// lockup end time.
@@ -146,18 +150,33 @@ func (bva *BaseLockup) Delegate(
 		return nil, err
 	}
 
-	err = bva.TrackDelegation(
-		ctx,
-		sdk.Coins{*balance},
-		lockedCoins,
-		sdk.Coins{msg.Amount},
-	)
-	if err != nil {
-		return nil, err
-	}
+	responses := []*codectypes.Any{}
 
-	msgDelegate := makeMsgDelegate(delegatorAddress, msg.ValidatorAddress, msg.Amount)
-	responses, err := sendMessage(ctx, msgDelegate)
+	err = bva.branchService.Execute(ctx, func(ctx context.Context) error {
+		err = bva.TrackDelegation(
+			ctx,
+			sdk.Coins{*balance},
+			lockedCoins,
+			sdk.Coins{msg.Amount},
+		)
+		if err != nil {
+			return err
+		}
+
+		msgDelegate := &stakingtypes.MsgDelegate{
+			DelegatorAddress: delegatorAddress,
+			ValidatorAddress: msg.ValidatorAddress,
+			Amount:           msg.Amount,
+		}
+		resp, err := sendMessage(ctx, msgDelegate)
+		if err != nil {
+			return err
+		}
+
+		responses = append(responses, resp...)
+
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -180,13 +199,28 @@ func (bva *BaseLockup) Undelegate(
 		return nil, err
 	}
 
-	err = bva.TrackUndelegation(ctx, sdk.Coins{msg.Amount})
-	if err != nil {
-		return nil, err
-	}
+	responses := []*codectypes.Any{}
+
+	err = bva.branchService.Execute(ctx, func(ctx context.Context) error {
+		err = bva.TrackUndelegation(ctx, sdk.Coins{msg.Amount})
+		if err != nil {
+			return err
+		}
 
-	msgUndelegate := makeMsgUndelegate(delegatorAddress, msg.ValidatorAddress, msg.Amount)
-	responses, err := sendMessage(ctx, msgUndelegate)
+		msgUndelegate := &stakingtypes.MsgUndelegate{
+			DelegatorAddress: delegatorAddress,
+			ValidatorAddress: msg.ValidatorAddress,
+			Amount:           msg.Amount,
+		}
+		resp, err := sendMessage(ctx, msgUndelegate)
+		if err != nil {
+			return err
+		}
+
+		responses = append(responses, resp...)
+
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -221,8 +255,23 @@ func (bva *BaseLockup) SendCoins(
 		return nil, err
 	}
 
-	msgSend := makeMsgSend(fromAddress, msg.ToAddress, msg.Amount)
-	responses, err := sendMessage(ctx, msgSend)
+	responses := []*codectypes.Any{}
+
+	err = bva.branchService.Execute(ctx, func(ctx context.Context) error {
+		msgSend := &banktypes.MsgSend{
+			FromAddress: fromAddress,
+			ToAddress:   msg.ToAddress,
+			Amount:      msg.Amount,
+		}
+		resp, err := sendMessage(ctx, msgSend)
+		if err != nil {
+			return err
+		}
+
+		responses = append(responses, resp...)
+
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -254,59 +303,72 @@ func (bva *BaseLockup) WithdrawUnlockedCoins(
 	}
 
 	amount := sdk.Coins{}
-	for _, denom := range msg.Denoms {
-		balance, err := bva.getBalance(ctx, fromAddress, denom)
-		if err != nil {
-			return nil, err
-		}
-		lockedAmt := lockedCoins.AmountOf(denom)
 
-		// get lockedCoin from that are not bonded for the sent denom
-		notBondedLockedCoin, err := bva.GetNotBondedLockedCoin(ctx, sdk.NewCoin(denom, lockedAmt), denom)
-		if err != nil {
-			return nil, err
-		}
+	err = bva.branchService.Execute(ctx, func(ctx context.Context) error {
+		for _, denom := range msg.Denoms {
+			balance, err := bva.getBalance(ctx, fromAddress, denom)
+			if err != nil {
+				return err
+			}
+			lockedAmt := lockedCoins.AmountOf(denom)
 
-		spendable, err := balance.SafeSub(notBondedLockedCoin)
-		if err != nil {
-			return nil, errorsmod.Wrapf(sdkerrors.ErrInsufficientFunds,
-				"locked amount exceeds account balance funds: %s > %s", notBondedLockedCoin, balance)
-		}
+			// get lockedCoin from that are not bonded for the sent denom
+			notBondedLockedCoin, err := bva.GetNotBondedLockedCoin(ctx, sdk.NewCoin(denom, lockedAmt), denom)
+			if err != nil {
+				return err
+			}
 
-		withdrawedAmt, err := bva.WithdrawedCoins.Get(ctx, denom)
-		if err != nil {
-			return nil, err
-		}
-		originalLockingAmt, err := bva.OriginalLocking.Get(ctx, denom)
-		if err != nil {
-			return nil, err
-		}
+			spendable, err := balance.SafeSub(notBondedLockedCoin)
+			if err != nil {
+				return errorsmod.Wrapf(sdkerrors.ErrInsufficientFunds,
+					"locked amount exceeds account balance funds: %s > %s", notBondedLockedCoin, balance)
+			}
 
-		// withdrawable amount is equal to original locking amount subtract already withdrawed amount
-		withdrawableAmt, err := originalLockingAmt.SafeSub(withdrawedAmt)
-		if err != nil {
-			return nil, err
+			withdrawedAmt, err := bva.WithdrawedCoins.Get(ctx, denom)
+			if err != nil {
+				return err
+			}
+			originalLockingAmt, err := bva.OriginalLocking.Get(ctx, denom)
+			if err != nil {
+				return err
+			}
+
+			// withdrawable amount is equal to original locking amount subtract already withdrawed amount
+			withdrawableAmt, err := originalLockingAmt.SafeSub(withdrawedAmt)
+			if err != nil {
+				return err
+			}
+
+			withdrawAmt := math.MinInt(withdrawableAmt, spendable.Amount)
+			// if zero amount go to the next iteration
+			if withdrawAmt.IsZero() {
+				continue
+			}
+			amount = append(amount, sdk.NewCoin(denom, withdrawAmt))
+
+			// update the withdrawed amount
+			err = bva.WithdrawedCoins.Set(ctx, denom, withdrawedAmt.Add(withdrawAmt))
+			if err != nil {
+				return err
+			}
+		}
+		if len(amount) == 0 {
+			return fmt.Errorf("no tokens available for withdrawing")
 		}
 
-		withdrawAmt := math.MinInt(withdrawableAmt, spendable.Amount)
-		// if zero amount go to the next iteration
-		if withdrawAmt.IsZero() {
-			continue
+		msgSend := &banktypes.MsgSend{
+			FromAddress: fromAddress,
+			ToAddress:   msg.ToAddress,
+			Amount:      amount,
 		}
-		amount = append(amount, sdk.NewCoin(denom, withdrawAmt))
 
-		// update the withdrawed amount
-		err = bva.WithdrawedCoins.Set(ctx, denom, withdrawedAmt.Add(withdrawAmt))
+		_, err = sendMessage(ctx, msgSend)
 		if err != nil {
-			return nil, err
+			return err
 		}
-	}
-	if len(amount) == 0 {
-		return nil, fmt.Errorf("no tokens available for withdrawing")
-	}
 
-	msgSend := makeMsgSend(fromAddress, msg.ToAddress, amount)
-	_, err = sendMessage(ctx, msgSend)
+		return nil
+	})
 	if err != nil {
 		return nil, err
 	}
@@ -461,7 +523,7 @@ func (bva *BaseLockup) TrackUndelegation(ctx context.Context, amount sdk.Coins)
 
 func (bva BaseLockup) getBalance(ctx context.Context, sender, denom string) (*sdk.Coin, error) {
 	// Query account balance for the sent denom
-	balanceQueryReq := banktypes.NewQueryBalanceRequest(sender, denom)
+	balanceQueryReq := &banktypes.QueryBalanceRequest{Address: sender, Denom: denom}
 	resp, err := accountstd.QueryModule[banktypes.QueryBalanceResponse](ctx, balanceQueryReq)
 	if err != nil {
 		return nil, err
@@ -506,7 +568,7 @@ func (bva BaseLockup) checkTokensSendable(ctx context.Context, sender string, am
 	return nil
 }
 
-// IterateSendEnabledEntries iterates over all the SendEnabled entries.
+// IterateCoinEntries iterates over all the CoinEntries entries.
 func (bva BaseLockup) IterateCoinEntries(
 	ctx context.Context,
 	entries collections.Map[string, math.Int],

x/accounts/defaults/lockup/periodic_locking_account.go

@@ -81,6 +81,24 @@ func (pva PeriodicLockingAccount) Init(ctx context.Context, msg *lockuptypes.Msg
 		if err != nil {
 			return nil, err
 		}
+
+		// Set initial value for all withdrawed token
+		err = pva.WithdrawedCoins.Set(ctx, coin.Denom, math.ZeroInt())
+		if err != nil {
+			return nil, err
+		}
+
+		// Set initial value for all delegated free token
+		err = pva.DelegatedFree.Set(ctx, coin.Denom, math.ZeroInt())
+		if err != nil {
+			return nil, err
+		}
+
+		// Set initial value for all delegated locking token
+		err = pva.DelegatedLocking.Set(ctx, coin.Denom, math.ZeroInt())
+		if err != nil {
+			return nil, err
+		}
 	}
 
 	err = pva.StartTime.Set(ctx, msg.StartTime)
@@ -123,7 +141,7 @@ func (pva *PeriodicLockingAccount) WithdrawUnlockedCoins(ctx context.Context, ms
 	return pva.BaseLockup.WithdrawUnlockedCoins(ctx, msg, pva.GetLockedCoinsWithDenoms)
 }
 
-// IterateSendEnabledEntries iterates over all the SendEnabled entries.
+// IteratePeriods iterates over all the Periods entries.
 func (pva PeriodicLockingAccount) IteratePeriods(
 	ctx context.Context,
 	cb func(value lockuptypes.Period) (bool, error),