[LFI][MC] Call setLFIRewriter during LFIMCStreamer initialization (#188625)

Calls `Streamer.setLFIRewriter` during generic LFIMCStreamer
initialization rather than requiring it to be done during
backend-specific initialization. This better follows the existing
conventions in `create*` functions in `TargetRegistry.h`.

Also re-adds the call to initSections for LFI in `llvm-mc.cpp`
(necessary in order to emit the ABI Note section), along with a test to
make sure ABI note emission with the rewriter is working.

llvm/include/llvm/MC/MCLFI.h

@@ -22,4 +22,6 @@ class Triple;
 LLVM_ABI void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
                                       const Triple &TheTriple);
 
+LLVM_ABI void emitLFINoteSection(MCStreamer &Streamer, MCContext &Ctx);
+
 } // namespace llvm

llvm/include/llvm/MC/TargetRegistry.h

@@ -239,7 +239,7 @@ public:
                                   const MCInstrInfo &MCII);
 
   using MCLFIRewriterCtorTy =
-      MCLFIRewriter *(*)(MCStreamer & S,
+      MCLFIRewriter *(*)(MCContext & Ctx,
                          std::unique_ptr<MCRegisterInfo> &&RegInfo,
                          std::unique_ptr<MCInstrInfo> &&InstInfo);
 
@@ -576,11 +576,12 @@ public:
     return nullptr;
   }
 
-  void createMCLFIRewriter(MCStreamer &S,
-                           std::unique_ptr<MCRegisterInfo> &&RegInfo,
-                           std::unique_ptr<MCInstrInfo> &&InstInfo) const {
+  MCLFIRewriter *
+  createMCLFIRewriter(MCContext &Ctx, std::unique_ptr<MCRegisterInfo> &&RegInfo,
+                      std::unique_ptr<MCInstrInfo> &&InstInfo) const {
     if (MCLFIRewriterCtorFn)
-      MCLFIRewriterCtorFn(S, std::move(RegInfo), std::move(InstInfo));
+      return MCLFIRewriterCtorFn(Ctx, std::move(RegInfo), std::move(InstInfo));
+    return nullptr;
   }
 
   /// createMCRelocationInfo - Create a target specific MCRelocationInfo.

llvm/lib/MC/MCAsmStreamer.cpp

@@ -20,6 +20,7 @@
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
@@ -2596,6 +2597,9 @@ void MCAsmStreamer::emitRawTextImpl(StringRef String) {
 }
 
 void MCAsmStreamer::finishImpl() {
+  if (getContext().getTargetTriple().isLFI())
+    emitLFINoteSection(*this, getContext());
+
   // If we are generating dwarf for assembly source files dump out the sections.
   if (getContext().getGenDwarfForAssembly())
     MCGenDwarfInfo::Emit(this);

llvm/lib/MC/MCELFStreamer.cpp

@@ -22,6 +22,7 @@
 #include "llvm/MC/MCELFObjectWriter.h"
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCFixup.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCSection.h"
@@ -373,6 +374,9 @@ void MCELFStreamer::finishImpl() {
                             DummyAttributeSection, GNUAttributes);
   }
 
+  if (Ctx.getTargetTriple().isLFI())
+    emitLFINoteSection(*this, Ctx);
+
   finalizeCGProfile();
   emitFrames();
 

llvm/lib/MC/MCLFI.cpp

@@ -15,6 +15,7 @@
 #include "llvm/BinaryFormat/ELF.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCRegisterInfo.h"
 #include "llvm/MC/MCSectionELF.h"
 #include "llvm/MC/MCStreamer.h"
@@ -34,6 +35,25 @@ cl::opt<bool> FlagEnableRewriting("lfi-enable-rewriter",
 void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
                              const Triple &TheTriple) {
   assert(TheTriple.isLFI());
+
+  std::string Error;
+  const Target *TheTarget = TargetRegistry::lookupTarget(TheTriple, Error);
+
+  // Create the target-specific MCLFIRewriter.
+  assert(TheTarget != nullptr);
+  if (FlagEnableRewriting) {
+    auto MRI =
+        std::unique_ptr<MCRegisterInfo>(TheTarget->createMCRegInfo(TheTriple));
+    auto MII = std::unique_ptr<MCInstrInfo>(TheTarget->createMCInstrInfo());
+    Streamer.setLFIRewriter(std::unique_ptr<MCLFIRewriter>(
+        TheTarget->createMCLFIRewriter(Ctx, std::move(MRI), std::move(MII))));
+  }
+}
+
+void emitLFINoteSection(MCStreamer &Streamer, MCContext &Ctx) {
+  const Triple &TheTriple = Ctx.getTargetTriple();
+  assert(TheTriple.isLFI());
+
   const char *NoteName;
   const char *NoteArch;
   switch (TheTriple.getArch()) {
@@ -45,25 +65,12 @@ void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
     reportFatalUsageError("Unsupported architecture for LFI");
   }
 
-  std::string Error; // empty
-  const Target *TheTarget = TargetRegistry::lookupTarget(TheTriple, Error);
-
-  // Create the Target specific MCLFIRewriter.
-  assert(TheTarget != nullptr);
-  if (FlagEnableRewriting) {
-    TheTarget->createMCLFIRewriter(
-        Streamer,
-        std::unique_ptr<MCRegisterInfo>(TheTarget->createMCRegInfo(TheTriple)),
-        std::unique_ptr<MCInstrInfo>(TheTarget->createMCInstrInfo()));
-  }
-
   // Emit an ELF Note section in its own COMDAT group which identifies LFI
   // object files.
   MCSectionELF *Note = Ctx.getELFSection(NoteName, ELF::SHT_NOTE,
                                          ELF::SHF_ALLOC | ELF::SHF_GROUP, 0,
                                          NoteName, /*IsComdat=*/true);
 
-  Streamer.pushSection();
   Streamer.switchSection(Note);
   Streamer.emitIntValue(strlen(NoteNamespace) + 1, 4);
   Streamer.emitIntValue(strlen(NoteArch) + 1, 4);
@@ -74,7 +81,6 @@ void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
   Streamer.emitBytes(NoteArch);
   Streamer.emitIntValue(0, 1); // NUL terminator
   Streamer.emitValueToAlignment(Align(4));
-  Streamer.popSection();
 }
 
 } // namespace llvm

llvm/test/MC/AArch64/LFI/abi-note.s

@@ -0,0 +1,15 @@
+// RUN: llvm-mc -triple aarch64_lfi %s | FileCheck %s
+// RUN: llvm-mc -filetype=obj -triple aarch64_lfi %s | llvm-readelf -S - | FileCheck %s --check-prefix=ELF
+
+// CHECK:      .section .note.LFI.ABI.aarch64,"aG",@note,.note.LFI.ABI.aarch64,comdat
+// CHECK-NEXT: .word 4
+// CHECK-NEXT: .word 8
+// CHECK-NEXT: .word 1
+// CHECK-NEXT: .ascii "LFI"
+// CHECK-NEXT: .byte 0
+// CHECK-NEXT: .p2align 2, 0x0
+// CHECK-NEXT: .ascii "aarch64"
+// CHECK-NEXT: .byte 0
+// CHECK-NEXT: .p2align 2, 0x0
+
+// ELF: .note.LFI.ABI.aarch64 NOTE {{.*}} AG