2 files changed, 72 insertions, 6 deletions

diff --git a/src/ir/LocalGraph.cpp b/src/ir/LocalGraph.cpp
index af1370731..a5495b5b3 100644
--- a/src/ir/LocalGraph.cpp
+++ b/src/ir/LocalGraph.cpp
@@ -281,6 +281,8 @@ struct LocalGraphFlower
                                       });
           if (lastSet != pred->lastSets.end()) {
             // There is a set here, apply it, and stop the flow.
+            // TODO: If we find a computed get, apply its sets and stop? That
+            //       could help but it requires more info on FlowBlock.
             for (auto* get : gets) {
               getSetsMap[get].insert(lastSet->second);
             }
@@ -512,7 +514,9 @@ void LocalGraph::computeSetInfluences() {
   }
 }
 
-void LocalGraph::computeGetInfluences() {
+static void
+doComputeGetInfluences(const LocalGraphBase::Locations& locations,
+                       LocalGraphBase::GetInfluencesMap& getInfluences) {
   for (auto& [curr, _] : locations) {
     if (auto* set = curr->dynCast<LocalSet>()) {
       FindAll<LocalGet> findAll(set->value);
@@ -523,6 +527,10 @@ void LocalGraph::computeGetInfluences() {
   }
 }
 
+void LocalGraph::computeGetInfluences() {
+  doComputeGetInfluences(locations, getInfluences);
+}
+
 void LocalGraph::computeSSAIndexes() {
   std::unordered_map<Index, std::set<LocalSet*>> indexSets;
   for (auto& [get, sets] : getSetsMap) {
@@ -555,13 +563,12 @@ LazyLocalGraph::LazyLocalGraph(Function* func, Module* module)
   : LocalGraphBase(func, module) {}
 
 void LazyLocalGraph::makeFlower() const {
-  // Lazy graphs do not provide |locations| publicly. TODO: perhaps refactor to
-  // avoid filling in this dummy data structure, but we may want to add a lazy
-  // version of it too, so see which makes sense first.
-  LocalGraph::Locations locations;
+  // |locations| is set here and filled in by |flower|.
+  assert(!locations);
+  locations.emplace();
 
   flower =
-    std::make_unique<LocalGraphFlower>(getSetsMap, locations, func, module);
+    std::make_unique<LocalGraphFlower>(getSetsMap, *locations, func, module);
 
   flower->prepareLaziness();
 
@@ -585,6 +592,9 @@ LazyLocalGraph::~LazyLocalGraph() {
 }
 
 void LazyLocalGraph::computeGetSets(LocalGet* get) const {
+  // We must never repeat work.
+  assert(!getSetsMap.count(get));
+
   if (!flower) {
     makeFlower();
   }
@@ -592,10 +602,40 @@ void LazyLocalGraph::computeGetSets(LocalGet* get) const {
 }
 
 void LazyLocalGraph::computeSetInfluences(LocalSet* set) const {
+  // We must never repeat work.
+  assert(!setInfluences.count(set));
+
   if (!flower) {
     makeFlower();
   }
   flower->computeSetInfluences(set, setInfluences);
 }
 
+void LazyLocalGraph::computeGetInfluences() const {
+  // We must never repeat work.
+  assert(!getInfluences);
+
+  // We do not need any flow for this, but we do need |locations| to be filled
+  // in.
+  getLocations();
+  assert(locations);
+
+  getInfluences.emplace();
+  doComputeGetInfluences(*locations, *getInfluences);
+}
+
+void LazyLocalGraph::computeLocations() const {
+  // We must never repeat work.
+  assert(!locations);
+
+  // |flower| fills in |locations| as it scans the function.
+  //
+  // In theory we could be even lazier here, but it is nice that flower will
+  // fill in the locations as it goes, avoiding an additional pass. And, in
+  // practice, if we ask for locations then we likely need other things anyhow.
+  if (!flower) {
+    makeFlower();
+  }
+}
+
 } // namespace wasm
diff --git a/src/ir/local-graph.h b/src/ir/local-graph.h
index 46bc135e0..6f1e621cf 100644
--- a/src/ir/local-graph.h
+++ b/src/ir/local-graph.h
@@ -202,17 +202,43 @@ struct LazyLocalGraph : public LocalGraphBase {
     }
     return iter->second;
   }
+  const GetInfluences& getGetInfluences(LocalGet* get) const {
+    if (!getInfluences) {
+      computeGetInfluences();
+      assert(getInfluences);
+    }
+    return (*getInfluences)[get];
+  }
+
+  const Locations& getLocations() const {
+    if (!locations) {
+      computeLocations();
+      assert(locations);
+    }
+    return *locations;
+  }
 
 private:
   // These data structures are mutable so that we can memoize.
   mutable GetSetsMap getSetsMap;
 
   mutable SetInfluencesMap setInfluences;
+  // The entire |getInfluences| is computed once the first request for one
+  // arrives, so the entire thing is either present or not, unlike setInfluences
+  // which is fine-grained. The difference is that the influences of a get may
+  // include sets of other indexes, so there is no simple way to lazify that
+  // computation.
+  mutable std::optional<GetInfluencesMap> getInfluences;
+  mutable std::optional<Locations> locations;
 
   // Compute the sets for a get and store them on getSetsMap.
   void computeGetSets(LocalGet* get) const;
   // Compute influences for a set and store them on setInfluences.
   void computeSetInfluences(LocalSet* set) const;
+  // Compute influences for all gets and store them on getInfluences.
+  void computeGetInfluences() const;
+  // Compute locations and store them on getInfluences.
+  void computeLocations() const;
 
   // This remains alive as long as we are, so that we can compute things lazily.
   // It is mutable as when we construct this is an internal detail, that does