Optimize Asyncify to not flatten/optimize unnecessarily (#5293)

Add a way to proxy passes and the addition of passes in pass runners. With
that we can make Asyncify only modify functions it actually needs to. On a
project that Asyncify only needs to modify a few functions on, this can save
a huge amount of time as it avoids flattening+optimizing the majority of
the module.

Fixes #4822

2 files changed, 55 insertions, 4 deletions

diff --git a/src/pass.h b/src/pass.h
index 02cea7879..60d6692a1 100644
--- a/src/pass.h
+++ b/src/pass.h
@@ -259,6 +259,8 @@ struct PassRunner {
   PassRunner(const PassRunner&) = delete;
   PassRunner& operator=(const PassRunner&) = delete;
 
+  virtual ~PassRunner() = default;
+
   // But we can make it easy to create a nested runner
   // TODO: Go through and use this in more places
   explicit PassRunner(const PassRunner* runner)
@@ -341,6 +343,9 @@ struct PassRunner {
   // Returns whether a pass by that name will remove debug info.
   static bool passRemovesDebugInfo(const std::string& name);
 
+protected:
+  virtual void doAdd(std::unique_ptr<Pass> pass);
+
 private:
   // Whether this is a nested pass runner.
   bool isNested = false;
@@ -352,8 +357,6 @@ private:
   // Whether this pass runner has run. A pass runner should only be run once.
   bool ran = false;
 
-  void doAdd(std::unique_ptr<Pass> pass);
-
   void runPass(Pass* pass);
   void runPassOnFunction(Pass* pass, Function* func);
 
diff --git a/src/passes/Asyncify.cpp b/src/passes/Asyncify.cpp
index 342cd013d..365891db4 100644
--- a/src/passes/Asyncify.cpp
+++ b/src/passes/Asyncify.cpp
@@ -848,6 +848,54 @@ public:
   }
 };
 
+// Proxy that runs wrapped pass for instrumented functions only
+struct InstrumentedProxy : public Pass {
+  std::unique_ptr<Pass> create() override {
+    return std::make_unique<InstrumentedProxy>(analyzer, pass->create());
+  }
+
+  InstrumentedProxy(ModuleAnalyzer* analyzer, std::unique_ptr<Pass> pass)
+    : analyzer(analyzer), pass(std::move(pass)) {}
+
+  bool isFunctionParallel() override { return pass->isFunctionParallel(); }
+
+  void runOnFunction(Module* module, Function* func) override {
+    if (!analyzer->needsInstrumentation(func)) {
+      return;
+    }
+    if (pass->getPassRunner() == nullptr) {
+      pass->setPassRunner(getPassRunner());
+    }
+    pass->runOnFunction(module, func);
+  }
+
+  bool modifiesBinaryenIR() override { return pass->modifiesBinaryenIR(); }
+
+  bool invalidatesDWARF() override { return pass->invalidatesDWARF(); }
+
+  bool requiresNonNullableLocalFixups() override {
+    return pass->requiresNonNullableLocalFixups();
+  }
+
+private:
+  ModuleAnalyzer* analyzer;
+  std::unique_ptr<Pass> pass;
+};
+
+struct InstrumentedPassRunner : public PassRunner {
+  InstrumentedPassRunner(Module* wasm, ModuleAnalyzer* analyzer)
+    : PassRunner(wasm), analyzer(analyzer) {}
+
+protected:
+  void doAdd(std::unique_ptr<Pass> pass) override {
+    PassRunner::doAdd(
+      std::unique_ptr<Pass>(new InstrumentedProxy(analyzer, std::move(pass))));
+  }
+
+private:
+  ModuleAnalyzer* analyzer;
+};
+
 // Instrument control flow, around calls and adding skips for rewinding.
 struct AsyncifyFlow : public Pass {
   bool isFunctionParallel() override { return true; }
@@ -1616,7 +1664,7 @@ struct Asyncify : public Pass {
     // practical to add code around each call, without affecting
     // anything else.
     {
-      PassRunner runner(module);
+      InstrumentedPassRunner runner(module, &analyzer);
       runner.add("flatten");
       // Dce is useful here, since AsyncifyFlow makes control flow conditional,
       // which may make unreachable code look reachable. It also lets us ignore
@@ -1647,7 +1695,7 @@ struct Asyncify : public Pass {
     // restore those locals). We also and optimize after as well to simplify
     // the code as much as possible.
     {
-      PassRunner runner(module);
+      InstrumentedPassRunner runner(module, &analyzer);
       if (optimize) {
         runner.addDefaultFunctionOptimizationPasses();
       }