[llvm] 39a56f7 - [SimplifyCFG] Teach SimplifyTerminatorOnSelect() to preserve DomTree

[Roman Lebedev via llvm-commits]

Author: Roman Lebedev
Date: 2020-12-30T00:48:12+03:00
New Revision: 39a56f7f1722b1e917a3bd5c829ec1d7effd5a11

URL: https://github.com/llvm/llvm-project/commit/39a56f7f1722b1e917a3bd5c829ec1d7effd5a11
DIFF: https://github.com/llvm/llvm-project/commit/39a56f7f1722b1e917a3bd5c829ec1d7effd5a11.diff

LOG: [SimplifyCFG] Teach SimplifyTerminatorOnSelect() to preserve DomTree

Added: 
    

Modified: 
    llvm/lib/Transforms/Utils/SimplifyCFG.cpp
    llvm/test/Transforms/SimplifyCFG/switch-on-const-select.ll

Removed: 
    


################################################################################
diff  --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index d1851da69138..24572810be70 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -3805,6 +3805,8 @@ bool SimplifyCFGOpt::SimplifyTerminatorOnSelect(Instruction *OldTerm,
   BasicBlock *KeepEdge1 = TrueBB;
   BasicBlock *KeepEdge2 = TrueBB != FalseBB ? FalseBB : nullptr;
 
+  SmallVector<DominatorTree::UpdateType, 4> Updates;
+
   // Then remove the rest.
   for (BasicBlock *Succ : successors(OldTerm)) {
     // Make sure only to keep exactly one copy of each edge.
@@ -3812,9 +3814,11 @@ bool SimplifyCFGOpt::SimplifyTerminatorOnSelect(Instruction *OldTerm,
       KeepEdge1 = nullptr;
     else if (Succ == KeepEdge2)
       KeepEdge2 = nullptr;
-    else
+    else {
       Succ->removePredecessor(OldTerm->getParent(),
                               /*KeepOneInputPHIs=*/true);
+      Updates.push_back({DominatorTree::Delete, OldTerm->getParent(), Succ});
+    }
   }
 
   IRBuilder<> Builder(OldTerm);
@@ -3822,14 +3826,17 @@ bool SimplifyCFGOpt::SimplifyTerminatorOnSelect(Instruction *OldTerm,
 
   // Insert an appropriate new terminator.
   if (!KeepEdge1 && !KeepEdge2) {
-    if (TrueBB == FalseBB)
+    if (TrueBB == FalseBB) {
       // We were only looking for one successor, and it was present.
       // Create an unconditional branch to it.
       Builder.CreateBr(TrueBB);
-    else {
+      Updates.push_back({DominatorTree::Insert, OldTerm->getParent(), TrueBB});
+    } else {
       // We found both of the successors we were looking for.
       // Create a conditional branch sharing the condition of the select.
       BranchInst *NewBI = Builder.CreateCondBr(Cond, TrueBB, FalseBB);
+      Updates.push_back({DominatorTree::Insert, OldTerm->getParent(), TrueBB});
+      Updates.push_back({DominatorTree::Insert, OldTerm->getParent(), FalseBB});
       if (TrueWeight != FalseWeight)
         setBranchWeights(NewBI, TrueWeight, FalseWeight);
     }
@@ -3841,15 +3848,20 @@ bool SimplifyCFGOpt::SimplifyTerminatorOnSelect(Instruction *OldTerm,
     // One of the selected values was a successor, but the other wasn't.
     // Insert an unconditional branch to the one that was found;
     // the edge to the one that wasn't must be unreachable.
-    if (!KeepEdge1)
+    if (!KeepEdge1) {
       // Only TrueBB was found.
       Builder.CreateBr(TrueBB);
-    else
+      Updates.push_back({DominatorTree::Insert, OldTerm->getParent(), TrueBB});
+    } else {
       // Only FalseBB was found.
       Builder.CreateBr(FalseBB);
+      Updates.push_back({DominatorTree::Insert, OldTerm->getParent(), FalseBB});
+    }
   }
 
   EraseTerminatorAndDCECond(OldTerm);
+  if (DTU)
+    DTU->applyUpdatesPermissive(Updates);
   return true;
 }
 

diff  --git a/llvm/test/Transforms/SimplifyCFG/switch-on-const-select.ll b/llvm/test/Transforms/SimplifyCFG/switch-on-const-select.ll
index 98c434a5a0ec..013096625b1c 100644
--- a/llvm/test/Transforms/SimplifyCFG/switch-on-const-select.ll
+++ b/llvm/test/Transforms/SimplifyCFG/switch-on-const-select.ll
@@ -1,5 +1,5 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt < %s -simplifycfg -S | FileCheck -enable-var-scope %s
+; RUN: opt < %s -simplifycfg -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck -enable-var-scope %s
 
 ; Test basic folding to a conditional branch.
 define i32 @foo(i64 %x, i64 %y) nounwind {