3 files changed, 52 insertions, 7 deletions

diff --git a/src/ir/LocalGraph.cpp b/src/ir/LocalGraph.cpp
index a5495b5b3..7e877e65e 100644
--- a/src/ir/LocalGraph.cpp
+++ b/src/ir/LocalGraph.cpp
@@ -624,6 +624,38 @@ void LazyLocalGraph::computeGetInfluences() const {
   doComputeGetInfluences(*locations, *getInfluences);
 }
 
+bool LazyLocalGraph::computeSSA(Index index) const {
+  // We must never repeat work.
+  assert(!SSAIndexes.count(index));
+
+  if (!flower) {
+    makeFlower();
+  }
+
+  // Similar logic to LocalGraph::computeSSAIndexes(), but optimized for the
+  // case of a single index.
+
+  // All the sets for this index that we've seen. We'll add all relevant ones,
+  // and exit if we see more than one.
+  SmallUnorderedSet<LocalSet*, 2> sets;
+  for (auto* set : flower->setsByIndex[index]) {
+    sets.insert(set);
+    if (sets.size() > 1) {
+      return SSAIndexes[index] = false;
+    }
+  }
+  for (auto* get : flower->getsByIndex[index]) {
+    for (auto* set : getSets(get)) {
+      sets.insert(set);
+      if (sets.size() > 1) {
+        return SSAIndexes[index] = false;
+      }
+    }
+  }
+  // Finally, check that we have 1 and not 0 sets.
+  return SSAIndexes[index] = (sets.size() == 1);
+}
+
 void LazyLocalGraph::computeLocations() const {
   // We must never repeat work.
   assert(!locations);
diff --git a/src/ir/local-graph.h b/src/ir/local-graph.h
index 6f1e621cf..7b8001be6 100644
--- a/src/ir/local-graph.h
+++ b/src/ir/local-graph.h
@@ -209,6 +209,16 @@ struct LazyLocalGraph : public LocalGraphBase {
     }
     return (*getInfluences)[get];
   }
+  bool isSSA(Index index) const {
+    auto iter = SSAIndexes.find(index);
+    if (iter == SSAIndexes.end()) {
+      auto ret = computeSSA(index);
+      // The result must have been memoized.
+      assert(SSAIndexes.count(index));
+      return ret;
+    }
+    return iter->second;
+  }
 
   const Locations& getLocations() const {
     if (!locations) {
@@ -229,6 +239,9 @@ private:
   // include sets of other indexes, so there is no simple way to lazify that
   // computation.
   mutable std::optional<GetInfluencesMap> getInfluences;
+  // A map if indexes to a bool indicating if the index is SSA. If there is no
+  // entry, it has not yet been computed.
+  mutable std::unordered_map<Index, bool> SSAIndexes;
   mutable std::optional<Locations> locations;
 
   // Compute the sets for a get and store them on getSetsMap.
@@ -237,6 +250,8 @@ private:
   void computeSetInfluences(LocalSet* set) const;
   // Compute influences for all gets and store them on getInfluences.
   void computeGetInfluences() const;
+  // Compute whether an index is SSA and store that on SSAIndexes.
+  bool computeSSA(Index index) const;
   // Compute locations and store them on getInfluences.
   void computeLocations() const;
 
diff --git a/src/passes/OptimizeAddedConstants.cpp b/src/passes/OptimizeAddedConstants.cpp
index 142f10024..fd5840836 100644
--- a/src/passes/OptimizeAddedConstants.cpp
+++ b/src/passes/OptimizeAddedConstants.cpp
@@ -66,7 +66,7 @@ public:
   MemoryAccessOptimizer(P* parent,
                         T* curr,
                         Module* module,
-                        LocalGraph* localGraph)
+                        LazyLocalGraph* localGraph)
     : parent(parent), curr(curr), module(module), localGraph(localGraph) {
     memory64 = module->getMemory(curr->memory)->is64();
   }
@@ -133,7 +133,7 @@ private:
   P* parent;
   T* curr;
   Module* module;
-  LocalGraph* localGraph;
+  LazyLocalGraph* localGraph;
   bool memory64;
 
   void optimizeConstantPointer() {
@@ -326,9 +326,7 @@ struct OptimizeAddedConstants
       helperIndexes.clear();
       propagatable.clear();
       if (propagate) {
-        localGraph = std::make_unique<LocalGraph>(func, getModule());
-        localGraph->computeSetInfluences();
-        localGraph->computeSSAIndexes();
+        localGraph = std::make_unique<LazyLocalGraph>(func, getModule());
         findPropagatable();
       }
       Super::doWalkFunction(func);
@@ -362,7 +360,7 @@ struct OptimizeAddedConstants
 private:
   bool propagated;
 
-  std::unique_ptr<LocalGraph> localGraph;
+  std::unique_ptr<LazyLocalGraph> localGraph;
 
   // Whether a set is propagatable.
   std::set<LocalSet*> propagatable;
@@ -379,7 +377,7 @@ private:
     // but if x has other uses, then avoid doing so - we'll be doing that add
     // anyhow, so the load/store offset trick won't actually help.
     GetParents parents(getFunction()->body);
-    for (auto& [location, _] : localGraph->locations) {
+    for (auto& [location, _] : localGraph->getLocations()) {
       if (auto* set = location->dynCast<LocalSet>()) {
         if (auto* add = set->value->dynCast<Binary>()) {
           if (add->op == AddInt32) {