Fix references passed to async functions (#3188)

* Fix references passed to async functions

Fixes #3167

#3167 was happening because references to `JsValue`s work by pushing the JS value onto the JS stack, and then popping it off again when the function returns. This breaks for async functions, because the `JsValue` needs to remain valid for the lifetime of the future, not just the initial synchronous call.

To fix this, I introduced a new descriptor, `LONGREF`, which is used instead of `REF` in async functions and marks that the reference can't be invalidated after the function returns. I then made a `LONGREF` to a `JsValue` behave the same as an owned `JsValue`, while working like a regular `REF` for everything else.

I also had to add a new version of `RefFromWasmAbi`, `LongRefFromWasmAbi`, to accomodate for `LONGREF`s to `JsValue`s requiring that the `JsValue` is destroyed from the Rust side.

While working on that, I also noticed that mutable references to slices were broken in async functions. They work by copying the data from JS to Rust, and then copying it back after the function returns to apply the changes. For async functions, that means that it's copied back before the future is actually finished, and so the changes won't be applied properly. To fix that, I changed it so that the data is copied back from the Rust side when the anchor is dropped, which happens when the future resolves for async functions.

I also bumped the schema version, since this is a breaking change in the interface between the macro and the CLI, even though the schema itself is unchanged.

* Update UI tests

* Use a separate instruction to lower the typed array reference

`Instruction::MutableSliceToMemory` previously manually lowered the reference to the type array using `addHeapObject`, but that broke when reference types were enabled. The most direct solution would have been to check whether reference types were enabled and use `addToExternrefTable` instead if so, but using a separate instruction means that it gets picked up by `wasm-bindgen-externref-xform`, and an `externref` is passed directly to wasm.

Author: Liamolucko

crates/backend/src/codegen.rs

@@ -242,6 +242,16 @@ impl ToTokens for ast::Struct {
                 }
             }
 
+            #[automatically_derived]
+            impl wasm_bindgen::convert::LongRefFromWasmAbi for #name {
+                type Abi = u32;
+                type Anchor = wasm_bindgen::__rt::Ref<'static, #name>;
+
+                unsafe fn long_ref_from_abi(js: Self::Abi) -> Self::Anchor {
+                    <Self as wasm_bindgen::convert::RefFromWasmAbi>::ref_from_abi(js)
+                }
+            }
+
             #[automatically_derived]
             impl wasm_bindgen::convert::OptionIntoWasmAbi for #name {
                 #[inline]
