Merge pull request #1991 from fitzgen/func-wrap-and-ref-types

Support reference types in `Func::wrap`

crates/wasmtime/src/func.rs

@@ -1,6 +1,6 @@
 use crate::runtime::StoreInner;
 use crate::trampoline::StoreInstanceHandle;
-use crate::{Extern, FuncType, Memory, Store, Trap, Val, ValType};
+use crate::{Extern, ExternRef, FuncType, Memory, Store, Trap, Val, ValType};
 use anyhow::{bail, ensure, Context as _, Result};
 use smallvec::{smallvec, SmallVec};
 use std::cmp::max;
@@ -193,17 +193,37 @@ macro_rules! getters {
                         unsafe extern "C" fn(
                             *mut VMContext,
                             *mut VMContext,
-                            $($args,)*
+                            $( $args::Abi, )*
-                        ) -> R,
+                        ) -> R::Abi,
                     >(anyfunc.as_ref().func_ptr.as_ptr());
+
                     let mut ret = None;
-                    $(let $args = $args.into_abi();)*
+
+                    let weak_store = instance.store.weak();
+                    let weak_store = WeakStore(&weak_store);
+
+                    $(
+                        // Because this returned closure is not marked `unsafe`,
+                        // we have to check that incoming values are compatible
+                        // with our store.
+                        if !$args.compatible_with_store(weak_store) {
+                            return Err(Trap::new(
+                                "attempt to pass cross-`Store` value to Wasm as function argument"
+                            ));
+                        }
+
+                        let $args = $args.into_abi_for_arg(weak_store);
+                    )*
 
                     invoke_wasm_and_catch_traps(anyfunc.as_ref().vmctx, &instance.store, || {
-                        ret = Some(fnptr(anyfunc.as_ref().vmctx, ptr::null_mut(), $($args,)*));
+                        ret = Some(fnptr(
+                            anyfunc.as_ref().vmctx,
+                            ptr::null_mut(),
+                            $( $args, )*
+                        ));
                     })?;
 
-                    Ok(ret.unwrap())
+                    Ok(R::from_abi(ret.unwrap(), weak_store))
                 }
             })
         }
@@ -298,6 +318,22 @@ impl Func {
         }
     }
 
+    pub(crate) unsafe fn from_caller_checked_anyfunc(
+        store: &Store,
+        anyfunc: *mut wasmtime_runtime::VMCallerCheckedAnyfunc,
+    ) -> Option<Self> {
+        let anyfunc = NonNull::new(anyfunc)?;
+        debug_assert!(
+            anyfunc.as_ref().type_index != wasmtime_runtime::VMSharedSignatureIndex::default()
+        );
+
+        let instance_handle = wasmtime_runtime::InstanceHandle::from_vmctx(anyfunc.as_ref().vmctx);
+        let export = wasmtime_runtime::ExportFunction { anyfunc };
+        let instance = store.existing_instance_handle(instance_handle);
+        let f = Func::from_wasmtime_function(export, instance);
+        Some(f)
+    }
+
     /// Creates a new `Func` from the given Rust closure.
     ///
     /// This function will create a new `Func` which, when called, will
@@ -305,16 +341,17 @@ impl Func {
     /// function being called is known statically so the type signature can
     /// be inferred. Rust types will map to WebAssembly types as follows:
     ///
-    /// | Rust Argument Type | WebAssembly Type |
+    /// | Rust Argument Type  | WebAssembly Type |
-    /// |--------------------|------------------|
+    /// |---------------------|------------------|
-    /// | `i32`              | `i32`            |
+    /// | `i32`               | `i32`            |
-    /// | `u32`              | `i32`            |
+    /// | `u32`               | `i32`            |
-    /// | `i64`              | `i64`            |
+    /// | `i64`               | `i64`            |
-    /// | `u64`              | `i64`            |
+    /// | `u64`               | `i64`            |
-    /// | `f32`              | `f32`            |
+    /// | `f32`               | `f32`            |
-    /// | `f64`              | `f64`            |
+    /// | `f64`               | `f64`            |
-    /// | (not supported)    | `v128`           |
+    /// | (not supported)     | `v128`           |
-    /// | (not supported)    | `externref`         |
+    /// | `Option<Func>`      | `funcref`        |
+    /// | `Option<ExternRef>` | `externref`      |
     ///
     /// Any of the Rust types can be returned from the closure as well, in
     /// addition to some extra types
@@ -783,6 +820,12 @@ pub(crate) fn invoke_wasm_and_catch_traps(
     }
 }
 
+// Public (but hidden) wrapper around a `Weak<StoreInner>` so that we can use it
+// in public (but hidden) trait methods.
+#[doc(hidden)]
+#[derive(Clone, Copy)]
+pub struct WeakStore<'a>(&'a Weak<StoreInner>);
+
 /// A trait implemented for types which can be arguments to closures passed to
 /// [`Func::wrap`] and friends.
 ///
