NFC: Propagate global effects via strongly-connected components

src/passes/GlobalEffects.cpp

@@ -22,7 +22,7 @@
 #include "ir/effects.h"
 #include "ir/module-utils.h"
 #include "pass.h"
-#include "support/unique_deferring_queue.h"
+#include "support/strongly_connected_components.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -107,61 +107,158 @@ std::map<Function*, FuncInfo> analyzeFuncs(Module& module,
   return std::move(analysis.map);
 }
 
+using CallGraph = std::unordered_map<Function*, std::unordered_set<Function*>>;
+
+CallGraph buildCallGraph(const Module& module,
+                         const std::map<Function*, FuncInfo>& funcInfos) {
+  CallGraph callGraph;
+  for (const auto& [func, info] : funcInfos) {
+    if (info.calledFunctions.empty()) {
+      continue;
+    }
+
+    auto& callees = callGraph[func];
+    for (Name callee : info.calledFunctions) {
+      callees.insert(module.getFunction(callee));
+    }
+  }
+
+  return callGraph;
+}
+
+void mergeMaybeEffects(std::optional<EffectAnalyzer>& dest,
+                       const std::optional<EffectAnalyzer>& src) {
+  if (dest == UnknownEffects) {
+    return;
+  }
+  if (src == UnknownEffects) {
+    dest = UnknownEffects;
+    return;
+  }
+
+  dest->mergeIn(*src);
+}
+
 // Propagate effects from callees to callers transitively
 // e.g. if A -> B -> C (A calls B which calls C)
 // Then B inherits effects from C and A inherits effects from both B and C.
