Change VMArrayCallFunction to be opaque

it's not always a native pointer with pulley
alexcrichton · Nov 19, 2024 · 5a7e902 · 5a7e902
1 parent b51cb66
commit 5a7e902
Show file tree

Hide file tree

Showing 7 changed files with 44 additions and 35 deletions.
diff --git a/crates/wasmtime/src/runtime/component/component.rs b/crates/wasmtime/src/runtime/component/component.rs
@@ -98,7 +98,7 @@ struct ComponentInner {
 
 pub(crate) struct AllCallFuncPointers {
     pub wasm_call: NonNull<VMWasmCallFunction>,
-    pub array_call: VMArrayCallFunction,
+    pub array_call: NonNull<VMArrayCallFunction>,
 }
 
 impl Component {
@@ -469,11 +469,7 @@ impl Component {
         } = &self.inner.info.trampolines[index];
         AllCallFuncPointers {
             wasm_call: self.func(wasm_call).cast(),
-            array_call: unsafe {
-                mem::transmute::<NonNull<VMFunctionBody>, VMArrayCallFunction>(
-                    self.func(array_call),
-                )
-            },
+            array_call: self.func(array_call).cast(),
         }
     }
 

diff --git a/crates/wasmtime/src/runtime/func.rs b/crates/wasmtime/src/runtime/func.rs
@@ -2287,12 +2287,8 @@ impl HostContext {
 
         let ctx = unsafe {
             VMArrayCallHostFuncContext::new(
-                VMFuncRef {
-                    array_call,
-                    wasm_call: None,
-                    type_index,
-                    vmctx: ptr::null_mut(),
-                },
+                array_call,
+                type_index,
                 Box::new(HostFuncState {
                     func,
                     ty: ty.into_registered_type(),

diff --git a/crates/wasmtime/src/runtime/trampoline/func.rs b/crates/wasmtime/src/runtime/trampoline/func.rs
@@ -81,12 +81,8 @@ where
 
     unsafe {
         Ok(VMArrayCallHostFuncContext::new(
-            VMFuncRef {
-                array_call,
-                wasm_call: None,
-                type_index: sig.index(),
-                vmctx: ptr::null_mut(),
-            },
+            array_call,
+            sig.index(),
             Box::new(TrampolineState { func, sig }),
         ))
     }

diff --git a/crates/wasmtime/src/runtime/vm.rs b/crates/wasmtime/src/runtime/vm.rs
@@ -265,16 +265,16 @@ impl ModuleRuntimeInfo {
     ///
     /// Returns `None` for Wasm functions which do not escape, and therefore are
     /// not callable from outside the Wasm module itself.
-    fn array_to_wasm_trampoline(&self, index: DefinedFuncIndex) -> Option<VMArrayCallFunction> {
+    fn array_to_wasm_trampoline(
+        &self,
+        index: DefinedFuncIndex,
+    ) -> Option<NonNull<VMArrayCallFunction>> {
         let m = match self {
             ModuleRuntimeInfo::Module(m) => m,
             ModuleRuntimeInfo::Bare(_) => unreachable!(),
         };
-        let ptr = m
-            .compiled_module()
-            .array_to_wasm_trampoline(index)?
-            .as_ptr();
-        Some(unsafe { mem::transmute::<*const u8, VMArrayCallFunction>(ptr) })
+        let ptr = NonNull::from(m.compiled_module().array_to_wasm_trampoline(index)?);
+        Some(ptr.cast())
     }
 
     /// Returns the `MemoryImage` structure used for copy-on-write

diff --git a/crates/wasmtime/src/runtime/vm/component.rs b/crates/wasmtime/src/runtime/vm/component.rs
@@ -391,7 +391,7 @@ impl ComponentInstance {
         &mut self,
         idx: TrampolineIndex,
         wasm_call: NonNull<VMWasmCallFunction>,
-        array_call: VMArrayCallFunction,
+        array_call: NonNull<VMArrayCallFunction>,
         type_index: VMSharedTypeIndex,
     ) {
         unsafe {
@@ -731,7 +731,7 @@ impl OwnedComponentInstance {
         &mut self,
         idx: TrampolineIndex,
         wasm_call: NonNull<VMWasmCallFunction>,
-        array_call: VMArrayCallFunction,
+        array_call: NonNull<VMArrayCallFunction>,
         type_index: VMSharedTypeIndex,
     ) {
         unsafe {

diff --git a/crates/wasmtime/src/runtime/vm/vmcontext.rs b/crates/wasmtime/src/runtime/vm/vmcontext.rs
@@ -37,9 +37,15 @@ use wasmtime_environ::{
 ///
 /// * The capacity of the `ValRaw` buffer. Must always be at least
 ///   `max(len(wasm_params), len(wasm_results))`.
-pub type VMArrayCallFunction =
+pub type VMArrayCallNative =
     unsafe extern "C" fn(*mut VMOpaqueContext, *mut VMOpaqueContext, *mut ValRaw, usize);
 
+/// An opaque function pointer which might be `VMArrayCallNative` or it might be
+/// pulley bytecode. Requires external knowledge to determine what kind of
+/// function pointer this is.
+#[repr(transparent)]
+pub struct VMArrayCallFunction(VMFunctionBody);
+
 /// A function pointer that exposes the Wasm calling convention.
 ///
 /// In practice, different Wasm function types end up mapping to different Rust
@@ -60,7 +66,7 @@ pub struct VMFunctionImport {
 
     /// Function pointer to use when calling this imported function with the
     /// "array" calling convention that `Func::new` et al use.
-    pub array_call: VMArrayCallFunction,
+    pub array_call: NonNull<VMArrayCallFunction>,
 
     /// The VM state associated with this function.
     ///
@@ -652,7 +658,7 @@ mod test_vmshared_type_index {
 pub struct VMFuncRef {
     /// Function pointer for this funcref if being called via the "array"
     /// calling convention that `Func::new` et al use.
-    pub array_call: VMArrayCallFunction,
+    pub array_call: NonNull<VMArrayCallFunction>,
 
     /// Function pointer for this funcref if being called via the calling
     /// convention we use when compiling Wasm.
@@ -709,7 +715,15 @@ impl VMFuncRef {
     /// This method is unsafe because it can be called with any pointers. They
     /// must all be valid for this wasm function call to proceed.
     pub unsafe fn array_call(&self, caller: *mut VMOpaqueContext, args_and_results: *mut [ValRaw]) {
-        (self.array_call)(
+        union GetNativePointer {
+            native: VMArrayCallNative,
+            ptr: NonNull<VMArrayCallFunction>,
+        }
+        let native = GetNativePointer {
+            ptr: self.array_call,
+        }
+        .native;
+        native(
             self.vmctx,
             caller,
             args_and_results.cast(),

diff --git a/crates/wasmtime/src/runtime/vm/vmcontext/vm_host_func_context.rs b/crates/wasmtime/src/runtime/vm/vmcontext/vm_host_func_context.rs
@@ -2,11 +2,13 @@
 //!
 //! Keep in sync with `wasmtime_environ::VMHostFuncOffsets`.
 
-use super::VMOpaqueContext;
+use super::{VMArrayCallNative, VMOpaqueContext};
 use crate::prelude::*;
 use crate::runtime::vm::{StoreBox, VMFuncRef};
+use crate::ValRaw;
 use core::any::Any;
-use wasmtime_environ::VM_ARRAY_CALL_HOST_FUNC_MAGIC;
+use core::ptr::{self, NonNull};
+use wasmtime_environ::{VMSharedTypeIndex, VM_ARRAY_CALL_HOST_FUNC_MAGIC};
 
 /// The `VM*Context` for array-call host functions.
 ///
@@ -30,13 +32,18 @@ impl VMArrayCallHostFuncContext {
     /// The `host_func` must be a pointer to a host (not Wasm) function and it
     /// must be `Send` and `Sync`.
     pub unsafe fn new(
-        func_ref: VMFuncRef,
+        host_func: VMArrayCallNative,
+        type_index: VMSharedTypeIndex,
         host_state: Box<dyn Any + Send + Sync>,
     ) -> StoreBox<VMArrayCallHostFuncContext> {
-        debug_assert!(func_ref.vmctx.is_null());
         let ctx = StoreBox::new(VMArrayCallHostFuncContext {
             magic: wasmtime_environ::VM_ARRAY_CALL_HOST_FUNC_MAGIC,
-            func_ref,
+            func_ref: VMFuncRef {
+                array_call: NonNull::new(host_func as *mut u8).unwrap().cast(),
+                type_index,
+                wasm_call: None,
+                vmctx: ptr::null_mut(),
+            },
             host_state,
         });
         let vmctx = VMOpaqueContext::from_vm_array_call_host_func_context(ctx.get());