1 files changed, 42 insertions, 2 deletions
diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
index 671f6f7..fe0419b 100644
--- a/llvm/lib/MC/MCAssembler.cpp
+++ b/llvm/lib/MC/MCAssembler.cpp
@@ -403,6 +403,46 @@ uint64_t MCAssembler::computeFragmentSize(const MCAsmLayout &Layout,
   llvm_unreachable("invalid fragment kind");
 }
 
+// Compute the amount of padding required before the fragment \p F to
+// obey bundling restrictions, where \p FOffset is the fragment's offset in
+// its section and \p FSize is the fragment's size.
+static uint64_t computeBundlePadding(unsigned BundleSize,
+                                     const MCEncodedFragment *F,
+                                     uint64_t FOffset, uint64_t FSize) {
+  uint64_t OffsetInBundle = FOffset & (BundleSize - 1);
+  uint64_t EndOfFragment = OffsetInBundle + FSize;
+
+  // There are two kinds of bundling restrictions:
+  //
+  // 1) For alignToBundleEnd(), add padding to ensure that the fragment will
+  //    *end* on a bundle boundary.
+  // 2) Otherwise, check if the fragment would cross a bundle boundary. If it
+  //    would, add padding until the end of the bundle so that the fragment
+  //    will start in a new one.
+  if (F->alignToBundleEnd()) {
+    // Three possibilities here:
+    //
+    // A) The fragment just happens to end at a bundle boundary, so we're good.
+    // B) The fragment ends before the current bundle boundary: pad it just
+    //    enough to reach the boundary.
+    // C) The fragment ends after the current bundle boundary: pad it until it
+    //    reaches the end of the next bundle boundary.
+    //
+    // Note: this code could be made shorter with some modulo trickery, but it's
+    // intentionally kept in its more explicit form for simplicity.
+    if (EndOfFragment == BundleSize)
+      return 0;
+    else if (EndOfFragment < BundleSize)
+      return BundleSize - EndOfFragment;
+    else { // EndOfFragment > BundleSize
+      return 2 * BundleSize - EndOfFragment;
+    }
+  } else if (OffsetInBundle > 0 && EndOfFragment > BundleSize)
+    return BundleSize - OffsetInBundle;
+  else
+    return 0;
+}
+
 void MCAsmLayout::layoutBundle(MCFragment *Prev, MCFragment *F) {
   // If bundling is enabled and this fragment has instructions in it, it has to
   // obey the bundling restrictions. With padding, we'll have:
@@ -433,8 +473,8 @@ void MCAsmLayout::layoutBundle(MCFragment *Prev, MCFragment *F) {
   if (FSize > Assembler.getBundleAlignSize())
     report_fatal_error("Fragment can't be larger than a bundle size");
 
-  uint64_t RequiredBundlePadding =
-      computeBundlePadding(Assembler, EF, EF->Offset, FSize);
+  uint64_t RequiredBundlePadding = computeBundlePadding(
+      Assembler.getBundleAlignSize(), EF, EF->Offset, FSize);
   if (RequiredBundlePadding > UINT8_MAX)
     report_fatal_error("Padding cannot exceed 255 bytes");
   EF->setBundlePadding(static_cast<uint8_t>(RequiredBundlePadding));