Index: llvm/test/tools/llvm-reduce/remove-bbs-illegal2.ll
===================================================================
--- /dev/null
+++ llvm/test/tools/llvm-reduce/remove-bbs-illegal2.ll
@@ -0,0 +1,28 @@
+; RUN: llvm-reduce --delta-passes=basic-blocks --test %python --test-arg %p/Inputs/remove-bbs.py -abort-on-invalid-reduction %s -o %t
+
+declare void @0()
+
+define internal void @1(ptr %x0, i1 %x1) {
+interesting0:
+  %x3 = alloca i32, align 4
+  store ptr null, ptr %x0, align 8
+  br label %interesting1
+
+interesting1:                                                ; preds = %x2
+  call void @0()
+  br label %x2
+
+x2:
+  br i1 %x1, label %interesting3, label %interesting4
+
+interesting3:
+  call void @0()
+  br label %x2
+
+interesting4:
+  br label %x5
+
+x5:
+  store i32 0, ptr %x3, align 4
+  ret void
+}
Index: llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp
===================================================================
--- llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp
+++ llvm/tools/llvm-reduce/deltas/ReduceBasicBlocks.cpp
@@ -102,35 +102,48 @@
     }
 }
 
-/// A BB is ok to remove if it's not the entry block, or else it is
-/// the entry block but the next block in the function has just one
-/// predecessor -- this property is required because that block is
-/// going to become the new entry block
-static bool okToRemove(BasicBlock &BB) {
-  if (!BB.isEntryBlock())
+/// It's OK to add a block to the set of removed blocks if the first
+/// basic block in the function that survives all of the deletions is
+/// a legal entry block
+static bool okToRemove(BasicBlock &Candidate, Function &F,
+                       const DenseSet<BasicBlock *> &BBsToDelete) {
+  for (auto &B : F) {
+    if (&B == &Candidate)
+      continue;
+    if (BBsToDelete.count(&B))
+      continue;
+    /// Ok we've found the first block that's not going to be deleted,
+    /// it will be the new entry block -- that's only legal if this
+    /// block has no predecessors among blocks that survive the
+    /// deletions
+    for (BasicBlock *Pred : predecessors(&B)) {
+      if (!BBsToDelete.contains(Pred))
+        return false;
+    }
     return true;
-  auto F = BB.getParent();
-  auto it = F->begin();
-  ++it;
-  return (it == F->end()) || (*it).hasNPredecessors(1);
+  }
+  return true;
 }
 
 /// Removes out-of-chunk arguments from functions, and modifies their calls
 /// accordingly. It also removes allocations of out-of-chunk arguments.
 static void extractBasicBlocksFromModule(Oracle &O, Module &Program) {
-  DenseSet<BasicBlock *> BBsToKeep;
+  DenseSet<BasicBlock *> BBsToKeep, BBsToDelete;
 
-  SmallVector<BasicBlock *> BBsToDelete;
   for (auto &F : Program) {
     for (auto &BB : F) {
-      if (!okToRemove(BB) || O.shouldKeep())
+      if (!okToRemove(BB, F, BBsToDelete) || O.shouldKeep())
         BBsToKeep.insert(&BB);
-      else {
-        BBsToDelete.push_back(&BB);
-        // Remove out-of-chunk BB from successor phi nodes
-        for (auto *Succ : successors(&BB))
-          Succ->removePredecessor(&BB);
-      }
+      else
+        BBsToDelete.insert(&BB);
+    }
+  }
+
+  // Remove out-of-chunk BB from successor phi nodes
+  for (auto &F : Program) {
+    for (auto &BB : F) {
+      for (auto *Succ : successors(&BB))
+        Succ->removePredecessor(&BB);
     }
   }