Minimal JS legalization (#1824)

Even when we don't want to fully legalize code for JS, we should still legalize things that only JS cares about. In particular, dynCall_* methods are used from JS to call into the wasm table, and if they exist they are only for JS, so we should only legalize them.

The use case motivating this is that in dynamic linking you may want to disable legalization, so that wasm=>wasm module calls are fast even with i64s, but you do still need dynCalls to be legalized even in that case, otherwise an invoke with an i64 parameter would fail.

9 files changed, 175 insertions, 52 deletions

diff --git a/src/abi/js.h b/src/abi/js.h
new file mode 100644
index 000000000..539b465fb
--- /dev/null
+++ b/src/abi/js.h
@@ -0,0 +1,43 @@
+/*
+ * Copyright 2018 WebAssembly Community Group participants
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef wasm_abi_abi_h
+#define wasm_abi_abi_h
+
+#include "wasm.h"
+
+namespace wasm {
+
+namespace ABI {
+
+enum LegalizationLevel {
+  Full = 0,
+  Minimal = 1
+};
+
+inline std::string getLegalizationPass(LegalizationLevel level) {
+  if (level == Full) {
+    return "legalize-js-interface";
+  } else {
+    return "legalize-js-interface-minimally";
+  }
+}
+
+} // namespace ABI
+
+} // namespace wasm
+
+#endif // wasm_abi_abi_h
diff --git a/src/asm2wasm.h b/src/asm2wasm.h
index dadb9bee3..522d53a2b 100644
--- a/src/asm2wasm.h
+++ b/src/asm2wasm.h
@@ -41,6 +41,7 @@
 #include "wasm-builder.h"
 #include "wasm-emscripten.h"
 #include "wasm-module-building.h"
+#include "abi/js.h"
 
 namespace wasm {
 
@@ -1452,9 +1453,9 @@ void Asm2WasmBuilder::processAsm(Ref ast) {
   // finalizeCalls also does autoDrop, which is crucial for the non-optimizing case,
   // so that the output of the first pass is valid
   passRunner.add<FinalizeCalls>(this);
-  if (legalizeJavaScriptFFI) {
-    passRunner.add("legalize-js-interface");
-  }
+  passRunner.add(ABI::getLegalizationPass(
+    legalizeJavaScriptFFI ? ABI::Full : ABI::Minimal
+  ));
   if (runOptimizationPasses) {
     // autodrop can add some garbage
     passRunner.add("vacuum");
diff --git a/src/passes/LegalizeJSInterface.cpp b/src/passes/LegalizeJSInterface.cpp
index b8d16894b..8a8bb3e27 100644
--- a/src/passes/LegalizeJSInterface.cpp
+++ b/src/passes/LegalizeJSInterface.cpp
@@ -22,6 +22,13 @@
 // stub methods added in this pass, that thunk i64s into i32, i32 and
 // vice versa as necessary.
 //
+// We can also legalize in a "minimal" way, that is, only JS-specific
+// components, that only JS will care about, such as dynCall methods
+// (wasm will never call them, as it can share the table directly). E.g.
+// is dynamic linking, where we can avoid legalizing wasm=>wasm calls
+// across modules, we still want to legalize dynCalls so JS can call into the
+// table even to a signature that is not legal.
+//
 // This pass also legalizes according to asm.js FFI rules, which
 // disallow f32s. TODO: an option to not do that, if it matters?
 //
@@ -40,68 +47,74 @@
 namespace wasm {
 
 struct LegalizeJSInterface : public Pass {
+  bool full;
+
+  LegalizeJSInterface(bool full) : full(full) {}
+
   void run(PassRunner* runner, Module* module) override {
     // for each illegal export, we must export a legalized stub instead
     for (auto& ex : module->exports) {
       if (ex->kind == ExternalKind::Function) {
         // if it's an import, ignore it
         auto* func = module->getFunction(ex->value);
-        if (isIllegal(func)) {
+        if (isIllegal(func) && isRelevant(ex.get(), func)) {
           auto legalName = makeLegalStub(func, module);
           ex->value = legalName;
         }
       }
     }
-    // Avoid iterator invalidation later.
-    std::vector<Function*> originalFunctions;
-    for (auto& func : module->functions) {
-      originalFunctions.push_back(func.get());
-    }
-    // for each illegal import, we must call a legalized stub instead
-    for (auto* im : originalFunctions) {
-      if (im->imported() && isIllegal(module->getFunctionType(im->type))) {
-        auto funcName = makeLegalStubForCalledImport(im, module);
-        illegalImportsToLegal[im->name] = funcName;
-        // we need to use the legalized version in the table, as the import from JS
-        // is legal for JS. Our stub makes it look like a native wasm function.
-        for (auto& segment : module->table.segments) {
-          for (auto& name : segment.data) {
-            if (name == im->name) {
-              name = funcName;
+    if (full) {
+      // Avoid iterator invalidation later.
+      std::vector<Function*> originalFunctions;
+      for (auto& func : module->functions) {
+        originalFunctions.push_back(func.get());
+      }
+      // for each illegal import, we must call a legalized stub instead
+      for (auto* im : originalFunctions) {
+        if (im->imported() && isIllegal(module->getFunctionType(im->type))) {
+          auto funcName = makeLegalStubForCalledImport(im, module);
+          illegalImportsToLegal[im->name] = funcName;
+          // we need to use the legalized version in the table, as the import from JS
+          // is legal for JS. Our stub makes it look like a native wasm function.
+          for (auto& segment : module->table.segments) {
+            for (auto& name : segment.data) {
+              if (name == im->name) {
+                name = funcName;
+              }
             }
           }
         }
       }
-    }
-    if (illegalImportsToLegal.size() > 0) {
-      for (auto& pair : illegalImportsToLegal) {
-        module->removeFunction(pair.first);
-      }
+      if (illegalImportsToLegal.size() > 0) {
+        for (auto& pair : illegalImportsToLegal) {
+          module->removeFunction(pair.first);
+        }
 
-      // fix up imports: call_import of an illegal must be turned to a call of a legal
+        // fix up imports: call_import of an illegal must be turned to a call of a legal
 
-      struct FixImports : public WalkerPass<PostWalker<FixImports>> {
-        bool isFunctionParallel() override { return true; }
+        struct FixImports : public WalkerPass<PostWalker<FixImports>> {
+          bool isFunctionParallel() override { return true; }
 
-        Pass* create() override { return new FixImports(illegalImportsToLegal); }
+          Pass* create() override { return new FixImports(illegalImportsToLegal); }
 
-        std::map<Name, Name>* illegalImportsToLegal;
+          std::map<Name, Name>* illegalImportsToLegal;
 
-        FixImports(std::map<Name, Name>* illegalImportsToLegal) : illegalImportsToLegal(illegalImportsToLegal) {}
+          FixImports(std::map<Name, Name>* illegalImportsToLegal) : illegalImportsToLegal(illegalImportsToLegal) {}
 
-        void visitCall(Call* curr) {
-          auto iter = illegalImportsToLegal->find(curr->target);
-          if (iter == illegalImportsToLegal->end()) return;
+          void visitCall(Call* curr) {
+            auto iter = illegalImportsToLegal->find(curr->target);
+            if (iter == illegalImportsToLegal->end()) return;
 
-          if (iter->second == getFunction()->name) return; // inside the stub function itself, is the one safe place to do the call
-          replaceCurrent(Builder(*getModule()).makeCall(iter->second, curr->operands, curr->type));
-        }
-      };
+            if (iter->second == getFunction()->name) return; // inside the stub function itself, is the one safe place to do the call
+            replaceCurrent(Builder(*getModule()).makeCall(iter->second, curr->operands, curr->type));
+          }
+        };
 
-      PassRunner passRunner(module);
-      passRunner.setIsNested(true);
-      passRunner.add<FixImports>(&illegalImportsToLegal);
-      passRunner.run();
+        PassRunner passRunner(module);
+        passRunner.setIsNested(true);
+        passRunner.add<FixImports>(&illegalImportsToLegal);
+        passRunner.run();
+      }
     }
   }
 
@@ -118,6 +131,11 @@ private:
     return false;
   }
 
+  bool isRelevant(Export* ex, Function* func) {
+    if (full) return true;
+    // We are doing minimal legalization - just what JS needs.
+    return ex->name.startsWith("dynCall_");
+  }
 
   // JS calls the export, so it must call a legal stub that calls the actual wasm function
   Name makeLegalStub(Function* func, Module* module) {
@@ -256,7 +274,11 @@ private:
 };
 
 Pass *createLegalizeJSInterfacePass() {
-  return new LegalizeJSInterface();
+  return new LegalizeJSInterface(true);
+}
+
+Pass *createLegalizeJSInterfaceMinimallyPass() {
+  return new LegalizeJSInterface(false);
 }
 
 } // namespace wasm
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index f250b7de6..63b73ac6a 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -85,6 +85,7 @@ void PassRegistry::registerPasses() {
   registerPass("inlining", "inline functions (you probably want inlining-optimizing)", createInliningPass);
   registerPass("inlining-optimizing", "inline functions and optimizes where we inlined", createInliningOptimizingPass);
   registerPass("legalize-js-interface", "legalizes i64 types on the import/export boundary", createLegalizeJSInterfacePass);
+  registerPass("legalize-js-interface-minimally", "legalizes i64 types on the import/export boundary in a minimal manner, only on things only JS will call", createLegalizeJSInterfaceMinimallyPass);
   registerPass("local-cse", "common subexpression elimination inside basic blocks", createLocalCSEPass);
   registerPass("log-execution", "instrument the build with logging of where execution goes", createLogExecutionPass);
   registerPass("i64-to-i32-lowering", "lower all uses of i64s to use i32s instead", createI64ToI32LoweringPass);
diff --git a/src/passes/passes.h b/src/passes/passes.h
index df4322eee..b04303429 100644
--- a/src/passes/passes.h
+++ b/src/passes/passes.h
@@ -42,6 +42,7 @@ Pass* createI64ToI32LoweringPass();
 Pass* createInliningPass();
 Pass* createInliningOptimizingPass();
 Pass* createLegalizeJSInterfacePass();
+Pass* createLegalizeJSInterfaceMinimallyPass();
 Pass* createLocalCSEPass();
 Pass* createLogExecutionPass();
 Pass* createInstrumentLocalsPass();
diff --git a/src/tools/asm2wasm.cpp b/src/tools/asm2wasm.cpp
index a001c08a6..4907033e5 100644
--- a/src/tools/asm2wasm.cpp
+++ b/src/tools/asm2wasm.cpp
@@ -96,7 +96,7 @@ int main(int argc, const char *argv[]) {
            [&wasmOnly](Options *o, const std::string& ) {
              wasmOnly = true;
            })
-      .add("--no-legalize-javascript-ffi", "-nj", "Do not legalize (i64->i32, f32->f64) the imports and exports for interfacing with JS", Options::Arguments::Zero,
+      .add("--no-legalize-javascript-ffi", "-nj", "Do not fully legalize (i64->i32, f32->f64) the imports and exports for interfacing with JS", Options::Arguments::Zero,
            [&legalizeJavaScriptFFI](Options *o, const std::string& ) {
              legalizeJavaScriptFFI = false;
            })
diff --git a/src/tools/wasm-emscripten-finalize.cpp b/src/tools/wasm-emscripten-finalize.cpp
index 6c5682703..e9e623e3c 100644
--- a/src/tools/wasm-emscripten-finalize.cpp
+++ b/src/tools/wasm-emscripten-finalize.cpp
@@ -30,6 +30,7 @@
 #include "wasm-io.h"
 #include "wasm-printing.h"
 #include "wasm-validator.h"
+#include "abi/js.h"
 
 using namespace cashew;
 using namespace wasm;
@@ -83,7 +84,7 @@ int main(int argc, const char *argv[]) {
       .add("--input-source-map", "-ism", "Consume source map from the specified file",
            Options::Arguments::One,
            [&inputSourceMapFilename](Options *o, const std::string& argument) { inputSourceMapFilename = argument; })
-      .add("--no-legalize-javascript-ffi", "-nj", "Do not legalize (i64->i32, "
+      .add("--no-legalize-javascript-ffi", "-nj", "Do not fully legalize (i64->i32, "
            "f32->f64) the imports and exports for interfacing with JS",
            Options::Arguments::Zero,
            [&legalizeJavaScriptFFI](Options *o, const std::string& ) {
@@ -158,13 +159,13 @@ int main(int argc, const char *argv[]) {
   EmscriptenGlueGenerator generator(wasm);
   generator.fixInvokeFunctionNames();
 
-  if (legalizeJavaScriptFFI) {
-    PassRunner passRunner(&wasm);
-    passRunner.setDebug(options.debug);
-    passRunner.setDebugInfo(debugInfo);
-    passRunner.add("legalize-js-interface");
-    passRunner.run();
-  }
+  PassRunner passRunner(&wasm);
+  passRunner.setDebug(options.debug);
+  passRunner.setDebugInfo(debugInfo);
+  passRunner.add(ABI::getLegalizationPass(
+    legalizeJavaScriptFFI ? ABI::Full : ABI::Minimal
+  ));
+  passRunner.run();
 
   std::vector<Name> initializerFunctions;
 
diff --git a/test/passes/legalize-js-interface-minimally.txt b/test/passes/legalize-js-interface-minimally.txt
new file mode 100644
index 000000000..6152847f4
--- /dev/null
+++ b/test/passes/legalize-js-interface-minimally.txt
@@ -0,0 +1,39 @@
+(module
+ (type $FUNCSIG$j (func (result i64)))
+ (type $FUNCSIG$vi (func (param i32)))
+ (import "env" "imported" (func $imported (result i64)))
+ (import "env" "setTempRet0" (func $setTempRet0 (param i32)))
+ (export "func" (func $func))
+ (export "dynCall_foo" (func $legalstub$dyn))
+ (func $func (; 2 ;) (type $FUNCSIG$j) (result i64)
+  (drop
+   (call $imported)
+  )
+  (unreachable)
+ )
+ (func $dyn (; 3 ;) (type $FUNCSIG$j) (result i64)
+  (drop
+   (call $imported)
+  )
+  (unreachable)
+ )
+ (func $legalstub$dyn (; 4 ;) (result i32)
+  (local $0 i64)
+  (set_local $0
+   (call $dyn)
+  )
+  (call $setTempRet0
+   (i32.wrap/i64
+    (i64.shr_u
+     (get_local $0)
+     (i64.const 32)
+    )
+   )
+  )
+  (i32.wrap/i64
+   (get_local $0)
+  )
+ )
+)
+(module
+)
diff --git a/test/passes/legalize-js-interface-minimally.wast b/test/passes/legalize-js-interface-minimally.wast
new file mode 100644
index 000000000..2e003a521
--- /dev/null
+++ b/test/passes/legalize-js-interface-minimally.wast
@@ -0,0 +1,15 @@
+(module
+  (import "env" "imported" (func $imported (result i64)))
+  (export "func" (func $func))
+  (export "dynCall_foo" (func $dyn))
+  (func $func (result i64)
+    (drop (call $imported))
+    (unreachable)
+  )
+  (func $dyn (result i64)
+    (drop (call $imported))
+    (unreachable)
+  )
+)
+(module)
+