[llvm] [ORC] Refactor some WaitingOnGraph loops into named methods. NFCI. (PR #184775)

[Lang Hames via llvm-commits]

https://github.com/lhames created https://github.com/llvm/llvm-project/pull/184775

Adds new methods to SuperNode for updating ElemToSuperNodeMap values, use it to tidy up some loops.

>From ea898779840f7a2d305c6eb30ad2a4c5fbe69f2d Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Thu, 5 Mar 2026 22:03:26 +1100
Subject: [PATCH] [ORC] Refactor some WaitingOnGraph loops into named methods.
 NFCI.

Adds new methods to SuperNode for updating ElemToSuperNodeMap values,
use it to tidy up some loops.
---
 .../llvm/ExecutionEngine/Orc/WaitingOnGraph.h | 82 +++++++++++++------
 1 file changed, 55 insertions(+), 27 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
index 88371b557fe7c..0f3f67f8c4fd5 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
@@ -88,6 +88,13 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
   using ElementSet = DenseSet<ElementId>;
   using ContainerElementsMap = DenseMap<ContainerId, ElementSet>;
 
+  class SuperNode;
+
+private:
+  using ElemToSuperNodeMap =
+      DenseMap<ContainerId, DenseMap<ElementId, SuperNode *>>;
+
+public:
   class SuperNode {
     friend class WaitingOnGraph;
     friend class WaitingOnGraphTest;
@@ -103,12 +110,50 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
   private:
     ContainerElementsMap Defs;
     ContainerElementsMap Deps;
+
+    ElemToSuperNodeMap *RegisteredElemToSN = nullptr;
+
+    /// Add a mapping from the Defs in this SuperNode to SN (which may or may
+    /// not be the same as this).
+    void mapDefsTo(ElemToSuperNodeMap &ElemToSN, SuperNode *SN) {
+      assert(!Defs.empty() && "Empty defs!?");
+      for (auto &[Container, Elements] : Defs) {
+        assert(!Elements.empty() && "Empty elements for container?");
+        auto &ContainerElemToSN = ElemToSN[Container];
+        for (auto &Elem : Elements)
+          ContainerElemToSN[Elem] = SN;
+      }
+      assert((!SN->RegisteredElemToSN || SN->RegisteredElemToSN == &ElemToSN) &&
+             "SN defs split across maps");
+      SN->RegisteredElemToSN = &ElemToSN;
+    }
+
+    /// Add a mapping from the Defs in this SuperNode to this.
+    /// (Equivalent to `SN.mapDefsTo(ElemToSN, &SN);`)
+    void mapDefsToThis(ElemToSuperNodeMap &ElemToSN) {
+      mapDefsTo(ElemToSN, this);
+    }
+
+    /// Remove a mapping from the Defs in this SuperNode from the registered
+    /// ElemToSuperNodeMap. The mapping must already exist.
+    void unmapDefsFromThis() {
+      assert(RegisteredElemToSN && "No registered ElemToSuperNodeMap");
+      for (auto &[Container, Elements] : Defs) {
+        auto I = RegisteredElemToSN->find(Container);
+        assert(I != RegisteredElemToSN->end() && "Container not in map");
+        auto &ContainerElemToSN = I->second;
+        for (auto &Elem : Elements) {
+          assert(ContainerElemToSN[Elem] == this && "Mapping not present");
+          ContainerElemToSN.erase(Elem);
+        }
+        if (ContainerElemToSN.empty())
+          RegisteredElemToSN->erase(I);
+      }
+      RegisteredElemToSN = nullptr;
+    }
   };
 
 private:
-  using ElemToSuperNodeMap =
-      DenseMap<ContainerId, DenseMap<ElementId, SuperNode *>>;
-
   using SuperNodeDepsMap = DenseMap<SuperNode *, DenseSet<SuperNode *>>;
 
   class Coalescer {
@@ -143,12 +188,9 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
                "Elements of SNs should be new to the coalescer");
         auto H = getHash(SN->Deps);
         if (auto *CanonicalSN = findCanonicalSuperNode(H, SN->Deps)) {
-          for (auto &[Container, Elems] : SN->Defs) {
+          SN->mapDefsTo(ElemToSN, CanonicalSN);
+          for (auto &[Container, Elems] : SN->Defs)
             CanonicalSN->Defs[Container].insert(Elems.begin(), Elems.end());
-            auto &ContainerElemToSN = ElemToSN[Container];
-            for (auto &Elem : Elems)
-              ContainerElemToSN[Elem] = CanonicalSN;
-          }
           std::swap(SN, SNs.back());
           SNs.pop_back();
         } else {
@@ -286,13 +328,8 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
   static SimplifyResult simplify(std::vector<std::unique_ptr<SuperNode>> SNs) {
     // Build ElemToSN map.
     ElemToSuperNodeMap ElemToSN;
-    for (auto &SN : SNs) {
-      for (auto &[Container, Elements] : SN->Defs) {
-        auto &ContainerElemToSN = ElemToSN[Container];
-        for (auto &E : Elements)
-          ContainerElemToSN[E] = SN.get();
-      }
-    }
+    for (auto &SN : SNs)
+      SN->mapDefsToThis(ElemToSN);
 
     SuperNodeDepsMap SuperNodeDeps;
     hoistDeps(SuperNodeDeps, SNs, ElemToSN);
@@ -380,11 +417,7 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
 
     // Update ElemToPendingSN for the remaining elements.
     for (auto &SN : NewSNs) {
-      for (auto &[Container, Elems] : SN->Defs) {
-        auto &Row = ElemToPendingSN[Container];
-        for (auto &Elem : Elems)
-          Row[Elem] = SN.get();
-      }
+      SN->mapDefsToThis(ElemToPendingSN);
       PendingSNs.push_back(std::move(SN));
     }
 
@@ -692,13 +725,8 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
     }
 
     // Update ElemToSNs (if passed) to remove elements pointing at SN.
-    for (auto *SN : ToRemoveFromElemToSNs) {
-      for (auto &[Container, Elems] : SN->defs()) {
-        auto &Row = (*ElemToSNs)[Container];
-        for (auto &Elem : Elems)
-          Row.erase(Elem);
-      }
-    }
+    for (auto *SN : ToRemoveFromElemToSNs)
+      SN->unmapDefsFromThis();
   }
 
   std::vector<std::unique_ptr<SuperNode>> PendingSNs;