From e6334a3313e3b6a2775f6cf5d89abe95b0789fad Mon Sep 17 00:00:00 2001
From: Christian Speckner <christian.speckner@mayflower.de>
Date: Sat, 29 Jun 2024 15:07:12 +0200
Subject: [PATCH] Support setting pass main argument again with --pass-arg

---
 src/pass.h                            |  1 +
 src/passes/pass.cpp                   | 10 ++++++++++
 src/tools/optimization-options.h      | 21 +++++++++++++++++++++
 src/tools/tool-options.h              |  9 ++++++++-
 test/lit/passes/extract-function.wast |  2 ++
 test/passes/set-globals.passes        |  2 +-
 6 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/src/pass.h b/src/pass.h
index e6e94d3908c..311b7d2f108 100644
--- a/src/pass.h
+++ b/src/pass.h
@@ -46,6 +46,7 @@ struct PassRegistry {
   registerTestPass(const char* name, const char* description, Creator create);
   std::unique_ptr<Pass> createPass(std::string name);
   std::vector<std::string> getRegisteredNames();
+  bool containsPass(const std::string& name);
   std::string getPassDescription(std::string name);
   bool isPassHidden(std::string name);
 
diff --git a/src/passes/pass.cpp b/src/passes/pass.cpp
index f5c5ac55a6b..408ce9592c5 100644
--- a/src/passes/pass.cpp
+++ b/src/passes/pass.cpp
@@ -73,6 +73,16 @@ std::vector<std::string> PassRegistry::getRegisteredNames() {
   return ret;
 }
 
+bool PassRegistry::containsPass(const std::string& name) {
+  for (auto& [passName, _] : passInfos) {
+    if (passName == name) {
+      return true;
+    }
+  }
+
+  return false;
+}
+
 std::string PassRegistry::getPassDescription(std::string name) {
   assert(passInfos.find(name) != passInfos.end());
   return passInfos[name].description;
diff --git a/src/tools/optimization-options.h b/src/tools/optimization-options.h
index a8d926baf56..2ad1d9fa329 100644
--- a/src/tools/optimization-options.h
+++ b/src/tools/optimization-options.h
@@ -328,6 +328,27 @@ struct OptimizationOptions : public ToolOptions {
     }
   }
 
+  void addPassArg(const std::string& key, const std::string& value) override {
+    for (auto iPass = passes.rbegin(); iPass != passes.rend(); iPass++) {
+      if (iPass->name != key) {
+        continue;
+      }
+
+      if (iPass->argument.has_value()) {
+        Fatal() << iPass->name << " already set to " << *(iPass->argument);
+      }
+
+      iPass->argument = value;
+      return;
+    }
+
+    if (!PassRegistry::get()->containsPass(key)) {
+      return ToolOptions::addPassArg(key, value);
+    }
+
+    Fatal() << "can't set " << key << ": pass not enabled";
+  }
+
   bool runningDefaultOptimizationPasses() {
     for (auto& pass : passes) {
       if (pass.name == DEFAULT_OPT_PASSES) {
diff --git a/src/tools/tool-options.h b/src/tools/tool-options.h
index 10a03acc785..10e6ec3c9f0 100644
--- a/src/tools/tool-options.h
+++ b/src/tools/tool-options.h
@@ -139,7 +139,8 @@ struct ToolOptions : public Options {
                key = argument.substr(0, colon);
                value = argument.substr(colon + 1);
              }
-             passOptions.arguments[key] = value;
+
+             addPassArg(key, value);
            })
       .add(
         "--closed-world",
@@ -213,6 +214,12 @@ struct ToolOptions : public Options {
     module.features.disable(disabledFeatures);
   }
 
+  virtual void addPassArg(const std::string& key, const std::string& value) {
+    passOptions.arguments[key] = value;
+  }
+
+  virtual ~ToolOptions() = default;
+
 private:
   FeatureSet enabledFeatures = FeatureSet::Default;
   FeatureSet disabledFeatures = FeatureSet::None;
diff --git a/test/lit/passes/extract-function.wast b/test/lit/passes/extract-function.wast
index 79e4c368ed4..9258f8c63cf 100644
--- a/test/lit/passes/extract-function.wast
+++ b/test/lit/passes/extract-function.wast
@@ -1,6 +1,8 @@
 ;; NOTE: Assertions have been generated by update_lit_checks.py --all-items and should not be edited.
 ;; RUN: foreach %s %t wasm-opt --extract-function=foo -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --extract-function --pass-arg=extract-function@foo -S -o - | filecheck %s
 ;; RUN: foreach %s %t wasm-opt --extract-function-index=0 -S -o - | filecheck %s
+;; RUN: foreach %s %t wasm-opt --extract-function-index --pass-arg=extract-function-index@0 -S -o - | filecheck %s
 
 (module
   ;; CHECK:      (type $0 (func))
diff --git a/test/passes/set-globals.passes b/test/passes/set-globals.passes
index f5c0f1c24a4..9f01c13a623 100644
--- a/test/passes/set-globals.passes
+++ b/test/passes/set-globals.passes
@@ -1 +1 @@
-set-globals=foo=1337,bar=42
+set-globals_pass-arg=set-globals@foo=1337,bar=42