Use per-tx escrow for Wei balance (#1314)

* Use per-tx escrow for Wei balance

* rebase

aclmapping/evm/mappings.go

@@ -42,8 +42,8 @@ func TransactionDependencyGenerator(_ aclkeeper.Keeper, evmKeeper evmkeeper.Keep
 		return []sdkacltypes.AccessOperation{}, err
 	}
 	ops := []sdkacltypes.AccessOperation{}
-	ops = appendRWBalanceOps(ops, state.GetMiddleManAddress(ctx))
-	ops = appendRWBalanceOps(ops, state.GetCoinbaseAddress(ctx))
+	ops = appendRWBalanceOps(ops, state.GetMiddleManAddress(ctx.TxIndex()))
+	ops = appendRWBalanceOps(ops, state.GetCoinbaseAddress(ctx.TxIndex()))
 	sender := evmMsg.Derived.SenderSeiAddr
 	ops = appendRWBalanceOps(ops, sender)
 	ops = append(ops,

x/evm/keeper/keeper.go

@@ -148,7 +148,11 @@ func (k *Keeper) GetHashFn(ctx sdk.Context) vm.GetHashFunc {
 func (k *Keeper) GetEVMTxDeferredInfo(ctx sdk.Context) (res []EvmTxDeferredInfo) {
 	k.deferredInfo.Range(func(key, value any) bool {
 		txIdx := key.(int)
-		if txIdx >= 0 && txIdx < len(k.txResults) && k.txResults[txIdx].Code == 0 {
+		if txIdx < 0 || txIdx >= len(k.txResults) {
+			ctx.Logger().Error(fmt.Sprintf("getting invalid tx index in EVM deferred info: %d, num of txs: %d", txIdx, len(k.txResults)))
+			return true
+		}
+		if k.txResults[txIdx].Code == 0 {
 			res = append(res, *(value.(*EvmTxDeferredInfo)))
 		}
 		return true

x/evm/keeper/msg_server_test.go

@@ -69,7 +69,7 @@ func TestEVMTransaction(t *testing.T) {
 	require.NotEmpty(t, res.ReturnData)
 	require.NotEmpty(t, res.Hash)
 	require.Equal(t, uint64(1000000)-res.GasUsed, k.BankKeeper().GetBalance(ctx, sdk.AccAddress(evmAddr[:]), "usei").Amount.Uint64())
-	require.Equal(t, res.GasUsed, k.BankKeeper().GetBalance(ctx, state.GetCoinbaseAddress(ctx), k.GetBaseDenom(ctx)).Amount.Uint64())
+	require.Equal(t, res.GasUsed, k.BankKeeper().GetBalance(ctx, state.GetCoinbaseAddress(ctx.TxIndex()), k.GetBaseDenom(ctx)).Amount.Uint64())
 	receipt, err := k.GetReceipt(ctx, common.HexToHash(res.Hash))
 	require.Nil(t, err)
 	require.NotNil(t, receipt)
@@ -301,7 +301,7 @@ func TestEVMPrecompiles(t *testing.T) {
 	require.NotEmpty(t, res.ReturnData)
 	require.NotEmpty(t, res.Hash)
 	require.Equal(t, uint64(1000000)-res.GasUsed, k.BankKeeper().GetBalance(ctx, sdk.AccAddress(evmAddr[:]), "usei").Amount.Uint64())
-	require.Equal(t, res.GasUsed, k.BankKeeper().GetBalance(ctx, state.GetCoinbaseAddress(ctx), k.GetBaseDenom(ctx)).Amount.Uint64())
+	require.Equal(t, res.GasUsed, k.BankKeeper().GetBalance(ctx, state.GetCoinbaseAddress(ctx.TxIndex()), k.GetBaseDenom(ctx)).Amount.Uint64())
 	receipt, err := k.GetReceipt(ctx, common.HexToHash(res.Hash))
 	require.Nil(t, err)
 	require.NotNil(t, receipt)
@@ -406,7 +406,7 @@ func TestEVMBlockEnv(t *testing.T) {
 	require.NotEmpty(t, res.ReturnData)
 	require.NotEmpty(t, res.Hash)
 	require.Equal(t, uint64(1000000)-res.GasUsed, k.BankKeeper().GetBalance(ctx, sdk.AccAddress(evmAddr[:]), "usei").Amount.Uint64())
-	require.Equal(t, res.GasUsed, k.BankKeeper().GetBalance(ctx, state.GetCoinbaseAddress(ctx), k.GetBaseDenom(ctx)).Amount.Uint64())
+	require.Equal(t, res.GasUsed, k.BankKeeper().GetBalance(ctx, state.GetCoinbaseAddress(ctx.TxIndex()), k.GetBaseDenom(ctx)).Amount.Uint64())
 	receipt, err := k.GetReceipt(ctx, common.HexToHash(res.Hash))
 	require.Nil(t, err)
 	require.NotNil(t, receipt)

x/evm/keeper/wei.go

@@ -0,0 +1,48 @@
+package keeper
+
+import (
+	"fmt"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	"github.com/sei-protocol/sei-chain/x/evm/state"
+)
+
+func (k *Keeper) SettleWeiEscrowAccounts(ctx sdk.Context, evmTxDeferredInfoList []EvmTxDeferredInfo) {
+	denom := k.GetBaseDenom(ctx)
+	// settle surplus escrow first
+	for _, info := range evmTxDeferredInfoList {
+		escrow := state.GetTempWeiEscrowAddress(info.TxIndx)
+		seiBalance := k.BankKeeper().GetBalance(ctx, escrow, denom)
+		if !seiBalance.Amount.IsPositive() {
+			continue
+		}
+		if err := k.BankKeeper().SendCoinsFromAccountToModule(ctx, escrow, banktypes.WeiEscrowName, sdk.NewCoins(seiBalance)); err != nil {
+			ctx.Logger().Error(fmt.Sprintf("failed to send %s from escrow %d to global escrow", seiBalance.String(), info.TxIndx))
+			// This should not happen in any case. We want to halt the chain if it does
+			panic(err)
+		}
+	}
+	// settle deficit escrows
+	for _, info := range evmTxDeferredInfoList {
+		escrow := state.GetTempWeiEscrowAddress(info.TxIndx)
+		seiBalance := k.BankKeeper().GetBalance(ctx, escrow, denom)
+		if !seiBalance.Amount.IsNegative() {
+			continue
+		}
+		settleAmt := sdk.NewCoin(seiBalance.Denom, seiBalance.Amount.Neg())
+		if err := k.BankKeeper().SendCoinsFromModuleToAccount(ctx, banktypes.WeiEscrowName, escrow, sdk.NewCoins(settleAmt)); err != nil {
+			ctx.Logger().Error(fmt.Sprintf("failed to send %s from global escrow to escrow %d", settleAmt.String(), info.TxIndx))
+			// This should not happen in any case. We want to halt the chain if it does
+			panic(err)
+		}
+	}
+	// sanity check
+	for _, info := range evmTxDeferredInfoList {
+		escrow := state.GetTempWeiEscrowAddress(info.TxIndx)
+		seiBalance := k.BankKeeper().GetBalance(ctx, escrow, denom)
+		if !seiBalance.Amount.IsZero() {
+			panic(fmt.Sprintf("failed to settle escrow account %d which still has a balance of %s", info.TxIndx, seiBalance.String()))
+		}
+	}
+}

x/evm/keeper/wei_test.go

@@ -0,0 +1,109 @@
+package keeper_test
+
+import (
+	"math/big"
+	"testing"
+
+	sdk "github.com/cosmos/cosmos-sdk/types"
+	banktypes "github.com/cosmos/cosmos-sdk/x/bank/types"
+	"github.com/ethereum/go-ethereum/common"
+	ethtypes "github.com/ethereum/go-ethereum/core/types"
+	testkeeper "github.com/sei-protocol/sei-chain/testutil/keeper"
+	"github.com/sei-protocol/sei-chain/x/evm/state"
+	"github.com/sei-protocol/sei-chain/x/evm/types"
+	"github.com/stretchr/testify/require"
+	abci "github.com/tendermint/tendermint/abci/types"
+)
+
+func TestSettleCommon(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper()
+
+	_, addr1 := testkeeper.MockAddressPair()
+	_, addr2 := testkeeper.MockAddressPair()
+	_, addr3 := testkeeper.MockAddressPair()
+	amt := sdk.NewCoins(sdk.NewCoin("usei", sdk.NewInt(10)))
+	k.BankKeeper().MintCoins(ctx, types.ModuleName, amt)
+	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, sdk.AccAddress(addr1[:]), amt)
+
+	// addr1 send one sei and one Wei to addr2
+	// escrow 1 would have a one-sei surplus
+	s := state.NewDBImpl(ctx.WithTxIndex(1), k, false)
+	s.SubBalance(addr1, big.NewInt(1_000_000_000_001))
+	s.AddBalance(addr2, big.NewInt(1_000_000_000_001))
+	require.Nil(t, s.Finalize())
+	k.AppendToEvmTxDeferredInfo(ctx.WithTxIndex(1), ethtypes.Bloom{}, common.Hash{})
+
+	// addr2 send two weis to addr3
+	// escrow 2 would have a one-sei surplus
+	s = state.NewDBImpl(ctx.WithTxIndex(2), k, false)
+	s.SubBalance(addr2, big.NewInt(2))
+	s.AddBalance(addr3, big.NewInt(2))
+	require.Nil(t, s.Finalize())
+	k.AppendToEvmTxDeferredInfo(ctx.WithTxIndex(2), ethtypes.Bloom{}, common.Hash{})
+
+	// addr1 send one wei to addr2
+	// escrow 3 would have a one-sei deficit
+	s = state.NewDBImpl(ctx.WithTxIndex(3), k, false)
+	s.SubBalance(addr1, big.NewInt(1))
+	s.AddBalance(addr2, big.NewInt(1))
+	require.Nil(t, s.Finalize())
+	k.AppendToEvmTxDeferredInfo(ctx.WithTxIndex(3), ethtypes.Bloom{}, common.Hash{})
+
+	globalEscrowBalance := k.BankKeeper().GetBalance(ctx, k.AccountKeeper().GetModuleAddress(banktypes.WeiEscrowName), "usei")
+	require.True(t, globalEscrowBalance.Amount.IsZero())
+
+	k.SetTxResults([]*abci.ExecTxResult{{Code: 0}, {Code: 0}, {Code: 0}, {Code: 0}})
+	deferredInfo := k.GetEVMTxDeferredInfo(ctx)
+	k.SettleWeiEscrowAccounts(ctx, deferredInfo)
+	globalEscrowBalance = k.BankKeeper().GetBalance(ctx, k.AccountKeeper().GetModuleAddress(banktypes.WeiEscrowName), "usei")
+	require.Equal(t, int64(1), globalEscrowBalance.Amount.Int64())
+}
+
+func TestSettleMultiRedeem(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper()
+
+	_, addr1 := testkeeper.MockAddressPair()
+	_, addr2 := testkeeper.MockAddressPair()
+	_, addr3 := testkeeper.MockAddressPair()
+	amt := sdk.NewCoins(sdk.NewCoin("usei", sdk.NewInt(1)))
+	k.BankKeeper().MintCoins(ctx, types.ModuleName, amt)
+	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, sdk.AccAddress(addr1[:]), amt)
+	k.BankKeeper().MintCoins(ctx, types.ModuleName, amt)
+	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, sdk.AccAddress(addr2[:]), amt)
+	k.BankKeeper().MintCoins(ctx, types.ModuleName, amt)
+	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, sdk.AccAddress(addr3[:]), amt)
+
+	// addr1 send one Wei to addr3
+	// addr2 send one Wei to addr3
+	// escrow 1 would have a two-sei surplus
+	s := state.NewDBImpl(ctx.WithTxIndex(1), k, false)
+	s.SubBalance(addr1, big.NewInt(1))
+	s.AddBalance(addr3, big.NewInt(1))
+	s.SubBalance(addr2, big.NewInt(1))
+	s.AddBalance(addr3, big.NewInt(1))
+	require.Nil(t, s.Finalize())
+	k.AppendToEvmTxDeferredInfo(ctx.WithTxIndex(1), ethtypes.Bloom{}, common.Hash{})
+
+	// addr3 send one wei to addr1
+	// addr3 send one wei to addr2
+	// addr3 send one wei to addr1
+	// escrow 2 would have a one-sei deficit
+	s = state.NewDBImpl(ctx.WithTxIndex(2), k, false)
+	s.SubBalance(addr3, big.NewInt(1))
+	s.AddBalance(addr1, big.NewInt(1))
+	s.SubBalance(addr3, big.NewInt(1))
+	s.AddBalance(addr2, big.NewInt(1))
+	s.SubBalance(addr3, big.NewInt(1))
+	s.AddBalance(addr1, big.NewInt(1))
+	require.Nil(t, s.Finalize())
+	k.AppendToEvmTxDeferredInfo(ctx.WithTxIndex(2), ethtypes.Bloom{}, common.Hash{})
+
+	globalEscrowBalance := k.BankKeeper().GetBalance(ctx, k.AccountKeeper().GetModuleAddress(banktypes.WeiEscrowName), "usei")
+	require.True(t, globalEscrowBalance.Amount.IsZero())
+
+	k.SetTxResults([]*abci.ExecTxResult{{Code: 0}, {Code: 0}, {Code: 0}})
+	deferredInfo := k.GetEVMTxDeferredInfo(ctx)
+	k.SettleWeiEscrowAccounts(ctx, deferredInfo)
+	globalEscrowBalance = k.BankKeeper().GetBalance(ctx, k.AccountKeeper().GetModuleAddress(banktypes.WeiEscrowName), "usei")
+	require.Equal(t, int64(1), globalEscrowBalance.Amount.Int64())
+}

x/evm/module.go

@@ -171,18 +171,19 @@ func (am AppModule) BeginBlock(sdk.Context, abci.RequestBeginBlock) {
 // returns no validator updates.
 func (am AppModule) EndBlock(ctx sdk.Context, _ abci.RequestEndBlock) []abci.ValidatorUpdate {
 	evmTxDeferredInfoList := am.keeper.GetEVMTxDeferredInfo(ctx)
+	am.keeper.SettleWeiEscrowAccounts(ctx, evmTxDeferredInfoList)
 	denom := am.keeper.GetBaseDenom(ctx)
 	for _, deferredInfo := range evmTxDeferredInfoList {
 		idx := deferredInfo.TxIndx
-		middleManAddress := state.GetMiddleManAddress(sdk.Context{}.WithTxIndex(idx))
+		middleManAddress := state.GetMiddleManAddress(idx)
 		balance := am.keeper.BankKeeper().GetBalance(ctx, middleManAddress, denom)
 		weiBalance := am.keeper.BankKeeper().GetWeiBalance(ctx, middleManAddress)
 		if !balance.Amount.IsZero() || !weiBalance.IsZero() {
 			if err := am.keeper.BankKeeper().SendCoinsAndWei(ctx, middleManAddress, am.keeper.AccountKeeper().GetModuleAddress(types.ModuleName), nil, denom, balance.Amount, weiBalance); err != nil {
 				panic(err)
 			}
 		}
-		coinbaseAddress := state.GetCoinbaseAddress(sdk.Context{}.WithTxIndex(idx))
+		coinbaseAddress := state.GetCoinbaseAddress(idx)
 		balance = am.keeper.BankKeeper().GetBalance(ctx, coinbaseAddress, denom)
 		weiBalance = am.keeper.BankKeeper().GetWeiBalance(ctx, coinbaseAddress)
 		if !balance.Amount.IsZero() || !weiBalance.IsZero() {

x/evm/state/balance.go

@@ -69,5 +69,5 @@ func (s *DBImpl) getSeiAddress(evmAddr common.Address) sdk.AccAddress {
 
 func (s *DBImpl) send(from sdk.AccAddress, to sdk.AccAddress, amt *big.Int) {
 	usei, wei := SplitUseiWeiAmount(amt)
-	s.err = s.k.BankKeeper().SendCoinsAndWei(s.ctx, from, to, nil, s.k.GetBaseDenom(s.ctx), sdk.NewIntFromBigInt(usei), sdk.NewIntFromBigInt(wei))
+	s.err = s.k.BankKeeper().SendCoinsAndWei(s.ctx, from, to, s.weiEscrowAddress, s.k.GetBaseDenom(s.ctx), sdk.NewIntFromBigInt(usei), sdk.NewIntFromBigInt(wei))
 }

x/evm/state/balance_test.go

@@ -15,7 +15,7 @@ func TestAddBalance(t *testing.T) {
 	k, ctx := testkeeper.MockEVMKeeper()
 	amt := sdk.NewCoins(sdk.NewCoin(k.GetBaseDenom(ctx), sdk.NewInt(15)))
 	k.BankKeeper().MintCoins(ctx, types.ModuleName, amt)
-	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, state.GetMiddleManAddress(ctx), amt)
+	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, state.GetMiddleManAddress(ctx.TxIndex()), amt)
 	db := state.NewDBImpl(ctx, k, false)
 	seiAddr, evmAddr := testkeeper.MockAddressPair()
 	require.Equal(t, big.NewInt(0), db.GetBalance(evmAddr))

x/evm/state/check_test.go

@@ -38,7 +38,7 @@ func TestEmpty(t *testing.T) {
 
 	// has balance
 	k.BankKeeper().MintCoins(statedb.Ctx(), types.ModuleName, sdk.NewCoins(sdk.NewCoin("usei", sdk.NewInt(1))))
-	k.BankKeeper().SendCoinsFromModuleToAccount(statedb.Ctx(), types.ModuleName, state.GetMiddleManAddress(ctx), sdk.NewCoins(sdk.NewCoin("usei", sdk.NewInt(1))))
+	k.BankKeeper().SendCoinsFromModuleToAccount(statedb.Ctx(), types.ModuleName, state.GetMiddleManAddress(ctx.TxIndex()), sdk.NewCoins(sdk.NewCoin("usei", sdk.NewInt(1))))
 	statedb.AddBalance(addr, big.NewInt(1000000000000))
 	require.False(t, statedb.Empty(addr))
 

x/evm/state/state_test.go

@@ -61,7 +61,7 @@ func TestCreate(t *testing.T) {
 	statedb.SetState(evmAddr, key, val)
 	statedb.SetTransientState(evmAddr, tkey, tval)
 	k.BankKeeper().MintCoins(ctx, types.ModuleName, sdk.NewCoins(sdk.NewCoin(k.GetBaseDenom(ctx), sdk.NewInt(10))))
-	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, state.GetMiddleManAddress(ctx), sdk.NewCoins(sdk.NewCoin(k.GetBaseDenom(ctx), sdk.NewInt(10))))
+	k.BankKeeper().SendCoinsFromModuleToAccount(ctx, types.ModuleName, state.GetMiddleManAddress(ctx.TxIndex()), sdk.NewCoins(sdk.NewCoin(k.GetBaseDenom(ctx), sdk.NewInt(10))))
 	statedb.AddBalance(evmAddr, big.NewInt(10000000000000))
 	// recreate an account should clear its state, but keep its balance and transient state
 	statedb.CreateAccount(evmAddr)

x/evm/state/statedb.go

@@ -27,6 +27,8 @@ type DBImpl struct {
 	// no single bottleneck for fee collection. Its account state and balance will be deleted
 	// before the block commits
 	coinbaseAddress sdk.AccAddress
+	// a temporary address that will serve as the Wei escrow account for this specific transaction.
+	weiEscrowAddress sdk.AccAddress
 
 	k          EVMKeeper
 	simulation bool
@@ -37,8 +39,9 @@ func NewDBImpl(ctx sdk.Context, k EVMKeeper, simulation bool) *DBImpl {
 		ctx:              ctx,
 		k:                k,
 		snapshottedCtxs:  []sdk.Context{},
-		middleManAddress: GetMiddleManAddress(ctx),
-		coinbaseAddress:  GetCoinbaseAddress(ctx),
+		middleManAddress: GetMiddleManAddress(ctx.TxIndex()),
+		coinbaseAddress:  GetCoinbaseAddress(ctx.TxIndex()),
+		weiEscrowAddress: GetTempWeiEscrowAddress(ctx.TxIndex()),
 		simulation:       simulation,
 		tempStateCurrent: NewTemporaryState(),
 	}
@@ -98,6 +101,7 @@ func (s *DBImpl) Copy() vm.StateDB {
 		k:                s.k,
 		middleManAddress: s.middleManAddress,
 		coinbaseAddress:  s.coinbaseAddress,
+		weiEscrowAddress: s.weiEscrowAddress,
 		simulation:       s.simulation,
 		err:              s.err,
 	}

x/evm/state/utils.go

@@ -13,19 +13,26 @@ var UseiToSweiMultiplier = big.NewInt(1_000_000_000_000)
 
 var MiddleManAddressPrefix = []byte("evm_middleman")
 var CoinbaseAddressPrefix = []byte("evm_coinbase")
+var WeiTmpEscrowPrefix = []byte("evm_weiescrow")
 
-func GetMiddleManAddress(ctx sdk.Context) sdk.AccAddress {
+func GetMiddleManAddress(txIdx int) sdk.AccAddress {
 	txIndexBz := make([]byte, 8)
-	binary.BigEndian.PutUint64(txIndexBz, uint64(ctx.TxIndex()))
+	binary.BigEndian.PutUint64(txIndexBz, uint64(txIdx))
 	return sdk.AccAddress(append(MiddleManAddressPrefix, txIndexBz...))
 }
 
-func GetCoinbaseAddress(ctx sdk.Context) sdk.AccAddress {
+func GetCoinbaseAddress(txIdx int) sdk.AccAddress {
 	txIndexBz := make([]byte, 8)
-	binary.BigEndian.PutUint64(txIndexBz, uint64(ctx.TxIndex()))
+	binary.BigEndian.PutUint64(txIndexBz, uint64(txIdx))
 	return sdk.AccAddress(append(CoinbaseAddressPrefix, txIndexBz...))
 }
 
+func GetTempWeiEscrowAddress(txIdx int) sdk.AccAddress {
+	txIndexBz := make([]byte, 8)
+	binary.BigEndian.PutUint64(txIndexBz, uint64(txIdx))
+	return sdk.AccAddress(append(WeiTmpEscrowPrefix, txIndexBz...))
+}
+
 func SplitUseiWeiAmount(amt *big.Int) (usei *big.Int, wei *big.Int) {
 	wei = new(big.Int).Mod(amt, UseiToSweiMultiplier)
 	usei = new(big.Int).Quo(amt, UseiToSweiMultiplier)