From b94520c9ee2ffdb15e8af2afa210e8eb76da9343 Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Sat, 19 Dec 2020 19:12:30 +0300
Subject: [PATCH] [SimplifyCFG] Teach SimplifyCondBranchToTwoReturns() to
 preserve DomTree, part 1

... for the general case of returning a value.
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     |  5 +++++
 .../LoopDeletion/simplify-then-delete.ll      |  2 +-
 .../SimplifyCFG/2008-09-08-MultiplePred.ll    |  2 +-
 .../merge-duplicate-conditional-ret-val.ll    | 21 +++++++++++++++++++
 4 files changed, 28 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/merge-duplicate-conditional-ret-val.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index fb282054e904d..8ee7e57a6135c 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -2688,6 +2688,11 @@ bool SimplifyCFGOpt::SimplifyCondBranchToTwoReturns(BranchInst *BI,
 
   EraseTerminatorAndDCECond(BI);
 
+  if (DTU) {
+    DTU->applyUpdatesPermissive({{DominatorTree::Delete, BB, TrueSucc},
+                                 {DominatorTree::Delete, BB, FalseSucc}});
+  }
+
   return true;
 }
 
diff --git a/llvm/test/Transforms/LoopDeletion/simplify-then-delete.ll b/llvm/test/Transforms/LoopDeletion/simplify-then-delete.ll
index 4278ef16d2145..294f610fa5155 100644
--- a/llvm/test/Transforms/LoopDeletion/simplify-then-delete.ll
+++ b/llvm/test/Transforms/LoopDeletion/simplify-then-delete.ll
@@ -1,4 +1,4 @@
-; RUN: opt < %s -S -indvars -loop-deletion -simplifycfg | FileCheck %s
+; RUN: opt < %s -S -indvars -loop-deletion -simplifycfg -simplifycfg-require-and-preserve-domtree=1 | FileCheck %s
 ; PR5794
 
 ; Indvars and loop deletion should be able to eliminate all looping
diff --git a/llvm/test/Transforms/SimplifyCFG/2008-09-08-MultiplePred.ll b/llvm/test/Transforms/SimplifyCFG/2008-09-08-MultiplePred.ll
index 6b216f598efb5..a2477f2f3a663 100644
--- a/llvm/test/Transforms/SimplifyCFG/2008-09-08-MultiplePred.ll
+++ b/llvm/test/Transforms/SimplifyCFG/2008-09-08-MultiplePred.ll
@@ -1,4 +1,4 @@
-; RUN: opt < %s -simplifycfg -disable-output
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -disable-output
 ; PR 2777
 @g_103 = common global i32 0		; <i32*> [#uses=1]
 
diff --git a/llvm/test/Transforms/SimplifyCFG/merge-duplicate-conditional-ret-val.ll b/llvm/test/Transforms/SimplifyCFG/merge-duplicate-conditional-ret-val.ll
new file mode 100644
index 0000000000000..55e02e99abea8
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/merge-duplicate-conditional-ret-val.ll
@@ -0,0 +1,21 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -simplifycfg-dup-ret -S | FileCheck %s
+
+declare void @bar()
+declare void @baz()
+
+define i8 @foo(i1 %c, i8 %v0, i8 %v1) {
+; CHECK-LABEL: @foo(
+; 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 %true, label %false
+
+true:
+  ret i8 %v0
+
+false:
+  ret i8 %v1
+}