@@ -791,160 +834,44 @@ pub(crate) fn invoke_wasm_and_catch_traps(
 /// stable over time.
 ///
 /// For more information see [`Func::wrap`]
-pub unsafe trait WasmTy: Copy {
+pub unsafe trait WasmTy {
+    // The raw ABI representation of this type inside Wasm.
+    #[doc(hidden)]
+    type Abi: Copy;
+
+    // Is this value compatible with the given store?
+    #[doc(hidden)]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool;
+
+    // Convert this value into its ABI representation, when passing a value into
+    // Wasm as an argument.
+    #[doc(hidden)]
+    fn into_abi_for_arg<'a>(self, store: WeakStore<'a>) -> Self::Abi;
+
+    // Convert from the raw ABI representation back into `Self`, when receiving
+    // a value from Wasm.
+    //
+    // Safety: The abi value *must* have be valid for this type (e.g. for
+    // `externref`, it must be a valid raw `VMExternRef` pointer, not some
+    // random, dangling pointer).
+    #[doc(hidden)]
+    unsafe fn from_abi<'a>(abi: Self::Abi, store: WeakStore<'a>) -> Self;
+
+    // Add this type to the given vec of expected valtypes.
     #[doc(hidden)]
     fn push(dst: &mut Vec<ValType>);
+
+    // Does the next valtype(s) match this type?
     #[doc(hidden)]
     fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()>;
+
+    // Load this type's raw ABI representation from an args array.
     #[doc(hidden)]
-    unsafe fn load(ptr: &mut *const u128) -> Self;
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi;
+
+    // Store this type's raw ABI representation into an args array.
     #[doc(hidden)]
-    unsafe fn store(abi: Self, ptr: *mut u128);
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128);
-}
-
-unsafe impl WasmTy for () {
-    fn push(_dst: &mut Vec<ValType>) {}
-    fn matches(_tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        Ok(())
-    }
-    #[inline]
-    unsafe fn load(_ptr: &mut *const u128) -> Self {}
-    #[inline]
-    unsafe fn store(_abi: Self, _ptr: *mut u128) {}
-}
-
-unsafe impl WasmTy for i32 {
-    fn push(dst: &mut Vec<ValType>) {
-        dst.push(ValType::I32);
-    }
-    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        let next = tys.next();
-        ensure!(
-            next == Some(ValType::I32),
-            "Type mismatch, expected i32, got {:?}",
-            next
-        );
-        Ok(())
-    }
-    #[inline]
-    unsafe fn load(ptr: &mut *const u128) -> Self {
-        let ret = **ptr as Self;
-        *ptr = (*ptr).add(1);
-        return ret;
-    }
-    #[inline]
-    unsafe fn store(abi: Self, ptr: *mut u128) {
-        *ptr = abi as u128;
-    }
-}
-
-unsafe impl WasmTy for u32 {
-    fn push(dst: &mut Vec<ValType>) {
-        <i32 as WasmTy>::push(dst)
-    }
-    fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        <i32 as WasmTy>::matches(tys)
-    }
-    #[inline]
-    unsafe fn load(ptr: &mut *const u128) -> Self {
-        <i32 as WasmTy>::load(ptr) as Self
-    }
-    #[inline]
-    unsafe fn store(abi: Self, ptr: *mut u128) {
-        <i32 as WasmTy>::store(abi as i32, ptr)
-    }
-}
-
-unsafe impl WasmTy for i64 {
-    fn push(dst: &mut Vec<ValType>) {
-        dst.push(ValType::I64);
-    }
-    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        let next = tys.next();
-        ensure!(
-            next == Some(ValType::I64),
-            "Type mismatch, expected i64, got {:?}",
-            next
-        );
-        Ok(())
-    }
-    #[inline]
-    unsafe fn load(ptr: &mut *const u128) -> Self {
-        let ret = **ptr as Self;
-        *ptr = (*ptr).add(1);
-        return ret;
-    }
-    #[inline]
-    unsafe fn store(abi: Self, ptr: *mut u128) {
-        *ptr = abi as u128;
-    }
-}
-
-unsafe impl WasmTy for u64 {
-    fn push(dst: &mut Vec<ValType>) {
-        <i64 as WasmTy>::push(dst)
-    }
-    fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        <i64 as WasmTy>::matches(tys)
-    }
-    #[inline]
-    unsafe fn load(ptr: &mut *const u128) -> Self {
-        <i64 as WasmTy>::load(ptr) as Self
-    }
-    #[inline]
-    unsafe fn store(abi: Self, ptr: *mut u128) {
-        <i64 as WasmTy>::store(abi as i64, ptr)
-    }
-}
-
-unsafe impl WasmTy for f32 {
-    fn push(dst: &mut Vec<ValType>) {
-        dst.push(ValType::F32);
-    }
-    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        let next = tys.next();
-        ensure!(
-            next == Some(ValType::F32),
-            "Type mismatch, expected f32, got {:?}",
-            next
-        );
-        Ok(())
-    }
-    #[inline]
-    unsafe fn load(ptr: &mut *const u128) -> Self {
-        let ret = f32::from_bits(**ptr as u32);
-        *ptr = (*ptr).add(1);
-        return ret;
-    }
-    #[inline]
-    unsafe fn store(abi: Self, ptr: *mut u128) {
-        *ptr = abi.to_bits() as u128;
-    }
-}
-
-unsafe impl WasmTy for f64 {
-    fn push(dst: &mut Vec<ValType>) {
-        dst.push(ValType::F64);
-    }
-    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        let next = tys.next();
-        ensure!(
-            next == Some(ValType::F64),
-            "Type mismatch, expected f64, got {:?}",
-            next
-        );
-        Ok(())
-    }
-    #[inline]
-    unsafe fn load(ptr: &mut *const u128) -> Self {
-        let ret = f64::from_bits(**ptr as u64);
-        *ptr = (*ptr).add(1);
-        return ret;
-    }
-    #[inline]
-    unsafe fn store(abi: Self, ptr: *mut u128) {
-        *ptr = abi.to_bits() as u128;
-    }
 }
 
 /// A trait implemented for types which can be returned from closures passed to
