Add hacked together vertical slice of host env updating idea

lib/api/src/externals/function.rs

@@ -11,8 +11,8 @@ use std::cmp::max;
 use std::fmt;
 use wasmer_vm::{
     raise_user_trap, resume_panic, wasmer_call_trampoline, Export, ExportFunction,
-    VMCallerCheckedAnyfunc, VMContext, VMDynamicFunctionContext, VMFunctionBody, VMFunctionKind,
+    FunctionExtraData, VMCallerCheckedAnyfunc, VMContext, VMDynamicFunctionContext, VMFunctionBody,
-    VMTrampoline,
+    VMFunctionKind, VMTrampoline,
 };
 
 /// A function defined in the Wasm module
@@ -85,7 +85,9 @@ impl Function {
         // The engine linker will replace the address with one pointing to a
         // generated dynamic trampoline.
         let address = std::ptr::null() as *const VMFunctionBody;
-        let vmctx = Box::into_raw(Box::new(dynamic_ctx)) as *mut VMContext;
+        let vmctx = FunctionExtraData {
+            host_env: Box::into_raw(Box::new(dynamic_ctx)) as *mut _,
+        };
 
         Self {
             store: store.clone(),
@@ -94,6 +96,7 @@ impl Function {
                 address,
                 kind: VMFunctionKind::Dynamic,
                 vmctx,
+                function_ptr: 0,
                 signature: ty.clone(),
             },
         }
@@ -134,7 +137,9 @@ impl Function {
         // The engine linker will replace the address with one pointing to a
         // generated dynamic trampoline.
         let address = std::ptr::null() as *const VMFunctionBody;
-        let vmctx = Box::into_raw(Box::new(dynamic_ctx)) as *mut VMContext;
+        let vmctx = FunctionExtraData {
+            host_env: Box::into_raw(Box::new(dynamic_ctx)) as *mut _,
+        };
 
         Self {
             store: store.clone(),
@@ -143,6 +148,8 @@ impl Function {
                 address,
                 kind: VMFunctionKind::Dynamic,
                 vmctx,
+                // TODO:
+                function_ptr: 0,
                 signature: ty.clone(),
             },
         }
@@ -174,7 +181,9 @@ impl Function {
     {
         let function = inner::Function::<Args, Rets>::new(func);
         let address = function.address() as *const VMFunctionBody;
-        let vmctx = std::ptr::null_mut() as *mut _ as *mut VMContext;
+        let vmctx = FunctionExtraData {
+            host_env: std::ptr::null_mut() as *mut _,
+        };
         let signature = function.ty();
 
         Self {
@@ -184,6 +193,8 @@ impl Function {
                 address,
                 vmctx,
                 signature,
+                // TODO:
+                function_ptr: 0,
                 kind: VMFunctionKind::Static,
             },
         }
@@ -216,7 +227,7 @@ impl Function {
         F: HostFunction<Args, Rets, WithEnv, Env>,
         Args: WasmTypeList,
         Rets: WasmTypeList,
-        Env: Sized + 'static,
+        Env: Sized + crate::WasmerPostInstantiate + 'static,
     {
         let function = inner::Function::<Args, Rets>::new(func);
         let address = function.address();
@@ -227,7 +238,11 @@ impl Function {
         // In the case of Host-defined functions `VMContext` is whatever environment
         // the user want to attach to the function.
         let box_env = Box::new(env);
-        let vmctx = Box::into_raw(box_env) as *mut _ as *mut VMContext;
+        let vmctx = FunctionExtraData {
+            host_env: Box::into_raw(box_env) as *mut _,
+        };
+        let function_ptr = Env::finish as usize;
+        dbg!(function_ptr);
         let signature = function.ty();
 
         Self {
@@ -237,6 +252,7 @@ impl Function {
                 address,
                 kind: VMFunctionKind::Static,
                 vmctx,
+                function_ptr,
                 signature,
             },
         }

lib/api/src/instance.rs

@@ -73,7 +73,7 @@ impl Instance {
     pub fn new(module: &Module, resolver: &dyn Resolver) -> Result<Self, InstantiationError> {
         let store = module.store();
 
-        let handle = module.instantiate(resolver)?;
+        let (handle, thunks) = module.instantiate(resolver)?;
 
         let exports = module
             .exports()
@@ -85,11 +85,20 @@ impl Instance {
             })
             .collect::<Exports>();
 
-        Ok(Self {
+        let instance = Self {
             handle,
             module: module.clone(),
             exports,
-        })
+        };
+
+        for (func, env) in thunks.iter() {
+            dbg!(func, env);
+            unsafe {
+                func(*env, (&instance) as *const _ as *const _);
+            }
+        }
+
+        Ok(instance)
     }
 
     /// Gets the [`Module`] associated with this instance.

lib/api/src/lib.rs

@@ -118,3 +118,17 @@ pub use wasmer_engine_native::{Native, NativeArtifact, NativeEngine};
 
 /// Version number of this crate.
 pub const VERSION: &str = env!("CARGO_PKG_VERSION");
+
+// --------------------------------------------------------------
+// TODO: put this in a proper location, just prototyping for now:
+// TODO: rename everything, all names are throw-away names
+
+/// Prototype trait for finishing envs.
+pub trait WasmerPostInstantiate {
+    /// The function that Wasmer will call on your type to let it finish
+    /// instantiating.
+    fn finish(&mut self, instance: &Instance);
+
+    /// Frees memory written to `self` so it can be dropped without any memory leaks.
+    fn free(&mut self);
+}

lib/api/src/module.rs

@@ -261,9 +261,18 @@ impl Module {
     pub(crate) fn instantiate(
         &self,
         resolver: &dyn Resolver,
-    ) -> Result<InstanceHandle, InstantiationError> {
+    ) -> Result<
+        (
+            InstanceHandle,
+            Vec<(
+                fn(*mut std::ffi::c_void, *const std::ffi::c_void),
+                *mut std::ffi::c_void,
+            )>,
+        ),
+        InstantiationError,
+    > {
         unsafe {
-            let instance_handle =
+            let (instance_handle, thunks) =
                 self.artifact
                     .instantiate(self.store.tunables(), resolver, Box::new(()))?;
 
@@ -274,7 +283,7 @@ impl Module {
             // instance tables.
             self.artifact.finish_instantiation(&instance_handle)?;
 
-            Ok(instance_handle)
+            Ok((instance_handle, thunks))
         }
     }
 

lib/api/src/native.rs

@@ -17,7 +17,8 @@ use crate::{FromToNativeWasmType, Function, FunctionType, RuntimeError, Store, W
 use std::panic::{catch_unwind, AssertUnwindSafe};
 use wasmer_types::NativeWasmType;
 use wasmer_vm::{
-    ExportFunction, VMContext, VMDynamicFunctionContext, VMFunctionBody, VMFunctionKind,
+    ExportFunction, FunctionExtraData, VMContext, VMDynamicFunctionContext, VMFunctionBody,
+    VMFunctionKind,
 };
 
 /// A WebAssembly function that can be called natively
@@ -26,7 +27,7 @@ pub struct NativeFunc<'a, Args = (), Rets = ()> {
     definition: FunctionDefinition,
     store: Store,
     address: *const VMFunctionBody,
-    vmctx: *mut VMContext,
+    vmctx: FunctionExtraData,
     arg_kind: VMFunctionKind,
     // exported: ExportFunction,
     _phantom: PhantomData<(&'a (), Args, Rets)>,
@@ -42,7 +43,7 @@ where
     pub(crate) fn new(
         store: Store,
         address: *const VMFunctionBody,
-        vmctx: *mut VMContext,
+        vmctx: FunctionExtraData,
         arg_kind: VMFunctionKind,
         definition: FunctionDefinition,
     ) -> Self {
@@ -68,6 +69,8 @@ where
             address: other.address,
             vmctx: other.vmctx,
             signature,
+            // TODO:
+            function_ptr: 0,
             kind: other.arg_kind,
         }
     }
@@ -87,6 +90,8 @@ where
                 address: other.address,
                 vmctx: other.vmctx,
                 signature,
+                // TODO:
+                function_ptr: 0,
                 kind: other.arg_kind,
             },
         }
@@ -163,7 +168,7 @@ macro_rules! impl_native_traits {
                         match self.arg_kind {
                             VMFunctionKind::Static => {
                                 let results = catch_unwind(AssertUnwindSafe(|| unsafe {
-                                    let f = std::mem::transmute::<_, unsafe extern "C" fn( *mut VMContext, $( $x, )*) -> Rets::CStruct>(self.address);
+                                    let f = std::mem::transmute::<_, unsafe extern "C" fn( FunctionExtraData, $( $x, )*) -> Rets::CStruct>(self.address);
                                     // We always pass the vmctx
                                     f( self.vmctx, $( $x, )* )
                                 })).map_err(|e| RuntimeError::new(format!("{:?}", e)))?;
@@ -173,12 +178,16 @@ macro_rules! impl_native_traits {
                                 let params_list = [ $( $x.to_native().to_value() ),* ];
                                 let results = if !has_env {
                                     type VMContextWithoutEnv = VMDynamicFunctionContext<VMDynamicFunctionWithoutEnv>;
-                                    let ctx = self.vmctx as *mut VMContextWithoutEnv;
+                                    unsafe {
-                                    unsafe { (*ctx).ctx.call(&params_list)? }
+                                        let ctx = self.vmctx.host_env as *mut VMContextWithoutEnv;
+                                        (*ctx).ctx.call(&params_list)?
+                                    }
                                 } else {
                                     type VMContextWithEnv = VMDynamicFunctionContext<VMDynamicFunctionWithEnv<std::ffi::c_void>>;
-                                    let ctx = self.vmctx as *mut VMContextWithEnv;
+                                    unsafe {
-                                    unsafe { (*ctx).ctx.call(&params_list)? }
+                                        let ctx = self.vmctx.host_env as *mut VMContextWithEnv;
+                                        (*ctx).ctx.call(&params_list)?
+                                    }
                                 };
                                 let mut rets_list_array = Rets::empty_array();
                                 let mut_rets = rets_list_array.as_mut() as *mut [i128] as *mut i128;

lib/api/src/types.rs

@@ -56,7 +56,9 @@ impl ValFuncRef for Val {
             Self::ExternRef(ExternRef::Null) => wasmer_vm::VMCallerCheckedAnyfunc {
                 func_ptr: ptr::null(),
                 type_index: wasmer_vm::VMSharedSignatureIndex::default(),
-                vmctx: ptr::null_mut(),
+                vmctx: wasmer_vm::FunctionExtraData {
+                    host_env: ptr::null_mut(),
+                },
             },
             Self::FuncRef(f) => f.checked_anyfunc(),
             _ => return Err(RuntimeError::new("val is not funcref")),
@@ -74,6 +76,8 @@ impl ValFuncRef for Val {
         let export = wasmer_vm::ExportFunction {
             address: item.func_ptr,
             signature,
+            // TODO:
+            function_ptr: 0,
             // All functions in tables are already Static (as dynamic functions
             // are converted to use the trampolines with static signatures).
             kind: wasmer_vm::VMFunctionKind::Static,

lib/cli/Cargo.toml

@@ -58,7 +58,7 @@ default = [
     "object-file",
     "cache",
     "wasi",
-    "emscripten",
+    # "emscripten",
 ]
 engine = []
 jit = [

lib/cli/src/commands/run/wasi.rs

@@ -58,7 +58,7 @@ impl Wasi {
         let import_object = wasi_env.import_object(&module)?;
         let instance = Instance::new(&module, &import_object)?;
 
-        wasi_env.set_memory(instance.exports.get_memory("memory")?.clone());
+        //wasi_env.set_memory(instance.exports.get_memory("memory")?.clone());
 
         let start = instance.exports.get_function("_start")?;
         let result = start.call(&[]);

lib/engine/src/artifact.rs

@@ -86,7 +86,16 @@ pub trait Artifact: Send + Sync + Upcastable {
         tunables: &dyn Tunables,
         resolver: &dyn Resolver,
         host_state: Box<dyn Any>,
-    ) -> Result<InstanceHandle, InstantiationError> {
+    ) -> Result<
+        (
+            InstanceHandle,
+            Vec<(
+                fn(*mut std::ffi::c_void, *const std::ffi::c_void),
+                *mut std::ffi::c_void,
+            )>,
+        ),
+        InstantiationError,
+    > {
         self.preinstantiate()?;
 
         let module = self.module();
@@ -98,6 +107,19 @@ pub trait Artifact: Send + Sync + Upcastable {
             self.table_styles(),
         )
         .map_err(InstantiationError::Link)?;
+
+        let mut thunks = vec![];
+        // ------------
+        for func in imports.functions.values() {
+            unsafe {
+                let finish: fn(*mut std::ffi::c_void, instance: *const std::ffi::c_void) =
+                    std::mem::transmute(func.function_ptr);
+                dbg!(func.extra_data.host_env);
+                thunks.push((finish, func.extra_data.host_env));
+            }
+        }
+        // ------------
+
         let finished_memories = tunables
             .create_memories(&module, self.memory_styles())
             .map_err(InstantiationError::Link)?
@@ -113,7 +135,7 @@ pub trait Artifact: Send + Sync + Upcastable {
 
         self.register_frame_info();
 
-        InstanceHandle::new(
+        let handle = InstanceHandle::new(
             module,
             self.finished_functions().clone(),
             finished_memories,
@@ -123,7 +145,8 @@ pub trait Artifact: Send + Sync + Upcastable {
             self.signatures().clone(),
             host_state,
         )
-        .map_err(|trap| InstantiationError::Start(RuntimeError::from_trap(trap)))
+        .map_err(|trap| InstantiationError::Start(RuntimeError::from_trap(trap)))?;
+        Ok((handle, thunks))
     }
 
     /// Finishes the instantiation of a just created `InstanceHandle`.

lib/engine/src/resolver.rs

@@ -167,7 +167,9 @@ pub fn resolve_imports(
                 };
                 function_imports.push(VMFunctionImport {
                     body: address,
-                    vmctx: f.vmctx,
+                    extra_data: f.vmctx,
+                    // TODO:
+                    function_ptr: f.function_ptr,
                 });
             }
             Export::Table(ref t) => {

lib/vm/src/export.rs

@@ -30,7 +30,9 @@ pub struct ExportFunction {
     /// The address of the native-code function.
     pub address: *const VMFunctionBody,
     /// Pointer to the containing `VMContext`.
-    pub vmctx: *mut VMContext,
+    pub vmctx: crate::vmcontext::FunctionExtraData,
+    /// temp code to set vmctx for host functions
+    pub function_ptr: usize,
     /// The function type, used for compatibility checking.
     pub signature: FunctionType,
     /// The function kind (it defines how it's the signature that provided `address` have)

lib/vm/src/instance.rs

@@ -291,13 +291,19 @@ impl Instance {
         match export {
             ExportIndex::Function(index) => {
                 let sig_index = &self.module.functions[*index];
-                let (address, vmctx) = if let Some(def_index) = self.module.local_func_index(*index)
+                let (address, vmctx, function_ptr) =
-                {
+                    if let Some(def_index) = self.module.local_func_index(*index) {
-                    (self.functions[def_index].0 as *const _, self.vmctx_ptr())
+                        (
-                } else {
+                            self.functions[def_index].0 as *const _,
-                    let import = self.imported_function(*index);
+                            crate::vmcontext::FunctionExtraData {
-                    (import.body, import.vmctx)
+                                vmctx: self.vmctx_ptr(),
-                };
+                            },
+                            0,
+                        )
+                    } else {
+                        let import = self.imported_function(*index);
+                        (import.body, import.extra_data, import.function_ptr)
+                    };
                 let signature = self.module.signatures[*sig_index].clone();
                 ExportFunction {
                     address,
@@ -308,6 +314,7 @@ impl Instance {
                     kind: VMFunctionKind::Static,
                     signature,
                     vmctx,
+                    function_ptr,
                 }
                 .into()
             }
@@ -365,28 +372,34 @@ impl Instance {
             None => return Ok(()),
         };
 
-        let (callee_address, callee_vmctx) = match self.module.local_func_index(start_index) {
+        let (callee_address, callee_extra_data) = match self.module.local_func_index(start_index) {
             Some(local_index) => {
                 let body = self
                     .functions
                     .get(local_index)
                     .expect("function index is out of bounds")
                     .0;
-                (body as *const _, self.vmctx_ptr())
+                (
+                    body as *const _,
+                    crate::FunctionExtraData {
+                        vmctx: self.vmctx_ptr(),
+                    },
+                )
             }
             None => {
                 assert_lt!(start_index.index(), self.module.num_imported_functions);
                 let import = self.imported_function(start_index);
-                (import.body, import.vmctx)
+                (import.body, import.extra_data)
             }
         };
 
         // Make the call.
         unsafe {
-            catch_traps(callee_vmctx, || {
+            catch_traps(callee_extra_data, || {
-                mem::transmute::<*const VMFunctionBody, unsafe extern "C" fn(*mut VMContext)>(
+                mem::transmute::<
-                    callee_address,
+                    *const VMFunctionBody,
-                )(callee_vmctx)
+                    unsafe extern "C" fn(crate::FunctionExtraData),
+                >(callee_address)(callee_extra_data)
             })
         }
     }
@@ -556,10 +569,15 @@ impl Instance {
         let type_index = self.signature_id(sig);
 
         let (func_ptr, vmctx) = if let Some(def_index) = self.module.local_func_index(index) {
-            (self.functions[def_index].0 as *const _, self.vmctx_ptr())
+            (
+                self.functions[def_index].0 as *const _,
+                crate::FunctionExtraData {
+                    vmctx: self.vmctx_ptr(),
+                },
+            )
         } else {
             let import = self.imported_function(index);
-            (import.body, import.vmctx)
+            (import.body, import.extra_data)
         };
         VMCallerCheckedAnyfunc {
             func_ptr,

lib/vm/src/lib.rs

@@ -48,10 +48,10 @@ pub use crate::sig_registry::SignatureRegistry;
 pub use crate::table::{LinearTable, Table, TableStyle};
 pub use crate::trap::*;
 pub use crate::vmcontext::{
-    VMBuiltinFunctionIndex, VMCallerCheckedAnyfunc, VMContext, VMDynamicFunctionContext,
+    FunctionExtraData, VMBuiltinFunctionIndex, VMCallerCheckedAnyfunc, VMContext,
-    VMFunctionBody, VMFunctionImport, VMFunctionKind, VMGlobalDefinition, VMGlobalImport,
+    VMDynamicFunctionContext, VMFunctionBody, VMFunctionImport, VMFunctionKind, VMGlobalDefinition,
-    VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex, VMTableDefinition, VMTableImport,
+    VMGlobalImport, VMMemoryDefinition, VMMemoryImport, VMSharedSignatureIndex, VMTableDefinition,
-    VMTrampoline,
+    VMTableImport, VMTrampoline,
 };
 pub use crate::vmoffsets::{TargetSharedSignatureIndex, VMOffsets};
 

lib/vm/src/trap/traphandlers.rs

@@ -429,13 +429,13 @@ impl Trap {
 /// Wildly unsafe because it calls raw function pointers and reads/writes raw
 /// function pointers.
 pub unsafe fn wasmer_call_trampoline(
-    vmctx: *mut VMContext,
+    vmctx: crate::FunctionExtraData,
     trampoline: VMTrampoline,
     callee: *const VMFunctionBody,
     values_vec: *mut u8,
 ) -> Result<(), Trap> {
     catch_traps(vmctx, || {
-        mem::transmute::<_, extern "C" fn(*mut VMContext, *const VMFunctionBody, *mut u8)>(
+        mem::transmute::<_, extern "C" fn(crate::FunctionExtraData, *const VMFunctionBody, *mut u8)>(
             trampoline,
         )(vmctx, callee, values_vec)
     })
@@ -447,7 +447,7 @@ pub unsafe fn wasmer_call_trampoline(
 /// # Safety
 ///
 /// Highly unsafe since `closure` won't have any destructors run.
-pub unsafe fn catch_traps<F>(vmctx: *mut VMContext, mut closure: F) -> Result<(), Trap>
+pub unsafe fn catch_traps<F>(vmctx: crate::FunctionExtraData, mut closure: F) -> Result<(), Trap>
 where
     F: FnMut(),
 {
@@ -481,7 +481,7 @@ where
 ///
 /// Check [`catch_traps`].
 pub unsafe fn catch_traps_with_result<F, R>(
-    vmctx: *mut VMContext,
+    vmctx: crate::FunctionExtraData,
     mut closure: F,
 ) -> Result<R, Trap>
 where
@@ -501,7 +501,7 @@ pub struct CallThreadState {
     jmp_buf: Cell<*const u8>,
     reset_guard_page: Cell<bool>,
     prev: Option<*const CallThreadState>,
-    vmctx: *mut VMContext,
+    vmctx: crate::FunctionExtraData,
     handling_trap: Cell<bool>,
 }
 
@@ -518,7 +518,7 @@ enum UnwindReason {
 }
 
 impl CallThreadState {
-    fn new(vmctx: *mut VMContext) -> Self {
+    fn new(vmctx: crate::FunctionExtraData) -> Self {
         Self {
             unwind: Cell::new(UnwindReason::None),
             vmctx,
@@ -561,7 +561,7 @@ impl CallThreadState {
 
     fn any_instance(&self, func: impl Fn(&InstanceHandle) -> bool) -> bool {
         unsafe {
-            if func(&InstanceHandle::from_vmctx(self.vmctx)) {
+            if func(&InstanceHandle::from_vmctx(self.vmctx.vmctx)) {
                 return true;
             }
             match self.prev {

lib/vm/src/vmcontext.rs

@@ -15,6 +15,29 @@ use std::ptr::{self, NonNull};
 use std::sync::Arc;
 use std::u32;
 
+/// We stop lying about what this daat is
+/// TODO:
+#[derive(Copy, Clone)]
+pub union FunctionExtraData {
+    /// Wasm function, it has a real VMContext:
+    pub vmctx: *mut VMContext,
+    /// Host functions can have custom environments
+    pub host_env: *mut std::ffi::c_void,
+}
+
+impl std::fmt::Debug for FunctionExtraData {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(f, "FunctionExtarData union")
+    }
+}
+
+impl std::cmp::PartialEq for FunctionExtraData {
+    fn eq(&self, rhs: &Self) -> bool {
+        // TODO
+        false
+    }
+}
+
 /// An imported function.
 #[derive(Debug, Copy, Clone)]
 #[repr(C)]
@@ -22,8 +45,12 @@ pub struct VMFunctionImport {
     /// A pointer to the imported function body.
     pub body: *const VMFunctionBody,
 
-    /// A pointer to the `VMContext` that owns the function.
+    /// A pointer to the `VMContext` that owns the function or host data.
-    pub vmctx: *mut VMContext,
+    pub extra_data: FunctionExtraData,
+
+    /// temporary hack
+    /// Only used by host env
+    pub function_ptr: usize,
 }
 
 #[cfg(test)]
@@ -46,7 +73,7 @@ mod test_vmfunction_import {
             usize::from(offsets.vmfunction_import_body())
         );
         assert_eq!(
-            offset_of!(VMFunctionImport, vmctx),
+            offset_of!(VMFunctionImport, extra_data),
             usize::from(offsets.vmfunction_import_vmctx())
         );
     }
@@ -729,7 +756,7 @@ pub struct VMCallerCheckedAnyfunc {
     /// Function signature id.
     pub type_index: VMSharedSignatureIndex,
     /// Function `VMContext`.
-    pub vmctx: *mut VMContext,
+    pub vmctx: FunctionExtraData,
     // If more elements are added here, remember to add offset_of tests below!
 }
 
@@ -768,7 +795,9 @@ impl Default for VMCallerCheckedAnyfunc {
         Self {
             func_ptr: ptr::null_mut(),
             type_index: Default::default(),
-            vmctx: ptr::null_mut(),
+            vmctx: FunctionExtraData {
+                vmctx: ptr::null_mut(),
+            },
         }
     }
 }

lib/vm/src/vmoffsets.rs

@@ -109,7 +109,7 @@ impl VMOffsets {
     ///
     /// [`VMFunctionImport`]: crate::vmcontext::VMFunctionImport
     pub fn size_of_vmfunction_import(&self) -> u8 {
-        2 * self.pointer_size
+        3 * self.pointer_size
     }
 }
 

lib/wasi/src/lib.rs

@@ -51,7 +51,25 @@ pub enum WasiError {
 #[derive(Debug, Clone)]
 pub struct WasiEnv {
     state: Arc<Mutex<WasiState>>,
-    memory: Arc<WasiMemory>,
+    memory: *mut Memory,
+}
+
+impl wasmer::WasmerPostInstantiate for WasiEnv {
+    fn finish(&mut self, instance: &wasmer::Instance) {
+        dbg!("in Wasi::Finish");
+        let memory = instance.exports.get_memory("memory").unwrap();
+        unsafe {
+            let heap_ptr = Box::into_raw(Box::new(memory.clone()));
+            self.memory = heap_ptr;
+        }
+    }
+
+    fn free(&mut self) {
+        unsafe {
+            Box::from_raw(self.memory);
+            self.memory = std::ptr::null_mut();
+        }
+    }
 }
 
 /// Wrapper type around `Memory` used to delay initialization of the memory.
@@ -140,7 +158,7 @@ impl WasiEnv {
     pub fn new(state: WasiState) -> Self {
         Self {
             state: Arc::new(Mutex::new(state)),
-            memory: Arc::new(WasiMemory::new()),
+            memory: std::ptr::null_mut(),
         }
     }
 
@@ -154,9 +172,9 @@ impl WasiEnv {
     }
 
     /// Set the memory
-    pub fn set_memory(&mut self, memory: Memory) -> bool {
+    /*pub fn set_memory(&mut self, memory: Memory) -> bool {
         self.memory.set_memory(memory)
-    }
+    }*/
 
     /// Get the WASI state
     pub fn state(&self) -> MutexGuard<WasiState> {
@@ -170,7 +188,8 @@ impl WasiEnv {
 
     /// Get a reference to the memory
     pub fn memory(&self) -> &Memory {
-        self.memory.get_memory().expect("The expected Memory is not attached to the `WasiEnv`. Did you forgot to call wasi_env.set_memory(...)?")
+        dbg!("Getting memory", &self.memory);
+        unsafe { &*self.memory }
     }
 
     pub(crate) fn get_memory_and_wasi_state(

tests/lib/wast/src/wasi_wast.rs

@@ -84,7 +84,7 @@ impl<'a> WasiTest<'a> {
         let imports = self.get_imports(store, &module, env.clone())?;
         let instance = Instance::new(&module, &imports)?;
         let memory: &Memory = instance.exports.get("memory")?;
-        env.set_memory(memory.clone());
+        //env.set_memory(memory.clone());
 
         let start = instance.exports.get_function("_start")?;
         // TODO: handle errors here when the error fix gets shipped