1 files changed, 80 insertions, 9 deletions

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 335ac03..b16fbd1 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -2960,6 +2960,81 @@ static bool FoldTwoEntryPHINode(PHINode *PN, const TargetTransformInfo &TTI,
   return true;
 }
 
+static bool SpeculativelyExecuteThenElseCode(BranchInst *BI,
+                                             const TargetTransformInfo &TTI,
+                                             DomTreeUpdater *DTU,
+                                             const DataLayout &DL) {
+  assert(BI->isConditional() && !isa<ConstantInt>(BI->getCondition()) &&
+         BI->getSuccessor(0) != BI->getSuccessor(1) &&
+         "Only for truly conditional branches.");
+  BasicBlock *BB = BI->getParent();
+
+  // Which ones of our successors end up with an unconditional branch?
+  SmallVector<BasicBlock *, 2> UncondSuccessors;
+  SmallVector<BasicBlock *, 2> OtherSuccessors;
+  for (BasicBlock *Succ : successors(BI)) {
+    auto *SuccBI = dyn_cast<BranchInst>(Succ->getTerminator());
+    if (SuccBI && SuccBI->isUnconditional())
+      UncondSuccessors.emplace_back(Succ);
+    else
+      OtherSuccessors.emplace_back(Succ);
+  }
+  assert(UncondSuccessors.size() + OtherSuccessors.size() == 2 &&
+         "Can not have more than two successors!");
+
+  // If none do, then we can't do anything.
+  if (UncondSuccessors.empty())
+    return false;
+
+  // We want to hoist code from the unconditional block[s] and eliminate them,
+  // but if they have their address taken, then we essentially can't do this.
+  for (BasicBlock *UncondSucc : UncondSuccessors)
+    if (UncondSucc->hasAddressTaken())
+      return false;
+
+  // All unconditional successors must have a single (and the same) predecessor.
+  // FIXME: lift this restriction.
+  for (BasicBlock *UncondSucc : UncondSuccessors)
+    if (!UncondSucc->getSinglePredecessor())
+      return false;
+
+  // Now, what is the merge point?
+  BasicBlock *MergeBB = nullptr;
+  // If there was only a single unconditional successor,
+  // then the other successor *must* be the merge point.
+  if (UncondSuccessors.size() == 1)
+    MergeBB = OtherSuccessors.front();
+
+  // All unconditional successors must have the same successor themselves.
+  for (BasicBlock *UncondSucc : UncondSuccessors) {
+    auto *SuccBI = cast<BranchInst>(UncondSucc->getTerminator());
+    assert(SuccBI->isUnconditional() && "Should be an unconditional branch.");
+    BasicBlock *SuccOfSucc = SuccBI->getSuccessor(0);
+    if (!MergeBB) // First unconditional successor, record it's successor.
+      MergeBB = SuccOfSucc;
+    else if (SuccOfSucc != MergeBB) // Do all succs have the same successor?
+      return false;
+  }
+
+  assert(MergeBB && "Should have found the merge point.");
+  assert(all_of(UncondSuccessors,
+                [MergeBB](BasicBlock *UncondSucc) {
+                  return is_contained(predecessors(MergeBB), UncondSucc);
+                }) &&
+         "All unconditional successors must be predecessors of merge block.");
+  assert((UncondSuccessors.size() != 1 ||
+          is_contained(predecessors(MergeBB), BB)) &&
+         "If there is only a single unconditional successor, then the dispatch "
+         "block must also be merge block's predecessor.");
+
+  if (auto *PN = dyn_cast<PHINode>(MergeBB->begin()))
+    // FIXME: lift this restriction.
+    if (PN->getNumIncomingValues() == 2)
+      return FoldTwoEntryPHINode(PN, TTI, DTU, DL);
+
+  return false;
+}
+
 static Value *createLogicalOp(IRBuilderBase &Builder,
                               Instruction::BinaryOps Opc, Value *LHS,
                               Value *RHS, const Twine &Name = "") {
@@ -6521,6 +6596,11 @@ bool SimplifyCFGOpt::simplifyCondBranch(BranchInst *BI, IRBuilder<> &Builder) {
         return requestResimplify();
   }
 
+  if (Options.FoldTwoEntryPHINode) {
+    if (SpeculativelyExecuteThenElseCode(BI, TTI, DTU, DL))
+      return true;
+  }
+
   // If this is a branch on a phi node in the current block, thread control
   // through this block if any PHI node entries are constants.
   if (PHINode *PN = dyn_cast<PHINode>(BI->getCondition()))
@@ -6736,15 +6816,6 @@ bool SimplifyCFGOpt::simplifyOnce(BasicBlock *BB) {
 
   IRBuilder<> Builder(BB);
 
-  if (Options.FoldTwoEntryPHINode) {
-    // If there is a trivial two-entry PHI node in this basic block, and we can
-    // eliminate it, do so now.
-    if (auto *PN = dyn_cast<PHINode>(BB->begin()))
-      if (PN->getNumIncomingValues() == 2)
-        if (FoldTwoEntryPHINode(PN, TTI, DTU, DL))
-          return true;
-  }
-
   Instruction *Terminator = BB->getTerminator();
   Builder.SetInsertPoint(Terminator);
   switch (Terminator->getOpcode()) {