Improve codegen for Closure<T>

This commit improves the codegen for `Closure<T>`, primarily for ZST
where the closure doesn't actually capture anything. Previously
`wasm-bindgen` would unconditionally allocate an `Rc` for a fat pointer,
meaning that it would always hit the allocator even when the `Box<T>`
didn't actually contain an allocation. Now the reference count for the
closure is stored on the JS object rather than in Rust.

Some more advanced tests were added along the way to ensure that
functionality didn't regress, and otherwise the calling convention for
`Closure` changed a good deal but should still be the same user-facing.
The primary change was that the reference count reaching zero may cause
JS to need to run the destructor. It simply returns this information in
`Drop for Closure` and otherwise when calling it now also retains a
function pointer that runs the destructor.

Closes #874

crates/cli-support/src/js/closures.rs

@@ -226,26 +226,32 @@ impl ClosureDescriptors {
 
             let (js, _ts, _js_doc) = {
                 let mut builder = Js2Rust::new("", input);
+                builder.prelude("this.cnt++;");
                 if closure.mutable {
                     builder
                         .prelude("let a = this.a;\n")
                         .prelude("this.a = 0;\n")
                         .rust_argument("a")
+                        .rust_argument("b")
                         .finally("this.a = a;\n");
                 } else {
-                    builder.rust_argument("this.a");
+                    builder.rust_argument("this.a")
+                        .rust_argument("b");
                 }
+                builder.finally("if (this.cnt-- == 1) d(this.a, b);");
                 builder
                     .process(&closure.function)?
-                    .finish("function", "this.f")
+                    .finish("function", "f")
             };
             input.expose_add_heap_object();
             input.function_table_needed = true;
             let body = format!(
-                "function(ptr, f, _ignored) {{
-                    let cb = {};
-                    cb.f = wasm.__wbg_function_table.get(f);
-                    cb.a = ptr;
+                "function(a, b, fi, di, _ignored) {{
+                    const f = wasm.__wbg_function_table.get(fi);
+                    const d = wasm.__wbg_function_table.get(di);
+                    const cb = {};
+                    cb.a = a;
+                    cb.cnt = 1;
                     let real = cb.bind(cb);
                     real.original = cb;
                     return addHeapObject(real);

crates/cli-support/src/js/mod.rs

@@ -325,9 +325,13 @@ impl<'a> Context<'a> {
             Ok(String::from(
                 "
                 function(i) {
-                    let obj = getObject(i).original;
-                    obj.a = obj.b = 0;
+                    const obj = getObject(i).original;
                     dropRef(i);
+                    if (obj.cnt-- == 1) {
+                        obj.a = 0;
+                        return 1;
+                    }
+                    return 0;
                 }
                 ",
             ))
@@ -335,13 +339,7 @@ impl<'a> Context<'a> {
 
         self.bind("__wbindgen_cb_forget", &|me| {
             me.expose_drop_ref();
-            Ok(String::from(
-                "
-                function(i) {
-                    dropRef(i);
-                }
-                ",
-            ))
+            Ok("dropRef".to_string())
         })?;
 
         self.bind("__wbindgen_json_parse", &|me| {

crates/wasm-interpreter/src/lib.rs

@@ -347,6 +347,8 @@ impl Interpreter {
                         self.descriptor_table_idx = Some(self.stack.pop().unwrap() as u32);
                         self.stack.pop();
                         self.stack.pop();
+                        self.stack.pop();
+                        self.stack.pop();
                         self.stack.push(0);
                     } else {
                         self.call(*idx, sections);

src/closure.rs

@@ -4,12 +4,10 @@
 //! closures" from Rust to JS. Some more details can be found on the `Closure`
 //! type itself.
 
-use std::cell::UnsafeCell;
 #[cfg(feature = "nightly")]
 use std::marker::Unsize;
 use std::mem::{self, ManuallyDrop};
 use std::prelude::v1::*;
-use std::rc::Rc;
 
 use JsValue;
 use convert::*;
@@ -102,7 +100,12 @@ use throw_str;
 /// ```
 pub struct Closure<T: ?Sized> {
     js: ManuallyDrop<JsValue>,
-    _keep_this_data_alive: Rc<UnsafeCell<Box<T>>>,
+    data: ManuallyDrop<Box<T>>,
+}
+
+union FatPtr<T: ?Sized> {
+    ptr: *mut T,
+    fields: (usize, usize),
 }
 
 impl<T> Closure<T>
@@ -135,9 +138,11 @@ impl<T> Closure<T>
     /// type.
     ///
     /// This is the function where the JS closure is manufactured.
-    pub fn wrap(t: Box<T>) -> Closure<T> {
-        let data = Rc::new(UnsafeCell::new(t));
-        let ptr = &*data as *const UnsafeCell<Box<T>>;
+    pub fn wrap(mut data: Box<T>) -> Closure<T> {
+        assert_eq!(mem::size_of::<*const T>(), mem::size_of::<FatPtr<T>>());
+        let (a, b) = unsafe {
+            FatPtr { ptr: &mut *data as *mut T }.fields
+        };
 
         // Here we need to create a `JsValue` with the data and `T::invoke()`
         // function pointer. To do that we... take a few unconventional turns.
@@ -190,19 +195,27 @@ impl<T> Closure<T>
 
         #[inline(never)]
         unsafe fn breaks_if_inlined<T: WasmClosure + ?Sized>(
-            ptr: usize,
+            a: usize,
+            b: usize,
             invoke: u32,
+            destroy: u32,
         ) -> u32 {
-            super::__wbindgen_describe_closure(ptr as u32, invoke, describe::<T> as u32)
+            super::__wbindgen_describe_closure(
+                a as u32,
+                b as u32,
+                invoke,
+                destroy,
+                describe::<T> as u32,
+            )
         }
 
         let idx = unsafe {
-            breaks_if_inlined::<T>(ptr as usize, T::invoke_fn())
+            breaks_if_inlined::<T>(a, b, T::invoke_fn(), T::destroy_fn())
         };
 
         Closure {
             js: ManuallyDrop::new(JsValue { idx }),
-            _keep_this_data_alive: data,
+            data: ManuallyDrop::new(data),
         }
     }
 
@@ -223,7 +236,6 @@ impl<T> Closure<T>
             if idx != !0 {
                 super::__wbindgen_cb_forget(idx);
             }
-            mem::forget(self);
         }
     }
 }
@@ -270,7 +282,9 @@ impl<T> Drop for Closure<T>
         unsafe {
             // this will implicitly drop our strong reference in addition to
             // invalidating all future invocations of the closure
-            super::__wbindgen_cb_drop(self.js.idx);
+            if super::__wbindgen_cb_drop(self.js.idx) != 0 {
+                ManuallyDrop::drop(&mut self.data);
+            }
         }
     }
 }
@@ -284,6 +298,7 @@ pub unsafe trait WasmClosure: 'static {
     fn describe();
 
     fn invoke_fn() -> u32;
+    fn destroy_fn() -> u32;
 }
 
 // The memory safety here in these implementations below is a bit tricky. We
@@ -315,30 +330,42 @@ macro_rules! doit {
             fn invoke_fn() -> u32 {
                 #[allow(non_snake_case)]
                 unsafe extern fn invoke<$($var: FromWasmAbi,)* R: ReturnWasmAbi>(
-                    a: *const UnsafeCell<Box<Fn($($var),*) -> R>>,
+                    a: usize,
+                    b: usize,
                     $($var: <$var as FromWasmAbi>::Abi),*
                 ) -> <R as ReturnWasmAbi>::Abi {
-                    if a.is_null() {
+                    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 a = Rc::from_raw(a);
-                        let my_handle = a.clone();
-                        drop(Rc::into_raw(a));
-                        let f: &Fn($($var),*) -> R = &**my_handle.get();
+                        let f: *const Fn($($var),*) -> R =
+                            FatPtr { fields: (a, b) }.ptr;
                         let mut _stack = GlobalStack::new();
                         $(
                             let $var = <$var as FromWasmAbi>::from_abi($var, &mut _stack);
                         )*
-                        f($($var),*)
+                        (*f)($($var),*)
                     };
                     ret.return_abi(&mut GlobalStack::new())
                 }
                 invoke::<$($var,)* R> as u32
             }
+
+            fn destroy_fn() -> u32 {
+                unsafe extern fn destroy<$($var: FromWasmAbi,)* R: ReturnWasmAbi>(
+                    a: usize,
+                    b: usize,
+                ) {
+                    debug_assert!(a != 0);
+                    drop(Box::from_raw(FatPtr::<Fn($($var,)*) -> R> {
+                        fields: (a, b)
+                    }.ptr));
+                }
+                destroy::<$($var,)* R> as u32
+            }
         }
 
         unsafe impl<$($var,)* R> WasmClosure for FnMut($($var),*) -> R
@@ -352,30 +379,43 @@ macro_rules! doit {
             fn invoke_fn() -> u32 {
                 #[allow(non_snake_case)]
                 unsafe extern fn invoke<$($var: FromWasmAbi,)* R: ReturnWasmAbi>(
-                    a: *const UnsafeCell<Box<FnMut($($var),*) -> R>>,
+                    a: usize,
+                    b: usize,
                     $($var: <$var as FromWasmAbi>::Abi),*
                 ) -> <R as ReturnWasmAbi>::Abi {
-                    if a.is_null() {
+                    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 a = Rc::from_raw(a);
-                        let my_handle = a.clone();
-                        drop(Rc::into_raw(a));
-                        let f: &mut FnMut($($var),*) -> R = &mut **my_handle.get();
+                        let f: *const FnMut($($var),*) -> R =
+                            FatPtr { fields: (a, b) }.ptr;
+                        let f = f as *mut FnMut($($var),*) -> R;
                         let mut _stack = GlobalStack::new();
                         $(
                             let $var = <$var as FromWasmAbi>::from_abi($var, &mut _stack);
                         )*
-                        f($($var),*)
+                        (*f)($($var),*)
                     };
                     ret.return_abi(&mut GlobalStack::new())
                 }
                 invoke::<$($var,)* R> as u32
             }
+
+            fn destroy_fn() -> u32 {
+                unsafe extern fn destroy<$($var: FromWasmAbi,)* R: ReturnWasmAbi>(
+                    a: usize,
+                    b: usize,
+                ) {
+                    debug_assert!(a != 0);
+                    drop(Box::from_raw(FatPtr::<FnMut($($var,)*) -> R> {
+                        fields: (a, b)
+                    }.ptr));
+                }
+                destroy::<$($var,)* R> as u32
+            }
         }
     )*)
 }

