[DwarfDebug] Move abstract entities tracking from DwarfCompileUnit to DwarfUnit

DwarfCompileUnit::constructVariableDIE() and constructLabelDIE() are
meant for constructing both abstract and concrete DIEs of a DbgEntity.
They use DwarfUnit::insertDIE() to store a freshly-created DIE. However,
insertDIE()/DwarfUnit::DITypeNodeToDieMap stores only single DIE per DINode.
If insertDIE() is called several times for the same instance of DINode,
only first DIE is saved in DwarfUnit::DITypeNodeToDieMap, as follows
from DenseMap::insert() specification.

It means, depending on what is called first,
DwarfCompileUnit::constructVariableDIE(LV, /* Abstract */ true) or
DwarfCompileUnit::constructVariableDIE(LV, /* Abstract */ false),
DwarfUnit::DITypeNodeToDieMap stores either abstract or concrete DIE of a node.

This behavior suggests an obscure API of DwarfCompileUnit, as it silently discards
one of the DIEs and makes it unclear what DwarfUnit::DITypeNodeToDieMap
is made for.

Also, DwarfFile and DwarfUnit have a tiny duplicate code piece.

To address that, DwarfInfoHolder class is introduced, which
stores DIEs for DILocalVariables and DILabels separately from
DIEs for other DINodes (as DILocalVariables and DILabels may have
concrete and abstract DIEs), and allows explicit access to
abstract/concrete DIEs of a debug info entity.

DwarfInfoHolder may later be used for tracking DIEs of abstract/concrete
lexical scopes.
Currently, concrete lexical block/subprogram DIEs are distinguished by
their DISubprogram/DILocalScope/DILocalScope+inlinedAt in DwarfCompileUnit.
As a result, the same DISubprogram can't be attached to two llvm::Functions
(https://lists.llvm.org/pipermail/llvm-dev/2020-September/145342.html).
Matching DISubprogram/DILocalScope DIEs with their LexicalScopes and
letting DwarfUnit members to access abstract scopes may enable linking
DISubprogram to several llvm::Functions, and allow the transition from
distinct to uniqued DISubprograms proposed here
llvm#142166 (comment).

Author: dzhidzhoev

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,7 @@ 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, VariableDie, Abstract);
   DV.setDIE(*VariableDie);
   // Abstract variables don't get common attributes later, so apply them now.
   if (Abstract) {
@@ -1010,7 +1010,7 @@ 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, LabelDie, Scope.isAbstractScope());
   DL.setDIE(*LabelDie);
 
   if (Scope.isAbstractScope())

llvm/lib/CodeGen/AsmPrinter/DwarfFile.h

@@ -15,6 +15,7 @@
 #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 <map>
 #include <memory>
@@ -26,9 +27,6 @@ class AsmPrinter;
 class DbgEntity;
 class DbgVariable;
 class DbgLabel;
-class DINode;
-class DILocalScope;
-class DISubprogram;
 class DwarfCompileUnit;
 class DwarfUnit;
 class LexicalScope;
