Fix LegalizeJSInterface with RefFuncs (#3749)

This code used to remove functions it no longer thinks are needed. That is,
if it adds a legalized version of an import, it would remove the illegal
one which is no longer needed. To avoid removing an illegal import that
is still used it checked for ref.func appearances.

But this was bad in two ways:

    We need to legalize the ref.funcs too. We can't call an illegal import
    in any way, not using a direct call, indirect call, or call by reference of
    a ref.func.

    It's silly to remove unneeded functions here. We have a pass for that.

This removes the removal of functions, and adds proper updating of
ref.calls, which means to call the stub function that looks like the
original import, but that calls the legalized one and connects things
up properly, exactly the same way as other calls.

Also remove code that checked if we were in the stub/thunk and to
not replace the call there. That code is not needed: no one will ever
call the illegal import, so we do not need to be careful about
preserving such calls.

1 files changed, 16 insertions, 45 deletions

diff --git a/src/passes/LegalizeJSInterface.cpp b/src/passes/LegalizeJSInterface.cpp
index a53481d65..e840e375b 100644
--- a/src/passes/LegalizeJSInterface.cpp
+++ b/src/passes/LegalizeJSInterface.cpp
@@ -107,51 +107,16 @@ struct LegalizeJSInterface : public Pass {
     }
 
     if (!illegalImportsToLegal.empty()) {
-      // Gather functions used in 'ref.func'. They should not be removed.
-      std::unordered_map<Name, std::atomic<bool>> usedInRefFunc;
-
-      // Fill in unordered_map, as we operate on it in parallel.
-      for (auto& func : module->functions) {
-        usedInRefFunc[func->name];
-      }
-
-      struct RefFuncScanner : public WalkerPass<PostWalker<RefFuncScanner>> {
-        Module& wasm;
-        std::unordered_map<Name, std::atomic<bool>>& usedInRefFunc;
-
-        bool isFunctionParallel() override { return true; }
-
-        Pass* create() override {
-          return new RefFuncScanner(wasm, usedInRefFunc);
-        }
-
-        RefFuncScanner(
-          Module& wasm,
-          std::unordered_map<Name, std::atomic<bool>>& usedInRefFunc)
-          : wasm(wasm), usedInRefFunc(usedInRefFunc) {}
-
-        void visitRefFunc(RefFunc* curr) { usedInRefFunc[curr->func] = true; }
-      };
-
-      RefFuncScanner(*module, usedInRefFunc).run(runner, module);
-      for (auto& pair : illegalImportsToLegal) {
-        if (!usedInRefFunc[pair.first]) {
-          module->removeFunction(pair.first);
-        }
-      }
-
       // fix up imports: call_import of an illegal must be turned to a call of a
-      // legal
-      struct FixImports : public WalkerPass<PostWalker<FixImports>> {
+      // legal. the same must be done with ref.funcs.
+      struct Fixer : public WalkerPass<PostWalker<Fixer>> {
         bool isFunctionParallel() override { return true; }
 
-        Pass* create() override {
-          return new FixImports(illegalImportsToLegal);
-        }
+        Pass* create() override { return new Fixer(illegalImportsToLegal); }
 
         std::map<Name, Name>* illegalImportsToLegal;
 
-        FixImports(std::map<Name, Name>* illegalImportsToLegal)
+        Fixer(std::map<Name, Name>* illegalImportsToLegal)
           : illegalImportsToLegal(illegalImportsToLegal) {}
 
         void visitCall(Call* curr) {
@@ -160,19 +125,25 @@ struct LegalizeJSInterface : public Pass {
             return;
           }
 
-          if (iter->second == getFunction()->name) {
-            // inside the stub function itself, is the one safe place to do the
-            // call
-            return;
-          }
           replaceCurrent(
             Builder(*getModule())
               .makeCall(
                 iter->second, curr->operands, curr->type, curr->isReturn));
         }
+
+        void visitRefFunc(RefFunc* curr) {
+          auto iter = illegalImportsToLegal->find(curr->func);
+          if (iter == illegalImportsToLegal->end()) {
+            return;
+          }
+
+          curr->func = iter->second;
+        }
       };
 
-      FixImports(&illegalImportsToLegal).run(runner, module);
+      Fixer fixer(&illegalImportsToLegal);
+      fixer.run(runner, module);
+      fixer.walkModuleCode(module);
     }
   }