[NFC] Optimize ParamUtils::getUsedParams()

src/passes/DeadArgumentElimination.cpp

@@ -158,7 +158,7 @@ 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) {
-      auto usedParams = ParamUtils::getUsedParams(func);
+      auto usedParams = ParamUtils::getUsedParams(func, getModule());
       for (Index i = 0; i < numParams; i++) {
         if (usedParams.count(i) == 0) {
           info->unusedParams.insert(i);

src/passes/SignaturePruning.cpp

@@ -96,7 +96,7 @@ struct SignaturePruning : public Pass {
 
         info.calls = std::move(FindAll<Call>(func->body).list);
         info.callRefs = std::move(FindAll<CallRef>(func->body).list);
-        info.usedParams = ParamUtils::getUsedParams(func);
+        info.usedParams = ParamUtils::getUsedParams(func, module);
       });
 
     // A map of types to all the information combined over all the functions

src/passes/param-utils.cpp

@@ -15,9 +15,9 @@
  */
 
 #include "passes/param-utils.h"
+#include "cfg/liveness-traversal.h"
 #include "ir/eh-utils.h"
 #include "ir/function-utils.h"
-#include "ir/local-graph.h"
 #include "ir/localize.h"
 #include "ir/possible-constant.h"
 #include "ir/type-updating.h"
@@ -28,25 +28,50 @@
 
 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;
+std::unordered_set<Index> getUsedParams(Function* func, Module* module) {
+  // To find which params are used, compute liveness at the entry.
+  // TODO: We could write bespoke code here rather than reuse LivenessWalker, as
+  //       we only need liveness at the entry. The code below computes it for
+  //       the param indexes in the entire function. However, there are usually
+  //       very few params (compared to locals, which we ignore here), so this
+  //       may be fast enough, and is very simple.
+  struct ParamLiveness
+    : public LivenessWalker<ParamLiveness, Visitor<ParamLiveness>> {
+    using Super = LivenessWalker<ParamLiveness, Visitor<ParamLiveness>>;
+
+    // Branches outside of the function can be ignored, as we only look at
+    // locals, which vanish when we leave.
+    bool ignoreBranchesOutsideOfFunc = true;
+
+    // Ignore unreachable code and non-params.
+    static void doVisitLocalGet(ParamLiveness* self, Expression** currp) {
+      auto* get = (*currp)->cast<LocalGet>();
+      if (self->currBasicBlock && self->getFunction()->isParam(get->index)) {
+        Super::doVisitLocalGet(self, currp);
+      }
     }
-
-    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);
+    static void doVisitLocalSet(ParamLiveness* self, Expression** currp) {
+      auto* set = (*currp)->cast<LocalSet>();
+      if (self->currBasicBlock && self->getFunction()->isParam(set->index)) {
+        Super::doVisitLocalSet(self, currp);
       }
     }
+  } walker;
+  walker.setModule(module);
+  walker.walkFunction(func);
+
+  if (!walker.entry) {
+    // Empty function: nothing is used.
+    return {};
   }
 
+  // We now have a sorted vector of the live params at the entry. Convert that
+  // to a set.
+  auto& sortedLiveness = walker.entry->contents.start;
+  std::unordered_set<Index> usedParams;
+  for (auto live : sortedLiveness) {
+    usedParams.insert(live);
+  }
   return usedParams;
 }
 

src/passes/param-utils.h

@@ -42,7 +42,7 @@ namespace wasm::ParamUtils {
 // function foo(x) {
 //   bar(x); // read of a param value
 // }
-std::unordered_set<Index> getUsedParams(Function* func);
+std::unordered_set<Index> getUsedParams(Function* func, Module* module);
 
 // The outcome of an attempt to remove a parameter(s).
 enum RemovalOutcome {

test/lit/passes/dae_all-features.wast

@@ -940,3 +940,34 @@
   (unreachable)
  )
 )
+
+(module
+ ;; CHECK:      (type $0 (func))
+
+ ;; CHECK:      (type $1 (func (param i32)))
+
+ ;; CHECK:      (func $target (type $0)
+ ;; CHECK-NEXT:  (local $0 i32)
+ ;; CHECK-NEXT:  (unreachable)
+ ;; CHECK-NEXT:  (drop
+ ;; CHECK-NEXT:   (local.get $0)
+ ;; CHECK-NEXT:  )
+ ;; CHECK-NEXT: )
+ (func $target (param $0 i32)
+  ;; The parameter here is unused: there is a get, but it is unreachable. We can
+  ;; remove the parameter here, and in the caller below.
+  (unreachable)
+  (drop
+   (local.get $0)
+  )
+ )
+
+ ;; CHECK:      (func $caller (type $1) (param $x i32)
+ ;; CHECK-NEXT:  (call $target)
+ ;; CHECK-NEXT: )
+ (func $caller (param $x i32)
+  (call $target
+   (local.get $x)
+  )
+ )
+)