src/lib.rs

@@ -445,11 +445,11 @@ externs! {
     fn __wbindgen_throw(a: *const u8, b: usize) -> !;
     fn __wbindgen_rethrow(a: u32) -> !;
 
-    fn __wbindgen_cb_drop(idx: u32) -> ();
+    fn __wbindgen_cb_drop(idx: u32) -> u32;
     fn __wbindgen_cb_forget(idx: u32) -> ();
 
     fn __wbindgen_describe(v: u32) -> ();
-    fn __wbindgen_describe_closure(a: u32, b: u32, c: u32) -> u32;
+    fn __wbindgen_describe_closure(a: u32, b: u32, c: u32, d: u32, e: u32) -> u32;
 
     fn __wbindgen_json_parse(ptr: *const u8, len: usize) -> u32;
     fn __wbindgen_json_serialize(idx: u32, ptr: *mut *mut u8) -> usize;

tests/wasm/closures.js

@@ -90,4 +90,10 @@ exports.string_arguments_call = a => {
 
 exports.string_ret_call = a => {
     assert.strictEqual(a('foo'), 'foobar');
-};
+};
+
+let DROP_DURING_CALL = null;
+exports.drop_during_call_save = f => {
+  DROP_DURING_CALL = f;
+};
+exports.drop_during_call_call = () => DROP_DURING_CALL();

tests/wasm/closures.rs

@@ -1,6 +1,6 @@
 #![cfg(feature = "nightly")]
 
-use std::cell::Cell;
+use std::cell::{Cell, RefCell};
 use std::rc::Rc;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen_test::*;
@@ -60,6 +60,9 @@ extern "C" {
     fn string_arguments_call(a: &mut FnMut(String));
 
     fn string_ret_call(a: &mut FnMut(String) -> String);
+
+    fn drop_during_call_save(a: &Closure<Fn()>);
+    fn drop_during_call_call();
 }
 
 #[wasm_bindgen_test]
@@ -206,3 +209,59 @@ fn string_ret() {
     });
     assert!(x);
 }
