1 files changed, 72 insertions, 29 deletions

diff --git a/src/ir/local-graph.h b/src/ir/local-graph.h
index 17cc6ff31..1c36fc1dc 100644
--- a/src/ir/local-graph.h
+++ b/src/ir/local-graph.h
@@ -38,12 +38,56 @@ namespace wasm {
 // debugging etc.; and it has no downside for optimization, since unreachable
 // code will be removed anyhow).
 //
-struct LocalGraph {
+// There are two options here, the normal LocalGraph which is eager and computes
+// everything up front, which is faster if most things end up needed, and a lazy
+// one which computes on demand, which can be much faster if we only need a
+// small subset of queries.
+//
+
+// Base class for both LocalGraph and LazyLocalGraph (not meant for direct use).
+struct LocalGraphBase {
+protected:
   // If a module is passed in, it is used to find which features are needed in
   // the computation (for example, if exception handling is disabled, then we
   // can generate a simpler CFG, as calls cannot throw).
+  LocalGraphBase(Function* func, Module* module = nullptr)
+    : func(func), module(module) {}
+
+public:
+  // A set of sets, returned from the query about which sets can be read from a
+  // get. Typically only one or two apply there, so this is a small set.
+  using Sets = SmallSet<LocalSet*, 2>;
+
+  // Where each get and set is. We compute this while doing the main computation
+  // and make it accessible for users, for easy replacing of things without
+  // extra work.
+  using Locations = std::map<Expression*, Expression**>;
+  Locations locations;
+
+  // Sets of gets or sets, that are influenced, returned from get*Influences().
+  using SetInfluences = std::unordered_set<LocalGet*>;
+  using GetInfluences = std::unordered_set<LocalSet*>;
+
+  // Defined publicly as other utilities need similar data layouts.
+  using GetSetsMap = std::unordered_map<LocalGet*, Sets>;
+
+protected:
+  Function* func;
+  Module* module;
+
+  std::set<Index> SSAIndexes;
+
+  // A map of each get to the sets relevant to it. This is mutable so that
+  // getSets() can be const in LazyLocalGraph (which does memoization, see
+  // below).
+  mutable GetSetsMap getSetsMap;
+
+  std::unordered_map<LocalSet*, SetInfluences> setInfluences;
+  std::unordered_map<LocalGet*, GetInfluences> getInfluences;
+};
+
+struct LocalGraph : public LocalGraphBase {
   LocalGraph(Function* func, Module* module = nullptr);
-  ~LocalGraph();
 
   // Get the sets relevant for a local.get.
   //
@@ -52,10 +96,12 @@ struct LocalGraph {
   // for a param.
   //
   // Often there is a single set, or a phi or two items, so we use a small set.
-  using Sets = SmallSet<LocalSet*, 2>;
   const Sets& getSets(LocalGet* get) const {
     auto iter = getSetsMap.find(get);
     if (iter == getSetsMap.end()) {
+      // A missing entry means there is nothing there (and we saved a little
+      // space by not putting something there).
+      //
       // Use a canonical constant empty set to avoid allocation.
       static const Sets empty;
       return empty;
@@ -63,12 +109,6 @@ struct LocalGraph {
     return iter->second;
   }
 
-  // Where each get and set is. We compute this while doing the main computation
-  // and make it accessible for users, for easy replacing of things without
-  // extra work.
-  using Locations = std::map<Expression*, Expression**>;
-  Locations locations;
-
   // Checks if two gets are equivalent, that is, definitely have the same
   // value.
   bool equivalent(LocalGet* a, LocalGet* b);
@@ -84,9 +124,6 @@ struct LocalGraph {
     computeGetInfluences();
   }
 
-  using SetInfluences = std::unordered_set<LocalGet*>;
-  using GetInfluences = std::unordered_set<LocalSet*>;
-
   const SetInfluences& getSetInfluences(LocalSet* set) const {
     auto iter = setInfluences.find(set);
     if (iter == setInfluences.end()) {
@@ -130,28 +167,34 @@ struct LocalGraph {
   void computeSSAIndexes();
 
   bool isSSA(Index x);
+};
 
-  // Defined publicly as other utilities need similar data layouts.
-  using GetSetsMap = std::unordered_map<LocalGet*, Sets>;
+// The internal implementation of the flow analysis used to compute things. This
+// must be declared in the header so that LazyLocalGraph can declare a unique
+// ptr to it, below.
+struct LocalGraphFlower;
 
-private:
-  Function* func;
-  std::set<Index> SSAIndexes;
+struct LazyLocalGraph : public LocalGraphBase {
+  LazyLocalGraph(Function* func, Module* module = nullptr);
+  ~LazyLocalGraph();
 
-  // A map of each get to the sets relevant to it. This is mutable so that
-  // getSets() can be const.
-  mutable GetSetsMap getSetsMap;
+  const Sets& getSets(LocalGet* get) const {
+    auto iter = getSetsMap.find(get);
+    if (iter == getSetsMap.end()) {
+      // A missing entry means we did not do the computation yet. Do it now.
+      computeGetSets(get);
+      iter = getSetsMap.find(get);
+      assert(iter != getSetsMap.end());
+    }
+    return iter->second;
+  }
 
-  // The internal implementation of the flow analysis used to compute
-  // getSetsMap.
-  struct LocalGraphFlower;
-  // This could be a unique_ptr, but the forward declaration is not compatible
-  // with that. It could alternatively be a shared_ptr, but that runs into what
-  // seems to be a false positive of clang's (but not gcc's) UBSan.
-  std::unique_ptr<LocalGraphFlower> flower;
+private:
+  // Compute the sets for a get and store them on getSetsMap.
+  void computeGetSets(LocalGet* get) const;
 
-  std::unordered_map<LocalSet*, SetInfluences> setInfluences;
-  std::unordered_map<LocalGet*, GetInfluences> getInfluences;
+  // This remains alive as long as we are, so that we can compute things lazily.
+  std::unique_ptr<LocalGraphFlower> flower;
 };
 
 } // namespace wasm