From 7423bf78eb53d81ce0c7b3a38e39a56341ca2a89 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Thu, 27 Jun 2024 19:12:35 -0700
Subject: [MC] Ensure subsections have a MCDataFragment

Similar to 21fac2d1d060b0f9b11a746718e58d4cd1ee97e5 for sections. This
makes it feasible to cache the current fragment in MCStreamer.
---
 llvm/include/llvm/MC/MCSection.h |  4 ++--
 llvm/lib/MC/MCAssembler.cpp      |  3 +--
 llvm/lib/MC/MCObjectStreamer.cpp | 18 +++++++++++++++---
 llvm/lib/MC/MCSection.cpp        | 11 -----------
 4 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/llvm/include/llvm/MC/MCSection.h b/llvm/include/llvm/MC/MCSection.h
index 22e5db3..e8e6575 100644
--- a/llvm/include/llvm/MC/MCSection.h
+++ b/llvm/include/llvm/MC/MCSection.h
@@ -26,6 +26,7 @@ class MCAsmInfo;
 class MCAssembler;
 class MCContext;
 class MCExpr;
+class MCObjectStreamer;
 class MCSymbol;
 class raw_ostream;
 class Triple;
@@ -35,6 +36,7 @@ class Triple;
 class MCSection {
 public:
   friend MCAssembler;
+  friend MCObjectStreamer;
   static constexpr unsigned NonUniqueID = ~0U;
 
   enum SectionVariant {
@@ -208,8 +210,6 @@ public:
     CurFragList->Tail = &F;
   }
 
-  void switchSubsection(unsigned Subsection);
-
   void dump() const;
 
   virtual void printSwitchToSection(const MCAsmInfo &MAI, const Triple &T,
diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
index 205a9b2..52374017 100644
--- a/llvm/lib/MC/MCAssembler.cpp
+++ b/llvm/lib/MC/MCAssembler.cpp
@@ -870,8 +870,7 @@ void MCAssembler::layout(MCAsmLayout &Layout) {
       MCDummyFragment Dummy;
       MCFragment *Tail = &Dummy;
       for (auto &[_, List] : Sec->Subsections) {
-        if (!List.Head)
-          continue;
+        assert(List.Head);
         Tail->Next = List.Head;
         Tail = List.Tail;
       }
diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index fec1cce..2357203 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -294,9 +294,21 @@ bool MCObjectStreamer::changeSectionImpl(MCSection *Section,
   assert(Section && "Cannot switch to a null section!");
   getContext().clearDwarfLocSeen();
 
-  bool Created = getAssembler().registerSection(*Section);
-  Section->switchSubsection(Subsection);
-  return Created;
+  auto &Subsections = Section->Subsections;
+  size_t I = 0, E = Subsections.size();
+  while (I != E && Subsections[I].first < Subsection)
+    ++I;
+  // If the subsection number is not in the sorted Subsections list, create a
+  // new fragment list.
+  if (I == E || Subsections[I].first != Subsection) {
+    auto *F = getContext().allocFragment<MCDataFragment>();
+    F->setParent(Section);
+    Subsections.insert(Subsections.begin() + I,
+                       {Subsection, MCSection::FragList{F, F}});
+  }
+  Section->CurFragList = &Subsections[I].second;
+
+  return getAssembler().registerSection(*Section);
 }
 
 void MCObjectStreamer::emitAssignment(MCSymbol *Symbol, const MCExpr *Value) {
diff --git a/llvm/lib/MC/MCSection.cpp b/llvm/lib/MC/MCSection.cpp
index 1e15b68..8c2ee56 100644
--- a/llvm/lib/MC/MCSection.cpp
+++ b/llvm/lib/MC/MCSection.cpp
@@ -66,17 +66,6 @@ void MCSection::setBundleLockState(BundleLockStateType NewState) {
   ++BundleLockNestingDepth;
 }
 
-void MCSection::switchSubsection(unsigned Subsection) {
-  size_t I = 0, E = Subsections.size();
-  while (I != E && Subsections[I].first < Subsection)
-    ++I;
-  // If the subsection number is not in the sorted Subsections list, create a
-  // new fragment list.
-  if (I == E || Subsections[I].first != Subsection)
-    Subsections.insert(Subsections.begin() + I, {Subsection, FragList{}});
-  CurFragList = &Subsections[I].second;
-}
-
 StringRef MCSection::getVirtualSectionKind() const { return "virtual"; }
 
 #if !defined(NDEBUG) || defined(LLVM_ENABLE_DUMP)
-- 
cgit v1.1