5 files changed, 179 insertions, 20 deletions

diff --git a/llvm/include/llvm/MC/MCObjectStreamer.h b/llvm/include/llvm/MC/MCObjectStreamer.h
index e10487988f6..d8ea75a0ee7 100644
--- a/llvm/include/llvm/MC/MCObjectStreamer.h
+++ b/llvm/include/llvm/MC/MCObjectStreamer.h
@@ -38,6 +38,8 @@ class MCObjectStreamer : public MCStreamer {
   bool EmitEHFrame;
   bool EmitDebugFrame;
   SmallVector<MCSymbol *, 2> PendingLabels;
+  SmallVector<MCSection*, 2> PendingLabelSections;
+  unsigned CurSubsectionIdx;
   struct PendingMCFixup {
     const MCSymbol *Sym;
     MCFixup Fixup;
@@ -87,17 +89,23 @@ public:
 protected:
   bool changeSectionImpl(MCSection *Section, const MCExpr *Subsection);
 
-  /// If any labels have been emitted but not assigned fragments, ensure that
-  /// they get assigned, either to F if possible or to a new data fragment.
-  /// Optionally, it is also possible to provide an offset \p FOffset, which
-  /// will be used as a symbol offset within the fragment.
+  /// Assign a label to the current Section and Subsection even though a
+  /// fragment is not yet present. Use flushPendingLabels(F) to associate
+  /// a fragment with this label.
+  void addPendingLabel(MCSymbol* label);
+
+  /// If any labels have been emitted but not assigned fragments in the current
+  /// Section and Subsection, ensure that they get assigned, either to fragment
+  /// F if possible or to a new data fragment. Optionally, one can provide an
+  /// offset \p FOffset as a symbol offset within the fragment.
   void flushPendingLabels(MCFragment *F, uint64_t FOffset = 0);
 
 public:
   void visitUsedSymbol(const MCSymbol &Sym) override;
 
-  /// Create a dummy fragment to assign any pending labels.
-  void flushPendingLabels() { flushPendingLabels(nullptr); }
+  /// Create a data fragment for any pending labels across all Sections
+  /// and Subsections.
+  void flushPendingLabels();
 
   MCAssembler &getAssembler() { return *Assembler; }
   MCAssembler *getAssemblerPtr() override;
diff --git a/llvm/include/llvm/MC/MCSection.h b/llvm/include/llvm/MC/MCSection.h
index d057feda87d..c3535b6deaf 100644
--- a/llvm/include/llvm/MC/MCSection.h
+++ b/llvm/include/llvm/MC/MCSection.h
@@ -92,6 +92,15 @@ private:
   /// below that number.
   SmallVector<std::pair<unsigned, MCFragment *>, 1> SubsectionFragmentMap;
 
+  /// State for tracking labels that don't yet have Fragments
+  struct PendingLabel {
+    MCSymbol* Sym;
+    unsigned Subsection;
+    PendingLabel(MCSymbol* Sym, unsigned Subsection = 0)
+      : Sym(Sym), Subsection(Subsection) {}
+  };
+  SmallVector<PendingLabel, 2> PendingLabels;
+
 protected:
   SectionVariant Variant;
   SectionKind Kind;
@@ -187,6 +196,18 @@ public:
   /// Check whether this section is "virtual", that is has no actual object
   /// file contents.
   virtual bool isVirtualSection() const = 0;
+
+  /// Add a pending label for the requested subsection. This label will be
+  /// associated with a fragment in flushPendingLabels()
+  void addPendingLabel(MCSymbol* label, unsigned Subsection = 0);
+
+  /// Associate all pending labels in a subsection with a fragment.
+  void flushPendingLabels(MCFragment *F, uint64_t FOffset = 0,
+			  unsigned Subsection = 0);
+
+  /// Associate all pending labels with empty data fragments. One fragment
+  /// will be created for each subsection as necessary.
+  void flushPendingLabels();
 };
 
 } // end namespace llvm
diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index 6c85f2296cf..1ed6ec1015b 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -42,20 +42,64 @@ MCAssembler *MCObjectStreamer::getAssemblerPtr() {
   return nullptr;
 }
 
