2 files changed, 50 insertions, 36 deletions

diff --git a/llvm/include/llvm/MC/MCAssembler.h b/llvm/include/llvm/MC/MCAssembler.h
index 095b1743ecb..6522b672d3e 100644
--- a/llvm/include/llvm/MC/MCAssembler.h
+++ b/llvm/include/llvm/MC/MCAssembler.h
@@ -650,6 +650,12 @@ private:
   bool FragmentNeedsRelaxation(const MCInstFragment *IF,
                                const MCAsmLayout &Layout) const;
 
+  /// Compute the effective fragment size assuming it is layed out at the given
+  /// \arg SectionAddress and \arg FragmentOffset.
+  uint64_t ComputeFragmentSize(MCAsmLayout &Layout, const MCFragment &F,
+                               uint64_t SectionAddress,
+                               uint64_t FragmentOffset) const;
+
   /// LayoutFragment - Performs layout of the given \arg Fragment; assuming that
   /// the previous fragment has already been layed out correctly, and the parent
   /// section has been initialized.
diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
index f41b25f15d3..dfc2663650e 100644
--- a/llvm/lib/MC/MCAssembler.cpp
+++ b/llvm/lib/MC/MCAssembler.cpp
@@ -375,52 +375,39 @@ bool MCAssembler::EvaluateFixup(const MCAsmLayout &Layout,
   return IsResolved;
 }
 
-void MCAssembler::LayoutFragment(MCAsmLayout &Layout, MCFragment &F) {
-  uint64_t StartAddress = Layout.getSectionAddress(F.getParent());
-
-  // Get the fragment start address.
-  uint64_t Address = StartAddress;
-  MCSectionData::iterator it = &F;
-  if (MCFragment *Prev = F.getPrevNode())
-    Address = (StartAddress + Layout.getFragmentOffset(Prev) +
-               Layout.getFragmentEffectiveSize(Prev));
-
-  ++stats::FragmentLayouts;
-
-  uint64_t FragmentOffset = Address - StartAddress;
-  Layout.setFragmentOffset(&F, FragmentOffset);
-
-  // Evaluate fragment size.
-  uint64_t EffectiveSize = 0;
+uint64_t MCAssembler::ComputeFragmentSize(MCAsmLayout &Layout,
+                                          const MCFragment &F,
+                                          uint64_t SectionAddress,
+                                          uint64_t FragmentOffset) const {
   switch (F.getKind()) {
+  case MCFragment::FT_Data:
+    return cast<MCDataFragment>(F).getContents().size();
+  case MCFragment::FT_Fill:
+    return cast<MCFillFragment>(F).getSize();
+  case MCFragment::FT_Inst:
+    return cast<MCInstFragment>(F).getInstSize();
+
   case MCFragment::FT_Align: {
-    MCAlignFragment &AF = cast<MCAlignFragment>(F);
+    const MCAlignFragment &AF = cast<MCAlignFragment>(F);
 
     assert((!AF.hasOnlyAlignAddress() || !AF.getNextNode()) &&
            "Invalid OnlyAlignAddress bit, not the last fragment!");
 
-    EffectiveSize = OffsetToAlignment(Address, AF.getAlignment());
-    if (EffectiveSize > AF.getMaxBytesToEmit())
-      EffectiveSize = 0;
-    break;
-  }
+    uint64_t Size = OffsetToAlignment(SectionAddress + FragmentOffset,
+                                      AF.getAlignment());
 
-  case MCFragment::FT_Data:
-    EffectiveSize = cast<MCDataFragment>(F).getContents().size();
-    break;
+    // Honor MaxBytesToEmit.
+    if (Size > AF.getMaxBytesToEmit())
+      return 0;
 
-  case MCFragment::FT_Fill: {
-    EffectiveSize = cast<MCFillFragment>(F).getSize();
-    break;
+    return Size;
   }
 
-  case MCFragment::FT_Inst:
-    EffectiveSize = cast<MCInstFragment>(F).getInstSize();
-    break;
-
   case MCFragment::FT_Org: {
-    MCOrgFragment &OF = cast<MCOrgFragment>(F);
+    const MCOrgFragment &OF = cast<MCOrgFragment>(F);
 
+    // FIXME: We should compute this sooner, we don't want to recurse here, and
+    // we would like to be more functional.
     int64_t TargetLocation;
     if (!OF.getOffset().EvaluateAsAbsolute(TargetLocation, &Layout))
       report_fatal_error("expected assembly-time absolute expression");
@@ -431,11 +418,32 @@ void MCAssembler::LayoutFragment(MCAsmLayout &Layout, MCFragment &F) {
       report_fatal_error("invalid .org offset '" + Twine(TargetLocation) +
                          "' (at offset '" + Twine(FragmentOffset) + "'");
 
-    EffectiveSize = Offset;
-    break;
+    return Offset;
   }
   }
 
+  assert(0 && "invalid fragment kind");
+  return 0;
+}
+
+void MCAssembler::LayoutFragment(MCAsmLayout &Layout, MCFragment &F) {
+  uint64_t StartAddress = Layout.getSectionAddress(F.getParent());
+
+  // Get the fragment start address.
+  uint64_t Address = StartAddress;
+  MCSectionData::iterator it = &F;
+  if (MCFragment *Prev = F.getPrevNode())
+    Address = (StartAddress + Layout.getFragmentOffset(Prev) +
+               Layout.getFragmentEffectiveSize(Prev));
+
+  ++stats::FragmentLayouts;
+
+  uint64_t FragmentOffset = Address - StartAddress;
+  Layout.setFragmentOffset(&F, FragmentOffset);
+
+  // Evaluate fragment size.
+  uint64_t EffectiveSize = ComputeFragmentSize(Layout, F, StartAddress,
+                                               FragmentOffset);
   Layout.setFragmentEffectiveSize(&F, EffectiveSize);
 }