[NFC] Move InstrumentedPass logic out and use it in another place (#6132)

Asyncify gained a way to wrap a pass so that it only runs on a given set of
functions, rather than on all functions, so the wrapper "filters" what the pass
operates on. That was useful in Asyncify as we wanted to only do work on
functions that Asyncify actually instrumented.

There is another place in the code that needs such functionality,
optimizeAfterInlining, which runs optimizations after we inline; again, we
only want to optimize on the functions we know are relevant because they
changed. To do that, move that logic out to a general place so it can be
reused. This makes the code there a lot less hackish.

While doing so make the logic only work on function-parallel passes. It
never did anyhow, but now it asserts on that. (It can't run on a general
pass because a general one does not provide an interface to affect which
functions it operates on; a general pass is entirely opaque in that way.)

3 files changed, 113 insertions, 75 deletions

diff --git a/src/passes/Asyncify.cpp b/src/passes/Asyncify.cpp
index b30af085e..838a290bb 100644
--- a/src/passes/Asyncify.cpp
+++ b/src/passes/Asyncify.cpp
@@ -315,6 +315,7 @@
 #include "ir/names.h"
 #include "ir/utils.h"
 #include "pass.h"
+#include "passes/pass-utils.h"
 #include "support/file.h"
 #include "support/string.h"
 #include "wasm-builder.h"
@@ -845,56 +846,6 @@ 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 addsEffects() override { return pass->addsEffects(); }
-
-  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; }
@@ -1706,12 +1657,21 @@ struct Asyncify : public Pass {
     // Add necessary globals before we emit code to use them.
     addGlobals(module, relocatable);
 
+    // Compute the set of functions we will instrument. All of the passes we run
+    // below only need to run there.
+    PassUtils::FuncSet instrumentedFuncs;
+    for (auto& func : module->functions) {
+      if (analyzer.needsInstrumentation(func.get())) {
+        instrumentedFuncs.insert(func.get());
+      }
+    }
+
     // Instrument the flow of code, adding code instrumentation and
     // skips for when rewinding. We do this on flat IR so that it is
     // practical to add code around each call, without affecting
     // anything else.
     {
-      InstrumentedPassRunner runner(module, &analyzer);
+      PassUtils::FilteredPassRunner runner(module, instrumentedFuncs);
       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
@@ -1752,7 +1712,7 @@ struct Asyncify : public Pass {
     // restore those locals). We also and optimize after as well to simplify
     // the code as much as possible.
     {
-      InstrumentedPassRunner runner(module, &analyzer);
+      PassUtils::FilteredPassRunner runner(module, instrumentedFuncs);
       if (optimize) {
         runner.addDefaultFunctionOptimizationPasses();
       }
diff --git a/src/passes/opt-utils.h b/src/passes/opt-utils.h
index 1945167b0..7bb733b03 100644
--- a/src/passes/opt-utils.h
+++ b/src/passes/opt-utils.h
@@ -23,37 +23,21 @@
 #include <ir/element-utils.h>
 #include <ir/module-utils.h>
 #include <pass.h>
+#include <passes/pass-utils.h>
 #include <wasm.h>
 
-namespace wasm {
+namespace wasm::OptUtils {
 
-namespace OptUtils {
-
-// Run useful optimizations after inlining new code into a set
-// of functions.
-inline void optimizeAfterInlining(const std::unordered_set<Function*>& funcs,
+// Run useful optimizations after inlining new code into a set of functions.
+inline void optimizeAfterInlining(const PassUtils::FuncSet& funcs,
                                   Module* module,
                                   PassRunner* parentRunner) {
-  // save the full list of functions on the side
-  std::vector<std::unique_ptr<Function>> all;
-  all.swap(module->functions);
-  module->updateFunctionsMap();
-  for (auto& func : funcs) {
-    module->addFunction(func);
-  }
-  PassRunner runner(module, parentRunner->options);
+  PassUtils::FilteredPassRunner runner(module, funcs, parentRunner->options);
   runner.setIsNested(true);
-  runner.setValidateGlobally(false); // not a full valid module
   // this is especially useful after inlining
   runner.add("precompute-propagate");
   runner.addDefaultFunctionOptimizationPasses(); // do all the usual stuff
   runner.run();
-  // restore all the funcs
-  for (auto& func : module->functions) {
-    func.release();
-  }
-  all.swap(module->functions);
-  module->updateFunctionsMap();
 }
 
 struct FunctionRefReplacer
@@ -102,7 +86,6 @@ inline void replaceFunctions(PassRunner* runner,
   }
 }
 
-} // namespace OptUtils
-} // namespace wasm
+} // namespace wasm::OptUtils
 
 #endif // wasm_passes_opt_utils_h
diff --git a/src/passes/pass-utils.h b/src/passes/pass-utils.h
new file mode 100644
index 000000000..1db4f5614
--- /dev/null
+++ b/src/passes/pass-utils.h
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2023 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_passes_pass_utils_h
+#define wasm_passes_pass_utils_h
+
+#include <unordered_set>
+
+#include <pass.h>
+#include <wasm.h>
+
+namespace wasm::PassUtils {
+
+using FuncSet = std::unordered_set<Function*>;
+
+// A wrapper around a parallel pass that filters it to run run only on select
+// functions.
+struct FilteredPass : public Pass {
+  std::unique_ptr<Pass> create() override {
+    // Function-parallel passes get a new instance per function. Create a copy
+    // of the wrapped pass along with ourselves.
+    return std::make_unique<FilteredPass>(pass->create(), relevantFuncs);
+  }
+
+  FilteredPass(std::unique_ptr<Pass>&& pass, const FuncSet& relevantFuncs)
+    : pass(std::move(pass)), relevantFuncs(relevantFuncs) {}
+
+  bool isFunctionParallel() override {
+    assert(pass->isFunctionParallel());
+    return true;
+  }
+
+  void runOnFunction(Module* module, Function* func) override {
+    if (!relevantFuncs.count(func)) {
+      return;
+    }
+
+    // The pass runner calling us set our pass runner, which we must do for the
+    // wrapped pass.
+    pass->setPassRunner(getPassRunner());
+    pass->runOnFunction(module, func);
+  }
+
+  bool modifiesBinaryenIR() override { return pass->modifiesBinaryenIR(); }
+
+  bool invalidatesDWARF() override { return pass->invalidatesDWARF(); }
+
+  bool addsEffects() override { return pass->addsEffects(); }
+
+  bool requiresNonNullableLocalFixups() override {
+    return pass->requiresNonNullableLocalFixups();
+  }
+
+private:
+  std::unique_ptr<Pass> pass;
+  const FuncSet& relevantFuncs;
+};
+
+// A pass runner that wraps all passes, filtering so that they only run on
+// select functions.
+struct FilteredPassRunner : public PassRunner {
+  FilteredPassRunner(Module* wasm, const FuncSet& relevantFuncs)
+    : PassRunner(wasm), relevantFuncs(relevantFuncs) {}
+
+  FilteredPassRunner(Module* wasm,
+                     const FuncSet& relevantFuncs,
+                     const PassOptions& options)
+    : PassRunner(wasm, options), relevantFuncs(relevantFuncs) {}
+
+protected:
+  void doAdd(std::unique_ptr<Pass> pass) override {
+    PassRunner::doAdd(
+      std::make_unique<FilteredPass>(std::move(pass), relevantFuncs));
+  }
+
+private:
+  const FuncSet& relevantFuncs;
+};
+
+} // namespace wasm::PassUtils
+
+#endif // wasm_passes_pass_utils_h