2 files changed, 47 insertions, 0 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;