[SimplifyCFG] Teach mergeEmptyReturnBlocks() to preserve DomTree

A first real transformation that didn't already knew how to do that,
but it's pretty tame - either change successor of all the predecessors
of a block and carefully delay deletion of the block until afterwards
the DomTree updates are appled, or add a successor to the block.

There wasn't a great test coverage for this, so i added extra, to be sure.

llvm/lib/Transforms/Scalar/SimplifyCFGPass.cpp

@@ -77,9 +77,12 @@ STATISTIC(NumSimpl, "Number of blocks simplified");
 
 /// If we have more than one empty (other than phi node) return blocks,
 /// merge them together to promote recursive block merging.
-static bool mergeEmptyReturnBlocks(Function &F) {
+static bool mergeEmptyReturnBlocks(Function &F, DomTreeUpdater *DTU) {
   bool Changed = false;
 
+  std::vector<DominatorTree::UpdateType> Updates;
+  SmallVector<BasicBlock *, 8> DeadBlocks;
+
   BasicBlock *RetBlock = nullptr;
 
   // Scan all the blocks in the function, looking for empty return blocks.
@@ -135,8 +138,15 @@ static bool mergeEmptyReturnBlocks(Function &F) {
     if (Ret->getNumOperands() == 0 ||
         Ret->getOperand(0) ==
           cast<ReturnInst>(RetBlock->getTerminator())->getOperand(0)) {
+      // All predecessors of BB should now branch to RetBlock instead.
+      if (DTU) {
+        for (auto *Predecessor : predecessors(&BB)) {
+          Updates.push_back({DominatorTree::Delete, Predecessor, &BB});
+          Updates.push_back({DominatorTree::Insert, Predecessor, RetBlock});
+        }
+      }
       BB.replaceAllUsesWith(RetBlock);
-      BB.eraseFromParent();
+      DeadBlocks.emplace_back(&BB);
       continue;
     }
 
@@ -160,6 +170,17 @@ static bool mergeEmptyReturnBlocks(Function &F) {
     RetBlockPHI->addIncoming(Ret->getOperand(0), &BB);
     BB.getTerminator()->eraseFromParent();
     BranchInst::Create(RetBlock, &BB);
+    if (DTU)
+      Updates.push_back({DominatorTree::Insert, &BB, RetBlock});
+  }
+
+  if (DTU) {
+    DTU->applyUpdatesPermissive(Updates);
+    for (auto *BB : DeadBlocks)
+      DTU->deleteBB(BB);
+  } else {
+    for (auto *BB : DeadBlocks)
+      BB->eraseFromParent();
   }
 
   return Changed;
@@ -200,7 +221,7 @@ static bool simplifyFunctionCFGImpl(Function &F, const TargetTransformInfo &TTI,
   DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
 
   bool EverChanged = removeUnreachableBlocks(F, DT ? &DTU : nullptr);
-  EverChanged |= mergeEmptyReturnBlocks(F);
+  EverChanged |= mergeEmptyReturnBlocks(F, DT ? &DTU : nullptr);
   EverChanged |= iterativelySimplifyCFG(F, TTI, DT ? &DTU : nullptr, Options);
 
   // If neither pass changed anything, we're done.

llvm/test/Transforms/SimplifyCFG/2005-10-02-InvokeSimplify.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -simplifycfg -disable-output
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -disable-output
 
 define i1 @foo() personality i8* bitcast (i32 (...)* @__gxx_personality_v0 to i8*) {
 	%X = invoke i1 @foo( )

llvm/test/Transforms/SimplifyCFG/merge-empty-return-blocks.ll

@@ -0,0 +1,48 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s
+
+define void @t0(i1 %c) {
+; CHECK-LABEL: @t0(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %c, label %end0, label %end1
+
+end0:
+  ret void
+
+end1:
+  ret void
+}
+
+define i8 @t1(i1 %c, i8 %v) {
+; CHECK-LABEL: @t1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    ret i8 [[V:%.*]]
+;
+entry:
+  br i1 %c, label %end0, label %end1
+
+end0:
+  ret i8 %v
+
+end1:
+  ret i8 %v
+}
+
+define i8 @t2(i1 %c, i8 %v0, i8 %v1) {
+; CHECK-LABEL: @t2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[C:%.*]], i8 [[V0:%.*]], i8 [[V1:%.*]]
+; CHECK-NEXT:    ret i8 [[SPEC_SELECT]]
+;
+entry:
+  br i1 %c, label %end0, label %end1
+
+end0:
+  ret i8 %v0
+
+end1:
+  ret i8 %v1
+}

llvm/test/Transforms/SimplifyCFG/return-merge.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -simplifycfg -S | not grep br
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | not grep br
 
 define i32 @test1(i1 %C) {
 entry: