Fix memory leak with imported funcrefs

lib/compiler-cranelift/src/func_environ.rs

@@ -825,7 +825,6 @@ impl<'module_environment> BaseFuncEnvironment for FuncEnvironment<'module_enviro
             self.get_table_set_func(&mut pos.func, table_index);
         let table_index = pos.ins().iconst(I32, table_index_arg as i64);
         let (vmctx, func_addr) = self.translate_load_builtin_function_address(&mut pos, func_idx);
-        let call_inst =
         pos.ins()
             .call_indirect(func_sig, func_addr, &[vmctx, table_index, index, value]);
         Ok(())

lib/engine-jit/src/artifact.rs

@@ -22,8 +22,8 @@ use wasmer_types::{
     TableIndex,
 };
 use wasmer_vm::{
-    FunctionBodyPtr, MemoryStyle, ModuleInfo, TableStyle, VMCallerCheckedAnyfunc, VMFuncRef,
+    FuncDataRegistry, FunctionBodyPtr, MemoryStyle, ModuleInfo, TableStyle, VMCallerCheckedAnyfunc,
-    VMFunctionEnvironment, VMSharedSignatureIndex, VMTrampoline,
+    VMFuncRef, VMFunctionEnvironment, VMSharedSignatureIndex, VMTrampoline,
 };
 
 /// A compiled wasm module, ready to be instantiated.
@@ -36,6 +36,7 @@ pub struct JITArtifact {
     // technically not "local", probably includes both local and imported
     // TODO: update name and docs
     local_func_data: BoxedSlice<LocalFunctionIndex, VMFuncRef>,
+    func_data_registry: Arc<FuncDataRegistry>,
     frame_info_registration: Mutex<Option<GlobalFrameInfoRegistration>>,
     finished_function_lengths: BoxedSlice<LocalFunctionIndex, usize>,
 }
@@ -227,6 +228,7 @@ impl JITArtifact {
         let finished_dynamic_function_trampolines =
             finished_dynamic_function_trampolines.into_boxed_slice();
         let signatures = signatures.into_boxed_slice();
+        let func_data_registry = inner_jit.func_data().clone();
         let local_func_data = finished_functions
             .iter()
             .map(|(k, v)| {
@@ -240,7 +242,7 @@ impl JITArtifact {
                         host_env: std::ptr::null_mut(),
                     },
                 };
-                inner_jit.func_data().register(metadata)
+                func_data_registry.register(metadata)
             })
             .collect::<PrimaryMap<LocalFunctionIndex, VMFuncRef>>();
         let local_func_data = local_func_data.into_boxed_slice();
@@ -254,6 +256,7 @@ impl JITArtifact {
             local_func_data,
             frame_info_registration: Mutex::new(None),
             finished_function_lengths,
+            func_data_registry,
         })
     }
 
@@ -337,6 +340,9 @@ impl Artifact for JITArtifact {
         &self.local_func_data
     }
 
+    fn func_data_registry(&self) -> &FuncDataRegistry {
+        &self.func_data_registry
+    }
     fn serialize(&self) -> Result<Vec<u8>, SerializeError> {
         // let mut s = flexbuffers::FlexbufferSerializer::new();
         // self.serializable.serialize(&mut s).map_err(|e| SerializeError::Generic(format!("{:?}", e)));

lib/engine-jit/src/engine.rs

@@ -34,7 +34,7 @@ impl JITEngine {
                 compiler: Some(compiler),
                 code_memory: vec![],
                 signatures: SignatureRegistry::new(),
-                func_data: FuncDataRegistry::new(),
+                func_data: Arc::new(FuncDataRegistry::new()),
                 features,
             })),
             target: Arc::new(target),
@@ -62,7 +62,7 @@ impl JITEngine {
                 compiler: None,
                 code_memory: vec![],
                 signatures: SignatureRegistry::new(),
-                func_data: FuncDataRegistry::new(),
+                func_data: Arc::new(FuncDataRegistry::new()),
                 features: Features::default(),
             })),
             target: Arc::new(Target::default()),
