Merge pull request #1857 from wasmerio/fix/dynamic-fn-env-wasmer-env

Fix dynamic function envs not working with `WasmerEnv`

lib/api/src/externals/function.rs

@@ -728,14 +728,16 @@ impl VMDynamicFunction for VMDynamicFunctionWithoutEnv {
     }
 }
 
+#[repr(C)]
 pub(crate) struct VMDynamicFunctionWithEnv<Env>
 where
     Env: Sized + 'static,
 {
+    // This field _must_ come first in this struct.
+    env: Box<Env>,
     function_type: FunctionType,
     #[allow(clippy::type_complexity)]
     func: Box<dyn Fn(&Env, &[Val]) -> Result<Vec<Val>, RuntimeError> + 'static>,
-    env: Box<Env>,
 }
 
 impl<Env> VMDynamicFunction for VMDynamicFunctionWithEnv<Env>

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<VMDynamicFunctionContext<VMDynamicFunctionWithEnv<Env>>>
+                            // Which we can interpret as `*const <field offset> *const Env` (due to
+                            // the precise layout of these types via `repr(C)`)
+                            // We extract the `*const Env`:
+                            unsafe {
+                                // Box<VMDynamicFunctionContext<...>>
+                                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,15 +198,17 @@ 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 {
                     body: address,
                     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 {

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<BoxedSlice<FunctionIndex, Option<ImportInitializerFuncPtr>>>,
+    pub host_function_env_initializers: Option<
+        BoxedSlice<FunctionIndex, (Option<ImportInitializerFuncPtr>, *mut std::ffi::c_void)>,
+    >,
 
     /// Resolved addresses for imported tables.
     pub tables: BoxedSlice<TableIndex, VMTableImport>,
@@ -34,7 +35,10 @@ impl Imports {
     /// Construct a new `Imports` instance.
     pub fn new(
         function_imports: PrimaryMap<FunctionIndex, VMFunctionImport>,
-        host_function_env_initializers: PrimaryMap<FunctionIndex, Option<ImportInitializerFuncPtr>>,
+        host_function_env_initializers: PrimaryMap<
+            FunctionIndex,
+            (Option<ImportInitializerFuncPtr>, *mut std::ffi::c_void),
+        >,
         table_imports: PrimaryMap<TableIndex, VMTableImport>,
         memory_imports: PrimaryMap<MemoryIndex, VMMemoryImport>,
         global_imports: PrimaryMap<GlobalIndex, VMGlobalImport>,
@@ -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()
                     };

tests/compilers/imports.rs

@@ -17,6 +17,8 @@ fn get_module(store: &Store) -> Result<Module> {
         (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<AtomicUsize>);
+    struct Env {
+        counter: Arc<AtomicUsize>,
+    };
 
     impl std::ops::Deref for Env {
         type Target = Arc<AtomicUsize>;
         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<Module> {
+    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<Memory>,
+    };
+
+    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(())
+}