MC CFG: Split MCBasicBlocks to mirror atom splitting.

When an MCTextAtom is split, all MCBasicBlocks backed by it are
automatically split, with a fallthrough between both blocks, and
the successors moved to the second block.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@188881 91177308-0d34-0410-b5e6-96231b3b80d8

5 files changed, 65 insertions, 2 deletions

diff --git a/include/llvm/MC/MCFunction.h b/include/llvm/MC/MCFunction.h
index 697b500..3fccaac 100644
--- a/include/llvm/MC/MCFunction.h
+++ b/include/llvm/MC/MCFunction.h
@@ -70,6 +70,14 @@ public:
 
   void addPredecessor(const MCBasicBlock *MCBB);
   bool isPredecessor(const MCBasicBlock *MCBB) const;
+
+  /// \brief Split block, mirrorring NewAtom = Insts->split(..).
+  /// This moves all successors to \p SplitBB, and
+  /// adds a fallthrough to it.
+  /// \p SplitBB The result of splitting Insts, a basic block directly following
+  /// this basic block.
+  /// \returns A new basic block, backed by \p SplitBB.
+  void splitBasicBlock(MCBasicBlock *SplitBB);
   /// @}
 };
 
diff --git a/include/llvm/MC/MCModule.h b/include/llvm/MC/MCModule.h
index c0c242c..63635c7 100644
--- a/include/llvm/MC/MCModule.h
+++ b/include/llvm/MC/MCModule.h
@@ -56,6 +56,21 @@ class MCModule {
   void map(MCAtom *NewAtom);
   /// @}
 
+  /// \name Basic block tracking
+  /// @{
+  typedef std::vector<MCBasicBlock*> BBsByAtomTy;
+  BBsByAtomTy BBsByAtom;
+
+  // For access to basic block > atom tracking.
+  friend class MCBasicBlock;
+  friend class MCTextAtom;
+
+  /// \brief Keep track of \p BBBackedByAtom as being backed by \p Atom.
+  /// This is used to update succs/preds when \p Atom is split.
+  void trackBBForAtom(const MCTextAtom *Atom, MCBasicBlock *BBBackedByAtom);
+  void splitBasicBlocksForAtom(const MCTextAtom *TA, const MCTextAtom *NewTA);
+  /// @}
+
   /// \name Function tracking
   /// @{
   typedef std::vector<MCFunction*> FunctionListTy;
diff --git a/lib/MC/MCAtom.cpp b/lib/MC/MCAtom.cpp
index f3ea6c3..17a7468 100644
--- a/lib/MC/MCAtom.cpp
+++ b/lib/MC/MCAtom.cpp
@@ -106,5 +106,6 @@ MCTextAtom *MCTextAtom::split(uint64_t SplitPt) {
 
   std::copy(I, Insts.end(), std::back_inserter(RightAtom->Insts));
   Insts.erase(I, Insts.end());
+  Parent->splitBasicBlocksForAtom(this, RightAtom);
   return RightAtom;
 }
diff --git a/lib/MC/MCFunction.cpp b/lib/MC/MCFunction.cpp
index 300ab5b..85e88e5 100644
--- a/lib/MC/MCFunction.cpp
+++ b/lib/MC/MCFunction.cpp
@@ -46,8 +46,9 @@ MCBasicBlock *MCFunction::find(uint64_t StartAddr) {
 // MCBasicBlock
 
 MCBasicBlock::MCBasicBlock(const MCTextAtom &Insts, MCFunction *Parent)
-  : Insts(&Insts), Parent(Parent)
-{}
+  : Insts(&Insts), Parent(Parent) {
+  getParent()->getParent()->trackBBForAtom(&Insts, this);
+}
 
 void MCBasicBlock::addSuccessor(const MCBasicBlock *MCBB) {
   if (!isSuccessor(MCBB))
@@ -68,3 +69,14 @@ bool MCBasicBlock::isPredecessor(const MCBasicBlock *MCBB) const {
   return std::find(Predecessors.begin(), Predecessors.end(),
                    MCBB) != Predecessors.end();
 }
+
+void MCBasicBlock::splitBasicBlock(MCBasicBlock *SplitBB) {
+  assert(Insts->getEndAddr() + 1 == SplitBB->Insts->getBeginAddr() &&
+         "Splitting unrelated basic blocks!");
+  SplitBB->addPredecessor(this);
+  assert(SplitBB->Successors.empty() &&
+         "Split basic block shouldn't already have successors!");
+  SplitBB->Successors = Successors;
+  Successors.clear();
+  addSuccessor(SplitBB);
+}
diff --git a/lib/MC/MCModule.cpp b/lib/MC/MCModule.cpp
index bdd5cc6..7e9e18a 100644
--- a/lib/MC/MCModule.cpp
+++ b/lib/MC/MCModule.cpp
@@ -103,6 +103,33 @@ MCFunction *MCModule::createFunction(StringRef Name) {
   return Functions.back();
 }
 
+static bool CompBBToAtom(MCBasicBlock *BB, const MCTextAtom *Atom) {
+  return BB->getInsts() < Atom;
+}
+
+void MCModule::splitBasicBlocksForAtom(const MCTextAtom *TA,
+                                       const MCTextAtom *NewTA) {
+  BBsByAtomTy::iterator
+    I = std::lower_bound(BBsByAtom.begin(), BBsByAtom.end(),
+                         TA, CompBBToAtom);
+  for (; I != BBsByAtom.end() && (*I)->getInsts() == TA; ++I) {
+    MCBasicBlock *BB = *I;
+    MCBasicBlock *NewBB = &BB->getParent()->createBlock(*NewTA);
+    BB->splitBasicBlock(NewBB);
+  }
+}
+
+void MCModule::trackBBForAtom(const MCTextAtom *Atom, MCBasicBlock *BB) {
+  assert(Atom == BB->getInsts() && "Text atom doesn't back the basic block!");
+  BBsByAtomTy::iterator I = std::lower_bound(BBsByAtom.begin(),
+                                             BBsByAtom.end(),
+                                             Atom, CompBBToAtom);
+  for (; I != BBsByAtom.end() && (*I)->getInsts() == Atom; ++I)
+    if (*I == BB)
+      return;
+  BBsByAtom.insert(I, BB);
+}
+
 MCModule::~MCModule() {
   for (AtomListTy::iterator AI = atom_begin(),
                             AE = atom_end();