nitro: use tag for stylus calls

arbos/programs/native.go

@@ -172,6 +172,7 @@ func callProgram(
 	evmData *EvmData,
 	stylusParams *ProgParams,
 	memoryModel *MemoryModel,
+	arbos_tag uint32,
 ) ([]byte, error) {
 	db := interpreter.Evm().StateDB
 	debug := stylusParams.DebugMode
@@ -198,6 +199,7 @@ func callProgram(
 		cbool(debug),
 		output,
 		(*u64)(&scope.Contract.Gas),
+		u32(arbos_tag),
 	))
 
 	depth := interpreter.Depth()
@@ -228,28 +230,30 @@ func cacheProgram(db vm.StateDB, module common.Hash, program Program, params *St
 		if err != nil {
 			panic("unable to recreate wasm")
 		}
-		state.CacheWasmRust(asm, module, program.version, debug)
-		db.RecordCacheWasm(state.CacheWasm{ModuleHash: module, Version: program.version, Debug: debug})
+		tag := db.Database().WasmCacheTag()
+		state.CacheWasmRust(asm, module, program.version, tag, debug)
+		db.RecordCacheWasm(state.CacheWasm{ModuleHash: module, Version: program.version, Tag: tag, Debug: debug})
 	}
 }
 
 // Evicts a program in Rust. We write a record so that we can undo on revert, unless we don't need to (e.g. expired)
 // For gas estimation and eth_call, we ignore permanent updates and rely on Rust's LRU.
 func evictProgram(db vm.StateDB, module common.Hash, version uint16, debug bool, runMode core.MessageRunMode, forever bool) {
 	if runMode == core.MessageCommitMode {
-		state.EvictWasmRust(module, version, debug)
+		tag := db.Database().WasmCacheTag()
+		state.EvictWasmRust(module, version, tag, debug)
 		if !forever {
-			db.RecordEvictWasm(state.EvictWasm{ModuleHash: module, Version: version, Debug: debug})
+			db.RecordEvictWasm(state.EvictWasm{ModuleHash: module, Version: version, Tag: tag, Debug: debug})
 		}
 	}
 }
 
 func init() {
-	state.CacheWasmRust = func(asm []byte, moduleHash common.Hash, version uint16, debug bool) {
-		C.stylus_cache_module(goSlice(asm), hashToBytes32(moduleHash), u16(version), cbool(debug))
+	state.CacheWasmRust = func(asm []byte, moduleHash common.Hash, version uint16, tag uint32, debug bool) {
+		C.stylus_cache_module(goSlice(asm), hashToBytes32(moduleHash), u16(version), u32(tag), cbool(debug))
 	}
-	state.EvictWasmRust = func(moduleHash common.Hash, version uint16, debug bool) {
-		C.stylus_evict_module(hashToBytes32(moduleHash), u16(version), cbool(debug))
+	state.EvictWasmRust = func(moduleHash common.Hash, version uint16, tag uint32, debug bool) {
+		C.stylus_evict_module(hashToBytes32(moduleHash), u16(version), u32(tag), cbool(debug))
 	}
 }
 

arbos/programs/programs.go

@@ -166,6 +166,7 @@ func (p Programs) CallProgram(
 	tracingInfo *util.TracingInfo,
 	calldata []byte,
 	reentrant bool,
+	runmode core.MessageRunMode,
 ) ([]byte, error) {
 	evm := interpreter.Evm()
 	contract := scope.Contract
@@ -237,7 +238,11 @@ func (p Programs) CallProgram(
 	if contract.CodeAddr != nil {
 		address = *contract.CodeAddr
 	}
-	return callProgram(address, moduleHash, localAsm, scope, interpreter, tracingInfo, calldata, evmData, goParams, model)
+	var arbos_tag uint32
+	if runmode == core.MessageCommitMode {
+		arbos_tag = statedb.Database().WasmCacheTag()
+	}
+	return callProgram(address, moduleHash, localAsm, scope, interpreter, tracingInfo, calldata, evmData, goParams, model, arbos_tag)
 }
 
 func getWasm(statedb vm.StateDB, program common.Address) ([]byte, error) {

arbos/programs/wasm.go

@@ -143,6 +143,7 @@ func callProgram(
 	evmData *EvmData,
 	params *ProgParams,
 	memoryModel *MemoryModel,
+	_arbos_tag uint32,
 ) ([]byte, error) {
 	reqHandler := newApiClosures(interpreter, tracingInfo, scope, memoryModel)
 	gasLeft, retData, err := CallProgramLoop(moduleHash, calldata, scope.Contract.Gas, evmData, params, reqHandler)

arbos/tx_processor.go

@@ -127,6 +127,7 @@ func (p *TxProcessor) ExecuteWASM(scope *vm.ScopeContext, input []byte, interpre
 		tracingInfo,
 		input,
 		reentrant,
+		p.RunMode(),
 	)
 }
 

execution/gethexec/executionengine.go

@@ -147,8 +147,9 @@ func (s *ExecutionEngine) Reorg(count arbutil.MessageIndex, newMessages []arbost
 		return nil, nil
 	}
 
+	tag := s.bc.StateCache().WasmCacheTag()
 	// reorg Rust-side VM state
-	C.stylus_reorg_vm(C.uint64_t(blockNum))
+	C.stylus_reorg_vm(C.uint64_t(blockNum), C.uint32_t(tag))
 
 	err := s.bc.ReorgToOldBlock(targetBlock)
 	if err != nil {