[NFC] DeadArgumentElimination: Compute callers once

src/passes/DeadArgumentElimination.cpp

@@ -34,6 +34,7 @@
 // watch for here).
 //
 
+#include <algorithm>
 #include <unordered_map>
 #include <unordered_set>
 
@@ -218,6 +219,20 @@ struct DAE : public Pass {
     }
   }
 
+  // For each function, the set of callers. This is used to propagate changes,
+  // e.g. if we remove a return value from a function, the calls might benefit
+  // from optimization. It is ok if this is an over-approximation, that is, if
+  // we think there are more callers than there are, as it would just lead to
+  // unneeded extra scanning of calling functions (in the example just given, if
+  // a caller did not actually call, they would not benefit from the extra
+  // optimization, but no harm is done, and no optimization missed). Such over-
+  // approximation can happen in later optimization iterations: We may manage to
+  // remove a call from a function to another (say, after applying a constant
+  // param, we see the call is not reached). This is somewhat rare, and the cost
+  // of computing this map is significant, so we compute it once at the start
+  // and then use that possibly-over-approximating data.
+  std::vector<std::vector<Name>> callers;
+
   bool iteration(Module* module, DAEFunctionInfoMap& infoMap) {
     allDroppedCalls.clear();
 
@@ -246,15 +261,10 @@ struct DAE : public Pass {
     std::vector<bool> tailCallees(numFunctions);
     std::vector<bool> hasUnseenCalls(numFunctions);
 
-    // For each function, the set of callers.
-    std::vector<std::unordered_set<Name>> callers(numFunctions);
-
     for (auto& [func, info] : infoMap) {
       for (auto& [name, calls] : info.calls) {
-        auto targetIndex = indexes[name];
-        auto& allCallsToName = allCalls[targetIndex];
+        auto& allCallsToName = allCalls[indexes[name]];
         allCallsToName.insert(allCallsToName.end(), calls.begin(), calls.end());
-        callers[targetIndex].insert(func);
       }
       for (auto& callee : info.tailCallees) {
         tailCallees[indexes[callee]] = true;
@@ -273,6 +283,23 @@ struct DAE : public Pass {
       }
     }
 
+    // See comment above, we compute callers once and never again.
+    if (callers.empty()) {
+      // Compute first as sets, to deduplicate.
+      std::vector<std::unordered_set<Name>> callersSets(numFunctions);
+      for (auto& [func, info] : infoMap) {
+        for (auto& [name, calls] : info.calls) {
+          callersSets[indexes[name]].insert(func);
+        }
+      }
+      // Copy into efficient vectors.
+      callers.resize(numFunctions);
+      for (Index i = 0; i < numFunctions; ++i) {
+        auto& set = callersSets[i];
+        callers[i] = std::vector<Name>(set.begin(), set.end());
+      }
+    }
+
     // Track which functions we changed that are worth re-optimizing at the end.
     std::unordered_set<Function*> worthOptimizing;