4 files changed, 88 insertions, 23 deletions

diff --git a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
index a7bbe2856c5..6ee649daa3f 100644
--- a/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/BasicBlockUtils.h
@@ -62,6 +62,12 @@ void DeleteDeadBlocks(ArrayRef <BasicBlock *> BBs,
                       DomTreeUpdater *DTU = nullptr,
                       bool KeepOneInputPHIs = false);
 
+/// Delete all basic blocks from \p F that are not reachable from its entry
+/// node. If \p KeepOneInputPHIs is true, one-input Phis in successors of
+/// blocks being deleted will be preserved.
+bool EliminateUnreachableBlocks(Function &F, DomTreeUpdater *DTU = nullptr,
+                                bool KeepOneInputPHIs = false);
+
 /// We know that BB has one predecessor. If there are any single-entry PHI nodes
 /// in it, fold them away. This handles the case when all entries to the PHI
 /// nodes in a block are guaranteed equal, such as when the block has exactly
diff --git a/llvm/lib/CodeGen/UnreachableBlockElim.cpp b/llvm/lib/CodeGen/UnreachableBlockElim.cpp
index 57e548c703a..177bab32bcc 100644
--- a/llvm/lib/CodeGen/UnreachableBlockElim.cpp
+++ b/llvm/lib/CodeGen/UnreachableBlockElim.cpp
@@ -40,31 +40,10 @@
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 using namespace llvm;
 
-static bool eliminateUnreachableBlock(Function &F) {
-  df_iterator_default_set<BasicBlock*> Reachable;
-
-  // Mark all reachable blocks.
-  for (BasicBlock *BB : depth_first_ext(&F, Reachable))
-    (void)BB/* Mark all reachable blocks */;
-
-  // Collect all dead blocks.
-  std::vector<BasicBlock*> DeadBlocks;
-  for (Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
-    if (!Reachable.count(&*I)) {
-      BasicBlock *BB = &*I;
-      DeadBlocks.push_back(BB);
-    }
-
-  // Delete the dead blocks.
-  DeleteDeadBlocks(DeadBlocks);
-
-  return !DeadBlocks.empty();
-}
-
 namespace {
 class UnreachableBlockElimLegacyPass : public FunctionPass {
   bool runOnFunction(Function &F) override {
-    return eliminateUnreachableBlock(F);
+    return llvm::EliminateUnreachableBlocks(F);
   }
 
 public:
@@ -89,7 +68,7 @@ FunctionPass *llvm::createUnreachableBlockEliminationPass() {
 
 PreservedAnalyses UnreachableBlockElimPass::run(Function &F,
                                                 FunctionAnalysisManager &AM) {
-  bool Changed = eliminateUnreachableBlock(F);
+  bool Changed = llvm::EliminateUnreachableBlocks(F);
   if (!Changed)
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index 81031b1a61e..aa7b933022c 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -110,6 +110,28 @@ void llvm::DeleteDeadBlocks(ArrayRef <BasicBlock *> BBs, DomTreeUpdater *DTU,
       BB->eraseFromParent();
 }
 
+bool llvm::EliminateUnreachableBlocks(Function &F, DomTreeUpdater *DTU,
+                                      bool KeepOneInputPHIs) {
+  df_iterator_default_set<BasicBlock*> Reachable;
+
+  // Mark all reachable blocks.
+  for (BasicBlock *BB : depth_first_ext(&F, Reachable))
+    (void)BB/* Mark all reachable blocks */;
+
+  // Collect all dead blocks.
+  std::vector<BasicBlock*> DeadBlocks;
+  for (Function::iterator I = F.begin(), E = F.end(); I != E; ++I)
+    if (!Reachable.count(&*I)) {
+      BasicBlock *BB = &*I;
+      DeadBlocks.push_back(BB);
+    }
+
+  // Delete the dead blocks.
+  DeleteDeadBlocks(DeadBlocks, DTU, KeepOneInputPHIs);
+
+  return !DeadBlocks.empty();
+}
+
 void llvm::FoldSingleEntryPHINodes(BasicBlock *BB,
                                    MemoryDependenceResults *MemDep) {
   if (!isa<PHINode>(BB->begin())) return;
diff --git a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
index 2d3731c03e0..bf8228e99d8 100644
--- a/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
+++ b/llvm/unittests/Transforms/Utils/BasicBlockUtilsTest.cpp
@@ -25,6 +25,64 @@ static std::unique_ptr<Module> parseIR(LLVMContext &C, const char *IR) {
   return Mod;
 }
 
+TEST(BasicBlockUtils, EliminateUnreachableBlocks) {
+  LLVMContext C;
+
+  std::unique_ptr<Module> M = parseIR(
+    C,
+    "define i32 @has_unreachable(i1 %cond) {\n"
+    "entry:\n"
+    "  br i1 %cond, label %bb0, label %bb1\n"
+    "bb0:\n"
+    "  br label %bb1\n"
+    "bb1:\n"
+    "  %phi = phi i32 [ 0, %entry ], [ 1, %bb0 ]"
+    "  ret i32 %phi\n"
+    "bb2:\n"
+    "  ret i32 42\n"
+    "}\n"
+    "\n"
+    );
+
+  auto *F = M->getFunction("has_unreachable");
+  DominatorTree DT(*F);
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+
+  EXPECT_EQ(F->size(), (size_t)4);
+  bool Result = EliminateUnreachableBlocks(*F, &DTU);
+  EXPECT_TRUE(Result);
+  EXPECT_EQ(F->size(), (size_t)3);
+  EXPECT_TRUE(DT.verify());
+}
+
+TEST(BasicBlockUtils, NoUnreachableBlocksToEliminate) {
+  LLVMContext C;
+
+  std::unique_ptr<Module> M = parseIR(
+    C,
+    "define i32 @no_unreachable(i1 %cond) {\n"
+    "entry:\n"
+    "  br i1 %cond, label %bb0, label %bb1\n"
+    "bb0:\n"
+    "  br label %bb1\n"
+    "bb1:\n"
+    "  %phi = phi i32 [ 0, %entry ], [ 1, %bb0 ]"
+    "  ret i32 %phi\n"
+    "}\n"
+    "\n"
+    );
+
+  auto *F = M->getFunction("no_unreachable");
+  DominatorTree DT(*F);
+  DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
+
+  EXPECT_EQ(F->size(), (size_t)3);
+  bool Result = EliminateUnreachableBlocks(*F, &DTU);
+  EXPECT_FALSE(Result);
+  EXPECT_EQ(F->size(), (size_t)3);
+  EXPECT_TRUE(DT.verify());
+}
+
 TEST(BasicBlockUtils, SplitBlockPredecessors) {
   LLVMContext C;