MC: Change MCAssembler::writeSectionData and writeFragmentPadding to take a raw_ostream. NFCI.

[pcc]

Summary

Also clean up a couple of hacks where we were writing the section
contents to another stream by setting the object writer's stream,
writing and setting it back.

Part of PR37466.

Depends on D47035

[pcc]

• Fix one more user of writeNopData

[echristo]

LGTM.

Thanks!

llvm/include/llvm/MC/MCAssembler.h

public:
   const MCSymbol *getAtom(const MCSymbol &S) const;
 
   /// Check whether a particular symbol is visible to the linker and is required
   /// in the symbol table, or whether it can be discarded by the assembler. This
   /// also effects whether the assembler treats the label as potentially
   /// defining a separate atom.
   bool isSymbolLinkerVisible(const MCSymbol &SD) const;
 
-  /// Emit the section contents using the given object writer.
-  void writeSectionData(const MCSection *Section,
+  /// Emit the section contents to \p OS.
+  void writeSectionData(raw_ostream &OS, const MCSection *Section,
                         const MCAsmLayout &Layout) const;
 
   /// Check whether a given symbol has been flagged with .thumb_func.
   bool isThumbFunc(const MCSymbol *Func) const;
 
   /// Flag a function symbol as the target of a .thumb_func directive.
   void setIsThumbFunc(const MCSymbol *Func) { ThumbFuncs.insert(Func); }
 
public:
 
   ArrayRef<std::string> getFileNames() { return FileNames; }
 
   void addFileName(StringRef FileName) {
     if (!is_contained(FileNames, FileName))
       FileNames.push_back(FileName);
   }
 
-  /// Write the necessary bundle padding to the given object writer.
+  /// Write the necessary bundle padding to \p OS.
   /// Expects a fragment \p F containing instructions and its size \p FSize.
-  void writeFragmentPadding(const MCFragment &F, uint64_t FSize,
-                            MCObjectWriter *OW) const;
+  void writeFragmentPadding(raw_ostream &OS, const MCFragment &F,
+                            uint64_t FSize) const;
 
   /// @}
 
   void dump() const;
 };
 
 /// Compute the amount of padding required before the fragment \p F to
 /// obey bundling restrictions, where \p FOffset is the fragment's offset in
 /// its section and \p FSize is the fragment's size.
 uint64_t computeBundlePadding(const MCAssembler &Assembler, const MCFragment *F,
                               uint64_t FOffset, uint64_t FSize);
 
 } // end namespace llvm
 
 #endif // LLVM_MC_MCASSEMBLER_H

llvm/lib/MC/ELFObjectWriter.cpp

