[CodeGen] Separate pass instrumentation ownership

3 files changed, 238 insertions, 136 deletions

diff --git a/backend/include/rcc/CodeGen/MachinePassInstrumentation.h b/backend/include/rcc/CodeGen/MachinePassInstrumentation.h
new file mode 100644
index 0000000..9cc5098
--- /dev/null
+++ b/backend/include/rcc/CodeGen/MachinePassInstrumentation.h
@@ -0,0 +1,153 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef RCC_CODEGEN_MACHINEPASSINSTRUMENTATION_H
+#define RCC_CODEGEN_MACHINEPASSINSTRUMENTATION_H
+
+#include "rcc/CodeGen/MachineFunctionAnalysisManager.h"
+#include "rcc/IR/PassInstrumentation.h"
+
+#include <cstdint>
+#include <functional>
+#include <mutex>
+#include <optional>
+#include <string>
+#include <string_view>
+#include <unordered_map>
+#include <vector>
+
+namespace rcc {
+
+class CompilationUnit;
+class MachineFunction;
+struct MachinePassResult;
+struct CompilationUnitPassResult;
+
+struct InstrumentationStatistic final {
+  std::string Name;
+  bool IsAnalysis = false;
+  std::uint64_t Invocations = 0;
+  std::uint64_t CacheHits = 0;
+  std::uint64_t Changes = 0;
+  std::uint64_t Failures = 0;
+  std::uint64_t TotalNanoseconds = 0;
+};
+
+struct MachineActionID final {
+  std::string Pass;
+  std::string Rule;
+  MachineFunctionStableID Function;
+  MachineInstrStableID Instruction;
+
+  [[nodiscard]] std::string toString() const;
+
+  friend bool operator==(const MachineActionID &, const MachineActionID &) = default;
+};
+
+struct MachineActionEvent final {
+  MachineActionID ID;
+  std::uint64_t Cost = 0;
+  bool Accepted = false;
+
+  friend bool operator==(const MachineActionEvent &, const MachineActionEvent &) = default;
+};
+
+class MachinePassInstrumentation final {
+  using FunctionCallbackRegistry = PassInstrumentationCallbacks<MachineFunction, MachinePassResult>;
+  using UnitCallbackRegistry =
+      PassInstrumentationCallbacks<CompilationUnit, CompilationUnitPassResult>;
+
+public:
+  using BeforeCallback = FunctionCallbackRegistry::BeforeCallback;
+  using AfterCallback = FunctionCallbackRegistry::AfterCallback;
+  using FailureCallback = FunctionCallbackRegistry::FailureCallback;
+  using BeforeAnalysisCallback = FunctionCallbackRegistry::BeforeAnalysisCallback;
+  using AfterAnalysisCallback = FunctionCallbackRegistry::AfterAnalysisCallback;
+  using AnalysisFailureCallback = FunctionCallbackRegistry::AnalysisFailureCallback;
+  using IRPrintCallback =
+      std::function<void(InstrumentationPhase, std::string_view, const MachineFunction &)>;
+  using ChangeCallback =
+      std::function<void(std::string_view, const MachineFunction &, MachinePassChanges, bool)>;
+  using ActionFilter = std::function<bool(const MachineActionID &, std::uint64_t)>;
+
+  using BeforeUnitCallback = UnitCallbackRegistry::BeforeCallback;
+  using AfterUnitCallback = UnitCallbackRegistry::AfterCallback;
+  using UnitFailureCallback = UnitCallbackRegistry::FailureCallback;
+  using BeforeUnitAnalysisCallback = UnitCallbackRegistry::BeforeAnalysisCallback;
+  using AfterUnitAnalysisCallback = UnitCallbackRegistry::AfterAnalysisCallback;
+  using UnitAnalysisFailureCallback = UnitCallbackRegistry::AnalysisFailureCallback;
+
+  void addBeforeCallback(BeforeCallback Callback);
+  void addAfterCallback(AfterCallback Callback);
+  void addFailureCallback(FailureCallback Callback);
+  void addBeforeAnalysisCallback(BeforeAnalysisCallback Callback);
+  void addAfterAnalysisCallback(AfterAnalysisCallback Callback);
+  void addAnalysisFailureCallback(AnalysisFailureCallback Callback);
+  void addIRPrintCallback(IRPrintCallback Callback);
+  void addChangeCallback(ChangeCallback Callback);
+  void setActionFilter(ActionFilter Filter);
+  void setPerFunctionFuel(std::optional<std::uint64_t> Fuel);
+
+  void addBeforeUnitCallback(BeforeUnitCallback Callback);
+  void addAfterUnitCallback(AfterUnitCallback Callback);
+  void addUnitFailureCallback(UnitFailureCallback Callback);
+  void addBeforeUnitAnalysisCallback(BeforeUnitAnalysisCallback Callback);
+  void addAfterUnitAnalysisCallback(AfterUnitAnalysisCallback Callback);
+  void addUnitAnalysisFailureCallback(UnitAnalysisFailureCallback Callback);
+
+  [[nodiscard]] std::vector<InstrumentationStatistic> statistics() const;
+  [[nodiscard]] std::vector<MachineActionEvent> actionEvents() const;
+  [[nodiscard]] bool requiresSerialFunctionExecution() const;
+
+  [[nodiscard]] Expected<bool> requestAction(std::string_view Pass, std::string_view Rule,
+                                             const MachineFunction &Function,
+                                             MachineInstrStableID Instruction, std::uint64_t Cost);
+  [[nodiscard]] Expected<bool> requestAction(std::string_view Pass, std::string_view Rule,
+                                             const MachineFunction &Function, std::nullptr_t,
+                                             std::uint64_t Cost);
+
+private:
+  void before(std::string_view Name, const MachineFunction &Function) const;
+  void after(std::string_view Name, const MachineFunction &Function,
+             const MachinePassResult &Result, std::uint64_t Nanoseconds) const;
+  void failure(std::string_view Name, const MachineFunction &Function, const Error &ErrorValue,
+               std::uint64_t Nanoseconds = 0) const;
+  void beforeAnalysis(std::string_view Name, const MachineFunction &Function) const;
+  void afterAnalysis(std::string_view Name, const MachineFunction &Function, bool CacheHit,
+                     std::uint64_t Nanoseconds) const;
+  void analysisFailure(std::string_view Name, const MachineFunction &Function,
+                       const Error &ErrorValue, std::uint64_t Nanoseconds) const;
+  void change(std::string_view Name, const MachineFunction &Function, MachinePassChanges Changes,
+              bool FingerprintChanged) const;
+
+  void beforeUnit(std::string_view Name, const CompilationUnit &Unit) const;
+  void afterUnit(std::string_view Name, const CompilationUnit &Unit,
+                 const CompilationUnitPassResult &Result, std::uint64_t Nanoseconds) const;
+  void unitFailure(std::string_view Name, const CompilationUnit &Unit, const Error &ErrorValue,
+                   std::uint64_t Nanoseconds = 0) const;
+  void beforeAnalysis(std::string_view Name, const CompilationUnit &Unit) const;
+  void afterAnalysis(std::string_view Name, const CompilationUnit &Unit, bool CacheHit,
+                     std::uint64_t Nanoseconds) const;
+  void analysisFailure(std::string_view Name, const CompilationUnit &Unit, const Error &ErrorValue,
+                       std::uint64_t Nanoseconds) const;
+
+  [[nodiscard]] InstrumentationStatistic &statistic(std::string_view Name, bool IsAnalysis) const;
+
+  FunctionCallbackRegistry FunctionCallbacks;
+  UnitCallbackRegistry UnitCallbacks;
+  mutable std::mutex Mutex;
+  mutable std::mutex ActionMutex;
+  std::vector<IRPrintCallback> IRPrint;
+  std::vector<ChangeCallback> Change;
+  ActionFilter Filter;
+  std::optional<std::uint64_t> PerFunctionFuel;
+  std::unordered_map<std::uint64_t, std::uint64_t> ConsumedFuel;
+  std::vector<MachineActionEvent> ActionEvents;
+  mutable std::unordered_map<std::string, InstrumentationStatistic> Statistics;
+
+  template <typename, typename> friend struct AnalysisManagerInstrumentation;
+  template <typename, typename, typename> friend struct PassManagerCallbacks;
+};
+
+} // namespace rcc
+
+#endif // RCC_CODEGEN_MACHINEPASSINSTRUMENTATION_H
diff --git a/backend/include/rcc/CodeGen/MachinePassManager.h b/backend/include/rcc/CodeGen/MachinePassManager.h
index 3587395..af7a5d3 100644
--- a/backend/include/rcc/CodeGen/MachinePassManager.h
+++ b/backend/include/rcc/CodeGen/MachinePassManager.h
@@ -6,7 +6,7 @@
 #include "rcc/CodeGen/InstructionDescriptor.h"
 #include "rcc/CodeGen/MachineFunctionAnalysisManager.h"
 #include "rcc/CodeGen/MachineFunctionProperties.h"
