diff --git a/lib/api/src/externals/function.rs b/lib/api/src/externals/function.rs index 832a1eb3e93..bd629d1b944 100644 --- a/lib/api/src/externals/function.rs +++ b/lib/api/src/externals/function.rs @@ -728,14 +728,16 @@ impl VMDynamicFunction for VMDynamicFunctionWithoutEnv { } } +#[repr(C)] pub(crate) struct VMDynamicFunctionWithEnv where Env: Sized + 'static, { + // This field _must_ come first in this struct. + env: Box, function_type: FunctionType, #[allow(clippy::type_complexity)] func: Box Result, RuntimeError> + 'static>, - env: Box, } impl VMDynamicFunction for VMDynamicFunctionWithEnv diff --git a/lib/engine/src/resolver.rs b/lib/engine/src/resolver.rs index 0150a2f4d25..f8d3ca40d88 100644 --- a/lib/engine/src/resolver.rs +++ b/lib/engine/src/resolver.rs @@ -7,8 +7,9 @@ use wasmer_types::entity::{BoxedSlice, EntityRef, PrimaryMap}; use wasmer_types::{ExternType, FunctionIndex, ImportIndex, MemoryIndex, TableIndex}; use wasmer_vm::{ - FunctionBodyPtr, Imports, MemoryStyle, ModuleInfo, TableStyle, VMFunctionBody, - VMFunctionImport, VMFunctionKind, VMGlobalImport, VMMemoryImport, VMTableImport, + FunctionBodyPtr, Imports, MemoryStyle, ModuleInfo, TableStyle, VMDynamicFunctionContext, + VMFunctionBody, VMFunctionImport, VMFunctionKind, VMGlobalImport, VMMemoryImport, + VMTableImport, }; /// Import resolver connects imports with available exported values. @@ -154,13 +155,36 @@ pub fn resolve_imports( } match resolved { Export::Function(ref f) => { - let address = match f.vm_function.kind { + let (address, env_ptr) = match f.vm_function.kind { VMFunctionKind::Dynamic => { // If this is a dynamic imported function, // the address of the function is the address of the // reverse trampoline. let index = FunctionIndex::new(function_imports.len()); - finished_dynamic_function_trampolines[index].0 as *mut VMFunctionBody as _ + let address = finished_dynamic_function_trampolines[index].0 + as *mut VMFunctionBody as _; + let env_ptr = if f.import_init_function_ptr.is_some() { + // Our function env looks like: + // Box>> + // Which we can interpret as `*const *const Env` (due to + // the precise layout of these types via `repr(C)`) + // We extract the `*const Env`: + unsafe { + // Box> + let dyn_func_ctx_ptr = f.vm_function.vmctx.host_env + as *mut VMDynamicFunctionContext<*mut std::ffi::c_void>; + // maybe report error here if it's null? + // invariants of these types are not enforced. + + // &VMDynamicFunctionContext<...> + let dyn_func_ctx = &*dyn_func_ctx_ptr; + dyn_func_ctx.ctx + } + } else { + std::ptr::null_mut() + }; + + (address, env_ptr) // TODO: We should check that the f.vmctx actually matches // the shape of `VMDynamicFunctionImportContext` @@ -174,7 +198,9 @@ pub fn resolve_imports( assert!(num_params < 9, "Only native functions with less than 9 arguments are allowed in Apple Silicon (for now). Received {} in the import {}.{}", num_params, module_name, field); } - f.vm_function.address + (f.vm_function.address, unsafe { + f.vm_function.vmctx.host_env + }) } }; function_imports.push(VMFunctionImport { @@ -182,7 +208,7 @@ pub fn resolve_imports( environment: f.vm_function.vmctx, }); - host_function_env_initializers.push(f.import_init_function_ptr); + host_function_env_initializers.push((f.import_init_function_ptr, env_ptr)); } Export::Table(ref t) => { table_imports.push(VMTableImport { diff --git a/lib/vm/src/imports.rs b/lib/vm/src/imports.rs index af1b4bf4c02..0f9482e4092 100644 --- a/lib/vm/src/imports.rs +++ b/lib/vm/src/imports.rs @@ -17,8 +17,9 @@ pub struct Imports { /// space may affect Wasm runtime performance due to increased cache pressure. /// /// We make it optional so that we can free the data after use. - pub host_function_env_initializers: - Option>>, + pub host_function_env_initializers: Option< + BoxedSlice, *mut std::ffi::c_void)>, + >, /// Resolved addresses for imported tables. pub tables: BoxedSlice, @@ -34,7 +35,10 @@ impl Imports { /// Construct a new `Imports` instance. pub fn new( function_imports: PrimaryMap, - host_function_env_initializers: PrimaryMap>, + host_function_env_initializers: PrimaryMap< + FunctionIndex, + (Option, *mut std::ffi::c_void), + >, table_imports: PrimaryMap, memory_imports: PrimaryMap, global_imports: PrimaryMap, @@ -69,12 +73,9 @@ impl Imports { inner .values() .cloned() - .zip(self.functions.values()) - .map(|(func_init, func)| { + .map(|(func_init, env_ptr)| { let host_env = if func_init.is_some() { - // this access is correct because we know that only functions with - // host envs have a value in `func_init`. - unsafe { func.environment.host_env } + env_ptr } else { std::ptr::null_mut() }; diff --git a/tests/compilers/imports.rs b/tests/compilers/imports.rs index 9c525a9a073..3d35944f4db 100644 --- a/tests/compilers/imports.rs +++ b/tests/compilers/imports.rs @@ -17,6 +17,8 @@ fn get_module(store: &Store) -> Result { (import "host" "1" (func (param i32) (result i32))) (import "host" "2" (func (param i32) (param i64))) (import "host" "3" (func (param i32 i64 i32 f32 f64))) + (memory $mem 1) + (export "memory" (memory $mem)) (func $foo call 0 @@ -87,16 +89,20 @@ fn dynamic_function_with_env() -> Result<()> { let module = get_module(&store)?; #[derive(WasmerEnv, Clone)] - struct Env(Arc); + struct Env { + counter: Arc, + }; impl std::ops::Deref for Env { type Target = Arc; fn deref(&self) -> &Self::Target { - &self.0 + &self.counter } } - let env: Env = Env(Arc::new(AtomicUsize::new(0))); + let env: Env = Env { + counter: Arc::new(AtomicUsize::new(0)), + }; Instance::new( &module, &imports! { @@ -292,3 +298,48 @@ fn static_function_that_fails() -> Result<()> { Ok(()) } + +fn get_module2(store: &Store) -> Result { + let wat = r#" + (import "host" "fn" (func)) + (memory $mem 1) + (export "memory" (memory $mem)) + (export "main" (func $main)) + (func $main (param) (result) + (call 0)) + "#; + + let module = Module::new(&store, &wat)?; + Ok(module) +} + +#[test] +fn dynamic_function_with_env_wasmer_env_init_works() -> Result<()> { + let store = get_store(false); + let module = get_module2(&store)?; + + #[allow(dead_code)] + #[derive(WasmerEnv, Clone)] + struct Env { + #[wasmer(export)] + memory: LazyInit, + }; + + let env: Env = Env { + memory: LazyInit::default(), + }; + let instance = Instance::new( + &module, + &imports! { + "host" => { + "fn" => Function::new_with_env(&store, &FunctionType::new(vec![], vec![]), env.clone(), |env, _values| { + assert!(env.memory_ref().is_some()); + Ok(vec![]) + }), + }, + }, + )?; + let f: NativeFunc<(), ()> = instance.exports.get_native_function("main")?; + f.call()?; + Ok(()) +}