@@ -956,64 +883,527 @@ unsafe impl WasmTy for f64 {
 ///
 /// For more information see [`Func::wrap`]
 pub unsafe trait WasmRet {
+    // Same as `WasmTy::Abi`.
     #[doc(hidden)]
-    type Abi;
+    type Abi: Copy;
+
+    // Same as `WasmTy::compatible_with_store`.
+    #[doc(hidden)]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool;
+
+    // Similar to `WasmTy::into_abi_for_arg` but used when host code is
+    // returning a value into Wasm, rather than host code passing an argument to
+    // a Wasm call. Unlike `into_abi_for_arg`, implementors of this method can
+    // raise traps, which means that callers must ensure that
+    // `invoke_wasm_and_catch_traps` is on the stack, and therefore this method
+    // is unsafe.
+    #[doc(hidden)]
+    unsafe fn into_abi_for_ret<'a>(self, store: WeakStore<'a>) -> Self::Abi;
+
+    // Same as `WasmTy::from_abi`.
+    #[doc(hidden)]
+    unsafe fn from_abi<'a>(abi: Self::Abi, store: WeakStore<'a>) -> Self;
+
+    // Same as `WasmTy::push`.
     #[doc(hidden)]
     fn push(dst: &mut Vec<ValType>);
+
+    // Same as `WasmTy::matches`.
     #[doc(hidden)]
     fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()>;
+
+    // Same as `WasmTy::load_from_args`.
     #[doc(hidden)]
-    fn into_abi(self) -> Self::Abi;
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi;
+
+    // Same as `WasmTy::store_to_args`.
     #[doc(hidden)]
-    unsafe fn store(abi: Self::Abi, ptr: *mut u128);
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128);
 }
 
-unsafe impl<T: WasmTy> WasmRet for T {
+unsafe impl WasmTy for () {
-    type Abi = T;
+    type Abi = Self;
-    fn push(dst: &mut Vec<ValType>) {
-        T::push(dst)
-    }
 
-    fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+    #[inline]
-        T::matches(tys)
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
     }
 
     #[inline]
-    fn into_abi(self) -> Self::Abi {
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {}
+
+    #[inline]
+    unsafe fn from_abi<'a>(_abi: Self::Abi, _store: WeakStore<'a>) -> Self {}
+
+    fn push(_dst: &mut Vec<ValType>) {}
+
+    fn matches(_tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        Ok(())
+    }
+
+    #[inline]
+    unsafe fn load_from_args(_ptr: &mut *const u128) -> Self::Abi {}
+
+    #[inline]
+    unsafe fn store_to_args(_abi: Self::Abi, _ptr: *mut u128) {}
+}
+
+unsafe impl WasmTy for i32 {
+    type Abi = Self;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
         self
     }
 
     #[inline]
-    unsafe fn store(abi: Self::Abi, ptr: *mut u128) {
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
-        T::store(abi, ptr);
+        abi
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        dst.push(ValType::I32);
+    }
+
+    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        let next = tys.next();
+        ensure!(
+            next == Some(ValType::I32),
+            "Type mismatch, expected i32, got {:?}",
+            next
+        );
+        Ok(())
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        let ret = **ptr as Self;
+        *ptr = (*ptr).add(1);
+        return ret;
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        *ptr = abi as u128;
     }
 }
 
-unsafe impl<T: WasmTy> WasmRet for Result<T, Trap> {
+unsafe impl WasmTy for u32 {
-    type Abi = T;
+    type Abi = <i32 as WasmTy>::Abi;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
+        self as i32
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
+        abi as Self
+    }
+
     fn push(dst: &mut Vec<ValType>) {
-        T::push(dst)
+        <i32 as WasmTy>::push(dst)
     }
 
     fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
-        T::matches(tys)
+        <i32 as WasmTy>::matches(tys)
     }
 
     #[inline]
