make the AsmPrinterHandler array public

This lets external consumers customize the output, similar to how
AssemblyAnnotationWriter lets the caller define callbacks when printing
IR. The array of handlers already existed, this just cleans up the code
so that it can be exposed publically.

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

llvm/include/llvm/CodeGen/AsmPrinter.h

@@ -139,6 +139,23 @@ public:
   using GOTEquivUsePair = std::pair<const GlobalVariable *, unsigned>;
   MapVector<const MCSymbol *, GOTEquivUsePair> GlobalGOTEquivs;
 
+  /// struct HandlerInfo and Handlers permit users or target extended
+  /// AsmPrinter to add their own handlers.
+  struct HandlerInfo {
+    std::unique_ptr<AsmPrinterHandler> Handler;
+    const char *TimerName;
+    const char *TimerDescription;
+    const char *TimerGroupName;
+    const char *TimerGroupDescription;
+
+    HandlerInfo(std::unique_ptr<AsmPrinterHandler> Handler,
+                const char *TimerName, const char *TimerDescription,
+                const char *TimerGroupName, const char *TimerGroupDescription)
+        : Handler(std::move(Handler)), TimerName(TimerName),
+          TimerDescription(TimerDescription), TimerGroupName(TimerGroupName),
+          TimerGroupDescription(TimerGroupDescription) {}
+  };
+
 private:
   MCSymbol *CurrentFnEnd = nullptr;
 
@@ -162,26 +179,10 @@ private:
 protected:
   MCSymbol *CurrentFnBegin = nullptr;
 
-  /// Protected struct HandlerInfo and Handlers permit target extended
-  /// AsmPrinter adds their own handlers.
-  struct HandlerInfo {
-    std::unique_ptr<AsmPrinterHandler> Handler;
-    const char *TimerName;
-    const char *TimerDescription;
-    const char *TimerGroupName;
-    const char *TimerGroupDescription;
-
-    HandlerInfo(std::unique_ptr<AsmPrinterHandler> Handler,
-                const char *TimerName, const char *TimerDescription,
-                const char *TimerGroupName, const char *TimerGroupDescription)
-        : Handler(std::move(Handler)), TimerName(TimerName),
-          TimerDescription(TimerDescription), TimerGroupName(TimerGroupName),
-          TimerGroupDescription(TimerGroupDescription) {}
-  };
-
   /// A vector of all debug/EH info emitters we should use. This vector
   /// maintains ownership of the emitters.
-  SmallVector<HandlerInfo, 1> Handlers;
+  std::vector<HandlerInfo> Handlers;
+  size_t NumUserHandlers = 0;
 
 public:
   struct SrcMgrDiagInfo {
@@ -446,6 +447,11 @@ public:
   // Overridable Hooks
   //===------------------------------------------------------------------===//
 
+  void addAsmPrinterHandler(HandlerInfo Handler) {
+    Handlers.insert(Handlers.begin(), std::move(Handler));
+    NumUserHandlers++;
+  }
+
   // Targets can, or in the case of EmitInstruction, must implement these to
   // customize output.
 

llvm/include/llvm/CodeGen/AsmPrinterHandler.h

@@ -23,6 +23,7 @@ class MachineBasicBlock;
 class MachineFunction;
 class MachineInstr;
 class MCSymbol;
+class Module;
 
 typedef MCSymbol *ExceptionSymbolProvider(AsmPrinter *Asm,
                                           const MachineBasicBlock *MBB);
@@ -37,6 +38,8 @@ public:
   /// this tracks that size.
   virtual void setSymbolSize(const MCSymbol *Sym, uint64_t Size) = 0;
 
+  virtual void beginModule(Module *M) {}
+
   /// Emit all sections that should come after the content.
   virtual void endModule() = 0;
 
@@ -75,6 +78,7 @@ public:
   /// Process end of a basic block during basic block sections.
   virtual void endBasicBlock(const MachineBasicBlock &MBB) {}
 };
+
 } // End of namespace llvm
 
 #endif

llvm/include/llvm/CodeGen/DebugHandlerBase.h

@@ -115,6 +115,8 @@ private:
 
   // AsmPrinterHandler overrides.
 public:
