perf(evm): cache repeated balance lookups

giga/deps/xbank/keeper/send.go

@@ -293,6 +293,7 @@ func (k BaseSendKeeper) initBalances(ctx sdk.Context, addr sdk.AccAddress, balan
 			accountStore.Set([]byte(balance.Denom), bz)
 		}
 	}
+	ctx.InvalidateCachedEVMBalance(addr)
 
 	return nil
 }
@@ -312,6 +313,7 @@ func (k BaseSendKeeper) setBalance(ctx sdk.Context, addr sdk.AccAddress, balance
 		bz := k.cdc.MustMarshal(&balance)
 		accountStore.Set([]byte(balance.Denom), bz)
 	}
+	ctx.InvalidateCachedEVMBalance(addr)
 
 	return nil
 }
@@ -320,13 +322,15 @@ func (k BaseSendKeeper) setWeiBalance(ctx sdk.Context, addr sdk.AccAddress, amt
 	store := prefix.NewStore(k.GetKVStore(ctx), types.WeiBalancesPrefix)
 	if amt.IsZero() {
 		store.Delete(addr)
+		ctx.InvalidateCachedEVMBalance(addr)
 		return nil
 	}
 	val, err := amt.Marshal()
 	if err != nil {
 		return err
 	}
 	store.Set(addr, val)
+	ctx.InvalidateCachedEVMBalance(addr)
 	return nil
 }
 

giga/deps/xevm/keeper/balance.go

@@ -8,10 +8,15 @@ import (
 )
 
 func (k *Keeper) GetBalance(ctx sdk.Context, addr sdk.AccAddress) *big.Int {
+	if balance, ok := ctx.GetCachedEVMBalance(addr); ok {
+		return balance
+	}
 	denom := k.GetBaseDenom(ctx)
 	allUsei := k.BankKeeper().GetBalance(ctx, addr, denom).Amount
 	lockedUsei := k.BankKeeper().LockedCoins(ctx, addr).AmountOf(denom) // LockedCoins doesn't use iterators
 	usei := allUsei.Sub(lockedUsei)
 	wei := k.BankKeeper().GetWeiBalance(ctx, addr)
-	return usei.Mul(state.SdkUseiToSweiMultiplier).Add(wei).BigInt()
+	balance := usei.Mul(state.SdkUseiToSweiMultiplier).Add(wei).BigInt()
+	ctx.SetCachedEVMBalance(addr, balance)
+	return balance
 }

giga/deps/xevm/state/balance_test.go

@@ -4,6 +4,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core/tracing"
 	"github.com/holiman/uint256"
 	testkeeper "github.com/sei-protocol/sei-chain/giga/deps/testutil/keeper"
