Merge pull request #1417 from alexcrichton/closure-ref

fitzgen · web-flow · commit f367a4247bfc · 2019-04-01T16:21:54.000-07:00
Support 1-reference argument closures
diff --git a/src/closure.rs b/src/closure.rs
@@ -519,7 +519,7 @@ where
 /// This trait is not stable and it's not recommended to use this in bounds or
 /// implement yourself.
 #[doc(hidden)]
-pub unsafe trait WasmClosure: 'static {
+pub unsafe trait WasmClosure {
     fn describe();
 }
 
@@ -541,7 +541,7 @@ macro_rules! doit {
     ($(
         ($($var:ident)*)
     )*) => ($(
-        unsafe impl<$($var,)* R> WasmClosure for Fn($($var),*) -> R
+        unsafe impl<$($var,)* R> WasmClosure for Fn($($var),*) -> R + 'static
             where $($var: FromWasmAbi + 'static,)*
                   R: ReturnWasmAbi + 'static,
         {
@@ -587,7 +587,7 @@ macro_rules! doit {
             }
         }
 
-        unsafe impl<$($var,)* R> WasmClosure for FnMut($($var),*) -> R
+        unsafe impl<$($var,)* R> WasmClosure for FnMut($($var),*) -> R + 'static
             where $($var: FromWasmAbi + 'static,)*
                   R: ReturnWasmAbi + 'static,
         {
@@ -696,3 +696,148 @@ doit! {
     (A B C D E F)
     (A B C D E F G)
 }
+
+// Copy the above impls down here for where there's only one argument and it's a
+// reference. We could add more impls for more kinds of references, but it
+// becomes a combinatorial explosion quickly. Let's see how far we can get with
+// just this one! Maybe someone else can figure out voodoo so we don't have to
+// duplicate.
+
+unsafe impl<A, R> WasmClosure for Fn(&A) -> R
+    where A: RefFromWasmAbi,
+          R: ReturnWasmAbi + 'static,
+{
+    fn describe() {
+        #[allow(non_snake_case)]
+        unsafe extern "C" fn invoke<A: RefFromWasmAbi, R: ReturnWasmAbi>(
+            a: usize,
+            b: usize,
+            arg: <A as RefFromWasmAbi>::Abi,
+        ) -> <R as ReturnWasmAbi>::Abi {
+            if a == 0 {
+                throw_str("closure invoked recursively or destroyed already");
+            }
+            // Make sure all stack variables are converted before we
+            // convert `ret` as it may throw (for `Result`, for
+            // example)
+            let ret = {
+                let f: *const Fn(&A) -> R =
+                    FatPtr { fields: (a, b) }.ptr;
+                let mut _stack = GlobalStack::new();
+                let arg = <A as RefFromWasmAbi>::ref_from_abi(arg, &mut _stack);
+                (*f)(&*arg)
+            };
+            ret.return_abi(&mut GlobalStack::new())
+        }
+
+        inform(invoke::<A, R> as u32);
+
+        unsafe extern fn destroy<A: RefFromWasmAbi, R: ReturnWasmAbi>(
+            a: usize,
+            b: usize,
+        ) {
+            debug_assert!(a != 0, "should never destroy a Fn whose pointer is 0");
+            drop(Box::from_raw(FatPtr::<Fn(&A) -> R> {
+                fields: (a, b)
+            }.ptr));
+        }
+        inform(destroy::<A, R> as u32);
+
+        <&Self>::describe();
+    }
+}
+
+unsafe impl<A, R> WasmClosure for FnMut(&A) -> R
+    where A: RefFromWasmAbi,
+          R: ReturnWasmAbi + 'static,
+{
+    fn describe() {
+        #[allow(non_snake_case)]
+        unsafe extern "C" fn invoke<A: RefFromWasmAbi, R: ReturnWasmAbi>(
+            a: usize,
+            b: usize,
+            arg: <A as RefFromWasmAbi>::Abi,
+        ) -> <R as ReturnWasmAbi>::Abi {
+            if a == 0 {
+                throw_str("closure invoked recursively or destroyed already");
+            }
+            // Make sure all stack variables are converted before we
+            // convert `ret` as it may throw (for `Result`, for
+            // example)
+            let ret = {
+                let f: *const FnMut(&A) -> R =
+                    FatPtr { fields: (a, b) }.ptr;
+                let f = f as *mut FnMut(&A) -> R;
+                let mut _stack = GlobalStack::new();
+                let arg = <A as RefFromWasmAbi>::ref_from_abi(arg, &mut _stack);
+                (*f)(&*arg)
+            };
+            ret.return_abi(&mut GlobalStack::new())
+        }
+
+        inform(invoke::<A, R> as u32);
+
+        unsafe extern fn destroy<A: RefFromWasmAbi, R: ReturnWasmAbi>(
+            a: usize,
+            b: usize,
+        ) {
+            debug_assert!(a != 0, "should never destroy a FnMut whose pointer is 0");
+            drop(Box::from_raw(FatPtr::<FnMut(&A) -> R> {
+                fields: (a, b)
+            }.ptr));
+        }
+        inform(destroy::<A, R> as u32);
+
+        <&mut Self>::describe();
+    }
+}
+
+#[allow(non_snake_case)]
+impl<T, A, R> WasmClosureFnOnce<(&A,), R> for T
+    where T: 'static + FnOnce(&A) -> R,
+          A: RefFromWasmAbi + 'static,
+          R: ReturnWasmAbi + 'static
+{
+    type FnMut = FnMut(&A) -> R;
+
+    fn into_fn_mut(self) -> Box<Self::FnMut> {
+        let mut me = Some(self);
+        Box::new(move |arg| {
+            let me = me.take().expect_throw("FnOnce called more than once");
+            me(arg)
+        })
+    }
+
+    fn into_js_function(self) -> JsValue {
+        use std::rc::Rc;
+        use crate::__rt::WasmRefCell;
+
+        let mut me = Some(self);
+
+        let rc1 = Rc::new(WasmRefCell::new(None));
+        let rc2 = rc1.clone();
+
+        let closure = Closure::wrap(Box::new(move |arg: &A| {
+            // Invoke ourself and get the result.
+            let me = me.take().expect_throw("FnOnce called more than once");
+            let result = me(arg);
+
+            // And then drop the `Rc` holding this function's `Closure`
+            // alive.
+            debug_assert_eq!(Rc::strong_count(&rc2), 1);
+            let option_closure = rc2.borrow_mut().take();
+            debug_assert!(option_closure.is_some());
+            drop(option_closure);
+
+            result
+        }) as Box<FnMut(&A) -> R>);
+
+        let js_val = closure.as_ref().clone();
+
+        *rc1.borrow_mut() = Some(closure);
+        debug_assert_eq!(Rc::strong_count(&rc1), 2);
+        drop(rc1);
+
+        js_val
+    }
+}
diff --git a/src/convert/closures.rs b/src/convert/closures.rs
@@ -2,6 +2,7 @@ use core::mem;
 
 use crate::convert::slices::WasmSlice;
 use crate::convert::{FromWasmAbi, GlobalStack, IntoWasmAbi, ReturnWasmAbi, Stack};
+use crate::convert::RefFromWasmAbi;
 use crate::describe::{inform, WasmDescribe, FUNCTION};
 use crate::throw_str;
 
@@ -117,3 +118,97 @@ stack_closures! {
     (6 invoke6 invoke6_mut A B C D E F)
     (7 invoke7 invoke7_mut A B C D E F G)
 }
+
+impl<'a, 'b, A, R> IntoWasmAbi for &'a (Fn(&A) -> R + 'b)
+    where A: RefFromWasmAbi,
+          R: ReturnWasmAbi
+{
+    type Abi = WasmSlice;
+
+    fn into_abi(self, _extra: &mut Stack) -> WasmSlice {
+        unsafe {
+            let (a, b): (usize, usize) = mem::transmute(self);
+            WasmSlice { ptr: a as u32, len: b as u32 }
+        }
+    }
+}
+
+#[allow(non_snake_case)]
+unsafe extern "C" fn invoke1_ref<A: RefFromWasmAbi, R: ReturnWasmAbi>(
+    a: usize,
+    b: usize,
+    arg: <A as RefFromWasmAbi>::Abi,
+) -> <R as ReturnWasmAbi>::Abi {
+    if a == 0 {
+        throw_str("closure invoked recursively or destroyed already");
+    }
+    // Scope all local variables before we call `return_abi` to
+    // ensure they're all destroyed as `return_abi` may throw
+    let ret = {
+        let f: &Fn(&A) -> R = mem::transmute((a, b));
+        let mut _stack = GlobalStack::new();
+        let arg = <A as RefFromWasmAbi>::ref_from_abi(arg, &mut _stack);
+        f(&*arg)
+    };
+    ret.return_abi(&mut GlobalStack::new())
+}
+
+impl<'a, A, R> WasmDescribe for Fn(&A) -> R + 'a
+    where A: RefFromWasmAbi,
+          R: ReturnWasmAbi,
+{
+    fn describe() {
+        inform(FUNCTION);
+        inform(invoke1_ref::<A, R> as u32);
+        inform(1);
+        <&A as WasmDescribe>::describe();
+        <R as WasmDescribe>::describe();
+    }
+}
+
+impl<'a, 'b, A, R> IntoWasmAbi for &'a mut (FnMut(&A) -> R + 'b)
+    where A: RefFromWasmAbi,
+          R: ReturnWasmAbi
+{
+    type Abi = WasmSlice;
+
+    fn into_abi(self, _extra: &mut Stack) -> WasmSlice {
+        unsafe {
+            let (a, b): (usize, usize) = mem::transmute(self);
+            WasmSlice { ptr: a as u32, len: b as u32 }
+        }
+    }
+}
+
+#[allow(non_snake_case)]
+unsafe extern "C" fn invoke1_mut_ref<A: RefFromWasmAbi, R: ReturnWasmAbi>(
+    a: usize,
+    b: usize,
+    arg: <A as RefFromWasmAbi>::Abi,
+) -> <R as ReturnWasmAbi>::Abi {
+    if a == 0 {
+        throw_str("closure invoked recursively or destroyed already");
+    }
+    // Scope all local variables before we call `return_abi` to
+    // ensure they're all destroyed as `return_abi` may throw
+    let ret = {
+        let f: &mut FnMut(&A) -> R = mem::transmute((a, b));
+        let mut _stack = GlobalStack::new();
+        let arg = <A as RefFromWasmAbi>::ref_from_abi(arg, &mut _stack);
+        f(&*arg)
+    };
+    ret.return_abi(&mut GlobalStack::new())
+}
+
+impl<'a, A, R> WasmDescribe for FnMut(&A) -> R + 'a
+    where A: RefFromWasmAbi,
+          R: ReturnWasmAbi
+{
+    fn describe() {
+        inform(FUNCTION);
+        inform(invoke1_mut_ref::<A, R> as u32);
+        inform(1);
+        <&A as WasmDescribe>::describe();
+        <R as WasmDescribe>::describe();
+    }
+}
diff --git a/tests/wasm/closures.js b/tests/wasm/closures.js
@@ -113,3 +113,9 @@ exports.calling_it_throws = a => {
 };
 
 exports.call_val = f => f();
+
+exports.pass_reference_first_arg_twice = (a, b, c) => {
+  b(a);
+  c(a);
+  a.free();
+};
diff --git a/tests/wasm/closures.rs b/tests/wasm/closures.rs
@@ -90,6 +90,18 @@ extern "C" {
 
     #[wasm_bindgen(js_name = calling_it_throws)]
     fn call_val_throws(f: &JsValue) -> bool;
+
+    fn pass_reference_first_arg_twice(
+        a: RefFirstArgument,
+        b: &Closure<FnMut(&RefFirstArgument)>,
+        c: &Closure<FnMut(&RefFirstArgument)>,
+    );
+    #[wasm_bindgen(js_name = pass_reference_first_arg_twice)]
+    fn pass_reference_first_arg_twice2(
+        a: RefFirstArgument,
+        b: &mut FnMut(&RefFirstArgument),
+        c: &mut FnMut(&RefFirstArgument),
+    );
 }
 
 #[wasm_bindgen_test]
@@ -439,3 +451,74 @@ fn test_closure_returner() {
         Ok(o)
     }
 }
+
+#[wasm_bindgen]
+pub struct RefFirstArgument {
+    contents: u32,
+}
+
+#[wasm_bindgen_test]
+fn reference_as_first_argument_builds_at_all() {
+    #[wasm_bindgen]
+    extern "C" {
+        fn ref_first_arg1(a: &Fn(&JsValue));
+        fn ref_first_arg2(a: &mut FnMut(&JsValue));
+        fn ref_first_arg3(a: &Closure<Fn(&JsValue)>);
+        fn ref_first_arg4(a: &Closure<FnMut(&JsValue)>);
+        fn ref_first_custom1(a: &Fn(&RefFirstArgument));
+        fn ref_first_custom2(a: &mut FnMut(&RefFirstArgument));
+        fn ref_first_custom3(a: &Closure<Fn(&RefFirstArgument)>);
+        fn ref_first_custom4(a: &Closure<FnMut(&RefFirstArgument)>);
+    }
+
+    Closure::wrap(Box::new(|_: &JsValue| ()) as Box<Fn(&JsValue)>);
+    Closure::wrap(Box::new(|_: &JsValue| ()) as Box<FnMut(&JsValue)>);
+    Closure::once(|_: &JsValue| ());
+    Closure::once_into_js(|_: &JsValue| ());
+    Closure::wrap(Box::new(|_: &RefFirstArgument| ()) as Box<Fn(&RefFirstArgument)>);
+    Closure::wrap(Box::new(|_: &RefFirstArgument| ()) as Box<FnMut(&RefFirstArgument)>);
+    Closure::once(|_: &RefFirstArgument| ());
+    Closure::once_into_js(|_: &RefFirstArgument| ());
+}
+
+#[wasm_bindgen_test]
+fn reference_as_first_argument_works() {
+    let a = Rc::new(Cell::new(0));
+    let b = {
+        let a = a.clone();
+        Closure::once(move |x: &RefFirstArgument| {
+            assert_eq!(a.get(), 0);
+            assert_eq!(x.contents, 3);
+            a.set(a.get() + 1);
+        })
+    };
+    let c = {
+        let a = a.clone();
+        Closure::once(move |x: &RefFirstArgument| {
+            assert_eq!(a.get(), 1);
+            assert_eq!(x.contents, 3);
+            a.set(a.get() + 1);
+        })
+    };
+    pass_reference_first_arg_twice(RefFirstArgument { contents: 3 }, &b, &c);
+    assert_eq!(a.get(), 2);
+}
+
+#[wasm_bindgen_test]
+fn reference_as_first_argument_works2() {
+    let a = Cell::new(0);
+    pass_reference_first_arg_twice2(
+        RefFirstArgument { contents: 3 },
+        &mut |x: &RefFirstArgument| {
+            assert_eq!(a.get(), 0);
+            assert_eq!(x.contents, 3);
+            a.set(a.get() + 1);
+        },
+        &mut |x: &RefFirstArgument| {
+            assert_eq!(a.get(), 1);
+            assert_eq!(x.contents, 3);
+            a.set(a.get() + 1);
+        },
+    );
+    assert_eq!(a.get(), 2);
+}
