[DebugMetadata][DwarfDebug] Fix DWARF emisson of function-local imported entities (3/7)

RFC https://discourse.llvm.org/t/rfc-dwarfdebug-fix-and-improve-handling-imported-entities-types-and-static-local-in-subprogram-and-lexical-block-scopes/68544

Fixed PR51501 (tests from D112337).

1. Reuse of DISubprogram's 'retainedNodes' to track other function-local
   entities together with local variables and labels (this patch cares about
   function-local import while D144006 and D144008 use the same approach for
   local types and static variables). So, effectively this patch moves ownership
   of tracking local import from DICompileUnit's 'imports' field to DISubprogram's
   'retainedNodes' and adjusts DWARF emitter for the new layout. The old layout
   is considered unsupported (DwarfDebug would assert on such debug metadata).

   DICompileUnit's 'imports' field is supposed to track global imported
   declarations as it does before.

   This addresses various FIXMEs and simplifies the next part of the patch.

2. Postpone emission of function-local imported entities from
   `DwarfDebug::endFunctionImpl()` to `DwarfDebug::endModule()`.
   While in `DwarfDebug::endFunctionImpl()` we do not have all the
   information about a parent subprogram or a referring subprogram
   (whether a subprogram inlined or not), so we can't guarantee we emit
   an imported entity correctly and place it in a proper subprogram tree.
   So now, we just gather needed details about the import itself and its
   parent entity (either a Subprogram or a LexicalBlock) during
   processing in `DwarfDebug::endFunctionImpl()`, but all the real work is
   done in `DwarfDebug::endModule()` when we have all the required
   information to make proper emission.

Authored-by: Kristina Bessonova <kbessonova@accesssoftek.com>

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

1 files changed, 84 insertions, 0 deletions

diff --git a/llvm/lib/Bitcode/Reader/MetadataLoader.cpp b/llvm/lib/Bitcode/Reader/MetadataLoader.cpp
index 4b5cfed..f13c26e 100644
--- a/llvm/lib/Bitcode/Reader/MetadataLoader.cpp
+++ b/llvm/lib/Bitcode/Reader/MetadataLoader.cpp
@@ -15,6 +15,7 @@
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLFunctionalExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
@@ -53,6 +54,7 @@
 #include <deque>
 #include <iterator>
 #include <limits>
+#include <map>
 #include <optional>
 #include <string>
 #include <tuple>
@@ -463,6 +465,9 @@ class MetadataLoader::MetadataLoaderImpl {
   bool NeedUpgradeToDIGlobalVariableExpression = false;
   bool NeedDeclareExpressionUpgrade = false;
 
+  /// Map DILocalScope to the enclosing DISubprogram, if any.
+  DenseMap<DILocalScope *, DISubprogram *> ParentSubprogram;
+
   /// True if metadata is being parsed for a module being ThinLTO imported.
   bool IsImporting = false;
 
@@ -521,6 +526,84 @@ class MetadataLoader::MetadataLoaderImpl {
     }
   }
 
+  DISubprogram *findEnclosingSubprogram(DILocalScope *S) {
+    if (!S)
+      return nullptr;
+    if (auto *SP = ParentSubprogram[S]) {
+      return SP;
+    }
+
+    DILocalScope *InitialScope = S;
+    DenseSet<DILocalScope *> Visited;
+    while (S && !isa<DISubprogram>(S)) {
+      S = dyn_cast_or_null<DILocalScope>(S->getScope());
+      if (Visited.contains(S))
+        break;
+      Visited.insert(S);
+    }
+    ParentSubprogram[InitialScope] = llvm::dyn_cast_or_null<DISubprogram>(S);
+
+    return ParentSubprogram[InitialScope];
+  }
+
+  /// Move local imports from DICompileUnit's 'imports' field to
+  /// DISubprogram's retainedNodes.
+  void upgradeCULocals() {
+    if (NamedMDNode *CUNodes = TheModule.getNamedMetadata("llvm.dbg.cu")) {
+      for (unsigned I = 0, E = CUNodes->getNumOperands(); I != E; ++I) {
+        auto *CU = dyn_cast<DICompileUnit>(CUNodes->getOperand(I));
+        if (!CU)
+          continue;
+
+        if (auto *RawImported = CU->getRawImportedEntities()) {
+          // Collect a set of imported entities to be moved.
+          SmallPtrSet<Metadata *, 8> EntitiesToRemove;
+          for (Metadata *Op : CU->getImportedEntities()->operands()) {
+            auto *IE = cast<DIImportedEntity>(Op);
+            if (auto *S = dyn_cast_or_null<DILocalScope>(IE->getScope())) {
+              EntitiesToRemove.insert(IE);
+            }
+          }
+
+          if (!EntitiesToRemove.empty()) {
+            // Make a new list of CU's 'imports'.
+            SmallVector<Metadata *> NewImports;
+            for (Metadata *Op : CU->getImportedEntities()->operands()) {
+              if (!EntitiesToRemove.contains(cast<DIImportedEntity>(Op))) {
+                NewImports.push_back(Op);
+              }
+            }
+
+            // Find DISubprogram corresponding to each entity.
+            std::map<DISubprogram *, SmallVector<Metadata *>> SPToEntities;
+            for (auto *I : EntitiesToRemove) {
+              auto *Entity = cast<DIImportedEntity>(I);
+              if (auto *SP = findEnclosingSubprogram(
+                      cast<DILocalScope>(Entity->getScope()))) {
+                SPToEntities[SP].push_back(Entity);
+              }
+            }
+
+            // Update DISubprograms' retainedNodes.
+            for (auto I = SPToEntities.begin(); I != SPToEntities.end(); ++I) {
+              auto *SP = I->first;
+              auto RetainedNodes = SP->getRetainedNodes();
+              SmallVector<Metadata *> MDs(RetainedNodes.begin(),
+                                          RetainedNodes.end());
+              MDs.append(I->second);
+              SP->replaceRetainedNodes(MDNode::get(Context, MDs));
+            }
+
+            // Remove entities with local scope from CU.
+            CU->replaceImportedEntities(MDTuple::get(Context, NewImports));
+          }
+        }
+      }
+    }
+
+    ParentSubprogram.clear();
+  }
+
   /// Remove a leading DW_OP_deref from DIExpressions in a dbg.declare that
   /// describes a function argument.
   void upgradeDeclareExpressions(Function &F) {
@@ -625,6 +708,7 @@ class MetadataLoader::MetadataLoaderImpl {
   void upgradeDebugInfo() {
     upgradeCUSubprograms();
     upgradeCUVariables();
+    upgradeCULocals();
   }
 
   void callMDTypeCallback(Metadata **Val, unsigned TypeID);