1 files changed, 10 insertions, 5 deletions
diff --git a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
index cb8614f..82b9979 100644
--- a/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
+++ b/llvm/lib/Transforms/Utils/BasicBlockUtils.cpp
@@ -47,12 +47,15 @@
 
 using namespace llvm;
 
-void llvm::DeleteDeadBlock(BasicBlock *BB, DomTreeUpdater *DTU) {
-  DeleteDeadBlocks({BB}, DTU);
+void llvm::DeleteDeadBlock(
+    BasicBlock *BB, DomTreeUpdater *DTU,
+    SmallVectorImpl<DominatorTree::UpdateType> *DTUpdates) {
+  DeleteDeadBlocks({BB}, DTU, DTUpdates);
 }
 
-void llvm::DeleteDeadBlocks(ArrayRef <BasicBlock *> BBs,
-                            DomTreeUpdater *DTU) {
+void llvm::DeleteDeadBlocks(
+    ArrayRef<BasicBlock *> BBs, DomTreeUpdater *DTU,
+    SmallVectorImpl<DominatorTree::UpdateType> *DTUpdates) {
 #ifndef NDEBUG
   // Make sure that all predecessors of each dead block is also dead.
   SmallPtrSet<BasicBlock *, 4> Dead(BBs.begin(), BBs.end());
@@ -68,7 +71,7 @@ void llvm::DeleteDeadBlocks(ArrayRef <BasicBlock *> BBs,
     // of their predecessors is going away.
     for (BasicBlock *Succ : successors(BB)) {
       Succ->removePredecessor(BB);
-      if (DTU)
+      if (DTU || DTUpdates)
         Updates.push_back({DominatorTree::Delete, BB, Succ});
     }
 
@@ -92,6 +95,8 @@ void llvm::DeleteDeadBlocks(ArrayRef <BasicBlock *> BBs,
   }
   if (DTU)
     DTU->applyUpdates(Updates, /*ForceRemoveDuplicates*/ true);
+  if (DTUpdates)
+    DTUpdates->append(Updates.begin(), Updates.end());
 
   for (BasicBlock *BB : BBs)
     if (DTU)