1 files changed, 59 insertions, 13 deletions

diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index 6c85f22..fc7fe4c 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();
 }