@@ -13,6 +14,24 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+func TestGetCodeHashCachesBalanceForNoCodeAddress(t *testing.T) {
+	k, ctx := testkeeper.MockEVMKeeper(t)
+	ctx = ctx.WithBlockTime(time.Now())
+	_, evmAddr := testkeeper.MockAddressPair()
+	db := state.NewDBImpl(ctx, k, false)
+	meter := db.Ctx().GasMeter()
+
+	before := meter.GasConsumed()
+	require.Equal(t, common.Hash{}, db.GetCodeHash(evmAddr))
+	firstReadCost := meter.GasConsumed() - before
+
+	before = meter.GasConsumed()
+	require.Equal(t, common.Hash{}, db.GetCodeHash(evmAddr))
+	secondReadCost := meter.GasConsumed() - before
+
+	require.Less(t, secondReadCost, firstReadCost)
+}
+
 func TestAddBalance(t *testing.T) {
 	k, ctx := testkeeper.MockEVMKeeper(t)
 	ctx = ctx.WithBlockTime(time.Now())

giga/deps/xevm/state/statedb.go

@@ -55,6 +55,11 @@ type DBImpl struct {
 
 func NewDBImpl(ctx sdk.Context, k EVMKeeper, simulation bool) *DBImpl {
 	feeCollector, _ := k.GetFeeCollectorAddress(ctx)
+	// Nested DBs reuse the request cache so writes can invalidate balances warmed
+	// by an outer DB. Each snapshot layer still has a distinct context identity.
+	if ctx.EVMBalanceCache() == nil {
+		ctx = ctx.WithEVMBalanceCache(sdk.NewEVMBalanceCache())
+	}
 	s := &DBImpl{
 		ctx:                ctx,
 		k:                  k,
@@ -98,6 +103,7 @@ func (s *DBImpl) Cleanup() {
 	s.tempState = nil
 	s.logger = nil
 	s.snapshottedCtxs = nil
+	s.ctx = s.ctx.WithEVMBalanceCache(nil)
 	clear(s.codeCache)
 }
 
@@ -107,6 +113,7 @@ func (s *DBImpl) CleanupForTracer() {
 		s.ctx = s.snapshottedCtxs[0]
 	}
 	feeCollector, _ := s.k.GetFeeCollectorAddress(s.Ctx())
+	s.ctx = s.ctx.WithEVMBalanceCache(sdk.NewEVMBalanceCache())
 	s.coinbaseEvmAddress = feeCollector
 	s.tempState = NewTemporaryState()
 	s.journal = []journalEntry{}
@@ -121,6 +128,11 @@ func (s *DBImpl) CleanupForTracer() {
 // CacheMultiStore.Write() on any shared store layer.
 func (s *DBImpl) ResetForTracer() {
 	feeCollector, _ := s.k.GetFeeCollectorAddress(s.Ctx())
+	balanceCache := sdk.NewEVMBalanceCache()
+	s.ctx = s.ctx.WithEVMBalanceCache(balanceCache)
+	for i := range s.snapshottedCtxs {
+		s.snapshottedCtxs[i] = s.snapshottedCtxs[i].WithEVMBalanceCache(balanceCache)
+	}
 	s.coinbaseEvmAddress = feeCollector
 	s.tempState = NewTemporaryState()
 	s.journal = []journalEntry{}
@@ -184,12 +196,16 @@ func (s *DBImpl) GetStorageRoot(common.Address) common.Hash {
 }
 
 func (s *DBImpl) Copy() vm.StateDB {
-	newCtx := s.ctx.WithMultiStore(s.ctx.MultiStore().CacheMultiStore()).WithEventManager(sdk.NewEventManager())
+	balanceCache := sdk.NewEVMBalanceCache()
+	newCtx := s.ctx.WithMultiStore(s.ctx.MultiStore().CacheMultiStore()).WithEventManager(sdk.NewEventManager()).WithEVMBalanceCache(balanceCache)
 	journal := make([]journalEntry, len(s.journal))
 	copy(journal, s.journal)
 	snapshots := make([]sdk.Context, len(s.snapshottedCtxs)+1)
 	copy(snapshots, s.snapshottedCtxs)
 	snapshots[len(s.snapshottedCtxs)] = s.ctx
+	for i := range snapshots {
+		snapshots[i] = snapshots[i].WithEVMBalanceCache(balanceCache)
+	}
 	copied := &DBImpl{
 		ctx:                newCtx,
 		snapshottedCtxs:    snapshots,

sei-cosmos/types/context.go

@@ -28,6 +28,7 @@ and standard additions here would be better just to add to the Context struct
 type Context struct {
 	ctx                context.Context
 	ms                 MultiStore
+	multiStoreIdentity *multiStoreIdentity
 	nextMs             MultiStore          // ms of the next height; only used in tracing
 	nextStoreKeys      map[string]struct{} // store key names that should use nextMs
 	header             tmproto.Header
@@ -62,6 +63,7 @@ type Context struct {
 	evmVmError                          string      // EVM VM error during execution
 	evmEntryViaWasmdPrecompile          bool        // EVM is entered via wasmd precompile directly
 	evmPrecompileCalledFromDelegateCall bool        // EVM precompile is called from a delegate call
+	evmBalanceCache                     *EVMBalanceCache
 
 	messageIndex int // Used to track current message being processed
 	txIndex      int
@@ -270,15 +272,16 @@ func NewContext(ms MultiStore, header tmproto.Header, isCheckTx bool) Context {
 	// https://github.com/gogo/protobuf/issues/519
 	header.Time = header.Time.UTC()
 	return Context{
-		ctx:             context.Background(),
-		ms:              ms,
-		header:          header,
-		chainID:         header.ChainID,
-		checkTx:         isCheckTx,
-		gasMeter:        NewInfiniteGasMeter(1, 1),
-		minGasPrice:     DecCoins{},
-		eventManager:    NewEventManager(),
-		evmEventManager: NewEVMEventManager(),
+		ctx:                context.Background(),
+		ms:                 ms,
+		multiStoreIdentity: newMultiStoreIdentity(),
+		header:             header,
+		chainID:            header.ChainID,
+		checkTx:            isCheckTx,
+		gasMeter:           NewInfiniteGasMeter(1, 1),
+		minGasPrice:        DecCoins{},
+		eventManager:       NewEventManager(),
+		evmEventManager:    NewEVMEventManager(),
 	}
 }
 
@@ -291,6 +294,7 @@ func (c Context) WithContext(ctx context.Context) Context {
 // WithMultiStore returns a Context with an updated MultiStore.
 func (c Context) WithMultiStore(ms MultiStore) Context {
 	c.ms = ms
+	c.multiStoreIdentity = newMultiStoreIdentity()
 	return c
 }
 

sei-cosmos/types/evm_balance_cache.go

@@ -0,0 +1,65 @@
+package types
+
+import "math/big"
+
+type multiStoreIdentity byte
+
+func newMultiStoreIdentity() *multiStoreIdentity {
+	return new(multiStoreIdentity)
+}
+
+// EVMBalanceCache stores spendable EVM balances for one execution request.
+type EVMBalanceCache struct {
+	entries map[string]cachedEVMBalance
+}
+
+type cachedEVMBalance struct {
+	multiStoreIdentity *multiStoreIdentity
+	balance            big.Int
+}
+
+// NewEVMBalanceCache returns an empty EVM balance cache.
+func NewEVMBalanceCache() *EVMBalanceCache {
+	return &EVMBalanceCache{entries: map[string]cachedEVMBalance{}}
+}
+
+// EVMBalanceCache returns the balance cache attached to the context, if any.
+func (c Context) EVMBalanceCache() *EVMBalanceCache {
+	return c.evmBalanceCache
+}
+
+// WithEVMBalanceCache returns a context that uses cache for EVM balance reads.
+func (c Context) WithEVMBalanceCache(cache *EVMBalanceCache) Context {
+	c.evmBalanceCache = cache
+	return c
+}
+
+// GetCachedEVMBalance returns the cached balance for addr in the active multistore layer.
+func (c Context) GetCachedEVMBalance(addr AccAddress) (*big.Int, bool) {
+	if c.evmBalanceCache == nil {
+		return nil, false
+	}
+	entry, ok := c.evmBalanceCache.entries[string(addr)]
+	if !ok || entry.multiStoreIdentity != c.multiStoreIdentity {
+		return nil, false
+	}
+	return new(big.Int).Set(&entry.balance), true
+}
+
+// SetCachedEVMBalance stores balance for addr in the active multistore layer.
+func (c Context) SetCachedEVMBalance(addr AccAddress, balance *big.Int) {
+	if c.evmBalanceCache == nil {
+		return
+	}
+	entry := cachedEVMBalance{multiStoreIdentity: c.multiStoreIdentity}
+	entry.balance.Set(balance)
+	c.evmBalanceCache.entries[string(addr)] = entry
+}
+
+// InvalidateCachedEVMBalance removes any cached balance for addr.
+func (c Context) InvalidateCachedEVMBalance(addr AccAddress) {
+	if c.evmBalanceCache == nil {
+		return
+	}
+	delete(c.evmBalanceCache.entries, string(addr))
+}

sei-cosmos/x/auth/keeper/account.go

@@ -63,6 +63,7 @@ func (ak AccountKeeper) SetAccount(ctx sdk.Context, acc types.AccountI) {
 	}
 
 	store.Set(types.AddressStoreKey(addr), bz)
+	ctx.InvalidateCachedEVMBalance(addr)
 }
 
 // RemoveAccount removes an account for the account mapper store.
@@ -71,6 +72,7 @@ func (ak AccountKeeper) RemoveAccount(ctx sdk.Context, acc types.AccountI) {
 	addr := acc.GetAddress()
 	store := ctx.KVStore(ak.key)
 	store.Delete(types.AddressStoreKey(addr))
+	ctx.InvalidateCachedEVMBalance(addr)
 }
 
 // IterateAccounts iterates over all the stored accounts and performs a callback function.

sei-cosmos/x/bank/keeper/send.go

@@ -315,6 +315,7 @@ func (k BaseSendKeeper) initBalances(ctx sdk.Context, addr sdk.AccAddress, balan
 			accountStore.Set([]byte(balance.Denom), bz)
 		}
 	}
+	ctx.InvalidateCachedEVMBalance(addr)
 
 	return nil
 }
@@ -334,6 +335,7 @@ func (k BaseSendKeeper) setBalance(ctx sdk.Context, addr sdk.AccAddress, balance
 		bz := k.cdc.MustMarshal(&balance)
 		accountStore.Set([]byte(balance.Denom), bz)
 	}
+	ctx.InvalidateCachedEVMBalance(addr)
 
 	return nil
 }
@@ -342,13 +344,15 @@ func (k BaseSendKeeper) setWeiBalance(ctx sdk.Context, addr sdk.AccAddress, amt
 	store := prefix.NewStore(ctx.KVStore(k.storeKey), types.WeiBalancesPrefix)
 	if amt.IsZero() {
 		store.Delete(addr)
+		ctx.InvalidateCachedEVMBalance(addr)
 		return nil
 	}
 	val, err := amt.Marshal()
 	if err != nil {
 		return err
 	}
 	store.Set(addr, val)
+	ctx.InvalidateCachedEVMBalance(addr)
 	return nil
 }
 

x/evm/keeper/balance.go

@@ -8,10 +8,15 @@ import (
 )
 
 func (k *Keeper) GetBalance(ctx sdk.Context, addr sdk.AccAddress) *big.Int {
+	if balance, ok := ctx.GetCachedEVMBalance(addr); ok {
+		return balance
+	}
 	denom := k.GetBaseDenom(ctx)
 	allUsei := k.BankKeeper().GetBalance(ctx, addr, denom).Amount
 	lockedUsei := k.BankKeeper().LockedCoins(ctx, addr).AmountOf(denom) // LockedCoins doesn't use iterators
 	usei := allUsei.Sub(lockedUsei)
 	wei := k.BankKeeper().GetWeiBalance(ctx, addr)
-	return usei.Mul(state.SdkUseiToSweiMultiplier).Add(wei).BigInt()
+	balance := usei.Mul(state.SdkUseiToSweiMultiplier).Add(wei).BigInt()
+	ctx.SetCachedEVMBalance(addr, balance)
+	return balance
 }