[MC] Move MCFragment::Atom to MCSectionMachO::Atoms

[MaskRay]

Mach-O's .subsections_via_symbols mechanism associates a fragment with
an atom (a non-temporary defined symbol). The current approach
(MCFragment::Atom) wastes space for other object file formats.

After #95077, MCFragment::LayoutOrder is only used by
AttemptToFoldSymbolOffsetDifference. While it could be removed, we
might explore future uses for LayoutOrder.

@aengelke suggests one use case: move Atom into MCSection. This works
because Mach-O doesn't support .subsection, and LayoutOrder, as the
index into the fragment list, is unchanged.

This patch moves MCFragment::Atom to MCSectionMachO::Atoms. getAtom
may be called at parse time before Atoms is initialized, so a bound
checking is needed to keep the hack working.

[llvmbot]

@llvm/pr-subscribers-mc

Author: Fangrui Song (MaskRay)

Changes

Mach-O's .subsections_via_symbols mechanism associates a fragment with
an atom (a non-temporary defined symbol). The current approach
(MCFragment::Atom) wastes space for other object file formats.

After #95077, MCFragment::LayoutOrder is only used by
AttemptToFoldSymbolOffsetDifference. While it could be removed, we
might explore future uses for LayoutOrder.

@aengelke suggests one use case: move Atom into MCSection. This works
because Mach-O doesn't support .subsection, and LayoutOrder, as the
index into the fragment list, is unchanged.

---

Full diff: https://github.com/llvm/llvm-project/pull/95341.diff

3 Files Affected:

• (modified) llvm/include/llvm/MC/MCFragment.h (+2-5)
• (modified) llvm/include/llvm/MC/MCSection.h (+5)
• (modified) llvm/lib/MC/MCFragment.cpp (+12)

diff --git a/llvm/include/llvm/MC/MCFragment.h b/llvm/include/llvm/MC/MCFragment.h
index ccfe6203514b0..1996de810c113 100644
--- a/llvm/include/llvm/MC/MCFragment.h
+++ b/llvm/include/llvm/MC/MCFragment.h
@@ -60,9 +60,6 @@ class MCFragment {
   /// The data for the section this fragment is in.
   MCSection *Parent;
 
-  /// The atom this fragment is in, as represented by its defining symbol.
-  const MCSymbol *Atom = nullptr;
-
   /// The offset of this fragment in its section.
   uint64_t Offset = 0;
 
@@ -96,8 +93,8 @@ class MCFragment {
   MCSection *getParent() const { return Parent; }
   void setParent(MCSection *Value) { Parent = Value; }
 
-  const MCSymbol *getAtom() const { return Atom; }
-  void setAtom(const MCSymbol *Value) { Atom = Value; }
+  const MCSymbol *getAtom() const;
+  void setAtom(const MCSymbol *Sym);
 
   unsigned getLayoutOrder() const { return LayoutOrder; }
   void setLayoutOrder(unsigned Value) { LayoutOrder = Value; }
diff --git a/llvm/include/llvm/MC/MCSection.h b/llvm/include/llvm/MC/MCSection.h
index e5455292d5c62..07632d369dab6 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 MCFragment;
 class MCSymbol;
 class raw_ostream;
 class Triple;
@@ -35,6 +36,7 @@ class Triple;
 class MCSection {
 public:
   friend MCAssembler;
+  friend MCFragment;
   static constexpr unsigned NonUniqueID = ~0U;
 
   enum SectionVariant {
@@ -111,6 +113,9 @@ class MCSection {
   // subsections.
   SmallVector<std::pair<unsigned, FragList>, 1> Subsections;
 
+  // Mach-O only: the defining non-temporary symbol for each fragment.
+  SmallVector<const MCSymbol *, 0> Atoms;
+
   /// State for tracking labels that don't yet have Fragments
   struct PendingLabel {
     MCSymbol* Sym;
diff --git a/llvm/lib/MC/MCFragment.cpp b/llvm/lib/MC/MCFragment.cpp
index e911fa21650f4..48b3976c9c614 100644
--- a/llvm/lib/MC/MCFragment.cpp
+++ b/llvm/lib/MC/MCFragment.cpp
@@ -263,6 +263,18 @@ void MCFragment::destroy() {
   }
 }
 
+const MCSymbol *MCFragment::getAtom() const {
+  if (LayoutOrder >= Parent->Atoms.size())
+    return nullptr;
+  return Parent->Atoms[LayoutOrder];
+}
+
+void MCFragment::setAtom(const MCSymbol *Sym) {
+  if (LayoutOrder >= Parent->Atoms.size())
+    Parent->Atoms.resize(LayoutOrder + 1);
+  Parent->Atoms[LayoutOrder] = Sym;
+}
+
 // Debugging methods
 
 namespace llvm {

[aengelke]

LGTM. Minor thing: We could resize the vector once in MCMachOStreamer::finishImpl with LayoutOrder+1 of the last fragment. This would avoid a resize call for every fragment (=> log n actual resizes) and getAtom/setAtom would not need index checks.