1 files changed, 35 insertions, 9 deletions
diff --git a/lib/MC/MCFunction.cpp b/lib/MC/MCFunction.cpp
index 2665d3e..767e1e0 100644
--- a/lib/MC/MCFunction.cpp
+++ b/lib/MC/MCFunction.cpp
@@ -9,15 +9,15 @@
 
 #include "llvm/MC/MCFunction.h"
 #include "llvm/MC/MCAtom.h"
-#include "llvm/Support/raw_ostream.h"
+#include "llvm/MC/MCModule.h"
 #include <algorithm>
 
 using namespace llvm;
 
 // MCFunction
 
-MCFunction::MCFunction(StringRef Name)
-  : Name(Name)
+MCFunction::MCFunction(StringRef Name, MCModule *Parent)
+  : Name(Name), ParentModule(Parent)
 {}
 
 MCFunction::~MCFunction() {
@@ -26,18 +26,32 @@ MCFunction::~MCFunction() {
 }
 
 MCBasicBlock &MCFunction::createBlock(const MCTextAtom &TA) {
-  Blocks.push_back(new MCBasicBlock(TA, this));
-  return *Blocks.back();
+  MCBasicBlock *MCBB = new MCBasicBlock(TA, this);
+  Blocks.push_back(MCBB);
+  return *MCBB;
+}
+
+MCBasicBlock *MCFunction::find(uint64_t StartAddr) {
+  for (const_iterator I = begin(), E = end(); I != E; ++I)
+    if ((*I)->getInsts()->getBeginAddr() == StartAddr)
+      return *I;
+  return 0;
+}
+
+const MCBasicBlock *MCFunction::find(uint64_t StartAddr) const {
+  return const_cast<MCFunction *>(this)->find(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) {
-  Successors.push_back(MCBB);
+  if (!isSuccessor(MCBB))
+    Successors.push_back(MCBB);
 }
 
 bool MCBasicBlock::isSuccessor(const MCBasicBlock *MCBB) const {
@@ -46,10 +60,22 @@ bool MCBasicBlock::isSuccessor(const MCBasicBlock *MCBB) const {
 }
 
 void MCBasicBlock::addPredecessor(const MCBasicBlock *MCBB) {
-  Predecessors.push_back(MCBB);
+  if (!isPredecessor(MCBB))
+    Predecessors.push_back(MCBB);
 }
 
 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);
+}