Let EnvState own its own compiler instead of unsafe lifetime cast.

Author: jimpo

core/executor/src/wasmtime/function_executor.rs

@@ -48,18 +48,14 @@ pub struct SupervisorFuncRef(*const VMCallerCheckedAnyfunc);
 ///
 /// This is stored as part of the host state of the "env" Wasmtime instance.
 pub struct FunctionExecutorState {
-	// The lifetime of the reference is not actually static, but is unsafely cast to static since
-	// the compiler is owned by the wasmtime-jit `Context` and must be accessed during host calls.
-	compiler: &'static mut Compiler,
 	sandbox_store: sandbox::Store<SupervisorFuncRef>,
 	heap: FreeingBumpHeapAllocator,
 }
 
 impl FunctionExecutorState {
 	/// Constructs a new `FunctionExecutorState`.
-	pub fn new(compiler: &'static mut Compiler, heap_base: u32) -> Self {
+	pub fn new(heap_base: u32) -> Self {
 		FunctionExecutorState {
-			compiler,
 			sandbox_store: sandbox::Store::new(),
 			heap: FreeingBumpHeapAllocator::new(heap_base),
 		}
@@ -87,8 +83,11 @@ impl<'a> FunctionExecutor<'a> {
 	///
 	/// The vmctx MUST come from a call to a function in the "env" module.
 	/// The state MUST be looked up from the host state of the "env" module.
-	pub unsafe fn new(vmctx: *mut VMContext, state: &'a mut FunctionExecutorState)
-		-> Result<Self>
+	pub unsafe fn new(
+		vmctx: *mut VMContext,
+		compiler: &'a mut Compiler,
+		state: &'a mut FunctionExecutorState,
+	) -> Result<Self>
 	{
 		let memory = match (*vmctx).lookup_global_export("memory") {
 			Some(Export::Memory { definition, vmctx: _, memory: _ }) =>
@@ -107,7 +106,7 @@ impl<'a> FunctionExecutor<'a> {
 			_ => None,
 		};
 		Ok(FunctionExecutor {
-			compiler: &mut state.compiler,
+			compiler,
 			sandbox_store: &mut state.sandbox_store,
 			heap: &mut state.heap,
 			memory,

core/executor/src/wasmtime/runtime.rs

@@ -33,7 +33,6 @@ use cranelift_wasm::DefinedFuncIndex;
 use std::cell::RefCell;
 use std::collections::HashMap;
 use std::convert::TryFrom;
-use std::mem;
 use std::rc::Rc;
 use wasm_interface::{HostFunctions, Pointer, WordSize};
 use wasmtime_environ::{Module, translate_signature};
@@ -128,15 +127,18 @@ pub fn create_instance<E: Externalities>(ext: &mut E, code: &[u8], heap_pages: u
 fn create_compiled_unit(code: &[u8])
 	-> std::result::Result<(CompiledModule, Context), WasmError>
 {
-	let isa = target_isa()?;
-	let mut context = Context::with_isa(isa, CompilationStrategy::Cranelift);
+	let compilation_strategy = CompilationStrategy::Cranelift;
+
+	let compiler = new_compiler(compilation_strategy)?;
+	let mut context = Context::new(Box::new(compiler));
 
 	// Enable/disable producing of debug info.
 	context.set_debug_info(false);
 
 	// Instantiate and link the env module.
 	let global_exports = context.get_global_exports();
-	let env_module = instantiate_env_module(global_exports)?;
+	let compiler = new_compiler(compilation_strategy)?;
+	let env_module = instantiate_env_module(global_exports, compiler)?;
 	context.name_instance("env".to_owned(), env_module);
 
 	// Compile the wasm module.
@@ -173,15 +175,8 @@ fn call_method(
 	grow_memory(&mut instance, heap_pages)?;
 
 	// Initialize the function executor state.
-	let executor_state = FunctionExecutorState::new(
-		// Unsafely extend the reference lifetime to static. This is necessary because the
-		// host state must be `Any`. This is OK as we will drop this object either before
-		// exiting the function in the happy case or in the worst case on the next runtime
-		// call, which also resets the host state. Thus this reference can never actually
-		// outlive the Context.
-		unsafe { mem::transmute::<_, &'static mut Compiler>(context.compiler_mut()) },
-		get_heap_base(&instance)?,
-	);
+	let heap_base = get_heap_base(&instance)?;
+	let executor_state = FunctionExecutorState::new(heap_base);
 	reset_env_state(context, Some(executor_state))?;
 
 	// Write the input data into guest memory.
@@ -222,8 +217,10 @@ fn call_method(
 }
 
 /// The implementation is based on wasmtime_wasi::instantiate_wasi.
-fn instantiate_env_module(global_exports: Rc<RefCell<HashMap<String, Option<Export>>>>)
-	-> std::result::Result<InstanceHandle, WasmError>
+fn instantiate_env_module(
+	global_exports: Rc<RefCell<HashMap<String, Option<Export>>>>,
+	compiler: Compiler,
+) -> std::result::Result<InstanceHandle, WasmError>
 {
 	let isa = target_isa()?;
 	let pointer_type = isa.pointer_type();
@@ -260,7 +257,7 @@ fn instantiate_env_module(global_exports: Rc<RefCell<HashMap<String, Option<Expo
 	let imports = Imports::none();
 	let data_initializers = Vec::new();
 	let signatures = PrimaryMap::new();
-	let env_state = EnvState::new::<SubstrateExternals>(code_memory);
+	let env_state = EnvState::new::<SubstrateExternals>(code_memory, compiler);
 
 	let result = InstanceHandle::new(
 		Rc::new(module),
@@ -283,6 +280,11 @@ fn target_isa() -> std::result::Result<Box<dyn TargetIsa>, WasmError> {
 	Ok(isa_builder.finish(cranelift_codegen::settings::Flags::new(flag_builder)))
 }
 
+fn new_compiler(strategy: CompilationStrategy) -> std::result::Result<Compiler, WasmError> {
+	let isa = target_isa()?;
+	Ok(Compiler::new(isa, strategy))
+}
+
 fn clear_globals(global_exports: &mut HashMap<String, Option<Export>>) {
 	global_exports.remove("memory");
 	global_exports.remove("__heap_base");

core/executor/src/wasmtime/trampoline.rs

@@ -22,10 +22,10 @@
 use cranelift_codegen::{Context, binemit, ir, isa};
 use cranelift_codegen::ir::{InstBuilder, StackSlotData, StackSlotKind, TrapCode};
 use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
-use wasmtime_jit::CodeMemory;
+use wasmtime_jit::{CodeMemory, Compiler};
 use wasmtime_runtime::{VMContext, VMFunctionBody};
 use wasm_interface::{HostFunctions, Function, Value, ValueType};
-use std::{cmp, mem, ptr};
+use std::{cmp, ptr};
 
 use crate::error::{Error, Result};
 use crate::wasmtime::function_executor::{FunctionExecutorState, FunctionExecutor};
@@ -34,23 +34,25 @@ use crate::wasmtime::function_executor::{FunctionExecutorState, FunctionExecutor
 /// construct a `FunctionExecutor` which can execute the host call.
 pub struct EnvState {
 	externals: &'static [&'static dyn Function],
+	compiler: Compiler,
+	// The code memory must be kept around on the state to prevent it from being dropped.
+	#[allow(dead_code)]
+	code_memory: CodeMemory,
 	trap: Option<Error>,
 	/// The executor state stored across host calls during a single Wasm runtime call.
 	/// During a runtime call, this MUST be `Some`.
 	pub executor_state: Option<FunctionExecutorState>,
-	// The code memory must be kept around on the state to prevent it from being dropped.
-	#[allow(dead_code)]
-	code_memory: CodeMemory,
 }
 
 impl EnvState {
 	/// Construct a new `EnvState` which owns the given code memory.
-	pub fn new<HF: HostFunctions>(code_memory: CodeMemory) -> Self {
+	pub fn new<HF: HostFunctions>(code_memory: CodeMemory, compiler: Compiler) -> Self {
 		EnvState {
 			externals: HF::functions(),
 			trap: None,
-			executor_state: None,
+			compiler,
 			code_memory,
+			executor_state: None,
 		}
 	}
 
@@ -67,8 +69,9 @@ unsafe extern "C" fn stub_fn(vmctx: *mut VMContext, func_index: u32, values_vec:
 		match stub_fn_inner(
 			vmctx,
 			state.externals,
-			func_index,
+			&mut state.compiler,
 			state.executor_state.as_mut(),
+			func_index,
 			values_vec
 		) {
 			Ok(()) => 0,
@@ -88,8 +91,9 @@ unsafe extern "C" fn stub_fn(vmctx: *mut VMContext, func_index: u32, values_vec:
 unsafe fn stub_fn_inner(
 	vmctx: *mut VMContext,
 	externals: &[&dyn Function],
-	func_index: u32,
+	compiler: &mut Compiler,
 	executor_state: Option<&mut FunctionExecutorState>,
+	func_index: u32,
 	values_vec: *mut i64,
 ) -> Result<()>
 {
@@ -99,7 +103,7 @@ unsafe fn stub_fn_inner(
 		.ok_or_else(|| "executor state is None during call to external function")?;
 
 	// Build the external function context.
-	let mut context = FunctionExecutor::new(vmctx, executor_state)?;
+	let mut context = FunctionExecutor::new(vmctx, compiler, executor_state)?;
 
 	let signature = func.signature();