@@ -404,7 +414,7 @@ impl TryToTokens for ast::Export {
 
         let mut argtys = Vec::new();
         for (i, arg) in self.function.arguments.iter().enumerate() {
-            argtys.push(&arg.ty);
+            argtys.push(&*arg.ty);
             let i = i + offset;
             let ident = Ident::new(&format!("arg{}", i), Span::call_site());
             let ty = &arg.ty;
@@ -426,16 +436,31 @@ impl TryToTokens for ast::Export {
                     });
                 }
                 syn::Type::Reference(syn::TypeReference { elem, .. }) => {
-                    args.push(quote! {
-                        #ident: <#elem as wasm_bindgen::convert::RefFromWasmAbi>::Abi
-                    });
-                    arg_conversions.push(quote! {
-                        let #ident = unsafe {
-                            <#elem as wasm_bindgen::convert::RefFromWasmAbi>
-                                ::ref_from_abi(#ident)
-                        };
-                        let #ident = &*#ident;
-                    });
+                    if self.function.r#async {
+                        args.push(quote! {
+                            #ident: <#elem as wasm_bindgen::convert::LongRefFromWasmAbi>::Abi
+                        });
+                        arg_conversions.push(quote! {
+                            let #ident = unsafe {
+                                <#elem as wasm_bindgen::convert::LongRefFromWasmAbi>
+                                    ::long_ref_from_abi(#ident)
+                            };
+                            let #ident = <<#elem as wasm_bindgen::convert::LongRefFromWasmAbi>
+                                ::Anchor as core::borrow::Borrow<#elem>>
+                                ::borrow(&#ident);
+                        });
+                    } else {
+                        args.push(quote! {
+                            #ident: <#elem as wasm_bindgen::convert::RefFromWasmAbi>::Abi
+                        });
+                        arg_conversions.push(quote! {
+                            let #ident = unsafe {
+                                <#elem as wasm_bindgen::convert::RefFromWasmAbi>
+                                    ::ref_from_abi(#ident)
+                            };
+                            let #ident = &*#ident;
+                        });
+                    }
                 }
                 _ => {
                     args.push(quote! {
@@ -550,6 +575,22 @@ impl TryToTokens for ast::Export {
         })
         .to_tokens(into);
 
+        let describe_args: TokenStream = argtys
+            .iter()
+            .map(|ty| match ty {
+                syn::Type::Reference(reference)
+                    if self.function.r#async && reference.mutability.is_none() =>
+                {
+                    let inner = &reference.elem;
+                    quote! {
+                        inform(LONGREF);
+                        <#inner as WasmDescribe>::describe();
+                    }
+                }
+                _ => quote! { <#ty as WasmDescribe>::describe(); },
+            })
+            .collect();
+
         // In addition to generating the shim function above which is what
         // our generated JS will invoke, we *also* generate a "descriptor"
         // shim. This descriptor shim uses the `WasmDescribe` trait to
@@ -573,7 +614,7 @@ impl TryToTokens for ast::Export {
                 inform(FUNCTION);
                 inform(0);
                 inform(#nargs);
-                #(<#argtys as WasmDescribe>::describe();)*
+                #describe_args
                 #describe_ret
             },
             attrs: attrs.clone(),
@@ -659,7 +700,7 @@ impl ToTokens for ast::ImportType {
             const #const_name: () = {
                 use wasm_bindgen::convert::{IntoWasmAbi, FromWasmAbi};
                 use wasm_bindgen::convert::{OptionIntoWasmAbi, OptionFromWasmAbi};
-                use wasm_bindgen::convert::RefFromWasmAbi;
+                use wasm_bindgen::convert::{RefFromWasmAbi, LongRefFromWasmAbi};
                 use wasm_bindgen::describe::WasmDescribe;
                 use wasm_bindgen::{JsValue, JsCast, JsObject};
                 use wasm_bindgen::__rt::core;
@@ -731,6 +772,17 @@ impl ToTokens for ast::ImportType {
                     }
                 }
 
+                impl LongRefFromWasmAbi for #rust_name {
+                    type Abi = <JsValue as LongRefFromWasmAbi>::Abi;
+                    type Anchor = #rust_name;
+
+                    #[inline]
+                    unsafe fn long_ref_from_abi(js: Self::Abi) -> Self::Anchor {
+                        let tmp = <JsValue as LongRefFromWasmAbi>::long_ref_from_abi(js);
+                        #rust_name { obj: tmp.into() }
+                    }
+                }
+
                 // TODO: remove this on the next major version
                 impl From<JsValue> for #rust_name {
                     #[inline]

crates/cli-support/src/descriptor.rs

@@ -28,6 +28,7 @@ tys! {
     STRING
     REF
     REFMUT
+    LONGREF
     SLICE
     VECTOR
     EXTERNREF
@@ -89,7 +90,7 @@ pub struct Closure {
     pub mutable: bool,
 }
 
-#[derive(Clone, Debug)]
+#[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub enum VectorKind {
     I8,
     U8,
@@ -132,6 +133,15 @@ impl Descriptor {
             CLOSURE => Descriptor::Closure(Box::new(Closure::decode(data))),
             REF => Descriptor::Ref(Box::new(Descriptor::_decode(data, clamped))),
             REFMUT => Descriptor::RefMut(Box::new(Descriptor::_decode(data, clamped))),
+            LONGREF => {
+                // This descriptor basically just serves as a macro, where most things
+                // become normal `Ref`s, but long refs to externrefs become owned.
+                let contents = Descriptor::_decode(data, clamped);
+                match contents {
+                    Descriptor::Externref | Descriptor::NamedExternref(_) => contents,
+                    _ => Descriptor::Ref(Box::new(contents)),
+                }
+            }
             SLICE => Descriptor::Slice(Box::new(Descriptor::_decode(data, clamped))),
             VECTOR => Descriptor::Vector(Box::new(Descriptor::_decode(data, clamped))),
             OPTIONAL => Descriptor::Option(Box::new(Descriptor::_decode(data, clamped))),

crates/cli-support/src/intrinsic.rs

@@ -84,6 +84,10 @@ fn opt_i64() -> Descriptor {
     Descriptor::Option(Box::new(Descriptor::I64))
 }
 
+fn slice(contents: Descriptor) -> Descriptor {
+    Descriptor::Ref(Box::new(Descriptor::Slice(Box::new(contents))))
+}
+
 intrinsics! {
     pub enum Intrinsic {
         #[symbol = "__wbindgen_jsval_eq"]
@@ -263,6 +267,9 @@ intrinsics! {
         #[symbol = "__wbindgen_json_serialize"]
         #[signature = fn(ref_externref()) -> String]
         JsonSerialize,
+        #[symbol = "__wbindgen_copy_to_typed_array"]
+        #[signature = fn(slice(U8), ref_externref()) -> Unit]
+        CopyToTypedArray,
         #[symbol = "__wbindgen_externref_heap_live_count"]
         #[signature = fn() -> I32]
         ExternrefHeapLiveCount,

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

@@ -929,15 +929,8 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
             js.push(format!("len{}", i));
         }
 
-        Instruction::MutableSliceToMemory {
-            kind,
-            malloc,
-            mem,
-            free,
-        } => {
-            // First up, pass the JS value into wasm, getting out a pointer and
-            // a length. These two pointer/length values get pushed onto the
-            // value stack.
+        Instruction::MutableSliceToMemory { kind, malloc, mem } => {
+            // Copy the contents of the typed array into wasm.
             let val = js.pop();
             let func = js.cx.pass_to_wasm_function(kind.clone(), *mem)?;
             let malloc = js.cx.export_name_of(*malloc);
@@ -950,25 +943,12 @@ fn instruction(js: &mut JsBuilder, instr: &Instruction, log_error: &mut bool) ->
                 malloc = malloc,
             ));
             js.prelude(&format!("var len{} = WASM_VECTOR_LEN;", i));
+            // Then pass it the pointer and the length of where we copied it.
             js.push(format!("ptr{}", i));
             js.push(format!("len{}", i));
-
-            // Next we set up a `finally` clause which will both update the
-            // original mutable slice with any modifications, and then free the
-            // Rust-backed memory.
-            let free = js.cx.export_name_of(*free);
-            let get = js.cx.memview_function(kind.clone(), *mem);
-            js.finally(&format!(
-                "
-                    {val}.set({get}().subarray(ptr{i} / {size}, ptr{i} / {size} + len{i}));
-                    wasm.{free}(ptr{i}, len{i} * {size});
-                ",
-                val = val,
-                get = get,
-                free = free,
-                size = kind.size(),
-                i = i,
-            ));
+            // Then we give wasm a reference to the original typed array, so that it can
+            // update it with modifications made on the wasm side before returning.
+            js.push(val);
         }
 
         Instruction::BoolFromI32 => {

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

@@ -3498,6 +3498,15 @@ impl<'a> Context<'a> {
                 "JSON.stringify(obj === undefined ? null : obj)".to_string()
             }
 
+            Intrinsic::CopyToTypedArray => {
+                assert_eq!(args.len(), 2);
+                format!(
+                    "new Uint8Array({dst}.buffer, {dst}.byteOffset, {dst}.byteLength).set({src})",
+                    src = args[0],
+                    dst = args[1]
+                )
+            }
+
             Intrinsic::ExternrefHeapLiveCount => {
                 assert_eq!(args.len(), 0);
                 self.expose_global_heap();

crates/cli-support/src/wit/incoming.rs

@@ -204,9 +204,13 @@ impl InstructionBuilder<'_, '_> {
                             kind,
                             malloc: self.cx.malloc()?,
                             mem: self.cx.memory()?,
-                            free: self.cx.free()?,
                         },
-                        &[AdapterType::I32, AdapterType::I32],
+                        &[AdapterType::I32, AdapterType::I32, AdapterType::Externref],
+                    );
+                    self.late_instruction(
+                        &[AdapterType::Externref],
+                        Instruction::I32FromExternrefOwned,
+                        &[AdapterType::I32],
                     );
                 } else {
                     self.instruction(
@@ -373,6 +377,27 @@ impl InstructionBuilder<'_, '_> {
         self.output.extend_from_slice(outputs);
     }
 
+    /// Add an instruction whose inputs are the results of previous instructions
+    /// instead of the parameters from JS / results from Rust.
+    pub fn late_instruction(
+        &mut self,
+        inputs: &[AdapterType],
+        instr: Instruction,
+        outputs: &[AdapterType],
+    ) {
+        for input in inputs {
+            assert_eq!(self.output.pop().unwrap(), *input);
+        }
+        self.instructions.push(InstructionData {
+            instr,
+            stack_change: StackChange::Modified {
+                popped: inputs.len(),
+                pushed: outputs.len(),
+            },
+        });
+        self.output.extend_from_slice(outputs);
+    }
+
     fn number(&mut self, input: wit_walrus::ValType, output: walrus::ValType) {
         let std = wit_walrus::Instruction::IntToWasm {
             input,

crates/cli-support/src/wit/standard.rs

@@ -65,7 +65,7 @@ pub enum AdapterJsImportKind {
     Normal,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub enum AdapterType {
     S8,
     S16,
@@ -185,7 +185,6 @@ pub enum Instruction {
     MutableSliceToMemory {
         kind: VectorKind,
         malloc: walrus::FunctionId,
-        free: walrus::FunctionId,
         mem: walrus::MemoryId,
     },
 
@@ -469,12 +468,9 @@ impl walrus::CustomSection for NonstandardWitSection {
                         roots.push_memory(mem);
                         roots.push_func(malloc);
                     }
-                    MutableSliceToMemory {
-                        free, malloc, mem, ..
-                    } => {
+                    MutableSliceToMemory { malloc, mem, .. } => {
                         roots.push_memory(mem);
                         roots.push_func(malloc);
-                        roots.push_func(free);
                     }
                     VectorLoad { free, mem, .. }
                     | OptionVectorLoad { free, mem, .. }

crates/macro/ui-tests/missing-catch.stderr

@@ -1,17 +1,17 @@
 error[E0277]: the trait bound `Result<wasm_bindgen::JsValue, wasm_bindgen::JsValue>: FromWasmAbi` is not satisfied
-  --> $DIR/missing-catch.rs:6:9
-   |
-6  |     pub fn foo() -> Result<JsValue, JsValue>;
-   |            ^^^ the trait `FromWasmAbi` is not implemented for `Result<wasm_bindgen::JsValue, wasm_bindgen::JsValue>`
-   |
+ --> ui-tests/missing-catch.rs:6:9
+  |
+6 |     pub fn foo() -> Result<JsValue, JsValue>;
+  |            ^^^ the trait `FromWasmAbi` is not implemented for `Result<wasm_bindgen::JsValue, wasm_bindgen::JsValue>`
+  |
 note: required by a bound in `FromWasmAbi`
-  --> $DIR/traits.rs:23:1
-   |
-23 | / pub trait FromWasmAbi: WasmDescribe {
-24 | |     /// The wasm ABI type that this converts from when coming back out from the
-25 | |     /// ABI boundary.
-26 | |     type Abi: WasmAbi;
-...  |
-35 | |     unsafe fn from_abi(js: Self::Abi) -> Self;
-36 | | }
-   | |_^ required by this bound in `FromWasmAbi`
+ --> $WORKSPACE/src/convert/traits.rs
+  |
+  | / pub trait FromWasmAbi: WasmDescribe {
+  | |     /// The wasm ABI type that this converts from when coming back out from the
+  | |     /// ABI boundary.
+  | |     type Abi: WasmAbi;
+... |
+  | |     unsafe fn from_abi(js: Self::Abi) -> Self;
+  | | }
+  | |_^ required by this bound in `FromWasmAbi`

crates/macro/ui-tests/traits-not-implemented.stderr

@@ -1,18 +1,18 @@
 error[E0277]: the trait bound `A: IntoWasmAbi` is not satisfied
-  --> $DIR/traits-not-implemented.rs:5:1
-   |
-5  | #[wasm_bindgen]
-   | ^^^^^^^^^^^^^^^ the trait `IntoWasmAbi` is not implemented for `A`
-   |
+ --> ui-tests/traits-not-implemented.rs:5:1
+  |
+5 | #[wasm_bindgen]
+  | ^^^^^^^^^^^^^^^ the trait `IntoWasmAbi` is not implemented for `A`
+  |
 note: required by a bound in `IntoWasmAbi`
-  --> $DIR/traits.rs:9:1
-   |
-9  | / pub trait IntoWasmAbi: WasmDescribe {
-10 | |     /// The wasm ABI type that this converts into when crossing the ABI
-11 | |     /// boundary.
-12 | |     type Abi: WasmAbi;
-...  |
-16 | |     fn into_abi(self) -> Self::Abi;
-17 | | }
-   | |_^ required by this bound in `IntoWasmAbi`
-   = note: this error originates in the attribute macro `wasm_bindgen` (in Nightly builds, run with -Z macro-backtrace for more info)
+ --> $WORKSPACE/src/convert/traits.rs
+  |
+  | / pub trait IntoWasmAbi: WasmDescribe {
+  | |     /// The wasm ABI type that this converts into when crossing the ABI
+  | |     /// boundary.
+  | |     type Abi: WasmAbi;
+... |
+  | |     fn into_abi(self) -> Self::Abi;
+  | | }
+  | |_^ required by this bound in `IntoWasmAbi`
+  = note: this error originates in the attribute macro `wasm_bindgen` (in Nightly builds, run with -Z macro-backtrace for more info)

crates/shared/src/lib.rs

@@ -6,7 +6,7 @@ mod schema_hash_approval;
 // This gets changed whenever our schema changes.
 // At this time versions of wasm-bindgen and wasm-bindgen-cli are required to have the exact same
 // SCHEMA_VERSION in order to work together.
-pub const SCHEMA_VERSION: &str = "0.2.83";
+pub const SCHEMA_VERSION: &str = "0.2.84";
 
 #[macro_export]
 macro_rules! shared_api {