Refactor to use graph traversal

Author: stevenfontanella

src/passes/GlobalEffects.cpp

@@ -22,25 +22,92 @@
 #include "ir/effects.h"
 #include "ir/module-utils.h"
 #include "pass.h"
+#include "support/hash.h"
 #include "support/unique_deferring_queue.h"
 #include "wasm.h"
 
 namespace wasm {
 
+namespace {
+
+struct FuncInfo {
+  // Effects in this function.
+  std::optional<EffectAnalyzer> effects;
+
+  // Directly-called functions from this function.
+  std::unordered_set<Name> calledFunctions;
+};
+
+// 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>>& in,
+  std::map<Function*, FuncInfo>& funcInfos) {
+
+  std::unordered_set<std::pair<Name, Name>> processed;
+  std::deque<std::pair<Name, Name>> work;
+
+  for (const auto& [callee, callers] : in) {
+    for (const auto& caller : callers) {
+      work.emplace_back(callee, caller);
+      processed.emplace(callee, caller);
+    }
+  }
+
+  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;
+    }
+
+    if (!calleeEffects) {
+      callerEffects.reset();
+      return;
+    }
+
+    callerEffects->mergeIn(*calleeEffects);
+  };
+
+  while (!work.empty()) {
+    auto [callee, caller] = work.back();
+    work.pop_back();
+
+    if (callee == caller) {
+      auto& callerEffects = funcInfos.at(module.getFunction(caller)).effects;
+      if (callerEffects) {
+        callerEffects->trap = true;
+      }
+    }
+
+    // TODO: can have early return here
+    propagate(callee, caller);
+
+    const auto& callerCallers = in.find(caller);
+    if (callerCallers == in.end()) {
+      continue;
+    }
+
+    for (const Name& callerCaller : callerCallers->second) {
+      if (processed.contains({callee, callerCaller})) {
+        continue;
+      }
+
+      processed.emplace(callee, callerCaller);
+      work.emplace_back(callee, callerCaller);
+    }
+  }
+}
+
 struct GenerateGlobalEffects : public Pass {
   void run(Module* module) override {
     // First, we do a scan of each function to see what effects they have,
     // including which functions they call directly (so that we can compute
     // transitive effects later).
 
-    struct FuncInfo {
-      // Effects in this function.
-      std::optional<EffectAnalyzer> effects;
-
-      // Directly-called functions from this function.
-      std::unordered_set<Name> calledFunctions;
-    };
-
     ModuleUtils::ParallelFunctionAnalysis<FuncInfo> analysis(
       *module, [&](Function* func, FuncInfo& funcInfo) {
         if (func->imported()) {
@@ -100,86 +167,15 @@ struct GenerateGlobalEffects : public Pass {
         }
       });
 
-    // Compute the transitive closure of effects. To do so, first construct for
-    // each function a list of the functions that it is called by (so we need to
-    // propagate its effects to them), and then we'll construct the closure of
-    // that.
-    //
-    // callers[foo] = [func that calls foo, another func that calls foo, ..]
-    //
+    // callee : caller
     std::unordered_map<Name, std::unordered_set<Name>> callers;
-
-    // Our work queue contains info about a new call pair: a call from a caller
-    // to a called function, that is information we then apply and propagate.
-    using CallPair = std::pair<Name, Name>; // { caller, called }
-    UniqueDeferredQueue<CallPair> work;
-    for (auto& [func, info] : analysis.map) {
-      for (auto& called : info.calledFunctions) {
-        work.push({func->name, called});
+    for (const auto& [func, info] : analysis.map) {
+      for (const auto& callee : info.calledFunctions) {
+        callers[callee].insert(func->name);
       }
     }
 
-    // Compute the transitive closure of the call graph, that is, fill out
-    // |callers| so that it contains the list of all callers - even through a
-    // chain - of each function.
-    while (!work.empty()) {
-      auto [caller, called] = work.pop();
-
-      // We must not already have an entry for this call (that would imply we
-      // are doing wasted work).
-      assert(!callers[called].contains(caller));
-
-      // Apply the new call information.
-      callers[called].insert(caller);
-
-      // We just learned that |caller| calls |called|. It also calls
-      // transitively, which we need to propagate to all places unaware of that
-      // information yet.
-      //
-      //   caller => called => called by called
-      //
-      auto& calledInfo = analysis.map[module->getFunction(called)];
-      for (auto calledByCalled : calledInfo.calledFunctions) {
-        if (!callers[calledByCalled].contains(caller)) {
-          work.push({caller, calledByCalled});
-        }
-      }
-    }
-
-    // Now that we have transitively propagated all static calls, apply that
-    // information. First, apply infinite recursion: if a function can call
-    // itself then it might recurse infinitely, which we consider an effect (a
-    // trap).
-    for (auto& [func, info] : analysis.map) {
-      if (callers[func->name].contains(func->name)) {
-        if (info.effects) {
-          info.effects->trap = true;
-        }
-      }
-    }
-
-    // Next, apply function effects to their callers.
-    for (auto& [func, info] : analysis.map) {
-      auto& funcEffects = info.effects;
-
-      for (auto& caller : callers[func->name]) {
-        auto& callerEffects = analysis.map[module->getFunction(caller)].effects;
-        if (!callerEffects) {
-          // Nothing is known for the caller, which is already the worst case.
-          continue;
-        }
-
-        if (!funcEffects) {
-          // Nothing is known for the called function, which means nothing is
-          // known for the caller either.
-          callerEffects.reset();
-          continue;
-        }
-
-        // Add func's effects to the caller.
-        callerEffects->mergeIn(*funcEffects);
-      }
-    }
+    propagateEffects(*module, callers, analysis.map);
 
     // Generate the final data, starting from a blank slate where nothing is
     // known.
@@ -202,6 +198,8 @@ struct DiscardGlobalEffects : public Pass {
   }
 };
 
+} // namespace
+
 Pass* createGenerateGlobalEffectsPass() { return new GenerateGlobalEffects(); }
 
 Pass* createDiscardGlobalEffectsPass() { return new DiscardGlobalEffects(); }