[MC] Factor MCObjectStreamer::addFragmentAtoms out of MachO streamer.

This code previously existed only in MCMachOStreamer but is
useful for WebAssembly too.  See: D46335

Differential Revision: https://reviews.llvm.org/D46297

llvm-svn: 331412

1 files changed, 26 insertions, 0 deletions

diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index ebedd9bf815..b2008b3e1bb 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -60,6 +60,32 @@ void MCObjectStreamer::flushPendingLabels(MCFragment *F, uint64_t FOffset) {
   PendingLabels.clear();
 }
 
+void MCObjectStreamer::addFragmentAtoms() {
+  // First, scan the symbol table to build a lookup table from fragments to
+  // defining symbols.
+  DenseMap<const MCFragment *, const MCSymbol *> DefiningSymbolMap;
+  for (const MCSymbol &Symbol : getAssembler().symbols()) {
+    if (getAssembler().isSymbolLinkerVisible(Symbol) && Symbol.isInSection() &&
+        !Symbol.isVariable()) {
+      // An atom defining symbol should never be internal to a fragment.
+      assert(Symbol.getOffset() == 0 &&
+             "Invalid offset in atom defining symbol!");
+      DefiningSymbolMap[Symbol.getFragment()] = &Symbol;
+    }
+  }
+
+  // Set the fragment atom associations by tracking the last seen atom defining
+  // symbol.
+  for (MCSection &Sec : getAssembler()) {
+    const MCSymbol *CurrentAtom = nullptr;
+    for (MCFragment &Frag : Sec) {
+      if (const MCSymbol *Symbol = DefiningSymbolMap.lookup(&Frag))
+        CurrentAtom = Symbol;
+      Frag.setAtom(CurrentAtom);
+    }
+  }
+}
+
 // As a compile-time optimization, avoid allocating and evaluating an MCExpr
 // tree for (Hi - Lo) when Hi and Lo are offsets into the same fragment.
 static Optional<uint64_t> absoluteSymbolDiff(const MCSymbol *Hi,