1 files changed, 70 insertions, 20 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 16ecd54b0..52765035c 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -20,6 +20,7 @@
 #include "ir/find_all.h"
 #include "ir/manipulation.h"
 #include "pass.h"
+#include "support/unique_deferring_queue.h"
 #include "wasm.h"
 
 namespace wasm {
@@ -288,22 +289,33 @@ template<typename T> inline void iterDefinedEvents(Module& wasm, T visitor) {
   }
 }
 
-// Performs a parallel map on function in the module, emitting a map object
-// of function => result.
-// TODO: use in inlining and elsewhere
-template<typename T> struct ParallelFunctionMap {
+// Helper class for analyzing the call graph.
+//
+// Provides hooks for running some initial calculation on each function (which
+// is done in parallel), writing to a FunctionInfo structure for each function.
+// Then you can call propagateBack() to propagate a property of interest to the
+// calling functions, transitively.
+//
+// For example, if some functions are known to call an import "foo", then you
+// can use this to find which functions call something that might eventually
+// reach foo, by initially marking the direct callers as "calling foo" and
+// propagating that backwards.
+template<typename T> struct CallGraphPropertyAnalysis {
+  Module& wasm;
+
+  // The basic information for each function about whom it calls and who is
+  // called by it.
+  struct FunctionInfo {
+    std::set<Function*> callsTo;
+    std::set<Function*> calledBy;
+  };
 
   typedef std::map<Function*, T> Map;
   Map map;
 
   typedef std::function<void(Function*, T&)> Func;
 
-  struct Info {
-    Map* map;
-    Func work;
-  };
-
-  ParallelFunctionMap(Module& wasm, Func work) {
+  CallGraphPropertyAnalysis(Module& wasm, Func work) : wasm(wasm) {
     // Fill in map, as we operate on it in parallel (each function to its own
     // entry).
     for (auto& func : wasm.functions) {
@@ -317,27 +329,65 @@ template<typename T> struct ParallelFunctionMap {
       }
     }
 
-    // Run on the implemented functions.
     struct Mapper : public WalkerPass<PostWalker<Mapper>> {
       bool isFunctionParallel() override { return true; }
 
-      Mapper(Info* info) : info(info) {}
+      Mapper(Module* module, Map* map, Func work)
+        : module(module), map(map), work(work) {}
+
+      Mapper* create() override { return new Mapper(module, map, work); }
 
-      Mapper* create() override { return new Mapper(info); }
+      void visitCall(Call* curr) {
+        (*map)[this->getFunction()].callsTo.insert(
+          module->getFunction(curr->target));
+      }
 
-      void doWalkFunction(Function* curr) {
-        assert((*info->map).count(curr));
-        info->work(curr, (*info->map)[curr]);
+      void visitFunction(Function* curr) {
+        assert((*map).count(curr));
+        work(curr, (*map)[curr]);
       }
 
     private:
-      Info* info;
+      Module* module;
+      Map* map;
+      Func work;
     };
 
-    Info info = {&map, work};
-
     PassRunner runner(&wasm);
-    Mapper(&info).run(&runner, &wasm);
+    Mapper(&wasm, &map, work).run(&runner, &wasm);
+
+    // Find what is called by what.
+    for (auto& pair : map) {
+      auto* func = pair.first;
+      auto& info = pair.second;
+      for (auto* target : info.callsTo) {
+        map[target].calledBy.insert(func);
+      }
+    }
+  }
+
+  // Propagate a property from a function to those that call it.
+  void propagateBack(std::function<bool(const T&)> hasProperty,
+                     std::function<bool(const T&)> canHaveProperty,
+                     std::function<void(T&)> addProperty) {
+    // The work queue contains items we just learned can change the state.
+    UniqueDeferredQueue<Function*> work;
+    for (auto& func : wasm.functions) {
+      if (hasProperty(map[func.get()])) {
+        work.push(func.get());
+      }
+    }
+    while (!work.empty()) {
+      auto* func = work.pop();
+      for (auto* caller : map[func].calledBy) {
+        // If we don't already have the property, and we are not forbidden
+        // from getting it, then it propagates back to us now.
+        if (!hasProperty(map[caller]) && canHaveProperty(map[caller])) {
+          addProperty(map[caller]);
+          work.push(caller);
+        }
+      }
+    }
   }
 };