From 46ef80bf2f744fb38f370381760105cdca5735be Mon Sep 17 00:00:00 2001
From: Nick Fitzgerald <fitzgen@gmail.com>
Date: Tue, 7 Jul 2020 10:16:28 -0700
Subject: [PATCH] wasmtime: Support `ExternRef`s in `Func::wrap`'d functions

Fixes #1868
---
 crates/wasmtime/src/func.rs | 57 ++++++++++++++++++++++++++++++++++++-
 tests/all/gc.rs             | 46 +++++++-----------------------
 2 files changed, 66 insertions(+), 37 deletions(-)

diff --git a/crates/wasmtime/src/func.rs b/crates/wasmtime/src/func.rs
index 3a775ee55d..70f5023850 100644
--- a/crates/wasmtime/src/func.rs
+++ b/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;
@@ -1131,6 +1131,61 @@ unsafe impl WasmTy for f64 {
     }
 }
 
+unsafe impl WasmTy for Option<ExternRef> {
+    type Abi = *mut u8;
+
+    #[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<T> WasmRet for T
 where
     T: WasmTy,
diff --git a/tests/all/gc.rs b/tests/all/gc.rs
index 3e19b31758..e616b2d0c8 100644
--- a/tests/all/gc.rs
+++ b/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(
-        &store,
-        FuncType::new(
-            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 make_ref = Func::wrap(&store, {
+        let live_refs = live_refs.clone();
+        move || Some(ExternRef::new(CountLiveRefs::new(live_refs.clone())))
+    });
 
-    let observe_ref = Func::new(
-        &store,
-        FuncType::new(
-            vec![ValType::ExternRef].into_boxed_slice(),
-            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 observe_ref = Func::wrap(&store, |r: Option<ExternRef>| {
+        let r = r.unwrap();
+        let r = r.data().downcast_ref::<CountLiveRefs>().unwrap();
+        assert!(r.live_refs.get() > 0);
+    });
 
     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(
-        &store,
-        FuncType::new(
-            vec![ValType::ExternRef].into_boxed_slice(),
-            vec![].into_boxed_slice(),
-        ),
-        |_caller, _params, _results| Ok(()),
-    );
+    let observe_ref = Func::wrap(&store, |_: Option<ExternRef>| {});
 
     let instance = Instance::new(&store, &module, &[observe_ref.into()])?;
     let init = instance.get_func("init").unwrap();