[llvm] [NFC] Move DroppedVariableStats to its own file and redesign it to be extensible. (PR #115563)

[Adrian Prantl via llvm-commits]

================
@@ -132,64 +75,122 @@ void DroppedVariableStats::calculateDroppedVarStatsOnFunction(
     PassDroppedVariables = false;
 }
 
-void DroppedVariableStats::runAfterPassInvalidated(
-    StringRef PassID, const PreservedAnalyses &PA) {
-  DebugVariablesStack.pop_back();
-  InlinedAts.pop_back();
+bool DroppedVariableStats::checkDroppedStatus(
+    DILocation *DbgLoc, const DIScope *Scope, const DIScope *DbgValScope,
+    DenseMap<VarID, DILocation *> &InlinedAtsMap, VarID Var) {
+
+  // If the Scope is a child of, or equal to the DbgValScope and is inlined at
+  // the Var's InlinedAt location, return true to signify that the Var has been
+  // dropped.
+  if (isScopeChildOfOrEqualTo(Scope, DbgValScope))
+    if (isInlinedAtChildOfOrEqualTo(DbgLoc->getInlinedAt(), InlinedAtsMap[Var]))
+      return true;
+  return false;
+}
+
+void DroppedVariableStats::run(DebugVariables &DbgVariables, StringRef FuncName,
+                               bool Before) {
+  auto &VarIDSet = (Before ? DbgVariables.DebugVariablesBefore
+                           : DbgVariables.DebugVariablesAfter);
+  auto &InlinedAtsMap = InlinedAts.back();
+  if (Before)
+    InlinedAtsMap.try_emplace(FuncName, DenseMap<VarID, DILocation *>());
+  VarIDSet = DenseSet<VarID>();
+  visitEveryDebugIntrinsic(VarIDSet, InlinedAtsMap, FuncName, Before);
+}
+
+void DroppedVariableStats::populateVarIDSetAndInlinedMap(
+    const DILocalVariable *DbgVar, DebugLoc DbgLoc, DenseSet<VarID> &VarIDSet,
+    DenseMap<StringRef, DenseMap<VarID, DILocation *>> &InlinedAtsMap,
+    StringRef FuncName, bool Before) {
+  VarID Key{DbgVar->getScope(), DbgLoc->getInlinedAtScope(), DbgVar};
+  VarIDSet.insert(Key);
+  if (Before)
+    InlinedAtsMap[FuncName].try_emplace(Key, DbgLoc.getInlinedAt());
 }
 
-void DroppedVariableStats::runAfterPass(StringRef PassID, Any IR,
-                                        const PreservedAnalyses &PA) {
-  std::string PassLevel;
-  std::string FuncOrModName;
-  if (auto *M = unwrapIR<Module>(IR)) {
-    this->runOnModule(M, false);
-    PassLevel = "Module";
-    FuncOrModName = M->getName();
-    calculateDroppedVarStatsOnModule(M, PassID, FuncOrModName, PassLevel);
-  } else if (auto *F = unwrapIR<Function>(IR)) {
-    this->runOnFunction(F, false);
-    PassLevel = "Function";
-    FuncOrModName = F->getName();
-    calculateDroppedVarStatsOnFunction(F, PassID, FuncOrModName, PassLevel);
+void DroppedVariableStatsIR::runOnFunction(const Function *F, bool Before) {
+  auto &DebugVariables = DebugVariablesStack.back()[F];
+  auto FuncName = F->getName();
+  Func = F;
+  run(DebugVariables, FuncName, Before);
+}
+
+void DroppedVariableStatsIR::calculateDroppedVarStatsOnFunction(
+    const Function *F, StringRef PassID, StringRef FuncOrModName,
+    StringRef PassLevel) {
+  Func = F;
+  StringRef FuncName = F->getName();
+  DebugVariables &DbgVariables = DebugVariablesStack.back()[F];
+  calculateDroppedStatsAndPrint(DbgVariables, FuncName, PassID, FuncOrModName,
+                                PassLevel);
+}
+
+void DroppedVariableStatsIR::runOnModule(const Module *M, bool Before) {
+  for (auto &F : *M)
+    runOnFunction(&F, Before);
+}
+
+void DroppedVariableStatsIR::calculateDroppedVarStatsOnModule(
+    const Module *M, StringRef PassID, StringRef FuncOrModName,
+    StringRef PassLevel) {
+  for (auto &F : *M) {
+    calculateDroppedVarStatsOnFunction(&F, PassID, FuncOrModName, PassLevel);
   }
+}
+
+void DroppedVariableStatsIR::registerCallbacks(
+    PassInstrumentationCallbacks &PIC) {
+  if (!DroppedVariableStatsEnabled)
+    return;
 
-  DebugVariablesStack.pop_back();
-  InlinedAts.pop_back();
-  return;
+  PIC.registerBeforeNonSkippedPassCallback(
+      [this](StringRef P, Any IR) { return runBeforePass(IR); });
+  PIC.registerAfterPassCallback(
+      [this](StringRef P, Any IR, const PreservedAnalyses &PA) {
+        return runAfterPass(P, IR);
+      });
+  PIC.registerAfterPassInvalidatedCallback(
+      [this](StringRef P, const PreservedAnalyses &PA) { return cleanup(); });
 }
 
-bool DroppedVariableStats::isScopeChildOfOrEqualTo(DIScope *Scope,
-                                                   const DIScope *DbgValScope) {
-  while (Scope != nullptr) {
-    if (VisitedScope.find(Scope) == VisitedScope.end()) {
-      VisitedScope.insert(Scope);
-      if (Scope == DbgValScope) {
-        VisitedScope.clear();
-        return true;
+void DroppedVariableStatsIR::visitEveryDebugVariable(
+    unsigned &DroppedCount, DenseSet<VarID> &DebugVariablesBeforeSet,
+    DenseSet<VarID> &DebugVariablesAfterSet,
+    DenseMap<VarID, DILocation *> &InlinedAtsMap) {
+  for (VarID Var : DebugVariablesBeforeSet) {
+    if (DebugVariablesAfterSet.contains(Var))
+      continue;
----------------
adrian-prantl wrote:

Why is this part of the virtual function? It seems like it wouldn't change.

https://github.com/llvm/llvm-project/pull/115563