-    fn into_abi(self) -> Self::Abi {
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        <i32 as WasmTy>::load_from_args(ptr)
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        <i32 as WasmTy>::store_to_args(abi, ptr)
+    }
+}
+
+unsafe impl WasmTy for i64 {
+    type Abi = Self;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
+        self
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
+        abi
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        dst.push(ValType::I64);
+    }
+
+    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        let next = tys.next();
+        ensure!(
+            next == Some(ValType::I64),
+            "Type mismatch, expected i64, got {:?}",
+            next
+        );
+        Ok(())
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        let ret = **ptr as Self;
+        *ptr = (*ptr).add(1);
+        return ret;
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        *ptr = abi as u128;
+    }
+}
+
+unsafe impl WasmTy for u64 {
+    type Abi = <i64 as WasmTy>::Abi;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
+        self as i64
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
+        abi as Self
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        <i64 as WasmTy>::push(dst)
+    }
+
+    fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        <i64 as WasmTy>::matches(tys)
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        <i64 as WasmTy>::load_from_args(ptr)
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        <i64 as WasmTy>::store_to_args(abi, ptr)
+    }
+}
+
+unsafe impl WasmTy for f32 {
+    type Abi = Self;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
+        self
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
+        abi
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        dst.push(ValType::F32);
+    }
+
+    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        let next = tys.next();
+        ensure!(
+            next == Some(ValType::F32),
+            "Type mismatch, expected f32, got {:?}",
+            next
+        );
+        Ok(())
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        let ret = f32::from_bits(**ptr as u32);
+        *ptr = (*ptr).add(1);
+        return ret;
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        *ptr = abi.to_bits() as u128;
+    }
+}
+
+unsafe impl WasmTy for f64 {
+    type Abi = Self;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
+        self
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
+        abi
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        dst.push(ValType::F64);
+    }
+
+    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        let next = tys.next();
+        ensure!(
+            next == Some(ValType::F64),
+            "Type mismatch, expected f64, got {:?}",
+            next
+        );
+        Ok(())
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        let ret = f64::from_bits(**ptr as u64);
+        *ptr = (*ptr).add(1);
+        return ret;
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        *ptr = abi.to_bits() as u128;
+    }
+}
+
+unsafe impl WasmTy for Option<ExternRef> {
+    type Abi = *mut u8;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, _store: WeakStore<'a>) -> bool {
+        true
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, store: WeakStore<'a>) -> Self::Abi {
+        if let Some(x) = self {
+            let store = Store::upgrade(store.0).unwrap();
+            let abi = x.inner.as_raw();
+            unsafe {
+                store
+                    .externref_activations_table()
+                    .insert_with_gc(x.inner, store.stack_map_registry());
+            }
+            abi
+        } else {
+            ptr::null_mut()
+        }
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, _store: WeakStore<'a>) -> Self {
+        if abi.is_null() {
+            None
+        } else {
+            Some(ExternRef {
+                inner: wasmtime_runtime::VMExternRef::clone_from_raw(abi),
+            })
+        }
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        dst.push(ValType::ExternRef);
+    }
+
+    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        let next = tys.next();
+        ensure!(
+            next == Some(ValType::ExternRef),
+            "Type mismatch, expected externref, got {:?}",
+            next
+        );
+        Ok(())
+    }
+
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        let ret = **ptr as usize as *mut u8;
+        *ptr = (*ptr).add(1);
+        ret
+    }
+
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        ptr::write(ptr, abi as usize as u128);
+    }
+}
+
+unsafe impl WasmTy for Option<Func> {
+    type Abi = *mut wasmtime_runtime::VMCallerCheckedAnyfunc;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool {
+        if let Some(f) = self {
+            let store = Store::upgrade(store.0).unwrap();
+            Store::same(&store, f.store())
+        } else {
+            true
+        }
+    }
+
+    #[inline]
+    fn into_abi_for_arg<'a>(self, _store: WeakStore<'a>) -> Self::Abi {
+        if let Some(f) = self {
+            f.caller_checked_anyfunc().as_ptr()
+        } else {
+            ptr::null_mut()
+        }
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, store: WeakStore<'a>) -> Self {
+        let store = Store::upgrade(store.0).unwrap();
+        Func::from_caller_checked_anyfunc(&store, abi)
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        dst.push(ValType::FuncRef);
+    }
+
+    fn matches(mut tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        let next = tys.next();
+        ensure!(
+            next == Some(ValType::FuncRef),
+            "Type mismatch, expected funcref, got {:?}",
+            next
+        );
+        Ok(())
+    }
+
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        let ret = **ptr as usize as *mut wasmtime_runtime::VMCallerCheckedAnyfunc;
+        *ptr = (*ptr).add(1);
+        ret
+    }
+
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        ptr::write(ptr, abi as usize as u128);
+    }
+}
+
+unsafe impl<T> WasmRet for T
+where
+    T: WasmTy,
+{
+    type Abi = <T as WasmTy>::Abi;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool {
+        <Self as WasmTy>::compatible_with_store(self, store)
+    }
+
+    #[inline]
+    unsafe fn into_abi_for_ret<'a>(self, store: WeakStore<'a>) -> Self::Abi {
+        <Self as WasmTy>::into_abi_for_arg(self, store)
+    }
+
+    #[inline]
+    unsafe fn from_abi<'a>(abi: Self::Abi, store: WeakStore<'a>) -> Self {
+        <Self as WasmTy>::from_abi(abi, store)
+    }
+
+    #[inline]
+    fn push(dst: &mut Vec<ValType>) {
+        <Self as WasmTy>::push(dst)
+    }
+
+    #[inline]
+    fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        <Self as WasmTy>::matches(tys)
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        <Self as WasmTy>::load_from_args(ptr)
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        <Self as WasmTy>::store_to_args(abi, ptr)
+    }
+}
+
+unsafe impl<T> WasmRet for Result<T, Trap>
+where
+    T: WasmTy,
+{
+    type Abi = <T as WasmTy>::Abi;
+
+    #[inline]
+    fn compatible_with_store<'a>(&self, store: WeakStore<'a>) -> bool {
         match self {
-            Ok(val) => return T::into_abi(val),
+            Ok(x) => <T as WasmTy>::compatible_with_store(x, store),
+            Err(_) => true,
+        }
+    }
+
+    #[inline]
+    unsafe fn into_abi_for_ret<'a>(self, store: WeakStore<'a>) -> Self::Abi {
+        match self {
+            Ok(val) => return <T as WasmTy>::into_abi_for_arg(val, store),
             Err(trap) => handle_trap(trap),
         }
 
-        fn handle_trap(trap: Trap) -> ! {
+        unsafe fn handle_trap(trap: Trap) -> ! {
-            unsafe { raise_user_trap(trap.into()) }
+            raise_user_trap(trap.into())
         }
     }
 
     #[inline]
-    unsafe fn store(abi: Self::Abi, ptr: *mut u128) {
+    unsafe fn from_abi<'a>(abi: Self::Abi, store: WeakStore<'a>) -> Self {
-        T::store(abi, ptr);
+        Ok(<T as WasmTy>::from_abi(abi, store))
+    }
+
+    fn push(dst: &mut Vec<ValType>) {
+        <T as WasmTy>::push(dst)
+    }
+
+    fn matches(tys: impl Iterator<Item = ValType>) -> anyhow::Result<()> {
+        <T as WasmTy>::matches(tys)
+    }
+
+    #[inline]
+    unsafe fn load_from_args(ptr: &mut *const u128) -> Self::Abi {
+        <T as WasmTy>::load_from_args(ptr)
+    }
+
+    #[inline]
+    unsafe fn store_to_args(abi: Self::Abi, ptr: *mut u128) {
+        <T as WasmTy>::store_to_args(abi, ptr);
     }
 }
 