+void MCObjectStreamer::addPendingLabel(MCSymbol* S) {
+  MCSection *CurSection = getCurrentSectionOnly();
+  if (CurSection) {
+    // Register labels that have not yet been assigned to a Section.
+    if (!PendingLabels.empty()) {
+      for (MCSymbol* Sym : PendingLabels)
+        CurSection->addPendingLabel(Sym);
+      PendingLabels.clear();
+    }
+
+    // Add this label to the current Section / Subsection.
+    CurSection->addPendingLabel(S, CurSubsectionIdx);
+
+    // Add this Section to the list of PendingLabelSections.
+    auto SecIt = std::find(PendingLabelSections.begin(),
+                           PendingLabelSections.end(), CurSection);
+    if (SecIt == PendingLabelSections.end())
+      PendingLabelSections.push_back(CurSection);
+  }
+  else
+    // There is no Section / Subsection for this label yet.
+    PendingLabels.push_back(S);
+}
+
 void MCObjectStreamer::flushPendingLabels(MCFragment *F, uint64_t FOffset) {
-  if (PendingLabels.empty())
+  MCSection *CurSection = getCurrentSectionOnly();
+  if (!CurSection) {
+    assert(PendingLabels.empty());
     return;
-  if (!F) {
-    F = new MCDataFragment();
-    MCSection *CurSection = getCurrentSectionOnly();
-    CurSection->getFragmentList().insert(CurInsertionPoint, F);
-    F->setParent(CurSection);
   }
-  for (MCSymbol *Sym : PendingLabels) {
-    Sym->setFragment(F);
-    Sym->setOffset(FOffset);
+  // Register labels that have not yet been assigned to a Section.
+  if (!PendingLabels.empty()) {
+    for (MCSymbol* Sym : PendingLabels)
+      CurSection->addPendingLabel(Sym, CurSubsectionIdx);
+    PendingLabels.clear();
   }
-  PendingLabels.clear();
+
+  // Associate a fragment with this label, either the supplied fragment
+  // or an empty data fragment.
+  if (F)
+    CurSection->flushPendingLabels(F, FOffset, CurSubsectionIdx);
+  else
+    CurSection->flushPendingLabels(nullptr, 0, CurSubsectionIdx);
+}
+
+void MCObjectStreamer::flushPendingLabels() {
+  // Register labels that have not yet been assigned to a Section.
+  if (!PendingLabels.empty()) {
+    MCSection *CurSection = getCurrentSectionOnly();
+    assert(CurSection);
+    for (MCSymbol* Sym : PendingLabels)
+      CurSection->addPendingLabel(Sym, CurSubsectionIdx);
+    PendingLabels.clear();
+  }
+
+  // Assign an empty data fragment to all remaining pending labels.
+  for (MCSection* Section : PendingLabelSections)
+    Section->flushPendingLabels();
 }
 
 // When fixup's offset is a forward declared label, e.g.:
@@ -120,6 +164,7 @@ void MCObjectStreamer::reset() {
   EmitEHFrame = true;
   EmitDebugFrame = false;
   PendingLabels.clear();
+  PendingLabelSections.clear();
   MCStreamer::reset();
 }
 
@@ -237,7 +282,7 @@ void MCObjectStreamer::EmitLabel(MCSymbol *Symbol, SMLoc Loc) {
     // fragment. (They will all be reassigned to a real fragment in
     // flushPendingLabels())
     Symbol->setOffset(0);
-    PendingLabels.push_back(Symbol);
+    addPendingLabel(Symbol);
   }
 }
 
@@ -257,7 +302,7 @@ void MCObjectStreamer::EmitLabelAtPos(MCSymbol *Symbol, SMLoc Loc,
     assert(isa<MCDummyFragment>(F) &&
            "F must either be an MCDataFragment or the pending MCDummyFragment");
     assert(Offset == 0);
-    PendingLabels.push_back(Symbol);
+    addPendingLabel(Symbol);
   }
 }
 
