[StripDebugInfo] Reuse DebugInfoFinder for findind Live CUs

[maksimsab]

Summary

Previous search does not take into account @llvm.dbg.* intrinsics
and debug types information while DebugInfoFinder takes into account
such information.

[aprantl]

Deleting specialized code while improving correctness? That's always welcome!

llvm/lib/Transforms/IPO/StripSymbols.cpp

if (GlobalVariable *GV = dyn_cast<GlobalVariable>(C)) {
         RemoveDeadConstant(GV);
     } else
       RemoveDeadConstant(C);
   }
 
   return true;
 }
 
-/// Collects compilation units referenced by functions or lexical scopes.
-/// Accepts any DIScope and uses recursive bottom-up approach to reach either
-/// DISubprogram or DILexicalBlockBase.
-static void
-collectCUsWithScope(const DIScope *Scope, std::set<DICompileUnit *> &LiveCUs,
-                    SmallPtrSet<const DIScope *, 8> &VisitedScopes) {
-  if (!Scope)
-    return;
-
-  auto InS = VisitedScopes.insert(Scope);
-  if (!InS.second)
-    return;
-
-  if (const auto *SP = dyn_cast<DISubprogram>(Scope)) {
-    if (SP->getUnit())
-      LiveCUs.insert(SP->getUnit());
-    return;
-  }
-  if (const auto *LB = dyn_cast<DILexicalBlockBase>(Scope)) {
-    const DISubprogram *SP = LB->getSubprogram();
-    if (SP && SP->getUnit())
-      LiveCUs.insert(SP->getUnit());
-    return;
-  }
-
-  collectCUsWithScope(Scope->getScope(), LiveCUs, VisitedScopes);
-}
-
-static void
-collectCUsForInlinedFuncs(const DILocation *Loc,
-                          std::set<DICompileUnit *> &LiveCUs,
-                          SmallPtrSet<const DIScope *, 8> &VisitedScopes) {
-  if (!Loc || !Loc->getInlinedAt())
-    return;
-  collectCUsWithScope(Loc->getScope(), LiveCUs, VisitedScopes);
-  collectCUsForInlinedFuncs(Loc->getInlinedAt(), LiveCUs, VisitedScopes);
-}
-
 static bool stripDeadDebugInfoImpl(Module &M) {
   bool Changed = false;
 
   LLVMContext &C = M.getContext();
 
   // Find all debug info in F. This is actually overkill in terms of what we
   // want to do, but we want to try and be as resilient as possible in the face
   // of potential debug info changes by using the formal interfaces given to us
static bool stripDeadDebugInfoImpl(Module &M) {
   for (GlobalVariable &GV : M.globals()) {
     SmallVector<DIGlobalVariableExpression *, 1> GVEs;
     GV.getDebugInfo(GVEs);
     for (auto *GVE : GVEs)
       LiveGVs.insert(GVE);
   }
 
   std::set<DICompileUnit *> LiveCUs;
-  SmallPtrSet<const DIScope *, 8> VisitedScopes;
-  // Any CU is live if is referenced from a subprogram metadata that is attached
-  // to a function defined or inlined in the module.
-  for (const Function &Fn : M.functions()) {
-    collectCUsWithScope(Fn.getSubprogram(), LiveCUs, VisitedScopes);
-    for (const_inst_iterator I = inst_begin(&Fn), E = inst_end(&Fn); I != E;
-         ++I) {
-      if (!I->getDebugLoc())
-        continue;
-      const DILocation *DILoc = I->getDebugLoc().get();
-      collectCUsForInlinedFuncs(DILoc, LiveCUs, VisitedScopes);
-    }
+  DebugInfoFinder LiveCUFinder;
+  for (const Function &F : M.functions()) {
+    if (auto *SP = cast_or_null<DISubprogram>(F.getSubprogram()))
+      LiveCUFinder.processSubprogram(SP);
+    for (const Instruction &I : instructions(F))
+      LiveCUFinder.processInstruction(M, I);
   }
+  auto FoundCUs = LiveCUFinder.compile_units();
+  LiveCUs.insert(FoundCUs.begin(), FoundCUs.end());
 
   bool HasDeadCUs = false;
   for (DICompileUnit *DIC : F.compile_units()) {
     // Create our live global variable list.
     bool GlobalVariableChange = false;
     for (auto *DIG : DIC->getGlobalVariables()) {
       if (DIG->getExpression() && DIG->getExpression()->isConstant())
         LiveGVs.insert(DIG);

llvm/test/Transforms/StripSymbols/find-live-cu.ll

+; This test checks that strip-dead-debug-info pass doesn't delete debug compile
+; units if they are used by @llvm.dbg.* intrinsics
+
+; RUN: opt -passes='strip-dead-debug-info,verify' %s -S | FileCheck %s
+
+; CHECK: !llvm.dbg.cu = !{!{{[0-9]+}}, !{{[0-9]+}}}
+; CHECK-COUNT-2: !DICompileUnit
+
+declare void @llvm.dbg.value(metadata, metadata, metadata)
+
+define void @func() {
+  %a = alloca i64
+  call void @llvm.dbg.value(metadata ptr %a, metadata !7, metadata !DIExpression()), !dbg !9
+  ret void
+}
+
+!llvm.dbg.cu = !{!0, !1}
+!llvm.module.flags = !{!10}
+
+
+; We have two different compile units to able to check different paths of
+; finding compile units (intrinsic argument and attached location to instruction)
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !11)
+!1 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !11)
+
+!2 = distinct !DISubprogram(name: "func", unit: !0)
+!3 = distinct !DICompositeType(tag: DW_TAG_class_type, scope: !2)
+!4 = !DIDerivedType(tag: DW_TAG_member, scope: !3, baseType: !5)
+!5 = !DIDerivedType(tag: DW_TAG_pointer_type, baseType: !6)
+!6 = !DIBasicType(tag: DW_TAG_base_type)
+!7 = !DILocalVariable(name: "a", type: !5, scope: !8)
+!8 = distinct !DISubprogram(name: "func", unit: !1)
+!9 = !DILocation(scope: !8)
+
+!10 = !{i32 2, !"Debug Info Version", i32 3}
+
+!11 = !DIFile(filename: "a", directory: "")