[CodeGen] Complete register allocation infrastructure

1 files changed, 93 insertions, 0 deletions

diff --git a/backend/include/rcc/CodeGen/VirtRegMap.h b/backend/include/rcc/CodeGen/VirtRegMap.h
new file mode 100644
index 0000000..f37af84
--- /dev/null
+++ b/backend/include/rcc/CodeGen/VirtRegMap.h
@@ -0,0 +1,93 @@
+// SPDX-License-Identifier: GPL-3.0-or-later
+
+#ifndef RCC_CODEGEN_VIRTREGMAP_H
+#define RCC_CODEGEN_VIRTREGMAP_H
+
+#include "rcc/CodeGen/MachineFrameInfo.h"
+#include "rcc/CodeGen/MachinePass.h"
+#include "rcc/CodeGen/TargetRegisterInfo.h"
+
+#include <cstdint>
+#include <optional>
+#include <unordered_map>
+#include <vector>
+
+namespace rcc {
+
+class MachineFunction;
+class TargetDescription;
+
+enum class VirtRegAllocationState : std::uint8_t {
+  Unassigned,
+  Assigned,
+  Spilled,
+  Split,
+};
+
+struct VirtRegSpillMetadata final {
+  std::optional<FrameIndex> StackSlot;
+  LaneMask Lanes = LaneMask::all();
+  std::uint64_t SpillCount = 0;
+  std::uint64_t ReloadCount = 0;
+  std::uint64_t RematerializationCount = 0;
+};
+
+struct VirtRegMapEntry final {
+  Register VirtualReg;
+  VirtRegAllocationState State = VirtRegAllocationState::Unassigned;
+  std::optional<PhysicalRegister> PhysicalReg;
+  std::optional<Register> SplitParent;
+  std::optional<Register> Origin;
+  bool Spillable = true;
+  VirtRegSpillMetadata Spill;
+};
+
+class VirtRegMap final {
+public:
+  explicit VirtRegMap(MachineFunction &Function) noexcept : Function(&Function) {}
+
+  [[nodiscard]] MachineFunction &function() const noexcept {
+    return *Function;
+  }
+  [[nodiscard]] Expected<void> add(Register VirtualReg, bool Spillable = true,
+                                   std::optional<Register> SplitParent = std::nullopt,
+                                   std::optional<Register> Origin = std::nullopt);
+  [[nodiscard]] Expected<void> assign(Register VirtualReg, PhysicalRegister PhysicalReg);
+  [[nodiscard]] Expected<void> assignStackSlot(Register VirtualReg, FrameIndex StackSlot,
+                                               LaneMask Lanes = LaneMask::all());
+  [[nodiscard]] Expected<void> markSplit(Register VirtualReg, Register Parent, Register Origin);
+  [[nodiscard]] Expected<void> setSpillable(Register VirtualReg, bool Spillable);
+  [[nodiscard]] VirtRegMapEntry *entry(Register VirtualReg) noexcept;
+  [[nodiscard]] const VirtRegMapEntry *entry(Register VirtualReg) const noexcept;
+  [[nodiscard]] std::optional<PhysicalRegister>
+  physicalRegister(Register VirtualReg) const noexcept;
+  [[nodiscard]] std::optional<FrameIndex> stackSlot(Register VirtualReg) const noexcept;
+  [[nodiscard]] bool isSpillable(Register VirtualReg) const noexcept;
+  [[nodiscard]] std::vector<const VirtRegMapEntry *> orderedEntries() const;
+  [[nodiscard]] Expected<void> verify() const;
+
+private:
+  MachineFunction *Function;
+  std::unordered_map<Register, VirtRegMapEntry> Entries;
+};
+
+[[nodiscard]] Expected<bool> rewriteVirtualRegisters(MachineFunction &Function,
+                                                     const TargetDescription &Target,
+                                                     const VirtRegMap &Mapping);
+
+class VirtualRegisterRewritePass final : public MachineFunctionPass {
+public:
+  explicit VirtualRegisterRewritePass(const VirtRegMap &Mapping) noexcept : Mapping(&Mapping) {}
+
+  [[nodiscard]] const MachinePassContract &contract() const noexcept override;
+  [[nodiscard]] Expected<MachinePassResult> run(MachineFunction &Function,
+                                                const TargetDescription &Target,
+                                                MachineAnalysisManager &Analyses) override;
+
+private:
+  const VirtRegMap *Mapping;
+};
+
+} // namespace rcc
+
+#endif // RCC_CODEGEN_VIRTREGMAP_H