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 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()
                     };