void ELFObjectWriter::writeSectionData(const MCAssembler &Asm, MCSection &Sec,
 
   // Compressing debug_frame requires handling alignment fragments which is
   // more work (possibly generalizing MCAssembler.cpp:writeFragment to allow
   // for writing to arbitrary buffers) for little benefit.
   bool CompressionEnabled =
       MAI->compressDebugSections() != DebugCompressionType::None;
   if (!CompressionEnabled || !SectionName.startswith(".debug_") ||
       SectionName == ".debug_frame") {
-    Asm.writeSectionData(&Section, Layout);
+    Asm.writeSectionData(getStream(), &Section, Layout);
     return;
   }
 
   assert((MAI->compressDebugSections() == DebugCompressionType::Z ||
           MAI->compressDebugSections() == DebugCompressionType::GNU) &&
          "expected zlib or zlib-gnu style compression");
 
   SmallVector<char, 128> UncompressedData;
   raw_svector_ostream VecOS(UncompressedData);
-  raw_pwrite_stream &OldStream = getStream();
-  setStream(VecOS);
-  Asm.writeSectionData(&Section, Layout);
-  setStream(OldStream);
+  Asm.writeSectionData(VecOS, &Section, Layout);
 
   SmallVector<char, 128> CompressedContents;
   if (Error E = zlib::compress(
           StringRef(UncompressedData.data(), UncompressedData.size()),
           CompressedContents)) {
     consumeError(std::move(E));
     getStream() << UncompressedData;
     return;

llvm/lib/MC/MCAssembler.cpp

void MCAssembler::registerSymbol(const MCSymbol &Symbol, bool *Created) {
   if (Created)
     *Created = New;
   if (New) {
     Symbol.setIsRegistered(true);
     Symbols.push_back(&Symbol);
   }
 }
 
-void MCAssembler::writeFragmentPadding(const MCFragment &F, uint64_t FSize,
-                                       MCObjectWriter *OW) const {
+void MCAssembler::writeFragmentPadding(raw_ostream &OS, const MCFragment &F,
+                                       uint64_t FSize) const {
   assert(getBackendPtr() && "Expected assembler backend");
   // Should NOP padding be written out before this fragment?
   unsigned BundlePadding = F.getBundlePadding();
   if (BundlePadding > 0) {
     assert(isBundlingEnabled() &&
            "Writing bundle padding with disabled bundling");
     assert(F.hasInstructions() &&
            "Writing bundle padding for a fragment without instructions");
 
     unsigned TotalLength = BundlePadding + static_cast<unsigned>(FSize);
     if (F.alignToBundleEnd() && TotalLength > getBundleAlignSize()) {
       // If the padding itself crosses a bundle boundary, it must be emitted
       // in 2 pieces, since even nop instructions must not cross boundaries.
       //             v--------------v   <- BundleAlignSize
       //        v---------v             <- BundlePadding
       // ----------------------------
       // | Prev |####|####|    F    |
       // ----------------------------
       //        ^-------------------^   <- TotalLength
       unsigned DistanceToBoundary = TotalLength - getBundleAlignSize();
-      if (!getBackend().writeNopData(OW->getStream(), DistanceToBoundary))
+      if (!getBackend().writeNopData(OS, DistanceToBoundary))
           report_fatal_error("unable to write NOP sequence of " +
                              Twine(DistanceToBoundary) + " bytes");
       BundlePadding -= DistanceToBoundary;
     }
-    if (!getBackend().writeNopData(OW->getStream(), BundlePadding))
+    if (!getBackend().writeNopData(OS, BundlePadding))
       report_fatal_error("unable to write NOP sequence of " +
                          Twine(BundlePadding) + " bytes");
   }
 }
 
 /// Write the fragment \p F to the output file.
-static void writeFragment(const MCAssembler &Asm, const MCAsmLayout &Layout,
-                          const MCFragment &F) {
+static void writeFragment(raw_ostream &OS, const MCAssembler &Asm,
+                          const MCAsmLayout &Layout, const MCFragment &F) {
   MCObjectWriter *OW = Asm.getWriterPtr();
   assert(OW && "Need ObjectWriter to write fragment");
 
   // FIXME: Embed in fragments instead?
   uint64_t FragmentSize = Asm.computeFragmentSize(Layout, F);
 
-  Asm.writeFragmentPadding(F, FragmentSize, OW);
+  support::endianness Endian = Asm.getBackend().Endian;
+
+  Asm.writeFragmentPadding(OS, F, FragmentSize);
 
   // This variable (and its dummy usage) is to participate in the assert at
   // the end of the function.
-  uint64_t Start = OW->getStream().tell();
+  uint64_t Start = OS.tell();
   (void) Start;
 
   ++stats::EmittedFragments;
 
   switch (F.getKind()) {
   case MCFragment::FT_Align: {
     ++stats::EmittedAlignFragments;
     const MCAlignFragment &AF = cast<MCAlignFragment>(F);
if (Count * AF.getValueSize() != FragmentSize)
                         "' is not a divisor of padding size '" +
                         Twine(FragmentSize) + "'");
 
     // See if we are aligning with nops, and if so do that first to try to fill
     // the Count bytes.  Then if that did not fill any bytes or there are any
     // bytes left to fill use the Value and ValueSize to fill the rest.
     // If we are aligning with nops, ask that target to emit the right data.
     if (AF.hasEmitNops()) {
-      if (!Asm.getBackend().writeNopData(OW->getStream(), Count))
+      if (!Asm.getBackend().writeNopData(OS, Count))
         report_fatal_error("unable to write nop sequence of " +
                           Twine(Count) + " bytes");
       break;
     }
 
     // Otherwise, write out in multiples of the value size.
     for (uint64_t i = 0; i != Count; ++i) {
       switch (AF.getValueSize()) {
       default: llvm_unreachable("Invalid size!");
-      case 1: OW->write8 (uint8_t (AF.getValue())); break;
-      case 2: OW->write16(uint16_t(AF.getValue())); break;
-      case 4: OW->write32(uint32_t(AF.getValue())); break;
-      case 8: OW->write64(uint64_t(AF.getValue())); break;
+      case 1: OS << char(AF.getValue()); break;
+      case 2:
+        support::endian::write<uint16_t>(OS, AF.getValue(), Endian);
+        break;
+      case 4:
+        support::endian::write<uint32_t>(OS, AF.getValue(), Endian);
+        break;
+      case 8:
+        support::endian::write<uint64_t>(OS, AF.getValue(), Endian);
+        break;
       }
     }
     break;
   }
 
   case MCFragment::FT_Data: 
     ++stats::EmittedDataFragments;
-    OW->writeBytes(cast<MCDataFragment>(F).getContents());
+    OS << cast<MCDataFragment>(F).getContents();
     break;
 
   case MCFragment::FT_Relaxable:
     ++stats::EmittedRelaxableFragments;
-    OW->writeBytes(cast<MCRelaxableFragment>(F).getContents());
+    OS << cast<MCRelaxableFragment>(F).getContents();
     break;
 
   case MCFragment::FT_CompactEncodedInst:
     ++stats::EmittedCompactEncodedInstFragments;
-    OW->writeBytes(cast<MCCompactEncodedInstFragment>(F).getContents());
+    OS << cast<MCCompactEncodedInstFragment>(F).getContents();
     break;
 
   case MCFragment::FT_Fill: {
     ++stats::EmittedFillFragments;
     const MCFillFragment &FF = cast<MCFillFragment>(F);
     uint8_t V = FF.getValue();
     const unsigned MaxChunkSize = 16;
     char Data[MaxChunkSize];
     memcpy(Data, &V, 1);
     for (unsigned I = 1; I < MaxChunkSize; ++I)
       Data[I] = Data[0];
 
     uint64_t Size = FragmentSize;
     for (unsigned ChunkSize = MaxChunkSize; ChunkSize; ChunkSize /= 2) {
       StringRef Ref(Data, ChunkSize);
       for (uint64_t I = 0, E = Size / ChunkSize; I != E; ++I)
-        OW->writeBytes(Ref);
+        OS << Ref;
       Size = Size % ChunkSize;
     }
     break;
   }
 
   case MCFragment::FT_LEB: {
     const MCLEBFragment &LF = cast<MCLEBFragment>(F);
-    OW->writeBytes(LF.getContents());
+    OS << LF.getContents();
     break;
   }
 
   case MCFragment::FT_Padding: {
     if (!Asm.getBackend().writeNopData(OW->getStream(), FragmentSize))
       report_fatal_error("unable to write nop sequence of " +
                          Twine(FragmentSize) + " bytes");
     break;
   }
 
   case MCFragment::FT_SymbolId: {
     const MCSymbolIdFragment &SF = cast<MCSymbolIdFragment>(F);
-    OW->write32(SF.getSymbol()->getIndex());
+    support::endian::write<uint32_t>(OS, SF.getSymbol()->getIndex(), Endian);
     break;
   }
 
   case MCFragment::FT_Org: {
     ++stats::EmittedOrgFragments;
     const MCOrgFragment &OF = cast<MCOrgFragment>(F);
 
     for (uint64_t i = 0, e = FragmentSize; i != e; ++i)
-      OW->write8(uint8_t(OF.getValue()));
+      OS << char(OF.getValue());
 
     break;
   }
 
   case MCFragment::FT_Dwarf: {
     const MCDwarfLineAddrFragment &OF = cast<MCDwarfLineAddrFragment>(F);
-    OW->writeBytes(OF.getContents());
+    OS << OF.getContents();
     break;
   }
   case MCFragment::FT_DwarfFrame: {
     const MCDwarfCallFrameFragment &CF = cast<MCDwarfCallFrameFragment>(F);
-    OW->writeBytes(CF.getContents());
+    OS << CF.getContents();
     break;
   }
   case MCFragment::FT_CVInlineLines: {
     const auto &OF = cast<MCCVInlineLineTableFragment>(F);
-    OW->writeBytes(OF.getContents());
+    OS << OF.getContents();
     break;
   }
   case MCFragment::FT_CVDefRange: {
     const auto &DRF = cast<MCCVDefRangeFragment>(F);
-    OW->writeBytes(DRF.getContents());
+    OS << DRF.getContents();
     break;
   }
   case MCFragment::FT_Dummy:
     llvm_unreachable("Should not have been added");
   }
 
-  assert(OW->getStream().tell() - Start == FragmentSize &&
+  assert(OS.tell() - Start == FragmentSize &&
          "The stream should advance by fragment size");
 }
 
-void MCAssembler::writeSectionData(const MCSection *Sec,
+void MCAssembler::writeSectionData(raw_ostream &OS, const MCSection *Sec,
                                    const MCAsmLayout &Layout) const {
   assert(getBackendPtr() && "Expected assembler backend");
 
   // Ignore virtual sections.
   if (Sec->isVirtualSection()) {
     assert(Layout.getSectionFileSize(Sec) == 0 && "Invalid size for section!");
 
     // Check that contents are only things legal inside a virtual section.
for (const MCFragment &F : *Sec) {
                "Invalid fill in virtual section!");
         break;
       }
     }
 
     return;
   }
 
-  uint64_t Start = getWriter().getStream().tell();
+  uint64_t Start = OS.tell();
   (void)Start;
 
   for (const MCFragment &F : *Sec)
-    writeFragment(*this, Layout, F);
+    writeFragment(OS, *this, Layout, F);
 
-  assert(getWriter().getStream().tell() - Start ==
-         Layout.getSectionAddressSize(Sec));
+  assert(OS.tell() - Start == Layout.getSectionAddressSize(Sec));
 }
 
 std::tuple<MCValue, uint64_t, bool>
 MCAssembler::handleFixup(const MCAsmLayout &Layout, MCFragment &F,
                          const MCFixup &Fixup) {
   // Evaluate the fixup.
   MCValue Target;
   uint64_t FixedValue;

llvm/lib/MC/MCELFStreamer.cpp

uint64_t RequiredBundlePadding = computeBundlePadding(
         Assembler, EF, DF->getContents().size(), FSize);
 
     if (RequiredBundlePadding > UINT8_MAX)
       report_fatal_error("Padding cannot exceed 255 bytes");
 
     if (RequiredBundlePadding > 0) {
       SmallString<256> Code;
       raw_svector_ostream VecOS(Code);
-      {
-        auto OW = Assembler.getBackend().createObjectWriter(VecOS);
-
-        EF->setBundlePadding(static_cast<uint8_t>(RequiredBundlePadding));
-
-        Assembler.writeFragmentPadding(*EF, FSize, OW.get());
-      }
+      Assembler.writeFragmentPadding(VecOS, *EF, FSize);
 
       DF->getContents().append(Code.begin(), Code.end());
     }
   }
 
   flushPendingLabels(DF, DF->getContents().size());
 
   for (unsigned i = 0, e = EF->getFixups().size(); i != e; ++i) {

llvm/lib/MC/MachObjectWriter.cpp

void MachObjectWriter::writeObject(MCAssembler &Asm,
   }
 
   // Write the linker options load commands.
   for (const auto &Option : Asm.getLinkerOptions())
     writeLinkerOptionsLoadCommand(Option);
 
   // Write the actual section data.
   for (const MCSection &Sec : Asm) {
-    Asm.writeSectionData(&Sec, Layout);
+    Asm.writeSectionData(getStream(), &Sec, Layout);
 
     uint64_t Pad = getPaddingSize(&Sec, Layout);
     WriteZeros(Pad);
   }
 
   // Write the extra padding.
   WriteZeros(SectionDataPadding);
 

llvm/lib/MC/WasmObjectWriter.cpp

for (const WasmFunction &Func : Functions) {
     auto &FuncSection = static_cast<MCSectionWasm &>(Func.Sym->getSection());
 
     int64_t Size = 0;
     if (!Func.Sym->getSize()->evaluateAsAbsolute(Size, Layout))
       report_fatal_error(".size expression must be evaluatable");
 
     encodeULEB128(Size, getStream());
     FuncSection.setSectionOffset(getStream().tell() - Section.ContentsOffset);
-    Asm.writeSectionData(&FuncSection, Layout);
+    Asm.writeSectionData(getStream(), &FuncSection, Layout);
   }
 
   // Apply fixups.
   applyRelocations(CodeRelocations, Section.ContentsOffset);
 
   endSection(Section);
 }
 

 void WasmObjectWriter::writeCustomSections(const MCAssembler &Asm,
                                            const MCAsmLayout &Layout) {
   for (auto &CustomSection : CustomSections) {
     SectionBookkeeping Section;
     auto *Sec = CustomSection.Section;
     startCustomSection(Section, CustomSection.Name);
 
     Sec->setSectionOffset(getStream().tell() - Section.ContentsOffset);
-    Asm.writeSectionData(Sec, Layout);
+    Asm.writeSectionData(getStream(), Sec, Layout);
 
     CustomSection.OutputContentsOffset = Section.ContentsOffset;
     CustomSection.OutputIndex = Section.Index;
 
     endSection(Section);
 
     // Apply fixups.
     auto &Relocations = CustomSectionsRelocations[CustomSection.Section];

llvm/lib/MC/WinCOFFObjectWriter.cpp

 // because it needs to compute a CRC.
 uint32_t WinCOFFObjectWriter::writeSectionContents(MCAssembler &Asm,
                                                    const MCAsmLayout &Layout,
                                                    const MCSection &MCSec) {
   // Save the contents of the section to a temporary buffer, we need this
   // to CRC the data before we dump it into the object file.
   SmallVector<char, 128> Buf;
   raw_svector_ostream VecOS(Buf);
-  raw_pwrite_stream &OldStream = getStream();
-
-  // Redirect the output stream to our buffer and fill our buffer with
-  // the section data.
-  setStream(VecOS);
-  Asm.writeSectionData(&MCSec, Layout);
-
-  // Reset the stream back to what it was before.
-  setStream(OldStream);
+  Asm.writeSectionData(VecOS, &MCSec, Layout);
 
   // Write the section contents to the object file.
   getStream() << Buf;
 
   // Calculate our CRC with an initial value of '0', this is not how
   // JamCRC is specified but it aligns with the expected output.
   JamCRC JC(/*Init=*/0);
   JC.update(Buf);

llvm/tools/dsymutil/MachOUtils.cpp

bool generateDsymCompanion(const DebugMap &DM, MCStreamer &MS,
 
   // Emit the Dwarf sections contents.
   for (const MCSection &Sec : MCAsm) {
     if (Sec.begin() == Sec.end())
       continue;
 
     uint64_t Pos = OutFile.tell();
     Writer.WriteZeros(alignTo(Pos, Sec.getAlignment()) - Pos);
-    MCAsm.writeSectionData(&Sec, Layout);
+    MCAsm.writeSectionData(OutFile, &Sec, Layout);
   }
 
   return true;
 }
 } // namespace MachOUtils
 } // namespace dsymutil
 } // namespace llvm