From 83659c707675f81680e355d98bd3600cd59a9278 Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Sat, 19 Dec 2020 18:09:10 +0300
Subject: [PATCH] [SimplifyCFG] simplifySingleResume():
 FoldReturnIntoUncondBranch() already knows how to preserve DomTree

... so just ensure that we pass DomTreeUpdater it into it.

Apparently, there were no dedicated tests just for that functionality,
so i'm adding one here.
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp          |  7 +++--
 .../SimplifyCFG/duplicate-ret-into-uncond-br.ll    | 31 ++++++++++++++++++++++
 2 files changed, 36 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/duplicate-ret-into-uncond-br.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index fec2a67..5a70bca 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -4424,7 +4424,7 @@ bool SimplifyCFGOpt::simplifyReturn(ReturnInst *RI, IRBuilder<> &Builder) {
       BasicBlock *Pred = UncondBranchPreds.pop_back_val();
       LLVM_DEBUG(dbgs() << "FOLDING: " << *BB
                         << "INTO UNCOND BRANCH PRED: " << *Pred);
-      (void)FoldReturnIntoUncondBranch(RI, BB, Pred);
+      (void)FoldReturnIntoUncondBranch(RI, BB, Pred, DTU);
     }
 
     // If we eliminated all predecessors of the block, delete the block now.
@@ -4432,7 +4432,10 @@ bool SimplifyCFGOpt::simplifyReturn(ReturnInst *RI, IRBuilder<> &Builder) {
       // We know there are no successors, so just nuke the block.
       if (LoopHeaders)
         LoopHeaders->erase(BB);
-      BB->eraseFromParent();
+      if (DTU)
+        DTU->deleteBB(BB);
+      else
+        BB->eraseFromParent();
     }
 
     return true;
diff --git a/llvm/test/Transforms/SimplifyCFG/duplicate-ret-into-uncond-br.ll b/llvm/test/Transforms/SimplifyCFG/duplicate-ret-into-uncond-br.ll
new file mode 100644
index 0000000..a48a084
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/duplicate-ret-into-uncond-br.ll
@@ -0,0 +1,31 @@
+; 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 void @foo(i1 %c) {
+; CHECK-LABEL: @foo(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[TRUE:%.*]], label [[FALSE:%.*]]
+; CHECK:       true:
+; CHECK-NEXT:    call void @bar()
+; CHECK-NEXT:    ret void
+; CHECK:       false:
+; CHECK-NEXT:    call void @baz()
+; CHECK-NEXT:    ret void
+;
+entry:
+  br i1 %c, label %true, label %false
+
+true:
+  call void @bar()
+  br label %end
+
+false:
+  call void @baz()
+  br label %end
+
+end:
+  ret void
+}
-- 
2.7.4