From 8ba41cce6e278a4aa68d3fd3d3637d87a52fe866 Mon Sep 17 00:00:00 2001 From: Alex Crichton Date: Mon, 24 Sep 2018 15:10:58 -0700 Subject: [PATCH] Improve codegen for `Closure` This commit improves the codegen for `Closure`, 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` 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 | 18 ++++-- crates/cli-support/src/js/mod.rs | 16 ++--- crates/wasm-interpreter/src/lib.rs | 2 + src/closure.rs | 92 +++++++++++++++++++-------- src/lib.rs | 4 +- tests/wasm/closures.js | 8 ++- tests/wasm/closures.rs | 61 +++++++++++++++++- 7 files changed, 156 insertions(+), 45 deletions(-) diff --git a/crates/cli-support/src/js/closures.rs b/crates/cli-support/src/js/closures.rs index 77103c165..358889df9 100644 --- a/crates/cli-support/src/js/closures.rs +++ b/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); diff --git a/crates/cli-support/src/js/mod.rs b/crates/cli-support/src/js/mod.rs index 0b4fedd07..91d30bf6b 100644 --- a/crates/cli-support/src/js/mod.rs +++ b/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| { diff --git a/crates/wasm-interpreter/src/lib.rs b/crates/wasm-interpreter/src/lib.rs index 80ffe6291..375215ec7 100644 --- a/crates/wasm-interpreter/src/lib.rs +++ b/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); diff --git a/src/closure.rs b/src/closure.rs index e67c3e897..7ec3b9d2b 100644 --- a/src/closure.rs +++ b/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 { js: ManuallyDrop, - _keep_this_data_alive: Rc>>, + data: ManuallyDrop>, +} + +union FatPtr { + ptr: *mut T, + fields: (usize, usize), } impl Closure @@ -135,9 +138,11 @@ impl Closure /// type. /// /// This is the function where the JS closure is manufactured. - pub fn wrap(t: Box) -> Closure { - let data = Rc::new(UnsafeCell::new(t)); - let ptr = &*data as *const UnsafeCell>; + pub fn wrap(mut data: Box) -> Closure { + assert_eq!(mem::size_of::<*const T>(), mem::size_of::>()); + 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 Closure #[inline(never)] unsafe fn breaks_if_inlined( - ptr: usize, + a: usize, + b: usize, invoke: u32, + destroy: u32, ) -> u32 { - super::__wbindgen_describe_closure(ptr as u32, invoke, describe:: as u32) + super::__wbindgen_describe_closure( + a as u32, + b as u32, + invoke, + destroy, + describe:: as u32, + ) } let idx = unsafe { - breaks_if_inlined::(ptr as usize, T::invoke_fn()) + breaks_if_inlined::(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 Closure if idx != !0 { super::__wbindgen_cb_forget(idx); } - mem::forget(self); } } } @@ -270,7 +282,9 @@ impl Drop for Closure 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 R>>, + a: usize, + b: usize, $($var: <$var as FromWasmAbi>::Abi),* ) -> ::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:: 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 R>>, + a: usize, + b: usize, $($var: <$var as FromWasmAbi>::Abi),* ) -> ::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:: R> { + fields: (a, b) + }.ptr)); + } + destroy::<$($var,)* R> as u32 + } } )*) } diff --git a/src/lib.rs b/src/lib.rs index 3b938dfa1..9a79f12c5 100644 --- a/src/lib.rs +++ b/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; diff --git a/tests/wasm/closures.js b/tests/wasm/closures.js index afdac0e80..e1048bedc 100644 --- a/tests/wasm/closures.js +++ b/tests/wasm/closures.js @@ -90,4 +90,10 @@ exports.string_arguments_call = a => { exports.string_ret_call = a => { assert.strictEqual(a('foo'), 'foobar'); -}; \ No newline at end of file +}; + +let DROP_DURING_CALL = null; +exports.drop_during_call_save = f => { + DROP_DURING_CALL = f; +}; +exports.drop_during_call_call = () => DROP_DURING_CALL(); diff --git a/tests/wasm/closures.rs b/tests/wasm/closures.rs index fd62baf89..f5e1b9c29 100644 --- a/tests/wasm/closures.rs +++ b/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 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 = 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 = 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); } +}