-void propagateEffects(
-  const Module& module,
-  const std::unordered_map<Name, std::unordered_set<Name>>& reverseCallGraph,
-  std::map<Function*, FuncInfo>& funcInfos) {
-
-  UniqueNonrepeatingDeferredQueue<std::pair<Name, Name>> work;
+//
+// Generate SCC for the call graph, then traverse it in reverse topological
+// order processing each callee before its callers. When traversing:
+// - Merge all of the effects of functions within the CC
+// - Also merge the (already computed) effects of each callee CC
+// - Add trap effects for potentially recursive call chains
+void propagateEffects(const Module& module,
+                      const PassOptions& passOptions,
+                      std::map<Function*, FuncInfo>& funcInfos,
+                      const CallGraph& callGraph) {
+  struct CallGraphSCCs
+    : SCCs<std::vector<Function*>::const_iterator, CallGraphSCCs> {
+    const std::map<Function*, FuncInfo>& funcInfos;
+    const std::unordered_map<Function*, std::unordered_set<Function*>>&
+      callGraph;
+    const Module& module;
+
+    CallGraphSCCs(
+      const std::vector<Function*>& funcs,
+      const std::map<Function*, FuncInfo>& funcInfos,
+      const std::unordered_map<Function*, std::unordered_set<Function*>>&
+        callGraph,
+      const Module& module)
+      : SCCs<std::vector<Function*>::const_iterator, CallGraphSCCs>(
+          funcs.begin(), funcs.end()),
+        funcInfos(funcInfos), callGraph(callGraph), module(module) {}
+
+    void pushChildren(Function* f) {
+      auto callees = callGraph.find(f);
+      if (callees == callGraph.end()) {
+        return;
+      }
 
-  for (const auto& [callee, callers] : reverseCallGraph) {
-    for (const auto& caller : callers) {
-      work.push(std::pair(callee, caller));
+      for (auto* callee : callees->second) {
+        push(callee);
+      }
     }
+  };
+
+  std::vector<Function*> allFuncs;
+  for (auto& [func, info] : funcInfos) {
+    allFuncs.push_back(func);
   }
+  CallGraphSCCs sccs(allFuncs, funcInfos, callGraph, module);
+
+  std::vector<std::optional<EffectAnalyzer>> componentEffects;
+  // Points to an index in componentEffects
+  std::unordered_map<Function*, Index> funcComponents;
 
-  auto propagate = [&](Name callee, Name caller) {
-    auto& callerEffects = funcInfos.at(module.getFunction(caller)).effects;
-    const auto& calleeEffects =
-      funcInfos.at(module.getFunction(callee)).effects;
-    if (!callerEffects) {
-      return;
+  for (auto ccIterator : sccs) {
+    std::optional<EffectAnalyzer>& ccEffects =
+      componentEffects.emplace_back(std::in_place, passOptions, module);
+
+    std::vector<Function*> ccFuncs(ccIterator.begin(), ccIterator.end());
+
+    for (Function* f : ccFuncs) {
+      funcComponents.emplace(f, componentEffects.size() - 1);
     }
 
-    if (!calleeEffects) {
-      callerEffects = UnknownEffects;
-      return;
+    std::unordered_set<int> calleeSccs;
+    for (Function* caller : ccFuncs) {
+      auto callees = callGraph.find(caller);
+      if (callees == callGraph.end()) {
+        continue;
+      }
+      for (auto* callee : callees->second) {
+        calleeSccs.insert(funcComponents.at(callee));
+      }
     }
 
-    callerEffects->mergeIn(*calleeEffects);
-  };
+    // Merge in effects from callees
+    for (int calleeScc : calleeSccs) {
+      const auto& calleeComponentEffects = componentEffects.at(calleeScc);
+      mergeMaybeEffects(ccEffects, calleeComponentEffects);
+    }
 
-  while (!work.empty()) {
-    auto [callee, caller] = work.pop();
+    // Add trap effects for potential cycles.
+    if (ccFuncs.size() > 1) {
+      if (ccEffects != UnknownEffects) {
+        ccEffects->trap = true;
+      }
+    } else {
+      auto* func = ccFuncs[0];
+      if (funcInfos.at(func).calledFunctions.contains(func->name)) {
+        if (ccEffects != UnknownEffects) {
+          ccEffects->trap = true;
+        }
+      }
+    }
 
-    if (callee == caller) {
-      auto& callerEffects = funcInfos.at(module.getFunction(caller)).effects;
-      if (callerEffects) {
-        callerEffects->trap = true;
+    // Aggregate effects within this CC
+    if (ccEffects) {
+      for (Function* f : ccFuncs) {
+        const auto& effects = funcInfos.at(f).effects;
+        mergeMaybeEffects(ccEffects, effects);
       }
     }
 
-    // Even if nothing changed, we still need to keep traversing the callers
-    // to look for a potential cycle which adds a trap affect on the above
-    // lines.
-    propagate(callee, caller);
+    // Assign each function's effects to its CC effects.
+    for (Function* f : ccFuncs) {
+      if (!ccEffects) {
+        funcInfos.at(f).effects = UnknownEffects;
+      } else {
+        funcInfos.at(f).effects.emplace(*ccEffects);
+      }
+    }
+  }
+}
 
-    const auto& callerCallers = reverseCallGraph.find(caller);
-    if (callerCallers == reverseCallGraph.end()) {
+void copyEffectsToFunctions(const std::map<Function*, FuncInfo>& funcInfos) {
+  for (auto& [func, info] : funcInfos) {
+    func->effects.reset();
+    if (!info.effects) {
       continue;
     }
 
-    for (const Name& callerCaller : callerCallers->second) {
-      work.push(std::pair(callee, callerCaller));
-    }
+    func->effects = std::make_shared<EffectAnalyzer>(*info.effects);
   }
 }
 
@@ -170,26 +267,11 @@ struct GenerateGlobalEffects : public Pass {
     std::map<Function*, FuncInfo> funcInfos =
       analyzeFuncs(*module, getPassOptions());
 
-    // callee : caller
-    std::unordered_map<Name, std::unordered_set<Name>> callers;
-    for (const auto& [func, info] : funcInfos) {
-      for (const auto& callee : info.calledFunctions) {
-        callers[callee].insert(func->name);
-      }
-    }
+    auto callGraph = buildCallGraph(*module, funcInfos);
 
-    propagateEffects(*module, callers, funcInfos);
+    propagateEffects(*module, getPassOptions(), funcInfos, callGraph);
 
-    // Generate the final data, starting from a blank slate where nothing is
-    // known.
-    for (auto& [func, info] : funcInfos) {
-      func->effects.reset();
-      if (!info.effects) {
-        continue;
-      }
-
-      func->effects = std::make_shared<EffectAnalyzer>(*info.effects);
-    }
+    copyEffectsToFunctions(funcInfos);
   }
 };