Add support for linked modules (#3069)

* Add support for linked modules

* Use `wasm_bindgen::link_to!` in wasm-bindgen-futures

* Fix tests

* Update schema

* Return `String` instead of `Result<String, JsValue>`

* Add documentation

* Add tests

* Refactor: Return Diagnostic from module_from_opts

* Refactor: Use Option::filter

* Fix inline_js offsets

* Fully-qualified names in quote

* Return absolute URLs and add node tests

* Fix message

* Enable compile doctest for example

* Disallow module paths in `link_to!`

* Fix documentation

* Opt-in to linked modules in wasm-bindgen with `--allow-links`

* Fix tests

* Support embedding for local modules

* Remove linked module embed limit

* Fix escaping

* Add and refer to the documentation

* Rename option and improve documentation

* Improve documentation

* Update crates/macro/src/lib.rs

Co-authored-by: Liam Murphy <liampm32@gmail.com>

* Add paragraph break in docs

crates/backend/src/ast.rs

@@ -2,7 +2,7 @@
 //! with all the added metadata necessary to generate WASM bindings
 //! for it.
 
-use crate::Diagnostic;
+use crate::{util::ShortHash, Diagnostic};
 use proc_macro2::{Ident, Span};
 use std::hash::{Hash, Hasher};
 use wasm_bindgen_shared as shared;
@@ -16,6 +16,8 @@ pub struct Program {
     pub exports: Vec<Export>,
     /// js -> rust interfaces
     pub imports: Vec<Import>,
+    /// linked-to modules
+    pub linked_modules: Vec<ImportModule>,
     /// rust enums
     pub enums: Vec<Enum>,
     /// rust structs
@@ -36,8 +38,23 @@ impl Program {
             && self.typescript_custom_sections.is_empty()
             && self.inline_js.is_empty()
     }
+
+    /// Name of the link function for a specific linked module
+    pub fn link_function_name(&self, idx: usize) -> String {
+        let hash = match &self.linked_modules[idx] {
+            ImportModule::Inline(idx, _) => ShortHash((1, &self.inline_js[*idx])).to_string(),
+            other => ShortHash((0, other)).to_string(),
+        };
+        format!("__wbindgen_link_{}", hash)
+    }
 }
 
+/// An abstract syntax tree representing a link to a module in Rust.
+/// In contrast to Program, LinkToModule must expand to an expression.
+/// linked_modules of the inner Program must contain exactly one element
+/// whose link is produced by the expression.
+pub struct LinkToModule(pub Program);
+
 /// A rust to js interface. Allows interaction with rust objects/functions
 /// from javascript.
 #[cfg_attr(feature = "extra-traits", derive(Debug))]

crates/backend/src/codegen.rs

@@ -129,6 +129,29 @@ impl TryToTokens for ast::Program {
     }
 }
 
