DeadArgumentElimination/SignaturePruning: Prune params even if called with effects

src/passes/DeadArgumentElimination.cpp

@@ -216,11 +216,26 @@ struct DAE : public Pass {
         allDroppedCalls[name] = calls;
       }
     }
+
     // Track which functions we changed, and optimize them later if necessary.
     std::unordered_set<Function*> changed;
+
     // If we refine return types then we will need to do more type updating
     // at the end.
     bool refinedReturnTypes = false;
+
+    // If we find that localizing call arguments can help (by moving their
+    // effects outside, so ParamUtils::removeParameters can handle them), then
+    // we do that at the end and perform another cycle. It is simpler to just do
+    // another cycle than to track the locations of calls, which is tricky as
+    // localization might move a call (if a call happens to be another call's
+    // param). In practice it is rare to find call arguments we want to remove,
+    // and even more rare to find effects get in the way, so this should not
+    // cause much overhead.
+    //
+    // This set tracks the functions for whom calls to it should be modified.
+    std::unordered_set<Name> callTargetsToLocalize;
+
     // We now have a mapping of all call sites for each function, and can look
     // for optimization opportunities.
     for (auto& [name, calls] : allCalls) {
@@ -263,12 +278,15 @@ struct DAE : public Pass {
       if (numParams == 0) {
         continue;
       }
-      auto removedIndexes = ParamUtils::removeParameters(
+      auto [removedIndexes, outcome] = ParamUtils::removeParameters(
         {func}, infoMap[name].unusedParams, calls, {}, module, getPassRunner());
       if (!removedIndexes.empty()) {
         // Success!
         changed.insert(func);
       }
+      if (outcome == ParamUtils::RemovalOutcome::Failure) {
+        callTargetsToLocalize.insert(name);
+      }
     }
     // We can also tell which calls have all their return values dropped. Note
     // that we can't do this if we changed anything so far, as we may have
@@ -307,10 +325,15 @@ struct DAE : public Pass {
         changed.insert(func.get());
       }
     }
+    if (!callTargetsToLocalize.empty()) {
+      ParamUtils::localizeCallsTo(
+        callTargetsToLocalize, *module, getPassRunner());
+    }
     if (optimize && !changed.empty()) {
       OptUtils::optimizeAfterInlining(changed, module, getPassRunner());
     }
-    return !changed.empty() || refinedReturnTypes;
+    return !changed.empty() || refinedReturnTypes ||
+           !callTargetsToLocalize.empty();
   }
 
 private:

src/passes/SignaturePruning.cpp

@@ -67,6 +67,16 @@ struct SignaturePruning : public Pass {
       return;
     }
 
+    // The first iteration may suggest additional work is possible. If so, run
+    // another cycle. (Even more cycles may help, but limit ourselves to 2 for
+    // now.)
+    if (iteration(module)) {
+      iteration(module);
+    }
+  }
+
+  // Returns true if more work is possible.
+  bool iteration(Module* module) {
     // First, find all the information we need. Start by collecting inside each
     // function in parallel.
 
@@ -101,6 +111,16 @@ struct SignaturePruning : public Pass {
     // Map heap types to all functions with that type.
     InsertOrderedMap<HeapType, std::vector<Function*>> sigFuncs;
 
+    // Heap types of call targets that we found we should localize calls to, in
+    // order to fully handle them. (See similar code in DeadArgumentElimination
+    // for individual functions; here we handle a HeapType at a time.) A slight
+    // complication is that we cannot track heap types here: heap types are
+    // rewritten using |GlobalTypeRewriter::updateSignatures| below, and even
+    // types that we do not modify end up replaced (as the entire set of types
+    // becomes one new big rec group). We therefore need something more stable
+    // to track here, which we do using either a Call or a Call Ref.
+    std::unordered_set<Expression*> callTargetsToLocalize;
+
     // Combine all the information we gathered into that map, iterating in a
     // deterministic order as we build up vectors where the order matters.
     for (auto& f : module->functions) {
@@ -215,12 +235,23 @@ struct SignaturePruning : public Pass {
       }
 
       auto oldParams = sig.params;
-      auto removedIndexes = ParamUtils::removeParameters(funcs,
-                                                         unusedParams,
-                                                         info.calls,
-                                                         info.callRefs,
-                                                         module,
-                                                         getPassRunner());
+      auto [removedIndexes, outcome] =
+        ParamUtils::removeParameters(funcs,
+                                     unusedParams,
+                                     info.calls,
+                                     info.callRefs,
+                                     module,
+                                     getPassRunner());
+      if (outcome == ParamUtils::RemovalOutcome::Failure) {
+        // Use either a Call or a CallRef that has this type (see explanation
+        // above on |callTargetsToLocalize|.
+        if (!info.calls.empty()) {
+          callTargetsToLocalize.insert(info.calls[0]);
+        } else {
+          assert(!info.callRefs.empty());
+          callTargetsToLocalize.insert(info.callRefs[0]);
+        }
+      }
       if (removedIndexes.empty()) {
         continue;
       }
@@ -262,6 +293,31 @@ struct SignaturePruning : public Pass {
 
     // Rewrite the types.
     GlobalTypeRewriter::updateSignatures(newSignatures, *module);
+
+    if (callTargetsToLocalize.empty()) {
+      return false;
+    }
+
+    // Localize after updating signatures, to not interfere with that
+    // operation (localization adds locals, and the indexes of locals must be
+    // taken into account in |GlobalTypeRewriter::updateSignatures| (as var
+    // indexes change when params are pruned).
+    std::unordered_set<HeapType> callTargetTypes;
+    for (auto* call : callTargetsToLocalize) {
+      HeapType type;
+      if (auto* c = call->dynCast<Call>()) {
+        type = module->getFunction(c->target)->type;
+      } else if (auto* c = call->dynCast<CallRef>()) {
+        type = c->target->type.getHeapType();
+      } else {
+        WASM_UNREACHABLE("bad call");
+      }
+      callTargetTypes.insert(type);
+    }
+
+    ParamUtils::localizeCallsTo(callTargetTypes, *module, getPassRunner());
+
+    return true;
   }
 };