1 files changed, 227 insertions, 0 deletions
diff --git a/backend/include/rcc/IR/PassInstrumentation.h b/backend/include/rcc/IR/PassInstrumentation.h
new file mode 100644
index 0000000..082ec81
--- /dev/null
+++ b/backend/include/rcc/IR/PassInstrumentation.h
@@ -0,0 +1,227 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef RCC_IR_PASSINSTRUMENTATION_H
+#define RCC_IR_PASSINSTRUMENTATION_H
+
+#include "rcc/CodeGen/MachineFunctionAnalysisManager.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;
+class TargetDescription;
+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 {
+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 AfterAnalysisCallback =
+      std::function<void(std::string_view, const MachineFunction &, 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]] 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:
+  struct BufferedEvents final {
+    const TargetDescription *Target = nullptr;
+    std::vector<std::function<void()>> Events;
+  };
+
+  [[nodiscard]] bool bufferFunctionEvent(const MachineFunction &Function,
+                                         std::function<void(const MachineFunction &)> Replay) const;
+  void beginBuffering(BufferedEvents &Buffer, const TargetDescription &Target) const noexcept;
+  void endBuffering() const noexcept;
+  void replayBufferedEvents(BufferedEvents &Buffer) const;
+
+  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;
+
+  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;
+
+  static thread_local BufferedEvents *ActiveBuffer;
+
+  template <typename, typename> friend struct AnalysisManagerInstrumentation;
+  template <typename, typename, typename> friend struct PassManagerCallbacks;
+};
+
+class MachineActionReplay final {
+public:
+  template <typename AnalysisManagerT>
+  explicit MachineActionReplay(AnalysisManagerT &Analyses)
+      : Request([&Analyses](std::string_view Rule, const MachineFunction &Function,
+                            MachineInstrStableID Instruction, std::uint64_t Cost) {
+          return Analyses.requestAction(Rule, Function, Instruction, Cost);
+        }) {}
+
+  template <typename AnalysisManagerT>
+  MachineActionReplay(
+      AnalysisManagerT &Analyses,
+      const std::unordered_map<MachineInstrStableID, MachineInstrStableID> &SourceToClone)
+      : MachineActionReplay(Analyses) {
+    for (const auto &[Source, Clone] : SourceToClone)
+      RecordingInstructionIDs.try_emplace(Clone, Source);
+  }
+
+  [[nodiscard]] Expected<bool> request(std::string_view Rule, const MachineFunction &Function,
+                                       MachineInstrStableID Instruction = {},
+                                       std::uint64_t Cost = 1);
+  void beginReplay() noexcept;
+  [[nodiscard]] Expected<void> finishReplay() const;
+
+private:
+  struct Decision final {
+    std::string Rule;
+    MachineFunctionStableID Function;
+    MachineInstrStableID Instruction;
+    std::uint64_t Cost = 0;
+    bool Accepted = false;
+  };
+
+  std::function<Expected<bool>(std::string_view, const MachineFunction &, MachineInstrStableID,
+                               std::uint64_t)>
+      Request;
+  std::unordered_map<MachineInstrStableID, MachineInstrStableID> RecordingInstructionIDs;
+  std::vector<Decision> Decisions;
+  std::size_t NextDecision = 0;
+  bool IsReplaying = false;
+};
+
+} // namespace rcc
+
+#endif // RCC_IR_PASSINSTRUMENTATION_H