+impl TryToTokens for ast::LinkToModule {
+    fn try_to_tokens(&self, tokens: &mut TokenStream) -> Result<(), Diagnostic> {
+        let mut program = TokenStream::new();
+        self.0.try_to_tokens(&mut program)?;
+        let link_function_name = self.0.link_function_name(0);
+        let name = Ident::new(&link_function_name, Span::call_site());
+        let abi_ret = quote! { <std::string::String as wasm_bindgen::convert::FromWasmAbi>::Abi };
+        let extern_fn = extern_fn(&name, &[], &[], &[], abi_ret);
+        (quote! {
+            {
+                #program
+                #extern_fn
+
+                unsafe {
+                    <std::string::String as wasm_bindgen::convert::FromWasmAbi>::from_abi(#name())
+                }
+            }
+        })
+        .to_tokens(tokens);
+        Ok(())
+    }
+}
+
 impl ToTokens for ast::Struct {
     fn to_tokens(&self, tokens: &mut TokenStream) {
         let name = &self.rust_name;
@@ -1118,8 +1141,8 @@ impl TryToTokens for ast::ImportFunction {
         let import_name = &self.shim;
         let attrs = &self.function.rust_attrs;
         let arguments = &arguments;
-        let abi_arguments = &abi_arguments;
-        let abi_argument_names = &abi_argument_names;
+        let abi_arguments = &abi_arguments[..];
+        let abi_argument_names = &abi_argument_names[..];
 
         let doc_comment = &self.doc_comment;
         let me = if is_method {
@@ -1144,23 +1167,13 @@ impl TryToTokens for ast::ImportFunction {
         // like rustc itself doesn't do great in that regard so let's just do
         // the best we can in the meantime.
         let extern_fn = respan(
-            quote! {
-                #[cfg(all(target_arch = "wasm32", not(target_os = "emscripten")))]
-                #(#attrs)*
-                #[link(wasm_import_module = "__wbindgen_placeholder__")]
-                extern "C" {
-                    fn #import_name(#(#abi_arguments),*) -> #abi_ret;
-                }
-
-                #[cfg(not(all(target_arch = "wasm32", not(target_os = "emscripten"))))]
-                unsafe fn #import_name(#(#abi_arguments),*) -> #abi_ret {
-                    #(
-                        drop(#abi_argument_names);
-                    )*
-                    panic!("cannot call wasm-bindgen imported functions on \
-                            non-wasm targets");
-                }
-            },
+            extern_fn(
+                import_name,
+                attrs,
+                abi_arguments,
+                abi_argument_names,
+                abi_ret,
+            ),
             &self.rust_name,
         );
 
@@ -1399,6 +1412,32 @@ impl<'a, T: ToTokens> ToTokens for Descriptor<'a, T> {
     }
 }
 
+fn extern_fn(
+    import_name: &Ident,
+    attrs: &[syn::Attribute],
+    abi_arguments: &[TokenStream],
+    abi_argument_names: &[Ident],
+    abi_ret: TokenStream,
+) -> TokenStream {
+    quote! {
+        #[cfg(all(target_arch = "wasm32", not(target_os = "emscripten")))]
+        #(#attrs)*
+        #[link(wasm_import_module = "__wbindgen_placeholder__")]
+        extern "C" {
+            fn #import_name(#(#abi_arguments),*) -> #abi_ret;
+        }
+
+        #[cfg(not(all(target_arch = "wasm32", not(target_os = "emscripten"))))]
+        unsafe fn #import_name(#(#abi_arguments),*) -> #abi_ret {
+            #(
+                drop(#abi_argument_names);
+            )*
+            panic!("cannot call wasm-bindgen imported functions on \
+                    non-wasm targets");
+        }
+    }
+}
+
 /// Converts `span` into a stream of tokens, and attempts to ensure that `input`
 /// has all the appropriate span information so errors in it point to `span`.
 fn respan(input: TokenStream, span: &dyn ToTokens) -> TokenStream {

crates/backend/src/encode.rs

@@ -146,6 +146,12 @@ fn shared_program<'a>(
             .iter()
             .map(|x| -> &'a str { &x })
             .collect(),
+        linked_modules: prog
+            .linked_modules
+            .iter()
+            .enumerate()
+            .map(|(i, a)| shared_linked_module(&prog.link_function_name(i), a, intern))
+            .collect::<Result<Vec<_>, _>>()?,
         local_modules: intern
             .files
             .borrow()
@@ -249,6 +255,17 @@ fn shared_import<'a>(i: &'a ast::Import, intern: &'a Interner) -> Result<Import<
     })
 }
 
+fn shared_linked_module<'a>(
+    name: &str,
+    i: &'a ast::ImportModule,
+    intern: &'a Interner,
+) -> Result<LinkedModule<'a>, Diagnostic> {
+    Ok(LinkedModule {
+        module: shared_module(i, intern)?,
+        link_function_name: intern.intern_str(name),
+    })
+}
+
 fn shared_module<'a>(
     m: &'a ast::ImportModule,
     intern: &'a Interner,

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

@@ -384,7 +384,7 @@ impl<'a> Context<'a> {
                     if (typeof document === 'undefined') {
                         script_src = location.href;
                     } else {
-                        script_src = document.currentScript.src;
+                        script_src = new URL(document.currentScript.src, location.href).toString();
                     }\n",
                 );
                 js.push_str("let wasm;\n");
@@ -3143,6 +3143,41 @@ impl<'a> Context<'a> {
                 assert!(!variadic);
                 self.invoke_intrinsic(intrinsic, args, prelude)
             }
+
+            AuxImport::LinkTo(path, content) => {
+                assert!(kind == AdapterJsImportKind::Normal);
+                assert!(!variadic);
+                assert_eq!(args.len(), 0);
+                if self.config.split_linked_modules {
+                    let base = match self.config.mode {
+                        OutputMode::Web
+                        | OutputMode::Bundler { .. }
+                        | OutputMode::Deno
+                        | OutputMode::Node {
+                            experimental_modules: true,
+                        } => "import.meta.url",
+                        OutputMode::Node {
+                            experimental_modules: false,
+                        } => "require('url').pathToFileURL(__filename)",
+                        OutputMode::NoModules { .. } => "script_src",
+                    };
+                    Ok(format!("new URL('{}', {}).toString()", path, base))
+                } else {
+                    if let Some(content) = content {
+                        let mut escaped = String::with_capacity(content.len());
+                        content.chars().for_each(|c| match c {
+                            '`' | '\\' | '$' => escaped.extend(['\\', c]),
+                            _ => escaped.extend([c]),
+                        });
+                        Ok(format!(
+                            "\"data:application/javascript,\" + encodeURIComponent(`{escaped}`)"
+                        ))
+                    } else {
+                        Err(anyhow!("wasm-bindgen needs to be invoked with `--split-linked-modules`, because \"{}\" cannot be embedded.\n\
+                            See https://rustwasm.github.io/wasm-bindgen/reference/cli.html#--split-linked-modules for details.", path))
+                    }
+                }
+            }
         }
     }
 

