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);
     }
   }