From 0e59c4e3e8f8e105834d137cccb1e1bb731b5a13 Mon Sep 17 00:00:00 2001
From: Peizhao Ou <peizhaoo@uci.edu>
Date: Thu, 16 Nov 2017 17:47:14 -0800
Subject: [PATCH 1/1] Fixes the issue of removing manually added fake
 conditional branches

---
 lib/CodeGen/BranchFolding.cpp           | 29 +++++++++++++++++++++++--
 lib/Target/AArch64/AArch64InstrInfo.cpp | 29 +++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/lib/CodeGen/BranchFolding.cpp b/lib/CodeGen/BranchFolding.cpp
index e0b7d73339d..bd89ee36696 100644
--- a/lib/CodeGen/BranchFolding.cpp
+++ b/lib/CodeGen/BranchFolding.cpp
@@ -493,6 +493,15 @@ static void FixTail(MachineBasicBlock *CurMBB, MachineBasicBlock *SuccBB,
     MachineBasicBlock *NextBB = &*I;
     if (TBB == NextBB && !Cond.empty() && !FBB) {
       if (!TII->ReverseBranchCondition(Cond)) {
+        // XXX-disabled: Don't fold conditional branches that we added
+        // intentionally.
+        MachineBasicBlock::iterator I = CurMBB->getLastNonDebugInstr();
+        if (I != CurMBB->end()) {
+          if (I->isConditionalBranch()) {
+            return;
+          }
+        }
+
         TII->RemoveBranch(*CurMBB);
         TII->InsertBranch(*CurMBB, SuccBB, nullptr, Cond, dl);
         return;
@@ -1254,7 +1263,6 @@ ReoptimizeBlock:
       goto ReoptimizeBlock;
     }
 
-    /*
     // If the previous block unconditionally falls through to this block and
     // this block has no other predecessors, move the contents of this block
     // into the prior block. This doesn't usually happen when SimplifyCFG
@@ -1290,7 +1298,6 @@ ReoptimizeBlock:
       MadeChange = true;
       return MadeChange;
     }
-    */
 
     // If the previous branch *only* branches to *this* block (conditional or
     // not) remove the branch.
@@ -1313,6 +1320,15 @@ ReoptimizeBlock:
     // If the prior block branches somewhere else on the condition and here if
     // the condition is false, remove the uncond second branch.
     if (PriorFBB == MBB) {
+      // XXX-disabled: Don't fold conditional branches that we added
+      // intentionally.
+      MachineBasicBlock::iterator I = PrevBB.getLastNonDebugInstr();
+      if (I != PrevBB.end()) {
+        if (I->isConditionalBranch()) {
+          return MadeChange ;
+        }
+      }
+
       DebugLoc dl = getBranchDebugLoc(PrevBB);
       TII->RemoveBranch(PrevBB);
       TII->InsertBranch(PrevBB, PriorTBB, nullptr, PriorCond, dl);
@@ -1325,6 +1341,15 @@ ReoptimizeBlock:
     // if the branch condition is reversible, reverse the branch to create a
     // fall-through.
     if (PriorTBB == MBB) {
+      // XXX-disabled: Don't fold conditional branches that we added
+      // intentionally.
+      MachineBasicBlock::iterator I = PrevBB.getLastNonDebugInstr();
+      if (I != PrevBB.end()) {
+        if (I->isConditionalBranch()) {
+          return MadeChange ;
+        }
+      }
+
       SmallVector<MachineOperand, 4> NewPriorCond(PriorCond);
       if (!TII->ReverseBranchCondition(NewPriorCond)) {
         DebugLoc dl = getBranchDebugLoc(PrevBB);
diff --git a/lib/Target/AArch64/AArch64InstrInfo.cpp b/lib/Target/AArch64/AArch64InstrInfo.cpp
index f398117de95..c9c982ed7b5 100644
--- a/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -217,6 +217,25 @@ bool AArch64InstrInfo::ReverseBranchCondition(
   return false;
 }
 
+// XXX-update: Returns whether we can remove a conditional branch instruction.
+// If it's one that is mannually added by us, then don't remove it (return
+// false). All their successors are the same.
+static bool shouldRemoveConditionalBranch(MachineInstr* I) {
+  auto* MBB = I->getParent();
+  assert(isCondBranchOpcode(I->getOpcode()));
+  bool SameSuccessor = true;
+  MachineBasicBlock* BB = nullptr;
+  for (auto* Succ : MBB->successors()) {
+    if (!BB) {
+      BB = Succ;
+    }
+    if (BB != Succ) {
+      SameSuccessor = false;
+    }
+  }
+  return !SameSuccessor;
+}
+
 unsigned AArch64InstrInfo::RemoveBranch(MachineBasicBlock &MBB) const {
   MachineBasicBlock::iterator I = MBB.getLastNonDebugInstr();
   if (I == MBB.end())
@@ -226,6 +245,11 @@ unsigned AArch64InstrInfo::RemoveBranch(MachineBasicBlock &MBB) const {
       !isCondBranchOpcode(I->getOpcode()))
     return 0;
 
+  // XXX-update: Don't remove fake conditional branches.
+  if (isCondBranchOpcode(I->getOpcode()) && !shouldRemoveConditionalBranch(I)) {
+    return 0;
+  }
+
   // Remove the branch.
   I->eraseFromParent();
 
@@ -237,6 +261,11 @@ unsigned AArch64InstrInfo::RemoveBranch(MachineBasicBlock &MBB) const {
   if (!isCondBranchOpcode(I->getOpcode()))
     return 1;
 
+  // XXX-update: Don't remove fake conditional branches.
+  if (!shouldRemoveConditionalBranch(I)) {
+    return 1;
+  }
+
   // Remove the branch.
   I->eraseFromParent();
   return 2;
-- 
2.34.1