Merge pull request #1063 from alexcrichton/wasm2es6js-start

wasm2es6js: Fix handling of start function

crates/cli-support/src/wasm2es6js.rs

@@ -121,7 +121,7 @@ impl Output {
         return ts
     }
 
-    pub fn js(self) -> Result<String, Error> {
+    pub fn js_and_wasm(mut self) -> Result<(String, Option<Vec<u8>>), Error> {
         let mut js_imports = String::new();
         let mut exports = String::new();
         let mut set_exports = String::new();
@@ -136,7 +136,7 @@ impl Output {
 
                 let name = (b'a' + (set.len() as u8)) as char;
                 js_imports.push_str(&format!(
-                    "import * as import_{} from '{}';",
+                    "import * as import_{} from '{}';\n",
                     name,
                     entry.module()
                 ));
@@ -155,6 +155,29 @@ impl Output {
                 set_exports.push_str(";\n");
             }
         }
+
+        // This is sort of tricky, but the gist of it is that if there's a start
+        // function we want to defer execution of the start function until after
+        // all our module's exports are bound. That way we'll execute it as soon
+        // as we're ready, but the module's imports and such will be able to
+        // work as everything is wired up.
+        //
+        // This ends up helping out in situations such as:
+        //
+        // * The start function calls an imported function
+        // * That imported function in turn tries to access the wasm module
+        //
+        // If we don't do this then the second step won't work because the start
+        // function is automatically executed before the promise of
+        // instantiation resolves, meaning that we won't actually have anything
+        // bound for it to access.
+        //
+        // If we remove the start function here (via `unstart`) then we'll
+        // reexport it as `__wasm2es6js_start` so be manually executed here.
+        if self.unstart() {
+            set_exports.push_str("wasm.exports.__wasm2es6js_start();\n");
+        }
+
         let inst = format!(
             "
             WebAssembly.instantiate(bytes,{{ {imports} }})
@@ -166,8 +189,8 @@ impl Output {
             imports = imports,
             set_exports = set_exports,
         );
+        let wasm = serialize(self.module).expect("failed to serialize");
         let (bytes, booted) = if self.base64 {
-            let wasm = serialize(self.module).expect("failed to serialize");
             (
                 format!(
                     "
@@ -199,8 +222,8 @@ impl Output {
         } else {
             bail!("the option --base64 or --fetch is required");
         };
-        Ok(format!(
-            "
+        let js = format!(
+            "\
             {js_imports}
             {bytes}
             export const booted = {booted};
@@ -210,6 +233,32 @@ impl Output {
             booted = booted,
             js_imports = js_imports,
             exports = exports,
-        ))
+        );
+        let wasm = if self.base64 { None } else { Some(wasm) };
+        Ok((js, wasm))
+    }
+
+    /// See comments above for what this is doing, but in a nutshell this
+    /// removes the start section, if any, and moves it to an exported function.
+    /// Returns whether a start function was found and removed.
+    fn unstart(&mut self) -> bool {
+        let mut start = None;
+        for (i, section) in self.module.sections().iter().enumerate() {
+            if let Section::Start(idx) = section  {
+                start = Some((i, *idx));
+                break;
+            }
+        }
+        let (i, idx) = match start {
+            Some(p) => p,
+            None => return false,
+        };
+        self.module.sections_mut().remove(i);
+        let entry = ExportEntry::new(
+            "__wasm2es6js_start".to_string(),
+            Internal::Function(idx),
+        );
+        self.module.export_section_mut().unwrap().entries_mut().push(entry);
+        true
     }
 }

crates/cli/src/bin/wasm-bindgen-test-runner/server.rs

@@ -68,8 +68,10 @@ pub fn spawn(
     let output = Config::new()
         .fetch(Some(format!("/{}", wasm_name)))
         .generate(&wasm)?;
-    let js = output.js()?;
+    let (js, wasm) = output.js_and_wasm()?;
+    let wasm = wasm.unwrap();
     fs::write(tmpdir.join(format!("{}_bg.js", module)), js).context("failed to write JS file")?;
+    fs::write(tmpdir.join(format!("{}_bg.wasm", module)), wasm).context("failed to write wasm file")?;
 
     // For now, always run forever on this port. We may update this later!
     let tmpdir = tmpdir.to_path_buf();

crates/cli/src/bin/wasm2es6js.rs

@@ -25,6 +25,7 @@ Usage:
 Options:
     -h --help               Show this screen.
     -o --output FILE        File to place output in
+    --out-dir DIR           Directory to place ouptut in
     --typescript            Output a `*.d.ts` file next to the JS output
     --base64                Inline the wasm module using base64 encoding
     --fetch PATH            Load module by passing the PATH argument to `fetch()`
@@ -37,6 +38,7 @@ bundlers for working with wasm. Use this program with care!
 #[derive(Debug, Deserialize)]
 struct Args {
     flag_output: Option<PathBuf>,
+    flag_out_dir: Option<PathBuf>,
     flag_typescript: bool,
     flag_base64: bool,
     flag_fetch: Option<String>,
@@ -68,22 +70,34 @@ fn rmain(args: &Args) -> Result<(), Error> {
         .generate(&wasm)?;
 
     if args.flag_typescript {
-        if let Some(ref p) = args.flag_output {
-            let dst = p.with_extension("d.ts");
-            let ts = object.typescript();
-            fs::write(&dst, ts).with_context(|_| format!("failed to write `{}`", dst.display()))?;
-        }
+        let ts = object.typescript();
+        write(&args, "d.ts", ts.as_bytes(), false)?;
     }
 
-    let js = object.js()?;
+    let (js, wasm) = object.js_and_wasm()?;
 
-    match args.flag_output {
-        Some(ref p) => {
-            fs::write(p, js).with_context(|_| format!("failed to write `{}`", p.display()))?;
-        }
-        None => {
-            println!("{}", js);
-        }
+    write(args, "js", js.as_bytes(), false)?;
+    if let Some(wasm) = wasm {
+        write(args, "wasm", &wasm, false)?;
+    }
+    Ok(())
+}
+
+fn write(
+    args: &Args,
+    extension: &str,
+    contents: &[u8],
+    print_fallback: bool,
+) -> Result<(), Error> {
+    if let Some(p) = &args.flag_output {
+        let dst = p.with_extension(extension);
+        fs::write(&dst, contents).with_context(|_| format!("failed to write `{}`", dst.display()))?;
+    } else if let Some(p) = &args.flag_out_dir {
+        let filename = args.arg_input.file_name().unwrap();
+        let dst = p.join(filename).with_extension(extension);
+        fs::write(&dst, contents).with_context(|_| format!("failed to write `{}`", dst.display()))?;
+    } else if print_fallback {
+        println!("{}", String::from_utf8_lossy(contents))
     }
 
     Ok(())