From 22eb34d9ab17a11efaea1d61ad4258dcf345ca23 Mon Sep 17 00:00:00 2001
From: Alex Crichton <alex@alexcrichton.com>
Date: Wed, 1 May 2019 11:54:06 -0700
Subject: [PATCH] Fix direct imports in `--target web`

Currently the import object constructed for the `--target web` output
only ever includes the current module as an one of the modules included.
With `wasm-bindgen`'s optimization to import directly from modules,
however, it's possible to have more modules imported from in the
generated wasm file. This commit ensures that the imports are hooked up
in the `--target web` es6 emulation mode, ensuring there aren't
extraneous errors about import objects.
---
 crates/cli-support/src/js/mod.rs | 33 ++++++++++++++++++++++++++++++--
 tests/headless/main.rs           |  1 +
 tests/headless/modules.js        |  3 +++
 tests/headless/modules.rs        | 12 ++++++++++++
 4 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 tests/headless/modules.js
 create mode 100644 tests/headless/modules.rs

diff --git a/crates/cli-support/src/js/mod.rs b/crates/cli-support/src/js/mod.rs
index c00185516..49601dafe 100644
--- a/crates/cli-support/src/js/mod.rs
+++ b/crates/cli-support/src/js/mod.rs
@@ -13,7 +13,7 @@ use crate::{
 };
 use failure::{bail, Error, ResultExt};
 use std::{
-    collections::{BTreeMap, HashMap, HashSet},
+    collections::{BTreeMap, HashMap, HashSet, BTreeSet},
     env, fs,
 };
 use walrus::{MemoryId, Module};
@@ -918,13 +918,41 @@ impl<'a> Context<'a> {
         } else {
             ""
         };
-
         let ts = Self::ts_for_init_fn(mem.import.is_some());
+
+        // Generate extra initialization for the `imports` object if necessary
+        // based on the values in `direct_imports` we find. These functions are
+        // intended to be imported directly to the wasm module and we need to
+        // ensure that the modules are actually imported from and inserted into
+        // the object correctly.
+        let mut map = BTreeMap::new();
+        for &(module, name) in self.direct_imports.values() {
+            map.entry(module).or_insert(BTreeSet::new()).insert(name);
+        }
+        let mut imports_init = String::new();
+        for (module, names) in map {
+            imports_init.push_str("imports['");
+            imports_init.push_str(module);
+            imports_init.push_str("'] = { ");
+            for (i, name) in names.into_iter().enumerate() {
+                if i != 0 {
+                    imports_init.push_str(", ");
+                }
+                let import = Import::Module { module, name, field: None };
+                let identifier = self.import_identifier(import);
+                imports_init.push_str(name);
+                imports_init.push_str(": ");
+                imports_init.push_str(&identifier);
+            }
+            imports_init.push_str(" };\n");
+        }
+
         let js = format!(
             "\
                 function init(module{init_memory_arg}) {{
                     let result;
                     const imports = {{ './{module}': __exports }};
+                    {imports_init}
                     if (module instanceof URL || typeof module === 'string' || module instanceof Request) {{
                         {init_memory2}
                         const response = fetch(module);
@@ -973,6 +1001,7 @@ impl<'a> Context<'a> {
             } else {
                 ""
             },
+            imports_init = imports_init,
         );
 
         (js, ts)
diff --git a/tests/headless/main.rs b/tests/headless/main.rs
index 3f3f3718b..448a11950 100644
--- a/tests/headless/main.rs
+++ b/tests/headless/main.rs
@@ -39,3 +39,4 @@ fn can_log_html_strings() {
 }
 
 pub mod snippets;
+pub mod modules;
diff --git a/tests/headless/modules.js b/tests/headless/modules.js
new file mode 100644
index 000000000..c6ce3f0eb
--- /dev/null
+++ b/tests/headless/modules.js
@@ -0,0 +1,3 @@
+export function get_five() {
+  return 5;
+}
diff --git a/tests/headless/modules.rs b/tests/headless/modules.rs
new file mode 100644
index 000000000..9e894f340
--- /dev/null
+++ b/tests/headless/modules.rs
@@ -0,0 +1,12 @@
+use wasm_bindgen::prelude::*;
+use wasm_bindgen_test::*;
+
+#[wasm_bindgen(raw_module = "./tests/headless/modules.js")]
+extern "C" {
+    fn get_five() -> u32;
+}
+
+#[wasm_bindgen_test]
+fn test_get_five() {
+    assert_eq!(get_five(), 5);
+}