+  void beginModule(Module *M) override;
+
   void beginInstruction(const MachineInstr *MI) override;
   void endInstruction() override;
 

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -188,7 +188,8 @@ AsmPrinter::AsmPrinter(TargetMachine &tm, std::unique_ptr<MCStreamer> Streamer)
 }
 
 AsmPrinter::~AsmPrinter() {
-  assert(!DD && Handlers.empty() && "Debug/EH info didn't get finalized");
+  assert(!DD && Handlers.size() == NumUserHandlers &&
+         "Debug/EH info didn't get finalized");
 
   if (GCMetadataPrinters) {
     gcp_map_type &GCMap = getGCMap(GCMetadataPrinters);
@@ -252,6 +253,9 @@ bool AsmPrinter::doInitialization(Module &M) {
   auto *MMIWP = getAnalysisIfAvailable<MachineModuleInfoWrapperPass>();
   MMI = MMIWP ? &MMIWP->getMMI() : nullptr;
 
+  if (!M.debug_compile_units().empty())
+    MMI->setDebugInfoAvailability(true);
+
   // Initialize TargetLoweringObjectFile.
   const_cast<TargetLoweringObjectFile&>(getObjFileLowering())
     .Initialize(OutContext, TM);
@@ -313,8 +317,7 @@ bool AsmPrinter::doInitialization(Module &M) {
                             CodeViewLineTablesGroupDescription);
     }
     if (!EmitCodeView || M.getDwarfVersion()) {
-      DD = new DwarfDebug(this, &M);
-      DD->beginModule();
+      DD = new DwarfDebug(this);
       Handlers.emplace_back(std::unique_ptr<DwarfDebug>(DD), DbgTimerName,
                             DbgTimerDescription, DWARFGroupName,
                             DWARFGroupDescription);
@@ -379,6 +382,13 @@ bool AsmPrinter::doInitialization(Module &M) {
     Handlers.emplace_back(std::make_unique<WinCFGuard>(this), CFGuardName,
                           CFGuardDescription, DWARFGroupName,
                           DWARFGroupDescription);
+
+  for (const HandlerInfo &HI : Handlers) {
+    NamedRegionTimer T(HI.TimerName, HI.TimerDescription, HI.TimerGroupName,
+                       HI.TimerGroupDescription, TimePassesIsEnabled);
+    HI.Handler->beginModule(&M);
+  }
+
   return false;
 }
 
@@ -1110,8 +1120,6 @@ void AsmPrinter::emitFunctionBody() {
   // Emit target-specific gunk before the function body.
   emitFunctionBodyStart();
 
-  bool ShouldPrintDebugScopes = MMI->hasDebugInfo();
-
   if (isVerbose()) {
     // Get MachineDominatorTree or compute it on the fly if it's unavailable
     MDT = getAnalysisIfAvailable<MachineDominatorTree>();
@@ -1149,14 +1157,11 @@ void AsmPrinter::emitFunctionBody() {
       if (MCSymbol *S = MI.getPreInstrSymbol())
         OutStreamer->emitLabel(S);
 
-      if (ShouldPrintDebugScopes) {
       for (const HandlerInfo &HI : Handlers) {
-          NamedRegionTimer T(HI.TimerName, HI.TimerDescription,
-                             HI.TimerGroupName, HI.TimerGroupDescription,
-                             TimePassesIsEnabled);
+        NamedRegionTimer T(HI.TimerName, HI.TimerDescription, HI.TimerGroupName,
+                           HI.TimerGroupDescription, TimePassesIsEnabled);
         HI.Handler->beginInstruction(&MI);
       }
-      }
 
       if (isVerbose())
         emitComments(MI, OutStreamer->GetCommentOS());
@@ -1209,15 +1214,12 @@ void AsmPrinter::emitFunctionBody() {
       if (MCSymbol *S = MI.getPostInstrSymbol())
         OutStreamer->emitLabel(S);
 
-      if (ShouldPrintDebugScopes) {
       for (const HandlerInfo &HI : Handlers) {
-          NamedRegionTimer T(HI.TimerName, HI.TimerDescription,
-                             HI.TimerGroupName, HI.TimerGroupDescription,
-                             TimePassesIsEnabled);
+        NamedRegionTimer T(HI.TimerName, HI.TimerDescription, HI.TimerGroupName,
+                           HI.TimerGroupDescription, TimePassesIsEnabled);
         HI.Handler->endInstruction();
       }
     }
-    }
 
     // We must emit temporary symbol for the end of this basic block, if either
     // we have BBLabels enabled or if this basic blocks marks the end of a
@@ -1637,7 +1639,11 @@ bool AsmPrinter::doFinalization(Module &M) {
                        HI.TimerGroupDescription, TimePassesIsEnabled);
     HI.Handler->endModule();
   }
-  Handlers.clear();
+
+  // This deletes all the ephemeral handlers that AsmPrinter added, while
+  // keeping all the user-added handlers alive until the AsmPrinter is
+  // destroyed.
+  Handlers.erase(Handlers.begin() + NumUserHandlers, Handlers.end());
   DD = nullptr;
 
   // If the target wants to know about weak references, print them all.

llvm/lib/CodeGen/AsmPrinter/CodeViewDebug.cpp

@@ -141,7 +141,6 @@ CodeViewDebug::CodeViewDebug(AsmPrinter *AP)
   if (!MMI->getModule()->getNamedMetadata("llvm.dbg.cu") ||
       !AP->getObjFileLowering().getCOFFDebugSymbolsSection()) {
     Asm = nullptr;
-    MMI->setDebugInfoAvailability(false);
     return;
   }
   // Tell MMI that we have debug info.
@@ -564,8 +563,6 @@ void CodeViewDebug::endModule() {
   if (!Asm || !MMI->hasDebugInfo())
     return;
 
-  assert(Asm != nullptr);
-
   // The COFF .debug$S section consists of several subsections, each starting
   // with a 4-byte control code (e.g. 0xF1, 0xF2, etc) and then a 4-byte length
   // of the payload followed by the payload itself.  The subsections are 4-byte

llvm/lib/CodeGen/AsmPrinter/DebugHandlerBase.cpp

@@ -91,6 +91,11 @@ DbgVariableLocation::extractFromMachineInstruction(
 
 DebugHandlerBase::DebugHandlerBase(AsmPrinter *A) : Asm(A), MMI(Asm->MMI) {}
 
+void DebugHandlerBase::beginModule(Module *M) {
+  if (M->debug_compile_units().empty())
+    Asm = nullptr;
+}
+
 // Each LexicalScope has first instruction and last instruction to mark
 // beginning and end of a scope respectively. Create an inverse map that list
 // scopes starts (and ends) with an instruction. One instruction may start (or
@@ -276,7 +281,7 @@ void DebugHandlerBase::beginFunction(const MachineFunction *MF) {
 }
 
 void DebugHandlerBase::beginInstruction(const MachineInstr *MI) {
-  if (!MMI->hasDebugInfo())
+  if (!Asm || !MMI->hasDebugInfo())
     return;
 
   assert(CurMI == nullptr);
@@ -302,7 +307,7 @@ void DebugHandlerBase::beginInstruction(const MachineInstr *MI) {
 }
 
 void DebugHandlerBase::endInstruction() {
-  if (!MMI->hasDebugInfo())
+  if (!Asm || !MMI->hasDebugInfo())
     return;
 
   assert(CurMI != nullptr);
@@ -334,7 +339,7 @@ void DebugHandlerBase::endInstruction() {
 }
 
 void DebugHandlerBase::endFunction(const MachineFunction *MF) {
-  if (hasDebugInfo(MMI, MF))
+  if (Asm && hasDebugInfo(MMI, MF))
     endFunctionImpl(MF);
   DbgValues.clear();
   DbgLabels.clear();

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp

@@ -148,10 +148,6 @@ static cl::opt<LinkageNameOption>
                                             "Abstract subprograms")),
                       cl::init(DefaultLinkageNames));
 
-static const char *const DWARFGroupName = "dwarf";
-static const char *const DWARFGroupDescription = "DWARF Emission";
-static const char *const DbgTimerName = "writer";
-static const char *const DbgTimerDescription = "DWARF Debug Writer";
 static constexpr unsigned ULEB128PadSize = 4;
 
 void DebugLocDwarfExpression::emitOp(uint8_t Op, const char *Comment) {
@@ -330,7 +326,7 @@ static AccelTableKind computeAccelTableKind(unsigned DwarfVersion,
   return AccelTableKind::None;
 }
 
-DwarfDebug::DwarfDebug(AsmPrinter *A, Module *M)
+DwarfDebug::DwarfDebug(AsmPrinter *A)
     : DebugHandlerBase(A), DebugLocs(A->OutStreamer->isVerboseAsm()),
       InfoHolder(A, "info_string", DIEValueAllocator),
       SkeletonHolder(A, "skel_string", DIEValueAllocator),
@@ -1111,21 +1107,17 @@ sortGlobalExprs(SmallVectorImpl<DwarfCompileUnit::GlobalExpr> &GVEs) {
 // Emit all Dwarf sections that should come prior to the content. Create
 // global DIEs and emit initial debug info sections. This is invoked by
 // the target AsmPrinter.
-void DwarfDebug::beginModule() {
-  NamedRegionTimer T(DbgTimerName, DbgTimerDescription, DWARFGroupName,
-                     DWARFGroupDescription, TimePassesIsEnabled);
-  if (DisableDebugInfoPrinting) {
-    MMI->setDebugInfoAvailability(false);
-    return;
-  }
+void DwarfDebug::beginModule(Module *M) {
+  DebugHandlerBase::beginModule(M);
 
-  const Module *M = MMI->getModule();
+  if (!Asm || !MMI->hasDebugInfo() || DisableDebugInfoPrinting)
+    return;
 
   unsigned NumDebugCUs = std::distance(M->debug_compile_units_begin(),
                                        M->debug_compile_units_end());
-  // Tell MMI whether we have debug info.
-  assert(MMI->hasDebugInfo() == (NumDebugCUs > 0) &&
-         "DebugInfoAvailabilty initialized unexpectedly");
+  assert(NumDebugCUs > 0 && "Asm unexpectedly initialized");
+  assert(MMI->hasDebugInfo() &&
+         "DebugInfoAvailabilty unexpectedly not initialized");
   SingleCU = NumDebugCUs == 1;
   DenseMap<DIGlobalVariable *, SmallVector<DwarfCompileUnit::GlobalExpr, 1>>
       GVMap;
@@ -1388,7 +1380,7 @@ void DwarfDebug::endModule() {
   // If we aren't actually generating debug info (check beginModule -
   // conditionalized on !DisableDebugInfoPrinting and the presence of the
   // llvm.dbg.cu metadata node)
-  if (!MMI->hasDebugInfo())
+  if (!Asm || !MMI->hasDebugInfo() || DisableDebugInfoPrinting)
     return;
 
   // Finalize the debug info for the module.
@@ -1899,7 +1891,8 @@ void DwarfDebug::beginInstruction(const MachineInstr *MI) {
   }
 
   DebugHandlerBase::beginInstruction(MI);
-  assert(CurMI);
+  if (!CurMI)
+    return;
 
   if (NoDebug)
     return;

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.h

@@ -621,13 +621,13 @@ public:
   //===--------------------------------------------------------------------===//
   // Main entry points.
   //
-  DwarfDebug(AsmPrinter *A, Module *M);
+  DwarfDebug(AsmPrinter *A);
 
   ~DwarfDebug() override;
 
   /// Emit all Dwarf sections that should come prior to the
   /// content.
-  void beginModule();
+  void beginModule(Module *M) override;
 
   /// Emit all Dwarf sections that should come after the content.
   void endModule() override;

llvm/lib/CodeGen/MachineModuleInfo.cpp

@@ -304,7 +304,7 @@ char MachineModuleInfoWrapperPass::ID = 0;
 bool MachineModuleInfoWrapperPass::doInitialization(Module &M) {
   MMI.initialize();
   MMI.TheModule = &M;
-  MMI.DbgInfoAvailable = !M.debug_compile_units().empty();
+  MMI.DbgInfoAvailable = false;
   return false;
 }
 
@@ -319,6 +319,6 @@ MachineModuleInfo MachineModuleAnalysis::run(Module &M,
                                              ModuleAnalysisManager &) {
   MachineModuleInfo MMI(TM);
   MMI.TheModule = &M;
-  MMI.DbgInfoAvailable = !M.debug_compile_units().empty();
+  MMI.DbgInfoAvailable = false;
   return MMI;
 }

llvm/lib/Target/BPF/BTFDebug.cpp

@@ -1076,6 +1076,9 @@ void BTFDebug::beginInstruction(const MachineInstr *MI) {
     }
   }
 
+  if (!CurMI) // no debug info
+    return;
+
   // Skip this instruction if no DebugLoc or the DebugLoc
   // is the same as the previous instruction.
   const DebugLoc &DL = MI->getDebugLoc();

llvm/unittests/CodeGen/AsmPrinterDwarfTest.cpp

@@ -8,8 +8,13 @@
 
 #include "TestAsmPrinter.h"
 #include "llvm/CodeGen/AsmPrinter.h"
+#include "llvm/CodeGen/MachineModuleInfo.h"
+#include "llvm/IR/LegacyPassManager.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCSectionELF.h"
+#include "llvm/Target/TargetMachine.h"
 #include "llvm/Testing/Support/Error.h"
 
 using namespace llvm;
@@ -367,4 +372,58 @@ TEST_F(AsmPrinterEmitDwarfUnitLengthAsHiLoDiffTest, DWARF64) {
   TestPrinter->getAP()->emitDwarfUnitLength(Hi, Lo, "");
 }
 
+class AsmPrinterHandlerTest : public AsmPrinterFixtureBase {
+  class TestHandler : public AsmPrinterHandler {
+    AsmPrinterHandlerTest &Test;
+
+  public:
+    TestHandler(AsmPrinterHandlerTest &Test) : Test(Test) {}
+    virtual ~TestHandler() {}
+    virtual void setSymbolSize(const MCSymbol *Sym, uint64_t Size) override {}
+    virtual void beginModule(Module *M) override { Test.BeginCount++; }
+    virtual void endModule() override { Test.EndCount++; }
+    virtual void beginFunction(const MachineFunction *MF) override {}
+    virtual void endFunction(const MachineFunction *MF) override {}
+    virtual void beginInstruction(const MachineInstr *MI) override {}
+    virtual void endInstruction() override {}
+  };
+
+protected:
+  bool init(const std::string &TripleStr, unsigned DwarfVersion,
+            dwarf::DwarfFormat DwarfFormat) {
+    if (!AsmPrinterFixtureBase::init(TripleStr, DwarfVersion, DwarfFormat))
+      return false;
+
+    auto *AP = TestPrinter->getAP();
+    AP->addAsmPrinterHandler(AsmPrinter::HandlerInfo(
+        std::unique_ptr<AsmPrinterHandler>(new TestHandler(*this)),
+        "TestTimerName", "TestTimerDesc", "TestGroupName", "TestGroupDesc"));
+    LLVMTargetMachine *LLVMTM = static_cast<LLVMTargetMachine *>(&AP->TM);
+    legacy::PassManager PM;
+    PM.add(new MachineModuleInfoWrapperPass(LLVMTM));
+    PM.add(TestPrinter->releaseAP()); // Takes ownership of destroying AP
+    LLVMContext Context;
+    std::unique_ptr<Module> M(new Module("TestModule", Context));
+    M->setDataLayout(LLVMTM->createDataLayout());
+    PM.run(*M);
+    // Now check that we can run it twice.
+    AP->addAsmPrinterHandler(AsmPrinter::HandlerInfo(
+        std::unique_ptr<AsmPrinterHandler>(new TestHandler(*this)),
+        "TestTimerName", "TestTimerDesc", "TestGroupName", "TestGroupDesc"));
+    PM.run(*M);
+    return true;
+  }
+
+  int BeginCount = 0;
+  int EndCount = 0;
+};
+
+TEST_F(AsmPrinterHandlerTest, Basic) {
+  if (!init("x86_64-pc-linux", /*DwarfVersion=*/4, dwarf::DWARF32))
+    return;
+
+  ASSERT_EQ(BeginCount, 3);
+  ASSERT_EQ(EndCount, 3);
+}
+
 } // end namespace

llvm/unittests/CodeGen/TestAsmPrinter.h

@@ -73,6 +73,7 @@ public:
   void setDwarfUsesRelocationsAcrossSections(bool Enable);
 
   AsmPrinter *getAP() const { return Asm.get(); }
+  AsmPrinter *releaseAP() { return Asm.release(); }
   MCContext &getCtx() const { return *MC; }
   MockMCStreamer &getMS() const { return *MS; }
 };