[JumpThreading] Preserve DT and LVI across the pass

JumpThreading now preserves dominance and lazy value information across the
entire pass.  The pass manager is also informed of this preservation with
the goal of DT and LVI being recalculated fewer times overall during
compilation.

This change prepares JumpThreading for enhanced opportunities; particularly
those across loop boundaries.

Patch by: Brian Rzycki <b.rzycki@samsung.com>,
          Sebastian Pop <s.pop@samsung.com>

Differential revision: https://reviews.llvm.org/D37528

llvm-svn: 314435

3 files changed, 253 insertions, 65 deletions

diff --git a/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp b/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
index 28157783daa..99ba50afe79 100644
--- a/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
+++ b/llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp
@@ -55,6 +55,7 @@ namespace {
     bool runOnFunction(Function &F) override;
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
+      AU.addRequired<DominatorTreeWrapperPass>();
       AU.addRequired<LazyValueInfoWrapperPass>();
       AU.addPreserved<GlobalsAAWrapperPass>();
     }
@@ -64,6 +65,7 @@ namespace {
 char CorrelatedValuePropagation::ID = 0;
 INITIALIZE_PASS_BEGIN(CorrelatedValuePropagation, "correlated-propagation",
                 "Value Propagation", false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LazyValueInfoWrapperPass)
 INITIALIZE_PASS_END(CorrelatedValuePropagation, "correlated-propagation",
                 "Value Propagation", false, false)
diff --git a/llvm/lib/Transforms/Scalar/JumpThreading.cpp b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
index 33afc207a95..5d94eb2f8ec 100644
--- a/llvm/lib/Transforms/Scalar/JumpThreading.cpp
+++ b/llvm/lib/Transforms/Scalar/JumpThreading.cpp
@@ -131,10 +131,11 @@ namespace {
     bool runOnFunction(Function &F) override;
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
-      if (PrintLVIAfterJumpThreading)
-        AU.addRequired<DominatorTreeWrapperPass>();
+      AU.addRequired<DominatorTreeWrapperPass>();
+      AU.addPreserved<DominatorTreeWrapperPass>();
       AU.addRequired<AAResultsWrapperPass>();
       AU.addRequired<LazyValueInfoWrapperPass>();
+      AU.addPreserved<LazyValueInfoWrapperPass>();
       AU.addPreserved<GlobalsAAWrapperPass>();
       AU.addRequired<TargetLibraryInfoWrapperPass>();
     }
@@ -148,6 +149,7 @@ char JumpThreading::ID = 0;
 
 INITIALIZE_PASS_BEGIN(JumpThreading, "jump-threading",
                 "Jump Threading", false, false)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(LazyValueInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)
@@ -277,6 +279,9 @@ bool JumpThreading::runOnFunction(Function &F) {
   if (skipFunction(F))
     return false;
   auto TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
+  // Get DT analysis before LVI. When LVI is initialized it conditionally adds
+  // DT if it's available.
+  auto DT = &getAnalysis<DominatorTreeWrapperPass>().getDomTree();
   auto LVI = &getAnalysis<LazyValueInfoWrapperPass>().getLVI();
   auto AA = &getAnalysis<AAResultsWrapperPass>().getAAResults();
   std::unique_ptr<BlockFrequencyInfo> BFI;
@@ -288,12 +293,11 @@ bool JumpThreading::runOnFunction(Function &F) {
     BFI.reset(new BlockFrequencyInfo(F, *BPI, LI));
   }
 
-  bool Changed = Impl.runImpl(F, TLI, LVI, AA, HasProfileData, std::move(BFI),
-                              std::move(BPI));
+  bool Changed = Impl.runImpl(F, TLI, LVI, AA, DT, HasProfileData,
+                              std::move(BFI), std::move(BPI));
   if (PrintLVIAfterJumpThreading) {
     dbgs() << "LVI for function '" << F.getName() << "':\n";
-    LVI->printLVI(F, getAnalysis<DominatorTreeWrapperPass>().getDomTree(),
-                  dbgs());
+    LVI->printLVI(F, *DT, dbgs());
   }
   return Changed;
 }
@@ -301,6 +305,9 @@ bool JumpThreading::runOnFunction(Function &F) {
 PreservedAnalyses JumpThreadingPass::run(Function &F,
                                          FunctionAnalysisManager &AM) {
   auto &TLI = AM.getResult<TargetLibraryAnalysis>(F);
+  // Get DT analysis before LVI. When LVI is initialized it conditionally adds
+  // DT if it's available.
+  auto &DT = AM.getResult<DominatorTreeAnalysis>(F);
   auto &LVI = AM.getResult<LazyValueAnalysis>(F);
   auto &AA = AM.getResult<AAManager>(F);
 
@@ -313,25 +320,28 @@ PreservedAnalyses JumpThreadingPass::run(Function &F,
     BFI.reset(new BlockFrequencyInfo(F, *BPI, LI));
   }
 
-  bool Changed = runImpl(F, &TLI, &LVI, &AA, HasProfileData, std::move(BFI),
-                         std::move(BPI));
+  bool Changed = runImpl(F, &TLI, &LVI, &AA, &DT, HasProfileData,
+                         std::move(BFI), std::move(BPI));
 
   if (!Changed)
     return PreservedAnalyses::all();
   PreservedAnalyses PA;
   PA.preserve<GlobalsAA>();
+  PA.preserve<DominatorTreeAnalysis>();
+  PA.preserve<LazyValueAnalysis>();
   return PA;
 }
 
 bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
                                 LazyValueInfo *LVI_, AliasAnalysis *AA_,
-                                bool HasProfileData_,
+                                DominatorTree *DT_, bool HasProfileData_,
                                 std::unique_ptr<BlockFrequencyInfo> BFI_,
                                 std::unique_ptr<BranchProbabilityInfo> BPI_) {
   DEBUG(dbgs() << "Jump threading on function '" << F.getName() << "'\n");
   TLI = TLI_;
   LVI = LVI_;
   AA = AA_;
+  DT = DT_;
   BFI.reset();
   BPI.reset();
   // When profile data is available, we need to update edge weights after
@@ -353,8 +363,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
   // back edges. This works for normal cases but not for unreachable blocks as
   // they may have cycle with no back edge.
   bool EverChanged = false;
-  EverChanged |= removeUnreachableBlocks(F, LVI);
-
+  EverChanged |= removeUnreachableBlocks(F, LVI, DT);
   FindLoopHeaders(F);
 
   bool Changed;
@@ -400,7 +409,7 @@ bool JumpThreadingPass::runImpl(Function &F, TargetLibraryInfo *TLI_,
         // awesome, but it allows us to use AssertingVH to prevent nasty
         // dangling pointer issues within LazyValueInfo.
         LVI->eraseBlock(BB);
-        if (TryToSimplifyUncondBranchFromEmptyBlock(BB))
+        if (TryToSimplifyUncondBranchFromEmptyBlock(BB, DT))
           Changed = true;
       }
     }
@@ -948,7 +957,7 @@ bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
         LoopHeaders.insert(BB);
 
       LVI->eraseBlock(SinglePred);
-      MergeBasicBlockIntoOnlyPred(BB);
+      MergeBasicBlockIntoOnlyPred(BB, DT);
 
       // Now that BB is merged into SinglePred (i.e. SinglePred Code followed by
       // BB code within one basic block `BB`), we need to invalidate the LVI
@@ -1031,18 +1040,27 @@ bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
   // successors to branch to.  Let GetBestDestForJumpOnUndef decide.
   if (isa<UndefValue>(Condition)) {
     unsigned BestSucc = GetBestDestForJumpOnUndef(BB);
+    std::vector<DominatorTree::UpdateType> Updates;
 
     // Fold the branch/switch.
     TerminatorInst *BBTerm = BB->getTerminator();
     for (unsigned i = 0, e = BBTerm->getNumSuccessors(); i != e; ++i) {
       if (i == BestSucc) continue;
-      BBTerm->getSuccessor(i)->removePredecessor(BB, true);
+      BasicBlock *Succ = BBTerm->getSuccessor(i);
+      Succ->removePredecessor(BB, true);
+      if (Succ == BB)
+        continue;
+      DominatorTree::UpdateType UT = {DominatorTree::Delete, BB, Succ};
+      // Make sure to remove a DT edge exactly once and not an edge to itself.
+      if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end())
+        Updates.push_back(UT);
     }
 
     DEBUG(dbgs() << "  In block '" << BB->getName()
           << "' folding undef terminator: " << *BBTerm << '\n');
     BranchInst::Create(BBTerm->getSuccessor(BestSucc), BBTerm);
     BBTerm->eraseFromParent();
+    DT->applyUpdates(Updates);
     return true;
   }
 
