Index: lib/Target/AArch64/AArch64BranchRelaxation.cpp
===================================================================
--- lib/Target/AArch64/AArch64BranchRelaxation.cpp
+++ lib/Target/AArch64/AArch64BranchRelaxation.cpp
@@ -80,6 +80,13 @@
   void adjustBlockOffsets(MachineBasicBlock &MBB);
   bool isBlockInRange(const MachineInstr &MI, const MachineBasicBlock &BB) const;
   void invertConditionalBranch(MachineInstr &MI) const;
+  unsigned insertInvertedConditionalBranch(MachineBasicBlock &MBB,
+                                           const MachineInstr &OldBr,
+                                           MachineBasicBlock &NewDestBB) const;
+  unsigned insertUnconditionalBranch(MachineBasicBlock &MBB,
+                                     MachineBasicBlock &NewDestBB,
+                                     const DebugLoc &DL) const;
+
   bool fixupConditionalBranch(MachineInstr &MI);
   void computeBlockSize(const MachineBasicBlock &MBB);
   unsigned getInstrOffset(const MachineInstr &MI) const;
@@ -127,21 +134,20 @@
   }
 }
 
-/// BBHasFallthrough - Return true if the specified basic block can fallthrough
+// FIXME: This is a less precise version of MachineBasicBlock::canFallThrough?
+
+/// \returns true if the specified basic block can fallthrough
 /// into the block immediately after it.
-static bool BBHasFallthrough(MachineBasicBlock *MBB) {
+static bool hasFallthrough(const MachineBasicBlock &MBB) {
   // Get the next machine basic block in the function.
-  MachineFunction::iterator MBBI(MBB);
+  MachineFunction::const_iterator MBBI(MBB);
+
   // Can't fall off end of function.
   auto NextBB = std::next(MBBI);
-  if (NextBB == MBB->getParent()->end())
+  if (NextBB == MBB.getParent()->end())
     return false;
 
-  for (MachineBasicBlock *S : MBB->successors())
-    if (S == &*NextBB)
-      return true;
-
-  return false;
+  return MBB.isSuccessor(&*NextBB);
 }
 
 /// scanFunction - Do the initial scan of the function, building up
@@ -321,6 +327,45 @@
     invertBccCondition(MI);
 }
 
+/// Insert a conditional branch at the end of \p MBB to \p NewDestBB, using the
+/// inverse condition of branch \p OldBr.
+/// \returns The number of bytes added to the block.
+unsigned AArch64BranchRelaxation::insertInvertedConditionalBranch(
+  MachineBasicBlock &MBB,
+  const MachineInstr &OldBr,
+  MachineBasicBlock &NewDestBB) const {
+  unsigned OppositeCondOpc = getOppositeConditionOpcode(OldBr.getOpcode());
+
+  MachineInstrBuilder MIB =
+    BuildMI(&MBB, OldBr.getDebugLoc(), TII->get(OppositeCondOpc))
+    .addOperand(OldBr.getOperand(0));
+
+  unsigned Opc = OldBr.getOpcode();
+
+  if (Opc == AArch64::TBZW || Opc == AArch64::TBNZW ||
+      Opc == AArch64::TBZX || Opc == AArch64::TBNZX)
+    MIB.addOperand(OldBr.getOperand(1));
+
+  if (OldBr.getOpcode() == AArch64::Bcc)
+    invertBccCondition(*MIB);
+
+  MIB.addMBB(&NewDestBB);
+
+  return TII->GetInstSizeInBytes(*MIB);
+}
+
+/// Insert an unconditional branch at the end of \p MBB to \p DestBB.
+/// \returns the number of bytes emitted.
+unsigned AArch64BranchRelaxation::insertUnconditionalBranch(
+  MachineBasicBlock &MBB,
+  MachineBasicBlock &DestBB,
+  const DebugLoc &DL) const {
+  MachineInstr *MI = BuildMI(&MBB, DL, TII->get(AArch64::B))
+    .addMBB(&DestBB);
+
+  return TII->GetInstSizeInBytes(*MI);
+}
+
 static void changeBranchDestBlock(MachineInstr &MI,
                                   MachineBasicBlock &NewDestBB) {
   unsigned Opc = MI.getOpcode();
@@ -350,7 +395,7 @@
   // split the MBB before the next instruction.
   MachineBasicBlock *MBB = MI.getParent();
   MachineInstr *BMI = &MBB->back();
-  bool NeedSplit = (BMI != &MI) || !BBHasFallthrough(MBB);
+  bool NeedSplit = (BMI != &MI) || !hasFallthrough(*MBB);
 
   if (BMI != &MI) {
     if (std::next(MachineBasicBlock::iterator(MI)) ==
@@ -394,26 +439,19 @@
     MBB->replaceSuccessor(FBB, NewBB);
     NewBB->addSuccessor(FBB);
   }
+
   MachineBasicBlock &NextBB = *std::next(MachineFunction::iterator(MBB));
 
   DEBUG(dbgs() << "  Insert B to BB#" << DestBB->getNumber()
                << ", invert condition and change dest. to BB#"
                << NextBB.getNumber() << '\n');
 
-  unsigned OppositeCondOpc = getOppositeConditionOpcode(MI.getOpcode());
   // Insert a new conditional branch and a new unconditional branch.
-  MachineInstrBuilder MIB = BuildMI(MBB, DebugLoc(), TII->get(OppositeCondOpc))
-    .addOperand(MI.getOperand(0));
+  BlockInfo[MBB->getNumber()].Size
+    += insertInvertedConditionalBranch(*MBB, MI, NextBB);
 
-  if (MI.getOpcode() == AArch64::TBZW || MI.getOpcode() == AArch64::TBNZW ||
-      MI.getOpcode() == AArch64::TBZX || MI.getOpcode() == AArch64::TBNZX)
-    MIB.addOperand(MI.getOperand(1));
-  if (MI.getOpcode() == AArch64::Bcc)
-    invertBccCondition(*MIB);
-  MIB.addMBB(&NextBB);
-  BlockInfo[MBB->getNumber()].Size += TII->GetInstSizeInBytes(MBB->back());
-  BuildMI(MBB, DebugLoc(), TII->get(AArch64::B)).addMBB(DestBB);
-  BlockInfo[MBB->getNumber()].Size += TII->GetInstSizeInBytes(MBB->back());
+  BlockInfo[MBB->getNumber()].Size +=
+    insertUnconditionalBranch(*MBB, *DestBB, MI.getDebugLoc());
 
   // Remove the old conditional branch.  It may or may not still be in MBB.
   BlockInfo[MBB->getNumber()].Size -= TII->GetInstSizeInBytes(MI);