-#include "rcc/IR/PassInstrumentation.h"
+#include "rcc/CodeGen/MachinePassInstrumentation.h"
 
 #include <cstddef>
 #include <memory>
diff --git a/backend/include/rcc/IR/PassInstrumentation.h b/backend/include/rcc/IR/PassInstrumentation.h
index 8ba2d2e..b0a673a 100644
--- a/backend/include/rcc/IR/PassInstrumentation.h
+++ b/backend/include/rcc/IR/PassInstrumentation.h
@@ -3,168 +3,117 @@
 #ifndef RCC_IR_PASSINSTRUMENTATION_H
 #define RCC_IR_PASSINSTRUMENTATION_H
 
-#include "rcc/CodeGen/MachineFunctionAnalysisManager.h"
+#include "rcc/Support/Error.h"
 
 #include <cstdint>
 #include <functional>
 #include <mutex>
-#include <optional>
-#include <string>
 #include <string_view>
-#include <unordered_map>
+#include <utility>
 #include <vector>
 
 namespace rcc {
 
-class CompilationUnit;
-class MachineFunction;
-struct MachinePassResult;
-struct CompilationUnitPassResult;
-
 enum class InstrumentationPhase : std::uint8_t {
   Before,
   After,
 };
 
-struct InstrumentationStatistic final {
-  std::string Name;
-  bool IsAnalysis = false;
-  std::uint64_t Invocations = 0;
-  std::uint64_t CacheHits = 0;
-  std::uint64_t Changes = 0;
-  std::uint64_t Failures = 0;
-  std::uint64_t TotalNanoseconds = 0;
-};
-
-struct MachineActionID final {
-  std::string Pass;
-  std::string Rule;
-  MachineFunctionStableID Function;
-  MachineInstrStableID Instruction;
-
-  [[nodiscard]] std::string toString() const;
-
-  friend bool operator==(const MachineActionID &, const MachineActionID &) = default;
-};
-
-struct MachineActionEvent final {
-  MachineActionID ID;
-  std::uint64_t Cost = 0;
-  bool Accepted = false;
-
-  friend bool operator==(const MachineActionEvent &, const MachineActionEvent &) = default;
-};
-
-class MachinePassInstrumentation final {
+template <typename IRUnitT, typename PassResultT> class PassInstrumentationCallbacks final {
 public:
-  using BeforeCallback = std::function<void(std::string_view, const MachineFunction &)>;
-  using AfterCallback =
-      std::function<void(std::string_view, const MachineFunction &, const MachinePassResult &)>;
-  using FailureCallback =
-      std::function<void(std::string_view, const MachineFunction &, const Error &)>;
-  using BeforeAnalysisCallback = std::function<void(std::string_view, const MachineFunction &)>;
+  using BeforeCallback = std::function<void(std::string_view, const IRUnitT &)>;
+  using AfterCallback = std::function<void(std::string_view, const IRUnitT &, const PassResultT &)>;
+  using FailureCallback = std::function<void(std::string_view, const IRUnitT &, const Error &)>;
+  using BeforeAnalysisCallback = std::function<void(std::string_view, const IRUnitT &)>;
   using AfterAnalysisCallback =
-      std::function<void(std::string_view, const MachineFunction &, bool, std::uint64_t)>;
+      std::function<void(std::string_view, const IRUnitT &, bool, std::uint64_t)>;
   using AnalysisFailureCallback =
-      std::function<void(std::string_view, const MachineFunction &, const Error &)>;
-  using IRPrintCallback =
-      std::function<void(InstrumentationPhase, std::string_view, const MachineFunction &)>;
-  using ChangeCallback =
-      std::function<void(std::string_view, const MachineFunction &, MachinePassChanges, bool)>;
-  using ActionFilter = std::function<bool(const MachineActionID &, std::uint64_t)>;
-
-  using BeforeUnitCallback = std::function<void(std::string_view, const CompilationUnit &)>;
-  using AfterUnitCallback = std::function<void(std::string_view, const CompilationUnit &,
-                                               const CompilationUnitPassResult &)>;
-  using UnitFailureCallback =
-      std::function<void(std::string_view, const CompilationUnit &, const Error &)>;
-  using BeforeUnitAnalysisCallback = std::function<void(std::string_view, const CompilationUnit &)>;
-  using AfterUnitAnalysisCallback =
-      std::function<void(std::string_view, const CompilationUnit &, bool, std::uint64_t)>;
-  using UnitAnalysisFailureCallback =
-      std::function<void(std::string_view, const CompilationUnit &, const Error &)>;
-
-  void addBeforeCallback(BeforeCallback Callback);
-  void addAfterCallback(AfterCallback Callback);
-  void addFailureCallback(FailureCallback Callback);
-  void addBeforeAnalysisCallback(BeforeAnalysisCallback Callback);
-  void addAfterAnalysisCallback(AfterAnalysisCallback Callback);
-  void addAnalysisFailureCallback(AnalysisFailureCallback Callback);
-  void addIRPrintCallback(IRPrintCallback Callback);
-  void addChangeCallback(ChangeCallback Callback);
-  void setActionFilter(ActionFilter Filter);
-  void setPerFunctionFuel(std::optional<std::uint64_t> Fuel);
-
-  void addBeforeUnitCallback(BeforeUnitCallback Callback);
-  void addAfterUnitCallback(AfterUnitCallback Callback);
-  void addUnitFailureCallback(UnitFailureCallback Callback);
-  void addBeforeUnitAnalysisCallback(BeforeUnitAnalysisCallback Callback);
-  void addAfterUnitAnalysisCallback(AfterUnitAnalysisCallback Callback);
-  void addUnitAnalysisFailureCallback(UnitAnalysisFailureCallback Callback);
-
-  [[nodiscard]] std::vector<InstrumentationStatistic> statistics() const;
-  [[nodiscard]] std::vector<MachineActionEvent> actionEvents() const;
-  [[nodiscard]] bool requiresSerialFunctionExecution() const;
-
-  [[nodiscard]] Expected<bool> requestAction(std::string_view Pass, std::string_view Rule,
-                                             const MachineFunction &Function,
-                                             MachineInstrStableID Instruction, std::uint64_t Cost);
-  [[nodiscard]] Expected<bool> requestAction(std::string_view Pass, std::string_view Rule,
-                                             const MachineFunction &Function, std::nullptr_t,
-                                             std::uint64_t Cost);
+      std::function<void(std::string_view, const IRUnitT &, const Error &)>;
+
+  void addBeforeCallback(BeforeCallback Callback) {
+    addCallback(Before, std::move(Callback));
+  }
+
+  void addAfterCallback(AfterCallback Callback) {
+    addCallback(After, std::move(Callback));
+  }
+
+  void addFailureCallback(FailureCallback Callback) {
+    addCallback(Failure, std::move(Callback));
+  }
+
+  void addBeforeAnalysisCallback(BeforeAnalysisCallback Callback) {
+    addCallback(BeforeAnalysis, std::move(Callback));
+  }
+
+  void addAfterAnalysisCallback(AfterAnalysisCallback Callback) {
+    addCallback(AfterAnalysis, std::move(Callback));
+  }
+
+  void addAnalysisFailureCallback(AnalysisFailureCallback Callback) {
+    addCallback(AnalysisFailure, std::move(Callback));
+  }
+
+  void runBefore(std::string_view Name, const IRUnitT &Unit) const {
+    for (const BeforeCallback &Callback : snapshot(Before))
+      Callback(Name, Unit);
+  }
+
+  void runAfter(std::string_view Name, const IRUnitT &Unit, const PassResultT &Result) const {
+    for (const AfterCallback &Callback : snapshot(After))
+      Callback(Name, Unit, Result);
+  }
+
+  void runFailure(std::string_view Name, const IRUnitT &Unit, const Error &FailureValue) const {
+    for (const FailureCallback &Callback : snapshot(Failure))
+      Callback(Name, Unit, FailureValue);
+  }
+
+  void runBeforeAnalysis(std::string_view Name, const IRUnitT &Unit) const {
+    for (const BeforeAnalysisCallback &Callback : snapshot(BeforeAnalysis))
+      Callback(Name, Unit);
+  }
+
+  void runAfterAnalysis(std::string_view Name, const IRUnitT &Unit, bool CacheHit,
+                        std::uint64_t Nanoseconds) const {
+    for (const AfterAnalysisCallback &Callback : snapshot(AfterAnalysis))
+      Callback(Name, Unit, CacheHit, Nanoseconds);
+  }
+
+  void runAnalysisFailure(std::string_view Name, const IRUnitT &Unit,
+                          const Error &FailureValue) const {
+    for (const AnalysisFailureCallback &Callback : snapshot(AnalysisFailure))
+      Callback(Name, Unit, FailureValue);
+  }
+
+  [[nodiscard]] bool empty() const {
+    std::scoped_lock Lock(Mutex);
+    return Before.empty() && After.empty() && Failure.empty() && BeforeAnalysis.empty() &&
+           AfterAnalysis.empty() && AnalysisFailure.empty();
+  }
 
 private:
-  void before(std::string_view Name, const MachineFunction &Function) const;
-  void after(std::string_view Name, const MachineFunction &Function,
-             const MachinePassResult &Result, std::uint64_t Nanoseconds) const;
-  void failure(std::string_view Name, const MachineFunction &Function, const Error &ErrorValue,
-               std::uint64_t Nanoseconds = 0) const;
-  void beforeAnalysis(std::string_view Name, const MachineFunction &Function) const;
-  void afterAnalysis(std::string_view Name, const MachineFunction &Function, bool CacheHit,
-                     std::uint64_t Nanoseconds) const;
-  void analysisFailure(std::string_view Name, const MachineFunction &Function,
-                       const Error &ErrorValue, std::uint64_t Nanoseconds) const;
-  void change(std::string_view Name, const MachineFunction &Function, MachinePassChanges Changes,
-              bool FingerprintChanged) const;
-
-  void beforeUnit(std::string_view Name, const CompilationUnit &Unit) const;
-  void afterUnit(std::string_view Name, const CompilationUnit &Unit,
-                 const CompilationUnitPassResult &Result, std::uint64_t Nanoseconds) const;
-  void unitFailure(std::string_view Name, const CompilationUnit &Unit, const Error &ErrorValue,
-                   std::uint64_t Nanoseconds = 0) const;
-  void beforeAnalysis(std::string_view Name, const CompilationUnit &Unit) const;
-  void afterAnalysis(std::string_view Name, const CompilationUnit &Unit, bool CacheHit,
-                     std::uint64_t Nanoseconds) const;
-  void analysisFailure(std::string_view Name, const CompilationUnit &Unit, const Error &ErrorValue,
-                       std::uint64_t Nanoseconds) const;
-
-  [[nodiscard]] InstrumentationStatistic &statistic(std::string_view Name, bool IsAnalysis) const;
+  template <typename CallbackT>
+  void addCallback(std::vector<CallbackT> &Callbacks, CallbackT Callback) {
+    if (!Callback)
+      return;
+    std::scoped_lock Lock(Mutex);
+    Callbacks.push_back(std::move(Callback));
+  }
+
+  template <typename CallbackT>
+  [[nodiscard]] std::vector<CallbackT> snapshot(const std::vector<CallbackT> &Callbacks) const {
+    std::scoped_lock Lock(Mutex);
+    return Callbacks;
+  }
 
   mutable std::mutex Mutex;
-  mutable std::mutex ActionMutex;
   std::vector<BeforeCallback> Before;
   std::vector<AfterCallback> After;
   std::vector<FailureCallback> Failure;
   std::vector<BeforeAnalysisCallback> BeforeAnalysis;
   std::vector<AfterAnalysisCallback> AfterAnalysis;
   std::vector<AnalysisFailureCallback> AnalysisFailure;
-  std::vector<IRPrintCallback> IRPrint;
-  std::vector<ChangeCallback> Change;
-  ActionFilter Filter;
-  std::optional<std::uint64_t> PerFunctionFuel;
-  std::unordered_map<std::uint64_t, std::uint64_t> ConsumedFuel;
-  std::vector<MachineActionEvent> ActionEvents;
-
-  std::vector<BeforeUnitCallback> BeforeUnit;
-  std::vector<AfterUnitCallback> AfterUnit;
-  std::vector<UnitFailureCallback> UnitFailure;
-  std::vector<BeforeUnitAnalysisCallback> BeforeUnitAnalysis;
-  std::vector<AfterUnitAnalysisCallback> AfterUnitAnalysis;
-  std::vector<UnitAnalysisFailureCallback> UnitAnalysisFailure;
-  mutable std::unordered_map<std::string, InstrumentationStatistic> Statistics;
-
-  template <typename, typename> friend struct AnalysisManagerInstrumentation;
-  template <typename, typename, typename> friend struct PassManagerCallbacks;
 };
 
 } // namespace rcc