diff --git a/crates/cli-support/src/js/closures.rs b/crates/cli-support/src/js/closures.rs index eaa7a145a..54a354061 100644 --- a/crates/cli-support/src/js/closures.rs +++ b/crates/cli-support/src/js/closures.rs @@ -47,6 +47,7 @@ pub fn rewrite(input: &mut Context) -> Result<(), Error> { info.delete_function_table_entries(input); info.inject_imports(input)?; + info.rewrite_calls(input); Ok(()) } @@ -63,7 +64,13 @@ struct ClosureDescriptors { /// /// This map is later used to replace all calls to the keys of this map with /// calls to the value of the map. - code_idx_to_descriptor: BTreeMap, + code_idx_to_descriptor: BTreeMap, +} + +struct DescribeInstruction { + new_idx: u32, + instr_idx: usize, + descriptor: Descriptor, } impl ClosureDescriptors { @@ -93,19 +100,19 @@ impl ClosureDescriptors { None => return Default::default(), }; for (i, function) in code.bodies().iter().enumerate() { - let mut call_found = false; - for instruction in function.code().elements() { - match instruction { - Instruction::Call(idx) if *idx == wbindgen_describe_closure => { - call_found = true; - break + let call_pos = function.code() + .elements() + .iter() + .position(|i| { + match i { + Instruction::Call(i) => *i == wbindgen_describe_closure, + _ => false, } - _ => {} - } - } - if !call_found { - continue - } + }); + let call_pos = match call_pos { + Some(i) => i, + None => continue, + }; let descriptor = input.interpreter.interpret_closure_descriptor( i, input.module, @@ -118,7 +125,11 @@ impl ClosureDescriptors { let new_idx = (ret.code_idx_to_descriptor.len() + imports) as u32; ret.code_idx_to_descriptor.insert( i as u32, - (new_idx, Descriptor::decode(descriptor)), + DescribeInstruction { + new_idx, + instr_idx: call_pos, + descriptor: Descriptor::decode(descriptor), + }, ); } return ret @@ -189,23 +200,26 @@ impl ClosureDescriptors { /// described by the fields internally. These new imports will be closure /// factories and are freshly generated shim in JS. fn inject_imports(&self, input: &mut Context) -> Result<(), Error> { + let wbindgen_describe_closure = match input.interpreter.describe_closure_idx() { + Some(i) => i, + None => return Ok(()), + }; + // We'll be injecting new imports and we'll need to give them all a // type. The signature is all `(i32, i32) -> i32` currently and we know // that this signature already exists in the module as it's the // signature of our `#[inline(never)]` functions. Find the type // signature index so we can assign it below. - let type_idx = input.module.type_section() - .unwrap() - .types() - .iter() - .position(|ty| { - let fnty = match ty { - Type::Function(f) => f, - }; - fnty.params() == &[ValueType::I32, ValueType::I32] && - fnty.return_type() == Some(ValueType::I32) - }) - .unwrap(); + let type_idx = { + let kind = input.module.import_section() + .unwrap() + .entries()[wbindgen_describe_closure as usize] + .external(); + match kind { + External::Function(i) => *i, + _ => unreachable!(), + } + }; // The last piece of the magic. For all our descriptors we found we // inject a JS shim for the descriptor. This JS shim will manufacture a @@ -214,10 +228,10 @@ impl ClosureDescriptors { // Once all that's said and done we inject a new import into the wasm module // of our new wrapper, and the `Remap` step above already wrote calls to // this function within the module. - for (i, (_new_idx, descriptor)) in self.code_idx_to_descriptor.iter() { + for (i, instr) in self.code_idx_to_descriptor.iter() { let import_name = format!("__wbindgen_closure_wrapper{}", i); - let closure = descriptor.closure().unwrap(); + let closure = instr.descriptor.closure().unwrap(); let (js, _ts, _js_doc) = { let mut builder = Js2Rust::new("", input); @@ -237,7 +251,7 @@ impl ClosureDescriptors { input.expose_add_heap_object(); input.function_table_needed = true; let body = format!( - "function(ptr, f) {{ + "function(ptr, f, _ignored) {{ let cb = {}; cb.f = wasm.__wbg_function_table.get(f); cb.a = ptr; @@ -261,10 +275,34 @@ impl ClosureDescriptors { } Ok(()) } + + /// The final step, rewriting calls to `__wbindgen_describe_closure` to the + /// imported functions + fn rewrite_calls(&self, input: &mut Context) { + // FIXME: Ok so this is a bit sketchy in that it introduces overhead. + // What we're doing is taking a our #[inline(never)] shim and *not* + // removing it, only switching the one function that it calls internally. + // + // This isn't great because now we have this non-inlined function which + // would certainly benefit from getting inlined. It's a tiny function + // though and surrounded by allocation so it's probably not a huge + // problem in the long run. Note that `wasm-opt` also implements + // inlining, so we can likely rely on that too. + // + // Still though, it'd be great to not only delete calls to + // `__wbindgen_describe_closure`, it'd be great to remove all of the + // `breaks_if_inlined` functions entirely. + let code = input.module.code_section_mut().unwrap(); + for (i, instr) in self.code_idx_to_descriptor.iter() { + let func = &mut code.bodies_mut()[*i as usize]; + let new_instr = Instruction::Call(instr.new_idx); + func.code_mut().elements_mut()[instr.instr_idx] = new_instr; + } + } } struct Remap<'a> { - code_idx_to_descriptor: &'a BTreeMap, + code_idx_to_descriptor: &'a BTreeMap, old_num_imports: u32, } @@ -367,19 +405,8 @@ impl<'a> Remap<'a> { if *idx < self.old_num_imports { return false } - let code_idx = *idx - self.old_num_imports; - - // If this `idx` points to a function which was effectively a descriptor - // function, then we want to re-point it to our imported function which - // is actually the shim factory. - if let Some((new_idx, _)) = self.code_idx_to_descriptor.get(&code_idx) { - *idx = *new_idx; - return true - } - - // And finally, otherwise this is just a normal function reference we - // don't want to touch, but we're injecting imports which shifts all - // function indices. + // ... otherwise we're injecting a number of new imports, so offset + // everything. *idx += self.code_idx_to_descriptor.len() as u32; false } diff --git a/crates/wasm-interpreter/src/lib.rs b/crates/wasm-interpreter/src/lib.rs index 4f9401817..aa28314a7 100644 --- a/crates/wasm-interpreter/src/lib.rs +++ b/crates/wasm-interpreter/src/lib.rs @@ -226,14 +226,22 @@ impl Interpreter { ) -> Option<&[u32]> { // Call the `code_idx` function. This is an internal `#[inline(never)]` // whose code is completely controlled by the `wasm-bindgen` crate, so - // it should take two arguments and return one (all of which we don't - // care about here). What we're interested in is that while executing - // this function it'll call `__wbindgen_describe_closure` with an - // argument that we look for. + // it should take some arguments (the number of arguments depends on the + // optimization level) and return one (all of which we don't care about + // here). What we're interested in is that while executing this function + // it'll call `__wbindgen_describe_closure` with an argument that we + // look for. assert!(self.descriptor_table_idx.is_none()); let closure_descriptor_idx = (code_idx + self.imports) as u32; - self.stack.push(0); - self.stack.push(0); + + let code_sig = sections.functions.entries()[code_idx].type_ref(); + let function_ty = match §ions.types.types()[code_sig as usize] { + Type::Function(t) => t, + }; + for _ in 0..function_ty.params().len() { + self.stack.push(0); + } + self.call(closure_descriptor_idx, sections); assert_eq!(self.stack.len(), 1); self.stack.pop(); @@ -346,8 +354,8 @@ impl Interpreter { } else if Some(*idx) == self.describe_closure_idx { self.descriptor_table_idx = Some(self.stack.pop().unwrap() as u32); - assert_eq!(self.stack.pop(), Some(0)); - assert_eq!(self.stack.pop(), Some(0)); + self.stack.pop(); + self.stack.pop(); self.stack.push(0); } else { self.call(*idx, sections);