@@ -1112,6 +1502,27 @@ impl Caller<'_> {
     }
 }
 
+#[inline(never)]
+#[cold]
+unsafe fn raise_cross_store_trap() -> ! {
+    #[derive(Debug)]
+    struct CrossStoreError;
+
+    impl std::error::Error for CrossStoreError {}
+
+    impl fmt::Display for CrossStoreError {
+        fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+            write!(
+                f,
+                "host function attempted to return cross-`Store` \
+                 value to Wasm",
+            )
+        }
+    }
+
+    raise_user_trap(Box::new(CrossStoreError));
+}
+
 macro_rules! impl_into_func {
     ($(
         ($($args:ident)*)
@@ -1141,40 +1552,62 @@ macro_rules! impl_into_func {
             R: WasmRet,
         {
             fn into_func(self, store: &Store) -> Func {
-                // Note that this shim's ABI must match that expected by
+                /// This shim is called by Wasm code, constructs a `Caller`,
-                // cranelift, since cranelift is generating raw function calls
+                /// calls the wrapped host function, and returns the translated
-                // directly to this function.
+                /// result back to Wasm.
-                unsafe extern "C" fn shim<F, $($args,)* R>(
+                ///
+                /// Note that this shim's ABI must *exactly* match that expected
+                /// by Cranelift, since Cranelift is generating raw function
+                /// calls directly to this function.
+                unsafe extern "C" fn wasm_to_host_shim<F, $($args,)* R>(
                     vmctx: *mut VMContext,
                     caller_vmctx: *mut VMContext,
-                    $($args: $args,)*
+                    $( $args: $args::Abi, )*
                 ) -> R::Abi
                 where
-                    F: Fn(Caller<'_>, $($args),*) -> R + 'static,
+                    F: Fn(Caller<'_>, $( $args ),*) -> R + 'static,
-                    $($args: WasmTy,)*
+                    $( $args: WasmTy, )*
                     R: WasmRet,
                 {
+                    let state = (*vmctx).host_state();
+                    // Double-check ourselves in debug mode, but we control
+                    // the `Any` here so an unsafe downcast should also
+                    // work.
+                    debug_assert!(state.is::<(F, Weak<StoreInner>)>());
+                    let (func, store) = &*(state as *const _ as *const (F, Weak<StoreInner>));
+                    let weak_store = WeakStore(store);
+
                     let ret = {
-                        let state = (*vmctx).host_state();
-                        // Double-check ourselves in debug mode, but we control
-                        // the `Any` here so an unsafe downcast should also
-                        // work.
-                        debug_assert!(state.is::<(F, Weak<StoreInner>)>());
-                        let (func, store) = &*(state as *const _ as *const (F, Weak<StoreInner>));
                         panic::catch_unwind(AssertUnwindSafe(|| {
                             func(
                                 Caller { store, caller_vmctx },
-                                $($args,)*
+                                $( $args::from_abi($args, weak_store), )*
                             )
                         }))
                     };
                     match ret {
-                        Ok(ret) => ret.into_abi(),
                         Err(panic) => wasmtime_runtime::resume_panic(panic),
+                        Ok(ret) => {
+                            // Because the wrapped function is not `unsafe`, we
+                            // can't assume it returned a value that is
+                            // compatible with this store.
+                            if !ret.compatible_with_store(weak_store) {
+                                raise_cross_store_trap();
+                            }
+
+                            ret.into_abi_for_ret(weak_store)
+                        }
                     }
                 }
 
-                unsafe extern "C" fn trampoline<$($args,)* R>(
+                /// This trampoline allows host code to indirectly call the
+                /// wrapped function (e.g. via `Func::call` on a `funcref` that
+                /// happens to reference our wrapped function).
+                ///
+                /// It reads the arguments out of the incoming `args` array,
+                /// calls the given function pointer, and then stores the result
+                /// back into the `args` array.
+                unsafe extern "C" fn host_trampoline<$($args,)* R>(
                     callee_vmctx: *mut VMContext,
                     caller_vmctx: *mut VMContext,
                     ptr: *const VMFunctionBody,
@@ -1189,14 +1622,14 @@ macro_rules! impl_into_func {
                         unsafe extern "C" fn(
                             *mut VMContext,
                             *mut VMContext,
-                            $($args,)*
+                            $( $args::Abi, )*
                         ) -> R::Abi,
                     >(ptr);
 
                     let mut _next = args as *const u128;
-                    $(let $args = $args::load(&mut _next);)*
+                    $( let $args = $args::load_from_args(&mut _next); )*
-                    let ret = ptr(callee_vmctx, caller_vmctx, $($args),*);
+                    let ret = ptr(callee_vmctx, caller_vmctx, $( $args ),*);
-                    R::store(ret, args);
+                    R::store_to_args(ret, args);
                 }
 
                 let mut _args = Vec::new();
@@ -1204,13 +1637,14 @@ macro_rules! impl_into_func {
                 let mut ret = Vec::new();
                 R::push(&mut ret);
                 let ty = FuncType::new(_args.into(), ret.into());
+
                 let store_weak = store.weak();
-                let trampoline = trampoline::<$($args,)* R>;
+                let trampoline = host_trampoline::<$($args,)* R>;
                 let (instance, export) = unsafe {
                     crate::trampoline::generate_raw_func_export(
                         &ty,
                         std::slice::from_raw_parts_mut(
-                            shim::<F, $($args,)* R> as *mut _,
+                            wasm_to_host_shim::<F, $($args,)* R> as *mut _,
                             0,
                         ),
                         trampoline,
@@ -1219,6 +1653,7 @@ macro_rules! impl_into_func {
                     )
                     .expect("failed to generate export")
                 };
+
                 Func {
                     instance,
                     export,

crates/wasmtime/src/values.rs

@@ -1,7 +1,7 @@
 use crate::r#ref::ExternRef;
 use crate::{Func, Store, ValType};
 use anyhow::{bail, Result};
-use std::ptr::{self, NonNull};
+use std::ptr;
 use wasmtime_runtime::{self as runtime, VMExternRef};
 
 /// Possible runtime values that a WebAssembly module can either consume or
@@ -270,17 +270,5 @@ pub(crate) unsafe fn from_checked_anyfunc(
     anyfunc: *mut wasmtime_runtime::VMCallerCheckedAnyfunc,
     store: &Store,
 ) -> Val {
-    let anyfunc = match NonNull::new(anyfunc) {
+    Val::FuncRef(Func::from_caller_checked_anyfunc(store, anyfunc))
-        None => return Val::FuncRef(None),
-        Some(f) => f,
-    };
-
-    debug_assert!(
-        anyfunc.as_ref().type_index != wasmtime_runtime::VMSharedSignatureIndex::default()
-    );
-    let instance_handle = wasmtime_runtime::InstanceHandle::from_vmctx(anyfunc.as_ref().vmctx);
-    let export = wasmtime_runtime::ExportFunction { anyfunc };
-    let instance = store.existing_instance_handle(instance_handle);
-    let f = Func::from_wasmtime_function(export, instance);
-    Val::FuncRef(Some(f))
 }

tests/all/func.rs

@@ -13,12 +13,16 @@ fn func_constructors() {
     Func::wrap(&store, || -> i64 { 0 });
     Func::wrap(&store, || -> f32 { 0.0 });
     Func::wrap(&store, || -> f64 { 0.0 });
+    Func::wrap(&store, || -> Option<ExternRef> { None });
+    Func::wrap(&store, || -> Option<Func> { None });
 
     Func::wrap(&store, || -> Result<(), Trap> { loop {} });
     Func::wrap(&store, || -> Result<i32, Trap> { loop {} });
     Func::wrap(&store, || -> Result<i64, Trap> { loop {} });
     Func::wrap(&store, || -> Result<f32, Trap> { loop {} });
     Func::wrap(&store, || -> Result<f64, Trap> { loop {} });
+    Func::wrap(&store, || -> Result<Option<ExternRef>, Trap> { loop {} });
+    Func::wrap(&store, || -> Result<Option<Func>, Trap> { loop {} });
 }
 
 #[test]
@@ -95,9 +99,12 @@ fn signatures_match() {
     assert_eq!(f.ty().params(), &[]);
     assert_eq!(f.ty().results(), &[ValType::F64]);
 
-    let f = Func::wrap(&store, |_: f32, _: f64, _: i32, _: i64, _: i32| -> f64 {
+    let f = Func::wrap(
-        loop {}
+        &store,
-    });
+        |_: f32, _: f64, _: i32, _: i64, _: i32, _: Option<ExternRef>, _: Option<Func>| -> f64 {
+            loop {}
+        },
+    );
     assert_eq!(
         f.ty().params(),
         &[
@@ -105,13 +112,18 @@ fn signatures_match() {
             ValType::F64,
             ValType::I32,
             ValType::I64,
-            ValType::I32
+            ValType::I32,
+            ValType::ExternRef,
+            ValType::FuncRef,
         ]
     );
     assert_eq!(f.ty().results(), &[ValType::F64]);
 }
 
 #[test]
+// Note: Cranelift only supports refrerence types (used in the wasm in this
+// test) on x64.
+#[cfg(target_arch = "x86_64")]
 fn import_works() -> Result<()> {
     static HITS: AtomicUsize = AtomicUsize::new(0);
 
@@ -120,9 +132,9 @@ fn import_works() -> Result<()> {
             (import "" "" (func))
             (import "" "" (func (param i32) (result i32)))
             (import "" "" (func (param i32) (param i64)))
-            (import "" "" (func (param i32 i64 i32 f32 f64)))
+            (import "" "" (func (param i32 i64 i32 f32 f64 externref funcref)))
 
-            (func $foo
+            (func (export "run") (param externref funcref)
                 call 0
                 i32.const 0
                 call 1
@@ -136,14 +148,18 @@ fn import_works() -> Result<()> {
                 i32.const 300
                 f32.const 400
                 f64.const 500
+                local.get 0
+                local.get 1
                 call 3
             )
-            (start $foo)
         "#,
     )?;
-    let store = Store::default();
+    let mut config = Config::new();
-    let module = Module::new(store.engine(), &wasm)?;
+    config.wasm_reference_types(true);
-    Instance::new(
+    let engine = Engine::new(&config);
+    let store = Store::new(&engine);
+    let module = Module::new(&engine, &wasm)?;
+    let instance = Instance::new(
         &store,
         &module,
         &[
@@ -163,17 +179,31 @@ fn import_works() -> Result<()> {
                 assert_eq!(HITS.fetch_add(1, SeqCst), 2);
             })
             .into(),
-            Func::wrap(&store, |a: i32, b: i64, c: i32, d: f32, e: f64| {
+            Func::wrap(
-                assert_eq!(a, 100);
+                &store,
-                assert_eq!(b, 200);
+                |a: i32, b: i64, c: i32, d: f32, e: f64, f: Option<ExternRef>, g: Option<Func>| {
-                assert_eq!(c, 300);
+                    assert_eq!(a, 100);
-                assert_eq!(d, 400.0);
+                    assert_eq!(b, 200);
-                assert_eq!(e, 500.0);
+                    assert_eq!(c, 300);
-                assert_eq!(HITS.fetch_add(1, SeqCst), 3);
+                    assert_eq!(d, 400.0);
-            })
+                    assert_eq!(e, 500.0);
+                    assert_eq!(
+                        f.as_ref().unwrap().data().downcast_ref::<String>().unwrap(),
+                        "hello"
+                    );
+                    assert_eq!(g.as_ref().unwrap().call(&[]).unwrap()[0].unwrap_i32(), 42);
+                    assert_eq!(HITS.fetch_add(1, SeqCst), 3);
+                },
+            )
             .into(),
         ],
     )?;
+    let run = instance.get_func("run").unwrap();
+    run.call(&[
+        Val::ExternRef(Some(ExternRef::new("hello".to_string()))),
+        Val::FuncRef(Some(Func::wrap(&store, || -> i32 { 42 }))),
+    ])?;
+    assert_eq!(HITS.load(SeqCst), 4);
     Ok(())
 }
 
@@ -228,6 +258,10 @@ fn get_from_wrapper() {
     assert!(f.get0::<f32>().is_ok());
     let f = Func::wrap(&store, || -> f64 { loop {} });
     assert!(f.get0::<f64>().is_ok());
+    let f = Func::wrap(&store, || -> Option<ExternRef> { loop {} });
+    assert!(f.get0::<Option<ExternRef>>().is_ok());
+    let f = Func::wrap(&store, || -> Option<Func> { loop {} });
+    assert!(f.get0::<Option<Func>>().is_ok());
 
     let f = Func::wrap(&store, |_: i32| {});
     assert!(f.get1::<i32, ()>().is_ok());
@@ -240,6 +274,10 @@ fn get_from_wrapper() {
     assert!(f.get1::<f32, ()>().is_ok());
     let f = Func::wrap(&store, |_: f64| {});
     assert!(f.get1::<f64, ()>().is_ok());
+    let f = Func::wrap(&store, |_: Option<ExternRef>| {});
+    assert!(f.get1::<Option<ExternRef>, ()>().is_ok());
+    let f = Func::wrap(&store, |_: Option<Func>| {});
+    assert!(f.get1::<Option<Func>, ()>().is_ok());
 }
 
 #[test]
@@ -404,3 +442,68 @@ fn func_write_nothing() -> anyhow::Result<()> {
         .contains("function attempted to return an incompatible value"));
     Ok(())
 }
+
+#[test]
+// Note: Cranelift only supports refrerence types (used in the wasm in this
+// test) on x64.
+#[cfg(target_arch = "x86_64")]
+fn return_cross_store_value() -> anyhow::Result<()> {
+    let wasm = wat::parse_str(
+        r#"
+            (import "" "" (func (result funcref)))
+
+            (func (export "run") (result funcref)
+                call 0
+            )
+        "#,
+    )?;
+    let mut config = Config::new();
+    config.wasm_reference_types(true);
+    let engine = Engine::new(&config);
+    let module = Module::new(&engine, &wasm)?;
+
+    let store1 = Store::new(&engine);
+    let store2 = Store::new(&engine);
+
+    let store2_func = Func::wrap(&store2, || {});
+    let return_cross_store_func = Func::wrap(&store1, move || Some(store2_func.clone()));
+
+    let instance = Instance::new(&store1, &module, &[return_cross_store_func.into()])?;
+
+    let run = instance.get_func("run").unwrap();
+    let result = run.call(&[]);
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("cross-`Store`"));
+
+    Ok(())
+}
+
+#[test]
+// Note: Cranelift only supports refrerence types (used in the wasm in this
+// test) on x64.
+#[cfg(target_arch = "x86_64")]
+fn pass_cross_store_arg() -> anyhow::Result<()> {
+    let mut config = Config::new();
+    config.wasm_reference_types(true);
+    let engine = Engine::new(&config);
+
+    let store1 = Store::new(&engine);
+    let store2 = Store::new(&engine);
+
+    let store1_func = Func::wrap(&store1, |_: Option<Func>| {});
+    let store2_func = Func::wrap(&store2, || {});
+
+    // Using regular `.call` fails with cross-Store arguments.
+    assert!(store1_func
+        .call(&[Val::FuncRef(Some(store2_func.clone()))])
+        .is_err());
+
+    // And using `.get` followed by a function call also fails with cross-Store
+    // arguments.
+    let f = store1_func.get1::<Option<Func>, ()>()?;
+    let result = f(Some(store2_func));
+    assert!(result.is_err());
+    assert!(result.unwrap_err().to_string().contains("cross-`Store`"));
+
+    Ok(())
+}

tests/all/gc.rs

@@ -160,35 +160,16 @@ fn many_live_refs() -> anyhow::Result<()> {
 
     let live_refs = Rc::new(Cell::new(0));
 
-    let make_ref = Func::new(
+    let make_ref = Func::wrap(&store, {
-        &store,
+        let live_refs = live_refs.clone();
-        FuncType::new(
+        move || Some(ExternRef::new(CountLiveRefs::new(live_refs.clone())))
-            vec![].into_boxed_slice(),
+    });
-            vec![ValType::ExternRef].into_boxed_slice(),
-        ),
-        {
-            let live_refs = live_refs.clone();
-            move |_caller, _params, results| {
-                results[0] =
-                    Val::ExternRef(Some(ExternRef::new(CountLiveRefs::new(live_refs.clone()))));
-                Ok(())
-            }
-        },
-    );
 
-    let observe_ref = Func::new(
+    let observe_ref = Func::wrap(&store, |r: Option<ExternRef>| {
-        &store,
+        let r = r.unwrap();
-        FuncType::new(
+        let r = r.data().downcast_ref::<CountLiveRefs>().unwrap();
-            vec![ValType::ExternRef].into_boxed_slice(),
+        assert!(r.live_refs.get() > 0);
-            vec![].into_boxed_slice(),
+    });
-        ),
-        |_caller, params, _results| {
-            let r = params[0].externref().unwrap().unwrap();
-            let r = r.data().downcast_ref::<CountLiveRefs>().unwrap();
-            assert!(r.live_refs.get() > 0);
-            Ok(())
-        },
-    );
 
     let instance = Instance::new(&store, &module, &[make_ref.into(), observe_ref.into()])?;
     let many_live_refs = instance.get_func("many_live_refs").unwrap();
@@ -413,14 +394,7 @@ fn gc_during_gc_from_many_table_gets() -> anyhow::Result<()> {
         "#,
     )?;
 
-    let observe_ref = Func::new(
+    let observe_ref = Func::wrap(&store, |_: Option<ExternRef>| {});
-        &store,
-        FuncType::new(
-            vec![ValType::ExternRef].into_boxed_slice(),
-            vec![].into_boxed_slice(),
-        ),
-        |_caller, _params, _results| Ok(()),
-    );
 
     let instance = Instance::new(&store, &module, &[observe_ref.into()])?;
     let init = instance.get_func("init").unwrap();