[NFC] Move and generalize parameter-removing logic from DeadArgumentElimination (#4544)

In preparation for removing dead arguments from all functions sharing a heap
type (which seems useful for j2wasm output), first this PR refactors that code
so it is reusable. This moves the code out of the pass into FunctionUtils, and
also generalizes it slightly by

supporting a set of functions and not just a single one, and
receiving a list of call_refs and not just calls
(no other changes to anything).

4 files changed, 267 insertions, 101 deletions

diff --git a/src/passes/CMakeLists.txt b/src/passes/CMakeLists.txt
index 29c8cd526..ac3971690 100644
--- a/src/passes/CMakeLists.txt
+++ b/src/passes/CMakeLists.txt
@@ -13,6 +13,7 @@ configure_file(WasmIntrinsics.cpp.in WasmIntrinsics.cpp @ONLY)
 
 FILE(GLOB passes_HEADERS *.h)
 set(passes_SOURCES
+  param-utils.cpp
   pass.cpp
   test_passes.cpp
   AlignmentLowering.cpp
diff --git a/src/passes/DeadArgumentElimination.cpp b/src/passes/DeadArgumentElimination.cpp
index ec81639b0..ce4d17c74 100644
--- a/src/passes/DeadArgumentElimination.cpp
+++ b/src/passes/DeadArgumentElimination.cpp
@@ -40,11 +40,11 @@
 #include "ir/effects.h"
 #include "ir/element-utils.h"
 #include "ir/find_all.h"
-#include "ir/local-graph.h"
 #include "ir/lubs.h"
 #include "ir/module-utils.h"
 #include "ir/type-updating.h"
 #include "ir/utils.h"
+#include "param-utils.h"
 #include "pass.h"
 #include "passes/opt-utils.h"
 #include "support/sorted_vector.h"
@@ -155,41 +155,13 @@ struct DAEScanner
     // part of, say if we are exported, or if another parallel function finds a
     // RefFunc to us and updates it before we check it).
     if (numParams > 0 && !info->hasUnseenCalls) {
-      findUnusedParams(func);
-    }
-  }
-
-  void findUnusedParams(Function* func) {
-    LocalGraph localGraph(func);
-    std::unordered_set<Index> usedParams;
-    for (auto& [get, sets] : localGraph.getSetses) {
-      if (!func->isParam(get->index)) {
-        continue;
-      }
-
-      // Check if this get of a param index can read from the parameter value
-      // passed into the function. We want to ignore values set in the function
-      // like this:
-      //
-      // function foo(x) {
-      //   x = 10;
-      //   bar(x); // read of a param index, but not the param value passed in.
-      // }
-      for (auto* set : sets) {
-        // A nullptr value indicates there is no LocalSet* that sets the value,
-        // so it must be the parameter value.
-        if (!set) {
-          usedParams.insert(get->index);
+      auto usedParams = ParamUtils::getUsedParams(func);
+      for (Index i = 0; i < numParams; i++) {
+        if (usedParams.count(i) == 0) {
+          info->unusedParams.insert(i);
         }
       }
     }
-
-    // We can now compute the unused params.
-    for (Index i = 0; i < numParams; i++) {
-      if (usedParams.count(i) == 0) {
-        info->unusedParams.insert(i);
-      }
-    }
   }
 };
 
@@ -315,38 +287,11 @@ struct DAE : public Pass {
       if (numParams == 0) {
         continue;
       }
-      // Iterate downwards, as we may remove more than one.
-      Index i = numParams - 1;
-      while (1) {
-        if (infoMap[name].unusedParams.has(i)) {
-          // Great, it's not used. Check if none of the calls has a param with
-          // side effects that we cannot remove (as if we can remove them, we
-          // will simply do that when we remove the parameter). Note: flattening
-          // the IR beforehand can help here.
-          bool callParamsAreValid =
-            std::none_of(calls.begin(), calls.end(), [&](Call* call) {
-              auto* operand = call->operands[i];
-              return EffectAnalyzer(runner->options, *module, operand)
-                .hasUnremovableSideEffects();
-            });
-          // The type must be valid for us to handle as a local (since we
-          // replace the parameter with a local).
-          // TODO: if there are no references at all, we can avoid creating a
-          //       local
-          bool typeIsValid =
-            TypeUpdating::canHandleAsLocal(func->getLocalType(i));
-          if (callParamsAreValid && typeIsValid) {
-            // Wonderful, nothing stands in our way! Do it.
-            // TODO: parallelize this?
-            removeParameter(func, i, calls);
-            TypeUpdating::handleNonDefaultableLocals(func, *module);
-            changed.insert(func);
-          }
-        }
-        if (i == 0) {
-          break;
-        }
-        i--;
+      auto removedIndexes = ParamUtils::removeParameters(
+        {func}, infoMap[name].unusedParams, calls, {}, module, runner);
+      if (!removedIndexes.empty()) {
+        // Success!
+        changed.insert(func);
       }
     }
     // We can also tell which calls have all their return values dropped. Note
@@ -395,42 +340,6 @@ struct DAE : public Pass {
 private:
   std::unordered_map<Call*, Expression**> allDroppedCalls;
 
-  void removeParameter(Function* func, Index i, std::vector<Call*>& calls) {
-    // It's cumbersome to adjust local names - TODO don't clear them?
-    Builder::clearLocalNames(func);
-    // Remove the parameter from the function. We must add a new local
-    // for uses of the parameter, but cannot make it use the same index
-    // (in general).
-    auto paramsType = func->getParams();
-    std::vector<Type> params(paramsType.begin(), paramsType.end());
-    auto type = params[i];
-    params.erase(params.begin() + i);
-    func->setParams(Type(params));
-    Index newIndex = Builder::addVar(func, type);
-    // Update local operations.
-    struct LocalUpdater : public PostWalker<LocalUpdater> {
-      Index removedIndex;
-      Index newIndex;
-      LocalUpdater(Function* func, Index removedIndex, Index newIndex)
-        : removedIndex(removedIndex), newIndex(newIndex) {
-        walk(func->body);
-      }
-      void visitLocalGet(LocalGet* curr) { updateIndex(curr->index); }
-      void visitLocalSet(LocalSet* curr) { updateIndex(curr->index); }
-      void updateIndex(Index& index) {
-        if (index == removedIndex) {
-          index = newIndex;
-        } else if (index > removedIndex) {
-          index--;
-        }
-      }
-    } localUpdater(func, i, newIndex);
-    // Remove the arguments from the calls.
-    for (auto* call : calls) {
-      call->operands.erase(call->operands.begin() + i);
-    }
-  }
-
   void
   removeReturnValue(Function* func, std::vector<Call*>& calls, Module* module) {
     func->setResults(Type::none);
diff --git a/src/passes/param-utils.cpp b/src/passes/param-utils.cpp
new file mode 100644
index 000000000..019df6d7f
--- /dev/null
+++ b/src/passes/param-utils.cpp
@@ -0,0 +1,177 @@
+/*
+ * Copyright 2022 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.
+ */
+
+#include "ir/function-utils.h"
+#include "ir/local-graph.h"
+#include "ir/type-updating.h"
+#include "support/sorted_vector.h"
+#include "wasm.h"
+
+namespace wasm::ParamUtils {
+
+std::unordered_set<Index> getUsedParams(Function* func) {
+  LocalGraph localGraph(func);
+
+  std::unordered_set<Index> usedParams;
+
+  for (auto& [get, sets] : localGraph.getSetses) {
+    if (!func->isParam(get->index)) {
+      continue;
+    }
+
+    for (auto* set : sets) {
+      // A nullptr value indicates there is no LocalSet* that sets the value,
+      // so it must be the parameter value.
+      if (!set) {
+        usedParams.insert(get->index);
+      }
+    }
+  }
+
+  return usedParams;
+}
+
+bool removeParameter(const std::vector<Function*> funcs,
+                     Index index,
+                     const std::vector<Call*>& calls,
+                     const std::vector<CallRef*>& callRefs,
+                     Module* module,
+                     PassRunner* runner) {
+  assert(funcs.size() > 0);
+  auto* first = funcs[0];
+#ifndef NDEBUG
+  for (auto* func : funcs) {
+    assert(func->type == first->type);
+  }
+#endif
+
+  // Check if none of the calls has a param with side effects that we cannot
+  // remove (as if we can remove them, we will simply do that when we remove the
+  // parameter). Note: flattening the IR beforehand can help here.
+  bool callParamsAreValid =
+    std::none_of(calls.begin(), calls.end(), [&](Call* call) {
+      auto* operand = call->operands[index];
+      return EffectAnalyzer(runner->options, *module, operand)
+        .hasUnremovableSideEffects();
+    });
+  if (!callParamsAreValid) {
+    return false;
+  }
+
+  // The type must be valid for us to handle as a local (since we
+  // replace the parameter with a local).
+  // TODO: if there are no references at all, we can avoid creating a
+  //       local
+  bool typeIsValid = TypeUpdating::canHandleAsLocal(first->getLocalType(index));
+  if (!typeIsValid) {
+    return false;
+  }
+
+  // We can do it!
+
+  // Remove the parameter from the function. We must add a new local
+  // for uses of the parameter, but cannot make it use the same index
+  // (in general).
+  auto paramsType = first->getParams();
+  std::vector<Type> params(paramsType.begin(), paramsType.end());
+  auto type = params[index];
+  params.erase(params.begin() + index);
+  // TODO: parallelize some of these loops?
+  for (auto* func : funcs) {
+    func->setParams(Type(params));
+
+    // It's cumbersome to adjust local names - TODO don't clear them?
+    Builder::clearLocalNames(func);
+  }
+  std::vector<Index> newIndexes;
+  for (auto* func : funcs) {
+    newIndexes.push_back(Builder::addVar(func, type));
+  }
+  // Update local operations.
+  struct LocalUpdater : public PostWalker<LocalUpdater> {
+    Index removedIndex;
+    Index newIndex;
+    LocalUpdater(Function* func, Index removedIndex, Index newIndex)
+      : removedIndex(removedIndex), newIndex(newIndex) {
+      walk(func->body);
+    }
+    void visitLocalGet(LocalGet* curr) { updateIndex(curr->index); }
+    void visitLocalSet(LocalSet* curr) { updateIndex(curr->index); }
+    void updateIndex(Index& index) {
+      if (index == removedIndex) {
+        index = newIndex;
+      } else if (index > removedIndex) {
+        index--;
+      }
+    }
+  };
+  for (Index i = 0; i < funcs.size(); i++) {
+    auto* func = funcs[i];
+    if (!func->imported()) {
+      LocalUpdater(funcs[i], index, newIndexes[i]);
+      TypeUpdating::handleNonDefaultableLocals(func, *module);
+    }
+  }
+
+  // Remove the arguments from the calls.
+  for (auto* call : calls) {
+    call->operands.erase(call->operands.begin() + index);
+  }
+  for (auto* call : callRefs) {
+    call->operands.erase(call->operands.begin() + index);
+  }
+
+  return true;
+}
+
+SortedVector removeParameters(const std::vector<Function*> funcs,
+                              SortedVector indexes,
+                              const std::vector<Call*>& calls,
+                              const std::vector<CallRef*>& callRefs,
+                              Module* module,
+                              PassRunner* runner) {
+  if (indexes.empty()) {
+    return {};
+  }
+
+  assert(funcs.size() > 0);
+  auto* first = funcs[0];
+#ifndef NDEBUG
+  for (auto* func : funcs) {
+    assert(func->type == first->type);
+  }
+#endif
+
+  // Iterate downwards, as we may remove more than one, and going forwards would
+  // alter the indexes after us.
+  Index i = first->getNumParams() - 1;
+  SortedVector removed;
+  while (1) {
+    if (indexes.has(i)) {
+      if (removeParameter(funcs, i, calls, callRefs, module, runner)) {
+        // Success!
+        removed.insert(i);
+      }
+    }
+    if (i == 0) {
+      break;
+    }
+    i--;
+  }
+  return removed;
+}
+
+} // namespace wasm::ParamUtils
diff --git a/src/passes/param-utils.h b/src/passes/param-utils.h
new file mode 100644
index 000000000..9bfbf1aec
--- /dev/null
+++ b/src/passes/param-utils.h
@@ -0,0 +1,79 @@
+/*
+ * Copyright 2022 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_ir_function_h
+#define wasm_ir_function_h
+
+#include "pass.h"
+#include "support/sorted_vector.h"
+#include "wasm.h"
+
+// Helper code for passes that manipulate function parameters, specifically
+// checking if they are used and removing them if so. This is closely tied to
+// the internals of those passes, and so is not in /ir/ (it would be inside the
+// pass .cpp file, but there is more than one).
+
+namespace wasm::ParamUtils {
+
+// Find which parameters are actually used in the function, that is, that the
+// values arriving in the parameter are read. This ignores values set in the
+// function, like this:
+//
+// function foo(x) {
+//   x = 10;
+//   bar(x); // read of a param index, but not the param value passed in.
+// }
+//
+// This is an actual use:
+//
+// function foo(x) {
+//   bar(x); // read of a param value
+// }
+std::unordered_set<Index> getUsedParams(Function* func);
+
+// Try to remove a parameter from a set of functions and replace it with a local
+// instead. This may not succeed if the parameter type cannot be used in a
+// local, or if we hit another limitation, in which case this returns false and
+// does nothing. If we succeed then the parameter is removed both from the
+// functions and from the calls to it, which are passed in (the caller must
+// ensure to pass in all relevant calls and call_refs).
+//
+// This does not check if removing the parameter would change the semantics
+// (say, if the parameter's value is used), which the caller is assumed to do.
+//
+// This assumes that the set of functions all have the same signature. The main
+// use cases are either to send a single function, or to send a set of functions
+// that all have the same heap type (and so if they all do not use some
+// parameter, it can be removed from them all).
+bool removeParameter(const std::vector<Function*> funcs,
+                     Index index,
+                     const std::vector<Call*>& calls,
+                     const std::vector<CallRef*>& callRefs,
+                     Module* module,
+                     PassRunner* runner);
+
+// The same as removeParameter, but gets a sorted list of indexes. It tries to
+// remove them all, and returns which we removed.
+SortedVector removeParameters(const std::vector<Function*> funcs,
+                              SortedVector indexes,
+                              const std::vector<Call*>& calls,
+                              const std::vector<CallRef*>& callRefs,
+                              Module* module,
+                              PassRunner* runner);
+
+} // namespace wasm::ParamUtils
+
+#endif // wasm_ir_function_h