[SimplifyCFG] Teach SimplifyCondBranchToCondBranch() to preserve DomTree

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -3668,6 +3668,8 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
   LLVM_DEBUG(dbgs() << "FOLDING BRs:" << *PBI->getParent()
                     << "AND: " << *BI->getParent());
 
+  SmallVector<DominatorTree::UpdateType, 5> Updates;
+
   // If OtherDest *is* BB, then BB is a basic block with a single conditional
   // branch in it, where one edge (OtherDest) goes back to itself but the other
   // exits.  We don't *know* that the program avoids the infinite loop
@@ -3681,6 +3683,7 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
     BasicBlock *InfLoopBlock =
         BasicBlock::Create(BB->getContext(), "infloop", BB->getParent());
     BranchInst::Create(InfLoopBlock, InfLoopBlock);
+    Updates.push_back({DominatorTree::Insert, InfLoopBlock, InfLoopBlock});
     OtherDest = InfLoopBlock;
   }
 
@@ -3702,11 +3705,20 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
   // Merge the conditions.
   Value *Cond = Builder.CreateOr(PBICond, BICond, "brmerge");
 
+  for (auto *Successor : successors(PBI->getParent()))
+    Updates.push_back({DominatorTree::Delete, PBI->getParent(), Successor});
+
   // Modify PBI to branch on the new condition to the new dests.
   PBI->setCondition(Cond);
   PBI->setSuccessor(0, CommonDest);
   PBI->setSuccessor(1, OtherDest);
 
+  for (auto *Successor : successors(PBI->getParent()))
+    Updates.push_back({DominatorTree::Insert, PBI->getParent(), Successor});
+
+  if (DTU)
+    DTU->applyUpdatesPermissive(Updates);
+
   // Update branch weight for PBI.
   uint64_t PredTrueWeight, PredFalseWeight, SuccTrueWeight, SuccFalseWeight;
   uint64_t PredCommon, PredOther, SuccCommon, SuccOther;

llvm/test/Transforms/SimplifyCFG/2008-07-13-InfLoopMiscompile.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -simplifycfg -S | FileCheck %s
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s
 ; PR2540
 ; Outval should end up with a select from 0/2, not all constants.
 

llvm/test/Transforms/SimplifyCFG/extract-cost.ll

@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -simplifycfg -S  < %s | FileCheck %s
+; RUN: opt -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S  < %s | FileCheck %s
 
 declare { i32, i1 } @llvm.uadd.with.overflow.i32(i32, i32) #1
 

llvm/test/Transforms/SimplifyCFG/pr34131.ll

@@ -1,4 +1,4 @@
-; RUN: opt -simplifycfg -S < %s | FileCheck %s
+; RUN: opt -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s
 
 ; Just checking for lack of crash here, but we should be able to check the IR?
 ; Earlier version using auto-generated checks from utils/update_test_checks.py