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

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.

Pull Request: https://github.com/llvm/llvm-project/pull/95341

5 files changed, 28 insertions, 6 deletions

diff --git a/llvm/include/llvm/MC/MCFragment.h b/llvm/include/llvm/MC/MCFragment.h
index 2f62bdb..55573082 100644
--- a/llvm/include/llvm/MC/MCFragment.h
+++ b/llvm/include/llvm/MC/MCFragment.h
@@ -60,9 +60,6 @@ private:
   /// 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,7 @@ public:
   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;
 
   unsigned getLayoutOrder() const { return LayoutOrder; }
   void setLayoutOrder(unsigned Value) { LayoutOrder = Value; }
diff --git a/llvm/include/llvm/MC/MCSectionMachO.h b/llvm/include/llvm/MC/MCSectionMachO.h
index fdf1773..e268562 100644
--- a/llvm/include/llvm/MC/MCSectionMachO.h
+++ b/llvm/include/llvm/MC/MCSectionMachO.h
@@ -32,6 +32,9 @@ class MCSectionMachO final : public MCSection {
   /// for example.
   unsigned Reserved2;
 
+  // The defining non-temporary symbol for each fragment.
+  SmallVector<const MCSymbol *, 0> Atoms;
+
   MCSectionMachO(StringRef Segment, StringRef Section, unsigned TAA,
                  unsigned reserved2, SectionKind K, MCSymbol *Begin);
   friend class MCContext;
@@ -74,6 +77,10 @@ public:
   bool useCodeAlign() const override;
   bool isVirtualSection() const override;
 
+  void allocAtoms();
+  const MCSymbol *getAtom(size_t I) const;
+  void setAtom(size_t I, const MCSymbol *Sym);
+
   static bool classof(const MCSection *S) {
     return S->getVariant() == SV_MachO;
   }
diff --git a/llvm/lib/MC/MCFragment.cpp b/llvm/lib/MC/MCFragment.cpp
index ffd3a50..0d8d639 100644
--- a/llvm/lib/MC/MCFragment.cpp
+++ b/llvm/lib/MC/MCFragment.cpp
@@ -17,6 +17,7 @@
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCFixup.h"
 #include "llvm/MC/MCSection.h"
+#include "llvm/MC/MCSectionMachO.h"
 #include "llvm/MC/MCSymbol.h"
 #include "llvm/MC/MCValue.h"
 #include "llvm/Support/Casting.h"
@@ -264,6 +265,10 @@ void MCFragment::destroy() {
   }
 }
 
+const MCSymbol *MCFragment::getAtom() const {
+  return cast<MCSectionMachO>(Parent)->getAtom(LayoutOrder);
+}
+
 // Debugging methods
 
 namespace llvm {
diff --git a/llvm/lib/MC/MCMachOStreamer.cpp b/llvm/lib/MC/MCMachOStreamer.cpp
index 10f9988..466aa63 100644
--- a/llvm/lib/MC/MCMachOStreamer.cpp
+++ b/llvm/lib/MC/MCMachOStreamer.cpp
@@ -519,11 +519,13 @@ void MCMachOStreamer::finishImpl() {
   // Set the fragment atom associations by tracking the last seen atom defining
   // symbol.
   for (MCSection &Sec : getAssembler()) {
+    cast<MCSectionMachO>(Sec).allocAtoms();
     const MCSymbol *CurrentAtom = nullptr;
+    size_t I = 0;
     for (MCFragment &Frag : Sec) {
       if (const MCSymbol *Symbol = DefiningSymbolMap.lookup(&Frag))
         CurrentAtom = Symbol;
-      Frag.setAtom(CurrentAtom);
+      cast<MCSectionMachO>(Sec).setAtom(I++, CurrentAtom);
     }
   }
 
diff --git a/llvm/lib/MC/MCSectionMachO.cpp b/llvm/lib/MC/MCSectionMachO.cpp
index f7eedac..53b7666 100644
--- a/llvm/lib/MC/MCSectionMachO.cpp
+++ b/llvm/lib/MC/MCSectionMachO.cpp
@@ -291,3 +291,15 @@ Error MCSectionMachO::ParseSectionSpecifier(StringRef Spec,       // In.
 
   return Error::success();
 }
+
+void MCSectionMachO::allocAtoms() {
+  auto *L = curFragList();
+  if (L->Tail)
+    Atoms.resize(L->Tail->getLayoutOrder() + 1);
+}
+
+const MCSymbol *MCSectionMachO::getAtom(size_t I) const {
+  return I < Atoms.size() ? Atoms[I] : nullptr;
+}
+
+void MCSectionMachO::setAtom(size_t I, const MCSymbol *Sym) { Atoms[I] = Sym; }