2 files changed, 75 insertions, 19 deletions

diff --git a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
index 3cc2c651f6c..57799c8cd8c 100644
--- a/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
+++ b/llvm/lib/Transforms/Scalar/SimpleLoopUnswitch.cpp
@@ -1378,25 +1378,25 @@ static void
 deleteDeadBlocksFromLoop(Loop &L,
                          SmallVectorImpl<BasicBlock *> &ExitBlocks,
                          DominatorTree &DT, LoopInfo &LI) {
-  // Find all the dead blocks, and remove them from their successors.
-  SmallVector<BasicBlock *, 16> DeadBlocks;
-  for (BasicBlock *BB : ExitBlocks)
-    if (!DT.isReachableFromEntry(BB)) {
-      for (BasicBlock *SuccBB : successors(BB))
+  // Find all the dead blocks tied to this loop, and remove them from their
+  // successors.
+  SmallPtrSet<BasicBlock *, 16> DeadBlockSet;
+
+  // Start with loop/exit blocks and get a transitive closure of reachable dead
+  // blocks.
+  SmallVector<BasicBlock *, 16> DeathCandidates(ExitBlocks.begin(),
+                                                ExitBlocks.end());
+  DeathCandidates.append(L.blocks().begin(), L.blocks().end());
+  while (!DeathCandidates.empty()) {
+    auto *BB = DeathCandidates.pop_back_val();
+    if (!DeadBlockSet.count(BB) && !DT.isReachableFromEntry(BB)) {
+      for (BasicBlock *SuccBB : successors(BB)) {
         SuccBB->removePredecessor(BB);
-      DeadBlocks.push_back(BB);
-    }
-
-  for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop())
-    for (BasicBlock *BB : ParentL->blocks())
-      if (!DT.isReachableFromEntry(BB)) {
-        for (BasicBlock *SuccBB : successors(BB))
-          SuccBB->removePredecessor(BB);
-        DeadBlocks.push_back(BB);
+        DeathCandidates.push_back(SuccBB);
       }
-
-  SmallPtrSet<BasicBlock *, 16> DeadBlockSet(DeadBlocks.begin(),
-                                             DeadBlocks.end());
+      DeadBlockSet.insert(BB);
+    }
+  }
 
   // Filter out the dead blocks from the exit blocks list so that it can be
   // used in the caller.
@@ -1405,7 +1405,7 @@ deleteDeadBlocksFromLoop(Loop &L,
 
   // Walk from this loop up through its parents removing all of the dead blocks.
   for (Loop *ParentL = &L; ParentL; ParentL = ParentL->getParentLoop()) {
-    for (auto *BB : DeadBlocks)
+    for (auto *BB : DeadBlockSet)
       ParentL->getBlocksSet().erase(BB);
     llvm::erase_if(ParentL->getBlocksVector(),
                    [&](BasicBlock *BB) { return DeadBlockSet.count(BB); });
@@ -1430,7 +1430,7 @@ deleteDeadBlocksFromLoop(Loop &L,
   // Remove the loop mappings for the dead blocks and drop all the references
   // from these blocks to others to handle cyclic references as we start
   // deleting the blocks themselves.
-  for (auto *BB : DeadBlocks) {
+  for (auto *BB : DeadBlockSet) {
     // Check that the dominator tree has already been updated.
     assert(!DT.getNode(BB) && "Should already have cleared domtree!");
     LI.changeLoopFor(BB, nullptr);
diff --git a/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll b/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll
index 0be1c8549ed..aa10670cb4e 100644
--- a/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll
+++ b/llvm/test/Transforms/SimpleLoopUnswitch/delete-dead-blocks.ll
@@ -43,3 +43,59 @@ get_out:
 get_out2:
   unreachable
 }
+
+;
+; This comes from PR38778
+; CHECK-LABEL: @Test2
+define void @Test2(i32) {
+header:
+  br label %loop
+loop:
+  switch i32 %0, label %continue [
+    i32 -2147483648, label %check
+    i32 98, label %guarded1
+    i32 99, label %guarded2
+  ]
+; CHECK-NOT: {{^}}guarded1:
+guarded1:
+  br i1 undef, label %continue, label %leave
+guarded2:
+  br label %continue
+check:
+  %val = add i32 0, 1
+  br i1 undef, label %continue, label %leave
+continue:
+  br label %loop
+leave:
+  %local = phi i32 [ 0, %guarded1 ], [ %val, %check ]
+  ret void
+}
+
+;
+; Yet another test from PR38778
+;
+; CHECK-LABEL: @Test3
+define void @Test3(i32) {
+header:
+  br label %outer
+outer:
+  %bad_input.i = icmp eq i32 %0, -2147483648
+  br label %inner
+inner:
+  br i1 %bad_input.i, label %overflow, label %switchme
+overflow:
+  br label %continue
+switchme:
+  switch i32 %0, label %continue [
+    i32 88, label %go_out
+    i32 99, label %case2
+  ]
+; CHECK-NOT: {{^}}case2:
+case2:
+  br label %continue
+continue:
+  %local_11_92 = phi i32 [ 0, %switchme ], [ 18, %case2 ], [ 0, %overflow ]
+  br i1 undef, label %outer, label %inner
+go_out:
+  unreachable
+}