[MC,COFF] Register .llvm.call-graph-profile in finalizeImpl

All sections should have been created before MCAssembler::layout so that
every section has an ordinal. Registering the section in
WinCOFFObjectWriter::executePostLayoutBinding is a hack.

Author: MaskRay

llvm/include/llvm/MC/MCWinCOFFStreamer.h

@@ -77,7 +77,6 @@ class MCWinCOFFStreamer : public MCObjectStreamer {
   void emitInstToData(const MCInst &Inst, const MCSubtargetInfo &STI) override;
 
   void finalizeCGProfileEntry(const MCSymbolRefExpr *&S);
-  void finalizeCGProfile();
 
 private:
   void Error(const Twine &Msg) const;

llvm/lib/MC/MCWinCOFFStreamer.cpp

@@ -25,7 +25,7 @@
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectStreamer.h"
 #include "llvm/MC/MCObjectWriter.h"
-#include "llvm/MC/MCSection.h"
+#include "llvm/MC/MCSectionCOFF.h"
 #include "llvm/MC/MCSymbolCOFF.h"
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
@@ -353,15 +353,21 @@ void MCWinCOFFStreamer::finalizeCGProfileEntry(const MCSymbolRefExpr *&SRE) {
     cast<MCSymbolCOFF>(S)->setExternal(true);
 }
 
-void MCWinCOFFStreamer::finalizeCGProfile() {
-  for (MCAssembler::CGProfileEntry &E : getAssembler().CGProfile) {
-    finalizeCGProfileEntry(E.From);
-    finalizeCGProfileEntry(E.To);
-  }
-}
-
 void MCWinCOFFStreamer::finishImpl() {
-  finalizeCGProfile();
+  MCAssembler &Asm = getAssembler();
+  if (Asm.getWriter().getEmitAddrsigSection()) {
+    // Register the section.
+    switchSection(Asm.getContext().getCOFFSection(".llvm_addrsig",
+                                                  COFF::IMAGE_SCN_LNK_REMOVE));
+  }
+  if (!Asm.CGProfile.empty()) {
+    for (MCAssembler::CGProfileEntry &E : Asm.CGProfile) {
+      finalizeCGProfileEntry(E.From);
+      finalizeCGProfileEntry(E.To);
+    }
+    switchSection(Asm.getContext().getCOFFSection(".llvm.call-graph-profile",
+                                                  COFF::IMAGE_SCN_LNK_REMOVE));
+  }
 
   MCObjectStreamer::finishImpl();
 }

llvm/lib/MC/WinCOFFObjectWriter.cpp

@@ -151,9 +151,6 @@ class WinCOFFWriter {
   bool UseOffsetLabels = false;
 
 public:
-  MCSectionCOFF *AddrsigSection = nullptr;
-  MCSectionCOFF *CGProfileSection = nullptr;
-
   enum DwoMode {
     AllSections,
     NonDwoOnly,
@@ -1096,10 +1093,10 @@ uint64_t WinCOFFWriter::writeObject(MCAssembler &Asm,
   }
 
   // Create the contents of the .llvm_addrsig section.
-  if (Mode != DwoOnly && OWriter.EmitAddrsigSection) {
-    auto *Frag = Asm.getContext().allocFragment<MCDataFragment>();
-    Frag->setParent(AddrsigSection);
-    AddrsigSection->addFragment(*Frag);
+  if (Mode != DwoOnly && OWriter.getEmitAddrsigSection()) {
+    auto *Sec = Asm.getContext().getCOFFSection(
+        ".llvm_addrsig", COFF::IMAGE_SCN_LNK_REMOVE);
+    auto *Frag = cast<MCDataFragment>(Sec->curFragList()->Head);
     raw_svector_ostream OS(Frag->getContents());
     for (const MCSymbol *S : OWriter.AddrsigSyms) {
       if (!S->isRegistered())
@@ -1118,10 +1115,10 @@ uint64_t WinCOFFWriter::writeObject(MCAssembler &Asm,
   }
 
   // Create the contents of the .llvm.call-graph-profile section.
-  if (Mode != DwoOnly && CGProfileSection) {
-    auto *Frag = Asm.getContext().allocFragment<MCDataFragment>();
-    Frag->setParent(CGProfileSection);
-    CGProfileSection->addFragment(*Frag);
+  if (Mode != DwoOnly && !Asm.CGProfile.empty()) {
+    auto *Sec = Asm.getContext().getCOFFSection(
+        ".llvm.call-graph-profile", COFF::IMAGE_SCN_LNK_REMOVE);
+    auto *Frag = cast<MCDataFragment>(Sec->curFragList()->Head);
     raw_svector_ostream OS(Frag->getContents());
     for (const MCAssembler::CGProfileEntry &CGPE : Asm.CGProfile) {
       uint32_t FromIndex = CGPE.From->getSymbol().getIndex();
@@ -1209,18 +1206,6 @@ bool WinCOFFObjectWriter::isSymbolRefDifferenceFullyResolvedImpl(
 
 void WinCOFFObjectWriter::executePostLayoutBinding(MCAssembler &Asm,
                                                    const MCAsmLayout &Layout) {
-  if (EmitAddrsigSection) {
-    ObjWriter->AddrsigSection = Asm.getContext().getCOFFSection(
-        ".llvm_addrsig", COFF::IMAGE_SCN_LNK_REMOVE);
-    Asm.registerSection(*ObjWriter->AddrsigSection);
-  }
-
-  if (!Asm.CGProfile.empty()) {
-    ObjWriter->CGProfileSection = Asm.getContext().getCOFFSection(
-        ".llvm.call-graph-profile", COFF::IMAGE_SCN_LNK_REMOVE);
-    Asm.registerSection(*ObjWriter->CGProfileSection);
-  }
-
   ObjWriter->executePostLayoutBinding(Asm, Layout);
   if (DwoWriter)
     DwoWriter->executePostLayoutBinding(Asm, Layout);