2 files changed, 64 insertions, 4 deletions
diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index c28311fe4b0..fed5e62157a 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1333,14 +1333,15 @@ static SmallPtrSet<const BasicBlock *, 16> recomputeLoopBlockSet(Loop &L,
   if (LoopBlockSet.empty())
     return LoopBlockSet;
 
-  // Add the loop header to the set.
-  LoopBlockSet.insert(Header);
-
   // We found backedges, recurse through them to identify the loop blocks.
   while (!Worklist.empty()) {
     BasicBlock *BB = Worklist.pop_back_val();
     assert(LoopBlockSet.count(BB) && "Didn't put block into the loop set!");
 
+    // No need to walk past the header.
+    if (BB == Header)
+      continue;
+
     // Because we know the inner loop structure remains valid we can use the
     // loop structure to jump immediately across the entire nested loop.
     // Further, because it is in loop simplified form, we can directly jump
@@ -1388,6 +1389,8 @@ static SmallPtrSet<const BasicBlock *, 16> recomputeLoopBlockSet(Loop &L,
         Worklist.push_back(Pred);
   }
 
+  assert(LoopBlockSet.count(Header) && "Cannot fail to add the header!");
+
   // We've found all the blocks participating in the loop, return our completed
   // set.
   return LoopBlockSet;
@@ -1792,9 +1795,12 @@ static bool unswitchInvariantBranch(
   // unnecessary loops.
   auto UpdateLCSSA = [&](Loop &UpdateL) {
 #ifndef NDEBUG
-    for (Loop *ChildL : UpdateL)
+    UpdateL.verifyLoop();
+    for (Loop *ChildL : UpdateL) {
+      ChildL->verifyLoop();
       assert(ChildL->isRecursivelyLCSSAForm(DT, LI) &&
              "Perturbed a child loop's LCSSA form!");
+    }
 #endif
     formLCSSA(UpdateL, DT, &LI, nullptr);
   };
diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
index 48075172b10..51afdf386e5 100644
--- a/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/nontrivial-unswitch.ll
@@ -2508,3 +2508,57 @@ loop1.exit:
   call void @g()
   ret void
 }
+
+; Test that when we are unswitching and need to rebuild the loop block set we
+; correctly skip past inner loops. We want to use the inner loop to efficiently
+; skip whole subregions of the outer loop blocks but just because the header of
+; the outer loop is also the preheader of an inner loop shouldn't confuse this
+; walk.
+define void @test23(i1 %arg, i1* %ptr) {
+; CHECK-LABEL: define void @test23(
+entry:
+  br label %outer.header
+; CHECK:       entry:
+; CHECK-NEXT:    br i1 %arg,
+;
+; Just verify that we unswitched the correct bits. We should call `@f` twice in
+; one unswitch and `@f` and then `@g` in the other.
+; CHECK:         call void
+; CHECK-SAME:              @f
+; CHECK:         call void
+; CHECK-SAME:              @f
+;
+; CHECK:         call void
+; CHECK-SAME:              @f
+; CHECK:         call void
+; CHECK-SAME:              @g
+
+outer.header:
+  br label %inner.header
+
+inner.header:
+  call void @f()
+  br label %inner.latch
+
+inner.latch:
+  %inner.cond = load i1, i1* %ptr
+  br i1 %inner.cond, label %inner.header, label %outer.body
+
+outer.body:
+  br i1 %arg, label %outer.body.left, label %outer.body.right
+
+outer.body.left:
+  call void @f()
+  br label %outer.latch
+
+outer.body.right:
+  call void @g()
+  br label %outer.latch
+
+outer.latch:
+  %outer.cond = load i1, i1* %ptr
+  br i1 %outer.cond, label %outer.header, label %exit
+
+exit:
+  ret void
+}