stage2 regalloc: Add getReg and getRegWithoutTracking

2 files changed, 77 insertions, 12 deletions

diff --git a/src/codegen.zig b/src/codegen.zig
index fbd412ceba..cacf51730a 100644
--- a/src/codegen.zig
+++ b/src/codegen.zig
@@ -1783,8 +1783,8 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                             switch (mc_arg) {
                                 .none => continue,
                                 .register => |reg| {
+                                    try self.register_manager.getRegWithoutTracking(reg);
                                     try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
-                                    // TODO interact with the register allocator to mark the instruction as moved.
                                 },
                                 .stack_offset => {
                                     // Here we need to emit instructions like this:
@@ -1925,8 +1925,8 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                                 .compare_flags_signed => unreachable,
                                 .compare_flags_unsigned => unreachable,
                                 .register => |reg| {
+                                    try self.register_manager.getRegWithoutTracking(reg);
                                     try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
-                                    // TODO interact with the register allocator to mark the instruction as moved.
                                 },
                                 .stack_offset => {
                                     return self.fail(inst.base.src, "TODO implement calling with parameters in memory", .{});
@@ -1988,8 +1988,8 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                                 .compare_flags_signed => unreachable,
                                 .compare_flags_unsigned => unreachable,
                                 .register => |reg| {
+                                    try self.register_manager.getRegWithoutTracking(reg);
                                     try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
-                                    // TODO interact with the register allocator to mark the instruction as moved.
                                 },
                                 .stack_offset => {
                                     return self.fail(inst.base.src, "TODO implement calling with parameters in memory", .{});
@@ -2039,8 +2039,8 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                     switch (mc_arg) {
                         .none => continue,
                         .register => |reg| {
+                            try self.register_manager.getRegWithoutTracking(reg);
                             try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
-                            // TODO interact with the register allocator to mark the instruction as moved.
                         },
                         .stack_offset => {
                             // Here we need to emit instructions like this:
@@ -2704,8 +2704,11 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                         const reg_name = input[1 .. input.len - 1];
                         const reg = parseRegName(reg_name) orelse
                             return self.fail(inst.base.src, "unrecognized register: '{s}'", .{reg_name});
-                        const arg = try self.resolveInst(inst.args[i]);
-                        try self.genSetReg(inst.base.src, inst.args[i].ty, reg, arg);
+
+                        const arg = inst.args[i];
+                        const arg_mcv = try self.resolveInst(arg);
+                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
                     if (mem.eql(u8, inst.asm_source, "svc #0")) {
@@ -2734,8 +2737,11 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                         const reg_name = input[1 .. input.len - 1];
                         const reg = parseRegName(reg_name) orelse
                             return self.fail(inst.base.src, "unrecognized register: '{s}'", .{reg_name});
-                        const arg = try self.resolveInst(inst.args[i]);
-                        try self.genSetReg(inst.base.src, inst.args[i].ty, reg, arg);
+
+                        const arg = inst.args[i];
+                        const arg_mcv = try self.resolveInst(arg);
+                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
                     if (mem.eql(u8, inst.asm_source, "svc #0")) {
@@ -2766,8 +2772,11 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                         const reg_name = input[1 .. input.len - 1];
                         const reg = parseRegName(reg_name) orelse
                             return self.fail(inst.base.src, "unrecognized register: '{s}'", .{reg_name});
-                        const arg = try self.resolveInst(inst.args[i]);
-                        try self.genSetReg(inst.base.src, inst.args[i].ty, reg, arg);
+
+                        const arg = inst.args[i];
+                        const arg_mcv = try self.resolveInst(arg);
+                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
                     if (mem.eql(u8, inst.asm_source, "ecall")) {
@@ -2796,8 +2805,11 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                         const reg_name = input[1 .. input.len - 1];
                         const reg = parseRegName(reg_name) orelse
                             return self.fail(inst.base.src, "unrecognized register: '{s}'", .{reg_name});
-                        const arg = try self.resolveInst(inst.args[i]);
-                        try self.genSetReg(inst.base.src, inst.args[i].ty, reg, arg);
+
+                        const arg = inst.args[i];
+                        const arg_mcv = try self.resolveInst(arg);
+                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
                     if (mem.eql(u8, inst.asm_source, "syscall")) {
diff --git a/src/register_manager.zig b/src/register_manager.zig
index abb1632165..356251eed8 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,59 @@ 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;
             };
         }
 
+        /// Allocates the specified register with the specified
+        /// instruction. Spills the register if it is currently
+        /// allocated.
+        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 {
+                try 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.getEntry(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.
         pub fn getRegAssumeFree(self: *Self, reg: Register, inst: *ir.Inst) !void {
+            if (!isTracked(reg)) return;
+
             try self.registers.putNoClobber(self.getFunction().gpa, 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);
         }