Fix anyref closure transformations

* Catch all closures by walking all `Descriptor` values and looking for
  either `Function` or `Closure`.
* Update the correct arguments for wasm by ensuring that the closure
  modifications skip the first two arguments.

crates/cli-support/src/anyref.rs

@@ -1,30 +1,33 @@
-use crate::descriptor::Function;
+use crate::descriptor::{Closure, Descriptor, Function};
-use crate::webidl::{ImportBinding, WasmBindgenAux, WebidlCustomSection, AuxImport};
+use crate::webidl::{AuxImport, ImportBinding, WasmBindgenAux, WebidlCustomSection};
 use failure::Error;
 use std::collections::HashSet;
 use walrus::Module;
+use wasm_bindgen_anyref_xform::Context;
 
 pub fn process(module: &mut Module) -> Result<(), Error> {
-    let mut cfg = wasm_bindgen_anyref_xform::Context::default();
+    let mut cfg = Context::default();
     cfg.prepare(module)?;
     let bindings = module
         .customs
-        .get_typed::<WebidlCustomSection>()
+        .get_typed_mut::<WebidlCustomSection>()
         .expect("webidl custom section should exist");
 
-    for (export, binding) in bindings.exports.iter() {
+    for (export, binding) in bindings.exports.iter_mut() {
-        let (args, ret) = extract_anyrefs(binding);
+        let (args, ret) = extract_anyrefs(binding, 0);
         cfg.export_xform(*export, &args, ret);
+        process_closure_arguments(&mut cfg, binding);
     }
 
-    for (import, kind) in bindings.imports.iter() {
+    for (import, kind) in bindings.imports.iter_mut() {
         let binding = match kind {
             ImportBinding::Function(f) => f,
             ImportBinding::Constructor(f) => f,
             ImportBinding::Method(f) => f,
         };
-        let (args, ret) = extract_anyrefs(binding);
+        let (args, ret) = extract_anyrefs(binding, 0);
         cfg.import_xform(*import, &args, ret);
+        process_closure_arguments(&mut cfg, binding);
     }
 
     let aux = module
@@ -32,13 +35,9 @@ pub fn process(module: &mut Module) -> Result<(), Error> {
         .get_typed_mut::<WasmBindgenAux>()
         .expect("webidl custom section should exist");
     for import in aux.import_map.values_mut() {
-        let closure = match import {
+        match import {
-            AuxImport::Closure(f) => f,
+            AuxImport::Closure(f) => process_closure(&mut cfg, f),
-            _ => continue,
+            _ => {}
-        };
-        let (args, ret) = extract_anyrefs(&closure.function);
-        if let Some(new) = cfg.table_element_xform(closure.shim_idx, &args, ret) {
-            closure.shim_idx = new;
         }
     }
 
@@ -69,9 +68,62 @@ pub fn process(module: &mut Module) -> Result<(), Error> {
     Ok(())
 }
 
-fn extract_anyrefs(f: &Function) -> (Vec<(usize, bool)>, bool) {
+/// Process the `function` provided to ensure that all references to `Closure`
+/// descriptors are processed below.
+fn process_closure_arguments(cfg: &mut Context, function: &mut Function) {
+    for arg in function.arguments.iter_mut() {
+        process_descriptor(cfg, arg);
+    }
+    process_descriptor(cfg, &mut function.ret);
+
+    fn process_descriptor(cfg: &mut Context, descriptor: &mut Descriptor) {
+        match descriptor {
+            Descriptor::Ref(d)
+            | Descriptor::RefMut(d)
+            | Descriptor::Option(d)
+            | Descriptor::Slice(d)
+            | Descriptor::Clamped(d)
+            | Descriptor::Vector(d) => process_descriptor(cfg, d),
+            Descriptor::Closure(c) => process_closure(cfg, c),
+            Descriptor::Function(c) => process_function(cfg, c),
+            _ => {}
+        }
+    }
+
+    fn process_function(cfg: &mut Context, function: &mut Function) {
+        let (args, ret) = extract_anyrefs(&function, 2);
+        if let Some(new) = cfg.table_element_xform(function.shim_idx, &args, ret) {
+            function.shim_idx = new;
+        }
+        process_closure_arguments(cfg, function);
+    }
+}
+
+/// Ensure that the `Closure` is processed in case any of its arguments
+/// recursively contain `anyref` and such.
+fn process_closure(cfg: &mut Context, closure: &mut Closure) {
+    let (args, ret) = extract_anyrefs(&closure.function, 2);
+    if let Some(new) = cfg.table_element_xform(closure.shim_idx, &args, ret) {
+        closure.shim_idx = new;
+    }
+    process_closure_arguments(cfg, &mut closure.function);
+}
+
+/// Extract a description of the anyref arguments from the function signature
+/// described by `f`.
+///
+/// The returned values are expected to be passed to the anyref transformation
+/// pass, and indicate which arguments (by index) in the wasm signature should
+/// be transformed from `i32` to `anyref` as well as whether the returned value
+/// is an `anyref` or not.
+///
+/// The `offset` argument here is typically 0 and indicates the offset at which
+/// the wasm abi arguments described by `f` start at. For closures this is 2
+/// because two synthetic arguments are injected into the wasm signature which
+/// aren't present in the `Function` signature.
+fn extract_anyrefs(f: &Function, offset: usize) -> (Vec<(usize, bool)>, bool) {
     let mut args = Vec::new();
-    let mut cur = 0;
+    let mut cur = offset;
     if f.ret.abi_returned_through_pointer() {
         cur += 1;
     }

crates/cli-support/src/js/rust2js.rs

@@ -1178,6 +1178,7 @@ impl<'a, 'b> Rust2Js<'a, 'b> {
                     //
                     // But for now, we just bounce wasm -> js -> wasm because it is
                     // easy.
+                    self.cx.require_internal_export("__wbindgen_anyref_heap_live_count_impl")?;
                     "wasm.__wbindgen_anyref_heap_live_count_impl()".into()
                 } else {
                     self.cx.expose_global_heap();

src/anyref.rs

@@ -109,7 +109,7 @@ impl Slab {
                 None => internal_error("slot out of bounds"),
             };
         }
-        self.data.len() as u32 - free_count - super::JSIDX_RESERVED
+        self.data.len() as u32 - free_count
     }
 }
 

src/lib.rs

@@ -553,7 +553,7 @@ impl Drop for JsValue {
     fn drop(&mut self) {
         unsafe {
             // We definitely should never drop anything in the stack area
-            debug_assert!(self.idx >= JSIDX_OFFSET);
+            debug_assert!(self.idx >= JSIDX_OFFSET, "free of stack slot {}", self.idx);
 
             // Otherwise if we're not dropping one of our reserved values,
             // actually call the intrinsic. See #1054 for eventually removing