@@ -159,7 +159,7 @@ pub struct JITEngineInner {
     signatures: SignatureRegistry,
     /// TODO:
     /// func refs
-    func_data: FuncDataRegistry,
+    func_data: Arc<FuncDataRegistry>,
 }
 
 impl JITEngineInner {
@@ -309,7 +309,7 @@ impl JITEngineInner {
     }
 
     /// Shared func metadata registry.
-    pub(crate) fn func_data(&self) -> &FuncDataRegistry {
+    pub(crate) fn func_data(&self) -> &Arc<FuncDataRegistry> {
         &self.func_data
     }
 }

lib/engine/src/artifact.rs

@@ -12,8 +12,8 @@ use wasmer_types::{
     SignatureIndex, TableIndex,
 };
 use wasmer_vm::{
-    FunctionBodyPtr, InstanceAllocator, InstanceHandle, MemoryStyle, ModuleInfo, TableStyle,
+    FuncDataRegistry, FunctionBodyPtr, InstanceAllocator, InstanceHandle, MemoryStyle, ModuleInfo,
-    VMFuncRef, VMSharedSignatureIndex, VMTrampoline,
+    TableStyle, VMFuncRef, VMSharedSignatureIndex, VMTrampoline,
 };
 
 /// An `Artifact` is the product that the `Engine`
@@ -71,6 +71,11 @@ pub trait Artifact: Send + Sync + Upcastable {
         todo!("Implement this")
     }
 
+    /// Get the func data registry
+    fn func_data_registry(&self) -> &FuncDataRegistry {
+        todo!("Implement this")
+    }
+
     /// Serializes an artifact into bytes
     fn serialize(&self) -> Result<Vec<u8>, SerializeError>;
 
@@ -148,6 +153,7 @@ pub trait Artifact: Send + Sync + Upcastable {
             imports,
             self.signatures().clone(),
             self.func_metadata().clone(),
+            self.func_data_registry(),
             host_state,
             import_function_envs,
         )

lib/vm/src/func_data_registry.rs

@@ -30,7 +30,7 @@ unsafe impl Sync for FuncDataRegistry {}
 pub struct VMFuncRef(pub(crate) *const VMCallerCheckedAnyfunc);
 
 impl VMFuncRef {
-    /// TODO: we probabyl don't want this function, need to do something about this,
+    /// TODO: we probably don't want this function, need to do something about this,
     /// it definitely needs to be unsafe
     /// Hack to unblock myself:
     pub fn new(inner: *const VMCallerCheckedAnyfunc) -> Self {

lib/vm/src/instance/mod.rs

@@ -14,7 +14,7 @@ pub use allocator::InstanceAllocator;
 pub use r#ref::InstanceRef;
 
 use crate::export::VMExport;
-use crate::func_data_registry::VMFuncRef;
+use crate::func_data_registry::{FuncDataRegistry, VMFuncRef};
 use crate::global::Global;
 use crate::imports::Imports;
 use crate::memory::{Memory, MemoryError};
@@ -89,8 +89,8 @@ pub(crate) struct Instance {
     /// get removed. A missing entry is considered equivalent to an empty slice.
     passive_data: RefCell<HashMap<DataIndex, Arc<[u8]>>>,
 
-    /// mapping of local function indices to their func ref backing data.
+    /// mapping of function indices to their func ref backing data.
-    func_metadata: BoxedSlice<LocalFunctionIndex, VMFuncRef>,
+    func_metadata: BoxedSlice<FunctionIndex, VMFuncRef>,
 
     /// Hosts can store arbitrary per-instance information here.
     host_state: Box<dyn Any>,
@@ -648,53 +648,7 @@ impl Instance {
         if index == FunctionIndex::reserved_value() {
             return VMFuncRef::null();
         }
-
+        self.func_metadata[index]
-        if let Some(def_index) = self.module.local_func_index(index) {
-            let ptr = self.func_metadata[def_index];
-            unsafe {
-                // TODO: review use of `&mut` here: this is probably unsound, we probably
-                // need to wrap one or more of these types in UnsafeCell somewhere or something.
-                let item: &mut VMCallerCheckedAnyfunc = &mut *(*ptr as *mut _);
-                item.vmctx.vmctx = self.vmctx_ptr();
-            }
-            ptr
-        } else {
-            // TODO: actually implement this for imported functions
-            // hack to test that the idea is sound:
-            let sig = self.module.functions[index];
-            let type_index = self.signature_id(sig);
-            let import = self.imported_function(index);
-            let ptr = VMFuncRef(Box::into_raw(Box::new(VMCallerCheckedAnyfunc {
-                func_ptr: import.body,
-                type_index,
-                vmctx: import.environment,
-            })));
-            ptr
-            //todo!("Handle imported funcrefs!")
-        }
-        //todo!("Read this data from somewhere; need to store it first!")
-        /*
-        let sig = self.module.functions[index];
-        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 _,
-                VMFunctionEnvironment {
-                    vmctx: self.vmctx_ptr(),
-                },
-            )
-        } else {
-            let import = self.imported_function(index);
-            (import.body, import.environment)
-        };
-
-        todo!("store this data somewhere and then look it up and pass it back");
-        VMCallerCheckedAnyfunc {
-            func_ptr,
-            type_index,
-            vmctx,
-        }*/
     }
 
     /// The `table.init` operation: initializes a portion of a table with a
@@ -936,9 +890,41 @@ pub struct InstanceHandle {
 fn build_imported_function_func_refs(
     module_info: &ModuleInfo,
     imports: &Imports,
+    finished_functions: &BoxedSlice<LocalFunctionIndex, FunctionBodyPtr>,
+    func_data_registry: &FuncDataRegistry,
     vmshared_signatures: &BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
+    vmctx_ptr: *mut VMContext,
 ) -> BoxedSlice<FunctionIndex, VMFuncRef> {
-    todo!()
+    let mut func_refs = PrimaryMap::with_capacity(module_info.functions.len());
+
+    // do imported functions
+    for (index, import) in imports.functions.iter() {
+        let sig_index = module_info.functions[index];
+        let type_index = vmshared_signatures[sig_index];
+        let anyfunc = VMCallerCheckedAnyfunc {
+            func_ptr: import.body,
+            type_index,
+            vmctx: import.environment,
+        };
+        let func_ref = func_data_registry.register(anyfunc);
+        func_refs.push(func_ref);
+    }
+
+    // do local functions
+    for (local_index, func_ptr) in finished_functions.iter() {
+        let index = module_info.func_index(local_index);
+        let sig_index = module_info.functions[index];
+        let type_index = vmshared_signatures[sig_index];
+        let anyfunc = VMCallerCheckedAnyfunc {
+            func_ptr: func_ptr.0,
+            type_index,
+            vmctx: VMFunctionEnvironment { vmctx: vmctx_ptr },
+        };
+        let func_ref = func_data_registry.register(anyfunc);
+        func_refs.push(func_ref);
+    }
+
+    func_refs.into_boxed_slice()
 }
 
 impl InstanceHandle {
@@ -974,8 +960,9 @@ impl InstanceHandle {
         finished_globals: BoxedSlice<LocalGlobalIndex, Arc<Global>>,
         imports: Imports,
         vmshared_signatures: BoxedSlice<SignatureIndex, VMSharedSignatureIndex>,
-        // TODO: come up with a better name
+        // TODO: delete this and delete where it comes from
-        func_metadata: BoxedSlice<LocalFunctionIndex, VMFuncRef>,
+        _func_metadata: BoxedSlice<LocalFunctionIndex, VMFuncRef>,
+        func_data_registry: &FuncDataRegistry,
         host_state: Box<dyn Any>,
         imported_function_envs: BoxedSlice<FunctionIndex, ImportFunctionEnv>,
     ) -> Result<Self, Trap> {
@@ -988,6 +975,9 @@ impl InstanceHandle {
 
         let handle = {
             let offsets = allocator.offsets().clone();
+            // TODO: come up with a better name
+            // use dummy value to create an instance so we can get the vmctx pointer
+            let func_metadata = PrimaryMap::new().into_boxed_slice();
             // Create the `Instance`. The unique, the One.
             let instance = Instance {
                 module,
@@ -1006,7 +996,21 @@ impl InstanceHandle {
                 vmctx: VMContext {},
             };
 
-            let instance_ref = allocator.write_instance(instance);
+            let mut instance_ref = allocator.write_instance(instance);
+
+            // TODO: clean up this code
+            {
+                let instance = instance_ref.as_mut();
+                let vmctx_ptr = instance.vmctx_ptr();
+                instance.func_metadata = build_imported_function_func_refs(
+                    &*instance.module,
+                    &imports,
+                    &instance.functions,
+                    func_data_registry,
+                    &vmshared_signatures,
+                    vmctx_ptr,
+                );
+            }
 
             Self {
                 instance: instance_ref,