Collapse inner loops with union-find during WTO construction

src/cfg/wto.h

@@ -66,7 +66,10 @@
 //      visit every block in `h`'s natural loop. Because inner loop headers have
 //      larger RPO indices than outer loop headers and are processed first, the
 //      first loop that visits a block `b != h` is its immediately enclosing
-//      loop (`loopParent[b] = h`).
+//      loop (`loopParent[b] = h`). As each loop body is discovered, we collapse
+//      its blocks into `h` using union-find so that outer loops skip over
+//      already-collapsed inner loop bodies instead of re-traversing them (both
+//      during the backward DFS and when walking the dominator tree).
 //   3. Link each reachable block into the child list of its `loopParent` in
 //      increasing RPO order, then walk the resulting loop nesting forest to
 //      emit each loop header `h` and its children as a nested `Cycle`.
@@ -118,63 +121,81 @@ WeakTopologicalOrdering<BasicBlock>::WeakTopologicalOrdering(
     return i == 0 || domTree.iDoms[i] != domTree.nonsense;
   };
 
+  static constexpr Index NoIndex = Index(-1);
+  struct Node {
+    Index loopParent = NoIndex;
+    Index firstChild = NoIndex;
+    Index nextSibling = NoIndex;
+    Index ufParent = NoIndex;
+    bool isLoopHeader = false;
+  };
+  std::vector<Node> nodes(numBlocks);
+
+  auto find = [&](Index x) {
+    Index root = x;
+    while (nodes[root].ufParent != NoIndex) {
+      root = nodes[root].ufParent;
+    }
+    // Path compression.
+    while (x != root) {
+      Index next = nodes[x].ufParent;
+      nodes[x].ufParent = root;
+      x = next;
+    }
+    return root;
+  };
+
   auto dominates = [&](Index dom, Index node) {
     assert(isReachable(dom));
-    if (!isReachable(node)) {
+    // Since blocks are indexed in RPO, dominators always precede the blocks
+    // they dominate.
+    if (node < dom || !isReachable(node)) {
       return false;
     }
+    // Walk up the dominator tree, using `find` to skip over already-collapsed
+    // inner loops.
     Index curr = node;
     while (curr > dom) {
-      curr = domTree.iDoms[curr];
+      curr = find(domTree.iDoms[curr]);
     }
     return curr == dom;
   };
 
-  static constexpr Index NoIndex = Index(-1);
-  struct Node {
-    Index loopParent = NoIndex;
-    Index firstChild = NoIndex;
-    Index nextSibling = NoIndex;
-    Index lastVisitedBy = NoIndex;
-    bool isLoopHeader = false;
-  };
-  std::vector<Node> nodes(numBlocks);
-
   // Discover natural loops from innermost to outermost (reverse RPO order).
-  // Because inner loops are processed before outer loops, the first loop whose
-  // natural loop body contains a block is its immediately enclosing loop.
-  //
-  // TODO: Collapse inner loops with union-find during natural loop discovery so
-  // outer loops do not re-traverse inner loop bodies.
+  // Because inner loops are processed before outer loops, collapsing each loop
+  // body into its header with union-find records each block's immediately
+  // enclosing loop while avoiding re-traversing inner loop bodies.
   std::vector<Index> worklist;
   for (Index i = numBlocks; i > 0; --i) {
     Index h = i - 1;
     if (!isReachable(h)) {
       continue;
     }
-    nodes[h].lastVisitedBy = h;
     for (auto* pred : blocks[h]->in) {
       Index p = pred->contents.index;
       if (dominates(h, p)) {
         nodes[h].isLoopHeader = true;
-        if (nodes[p].lastVisitedBy != h) {
-          nodes[p].lastVisitedBy = h;
-          worklist.push_back(p);
+        Index rep = find(p);
+        if (rep != h) {
+          nodes[rep].loopParent = h;
+          nodes[rep].ufParent = h;
+          worklist.push_back(rep);
         }
       }
     }
     while (!worklist.empty()) {
       Index curr = worklist.back();
       worklist.pop_back();
-      if (nodes[curr].loopParent == NoIndex) {
-        nodes[curr].loopParent = h;
-      }
       for (auto* pred : blocks[curr]->in) {
         Index p = pred->contents.index;
-        if (isReachable(p) && nodes[p].lastVisitedBy != h) {
-          assert(dominates(h, p) && "Expected reducible CFG");
-          nodes[p].lastVisitedBy = h;
-          worklist.push_back(p);
+        if (isReachable(p)) {
+          Index rep = find(p);
+          if (rep != h) {
+            assert(dominates(h, rep) && "Expected reducible CFG");
+            nodes[rep].loopParent = h;
+            nodes[rep].ufParent = h;
+            worklist.push_back(rep);
+          }
         }
       }
     }