statedb: add cacheing for multistore before precompile runs

x/evm/keeper/erc20.go

@@ -218,11 +218,8 @@ func (k Keeper) CallContractWithInput(
 	// sent by a user
 	txConfig := k.TxConfig(ctx, gethcommon.BigToHash(big.NewInt(0)))
 
-	// Using tmp context to not modify the state in case of evm revert
-	tmpCtx, commitCtx := ctx.CacheContext()
-
 	evmResp, evmObj, err = k.ApplyEvmMsg(
-		tmpCtx, evmMsg, evm.NewNoOpTracer(), commit, evmCfg, txConfig,
+		ctx, evmMsg, evm.NewNoOpTracer(), commit, evmCfg, txConfig,
 	)
 	if err != nil {
 		// We don't know the actual gas used, so consuming the gas limit
@@ -245,7 +242,6 @@ func (k Keeper) CallContractWithInput(
 	} else {
 		// Success, committing the state to ctx
 		if commit {
-			commitCtx()
 			totalGasUsed, err := k.AddToBlockGasUsed(ctx, evmResp.GasUsed)
 			if err != nil {
 				k.ResetGasMeterAndConsumeGas(ctx, ctx.GasMeter().Limit())

x/evm/keeper/precompiles.go

@@ -21,3 +21,7 @@ func (k *Keeper) AddPrecompiles(
 		}
 	}
 }
+
+func (k *Keeper) IsPrecompile(addr gethcommon.Address) bool {
+	return k.precompiles.Has(addr)
+}

x/evm/precompile/funtoken.go

@@ -72,8 +72,7 @@ func (p precompileFunToken) Run(
 	if err != nil {
 		return nil, err
 	}
-	// Dirty journal entries in `StateDB` must be committed
-	return bz, start.StateDB.Commit()
+	return bz, OnRunEnd(start.StateDB, start.SnapshotBeforeRun, p.Address())
 }
 
 func PrecompileFunToken(keepers keepers.PublicKeepers) vm.PrecompiledContract {

x/evm/precompile/precompile.go

@@ -147,6 +147,11 @@ type OnRunStartResult struct {
 	Method *gethabi.Method
 
 	StateDB *statedb.StateDB
+
+	// SnapshotBeforeRun captures the state before precompile execution to enable
+	// proper state reversal if the call fails or if [statedb.JournalChange]
+	// is reverted in general.
+	SnapshotBeforeRun statedb.PrecompileSnapshotBeforeRun
 }
 
 // OnRunStart prepares the execution environment for a precompiled contract call.
@@ -188,19 +193,40 @@ func OnRunStart(
 		err = fmt.Errorf("failed to load the sdk.Context from the EVM StateDB")
 		return
 	}
-	ctx := stateDB.GetContext()
-	if err = stateDB.Commit(); err != nil {
+	cacheCtx, snapshot := stateDB.CacheCtxForPrecompile(contract.Address())
+	if err = stateDB.CommitCacheCtx(); err != nil {
 		return res, fmt.Errorf("error committing dirty journal entries: %w", err)
 	}
 
 	return OnRunStartResult{
-		Args:    args,
-		Ctx:     ctx,
-		Method:  method,
-		StateDB: stateDB,
+		Args:              args,
+		Ctx:               cacheCtx,
+		Method:            method,
+		StateDB:           stateDB,
+		SnapshotBeforeRun: snapshot,
 	}, nil
 }
 
+// OnRunEnd finalizes a precompile execution by saving its state snapshot to the
+// journal. This ensures that any state changes can be properly reverted if needed.
+//
+// Args:
+//   - stateDB: The EVM state database
+//   - snapshot: The state snapshot taken before the precompile executed
+//   - precompileAddr: The address of the precompiled contract
+//
+// The snapshot is critical for maintaining state consistency when:
+//   - The operation gets reverted ([statedb.JournalChange] Revert).
+//   - The precompile modifies state in other modules (e.g., bank, wasm)
+//   - Multiple precompiles are called within a single transaction
+func OnRunEnd(
+	stateDB *statedb.StateDB,
+	snapshot statedb.PrecompileSnapshotBeforeRun,
+	precompileAddr gethcommon.Address,
+) error {
+	return stateDB.SavePrecompileSnapshotToJournal(precompileAddr, snapshot)
+}
+
 var precompileMethodIsTxMap map[PrecompileMethod]bool = map[PrecompileMethod]bool{
 	WasmMethod_execute:      true,
 	WasmMethod_instantiate:  true,

x/evm/precompile/test/export.go

@@ -267,6 +267,7 @@ func IncrementWasmCounterWithExecuteMulti(
 	deps *evmtest.TestDeps,
 	wasmContract sdk.AccAddress,
 	times uint,
+	finalizeTx bool,
 ) (evmObj *vm.EVM) {
 	msgArgsBz := []byte(`
 	{ 
@@ -308,7 +309,7 @@ func IncrementWasmCounterWithExecuteMulti(
 	s.Require().NoError(err)
 
 	ethTxResp, evmObj, err := deps.EvmKeeper.CallContractWithInput(
-		deps.Ctx, deps.Sender.EthAddr, &precompile.PrecompileAddr_Wasm, true, input,
+		deps.Ctx, deps.Sender.EthAddr, &precompile.PrecompileAddr_Wasm, finalizeTx, input,
 	)
 	s.Require().NoError(err)
 	s.Require().NotEmpty(ethTxResp.Ret)

x/evm/precompile/wasm.go

@@ -62,9 +62,7 @@ func (p precompileWasm) Run(
 	if err != nil {
 		return nil, err
 	}
-
-	// Dirty journal entries in `StateDB` must be committed
-	return bz, start.StateDB.Commit()
+	return bz, OnRunEnd(start.StateDB, start.SnapshotBeforeRun, p.Address())
 }
 
 type precompileWasm struct {

x/evm/precompile/wasm_test.go

@@ -100,13 +100,13 @@ func (s *WasmSuite) TestExecuteMultiHappy() {
 	test.AssertWasmCounterState(&s.Suite, deps, wasmContract, 0)
 	// count += 2
 	test.IncrementWasmCounterWithExecuteMulti(
-		&s.Suite, &deps, wasmContract, 2)
+		&s.Suite, &deps, wasmContract, 2, true)
 	// count = 2
 	test.AssertWasmCounterState(&s.Suite, deps, wasmContract, 2)
 	s.assertWasmCounterStateRaw(deps, wasmContract, 2)
 	// count += 67
 	test.IncrementWasmCounterWithExecuteMulti(
-		&s.Suite, &deps, wasmContract, 67)
+		&s.Suite, &deps, wasmContract, 67, true)
 	// count = 69
 	test.AssertWasmCounterState(&s.Suite, deps, wasmContract, 69)
 	s.assertWasmCounterStateRaw(deps, wasmContract, 69)

x/evm/statedb/interfaces.go

@@ -38,4 +38,6 @@ type Keeper interface {
 	// DeleteAccount handles contract's suicide call, clearing the balance,
 	// contract bytecode, contract state, and its native account.
 	DeleteAccount(ctx sdk.Context, addr common.Address) error
+
+	IsPrecompile(addr common.Address) bool
 }

x/evm/statedb/journal.go

@@ -21,6 +21,8 @@ import (
 	"math/big"
 	"sort"
 
+	store "github.com/cosmos/cosmos-sdk/store/types"
+	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/ethereum/go-ethereum/common"
 )
 
@@ -335,3 +337,42 @@ func (ch accessListAddSlotChange) Revert(s *StateDB) {
 func (ch accessListAddSlotChange) Dirtied() *common.Address {
 	return nil
 }
+
+// ------------------------------------------------------
+// PrecompileSnapshotBeforeRun
+
+// PrecompileSnapshotBeforeRun: Precompiles can alter persistent storage of other
+// modules. These changes to persistent storage are not reverted by a `Revert` of
+// [JournalChange] by default, as it generally manages only changes to accounts
+// and Bank balances for ether (NIBI).
+//
+// As a workaround to make state changes from precompiles reversible, we store
+// [PrecompileSnapshotBeforeRun] snapshots that sync and record the prior state
+// of the other modules, allowing precompile calls to truly be reverted.
+//
+// As a simple example, suppose that a transaction calls a precompile.
+//  1. If the precompile changes the state in the Bank Module or Wasm module
+//  2. The call gets reverted (`revert()` in Solidity), which shoud restore the
+//     state to a in-memory snapshot recorded on the StateDB journal.
+//  3. This could cause a problem where changes to the rest of the blockchain state
+//     are still in effect following the reversion in the EVM state DB.
+type PrecompileSnapshotBeforeRun struct {
+	MultiStore store.CacheMultiStore
+	Events     sdk.Events
+	Precompile common.Address
+}
+
+var _ JournalChange = PrecompileSnapshotBeforeRun{}
+
+func (ch PrecompileSnapshotBeforeRun) Revert(s *StateDB) {
+	s.cacheCtx = s.cacheCtx.WithMultiStore(ch.MultiStore)
+	// Rewrite the `writeCacheCtxFn` using the same logic as sdk.Context.CacheCtx
+	s.writeToCommitCtxFromCacheCtx = func() {
+		s.ctx.EventManager().EmitEvents(ch.Events)
+		ch.MultiStore.Write()
+	}
+}
+
+func (ch PrecompileSnapshotBeforeRun) Dirtied() *common.Address {
+	return &ch.Precompile
+}

x/evm/statedb/journal_test.go

@@ -18,7 +18,7 @@ import (
 	"github.com/NibiruChain/nibiru/v2/x/evm/statedb"
 )
 
-func (s *Suite) TestPrecompileSnapshots() {
+func (s *Suite) TestComplexJournalChanges() {
 	deps := evmtest.NewTestDeps()
 	bankDenom := evm.EVMBankDenom
 	s.Require().NoError(testapp.FundAccount(
@@ -32,25 +32,11 @@ func (s *Suite) TestPrecompileSnapshots() {
 
 	wasmContract := test.SetupWasmContracts(&deps, &s.Suite)[1]
 	fmt.Printf("wasmContract: %s\n", wasmContract)
-	assertionsBeforeRun := func(deps *evmtest.TestDeps) {
-		test.AssertWasmCounterState(
-			&s.Suite, *deps, wasmContract, 0,
-		)
-	}
-	run := func(deps *evmtest.TestDeps) *vm.EVM {
-		return test.IncrementWasmCounterWithExecuteMulti(
-			&s.Suite, deps, wasmContract, 7,
-		)
-	}
-	assertionsAfterRun := func(deps *evmtest.TestDeps) {
-		test.AssertWasmCounterState(
-			&s.Suite, *deps, wasmContract, 7,
-		)
-	}
 
 	s.T().Log("Assert before transition")
-
-	assertionsBeforeRun(&deps)
+	test.AssertWasmCounterState(
+		&s.Suite, deps, wasmContract, 0,
+	)
 
 	deployArgs := []any{"name", "SYMBOL", uint8(18)}
 	deployResp, err := evmtest.DeployContract(
@@ -136,15 +122,48 @@ func (s *Suite) TestPrecompileSnapshots() {
 		s.Require().ErrorContains(err, vm.ErrExecutionReverted.Error())
 	})
 
-	s.Run("Precompile calls also start and end clean (no dirty changes)", func() {
-		evmObj = run(&deps)
-		assertionsAfterRun(&deps)
+	s.Run("Precompile calls populate snapshots", func() {
+		s.T().Log("commitEvmTx=true, expect 0 dirty journal entries")
+		commitEvmTx := true
+		evmObj = test.IncrementWasmCounterWithExecuteMulti(
+			&s.Suite, &deps, wasmContract, 7, commitEvmTx,
+		)
+		// assertions after run
+		test.AssertWasmCounterState(
+			&s.Suite, deps, wasmContract, 7,
+		)
 		stateDB, ok := evmObj.StateDB.(*statedb.StateDB)
 		s.Require().True(ok, "error retrieving StateDB from the EVM")
 		if stateDB.DirtiesCount() != 0 {
 			debugDirtiesCountMismatch(stateDB, s.T())
 			s.FailNow("expected 0 dirty journal changes")
 		}
+
+		s.T().Log("commitEvmTx=false, expect dirty journal entries")
+		commitEvmTx = false
+		evmObj = test.IncrementWasmCounterWithExecuteMulti(
+			&s.Suite, &deps, wasmContract, 5, commitEvmTx,
+		)
+		stateDB, ok = evmObj.StateDB.(*statedb.StateDB)
+		s.Require().True(ok, "error retrieving StateDB from the EVM")
+
+		s.T().Log("Expect exactly 1 dirty journal entry for the precompile snapshot")
+		if stateDB.DirtiesCount() != 1 {
+			debugDirtiesCountMismatch(stateDB, s.T())
+			s.FailNow("expected 1 dirty journal changes")
+		}
+
+		s.T().Log("Expect no change since the StateDB has not been committed")
+		test.AssertWasmCounterState(
+			&s.Suite, deps, wasmContract, 7, // 7 = 7 + 0
+		)
+
+		s.T().Log("Expect change after the StateDB gets committed")
+		err = stateDB.Commit()
+		s.Require().NoError(err)
+		test.AssertWasmCounterState(
+			&s.Suite, deps, wasmContract, 12, // 12 = 7 + 5
+		)
 	})
 }
 

x/evm/statedb/state_object.go

@@ -121,8 +121,9 @@ type stateObject struct {
 	address common.Address
 
 	// flags
-	DirtyCode bool
-	Suicided  bool
+	DirtyCode    bool
+	Suicided     bool
+	IsPrecompile bool
 }
 
 // newObject creates a state object.