[DwarfDebug] Associate subprogram DIEs with their Functions

llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.cpp

@@ -178,7 +178,7 @@ unsigned DwarfCompileUnit::getOrCreateSourceID(const DIFile *File) {
 DIE *DwarfCompileUnit::getOrCreateGlobalVariableDIE(
     const DIGlobalVariable *GV, ArrayRef<GlobalExpr> GlobalExprs) {
   // Check for pre-existence.
-  if (DIE *Die = getDIE(GV))
+  if (DIE *Die = getDIEs(GV).getVariableDIE(GV))
     return Die;
 
   assert(GV);
@@ -795,7 +795,9 @@ DIE *DwarfCompileUnit::constructLexicalScopeDIE(LexicalScope *Scope) {
 
 DIE *DwarfCompileUnit::constructVariableDIE(DbgVariable &DV, bool Abstract) {
   auto *VariableDie = DIE::get(DIEValueAllocator, DV.getTag());
-  insertDIE(DV.getVariable(), VariableDie);
+  getDIEs(DV.getVariable())
+      .getLVs()
+      .insertDIE(DV.getVariable(), &DV, VariableDie, Abstract);
   DV.setDIE(*VariableDie);
   // Abstract variables don't get common attributes later, so apply them now.
   if (Abstract) {
@@ -1010,7 +1012,9 @@ DIE *DwarfCompileUnit::constructVariableDIE(DbgVariable &DV,
 DIE *DwarfCompileUnit::constructLabelDIE(DbgLabel &DL,
                                          const LexicalScope &Scope) {
   auto LabelDie = DIE::get(DIEValueAllocator, DL.getTag());
-  insertDIE(DL.getLabel(), LabelDie);
+  getDIEs(DL.getLabel())
+      .getLabels()
+      .insertDIE(DL.getLabel(), &DL, LabelDie, Scope.isAbstractScope());
   DL.setDIE(*LabelDie);
 
   if (Scope.isAbstractScope())
@@ -1472,8 +1476,9 @@ DIE *DwarfCompileUnit::getOrCreateImportedEntityDIE(
   return IMDie;
 }
 
-void DwarfCompileUnit::finishSubprogramDefinition(const DISubprogram *SP) {
-  DIE *D = getDIE(SP);
+void DwarfCompileUnit::finishSubprogramDefinition(const DISubprogram *SP,
+                                                  const Function *F) {
+  DIE *D = getDIEs(SP).getLocalScopes().getConcreteDIE(SP, F);
   if (DIE *AbsSPDIE = getAbstractScopeDIEs().lookup(SP)) {
     if (D)
       // If this subprogram has an abstract definition, reference that

llvm/lib/CodeGen/AsmPrinter/DwarfCompileUnit.h

@@ -79,16 +79,10 @@ class DwarfCompileUnit final : public DwarfUnit {
   // List of concrete lexical block scopes belong to subprograms within this CU.
   DenseMap<const DILocalScope *, DIE *> LexicalBlockDIEs;
 
-  // List of abstract local scopes (either DISubprogram or DILexicalBlock).
-  DenseMap<const DILocalScope *, DIE *> AbstractLocalScopeDIEs;
-  SmallPtrSet<const DISubprogram *, 8> FinalizedAbstractSubprograms;
-
   // List of inlined lexical block scopes that belong to subprograms within this
   // CU.
   DenseMap<const DILocalScope *, SmallVector<DIE *, 2>> InlinedLocalScopeDIEs;
 
-  DenseMap<const DINode *, std::unique_ptr<DbgEntity>> AbstractEntities;
-
   /// DWO ID for correlating skeleton and split units.
   uint64_t DWOId = 0;
 
@@ -126,22 +120,20 @@ class DwarfCompileUnit final : public DwarfUnit {
 
   bool isDwoUnit() const override;
 
-  DenseMap<const DILocalScope *, DIE *> &getAbstractScopeDIEs() {
-    if (isDwoUnit() && !DD->shareAcrossDWOCUs())
-      return AbstractLocalScopeDIEs;
-    return DU->getAbstractScopeDIEs();
+  DwarfInfoHolder &getDIEs(const DINode *N) { return DwarfUnit::getDIEs(N); }
+
+  DwarfInfoHolder &getDIEs() { return getDIEs(nullptr); }
+
+  DwarfInfoHolder::AbstractScopeMapT &getAbstractScopeDIEs() {
+    return getDIEs().getLocalScopes().getAbstractDIEs();
   }
 
   DenseMap<const DINode *, std::unique_ptr<DbgEntity>> &getAbstractEntities() {
-    if (isDwoUnit() && !DD->shareAcrossDWOCUs())
-      return AbstractEntities;
-    return DU->getAbstractEntities();
+    return getDIEs().getAbstractEntities();
   }
 
   auto &getFinalizedAbstractSubprograms() {
-    if (isDwoUnit() && !DD->shareAcrossDWOCUs())
-      return FinalizedAbstractSubprograms;
-    return DU->getFinalizedAbstractSubprograms();
+    return getDIEs().getFinalizedAbstractSubprograms();
   }
 
   void finishNonUnitTypeDIE(DIE& D, const DICompositeType *CTy) override;
@@ -327,7 +319,7 @@ class DwarfCompileUnit final : public DwarfUnit {
   DIE *getOrCreateImportedEntityDIE(const DIImportedEntity *IE);
   DIE *constructImportedEntityDIE(const DIImportedEntity *IE);
 
-  void finishSubprogramDefinition(const DISubprogram *SP);
+  void finishSubprogramDefinition(const DISubprogram *SP, const Function *F);
   void finishEntityDefinition(const DbgEntity *Entity);
   void attachLexicalScopesAbstractOrigins();
 

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.cpp

@@ -502,7 +502,8 @@ void DwarfDebug::addSubprogramNames(
   // well into the name table. Only do that if we are going to actually emit
   // that name.
   if (LinkageName != "" && SP->getName() != LinkageName &&
-      (useAllLinkageNames() || InfoHolder.getAbstractScopeDIEs().lookup(SP)))
+      (useAllLinkageNames() ||
+       InfoHolder.getDIEs().getLocalScopes().getAbstractDIEs().lookup(SP)))
     addAccelName(Unit, NameTableKind, LinkageName, Die);
 
   // If this is an Objective-C selector name add it to the ObjC accelerator
@@ -1263,11 +1264,13 @@ void DwarfDebug::finishEntityDefinitions() {
 }
 
 void DwarfDebug::finishSubprogramDefinitions() {
-  for (const DISubprogram *SP : ProcessedSPNodes) {
+  for (auto SPF : ProcessedSPNodes) {
+    const DISubprogram *SP = SPF.first;
     assert(SP->getUnit()->getEmissionKind() != DICompileUnit::NoDebug);
-    forBothCUs(
-        getOrCreateDwarfCompileUnit(SP->getUnit()),
-        [&](DwarfCompileUnit &CU) { CU.finishSubprogramDefinition(SP); });
+    forBothCUs(getOrCreateDwarfCompileUnit(SP->getUnit()),
+               [&](DwarfCompileUnit &CU) {
+                 CU.finishSubprogramDefinition(SP, SPF.second);
+               });
   }
 }
 
@@ -2784,7 +2787,7 @@ void DwarfDebug::endFunctionImpl(const MachineFunction *MF) {
     constructAbstractSubprogramScopeDIE(TheCU, AScope);
   }
 
-  ProcessedSPNodes.insert(SP);
+  ProcessedSPNodes.insert(std::make_pair(SP, &F));
   DIE &ScopeDIE =
       TheCU.constructSubprogramScopeDIE(SP, F, FnScope, FunctionLineTableLabel);
   if (auto *SkelCU = TheCU.getSkeleton())

llvm/lib/CodeGen/AsmPrinter/DwarfDebug.h

@@ -373,7 +373,8 @@ class DwarfDebug : public DebugHandlerBase {
 
   /// This is a collection of subprogram MDNodes that are processed to
   /// create DIEs.
-  SmallSetVector<const DISubprogram *, 16> ProcessedSPNodes;
+  SmallSetVector<std::pair<const DISubprogram *, const Function *>, 16>
+      ProcessedSPNodes;
 
   /// Map function-local imported entities to their parent local scope
   /// (either DILexicalBlock or DISubprogram) for a processed function

llvm/lib/CodeGen/AsmPrinter/DwarfFile.h

@@ -15,9 +15,12 @@
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/DIE.h"
+#include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/Support/Allocator.h"
+#include <functional>
 #include <map>
 #include <memory>
+#include <optional>
 #include <utility>
 
 namespace llvm {
@@ -26,9 +29,6 @@ class AsmPrinter;
 class DbgEntity;
 class DbgVariable;
 class DbgLabel;
-class DINode;
-class DILocalScope;
-class DISubprogram;
 class DwarfCompileUnit;
 class DwarfUnit;
 class LexicalScope;
@@ -53,6 +53,144 @@ struct RangeSpanList {
   SmallVector<RangeSpan, 2> Ranges;
 };
 
+/// Tracks abstract and concrete DIEs for debug info entities of a certain type.
+template <typename DINodeT, typename DbgEntityT> class DINodeInfoHolder {
+public:
+  using AbstractMapT = DenseMap<const DINodeT *, DIE *>;
+  using ConcreteMapT =
+      DenseMap<const DINodeT *, SmallDenseMap<const DbgEntityT *, DIE *, 2>>;
+
+private:
+  AbstractMapT AbstractMap;
+  ConcreteMapT ConcreteMap;
+
+public:
+  void insertAbstractDIE(const DINodeT *N, DIE *D) {
+    auto [_, Inserted] = AbstractMap.try_emplace(N, D);
+    assert(Inserted && "Duplicate abstract DIE for debug info node");
+  }
+
+  void insertConcreteDIE(const DINodeT *N, const DbgEntityT *E, DIE *D) {
+    auto [_, Inserted] = ConcreteMap[N].try_emplace(E, D);
+    assert(Inserted && "Duplicate concrete DIE for debug info node");
+  }
+
+  void insertDIE(const DINodeT *N, const DbgEntityT *E, DIE *D, bool Abstract) {
+    if (Abstract)
+      insertAbstractDIE(N, D);
+    else
+      insertConcreteDIE(N, E, D);
+  }
+
+  DIE *getAbstractDIE(const DINodeT *N) const { return AbstractMap.lookup(N); }
+
+  std::optional<
+      std::reference_wrapper<const typename ConcreteMapT::mapped_type>>
+  getConcreteDIEs(const DINodeT *N) const {
+    if (auto I = ConcreteMap.find(N); I != ConcreteMap.end())
+      return std::make_optional(std::ref(I->second));
+    return std::nullopt;
+  }
+
+  DIE *getConcreteDIE(const DINodeT *N, const DbgEntityT *E) const {
+    if (auto I = getConcreteDIEs(N))
+      return I->get().lookup(E);
+    return nullptr;
+  }
+
+  DIE *getAnyConcreteDIE(const DINodeT *N) const {
+    if (auto I = getConcreteDIEs(N))
+      return I->get().empty() ? nullptr : I->get().begin()->second;
+    return nullptr;
+  }
+
+  /// Returns abstract DIE for the entity.
+  /// If no abstract DIE was created, returns any concrete DIE for the entity.
+  DIE *getDIE(const DINodeT *N) const {
+    if (DIE *D = getAbstractDIE(N))
+      return D;
+
+    return getAnyConcreteDIE(N);
+  }
+
+  AbstractMapT &getAbstractDIEs() { return AbstractMap; }
+};
+
+/// Tracks DIEs for debug info entites.
+/// These DIEs can be shared across CUs, that is why we keep the map here
+/// instead of in DwarfCompileUnit.
+class DwarfInfoHolder {
+public:
+  using LocalScopeHolderT = DINodeInfoHolder<DILocalScope, Function>;
+  using AbstractScopeMapT = LocalScopeHolderT::AbstractMapT;
+
+private:
+  /// DIEs of local DbgVariables.
+  DINodeInfoHolder<DILocalVariable, DbgVariable> LVHolder;
+  /// DIEs of labels.
+  DINodeInfoHolder<DILabel, DbgLabel> LabelHolder;
+  DenseMap<const DINode *, std::unique_ptr<DbgEntity>> AbstractEntities;
+  /// DIEs of abstract local scopes and concrete non-inlined subprograms.
+  /// Inlined subprograms and concrete lexical blocks are not stored here.
+  LocalScopeHolderT LSHolder;
+  /// Keeps track of abstract subprograms to populate them only once.
+  // FIXME: merge creation and population of abstract scopes.
+  SmallPtrSet<const DISubprogram *, 8> FinalizedAbstractSubprograms;
+
+  /// Other DINodes with the corresponding DIEs.
+  DenseMap<const DINode *, DIE *> MDNodeToDieMap;
+
+public:
+  void insertDIE(const DINode *N, DIE *Die) {
+    assert((!isa<DILabel>(N) && !isa<DILocalVariable>(N) &&
+            !isa<DILocalScope>(N)) &&
+           "Use getLabels().insertDIE() for labels or getLVs().insertDIE() for "
+           "local variables, or getSubprogram().insertDIE() for subprograms.");
+    auto [_, Inserted] = MDNodeToDieMap.try_emplace(N, Die);
+    assert((Inserted || isa<DIType>(N)) &&
+           "DIE for this DINode has already been added");
+  }
+
+  void insertDIE(DIE *D) { MDNodeToDieMap.try_emplace(nullptr, D); }
+
+  DIE *getDIE(const DINode *N) const {
+    DIE *D = MDNodeToDieMap.lookup(N);
+    assert((!D || (!isa<DILabel>(N) && !isa<DILocalVariable>(N) &&
+                   !isa<DILocalScope>(N))) &&
+           "Use getLabels().getDIE() for labels or getLVs().getDIE() for "
+           "local variables, or getLocalScopes().getDIE() for local scopes.");
+    return D;
+  }
+
+  auto &getLVs() { return LVHolder; }
+  auto &getLVs() const { return LVHolder; }
+
+  auto &getLabels() { return LabelHolder; }
+  auto &getLabels() const { return LabelHolder; }
+
+  auto &getLocalScopes() { return LSHolder; }
+  auto &getLocalScopes() const { return LSHolder; }
+
+  /// For a global variable, returns DIE of the variable.
+  ///
+  /// For a local variable, returns abstract DIE of the variable.
+  /// If no abstract DIE was created, returns any concrete DIE of the variable.
+  DIE *getVariableDIE(const DIVariable *V) const {
+    if (auto *LV = dyn_cast<DILocalVariable>(V))
+      if (DIE *D = getLVs().getDIE(LV))
+        return D;
+    return getDIE(V);
+  }
+
+  DenseMap<const DINode *, std::unique_ptr<DbgEntity>> &getAbstractEntities() {
+    return AbstractEntities;
+  }
+
+  auto &getFinalizedAbstractSubprograms() {
+    return FinalizedAbstractSubprograms;
+  }
+};
+
 class DwarfFile {
   // Target of Dwarf emission, used for sizing of abbreviations.
   AsmPrinter *Asm;
@@ -93,17 +231,7 @@ class DwarfFile {
   using LabelList = SmallVector<DbgLabel *, 4>;
   DenseMap<LexicalScope *, LabelList> ScopeLabels;
 
-  // Collection of abstract subprogram DIEs.
-  DenseMap<const DILocalScope *, DIE *> AbstractLocalScopeDIEs;
-  DenseMap<const DINode *, std::unique_ptr<DbgEntity>> AbstractEntities;
-  /// Keeps track of abstract subprograms to populate them only once.
-  // FIXME: merge creation and population of abstract scopes.
-  SmallPtrSet<const DISubprogram *, 8> FinalizedAbstractSubprograms;
-
-  /// Maps MDNodes for type system with the corresponding DIEs. These DIEs can
-  /// be shared across CUs, that is why we keep the map here instead
-  /// of in DwarfCompileUnit.
-  DenseMap<const MDNode *, DIE *> DITypeNodeToDieMap;
+  DwarfInfoHolder InfoHolder;
 
 public:
   DwarfFile(AsmPrinter *AP, StringRef Pref, BumpPtrAllocator &DA);
@@ -171,25 +299,7 @@ class DwarfFile {
     return ScopeLabels;
   }
 
-  DenseMap<const DILocalScope *, DIE *> &getAbstractScopeDIEs() {
-    return AbstractLocalScopeDIEs;
-  }
-
-  DenseMap<const DINode *, std::unique_ptr<DbgEntity>> &getAbstractEntities() {
-    return AbstractEntities;
-  }
-
-  auto &getFinalizedAbstractSubprograms() {
-    return FinalizedAbstractSubprograms;
-  }
-
-  void insertDIE(const MDNode *TypeMD, DIE *Die) {
-    DITypeNodeToDieMap.insert(std::make_pair(TypeMD, Die));
-  }
-
-  DIE *getDIE(const MDNode *TypeMD) {
-    return DITypeNodeToDieMap.lookup(TypeMD);
-  }
+  DwarfInfoHolder &getDIEs() { return InfoHolder; }
 };
 
 } // end namespace llvm