Removing duplicate closure wrappers in the JS glue (#2002)

* Removing duplicate closure wrappers in the JS glue

* Fixing build error

* Adding in explanatory comment

Author: Pauan

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<Descriptor>,
     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,

crates/cli-support/src/descriptors.rs

@@ -22,6 +22,7 @@ use wasm_bindgen_wasm_interpreter::Interpreter;
 pub struct WasmBindgenDescriptorsSection {
     pub descriptors: HashMap<String, Descriptor>,
     pub closure_imports: HashMap<ImportId, Closure>,
+    cached_closures: HashMap<Descriptor, FunctionId>,
 }
 
 pub type WasmBindgenDescriptorsSectionId = TypedCustomSectionId<WasmBindgenDescriptorsSection>;
@@ -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(());
 

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::<Vec<_>>()
-                    .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) => {

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.