@@ -292,7 +337,6 @@ void MCObjectStreamer::ChangeSection(MCSection *Section,
 bool MCObjectStreamer::changeSectionImpl(MCSection *Section,
                                          const MCExpr *Subsection) {
   assert(Section && "Cannot switch to a null section!");
-  flushPendingLabels(nullptr);
   getContext().clearDwarfLocSeen();
 
   bool Created = getAssembler().registerSection(*Section);
@@ -303,8 +347,9 @@ bool MCObjectStreamer::changeSectionImpl(MCSection *Section,
     report_fatal_error("Cannot evaluate subsection number");
   if (IntSubsection < 0 || IntSubsection > 8192)
     report_fatal_error("Subsection number out of range");
+  CurSubsectionIdx = unsigned(IntSubsection);
   CurInsertionPoint =
-      Section->getSubsectionInsertionPoint(unsigned(IntSubsection));
+      Section->getSubsectionInsertionPoint(CurSubsectionIdx);
   return Created;
 }
 
@@ -712,7 +757,9 @@ void MCObjectStreamer::FinishImpl() {
   // Dump out the dwarf file & directory tables and line tables.
   MCDwarfLineTable::Emit(this, getAssembler().getDWARFLinetableParams());
 
+  // Update any remaining pending labels with empty data fragments.
   flushPendingLabels();
+
   resolvePendingFixups();
   getAssembler().Finish();
 }
diff --git a/llvm/lib/MC/MCSection.cpp b/llvm/lib/MC/MCSection.cpp
index 2c892ab8160..7a953054e72 100644
--- a/llvm/lib/MC/MCSection.cpp
+++ b/llvm/lib/MC/MCSection.cpp
@@ -86,6 +86,41 @@ MCSection::getSubsectionInsertionPoint(unsigned Subsection) {
   return IP;
 }
 
+void MCSection::addPendingLabel(MCSymbol* label, unsigned Subsection) {
+  PendingLabels.push_back(PendingLabel(label, Subsection));
+}
+
+void MCSection::flushPendingLabels(MCFragment *F, uint64_t FOffset,
+				   unsigned Subsection) {
+  if (PendingLabels.empty())
+    return;
+
+  // Set the fragment and fragment offset for all pending symbols in the
+  // specified Subsection, and remove those symbols from the pending list.
+  for (auto It = PendingLabels.begin(); It != PendingLabels.end(); ++It) {
+    PendingLabel& Label = *It;
+    if (Label.Subsection == Subsection) {
+      Label.Sym->setFragment(F);
+      Label.Sym->setOffset(FOffset);
+      PendingLabels.erase(It--);
+    }
+  }
+}
+
+void MCSection::flushPendingLabels() {
+  // Make sure all remaining pending labels point to data fragments, by
+  // creating new empty data fragments for each Subsection with labels pending.
+  while (!PendingLabels.empty()) {
+    PendingLabel& Label = PendingLabels[0];
+    iterator CurInsertionPoint =
+      this->getSubsectionInsertionPoint(Label.Subsection);
+    MCFragment *F = new MCDataFragment();
+    getFragmentList().insert(CurInsertionPoint, F);
+    F->setParent(this);
+    flushPendingLabels(F, 0, Label.Subsection);
+  }
+}
+
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
 LLVM_DUMP_METHOD void MCSection::dump() const {
   raw_ostream &OS = errs();
diff --git a/llvm/test/MC/MachO/pending-labels.s b/llvm/test/MC/MachO/pending-labels.s
new file mode 100644
index 00000000000..a4644357ca4
--- /dev/null
+++ b/llvm/test/MC/MachO/pending-labels.s
@@ -0,0 +1,48 @@
+// Verify relocations for temporary labels are referenced by real symbols
+// at the same address.
+//
+// RUN: llvm-mc -triple x86_64-apple-darwin -filetype=obj -o - %s | llvm-objdump -r - | FileCheck %s
+
+L1:
+	.section __TEXT,__text_cold,regular,pure_instructions
+L2:
+	.text
+L3:
+	.section __TEXT,__text_cold,regular,pure_instructions
+L4:
+_function2:
+L5:
+	nop
+L6:
+	.text
+L7:
+_function1:
+L8:
+	nop
+
+	.data
+__data:
+	.quad L1-.
+	.quad L2-.
+	.quad L3-.
+	.quad L4-.
+	.quad L5-.
+	.quad L6-.
+	.quad L7-.
+	.quad L8-.
+// CHECK: 0000000000000038 X86_64_RELOC_SUBTRACTOR _function1-__data
+// CHECK: 0000000000000038 X86_64_RELOC_UNSIGNED _function1
+// CHECK: 0000000000000030 X86_64_RELOC_SUBTRACTOR _function1-__data
+// CHECK: 0000000000000030 X86_64_RELOC_UNSIGNED _function1
+// CHECK: 0000000000000028 X86_64_RELOC_SUBTRACTOR _function2-__data
+// CHECK: 0000000000000028 X86_64_RELOC_UNSIGNED _function2
+// CHECK: 0000000000000020 X86_64_RELOC_SUBTRACTOR _function2-__data
+// CHECK: 0000000000000020 X86_64_RELOC_UNSIGNED _function2
+// CHECK: 0000000000000018 X86_64_RELOC_SUBTRACTOR _function2-__data
+// CHECK: 0000000000000018 X86_64_RELOC_UNSIGNED _function2
+// CHECK: 0000000000000010 X86_64_RELOC_SUBTRACTOR _function1-__data
+// CHECK: 0000000000000010 X86_64_RELOC_UNSIGNED _function1
+// CHECK: 0000000000000008 X86_64_RELOC_SUBTRACTOR _function2-__data
+// CHECK: 0000000000000008 X86_64_RELOC_UNSIGNED _function2
+// CHECK: 0000000000000000 X86_64_RELOC_SUBTRACTOR _function1-__data
+// CHECK: 0000000000000000 X86_64_RELOC_UNSIGNED _function1