@@ -1053,7 +1071,7 @@ bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
     DEBUG(dbgs() << "  In block '" << BB->getName()
           << "' folding terminator: " << *BB->getTerminator() << '\n');
     ++NumFolds;
-    ConstantFoldTerminator(BB, true);
+    ConstantFoldTerminator(BB, true, nullptr, DT);
     return true;
   }
 
@@ -1086,9 +1104,12 @@ bool JumpThreadingPass::ProcessBlock(BasicBlock *BB) {
       if (Ret != LazyValueInfo::Unknown) {
         unsigned ToRemove = Ret == LazyValueInfo::True ? 1 : 0;
         unsigned ToKeep = Ret == LazyValueInfo::True ? 0 : 1;
-        CondBr->getSuccessor(ToRemove)->removePredecessor(BB, true);
+        BasicBlock *ToRemoveSucc = CondBr->getSuccessor(ToRemove);
+        ToRemoveSucc->removePredecessor(BB, true);
         BranchInst::Create(CondBr->getSuccessor(ToKeep), CondBr);
         CondBr->eraseFromParent();
+        if (BB != ToRemoveSucc)
+          DT->deleteEdge(BB, ToRemoveSucc);
         if (CondCmp->use_empty())
           CondCmp->eraseFromParent();
         // We can safely replace *some* uses of the CondInst if it has
@@ -1182,9 +1203,12 @@ bool JumpThreadingPass::ProcessImpliedCondition(BasicBlock *BB) {
     Optional<bool> Implication =
         isImpliedCondition(PBI->getCondition(), Cond, DL, CondIsTrue);
     if (Implication) {
-      BI->getSuccessor(*Implication ? 1 : 0)->removePredecessor(BB);
+      BasicBlock *RemoveSucc = BI->getSuccessor(*Implication ? 1 : 0);
+      RemoveSucc->removePredecessor(BB);
       BranchInst::Create(BI->getSuccessor(*Implication ? 0 : 1), BI);
       BI->eraseFromParent();
+      if (BB != RemoveSucc)
+        DT->deleteEdge(BB, RemoveSucc);
       return true;
     }
     CurrentBB = CurrentPred;
@@ -1576,18 +1600,27 @@ bool JumpThreadingPass::ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
   if (OnlyDest && OnlyDest != MultipleDestSentinel) {
     if (PredWithKnownDest ==
         (size_t)std::distance(pred_begin(BB), pred_end(BB))) {
+      std::vector<DominatorTree::UpdateType> Updates;
       bool SeenFirstBranchToOnlyDest = false;
       for (BasicBlock *SuccBB : successors(BB)) {
-        if (SuccBB == OnlyDest && !SeenFirstBranchToOnlyDest)
+        if (SuccBB == OnlyDest && !SeenFirstBranchToOnlyDest) {
           SeenFirstBranchToOnlyDest = true; // Don't modify the first branch.
-        else
+        } else {
           SuccBB->removePredecessor(BB, true); // This is unreachable successor.
+          if (SuccBB != OnlyDest && SuccBB != BB) {
+            DominatorTree::UpdateType UT = {DominatorTree::Delete, BB, SuccBB};
+            // Make sure to remove a DT edge exactly once.
+            if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end())
+              Updates.push_back(UT);
+          }
+        }
       }
 
       // Finally update the terminator.
       TerminatorInst *Term = BB->getTerminator();
       BranchInst::Create(OnlyDest, Term);
       Term->eraseFromParent();
+      DT->applyUpdates(Updates);
 
       // If the condition is now dead due to the removal of the old terminator,
       // erase it.
@@ -1922,7 +1955,6 @@ bool JumpThreadingPass::ThreadEdge(BasicBlock *BB,
 
       UsesToRename.push_back(&U);
     }
-
     // If there are no uses outside the block, we're done with this instruction.
     if (UsesToRename.empty())
       continue;
@@ -1951,6 +1983,10 @@ bool JumpThreadingPass::ThreadEdge(BasicBlock *BB,
       PredTerm->setSuccessor(i, NewBB);
     }
 
+  DT->applyUpdates({{DominatorTree::Insert, NewBB, SuccBB},
+                    {DominatorTree::Insert, PredBB, NewBB},
+                    {DominatorTree::Delete, PredBB, BB}});
+
   // At this point, the IR is fully up to date and consistent.  Do a quick scan
   // over the new instructions and zap any that are constants or dead.  This
   // frequently happens because of phi translation.
@@ -1977,7 +2013,7 @@ BasicBlock *JumpThreadingPass::SplitBlockPreds(BasicBlock *BB,
     for (auto Pred : Preds)
       PredBBFreq += BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, BB);
 
-  BasicBlock *PredBB = SplitBlockPredecessors(BB, Preds, Suffix);
+  BasicBlock *PredBB = SplitBlockPredecessors(BB, Preds, Suffix, DT);
 
   // Set the block frequency of the newly created PredBB, which is the sum of
   // frequencies of Preds.
@@ -2147,7 +2183,7 @@ bool JumpThreadingPass::DuplicateCondBranchOnPHIIntoPred(
   BranchInst *OldPredBranch = dyn_cast<BranchInst>(PredBB->getTerminator());
 
   if (!OldPredBranch || !OldPredBranch->isUnconditional()) {
-    PredBB = SplitEdge(PredBB, BB);
+    PredBB = SplitEdge(PredBB, BB, DT);
     OldPredBranch = cast<BranchInst>(PredBB->getTerminator());
   }
 
@@ -2244,6 +2280,8 @@ bool JumpThreadingPass::DuplicateCondBranchOnPHIIntoPred(
 
   // Remove the unconditional branch at the end of the PredBB block.
   OldPredBranch->eraseFromParent();
+  if (BB != PredBB)
+    DT->deleteEdge(PredBB, BB);
 
   ++NumDupes;
   return true;
@@ -2309,6 +2347,7 @@ bool JumpThreadingPass::TryToUnfoldSelect(CmpInst *CondCmp, BasicBlock *BB) {
       // Move the unconditional branch to NewBB.
       PredTerm->removeFromParent();
       NewBB->getInstList().insert(NewBB->end(), PredTerm);
+      DT->insertEdge(NewBB, BB);
       // Create a conditional branch and update PHI nodes.
       BranchInst::Create(NewBB, BB, SI->getCondition(), Pred);
       CondLHS->setIncomingValue(I, SI->getFalseValue());
@@ -2316,6 +2355,7 @@ bool JumpThreadingPass::TryToUnfoldSelect(CmpInst *CondCmp, BasicBlock *BB) {
       // The select is now dead.
       SI->eraseFromParent();
 
+      DT->insertEdge(Pred, NewBB);
       // Update any other PHI nodes in BB.
       for (BasicBlock::iterator BI = BB->begin();
            PHINode *Phi = dyn_cast<PHINode>(BI); ++BI)
@@ -2393,7 +2433,7 @@ bool JumpThreadingPass::TryToUnfoldSelectInCurrBB(BasicBlock *BB) {
       continue;
     // Expand the select.
     TerminatorInst *Term =
-        SplitBlockAndInsertIfThen(SI->getCondition(), SI, false);
+        SplitBlockAndInsertIfThen(SI->getCondition(), SI, false, nullptr, DT);
     PHINode *NewPN = PHINode::Create(SI->getType(), 2, "", SI);
     NewPN->addIncoming(SI->getTrueValue(), Term->getParent());
     NewPN->addIncoming(SI->getFalseValue(), BB);
@@ -2485,8 +2525,8 @@ bool JumpThreadingPass::ThreadGuard(BasicBlock *BB, IntrinsicInst *Guard,
   if (!TrueDestIsSafe && !FalseDestIsSafe)
     return false;
 
-  BasicBlock *UnguardedBlock = TrueDestIsSafe ? TrueDest : FalseDest;
-  BasicBlock *GuardedBlock = FalseDestIsSafe ? TrueDest : FalseDest;
+  BasicBlock *PredUnguardedBlock = TrueDestIsSafe ? TrueDest : FalseDest;
+  BasicBlock *PredGuardedBlock = FalseDestIsSafe ? TrueDest : FalseDest;
 
   ValueToValueMapTy UnguardedMapping, GuardedMapping;
   Instruction *AfterGuard = Guard->getNextNode();
@@ -2495,17 +2535,28 @@ bool JumpThreadingPass::ThreadGuard(BasicBlock *BB, IntrinsicInst *Guard,
     return false;
   // Duplicate all instructions before the guard and the guard itself to the
   // branch where implication is not proved.
-  GuardedBlock = DuplicateInstructionsInSplitBetween(
-      BB, GuardedBlock, AfterGuard, GuardedMapping);
+  BasicBlock *GuardedBlock = DuplicateInstructionsInSplitBetween(
+      BB, PredGuardedBlock, AfterGuard, GuardedMapping);
   assert(GuardedBlock && "Could not create the guarded block?");
   // Duplicate all instructions before the guard in the unguarded branch.
   // Since we have successfully duplicated the guarded block and this block
   // has fewer instructions, we expect it to succeed.
-  UnguardedBlock = DuplicateInstructionsInSplitBetween(BB, UnguardedBlock,
-                                                       Guard, UnguardedMapping);
+  BasicBlock *UnguardedBlock = DuplicateInstructionsInSplitBetween(
+      BB, PredUnguardedBlock, Guard, UnguardedMapping);
   assert(UnguardedBlock && "Could not create the unguarded block?");
   DEBUG(dbgs() << "Moved guard " << *Guard << " to block "
                << GuardedBlock->getName() << "\n");
+  // DuplicateInstructionsInSplitBetween inserts a new block, BB.split, between
+  // PredBB and BB. We need to perform two inserts and one delete in DT for each
+  // of the above calls.
+  DT->applyUpdates({// Guarded block split.
+                    {DominatorTree::Delete, PredGuardedBlock, BB},
+                    {DominatorTree::Insert, PredGuardedBlock, GuardedBlock},
+                    {DominatorTree::Insert, GuardedBlock, BB},
+                    // Unguarded block split.
+                    {DominatorTree::Delete, PredUnguardedBlock, BB},
+                    {DominatorTree::Insert, PredUnguardedBlock, UnguardedBlock},
+                    {DominatorTree::Insert, UnguardedBlock, BB}});
 
   // Some instructions before the guard may still have uses. For them, we need
   // to create Phi nodes merging their copies in both guarded and unguarded
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index 21412dcf68e..8b580f5fc2d 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -68,7 +68,8 @@ STATISTIC(NumRemoved, "Number of unreachable basic blocks removed");
 /// conditions and indirectbr addresses this might make dead if
 /// DeleteDeadConditions is true.
 bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
-                                  const TargetLibraryInfo *TLI) {
+                                  const TargetLibraryInfo *TLI,
+                                  DominatorTree *DT) {
   TerminatorInst *T = BB->getTerminator();
   IRBuilder<> Builder(T);
 
@@ -95,6 +96,8 @@ bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
       // Replace the conditional branch with an unconditional one.
       Builder.CreateBr(Destination);
       BI->eraseFromParent();
+      if (DT && OldDest != Destination && OldDest != BB)
+        DT->deleteEdge(BB, OldDest);
       return true;
     }
 
@@ -165,9 +168,17 @@ bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
                           createBranchWeights(Weights));
         }
         // Remove this entry.
-        DefaultDest->removePredecessor(SI->getParent());
+        BasicBlock *ParentBB = SI->getParent();
+        DefaultDest->removePredecessor(ParentBB);
         i = SI->removeCase(i);
         e = SI->case_end();
+        if (DT && DefaultDest != ParentBB) {
+          // DefaultDest may still be a successor of a non-default case.
+          if (none_of(successors(ParentBB), [DefaultDest](BasicBlock *S) {
+                return S == DefaultDest;
+              }))
+            DT->deleteEdge(ParentBB, DefaultDest);
+        }
         continue;
       }
 
@@ -193,19 +204,29 @@ bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
       // Insert the new branch.
       Builder.CreateBr(TheOnlyDest);
       BasicBlock *BB = SI->getParent();
+      BasicBlock *TheOnlyDestCheck = TheOnlyDest;
+      std::vector<DominatorTree::UpdateType> Updates;
 
       // Remove entries from PHI nodes which we no longer branch to...
       for (BasicBlock *Succ : SI->successors()) {
         // Found case matching a constant operand?
-        if (Succ == TheOnlyDest)
+        if (Succ == TheOnlyDest) {
           TheOnlyDest = nullptr; // Don't modify the first branch to TheOnlyDest
-        else
+        } else {
           Succ->removePredecessor(BB);
+          if (DT && Succ != TheOnlyDestCheck && Succ != BB) {
+            DominatorTree::UpdateType UT = {DominatorTree::Delete, BB, Succ};
+            if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end())
+              Updates.push_back(UT);
+          }
+        }
       }
 
       // Delete the old switch.
       Value *Cond = SI->getCondition();
       SI->eraseFromParent();
+      if (DT && !Updates.empty())
+        DT->applyUpdates(Updates);
       if (DeleteDeadConditions)
         RecursivelyDeleteTriviallyDeadInstructions(Cond, TLI);
       return true;
@@ -253,17 +274,30 @@ bool llvm::ConstantFoldTerminator(BasicBlock *BB, bool DeleteDeadConditions,
     if (BlockAddress *BA =
           dyn_cast<BlockAddress>(IBI->getAddress()->stripPointerCasts())) {
       BasicBlock *TheOnlyDest = BA->getBasicBlock();
+      BasicBlock *TheOnlyDestCheck = TheOnlyDest;
+      std::vector<DominatorTree::UpdateType> Updates;
       // Insert the new branch.
       Builder.CreateBr(TheOnlyDest);
 
       for (unsigned i = 0, e = IBI->getNumDestinations(); i != e; ++i) {
-        if (IBI->getDestination(i) == TheOnlyDest)
+        if (IBI->getDestination(i) == TheOnlyDest) {
           TheOnlyDest = nullptr;
-        else
-          IBI->getDestination(i)->removePredecessor(IBI->getParent());
+        } else {
+          BasicBlock *ParentBB = IBI->getParent();
+          BasicBlock *DestBB = IBI->getDestination(i);
+          DestBB->removePredecessor(ParentBB);
+          if (DT && DestBB != TheOnlyDestCheck && DestBB != ParentBB) {
+            DominatorTree::UpdateType UT = {DominatorTree::Delete, ParentBB,
+                                            DestBB};
+            if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end())
+              Updates.push_back(UT);
+          }
+        }
       }
       Value *Address = IBI->getAddress();
       IBI->eraseFromParent();
+      if (DT && !Updates.empty())
+        DT->applyUpdates(Updates);
       if (DeleteDeadConditions)
         RecursivelyDeleteTriviallyDeadInstructions(Address, TLI);
 
@@ -553,7 +587,8 @@ bool llvm::SimplifyInstructionsInBlock(BasicBlock *BB,
 ///
 /// .. and delete the predecessor corresponding to the '1', this will attempt to
 /// recursively fold the and to 0.
-void llvm::RemovePredecessorAndSimplify(BasicBlock *BB, BasicBlock *Pred) {
+void llvm::RemovePredecessorAndSimplify(BasicBlock *BB, BasicBlock *Pred,
+                                        DominatorTree *DT) {
   // This only adjusts blocks with PHI nodes.
   if (!isa<PHINode>(BB->begin()))
     return;
@@ -576,6 +611,8 @@ void llvm::RemovePredecessorAndSimplify(BasicBlock *BB, BasicBlock *Pred) {
     // of the block.
     if (PhiIt != OldPhiIt) PhiIt = &BB->front();
   }
+  if (DT && BB != Pred)
+    DT->deleteEdge(Pred, BB);
 }
 
 
@@ -597,6 +634,23 @@ void llvm::MergeBasicBlockIntoOnlyPred(BasicBlock *DestBB, DominatorTree *DT) {
   BasicBlock *PredBB = DestBB->getSinglePredecessor();
   assert(PredBB && "Block doesn't have a single predecessor!");
 
+  // Collect all the edges that enter PredBB, discarding edges to itself and
+  // duplicates. These dominator edges will be redirected to DestBB.
+  std::vector<DominatorTree::UpdateType> Updates;
+  if (DT)
+    for (pred_iterator PI = pred_begin(PredBB), E = pred_end(PredBB); PI != E;
+         ++PI) {
+      if (*PI == PredBB)
+        continue;
+      DominatorTree::UpdateType UT = {DominatorTree::Delete, *PI, PredBB};
+      if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end()) {
+        Updates.push_back(UT);
+        // DestBB cannot dominate itself.
+        if (*PI != DestBB)
+          Updates.push_back({DominatorTree::Insert, *PI, DestBB});
+      }
+    }
+
   // Zap anything that took the address of DestBB.  Not doing this will give the
   // address an invalid value.
   if (DestBB->hasAddressTaken()) {
@@ -617,16 +671,25 @@ void llvm::MergeBasicBlockIntoOnlyPred(BasicBlock *DestBB, DominatorTree *DT) {
 
   // If the PredBB is the entry block of the function, move DestBB up to
   // become the entry block after we erase PredBB.
-  if (PredBB == &DestBB->getParent()->getEntryBlock())
+  bool ReplacedEntryBB = false;
+  if (PredBB == &DestBB->getParent()->getEntryBlock()) {
     DestBB->moveAfter(PredBB);
+    ReplacedEntryBB = true;
+  }
 
-  if (DT) {
-    BasicBlock *PredBBIDom = DT->getNode(PredBB)->getIDom()->getBlock();
-    DT->changeImmediateDominator(DestBB, PredBBIDom);
-    DT->eraseNode(PredBB);
+  if (DT && !ReplacedEntryBB) {
+    Updates.push_back({DominatorTree::Delete, PredBB, DestBB});
+    DT->applyUpdates(Updates);
   }
+
   // Nuke BB.
   PredBB->eraseFromParent();
+
+  // The entry block was removed and there is no external interface for the
+  // dominator tree to be notified of this change. In this corner-case we
+  // recalculate the entire tree.
+  if (DT && ReplacedEntryBB)
+    DT->recalculate(*(DestBB->getParent()));
 }
 
 /// CanMergeValues - Return true if we can choose one of these values to use
@@ -834,7 +897,8 @@ static void redirectValuesFromPredecessorsToPhi(BasicBlock *BB,
 /// potential side-effect free intrinsics and the branch.  If possible,
 /// eliminate BB by rewriting all the predecessors to branch to the successor
 /// block and return true.  If we can't transform, return false.
-bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB) {
+bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB,
+                                                   DominatorTree *DT) {
   assert(BB != &BB->getParent()->getEntryBlock() &&
          "TryToSimplifyUncondBranchFromEmptyBlock called on entry block!");
 
@@ -875,6 +939,20 @@ bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB) {
 
   DEBUG(dbgs() << "Killing Trivial BB: \n" << *BB);
 
+  // Collect all the edges that enter BB, discarding edges to itself and
+  // duplicates. These dominator edges will be redirected to Succ.
+  std::vector<DominatorTree::UpdateType> Updates;
+  if (DT)
+    for (pred_iterator PI = pred_begin(BB), E = pred_end(BB); PI != E; ++PI) {
+      if (*PI == BB)
+        continue;
+      DominatorTree::UpdateType UT = {DominatorTree::Delete, *PI, BB};
+      if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end()) {
+        Updates.push_back(UT);
+        Updates.push_back({DominatorTree::Insert, *PI, Succ});
+      }
+    }
+
   if (isa<PHINode>(Succ->begin())) {
     // If there is more than one pred of succ, and there are PHI nodes in
     // the successor, then we need to add incoming edges for the PHI nodes
@@ -909,16 +987,27 @@ bool llvm::TryToSimplifyUncondBranchFromEmptyBlock(BasicBlock *BB) {
   // add the metadata to the branch instructions in the predecessors.
   unsigned LoopMDKind = BB->getContext().getMDKindID("llvm.loop");
   Instruction *TI = BB->getTerminator();
-  if (TI)
+  if (TI) {
     if (MDNode *LoopMD = TI->getMetadata(LoopMDKind))
       for (pred_iterator PI = pred_begin(BB), E = pred_end(BB); PI != E; ++PI) {
         BasicBlock *Pred = *PI;
         Pred->getTerminator()->setMetadata(LoopMDKind, LoopMD);
       }
+    // Replace the terminator instruction with unreachable to ensure the CFG is
+    // consistent. This is necessary for dominance to be correctly calculated.
+    new UnreachableInst(BB->getContext(), TI);
+    TI->eraseFromParent();
+  }
 
   // Everything that jumped to BB now goes to Succ.
   BB->replaceAllUsesWith(Succ);
   if (!Succ->hasName()) Succ->takeName(BB);
+
+  if (DT) {
+    Updates.push_back({DominatorTree::Delete, BB, Succ});
+    DT->applyUpdates(Updates);
+  }
+
   BB->eraseFromParent();              // Delete the old basic block.
   return true;
 }
@@ -1399,13 +1488,19 @@ unsigned llvm::removeAllNonTerminatorAndEHPadInstructions(BasicBlock *BB) {
 }
 
 unsigned llvm::changeToUnreachable(Instruction *I, bool UseLLVMTrap,
-                                   bool PreserveLCSSA) {
+                                   bool PreserveLCSSA, DominatorTree *DT) {
   BasicBlock *BB = I->getParent();
+  std::vector<DominatorTree::UpdateType> Updates;
   // Loop over all of the successors, removing BB's entry from any PHI
   // nodes.
-  for (BasicBlock *Successor : successors(BB))
+  for (BasicBlock *Successor : successors(BB)) {
     Successor->removePredecessor(BB, PreserveLCSSA);
-
+    if (DT && BB != Successor) {
+      DominatorTree::UpdateType UT = {DominatorTree::Delete, BB, Successor};
+      if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end())
+        Updates.push_back(UT);
+    }
+  }
   // Insert a call to llvm.trap right before this.  This turns the undefined
   // behavior into a hard fail instead of falling through into random code.
   if (UseLLVMTrap) {
@@ -1425,11 +1520,13 @@ unsigned llvm::changeToUnreachable(Instruction *I, bool UseLLVMTrap,
     BB->getInstList().erase(BBI++);
     ++NumInstrsRemoved;
   }
+  if (DT && !Updates.empty())
+    DT->applyUpdates(Updates);
   return NumInstrsRemoved;
 }
 
 /// changeToCall - Convert the specified invoke into a normal call.
-static void changeToCall(InvokeInst *II) {
+static void changeToCall(InvokeInst *II, DominatorTree *DT = nullptr) {
   SmallVector<Value*, 8> Args(II->arg_begin(), II->arg_end());
   SmallVector<OperandBundleDef, 1> OpBundles;
   II->getOperandBundlesAsDefs(OpBundles);
@@ -1442,11 +1539,16 @@ static void changeToCall(InvokeInst *II) {
   II->replaceAllUsesWith(NewCall);
 
   // Follow the call by a branch to the normal destination.
-  BranchInst::Create(II->getNormalDest(), II);
+  BasicBlock *NormalDestBB = II->getNormalDest();
+  BranchInst::Create(NormalDestBB, II);
 
   // Update PHI nodes in the unwind destination
-  II->getUnwindDest()->removePredecessor(II->getParent());
+  BasicBlock *BB = II->getParent();
+  BasicBlock *UnwindDestBB = II->getUnwindDest();
+  UnwindDestBB->removePredecessor(BB);
   II->eraseFromParent();
+  if (DT && BB != UnwindDestBB && NormalDestBB != UnwindDestBB)
+    DT->deleteEdge(BB, UnwindDestBB);
 }
 
 BasicBlock *llvm::changeToInvokeAndSplitBasicBlock(CallInst *CI,
@@ -1487,7 +1589,8 @@ BasicBlock *llvm::changeToInvokeAndSplitBasicBlock(CallInst *CI,
 }
 
 static bool markAliveBlocks(Function &F,
-                            SmallPtrSetImpl<BasicBlock*> &Reachable) {
+                            SmallPtrSetImpl<BasicBlock *> &Reachable,
+                            DominatorTree *DT = nullptr) {
 
   SmallVector<BasicBlock*, 128> Worklist;
   BasicBlock *BB = &F.front();
@@ -1508,7 +1611,7 @@ static bool markAliveBlocks(Function &F,
         if (II->getIntrinsicID() == Intrinsic::assume) {
           if (match(II->getArgOperand(0), m_CombineOr(m_Zero(), m_Undef()))) {
             // Don't insert a call to llvm.trap right before the unreachable.
-            changeToUnreachable(II, false);
+            changeToUnreachable(II, false, false, DT);
             Changed = true;
             break;
           }
@@ -1525,7 +1628,8 @@ static bool markAliveBlocks(Function &F,
           // still be useful for widening.
           if (match(II->getArgOperand(0), m_Zero()))
             if (!isa<UnreachableInst>(II->getNextNode())) {
-              changeToUnreachable(II->getNextNode(), /*UseLLVMTrap=*/ false);
+              changeToUnreachable(II->getNextNode(), /*UseLLVMTrap=*/false,
+                                  false, DT);
               Changed = true;
               break;
             }
@@ -1535,7 +1639,7 @@ static bool markAliveBlocks(Function &F,
       if (auto *CI = dyn_cast<CallInst>(&I)) {
         Value *Callee = CI->getCalledValue();
         if (isa<ConstantPointerNull>(Callee) || isa<UndefValue>(Callee)) {
-          changeToUnreachable(CI, /*UseLLVMTrap=*/false);
+          changeToUnreachable(CI, /*UseLLVMTrap=*/false, false, DT);
           Changed = true;
           break;
         }
@@ -1545,7 +1649,7 @@ static bool markAliveBlocks(Function &F,
           // though.
           if (!isa<UnreachableInst>(CI->getNextNode())) {
             // Don't insert a call to llvm.trap right before the unreachable.
-            changeToUnreachable(CI->getNextNode(), false);
+            changeToUnreachable(CI->getNextNode(), false, false, DT);
             Changed = true;
           }
           break;
@@ -1564,7 +1668,7 @@ static bool markAliveBlocks(Function &F,
         if (isa<UndefValue>(Ptr) ||
             (isa<ConstantPointerNull>(Ptr) &&
              SI->getPointerAddressSpace() == 0)) {
-          changeToUnreachable(SI, true);
+          changeToUnreachable(SI, true, false, DT);
           Changed = true;
           break;
         }
@@ -1576,16 +1680,19 @@ static bool markAliveBlocks(Function &F,
       // Turn invokes that call 'nounwind' functions into ordinary calls.
       Value *Callee = II->getCalledValue();
       if (isa<ConstantPointerNull>(Callee) || isa<UndefValue>(Callee)) {
-        changeToUnreachable(II, true);
+        changeToUnreachable(II, true, false, DT);
         Changed = true;
       } else if (II->doesNotThrow() && canSimplifyInvokeNoUnwind(&F)) {
         if (II->use_empty() && II->onlyReadsMemory()) {
           // jump to the normal destination branch.
+          BasicBlock *UnwindDestBB = II->getUnwindDest();
           BranchInst::Create(II->getNormalDest(), II);
-          II->getUnwindDest()->removePredecessor(II->getParent());
+          UnwindDestBB->removePredecessor(II->getParent());
           II->eraseFromParent();
+          if (DT && BB != UnwindDestBB)
+            DT->deleteEdge(BB, UnwindDestBB);
         } else
-          changeToCall(II);
+          changeToCall(II, DT);
         Changed = true;
       }
     } else if (auto *CatchSwitch = dyn_cast<CatchSwitchInst>(Terminator)) {
@@ -1628,7 +1735,7 @@ static bool markAliveBlocks(Function &F,
       }
     }
 
-    Changed |= ConstantFoldTerminator(BB, true);
+    Changed |= ConstantFoldTerminator(BB, true, nullptr, DT);
     for (BasicBlock *Successor : successors(BB))
       if (Reachable.insert(Successor).second)
         Worklist.push_back(Successor);
@@ -1636,11 +1743,11 @@ static bool markAliveBlocks(Function &F,
   return Changed;
 }
 
-void llvm::removeUnwindEdge(BasicBlock *BB) {
+void llvm::removeUnwindEdge(BasicBlock *BB, DominatorTree *DT) {
   TerminatorInst *TI = BB->getTerminator();
 
   if (auto *II = dyn_cast<InvokeInst>(TI)) {
-    changeToCall(II);
+    changeToCall(II, DT);
     return;
   }
 
@@ -1668,15 +1775,19 @@ void llvm::removeUnwindEdge(BasicBlock *BB) {
   UnwindDest->removePredecessor(BB);
   TI->replaceAllUsesWith(NewTI);
   TI->eraseFromParent();
+  if (DT && BB != UnwindDest)
+    DT->deleteEdge(BB, UnwindDest);
 }
 
 /// removeUnreachableBlocks - Remove blocks that are not reachable, even
 /// if they are in a dead cycle.  Return true if a change was made, false
 /// otherwise. If `LVI` is passed, this function preserves LazyValueInfo
 /// after modifying the CFG.
-bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
+bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI,
+                                   DominatorTree *DT) {
   SmallPtrSet<BasicBlock*, 16> Reachable;
-  bool Changed = markAliveBlocks(F, Reachable);
+  bool Changed = markAliveBlocks(F, Reachable, DT);
+  std::vector<DominatorTree::UpdateType> Updates;
 
   // If there are unreachable blocks in the CFG...
   if (Reachable.size() == F.size())
@@ -1685,6 +1796,8 @@ bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
   assert(Reachable.size() < F.size());
   NumRemoved += F.size()-Reachable.size();
 
+  SmallPtrSet<TerminatorInst *, 4> TIRemoved;
+
   // Loop over all of the basic blocks that are not reachable, dropping all of
   // their internal references...
   for (Function::iterator BB = ++F.begin(), E = F.end(); BB != E; ++BB) {
@@ -1692,13 +1805,35 @@ bool llvm::removeUnreachableBlocks(Function &F, LazyValueInfo *LVI) {
       continue;
 
     for (BasicBlock *Successor : successors(&*BB))
-      if (Reachable.count(Successor))
+      if (Reachable.count(Successor)) {
         Successor->removePredecessor(&*BB);
+        if (DT && &*BB != Successor) {
+          DominatorTree::UpdateType UT = {DominatorTree::Delete, &*BB,
+                                          Successor};
+          if (std::find(Updates.begin(), Updates.end(), UT) == Updates.end())
+            Updates.push_back(UT);
+        }
+      }
+
     if (LVI)
       LVI->eraseBlock(&*BB);
+
+    TerminatorInst *TI = BB->getTerminator();
+    TIRemoved.insert(TI);
+    new UnreachableInst(BB->getContext(), TI);
+
     BB->dropAllReferences();
   }
 
+  // Remove all the terminator instructions after dropping all references. This
+  // keeps the state of the CFG consistent and prevents asserts from circular
+  // use counts in groups of unreachable basic blocks.
+  for (TerminatorInst *TI : TIRemoved)
+    TI->eraseFromParent();
+
+  if (DT && !Updates.empty())
+    DT->applyUpdates(Updates);
+
   for (Function::iterator I = ++F.begin(); I != F.end();)
     if (!Reachable.count(&*I))
       I = F.getBasicBlockList().erase(I);