+
+#[wasm_bindgen_test]
+fn drop_drops() {
+    static mut HIT: bool = false;
+
+    struct A;
+
+    impl Drop for A {
+        fn drop(&mut self) {
+            unsafe { HIT = true; }
+        }
+    }
+    let a = A;
+    let x: Closure<Fn()> = Closure::new(move || drop(&a));
+    drop(x);
+    unsafe { assert!(HIT); }
+}
+
+#[wasm_bindgen_test]
+fn drop_during_call_ok() {
+    static mut HIT: bool = false;
+    struct A;
+    impl Drop for A {
+        fn drop(&mut self) {
+            unsafe { HIT = true; }
+        }
+    }
+
+    let rc = Rc::new(RefCell::new(None));
+    let rc2 = rc.clone();
+    let x = 3;
+    let a = A;
+    let x: Closure<Fn()> = Closure::new(move || {
+        // "drop ourselves"
+        drop(rc2.borrow_mut().take().unwrap());
+
+        // `A` should not have been destroyed as a result
+        unsafe { assert!(!HIT); }
+
+        // allocate some heap memory to try to paper over our `3`
+        drop(String::from("1234567890"));
+
+        // make sure our closure memory is still valid
+        assert_eq!(x, 3);
+
+        // make sure `A` is bound to our closure environment.
+        drop(&a);
+        unsafe { assert!(!HIT); }
+    });
+    drop_during_call_save(&x);
+    *rc.borrow_mut() = Some(x);
+    drop(rc);
+    unsafe { assert!(!HIT); }
+    drop_during_call_call();
+    unsafe { assert!(HIT); }
+}