From 21fac2d1d060b0f9b11a746718e58d4cd1ee97e5 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sun, 23 Jun 2024 10:08:51 -0700
Subject: [MC] Ensure all new sections have a MCDataFragment

MCAssembler::layout ensures that every section has at least one
fragment, which simplifies MCAsmLayout::getSectionAddressSize (see
e73353c7201a3080851d99a16f5fe2c17f7697c6 from 2010). It's better to
ensure the condition is satisfied at create time (COFF, GOFF, Mach-O) to
simplify more fragment processing.
---
 llvm/lib/MC/MCAssembler.cpp        | 12 ++----------
 llvm/lib/MC/MCContext.cpp          | 10 ++++++----
 llvm/lib/MC/MCMachOStreamer.cpp    |  7 +++----
 llvm/tools/dsymutil/MachOUtils.cpp |  3 ---
 4 files changed, 11 insertions(+), 21 deletions(-)

diff --git a/llvm/lib/MC/MCAssembler.cpp b/llvm/lib/MC/MCAssembler.cpp
index fe0419b..f8b7e4b 100644
--- a/llvm/lib/MC/MCAssembler.cpp
+++ b/llvm/lib/MC/MCAssembler.cpp
@@ -126,6 +126,7 @@ void MCAssembler::reset() {
 bool MCAssembler::registerSection(MCSection &Section) {
   if (Section.isRegistered())
     return false;
+  assert(Section.curFragList()->Head && "allocInitialFragment not called");
   Sections.push_back(&Section);
   Section.setIsRegistered(true);
   return true;
@@ -856,17 +857,8 @@ void MCAssembler::layout(MCAsmLayout &Layout) {
 
   // Create dummy fragments and assign section ordinals.
   unsigned SectionIndex = 0;
-  for (MCSection &Sec : *this) {
-    // Create dummy fragments to eliminate any empty sections, this simplifies
-    // layout.
-    if (Sec.empty()) {
-      auto *F = getContext().allocFragment<MCDataFragment>();
-      F->setParent(&Sec);
-      Sec.addFragment(*F);
-    }
-
+  for (MCSection &Sec : *this)
     Sec.setOrdinal(SectionIndex++);
-  }
 
   // Assign layout order indices to sections and fragments.
   for (unsigned i = 0, e = Layout.getSectionOrder().size(); i != e; ++i) {
diff --git a/llvm/lib/MC/MCContext.cpp b/llvm/lib/MC/MCContext.cpp
index 5d1261a..36603d0 100644
--- a/llvm/lib/MC/MCContext.cpp
+++ b/llvm/lib/MC/MCContext.cpp
@@ -483,10 +483,12 @@ MCSectionMachO *MCContext::getMachOSection(StringRef Segment, StringRef Section,
 
   // Otherwise, return a new section.
   StringRef Name = R.first->first();
-  R.first->second = new (MachOAllocator.Allocate())
+  auto *Ret = new (MachOAllocator.Allocate())
       MCSectionMachO(Segment, Name.substr(Name.size() - Section.size()),
                      TypeAndAttributes, Reserved2, Kind, Begin);
-  return R.first->second;
+  R.first->second = Ret;
+  allocInitialFragment(*Ret);
+  return Ret;
 }
 
 MCSectionELF *MCContext::createELFSectionImpl(StringRef Section, unsigned Type,
@@ -672,7 +674,7 @@ MCSectionGOFF *MCContext::getGOFFSection(StringRef Section, SectionKind Kind,
   MCSectionGOFF *GOFFSection = new (GOFFAllocator.Allocate())
       MCSectionGOFF(CachedName, Kind, Parent, Subsection);
   Iter->second = GOFFSection;
-
+  allocInitialFragment(*GOFFSection);
   return GOFFSection;
 }
 
@@ -701,8 +703,8 @@ MCSectionCOFF *MCContext::getCOFFSection(StringRef Section,
   StringRef CachedName = Iter->first.SectionName;
   MCSectionCOFF *Result = new (COFFAllocator.Allocate()) MCSectionCOFF(
       CachedName, Characteristics, COMDATSymbol, Selection, Begin);
-
   Iter->second = Result;
+  allocInitialFragment(*Result);
   return Result;
 }
 
diff --git a/llvm/lib/MC/MCMachOStreamer.cpp b/llvm/lib/MC/MCMachOStreamer.cpp
index c0a36ae..18bfd08 100644
--- a/llvm/lib/MC/MCMachOStreamer.cpp
+++ b/llvm/lib/MC/MCMachOStreamer.cpp
@@ -554,13 +554,12 @@ void MCMachOStreamer::finalizeCGProfile() {
   MCSection *CGProfileSection = Asm.getContext().getMachOSection(
       "__LLVM", "__cg_profile", 0, SectionKind::getMetadata());
   Asm.registerSection(*CGProfileSection);
-  auto *Frag = getContext().allocFragment<MCDataFragment>();
-  Frag->setParent(CGProfileSection);
-  CGProfileSection->addFragment(*Frag);
   // For each entry, reserve space for 2 32-bit indices and a 64-bit count.
   size_t SectionBytes =
       Asm.CGProfile.size() * (2 * sizeof(uint32_t) + sizeof(uint64_t));
-  Frag->getContents().resize(SectionBytes);
+  cast<MCDataFragment>(*CGProfileSection->begin())
+      .getContents()
+      .resize(SectionBytes);
 }
 
 MCStreamer *llvm::createMachOStreamer(MCContext &Context,
diff --git a/llvm/tools/dsymutil/MachOUtils.cpp b/llvm/tools/dsymutil/MachOUtils.cpp
index b08ae4c..e2ec8ca 100644
--- a/llvm/tools/dsymutil/MachOUtils.cpp
+++ b/llvm/tools/dsymutil/MachOUtils.cpp
@@ -629,9 +629,6 @@ bool generateDsymCompanion(
 
   // Emit the Dwarf sections contents.
   for (const MCSection &Sec : MCAsm) {
-    if (Sec.empty())
-      continue;
-
     uint64_t Pos = OutFile.tell();
     OutFile.write_zeros(alignTo(Pos, Sec.getAlign()) - Pos);
     MCAsm.writeSectionData(OutFile, &Sec, Layout);
-- 
cgit v1.1