@@ -53,6 +51,103 @@ 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 {
+  DenseMap<const DINodeT *, DIE *> AbstractMap;
+  DenseMap<const DINodeT *, SmallDenseMap<const DbgEntityT *, DIE *, 2>>
+      ConcreteMap;
+
+  static const DINodeT *getNode(const DbgEntityT &N) {
+    return cast<DINodeT>(N.getEntity());
+  }
+
+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 DbgEntityT &N, DIE *D) {
+    auto [_, Inserted] = ConcreteMap[getNode(N)].try_emplace(&N, D);
+    assert(Inserted && "Duplicate concrete DIE for debug info node");
+  }
+
+  void insertDIE(const DbgEntityT &N, DIE *D, bool Abstract) {
+    if (Abstract)
+      insertAbstractDIE(getNode(N), D);
+    else
+      insertConcreteDIE(N, D);
+  }
+
+  DIE *getAbstractDIE(const DINodeT *N) const { return AbstractMap.lookup(N); }
+
+  DIE *getConcreteDIE(const DbgEntityT &N) const {
+    if (auto I = ConcreteMap.find(getNode(N)); I != ConcreteMap.end())
+      return I->second.lookup(&N);
+    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;
+
+    if (auto I = ConcreteMap.find(N); I != ConcreteMap.end())
+      return I->second.empty() ? nullptr : I->second.begin()->second;
+
+    return nullptr;
+  }
+};
+
+/// 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 {
+  // DIEs of local DbgVariables.
+  DINodeInfoHolder<DILocalVariable, DbgVariable> LVHolder;
+  DINodeInfoHolder<DILabel, DbgLabel> LabelHolder;
+
+  /// 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)) &&
+           "Use getLabels().insertDIE() for labels or getLVs().insertDIE() for "
+           "local variables");
+    auto [_, Inserted] = MDNodeToDieMap.try_emplace(N, Die);
+    assert((Inserted || isa<DISubprogram>(N) || 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))) &&
+           "Use getLabels().getDIE() for labels or getLVs().getDIE() for "
+           "local variables");
+    return D;
+  }
+
+  auto &getLVs() { return LVHolder; }
+  auto &getLVs() const { return LVHolder; }
+  auto &getLabels() { return LabelHolder; }
+  auto &getLabels() const { return LabelHolder; }
+
+  /// 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);
+  }
+};
+
 class DwarfFile {
   // Target of Dwarf emission, used for sizing of abbreviations.
   AsmPrinter *Asm;
@@ -94,16 +189,14 @@ class DwarfFile {
   DenseMap<LexicalScope *, LabelList> ScopeLabels;
 
   // Collection of abstract subprogram DIEs.
+  // TODO: move it to InfoHolder?
   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);
@@ -183,13 +276,7 @@ class DwarfFile {
     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

llvm/lib/CodeGen/AsmPrinter/DwarfUnit.cpp

@@ -193,23 +193,11 @@ bool DwarfUnit::isShareableAcrossCUs(const DINode *D) const {
          !DD->generateTypeUnits();
 }
 
-DIE *DwarfUnit::getDIE(const DINode *D) const {
-  if (isShareableAcrossCUs(D))
-    return DU->getDIE(D);
-  return MDNodeToDieMap.lookup(D);
-}
-
 void DwarfUnit::insertDIE(const DINode *Desc, DIE *D) {
-  if (isShareableAcrossCUs(Desc)) {
-    DU->insertDIE(Desc, D);
-    return;
-  }
-  MDNodeToDieMap.insert(std::make_pair(Desc, D));
+  getDIEs(Desc).insertDIE(Desc, D);
 }
 
-void DwarfUnit::insertDIE(DIE *D) {
-  MDNodeToDieMap.insert(std::make_pair(nullptr, D));
-}
+void DwarfUnit::insertDIE(DIE *D) { InfoHolder.insertDIE(D); }
 
 void DwarfUnit::addFlag(DIE &Die, dwarf::Attribute Attribute) {
   if (DD->getDwarfVersion() >= 4)
@@ -803,7 +791,7 @@ void DwarfUnit::constructTypeDIE(DIE &Buffer, const DIStringType *STy) {
     addString(Buffer, dwarf::DW_AT_name, Name);
 
   if (DIVariable *Var = STy->getStringLength()) {
-    if (auto *VarDIE = getDIE(Var))
+    if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
       addDIEEntry(Buffer, dwarf::DW_AT_string_length, *VarDIE);
   } else if (DIExpression *Expr = STy->getStringLengthExp()) {
     DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1122,8 +1110,8 @@ void DwarfUnit::constructTypeDIE(DIE &Buffer, const DICompositeType *CTy) {
           constructTypeDIE(VariantPart, Composite);
         }
       } else if (Tag == dwarf::DW_TAG_namelist) {
-        auto *Var = dyn_cast<DINode>(Element);
-        auto *VarDIE = getDIE(Var);
+        auto *Var = dyn_cast<DIVariable>(Element);
+        auto *VarDIE = getDIEs(Var).getVariableDIE(Var);
         if (VarDIE) {
           DIE &ItemDie = createAndAddDIE(dwarf::DW_TAG_namelist_item, Buffer);
           addDIEEntry(ItemDie, dwarf::DW_AT_namelist_item, *VarDIE);
@@ -1185,7 +1173,7 @@ void DwarfUnit::constructTypeDIE(DIE &Buffer, const DICompositeType *CTy) {
       Tag == dwarf::DW_TAG_class_type || Tag == dwarf::DW_TAG_structure_type ||
       Tag == dwarf::DW_TAG_union_type) {
     if (auto *Var = dyn_cast_or_null<DIVariable>(CTy->getRawSizeInBits())) {
-      if (auto *VarDIE = getDIE(Var))
+      if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
         addDIEEntry(Buffer, dwarf::DW_AT_bit_size, *VarDIE);
     } else if (auto *Exp =
                    dyn_cast_or_null<DIExpression>(CTy->getRawSizeInBits())) {
@@ -1586,7 +1574,7 @@ void DwarfUnit::constructSubrangeDIE(DIE &DW_Subrange, const DISubrangeType *SR,
   auto AddBoundTypeEntry = [&](dwarf::Attribute Attr,
                                DISubrangeType::BoundType Bound) -> void {
     if (auto *BV = dyn_cast_if_present<DIVariable *>(Bound)) {
-      if (auto *VarDIE = getDIE(BV))
+      if (auto *VarDIE = getDIEs(BV).getVariableDIE(BV))
         addDIEEntry(DW_Subrange, Attr, *VarDIE);
     } else if (auto *BE = dyn_cast_if_present<DIExpression *>(Bound)) {
       DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1628,7 +1616,7 @@ void DwarfUnit::constructSubrangeDIE(DIE &Buffer, const DISubrange *SR) {
   auto AddBoundTypeEntry = [&](dwarf::Attribute Attr,
                                DISubrange::BoundType Bound) -> void {
     if (auto *BV = dyn_cast_if_present<DIVariable *>(Bound)) {
-      if (auto *VarDIE = getDIE(BV))
+      if (auto *VarDIE = getDIEs(BV).getVariableDIE(BV))
         addDIEEntry(DW_Subrange, Attr, *VarDIE);
     } else if (auto *BE = dyn_cast_if_present<DIExpression *>(Bound)) {
       DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1670,7 +1658,7 @@ void DwarfUnit::constructGenericSubrangeDIE(DIE &Buffer,
   auto AddBoundTypeEntry = [&](dwarf::Attribute Attr,
                                DIGenericSubrange::BoundType Bound) -> void {
     if (auto *BV = dyn_cast_if_present<DIVariable *>(Bound)) {
-      if (auto *VarDIE = getDIE(BV))
+      if (auto *VarDIE = getDIEs(BV).getVariableDIE(BV))
         addDIEEntry(DwGenericSubrange, Attr, *VarDIE);
     } else if (auto *BE = dyn_cast_if_present<DIExpression *>(Bound)) {
       if (BE->isConstant() &&
@@ -1749,7 +1737,7 @@ void DwarfUnit::constructArrayTypeDIE(DIE &Buffer, const DICompositeType *CTy) {
   }
 
   if (DIVariable *Var = CTy->getDataLocation()) {
-    if (auto *VarDIE = getDIE(Var))
+    if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
       addDIEEntry(Buffer, dwarf::DW_AT_data_location, *VarDIE);
   } else if (DIExpression *Expr = CTy->getDataLocationExp()) {
     DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1760,7 +1748,7 @@ void DwarfUnit::constructArrayTypeDIE(DIE &Buffer, const DICompositeType *CTy) {
   }
 
   if (DIVariable *Var = CTy->getAssociated()) {
-    if (auto *VarDIE = getDIE(Var))
+    if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
       addDIEEntry(Buffer, dwarf::DW_AT_associated, *VarDIE);
   } else if (DIExpression *Expr = CTy->getAssociatedExp()) {
     DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1771,7 +1759,7 @@ void DwarfUnit::constructArrayTypeDIE(DIE &Buffer, const DICompositeType *CTy) {
   }
 
   if (DIVariable *Var = CTy->getAllocated()) {
-    if (auto *VarDIE = getDIE(Var))
+    if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
       addDIEEntry(Buffer, dwarf::DW_AT_allocated, *VarDIE);
   } else if (DIExpression *Expr = CTy->getAllocatedExp()) {
     DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1896,7 +1884,7 @@ DIE &DwarfUnit::constructMemberDIE(DIE &Buffer, const DIDerivedType *DT) {
     if (DT->getRawSizeInBits() == nullptr) {
       // No size, just ignore.
     } else if (auto *Var = dyn_cast<DIVariable>(DT->getRawSizeInBits())) {
-      if (auto *VarDIE = getDIE(Var))
+      if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
         addDIEEntry(MemberDie, dwarf::DW_AT_bit_size, *VarDIE);
     } else if (auto *Exp = dyn_cast<DIExpression>(DT->getRawSizeInBits())) {
       DIELoc *Loc = new (DIEValueAllocator) DIELoc;
@@ -1921,7 +1909,7 @@ DIE &DwarfUnit::constructMemberDIE(DIE &Buffer, const DIDerivedType *DT) {
     // See https://dwarfstd.org/issues/250501.1.html
     if (auto *Var = dyn_cast_or_null<DIVariable>(DT->getRawOffsetInBits())) {
       if (!Asm->TM.Options.DebugStrictDwarf || DD->getDwarfVersion() >= 6) {
-        if (auto *VarDIE = getDIE(Var))
+        if (auto *VarDIE = getDIEs(Var).getVariableDIE(Var))
           addDIEEntry(MemberDie, dwarf::DW_AT_data_bit_offset, *VarDIE);
       }
     } else if (auto *Expr =

llvm/lib/CodeGen/AsmPrinter/DwarfUnit.h

@@ -61,7 +61,7 @@ class DwarfUnit : public DIEUnit {
 
   /// Tracks the mapping of unit level debug information variables to debug
   /// information entries.
-  DenseMap<const MDNode *, DIE *> MDNodeToDieMap;
+  DwarfInfoHolder InfoHolder;
 
   /// A list of all the DIEBlocks in use.
   std::vector<DIEBlock *> DIEBlocks;
@@ -139,7 +139,9 @@ class DwarfUnit : public DIEUnit {
   /// We delegate the request to DwarfDebug when the MDNode can be part of the
   /// type system, since DIEs for the type system can be shared across CUs and
   /// the mappings are kept in DwarfDebug.
-  DIE *getDIE(const DINode *D) const;
+  DIE *getDIE(const DINode *D) const {
+    return getDIEs(D).getDIE(D);
+  }
 
   /// Returns a fresh newly allocated DIELoc.
   DIELoc *getDIELoc() { return new (DIEValueAllocator) DIELoc; }
@@ -153,6 +155,18 @@ class DwarfUnit : public DIEUnit {
 
   void insertDIE(DIE *D);
 
+  const DwarfInfoHolder &getDIEs(const DINode *N) const {
+    if (isShareableAcrossCUs(N))
+      return DU->getDIEs();
+
+    return InfoHolder;
+  }
+
+  DwarfInfoHolder &getDIEs(const DINode *N) {
+    return const_cast<DwarfInfoHolder &>(
+        const_cast<const DwarfUnit *>(this)->getDIEs(N));
+  }
+
   /// Add a flag that is true to the DIE.
   void addFlag(DIE &Die, dwarf::Attribute Attribute);