diff --git a/crates/cli-support/src/js/closures.rs b/crates/cli-support/src/js/closures.rs index 77103c16..358889df 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 710a0b5c..57490dd2 100644 --- a/crates/cli-support/src/js/mod.rs +++ b/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| { diff --git a/crates/wasm-interpreter/src/lib.rs b/crates/wasm-interpreter/src/lib.rs index 80ffe629..375215ec 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 e67c3e89..7ec3b9d2 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 3b938dfa..9a79f12c 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 afdac0e8..e1048bed 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 fd62baf8..f5e1b9c2 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); } +}