Change VMArrayCallFunction to be an opaque pointer

This pointer is technically not safe to be a function pointer because it will get unloaded when the module is dropped. Additionally with Pulley this isn't actually a native function pointer. This builds on bytecodealliance#9630 to prepare for future integration with Pulley to ensure that the number of locations that have to deal with bytecode-vs-native-code are minimized.
alexcrichton · Nov 20, 2024 · 5896de1 · 5896de1
1 parent 27ce0ba
commit 5896de1
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 42 deletions.
diff --git a/crates/wasmtime/src/runtime/component/component.rs b/crates/wasmtime/src/runtime/component/component.rs
@@ -15,7 +15,6 @@ use crate::{
 use crate::{FuncType, ValType};
 use alloc::sync::Arc;
 use core::any::Any;
-use core::mem;
 use core::ops::Range;
 use core::ptr::NonNull;
 #[cfg(feature = "std")]
@@ -98,7 +97,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 +468,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
@@ -16,7 +16,7 @@ use core::future::Future;
 use core::mem::{self, MaybeUninit};
 use core::num::NonZeroUsize;
 use core::pin::Pin;
-use core::ptr::{self, NonNull};
+use core::ptr::NonNull;
 use wasmtime_environ::VMSharedTypeIndex;
 
 /// A reference to the abstract `nofunc` heap value.
@@ -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
@@ -1,12 +1,9 @@
 //! Support for a calling of an imported function.
 
 use crate::prelude::*;
-use crate::runtime::vm::{
-    StoreBox, VMArrayCallHostFuncContext, VMContext, VMFuncRef, VMOpaqueContext,
-};
+use crate::runtime::vm::{StoreBox, VMArrayCallHostFuncContext, VMContext, VMOpaqueContext};
 use crate::type_registry::RegisteredType;
 use crate::{FuncType, ValRaw};
-use core::ptr;
 
 struct TrampolineState<F> {
     func: F,
@@ -81,12 +78,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
@@ -9,7 +9,6 @@ use crate::prelude::*;
 use crate::store::StoreOpaque;
 use alloc::sync::Arc;
 use core::fmt;
-use core::mem;
 use core::ops::Deref;
 use core::ops::DerefMut;
 use core::ptr::NonNull;
@@ -278,16 +277,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,12 @@
 //!
 //! Keep in sync with `wasmtime_environ::VMHostFuncOffsets`.
 
-use super::VMOpaqueContext;
+use super::{VMArrayCallNative, VMOpaqueContext};
 use crate::prelude::*;
 use crate::runtime::vm::{StoreBox, VMFuncRef};
 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 +31,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());