crates/cli-support/src/lib.rs

@@ -46,6 +46,7 @@ pub struct Bindgen {
     multi_value: bool,
     wasm_interface_types: bool,
     encode_into: EncodeInto,
+    split_linked_modules: bool,
 }
 
 pub struct Output {
@@ -120,6 +121,7 @@ impl Bindgen {
             wasm_interface_types,
             encode_into: EncodeInto::Test,
             omit_default_module_path: true,
+            split_linked_modules: true,
         }
     }
 
@@ -304,6 +306,11 @@ impl Bindgen {
         self
     }
 
+    pub fn split_linked_modules(&mut self, split_linked_modules: bool) -> &mut Bindgen {
+        self.split_linked_modules = split_linked_modules;
+        self
+    }
+
     pub fn generate<P: AsRef<Path>>(&mut self, path: P) -> Result<(), Error> {
         self.generate_output()?.emit(path.as_ref())
     }

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

@@ -1,3 +1,4 @@
+use crate::decode::LocalModule;
 use crate::descriptor::{Descriptor, Function};
 use crate::descriptors::WasmBindgenDescriptorsSection;
 use crate::intrinsic::Intrinsic;
@@ -340,6 +341,45 @@ impl<'a> Context<'a> {
         Ok(())
     }
 
+    fn link_module(
+        &mut self,
+        id: ImportId,
+        module: &decode::ImportModule,
+        offset: usize,
+        local_modules: &[LocalModule],
+        inline_js: &[&str],
+    ) -> Result<(), Error> {
+        let descriptor = Function {
+            shim_idx: 0,
+            arguments: Vec::new(),
+            ret: Descriptor::String,
+            inner_ret: None,
+        };
+        let id = self.import_adapter(id, descriptor, AdapterJsImportKind::Normal)?;
+        let (path, content) = match module {
+            decode::ImportModule::Named(n) => (
+                format!("snippets/{}", n),
+                local_modules
+                    .iter()
+                    .find(|m| m.identifier == *n)
+                    .map(|m| m.contents),
+            ),
+            decode::ImportModule::RawNamed(n) => (n.to_string(), None),
+            decode::ImportModule::Inline(idx) => (
+                format!(
+                    "snippets/{}/inline{}.js",
+                    self.unique_crate_identifier,
+                    *idx as usize + offset
+                ),
+                Some(inline_js[*idx as usize]),
+            ),
+        };
+        self.aux
+            .import_map
+            .insert(id, AuxImport::LinkTo(path, content.map(str::to_string)));
+        Ok(())
+    }
+
     fn program(&mut self, program: decode::Program<'a>) -> Result<(), Error> {
         self.unique_crate_identifier = program.unique_crate_identifier;
         let decode::Program {
@@ -352,9 +392,10 @@ impl<'a> Context<'a> {
             inline_js,
             unique_crate_identifier,
             package_json,
+            linked_modules,
         } = program;
 
-        for module in local_modules {
+        for module in &local_modules {
             // All local modules we find should be unique, but the same module
             // may have showed up in a few different blocks. If that's the case
             // all the same identifiers should have the same contents.
@@ -373,6 +414,25 @@ impl<'a> Context<'a> {
             self.export(export)?;
         }
 
+        let offset = self
+            .aux
+            .snippets
+            .get(unique_crate_identifier)
+            .map(|s| s.len())
+            .unwrap_or(0);
+        for module in linked_modules {
+            match self.function_imports.remove(module.link_function_name) {
+                Some((id, _)) => self.link_module(
+                    id,
+                    &module.module,
+                    offset,
+                    &local_modules[..],
+                    &inline_js[..],
+                )?,
+                None => (),
+            }
+        }
+
         // Register vendor prefixes for all types before we walk over all the
         // imports to ensure that if a vendor prefix is listed somewhere it'll
         // apply to all the imports.

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

@@ -328,6 +328,12 @@ pub enum AuxImport {
     /// This is an intrinsic function expected to be implemented with a JS glue
     /// shim. Each intrinsic has its own expected signature and implementation.
     Intrinsic(Intrinsic),
+
+    /// This is a function which returns a URL pointing to a specific file,
+    /// usually a JS snippet. The supplied path is relative to the JS glue shim.
+    /// The Option may contain the contents of the linked file, so it can be
+    /// embedded.
+    LinkTo(String, Option<String>),
 }
 
 /// Values that can be imported verbatim to hook up to an import.

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

@@ -355,6 +355,9 @@ fn check_standard_import(import: &AuxImport) -> Result<(), Error> {
         AuxImport::Intrinsic(intrinsic) => {
             format!("wasm-bindgen specific intrinsic `{}`", intrinsic.name())
         }
+        AuxImport::LinkTo(path, _) => {
+            format!("wasm-bindgen specific link function for `{}`", path)
+        }
         AuxImport::Closure { .. } => format!("creating a `Closure` wrapper"),
     };
     bail!("import of {} requires JS glue", item);