2 files changed, 85 insertions, 12 deletions

diff --git a/lib/MC/MCAssembler.cpp b/lib/MC/MCAssembler.cpp
index 321b2832b9..e29ea25102 100644
--- a/lib/MC/MCAssembler.cpp
+++ b/lib/MC/MCAssembler.cpp
@@ -25,6 +25,7 @@
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"  // @LOCALMOD
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Support/TargetRegistry.h"
 
@@ -186,13 +187,14 @@ MCFragment::~MCFragment() {
 }
 
 MCFragment::MCFragment(FragmentType _Kind, MCSectionData *_Parent)
-  : Kind(_Kind), Parent(_Parent), Atom(0), Offset(~UINT64_C(0)),
+  : Kind(_Kind),
     // @LOCALMOD-BEGIN
     BundleAlign(BundleAlignNone),
     BundleGroupStart(false),
     BundleGroupEnd(false),
-    BundlePadding(~UINT8_C(0))
+    BundlePadding(~UINT8_C(0)),
     // @LOCALMOD-END
+    Parent(_Parent), Atom(0), Offset(~UINT64_C(0))
 {
   if (Parent)
     Parent->getFragmentList().push_back(this);
@@ -226,14 +228,16 @@ MCSectionData::MCSectionData(const MCSection &_Section, MCAssembler *A)
     BundlingEnabled(false),
     BundleLocked(false),
     BundleGroupFirstFrag(false),
-    BundleAlignNext(MCFragment::BundleAlignNone)
+    BundleAlignNext(MCFragment::BundleAlignNone),
+    BundleOffsetKnown(false),
+    BundleOffset(0)
 // @LOCALMOD-END
 {
   if (A)
     A->getSectionList().push_back(this);
 
   // @LOCALMOD-BEGIN
-  unsigned BundleSize = A->getBackend().getBundleSize();
+  BundleSize = A->getBackend().getBundleSize();
   if (BundleSize && _Section.UseCodeAlign()) {
     BundlingEnabled = true;
     setAlignment(BundleSize);
@@ -241,6 +245,69 @@ MCSectionData::MCSectionData(const MCSection &_Section, MCAssembler *A)
   // @LOCALMOD-END
 }
 
+// @LOCALMOD-BEGIN
+void MCSectionData::MarkBundleOffsetUnknown() {
+  BundleOffsetKnown = false;
+  BundleOffset = 0;
+}
+
+// Only create a new fragment if:
+// 1) we are emitting the first instruction of a bundle locked sequence.
+// 2) we are not currently emitting a bundle locked sequence and we cannot
+//    guarantee the instruction would not span a bundle boundary.
+// Otherwise, append to the current fragment to reduce the number of fragments.
+bool MCSectionData::ShouldCreateNewFragment(size_t Size) {
+  // The first instruction of a bundle locked region starts a new fragment.
+  if (isBundleLocked() && isBundleGroupFirstFrag())
+    return true;
+  // Unless we know the relative offset of the end of the current fragment,
+  // we need to create a new fragment.
+  if (!isBundleLocked() && !BundleOffsetKnown)
+    return true;
+  assert(BundleSize != 0 && "BundleSize needs to be non-zero");
+  assert(Size < BundleSize && "Instruction size must be less than BundleSize");
+  // If inserting the instruction would overlap a bundle boundary, start a
+  // new fragment.
+  // TODO(sehr): we could still explicitly insert a NOP and continue here.
+  if (BundleOffset + (unsigned) Size > BundleSize)
+    return true;
+  return false;
+}
+
+void MCSectionData::UpdateBundleOffset(size_t Size) {
+  // A bundle locked fragment could move if it spans a bundle boundary.
+  if (isBundleLocked()) {
+    BundleOffsetKnown = false;
+    return;
+  }
+  // If inserting the instruction would overlap a bundle boundary, starting a
+  // new fragment moves the known offset to the end of the instruction in the
+  // next bundle.
+  // TODO(sehr): we could insert a NOP and continue the fragment.
+  if (BundleOffset + (unsigned) Size > BundleSize)
+    BundleOffset = Size;
+  else
+    BundleOffset = BundleOffset + Size;
+}
+
+void MCSectionData::AlignBundleOffsetTo(size_t AlignBase) {
+  // If BundleOffset is already known, an alignment just moves bundleOffset.
+  if (BundleOffsetKnown) {
+    BundleOffset = RoundUpToAlignment(BundleOffset, AlignBase);
+    return;
+  }
+  // Otherwise, if AlignBase is at least as big as a bundle, then we know the
+  // offset relative to a bundle start.
+  if (AlignBase >= BundleSize) {
+    BundleOffsetKnown = true;
+    BundleOffset = 0;
+  } else {
+    BundleOffsetKnown = false;
+    BundleOffset = 0;
+  }
+}
+// @LOCALMOD-END
+
 /* *** */
 
 MCSymbolData::MCSymbolData() : Symbol(0) {}
diff --git a/lib/MC/MCELFStreamer.cpp b/lib/MC/MCELFStreamer.cpp
index 26b171ed21..b25f8efcd4 100644
--- a/lib/MC/MCELFStreamer.cpp
+++ b/lib/MC/MCELFStreamer.cpp
@@ -359,6 +359,7 @@ void MCELFStreamer::EmitBytes(StringRef Data, unsigned AddrSpace) {
   // TODO: This is exactly the same as WinCOFFStreamer. Consider merging into
   // MCObjectStreamer.
   getOrCreateDataFragment()->getContents().append(Data.begin(), Data.end());
+  getCurrentSectionData()->MarkBundleOffsetUnknown();  // @LOCALMOD
 }
 
 void MCELFStreamer::EmitValueToAlignment(unsigned ByteAlignment,
@@ -371,6 +372,10 @@ void MCELFStreamer::EmitValueToAlignment(unsigned ByteAlignment,
   new MCAlignFragment(ByteAlignment, Value, ValueSize, MaxBytesToEmit,
                       getCurrentSectionData());
 
+  // @LOCALMOD-BEGIN
+  // Bump the bundle offset to account for alignment.
+  getCurrentSectionData()->AlignBundleOffsetTo(ByteAlignment);
+  // @LOCALMOD-END
   // Update the maximum alignment on the current section if necessary.
   if (ByteAlignment > getCurrentSectionData()->getAlignment())
     getCurrentSectionData()->setAlignment(ByteAlignment);
@@ -386,6 +391,10 @@ void MCELFStreamer::EmitCodeAlignment(unsigned ByteAlignment,
                                            getCurrentSectionData());
   F->setEmitNops(true);
 
+  // @LOCALMOD-BEGIN
+  // Bump the bundle offset to account for alignment.
+  getCurrentSectionData()->AlignBundleOffsetTo(ByteAlignment);
+  // @LOCALMOD-END
   // Update the maximum alignment on the current section if necessary.
   if (ByteAlignment > getCurrentSectionData()->getAlignment())
     getCurrentSectionData()->setAlignment(ByteAlignment);
@@ -395,6 +404,7 @@ void MCELFStreamer::EmitValueImpl(const MCExpr *Value, unsigned Size,
                                   unsigned AddrSpace) {
   fixSymbolsInTLSFixups(Value);
   MCObjectStreamer::EmitValueImpl(Value, Size, AddrSpace);
+  getCurrentSectionData()->MarkBundleOffsetUnknown();  // @LOCALMOD
 }
 
 
@@ -463,6 +473,7 @@ void MCELFStreamer::EmitInstToFragment(const MCInst &Inst) {
 
   for (unsigned i = 0, e = F.getFixups().size(); i != e; ++i)
     fixSymbolsInTLSFixups(F.getFixups()[i].getValue());
+  getCurrentSectionData()->MarkBundleOffsetUnknown();  // @LOCALMOD
 }
 
 void MCELFStreamer::EmitInstToData(const MCInst &Inst) {
@@ -486,19 +497,14 @@ void MCELFStreamer::EmitInstToData(const MCInst &Inst) {
       DF->addFixup(Fixups[i]);
     }
     DF->getContents().append(Code.begin(), Code.end());
+    getCurrentSectionData()->UpdateBundleOffset(Code.size());
   } else {
-    // Only create a new fragment if:
-    // 1) there is no current fragment,
-    // 2) we are not currently emitting a bundle locked sequence, or
-    // 3) we are emitting the first instruction of a bundle locked sequence.
-    // Otherwise, append to the current fragment to reduce the number of
-    // fragments.
     MCTinyFragment *TF = dyn_cast_or_null<MCTinyFragment>(getCurrentFragment());
     MCSectionData *SD = getCurrentSectionData();
-    if (!TF || !SD->isBundleLocked() || SD->isBundleGroupFirstFrag()) {
+    if (!TF || SD->ShouldCreateNewFragment(Code.size()))
       TF = new MCTinyFragment(SD);
-    }
     TF->getContents().append(Code.begin(), Code.end());
+    SD->UpdateBundleOffset(Code.size());
   }
   // @LOCALMOD-END
 }