From f1ce6960561bc28129fa3306f57a5c6df21258ab Mon Sep 17 00:00:00 2001
From: Roman Lebedev <lebedev.ri@gmail.com>
Date: Fri, 1 Jan 2021 16:33:19 +0300
Subject: [PATCH] [SimplifyCFG] Teach tryWidenCondBranchToCondBranch() to
 preserve DomTree

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 19 +++++++++++++++----
 .../Transforms/SimplifyCFG/wc-widen-block.ll  |  2 +-
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 53353cd5e923..30e3f0bc174c 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3530,7 +3530,8 @@ static bool mergeConditionalStores(BranchInst *PBI, BranchInst *QBI,
 /// If the previous block ended with a widenable branch, determine if reusing
 /// the target block is profitable and legal.  This will have the effect of
 /// "widening" PBI, but doesn't require us to reason about hosting safety.
-static bool tryWidenCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI) {
+static bool tryWidenCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
+                                           DomTreeUpdater *DTU) {
   // TODO: This can be generalized in two important ways:
   // 1) We can allow phi nodes in IfFalseBB and simply reuse all the input
   //    values from the PBI edge.
@@ -3553,15 +3554,25 @@ static bool tryWidenCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI) {
   if (BI->getSuccessor(1) != IfFalseBB && // no inf looping
       BI->getSuccessor(1)->getTerminatingDeoptimizeCall() && // profitability
       NoSideEffects(*BI->getParent())) {
-    BI->getSuccessor(1)->removePredecessor(BI->getParent());
+    auto *OldSuccessor = BI->getSuccessor(1);
+    OldSuccessor->removePredecessor(BI->getParent());
     BI->setSuccessor(1, IfFalseBB);
+    if (DTU)
+      DTU->applyUpdatesPermissive(
+          {{DominatorTree::Delete, BI->getParent(), OldSuccessor},
+           {DominatorTree::Insert, BI->getParent(), IfFalseBB}});
     return true;
   }
   if (BI->getSuccessor(0) != IfFalseBB && // no inf looping
       BI->getSuccessor(0)->getTerminatingDeoptimizeCall() && // profitability
       NoSideEffects(*BI->getParent())) {
-    BI->getSuccessor(0)->removePredecessor(BI->getParent());
+    auto *OldSuccessor = BI->getSuccessor(0);
+    OldSuccessor->removePredecessor(BI->getParent());
     BI->setSuccessor(0, IfFalseBB);
+    if (DTU)
+      DTU->applyUpdatesPermissive(
+          {{DominatorTree::Delete, BI->getParent(), OldSuccessor},
+           {DominatorTree::Insert, BI->getParent(), IfFalseBB}});
     return true;
   }
   return false;
@@ -3626,7 +3637,7 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
   // If the previous block ended with a widenable branch, determine if reusing
   // the target block is profitable and legal.  This will have the effect of
   // "widening" PBI, but doesn't require us to reason about hosting safety.
-  if (tryWidenCondBranchToCondBranch(PBI, BI))
+  if (tryWidenCondBranchToCondBranch(PBI, BI, DTU))
     return true;
 
   if (auto *CE = dyn_cast<ConstantExpr>(BI->getCondition()))
diff --git a/llvm/test/Transforms/SimplifyCFG/wc-widen-block.ll b/llvm/test/Transforms/SimplifyCFG/wc-widen-block.ll
index 6e959bc37ed1..aefadabf3bb0 100644
--- a/llvm/test/Transforms/SimplifyCFG/wc-widen-block.ll
+++ b/llvm/test/Transforms/SimplifyCFG/wc-widen-block.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -passes=simplify-cfg -S < %s | FileCheck %s
+; RUN: opt -passes=simplify-cfg -simplifycfg-require-and-preserve-domtree=1 -S < %s | FileCheck %s
 
 define i32 @basic(i1 %cond_0, i32* %p) {
 ; CHECK-LABEL: @basic(