1 files changed, 67 insertions, 32 deletions

diff --git a/src/ir/module-utils.h b/src/ir/module-utils.h
index 0fd478551..f2fda908c 100644
--- a/src/ir/module-utils.h
+++ b/src/ir/module-utils.h
@@ -262,33 +262,20 @@ template<typename T> inline void iterDefinedEvents(Module& wasm, T visitor) {
   }
 }
 
-// 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 {
+// Helper class for performing an operation on all the functions in the module,
+// in parallel, with an Info object for each one that can contain results of
+// some computation that the operation performs.
+// The operation performend should not modify the wasm module in any way.
+// TODO: enforce this
+template<typename T> struct ParallelFunctionAnalysis {
   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;
 
-  CallGraphPropertyAnalysis(Module& wasm, Func work) : wasm(wasm) {
+  ParallelFunctionAnalysis(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) {
@@ -304,30 +291,78 @@ template<typename T> struct CallGraphPropertyAnalysis {
 
     struct Mapper : public WalkerPass<PostWalker<Mapper>> {
       bool isFunctionParallel() override { return true; }
+      bool modifiesBinaryenIR() override { return false; }
 
-      Mapper(Module* module, Map* map, Func work)
+      Mapper(Module& module, Map& map, Func work)
         : module(module), map(map), work(work) {}
 
       Mapper* create() override { return new Mapper(module, map, work); }
 
-      void visitCall(Call* curr) {
-        (*map)[this->getFunction()].callsTo.insert(
-          module->getFunction(curr->target));
-      }
-
-      void visitFunction(Function* curr) {
-        assert((*map).count(curr));
-        work(curr, (*map)[curr]);
+      void doWalkFunction(Function* curr) {
+        assert(map.count(curr));
+        work(curr, map[curr]);
       }
 
     private:
-      Module* module;
-      Map* map;
+      Module& module;
+      Map& map;
       Func work;
     };
 
     PassRunner runner(&wasm);
-    Mapper(&wasm, &map, work).run(&runner, &wasm);
+    Mapper(wasm, map, work).run(&runner, &wasm);
+  }
+};
+
+// 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;
+
+  CallGraphPropertyAnalysis(Module& wasm, Func work) : wasm(wasm) {
+    ParallelFunctionAnalysis<T> analysis(wasm, [&](Function* func, T& info) {
+      work(func, info);
+      if (func->imported()) {
+        return;
+      }
+      struct Mapper : public PostWalker<Mapper> {
+        Mapper(Module* module, T& info, Func work)
+          : module(module), info(info), work(work) {}
+
+        void visitCall(Call* curr) {
+          info.callsTo.insert(module->getFunction(curr->target));
+        }
+
+      private:
+        Module* module;
+        T& info;
+        Func work;
+      } mapper(&wasm, info, work);
+      mapper.walk(func->body);
+    });
+
+    map.swap(analysis.map);
 
     // Find what is called by what.
     for (auto& pair : map) {