Merge pull request #879 from alexcrichton/closure-zst

Improve codegen for `Closure<T>`

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

@@ -327,9 +327,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;
                 }
                 ",
             ))
@@ -337,13 +341,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); }
+}