Merge pull request #8459 from joachimschmidt557/stage2-regalloc

Stage2 Register Manager: add getReg function

1 files changed, 89 insertions, 2 deletions

diff --git a/src/register_manager.zig b/src/register_manager.zig
index abb1632165..e11f2c3111 100644
--- a/src/register_manager.zig
+++ b/src/register_manager.zig
@@ -35,6 +35,10 @@ pub fn RegisterManager(
             self.registers.deinit(allocator);
         }
 
+        fn isTracked(reg: Register) bool {
+            return std.mem.indexOfScalar(Register, callee_preserved_regs, reg) != null;
+        }
+
         fn markRegUsed(self: *Self, reg: Register) void {
             if (FreeRegInt == u0) return;
             const index = reg.allocIndex() orelse return;
@@ -51,6 +55,13 @@ pub fn RegisterManager(
             self.free_registers |= @as(FreeRegInt, 1) << shift;
         }
 
+        pub fn isRegFree(self: Self, reg: Register) bool {
+            if (FreeRegInt == u0) return true;
+            const index = reg.allocIndex() orelse return true;
+            const shift = @intCast(ShiftInt, index);
+            return self.free_registers & @as(FreeRegInt, 1) << shift != 0;
+        }
+
         /// Returns whether this register was allocated in the course
         /// of this function
         pub fn isRegAllocated(self: Self, reg: Register) bool {
@@ -117,17 +128,61 @@ pub fn RegisterManager(
                 const regs_entry = self.registers.remove(reg).?;
                 const spilled_inst = regs_entry.value;
                 try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+                self.markRegFree(reg);
 
                 break :b reg;
             };
         }
 
-        pub fn getRegAssumeFree(self: *Self, reg: Register, inst: *ir.Inst) !void {
-            try self.registers.putNoClobber(self.getFunction().gpa, reg, inst);
+        /// Allocates the specified register with the specified
+        /// instruction. Spills the register if it is currently
+        /// allocated.
+        /// Before calling, must ensureCapacity + 1 on self.registers.
+        pub fn getReg(self: *Self, reg: Register, inst: *ir.Inst) !void {
+            if (!isTracked(reg)) return;
+
+            if (!self.isRegFree(reg)) {
+                // Move the instruction that was previously there to a
+                // stack allocation.
+                const regs_entry = self.registers.getEntry(reg).?;
+                const spilled_inst = regs_entry.value;
+                regs_entry.value = inst;
+                try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+            } else {
+                self.getRegAssumeFree(reg, inst);
+            }
+        }
+
+        /// Spills the register if it is currently allocated.
+        /// Does not track the register.
+        pub fn getRegWithoutTracking(self: *Self, reg: Register) !void {
+            if (!isTracked(reg)) return;
+
+            if (!self.isRegFree(reg)) {
+                // Move the instruction that was previously there to a
+                // stack allocation.
+                const regs_entry = self.registers.remove(reg).?;
+                const spilled_inst = regs_entry.value;
+                try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+                self.markRegFree(reg);
+            }
+        }
+
+        /// Allocates the specified register with the specified
+        /// instruction. Assumes that the register is free and no
+        /// spilling is necessary.
+        /// Before calling, must ensureCapacity + 1 on self.registers.
+        pub fn getRegAssumeFree(self: *Self, reg: Register, inst: *ir.Inst) void {
+            if (!isTracked(reg)) return;
+
+            self.registers.putAssumeCapacityNoClobber(reg, inst);
             self.markRegUsed(reg);
         }
 
+        /// Marks the specified register as free
         pub fn freeReg(self: *Self, reg: Register) void {
+            if (!isTracked(reg)) return;
+
             _ = self.registers.remove(reg);
             self.markRegFree(reg);
         }
@@ -226,3 +281,35 @@ test "allocReg: spilling" {
     std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction));
     std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items);
 }
+
+test "getReg" {
+    const allocator = std.testing.allocator;
+
+    var function = MockFunction{
+        .allocator = allocator,
+    };
+    defer function.deinit();
+
+    var mock_instruction = ir.Inst{
+        .tag = .breakpoint,
+        .ty = Type.initTag(.void),
+        .src = .unneeded,
+    };
+
+    std.testing.expect(!function.register_manager.isRegAllocated(.r2));
+    std.testing.expect(!function.register_manager.isRegAllocated(.r3));
+
+    try function.register_manager.registers.ensureCapacity(allocator, function.register_manager.registers.count() + 2);
+    try function.register_manager.getReg(.r3, &mock_instruction);
+
+    std.testing.expect(!function.register_manager.isRegAllocated(.r2));
+    std.testing.expect(function.register_manager.isRegAllocated(.r3));
+
+    // Spill r3
+    try function.register_manager.registers.ensureCapacity(allocator, function.register_manager.registers.count() + 2);
+    try function.register_manager.getReg(.r3, &mock_instruction);
+
+    std.testing.expect(!function.register_manager.isRegAllocated(.r2));
+    std.testing.expect(function.register_manager.isRegAllocated(.r3));
+    std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r3}, function.spilled.items);
+}