diff --git a/crates/cli-support/src/descriptor.rs b/crates/cli-support/src/descriptor.rs index 70ef467809b..e83bc7317f8 100644 --- a/crates/cli-support/src/descriptor.rs +++ b/crates/cli-support/src/descriptor.rs @@ -39,7 +39,7 @@ tys! { CLAMPED } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub enum Descriptor { I8, U8, @@ -69,14 +69,14 @@ pub enum Descriptor { Unit, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Function { pub arguments: Vec, pub shim_idx: u32, pub ret: Descriptor, } -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Closure { pub shim_idx: u32, pub dtor_idx: u32, diff --git a/crates/cli-support/src/descriptors.rs b/crates/cli-support/src/descriptors.rs index 74a7938037f..9fbe1915366 100644 --- a/crates/cli-support/src/descriptors.rs +++ b/crates/cli-support/src/descriptors.rs @@ -22,6 +22,7 @@ use wasm_bindgen_wasm_interpreter::Interpreter; pub struct WasmBindgenDescriptorsSection { pub descriptors: HashMap, pub closure_imports: HashMap, + cached_closures: HashMap, } pub type WasmBindgenDescriptorsSectionId = TypedCustomSectionId; @@ -126,10 +127,20 @@ impl WasmBindgenDescriptorsSection { // ourselves, and then we're good to go. let ty = module.funcs.get(wbindgen_describe_closure).ty(); for (func, descriptor) in func_to_descriptor { - let import_name = format!("__wbindgen_closure_wrapper{}", func.index()); - let (id, import_id) = - module.add_import_func("__wbindgen_placeholder__", &import_name, ty); - module.funcs.get_mut(id).name = Some(import_name); + // This uses a cache so that if the same closure exists multiple times it will + // deduplicate it so it only exists once. + let id = match self.cached_closures.get(&descriptor) { + Some(id) => *id, + None => { + let import_name = format!("__wbindgen_closure_wrapper{}", func.index()); + let (id, import_id) = + module.add_import_func("__wbindgen_placeholder__", &import_name, ty); + module.funcs.get_mut(id).name = Some(import_name); + self.closure_imports.insert(import_id, descriptor.clone().unwrap_closure()); + self.cached_closures.insert(descriptor, id); + id + }, + }; let local = match &mut module.funcs.get_mut(func).kind { walrus::FunctionKind::Local(l) => l, @@ -144,8 +155,6 @@ impl WasmBindgenDescriptorsSection { local, entry, ); - self.closure_imports - .insert(import_id, descriptor.unwrap_closure()); } return Ok(()); diff --git a/crates/cli-support/src/js/mod.rs b/crates/cli-support/src/js/mod.rs index 9d07120db8b..b6e9f80dd30 100644 --- a/crates/cli-support/src/js/mod.rs +++ b/crates/cli-support/src/js/mod.rs @@ -1724,6 +1724,86 @@ impl<'a> Context<'a> { ); } + fn expose_make_mut_closure(&mut self) -> Result<(), Error> { + if !self.should_write_global("make_mut_closure") { + return Ok(()); + } + + let table = self.export_function_table()?; + + // For mutable closures they can't be invoked recursively. + // To handle that we swap out the `this.a` pointer with zero + // while we invoke it. If we finish and the closure wasn't + // destroyed, then we put back the pointer so a future + // invocation can succeed. + self.global(&format!( + " + function makeMutClosure(arg0, arg1, dtor, f) {{ + const state = {{ a: arg0, b: arg1, cnt: 1 }}; + const real = (...args) => {{ + // First up with a closure we increment the internal reference + // count. This ensures that the Rust closure environment won't + // be deallocated while we're invoking it. + state.cnt++; + const a = state.a; + state.a = 0; + try {{ + return f(a, state.b, ...args); + }} finally {{ + if (--state.cnt === 0) wasm.{}.get(dtor)(a, state.b); + else state.a = a; + }} + }}; + real.original = state; + return real; + }} + ", + table + )); + + Ok(()) + } + + fn expose_make_closure(&mut self) -> Result<(), Error> { + if !self.should_write_global("make_closure") { + return Ok(()); + } + + let table = self.export_function_table()?; + + // For shared closures they can be invoked recursively so we + // just immediately pass through `this.a`. If we end up + // executing the destructor, however, we clear out the + // `this.a` pointer to prevent it being used again the + // future. + self.global(&format!( + " + function makeClosure(arg0, arg1, dtor, f) {{ + const state = {{ a: arg0, b: arg1, cnt: 1 }}; + const real = (...args) => {{ + // First up with a closure we increment the internal reference + // count. This ensures that the Rust closure environment won't + // be deallocated while we're invoking it. + state.cnt++; + try {{ + return f(state.a, state.b, ...args); + }} finally {{ + if (--state.cnt === 0) {{ + wasm.{}.get(dtor)(state.a, state.b); + state.a = 0; + }} + }} + }}; + real.original = state; + return real; + }} + ", + table + )); + + Ok(()) + } + fn global(&mut self, s: &str) { let s = s.trim(); @@ -2338,73 +2418,36 @@ impl<'a> Context<'a> { dtor, mutable, adapter, - nargs, + nargs: _, } => { assert!(kind == AdapterJsImportKind::Normal); assert!(!variadic); assert_eq!(args.len(), 3); - let arg_names = (0..*nargs) - .map(|i| format!("arg{}", i)) - .collect::>() - .join(", "); - let mut js = format!("({}) => {{\n", arg_names); - // First up with a closure we increment the internal reference - // count. This ensures that the Rust closure environment won't - // be deallocated while we're invoking it. - js.push_str("state.cnt++;\n"); - - let table = self.export_function_table()?; - let dtor = format!("wasm.{}.get({})", table, dtor); + let call = self.adapter_name(*adapter); if *mutable { - // For mutable closures they can't be invoked recursively. - // To handle that we swap out the `this.a` pointer with zero - // while we invoke it. If we finish and the closure wasn't - // destroyed, then we put back the pointer so a future - // invocation can succeed. - js.push_str("const a = state.a;\n"); - js.push_str("state.a = 0;\n"); - js.push_str("try {\n"); - js.push_str(&format!("return {}(a, state.b, {});\n", call, arg_names)); - js.push_str("} finally {\n"); - js.push_str("if (--state.cnt === 0) "); - js.push_str(&dtor); - js.push_str("(a, state.b);\n"); - js.push_str("else state.a = a;\n"); - js.push_str("}\n"); + self.expose_make_mut_closure()?; + + Ok(format!( + "makeMutClosure({arg0}, {arg1}, {dtor}, {call})", + arg0 = &args[0], + arg1 = &args[1], + dtor = dtor, + call = call, + )) + } else { - // For shared closures they can be invoked recursively so we - // just immediately pass through `this.a`. If we end up - // executing the destructor, however, we clear out the - // `this.a` pointer to prevent it being used again the - // future. - js.push_str("try {\n"); - js.push_str(&format!( - "return {}(state.a, state.b, {});\n", - call, arg_names - )); - js.push_str("} finally {\n"); - js.push_str("if (--state.cnt === 0) {\n"); - js.push_str(&dtor); - js.push_str("(state.a, state.b);\n"); - js.push_str("state.a = 0;\n"); - js.push_str("}\n"); - js.push_str("}\n"); + self.expose_make_closure()?; + + Ok(format!( + "makeClosure({arg0}, {arg1}, {dtor}, {call})", + arg0 = &args[0], + arg1 = &args[1], + dtor = dtor, + call = call, + )) } - js.push_str("}\n"); - - prelude.push_str(&format!( - " - const state = {{ a: {arg0}, b: {arg1}, cnt: 1 }}; - const real = {body}; - real.original = state; - ", - body = js, - arg0 = &args[0], - arg1 = &args[1], - )); - Ok("real".to_string()) } AuxImport::StructuralMethod(name) => { diff --git a/crates/cli-support/src/wit/mod.rs b/crates/cli-support/src/wit/mod.rs index 39e372c0592..97c5db27b9a 100644 --- a/crates/cli-support/src/wit/mod.rs +++ b/crates/cli-support/src/wit/mod.rs @@ -159,6 +159,7 @@ impl<'a> Context<'a> { let WasmBindgenDescriptorsSection { descriptors, closure_imports, + .. } = *custom; // Store all the executed descriptors in our own field so we have // access to them while processing programs.