1 files changed, 67 insertions, 54 deletions

diff --git a/lib/MC/MCAtom.cpp b/lib/MC/MCAtom.cpp
index d71444324f..2626b39db4 100644
--- a/lib/MC/MCAtom.cpp
+++ b/lib/MC/MCAtom.cpp
@@ -10,88 +10,101 @@
 #include "llvm/MC/MCAtom.h"
 #include "llvm/MC/MCModule.h"
 #include "llvm/Support/ErrorHandling.h"
+#include <iterator>
 
 using namespace llvm;
 
-void MCAtom::addInst(const MCInst &I, uint64_t Address, unsigned Size) {
-  assert(Type == TextAtom && "Trying to add MCInst to a non-text atom!");
-
-  assert(Address < End+Size &&
-         "Instruction not contiguous with end of atom!");
-  if (Address > End)
-    Parent->remap(this, Begin, End+Size);
-
-  Text.push_back(std::make_pair(Address, I));
+void MCAtom::remap(uint64_t NewBegin, uint64_t NewEnd) {
+  Parent->remap(this, NewBegin, NewEnd);
 }
 
-void MCAtom::addData(const MCData &D) {
-  assert(Type == DataAtom && "Trying to add MCData to a non-data atom!");
-  Parent->remap(this, Begin, End+1);
-
-  Data.push_back(D);
+void MCAtom::remapForTruncate(uint64_t TruncPt) {
+  assert((TruncPt >= Begin && TruncPt < End) &&
+         "Truncation point not contained in atom!");
+  remap(Begin, TruncPt);
 }
 
-MCAtom *MCAtom::split(uint64_t SplitPt) {
+void MCAtom::remapForSplit(uint64_t SplitPt,
+                           uint64_t &LBegin, uint64_t &LEnd,
+                           uint64_t &RBegin, uint64_t &REnd) {
   assert((SplitPt > Begin && SplitPt <= End) &&
          "Splitting at point not contained in atom!");
 
   // Compute the new begin/end points.
-  uint64_t LeftBegin = Begin;
-  uint64_t LeftEnd = SplitPt - 1;
-  uint64_t RightBegin = SplitPt;
-  uint64_t RightEnd = End;
+  LBegin = Begin;
+  LEnd = SplitPt - 1;
+  RBegin = SplitPt;
+  REnd = End;
 
   // Remap this atom to become the lower of the two new ones.
-  Parent->remap(this, LeftBegin, LeftEnd);
+  remap(LBegin, LEnd);
+}
 
-  // Create a new atom for the higher atom.
-  MCAtom *RightAtom = Parent->createAtom(Type, RightBegin, RightEnd);
+// MCDataAtom
 
-  // Split the contents of the original atom between it and the new one.  The
-  // precise method depends on whether this is a data or a text atom.
-  if (isDataAtom()) {
-    std::vector<MCData>::iterator I = Data.begin() + (RightBegin - LeftBegin);
+void MCDataAtom::addData(const MCData &D) {
+  Data.push_back(D);
+  if (Data.size() > Begin - End)
+    remap(Begin, End + 1);
+}
 
-    assert(I != Data.end() && "Split point not found in range!");
+void MCDataAtom::truncate(uint64_t TruncPt) {
+  remapForTruncate(TruncPt);
 
-    std::copy(I, Data.end(), RightAtom->Data.end());
-    Data.erase(I, Data.end());
-  } else if (isTextAtom()) {
-    std::vector<std::pair<uint64_t, MCInst> >::iterator I = Text.begin();
+  Data.resize(TruncPt - Begin + 1);
+}
 
-    while (I != Text.end() && I->first < SplitPt) ++I;
+MCDataAtom *MCDataAtom::split(uint64_t SplitPt) {
+  uint64_t LBegin, LEnd, RBegin, REnd;
+  remapForSplit(SplitPt, LBegin, LEnd, RBegin, REnd);
 
-    assert(I != Text.end() && "Split point not found in disassembly!");
-    assert(I->first == SplitPt &&
-           "Split point does not fall on instruction boundary!");
+  MCDataAtom *RightAtom = Parent->createDataAtom(RBegin, REnd);
+  RightAtom->setName(getName());
 
-    std::copy(I, Text.end(), RightAtom->Text.end());
-    Text.erase(I, Text.end());
-  } else
-    llvm_unreachable("Unknown atom type!");
+  std::vector<MCData>::iterator I = Data.begin() + (RBegin - LBegin);
+  assert(I != Data.end() && "Split point not found in range!");
 
+  std::copy(I, Data.end(), std::back_inserter(RightAtom->Data));
+  Data.erase(I, Data.end());
   return RightAtom;
 }
 
-void MCAtom::truncate(uint64_t TruncPt) {
-  assert((TruncPt >= Begin && TruncPt < End) &&
-         "Truncation point not contained in atom!");
+// MCTextAtom
 
-  Parent->remap(this, Begin, TruncPt);
+void MCTextAtom::addInst(const MCInst &I, uint64_t Size) {
+  if (NextInstAddress > End)
+    remap(Begin, NextInstAddress);
+  Insts.push_back(MCDecodedInst(I, NextInstAddress, Size));
+  NextInstAddress += Size;
+}
 
-  if (isDataAtom()) {
-    Data.resize(TruncPt - Begin + 1);
-  } else if (isTextAtom()) {
-    std::vector<std::pair<uint64_t, MCInst> >::iterator I = Text.begin();
+void MCTextAtom::truncate(uint64_t TruncPt) {
+  remapForTruncate(TruncPt);
 
-    while (I != Text.end() && I->first <= TruncPt) ++I;
+  InstListTy::iterator I = Insts.begin();
+  while (I != Insts.end() && I->Address <= TruncPt) ++I;
 
-    assert(I != Text.end() && "Truncation point not found in disassembly!");
-    assert(I->first == TruncPt+1 &&
-           "Truncation point does not fall on instruction boundary");
+  assert(I != Insts.end() && "Truncation point not found in disassembly!");
+  assert(I->Address == TruncPt + 1 &&
+         "Truncation point does not fall on instruction boundary");
 
-    Text.erase(I, Text.end());
-  } else
-    llvm_unreachable("Unknown atom type!");
+  Insts.erase(I, Insts.end());
 }
 
+MCTextAtom *MCTextAtom::split(uint64_t SplitPt) {
+  uint64_t LBegin, LEnd, RBegin, REnd;
+  remapForSplit(SplitPt, LBegin, LEnd, RBegin, REnd);
+
+  MCTextAtom *RightAtom = Parent->createTextAtom(RBegin, REnd);
+  RightAtom->setName(getName());
+
+  InstListTy::iterator I = Insts.begin();
+  while (I != Insts.end() && I->Address < SplitPt) ++I;
+  assert(I != Insts.end() && "Split point not found in disassembly!");
+  assert(I->Address == SplitPt &&
+         "Split point does not fall on instruction boundary!");
+
+  std::copy(I, Insts.end(), std::back_inserter(RightAtom->Insts));
+  Insts.erase(I, Insts.end());
+  return RightAtom;
+}