2 files changed, 94 insertions, 13 deletions

diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index 6c85f2296cf..fc7fe4cb76c 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);
@@ -712,7 +756,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();