FixIrreducible: don't crash when moving a child loop

Summary:
When an irreducible SCC is converted into a new natural loop, existing
loops included in that SCC now become children of the new loop. The
logic that moves these loops from the parent loop to the new loop
invoked undefined behaviour when it modified the container that it was
iterating over. Fixed this by first extracting all the loops that are
to be removed from the parent.

Fixes bug 45623.

Reviewed By: arsenm

Differential Revision: https://reviews.llvm.org/D78544

llvm/lib/Transforms/Utils/FixIrreducible.cpp

@@ -118,14 +118,17 @@ static void reconnectChildLoops(LoopInfo &LI, Loop *ParentLoop, Loop *NewLoop,
                                 SetVector<BasicBlock *> &Headers) {
   auto &CandidateLoops = ParentLoop ? ParentLoop->getSubLoopsVector()
                                     : LI.getTopLevelLoopsVector();
-  // Partition the candidate loops into two ranges. The first part
-  // contains loops that are not children of the new loop. The second
-  // part contains children that need to be moved to the new loop.
-  auto FirstChild =
-      std::partition(CandidateLoops.begin(), CandidateLoops.end(), [&](Loop *L) {
+  // The new loop cannot be its own child, and any candidate is a
+  // child iff its header is owned by the new loop. Move all the
+  // children to a new vector.
+  auto FirstChild = std::partition(
+      CandidateLoops.begin(), CandidateLoops.end(), [&](Loop *L) {
         return L == NewLoop || Blocks.count(L->getHeader()) == 0;
       });
-  for (auto II = FirstChild, IE = CandidateLoops.end(); II != IE; ++II) {
+  SmallVector<Loop *, 8> ChildLoops(FirstChild, CandidateLoops.end());
+  CandidateLoops.erase(FirstChild, CandidateLoops.end());
+
+  for (auto II = ChildLoops.begin(), IE = ChildLoops.end(); II != IE; ++II) {
     auto Child = *II;
     LLVM_DEBUG(dbgs() << "child loop: " << Child->getHeader()->getName()
                       << "\n");
@@ -143,14 +146,10 @@ static void reconnectChildLoops(LoopInfo &LI, Loop *ParentLoop, Loop *NewLoop,
       continue;
     }
 
-    if (ParentLoop) {
-      LLVM_DEBUG(dbgs() << "removed child loop from parent\n");
-      ParentLoop->removeChildLoop(Child);
-    }
-    LLVM_DEBUG(dbgs() << "added child loop to new loop\n");
+    Child->setParentLoop(nullptr);
     NewLoop->addChildLoop(Child);
+    LLVM_DEBUG(dbgs() << "added child loop to new loop\n");
   }
-  CandidateLoops.erase(FirstChild, CandidateLoops.end());
 }
 
 // Given a set of blocks and headers in an irreducible SCC, convert it into a

llvm/test/Transforms/FixIrreducible/bug45623.ll

@@ -0,0 +1,89 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -fix-irreducible -S | FileCheck %s
+
+define dso_local void @tre_tnfa_run_backtrack() {
+; CHECK-LABEL: @tre_tnfa_run_backtrack(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br label [[RETRY:%.*]]
+; CHECK:       retry:
+; CHECK-NEXT:    br label [[IRR_GUARD:%.*]]
+; CHECK:       while.body248:
+; CHECK-NEXT:    br i1 undef, label [[IF_THEN250:%.*]], label [[IF_END275:%.*]]
+; CHECK:       if.then250:
+; CHECK-NEXT:    br label [[FOR_COND264:%.*]]
+; CHECK:       for.cond264:
+; CHECK-NEXT:    br i1 undef, label [[FOR_BODY267:%.*]], label [[IRR_GUARD]]
+; CHECK:       for.body267:
+; CHECK-NEXT:    br label [[FOR_COND264]]
+; CHECK:       if.end275:
+; CHECK-NEXT:    br label [[FOR_COND342:%.*]]
+; CHECK:       for.cond342:
+; CHECK-NEXT:    br i1 undef, label [[FOR_BODY345:%.*]], label [[FOR_END580:%.*]]
+; CHECK:       for.body345:
+; CHECK-NEXT:    br label [[FOR_COND342]]
+; CHECK:       for.end580:
+; CHECK-NEXT:    br label [[IRR_GUARD]]
+; CHECK:       backtrack:
+; CHECK-NEXT:    br i1 undef, label [[IF_THEN595:%.*]], label [[IF_ELSE629:%.*]]
+; CHECK:       if.then595:
+; CHECK-NEXT:    br label [[FOR_COND616:%.*]]
+; CHECK:       for.cond616:
+; CHECK-NEXT:    br i1 undef, label [[FOR_BODY619:%.*]], label [[FOR_END626:%.*]]
+; CHECK:       for.body619:
+; CHECK-NEXT:    br label [[FOR_COND616]]
+; CHECK:       for.end626:
+; CHECK-NEXT:    br label [[IRR_GUARD]]
+; CHECK:       if.else629:
+; CHECK-NEXT:    br label [[RETRY]]
+; CHECK:       irr.guard:
+; CHECK-NEXT:    [[GUARD_BACKTRACK:%.*]] = phi i1 [ true, [[FOR_END580]] ], [ true, [[FOR_COND264]] ], [ undef, [[RETRY]] ], [ false, [[FOR_END626]] ]
+; CHECK-NEXT:    br i1 [[GUARD_BACKTRACK]], label [[BACKTRACK:%.*]], label [[WHILE_BODY248:%.*]]
+;
+entry:
+  br label %retry
+
+retry:
+  br i1 undef, label %backtrack, label %while.body248
+
+while.body248:                                    ; preds = %for.end626, %retry
+  br i1 undef, label %if.then250, label %if.end275
+
+if.then250:                                       ; preds = %while.body248
+  br label %for.cond264
+
+for.cond264:                                      ; preds = %for.body267, %if.then250
+  br i1 undef, label %for.body267, label %backtrack
+
+for.body267:                                      ; preds = %for.cond264
+  br label %for.cond264
+
+if.end275:                                        ; preds = %while.body248
+  br label %for.cond342
+
+for.cond342:                                      ; preds = %for.body345, %if.end275
+  br i1 undef, label %for.body345, label %for.end580
+
+for.body345:                                      ; preds = %for.cond342
+  br label %for.cond342
+
+for.end580:                                       ; preds = %for.cond342
+  br label %backtrack
+
+backtrack:                                        ; preds = %for.end580, %for.cond264, %retry
+  br i1 undef, label %if.then595, label %if.else629
+
+if.then595:                                       ; preds = %backtrack
+  br label %for.cond616
+
+for.cond616:                                      ; preds = %for.body619, %if.then595
+  br i1 undef, label %for.body619, label %for.end626
+
+for.body619:                                      ; preds = %for.cond616
+  br label %for.cond616
+
+for.end626:                                       ; preds = %for.cond616
+  br label %while.body248
+
+if.else629:                                       ; preds = %backtrack
+  br label %retry
+}