Change DroppedVariableStats to be extensible.

Change DroppedVariableStats to have an extensible design so that we can
use it to add dropped statistics to MIR passes and the instruction
selector.

llvm/include/llvm/CodeGen/DroppedVariableStats.h

@@ -23,9 +23,16 @@
 
 namespace llvm {
 
-/// A class to collect and print dropped debug information variable statistics.
-/// After every LLVM IR pass is run, it will print how many #dbg_values were
-/// dropped due to that pass.
+/// A unique key that represents a debug variable.
+/// First const DIScope *: Represents the scope of the debug variable.
+/// Second const DIScope *: Represents the InlinedAt scope of the debug
+/// variable. const DILocalVariable *: It is a pointer to the debug variable
+/// itself.
+using VarID =
+    std::tuple<const DIScope *, const DIScope *, const DILocalVariable *>;
+
+/// A base class to collect and print dropped debug information variable
+/// statistics.
 class DroppedVariableStats {
 public:
   DroppedVariableStats(bool DroppedVarStatsEnabled)
@@ -35,30 +42,40 @@ class DroppedVariableStats {
           << "Pass Level, Pass Name, Num of Dropped Variables, Func or "
              "Module Name\n";
   };
+
+  virtual ~DroppedVariableStats() = default;
+
   // We intend this to be unique per-compilation, thus no copies.
   DroppedVariableStats(const DroppedVariableStats &) = delete;
   void operator=(const DroppedVariableStats &) = delete;
 
-  void registerCallbacks(PassInstrumentationCallbacks &PIC);
-  void runBeforePass(StringRef PassID, Any IR);
-  void runAfterPass(StringRef PassID, Any IR, const PreservedAnalyses &PA);
-  void runAfterPassInvalidated(StringRef PassID, const PreservedAnalyses &PA);
   bool getPassDroppedVariables() { return PassDroppedVariables; }
 
-private:
-  bool PassDroppedVariables = false;
-  bool DroppedVariableStatsEnabled = false;
-  /// A unique key that represents a #dbg_value.
-  using VarID =
-      std::tuple<const DIScope *, const DIScope *, const DILocalVariable *>;
+protected:
+  void setup() {
+    DebugVariablesStack.push_back(
+        {DenseMap<const Function *, DebugVariables>()});
+    InlinedAts.push_back(
+        {DenseMap<StringRef, DenseMap<VarID, DILocation *>>()});
+  }
 
+  void cleanup() {
+    assert(!DebugVariablesStack.empty() &&
+           "DebugVariablesStack shouldn't be empty!");
+    assert(!InlinedAts.empty() && "InlinedAts shouldn't be empty!");
+    DebugVariablesStack.pop_back();
+    InlinedAts.pop_back();
+  }
+
+  bool DroppedVariableStatsEnabled = false;
   struct DebugVariables {
     /// DenseSet of VarIDs before an optimization pass has run.
     DenseSet<VarID> DebugVariablesBefore;
     /// DenseSet of VarIDs after an optimization pass has run.
     DenseSet<VarID> DebugVariablesAfter;
   };
 
+protected:
   /// A stack of a DenseMap, that maps DebugVariables for every pass to an
   /// llvm::Function. A stack is used because an optimization pass can call
   /// other passes.
@@ -70,36 +87,134 @@ class DroppedVariableStats {
   /// DenseMap of VarIDs and their inlinedAt locations before an optimization
   /// pass has run.
   SmallVector<DenseMap<StringRef, DenseMap<VarID, DILocation *>>> InlinedAts;
+  /// Calculate the number of dropped variables in an llvm::Function or
+  /// llvm::MachineFunction and print the relevant information to stdout.
+  void calculateDroppedStatsAndPrint(DebugVariables &DbgVariables,
+                                     StringRef FuncName, StringRef PassID,
+                                     StringRef FuncOrModName,
+                                     StringRef PassLevel, const Function *Func);
 
-  /// Iterate over all Functions in a Module and report any dropped debug
-  /// information. Will call calculateDroppedVarStatsOnFunction on every
-  /// Function.
-  void calculateDroppedVarStatsOnModule(const Module *M, StringRef PassID,
-                                        std::string FuncOrModName,
-                                        std::string PassLevel);
-  /// Iterate over all Instructions in a Function and report any dropped debug
-  /// information.
-  void calculateDroppedVarStatsOnFunction(const Function *F, StringRef PassID,
-                                          std::string FuncOrModName,
-                                          std::string PassLevel);
+  /// Check if a \p Var has been dropped or is a false positive. Also update the
+  /// \p DroppedCount if a debug variable is dropped.
+  bool updateDroppedCount(DILocation *DbgLoc, const DIScope *Scope,
+                          const DIScope *DbgValScope,
+                          DenseMap<VarID, DILocation *> &InlinedAtsMap,
+                          VarID Var, unsigned &DroppedCount);
+  /// Run code to populate relevant data structures over an llvm::Function or
+  /// llvm::MachineFunction.
+  void run(DebugVariables &DbgVariables, StringRef FuncName, bool Before);
+  /// Populate the VarIDSet and InlinedAtMap with the relevant information
+  /// needed for before and after pass analysis to determine dropped variable
+  /// status.
+  void populateVarIDSetAndInlinedMap(
+      const DILocalVariable *DbgVar, DebugLoc DbgLoc, DenseSet<VarID> &VarIDSet,
+      DenseMap<StringRef, DenseMap<VarID, DILocation *>> &InlinedAtsMap,
+      StringRef FuncName, bool Before);
+  /// Visit every llvm::Instruction or llvm::MachineInstruction and check if the
+  /// debug variable denoted by its ID \p Var may have been dropped by an
+  /// optimization pass.
+  virtual void
+  visitEveryInstruction(unsigned &DroppedCount,
+                        DenseMap<VarID, DILocation *> &InlinedAtsMap,
+                        VarID Var) = 0;
+  /// Visit every debug record in an llvm::Function or llvm::MachineFunction
+  /// and call populateVarIDSetAndInlinedMap on it.
+  virtual void visitEveryDebugRecord(
+      DenseSet<VarID> &VarIDSet,
+      DenseMap<StringRef, DenseMap<VarID, DILocation *>> &InlinedAtsMap,
+      StringRef FuncName, bool Before) = 0;
+
+private:
+  /// Remove a dropped debug variable's VarID from all Sets in the
+  /// DroppedVariablesBefore stack.
+  void removeVarFromAllSets(VarID Var, const Function *F) {
+    // Do not remove Var from the last element, it will be popped from the
+    // stack.
+    for (auto &DebugVariablesMap : llvm::drop_end(DebugVariablesStack))
+      DebugVariablesMap[F].DebugVariablesBefore.erase(Var);
+  }
+  /// Return true if \p Scope is the same as \p DbgValScope or a child scope of
+  /// \p DbgValScope, return false otherwise.
+  bool isScopeChildOfOrEqualTo(const DIScope *Scope,
+                               const DIScope *DbgValScope);
+  /// Return true if \p InlinedAt is the same as \p DbgValInlinedAt or part of
+  /// the InlinedAt chain, return false otherwise.
+  bool isInlinedAtChildOfOrEqualTo(const DILocation *InlinedAt,
+                                   const DILocation *DbgValInlinedAt);
+  bool PassDroppedVariables = false;
+};
+
+/// A class to collect and print dropped debug information due to LLVM IR
+/// optimization passes. After every LLVM IR pass is run, it will print how many
+/// #dbg_values were dropped due to that pass.
+class DroppedVariableStatsIR : public DroppedVariableStats {
+public:
+  DroppedVariableStatsIR(bool DroppedVarStatsEnabled)
+      : llvm::DroppedVariableStats(DroppedVarStatsEnabled) {}
+
+  virtual ~DroppedVariableStatsIR() = default;
+
+  void runBeforePass(Any IR) {
+    setup();
+    if (const auto *M = unwrapIR<Module>(IR))
+      return this->runOnModule(M, true);
+    if (const auto *F = unwrapIR<Function>(IR))
+      return this->runOnFunction(F, true);
+  }
+
+  void runAfterPass(StringRef P, Any IR) {
+    if (const auto *M = unwrapIR<Module>(IR))
+      runAfterPassModule(P, M);
+    else if (const auto *F = unwrapIR<Function>(IR))
+      runAfterPassFunction(P, F);
+    cleanup();
+  }
+
+  void registerCallbacks(PassInstrumentationCallbacks &PIC);
+
+private:
+  const Function *Func;
+
+  void runAfterPassFunction(StringRef PassID, const Function *F) {
+    runOnFunction(F, false);
+    calculateDroppedVarStatsOnFunction(F, PassID, F->getName().str(),
+                                       "Function");
+  }
+
+  void runAfterPassModule(StringRef PassID, const Module *M) {
+    runOnModule(M, false);
+    calculateDroppedVarStatsOnModule(M, PassID, M->getName().str(), "Module");
+  }
   /// Populate DebugVariablesBefore, DebugVariablesAfter, InlinedAts before or
   /// after a pass has run to facilitate dropped variable calculation for an
   /// llvm::Function.
   void runOnFunction(const Function *F, bool Before);
+  /// Iterate over all Instructions in a Function and report any dropped debug
+  /// information.
+  void calculateDroppedVarStatsOnFunction(const Function *F, StringRef PassID,
+                                          StringRef FuncOrModName,
+                                          StringRef PassLevel);
   /// Populate DebugVariablesBefore, DebugVariablesAfter, InlinedAts before or
   /// after a pass has run to facilitate dropped variable calculation for an
   /// llvm::Module. Calls runOnFunction on every Function in the Module.
   void runOnModule(const Module *M, bool Before);
-  /// Remove a dropped #dbg_value VarID from all Sets in the
-  /// DroppedVariablesBefore stack.
-  void removeVarFromAllSets(VarID Var, const Function *F);
-  /// Return true if \p Scope is the same as \p DbgValScope or a child scope of
-  /// \p DbgValScope, return false otherwise.
-  bool isScopeChildOfOrEqualTo(DIScope *Scope, const DIScope *DbgValScope);
-  /// Return true if \p InlinedAt is the same as \p DbgValInlinedAt or part of
-  /// the InlinedAt chain, return false otherwise.
-  bool isInlinedAtChildOfOrEqualTo(const DILocation *InlinedAt,
-                                   const DILocation *DbgValInlinedAt);
+  /// Iterate over all Functions in a Module and report any dropped debug
+  /// information. Will call calculateDroppedVarStatsOnFunction on every
+  /// Function.
+  void calculateDroppedVarStatsOnModule(const Module *M, StringRef PassID,
+                                        StringRef FuncOrModName,
+                                        StringRef PassLevel);
+  /// Override base class method to run on an llvm::Function specifically.
+  virtual void
+  visitEveryInstruction(unsigned &DroppedCount,
+                        DenseMap<VarID, DILocation *> &InlinedAtsMap,
+                        VarID Var) override;
+  /// Override base class method to run on #dbg_values specifically.
+  virtual void visitEveryDebugRecord(
+      DenseSet<VarID> &VarIDSet,
+      DenseMap<StringRef, DenseMap<VarID, DILocation *>> &InlinedAtsMap,
+      StringRef FuncName, bool Before) override;
+
   template <typename IRUnitT> static const IRUnitT *unwrapIR(Any IR) {
     const IRUnitT **IRPtr = llvm::any_cast<const IRUnitT *>(&IR);
     return IRPtr ? *IRPtr : nullptr;

llvm/include/llvm/Passes/StandardInstrumentations.h

@@ -597,7 +597,7 @@ class StandardInstrumentations {
   PrintCrashIRInstrumentation PrintCrashIR;
   IRChangedTester ChangeTester;
   VerifyInstrumentation Verify;
-  DroppedVariableStats DroppedStats;
+  DroppedVariableStatsIR DroppedStatsIR;
 
   bool VerifyEach;