[NFC] LazyLocalGraph: Add getSetInfluences() (#6909)

This new API on lazy local graphs allows us to use laziness in another place,
StackIR opts. This makes writing the binary (which includes StackIR opts, when
those are enabled), 10% faster.

1 files changed, 62 insertions, 6 deletions

diff --git a/src/ir/LocalGraph.cpp b/src/ir/LocalGraph.cpp
index fa7c976a4..11965040c 100644
--- a/src/ir/LocalGraph.cpp
+++ b/src/ir/LocalGraph.cpp
@@ -308,26 +308,41 @@ struct LocalGraphFlower
   };
   std::unordered_map<LocalGet*, BlockLocation> getLocations;
 
-  // Set up getLocations using the flow blocks, so that we are ready to handle
-  // later lazy requests for the sets of particular gets. This is done in lazy
-  // mode.
+  // In lazy mode we also need to categorize gets and sets by their index.
+  std::vector<std::vector<LocalGet*>> getsByIndex;
+  std::vector<std::vector<LocalSet*>> setsByIndex;
+
+  // Prepare for all later lazy work.
   void prepareLaziness() {
     prepareFlowBlocks();
 
+    // Set up getLocations, getsByIndex, and setsByIndex.
+    auto numLocals = func->getNumLocals();
+    getsByIndex.resize(numLocals);
+    setsByIndex.resize(numLocals);
+
     for (auto& block : flowBlocks) {
       const auto& actions = block.actions;
       for (Index i = 0; i < actions.size(); i++) {
         if (auto* get = actions[i]->dynCast<LocalGet>()) {
           getLocations[get] = BlockLocation{&block, i};
+          getsByIndex[get->index].push_back(get);
+        } else if (auto* set = actions[i]->dynCast<LocalSet>()) {
+          setsByIndex[set->index].push_back(set);
+        } else {
+          WASM_UNREACHABLE("bad action");
         }
       }
     }
   }
 
-  // Flow a specific get. This is done in lazy mode.
-  void flowGet(LocalGet* get) {
+  // Flow a specific get to its sets. This is done in lazy mode.
+  void computeGetSets(LocalGet* get) {
     auto index = get->index;
 
+    // We must never repeat work.
+    assert(!getSetsMap.count(get));
+
     // Regardless of what we do below, ensure an entry for this get, so that we
     // know we computed it.
     auto& sets = getSetsMap[get];
@@ -390,6 +405,43 @@ struct LocalGraphFlower
     // We must do an inter-block flow.
     flowBackFromStartOfBlock(block, index, gets);
   }
+
+  void computeSetInfluences(LocalSet* set,
+                            LocalGraphBase::SetInfluencesMap& setInfluences) {
+    auto index = set->index;
+
+    // We must never repeat work.
+    assert(!setInfluences.count(set));
+
+    // In theory we could flow the set forward, but to keep things simple we
+    // reuse the logic for flowing gets backwards: We flow all the gets of the
+    // set's index, thus fully computing that index and all its sets, including
+    // this one. This is not 100% lazy, but still avoids extra work by never
+    // doing work for local indexes we don't care about.
+    for (auto* get : getsByIndex[index]) {
+      // Don't repeat work.
+      if (!getSetsMap.count(get)) {
+        computeGetSets(get);
+      }
+    }
+
+    // Ensure empty entries for each set of this index, to mark them as
+    // computed.
+    for (auto* set : setsByIndex[index]) {
+      setInfluences[set];
+    }
+
+    // Also ensure |set| itself, that we were originally asked about. It may be
+    // in unreachable code, which means it is not listed in setsByIndex.
+    setInfluences[set];
+
+    // Apply the info from the gets to the sets.
+    for (auto* get : getsByIndex[index]) {
+      for (auto* set : getSetsMap[get]) {
+        setInfluences[set].insert(get);
+      }
+    }
+  }
 };
 
 // LocalGraph implementation
@@ -526,7 +578,11 @@ LazyLocalGraph::~LazyLocalGraph() {
 }
 
 void LazyLocalGraph::computeGetSets(LocalGet* get) const {
-  flower->flowGet(get);
+  flower->computeGetSets(get);
+}
+
+void LazyLocalGraph::computeSetInfluences(LocalSet* set) const {
+  flower->computeSetInfluences(set, setInfluences);
 }
 
 } // namespace wasm