[MC] Add MCObjectStreamer::addFragmentAtoms and use it with MachO and WebAssembly

[sbc100]

Summary

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

This is part of a larger WebAssembly change: D46335

[yurydelendik]

Looks good.

include/llvm/MC/MCObjectStreamer.h

protected:
   bool changeSectionImpl(MCSection *Section, const MCExpr *Subsection);
 
   /// If any labels have been emitted but not assigned fragments, ensure that
   /// they get assigned, either to F if possible or to a new data fragment.
   /// Optionally, it is also possible to provide an offset \p FOffset, which
   /// will be used as a symbol offset within the fragment.
   void flushPendingLabels(MCFragment *F, uint64_t FOffset = 0);
 
+  void addFragmentAtoms();
+
 public:
   void visitUsedSymbol(const MCSymbol &Sym) override;
 
   MCAssembler &getAssembler() { return *Assembler; }
   MCAssembler *getAssemblerPtr() override;
   /// \name MCStreamer Interface
   /// @{
 

lib/MC/MCMachOStreamer.cpp

void MCMachOStreamer::EmitInstToData(const MCInst &Inst,
   DF->getContents().append(Code.begin(), Code.end());
 }
 
 void MCMachOStreamer::FinishImpl() {
   EmitFrames(&getAssembler().getBackend());
 
   // We have to set the fragment atom associations so we can relax properly for
   // Mach-O.
-
-  // 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);
-    }
-  }
+  addFragmentAtoms();
 
   this->MCObjectStreamer::FinishImpl();
 }
 
 MCStreamer *llvm::createMachOStreamer(MCContext &Context,
                                       std::unique_ptr<MCAsmBackend> &&MAB,
                                       raw_pwrite_stream &OS,
                                       std::unique_ptr<MCCodeEmitter> &&CE,

lib/MC/MCObjectStreamer.cpp

void MCObjectStreamer::flushPendingLabels(MCFragment *F, uint64_t FOffset) {
   }
   for (MCSymbol *Sym : PendingLabels) {
     Sym->setFragment(F);
     Sym->setOffset(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,
                                              const MCSymbol *Lo) {
   assert(Hi && Lo);
   if (!Hi->getFragment() || Hi->getFragment() != Lo->getFragment() ||
       Hi->isVariable() || Lo->isVariable())
     return None;

lib/MC/MCWasmStreamer.cpp

void MCWasmStreamer::EmitInstToData(const MCInst &Inst,
   }
   DF->setHasInstructions(true);
   DF->getContents().append(Code.begin(), Code.end());
 }
 
 void MCWasmStreamer::FinishImpl() {
   EmitFrames(nullptr);
 
+  // Set fragment atoms so we can map from code fragment to defining symbol
+  addFragmentAtoms();
+
   this->MCObjectStreamer::FinishImpl();
 }
 
 MCStreamer *llvm::createWasmStreamer(MCContext &Context,
                                      std::unique_ptr<MCAsmBackend> &&MAB,
                                      raw_pwrite_stream &OS,
                                      std::unique_ptr<MCCodeEmitter> &&CE,
                                      bool RelaxAll) {