Commit 8b4e91e18c

joachimschmidt557 <joachim.schmidt557@outlook.com>
2021-05-14 11:09:11
stage2 register manager: clean up API and add more unit tests
1 parent 5185b56
Changed files (2)
src/codegen.zig
@@ -952,7 +952,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
         /// allocated. A second call to `copyToTmpRegister` may return the same register.
         /// This can have a side effect of spilling instructions to the stack to free up a register.
         fn copyToTmpRegister(self: *Self, src: LazySrcLoc, ty: Type, mcv: MCValue) !Register {
-            const reg = try self.register_manager.allocRegWithoutTracking(&.{});
+            const reg = try self.register_manager.allocReg(null, &.{});
             try self.genSetReg(src, ty, reg, mcv);
             return reg;
         }
@@ -2228,7 +2228,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                             switch (mc_arg) {
                                 .none => continue,
                                 .register => |reg| {
-                                    try self.register_manager.getRegWithoutTracking(reg);
+                                    try self.register_manager.getReg(reg, null);
                                     try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
                                 },
                                 .stack_offset => {
@@ -2370,7 +2370,7 @@ 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.register_manager.getReg(reg, null);
                                     try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
                                 },
                                 .stack_offset => {
@@ -2433,7 +2433,7 @@ 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.register_manager.getReg(reg, null);
                                     try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
                                 },
                                 .stack_offset => {
@@ -2486,7 +2486,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
                         .register => |reg| {
                             // TODO prevent this macho if block to be generated for all archs
                             switch (arch) {
-                                .x86_64, .aarch64 => try self.register_manager.getRegWithoutTracking(reg),
+                                .x86_64, .aarch64 => try self.register_manager.getReg(reg, null),
                                 else => unreachable,
                             }
                             try self.genSetReg(arg.src, arg.ty, reg, arg_mcv);
@@ -3190,7 +3190,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
 
                         const arg = inst.args[i];
                         const arg_mcv = try self.resolveInst(arg);
-                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.register_manager.getReg(reg, null);
                         try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
@@ -3223,7 +3223,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
 
                         const arg = inst.args[i];
                         const arg_mcv = try self.resolveInst(arg);
-                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.register_manager.getReg(reg, null);
                         try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
@@ -3258,7 +3258,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
 
                         const arg = inst.args[i];
                         const arg_mcv = try self.resolveInst(arg);
-                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.register_manager.getReg(reg, null);
                         try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
@@ -3291,7 +3291,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
 
                         const arg = inst.args[i];
                         const arg_mcv = try self.resolveInst(arg);
-                        try self.register_manager.getRegWithoutTracking(reg);
+                        try self.register_manager.getReg(reg, null);
                         try self.genSetReg(inst.base.src, arg.ty, reg, arg_mcv);
                     }
 
src/register_manager.zig
@@ -7,6 +7,9 @@ const ir = @import("ir.zig");
 const Type = @import("type.zig").Type;
 const Module = @import("Module.zig");
 const LazySrcLoc = Module.LazySrcLoc;
+const expect = std.testing.expect;
+const expectEqual = std.testing.expectEqual;
+const expectEqualSlices = std.testing.expectEqualSlices;
 
 const log = std.log.scoped(.register_manager);
 
@@ -66,77 +69,14 @@ pub fn RegisterManager(
             return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0;
         }
 
-        /// Returns `null` if all registers are allocated.
+        /// Allocates a specified number of registers, optionally
+        /// tracking them. Returns `null` if not enough registers are
+        /// free.
         pub fn tryAllocRegs(
             self: *Self,
             comptime count: comptime_int,
-            insts: [count]*ir.Inst,
-            exceptions: []Register,
-        ) ?[count]Register {
-            if (self.tryAllocRegsWithoutTracking(count, exceptions)) |regs| {
-                for (regs) |reg, i| {
-                    const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null
-                    self.registers[index] = insts[i];
-                    self.markRegUsed(reg);
-                }
-
-                return regs;
-            } else {
-                return null;
-            }
-        }
-
-        /// Returns `null` if all registers are allocated.
-        pub fn tryAllocReg(self: *Self, inst: *ir.Inst, exceptions: []Register) ?Register {
-            return if (tryAllocRegs(self, 1, .{inst}, exceptions)) |regs| regs[0] else null;
-        }
-
-        pub fn allocRegs(
-            self: *Self,
-            comptime count: comptime_int,
-            insts: [count]*ir.Inst,
-            exceptions: []Register,
-        ) ![count]Register {
-            comptime assert(count > 0 and count <= callee_preserved_regs.len);
-            assert(count + exceptions.len <= callee_preserved_regs.len);
-
-            return self.tryAllocRegs(count, insts, exceptions) orelse blk: {
-                // We'll take over the first count registers. Spill
-                // the instructions that were previously there to a
-                // stack allocations.
-                var regs: [count]Register = undefined;
-                var i: usize = 0;
-                for (callee_preserved_regs) |reg| {
-                    if (i >= count) break;
-                    if (mem.indexOfScalar(Register, exceptions, reg) != null) continue;
-                    regs[i] = reg;
-
-                    const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null
-                    if (self.isRegFree(reg)) {
-                        self.markRegUsed(reg);
-                    } else {
-                        const spilled_inst = self.registers[index].?;
-                        try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
-                    }
-                    self.registers[index] = insts[i];
-
-                    i += 1;
-                }
-
-                break :blk regs;
-            };
-        }
-
-        pub fn allocReg(self: *Self, inst: *ir.Inst, exceptions: []Register) !Register {
-            return (try self.allocRegs(1, .{inst}, exceptions))[0];
-        }
-
-        /// Does not track the registers.
-        /// Returns `null` if not enough registers are free.
-        pub fn tryAllocRegsWithoutTracking(
-            self: *Self,
-            comptime count: comptime_int,
-            exceptions: []Register,
+            insts: [count]?*ir.Inst,
+            exceptions: []const Register,
         ) ?[count]Register {
             comptime if (callee_preserved_regs.len == 0) return null;
             comptime assert(count > 0 and count <= callee_preserved_regs.len);
@@ -156,18 +96,40 @@ pub fn RegisterManager(
                 }
             }
 
-            return if (i < count) null else regs;
+            if (i == count) {
+                for (regs) |reg, j| {
+                    if (insts[j]) |inst| {
+                        // Track the register
+                        const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null
+                        self.registers[index] = inst;
+                        self.markRegUsed(reg);
+                    }
+                }
+
+                return regs;
+            } else return null;
         }
 
-        /// Does not track the register.
-        /// Returns `null` if all registers are allocated.
-        pub fn tryAllocRegWithoutTracking(self: *Self, exceptions: []Register) ?Register {
-            return if (self.tryAllocRegsWithoutTracking(1, exceptions)) |regs| regs[0] else null;
+        /// Allocates a register and optionally tracks it with a
+        /// corresponding instruction. Returns `null` if all registers
+        /// are allocated.
+        pub fn tryAllocReg(self: *Self, inst: ?*ir.Inst, exceptions: []const Register) ?Register {
+            return if (tryAllocRegs(self, 1, .{inst}, exceptions)) |regs| regs[0] else null;
         }
 
-        /// Does not track the registers
-        pub fn allocRegsWithoutTracking(self: *Self, comptime count: comptime_int, exceptions: []Register) ![count]Register {
-            return self.tryAllocRegsWithoutTracking(count, exceptions) orelse blk: {
+        /// Allocates a specified number of registers, optionally
+        /// tracking them. Asserts that count + exceptions.len is not
+        /// larger than the total number of registers available.
+        pub fn allocRegs(
+            self: *Self,
+            comptime count: comptime_int,
+            insts: [count]?*ir.Inst,
+            exceptions: []const Register,
+        ) ![count]Register {
+            comptime assert(count > 0 and count <= callee_preserved_regs.len);
+            assert(count + exceptions.len <= callee_preserved_regs.len);
+
+            return self.tryAllocRegs(count, insts, exceptions) orelse blk: {
                 // We'll take over the first count registers. Spill
                 // the instructions that were previously there to a
                 // stack allocations.
@@ -179,11 +141,22 @@ pub fn RegisterManager(
                     regs[i] = reg;
 
                     const index = reg.allocIndex().?; // allocIndex() on a callee-preserved reg should never return null
-                    if (!self.isRegFree(reg)) {
-                        const spilled_inst = self.registers[index].?;
-                        try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
-                        self.registers[index] = null;
-                        self.markRegFree(reg);
+                    if (insts[i]) |inst| {
+                        // Track the register
+                        if (self.isRegFree(reg)) {
+                            self.markRegUsed(reg);
+                        } else {
+                            const spilled_inst = self.registers[index].?;
+                            try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+                        }
+                        self.registers[index] = inst;
+                    } else {
+                        // Don't track the register
+                        if (!self.isRegFree(reg)) {
+                            const spilled_inst = self.registers[index].?;
+                            try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+                            self.freeReg(reg);
+                        }
                     }
 
                     i += 1;
@@ -193,39 +166,36 @@ pub fn RegisterManager(
             };
         }
 
-        /// Does not track the register.
-        pub fn allocRegWithoutTracking(self: *Self, exceptions: []Register) !Register {
-            return (try self.allocRegsWithoutTracking(1, exceptions))[0];
-        }
-
-        /// 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 {
-            const index = reg.allocIndex() orelse return;
-
-            if (!self.isRegFree(reg)) {
-                // Move the instruction that was previously there to a
-                // stack allocation.
-                const spilled_inst = self.registers[index].?;
-                self.registers[index] = inst;
-                try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
-            } else {
-                self.getRegAssumeFree(reg, inst);
-            }
+        /// Allocates a register and optionally tracks it with a
+        /// corresponding instruction.
+        pub fn allocReg(self: *Self, inst: ?*ir.Inst, exceptions: []const Register) !Register {
+            return (try self.allocRegs(1, .{inst}, exceptions))[0];
         }
 
-        /// Spills the register if it is currently allocated.
-        /// Does not track the register.
-        pub fn getRegWithoutTracking(self: *Self, reg: Register) !void {
+        /// Spills the register if it is currently allocated. If a
+        /// corresponding instruction is passed, will also track this
+        /// register.
+        pub fn getReg(self: *Self, reg: Register, inst: ?*ir.Inst) !void {
             const index = reg.allocIndex() orelse return;
 
-            if (!self.isRegFree(reg)) {
-                // Move the instruction that was previously there to a
-                // stack allocation.
-                const spilled_inst = self.registers[index].?;
-                try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
-                self.markRegFree(reg);
+            if (inst) |tracked_inst|
+                if (!self.isRegFree(reg)) {
+                    // Move the instruction that was previously there to a
+                    // stack allocation.
+                    const spilled_inst = self.registers[index].?;
+                    self.registers[index] = tracked_inst;
+                    try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+                } else {
+                    self.getRegAssumeFree(reg, tracked_inst);
+                }
+            else {
+                if (!self.isRegFree(reg)) {
+                    // Move the instruction that was previously there to a
+                    // stack allocation.
+                    const spilled_inst = self.registers[index].?;
+                    try self.getFunction().spillInstruction(spilled_inst.src, reg, spilled_inst);
+                    self.freeReg(reg);
+                }
             }
         }
 
@@ -250,42 +220,63 @@ pub fn RegisterManager(
     };
 }
 
-const MockRegister = enum(u2) {
+const MockRegister1 = enum(u2) {
     r0,
     r1,
     r2,
     r3,
 
-    pub fn allocIndex(self: MockRegister) ?u2 {
-        inline for (mock_callee_preserved_regs) |cpreg, i| {
+    pub fn allocIndex(self: MockRegister1) ?u2 {
+        inline for (callee_preserved_regs) |cpreg, i| {
             if (self == cpreg) return i;
         }
         return null;
     }
-};
-
-const mock_callee_preserved_regs = [_]MockRegister{ .r2, .r3 };
 
-const MockFunction = struct {
-    allocator: *Allocator,
-    register_manager: RegisterManager(Self, MockRegister, &mock_callee_preserved_regs) = .{},
-    spilled: std.ArrayListUnmanaged(MockRegister) = .{},
+    const callee_preserved_regs = [_]MockRegister1{ .r2, .r3 };
+};
 
-    const Self = @This();
+const MockRegister2 = enum(u2) {
+    r0,
+    r1,
+    r2,
+    r3,
 
-    pub fn deinit(self: *Self) void {
-        self.spilled.deinit(self.allocator);
+    pub fn allocIndex(self: MockRegister2) ?u2 {
+        inline for (callee_preserved_regs) |cpreg, i| {
+            if (self == cpreg) return i;
+        }
+        return null;
     }
 
-    pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: MockRegister, inst: *ir.Inst) !void {
-        try self.spilled.append(self.allocator, reg);
-    }
+    const callee_preserved_regs = [_]MockRegister2{ .r0, .r1, .r2, .r3 };
 };
 
-test "tryAllocReg: no spilling" {
+fn MockFunction(comptime Register: type) type {
+    return struct {
+        allocator: *Allocator,
+        register_manager: RegisterManager(Self, Register, &Register.callee_preserved_regs) = .{},
+        spilled: std.ArrayListUnmanaged(Register) = .{},
+
+        const Self = @This();
+
+        pub fn deinit(self: *Self) void {
+            self.spilled.deinit(self.allocator);
+        }
+
+        pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: Register, inst: *ir.Inst) !void {
+            try self.spilled.append(self.allocator, reg);
+        }
+    };
+}
+
+const MockFunction1 = MockFunction(MockRegister1);
+const MockFunction2 = MockFunction(MockRegister2);
+
+test "default state" {
     const allocator = std.testing.allocator;
 
-    var function = MockFunction{
+    var function = MockFunction1{
         .allocator = allocator,
     };
     defer function.deinit();
@@ -296,27 +287,48 @@ test "tryAllocReg: no spilling" {
         .src = .unneeded,
     };
 
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r3));
+    try expect(!function.register_manager.isRegAllocated(.r2));
+    try expect(!function.register_manager.isRegAllocated(.r3));
+    try expect(function.register_manager.isRegFree(.r2));
+    try expect(function.register_manager.isRegFree(.r3));
+}
+
+test "tryAllocReg: no spilling" {
+    const allocator = std.testing.allocator;
+
+    var function = MockFunction1{
+        .allocator = allocator,
+    };
+    defer function.deinit();
+
+    var mock_instruction = ir.Inst{
+        .tag = .breakpoint,
+        .ty = Type.initTag(.void),
+        .src = .unneeded,
+    };
 
-    try std.testing.expectEqual(@as(?MockRegister, .r2), function.register_manager.tryAllocReg(&mock_instruction, &.{}));
-    try std.testing.expectEqual(@as(?MockRegister, .r3), function.register_manager.tryAllocReg(&mock_instruction, &.{}));
-    try std.testing.expectEqual(@as(?MockRegister, null), function.register_manager.tryAllocReg(&mock_instruction, &.{}));
+    try expectEqual(@as(?MockRegister1, .r2), function.register_manager.tryAllocReg(&mock_instruction, &.{}));
+    try expectEqual(@as(?MockRegister1, .r3), function.register_manager.tryAllocReg(&mock_instruction, &.{}));
+    try expectEqual(@as(?MockRegister1, null), function.register_manager.tryAllocReg(&mock_instruction, &.{}));
 
-    try std.testing.expect(function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(function.register_manager.isRegAllocated(.r3));
+    try expect(function.register_manager.isRegAllocated(.r2));
+    try expect(function.register_manager.isRegAllocated(.r3));
+    try expect(!function.register_manager.isRegFree(.r2));
+    try expect(!function.register_manager.isRegFree(.r3));
 
     function.register_manager.freeReg(.r2);
     function.register_manager.freeReg(.r3);
 
-    try std.testing.expect(function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(function.register_manager.isRegAllocated(.r3));
+    try expect(function.register_manager.isRegAllocated(.r2));
+    try expect(function.register_manager.isRegAllocated(.r3));
+    try expect(function.register_manager.isRegFree(.r2));
+    try expect(function.register_manager.isRegFree(.r3));
 }
 
 test "allocReg: spilling" {
     const allocator = std.testing.allocator;
 
-    var function = MockFunction{
+    var function = MockFunction1{
         .allocator = allocator,
     };
     defer function.deinit();
@@ -327,26 +339,28 @@ test "allocReg: spilling" {
         .src = .unneeded,
     };
 
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r3));
-
-    try std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction, &.{}));
-    try std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction, &.{}));
+    try expectEqual(@as(?MockRegister1, .r2), try function.register_manager.allocReg(&mock_instruction, &.{}));
+    try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(&mock_instruction, &.{}));
 
     // Spill a register
-    try std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction, &.{}));
-    try std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items);
+    try expectEqual(@as(?MockRegister1, .r2), try function.register_manager.allocReg(&mock_instruction, &.{}));
+    try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r2}, function.spilled.items);
 
     // No spilling necessary
     function.register_manager.freeReg(.r3);
-    try std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction, &.{}));
-    try std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items);
+    try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(&mock_instruction, &.{}));
+    try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r2}, function.spilled.items);
+
+    // Exceptions
+    function.register_manager.freeReg(.r2);
+    function.register_manager.freeReg(.r3);
+    try expectEqual(@as(?MockRegister1, .r3), try function.register_manager.allocReg(&mock_instruction, &.{.r2}));
 }
 
-test "getReg" {
+test "tryAllocRegs" {
     const allocator = std.testing.allocator;
 
-    var function = MockFunction{
+    var function = MockFunction2{
         .allocator = allocator,
     };
     defer function.deinit();
@@ -357,18 +371,67 @@ test "getReg" {
         .src = .unneeded,
     };
 
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r3));
+    try expectEqual([_]MockRegister2{ .r0, .r1, .r2 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{}).?);
+
+    // Exceptions
+    function.register_manager.freeReg(.r0);
+    function.register_manager.freeReg(.r1);
+    function.register_manager.freeReg(.r2);
+    try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, function.register_manager.tryAllocRegs(3, .{ null, null, null }, &.{.r1}).?);
+}
+
+test "allocRegs" {
+    const allocator = std.testing.allocator;
+
+    var function = MockFunction2{
+        .allocator = allocator,
+    };
+    defer function.deinit();
+
+    var mock_instruction = ir.Inst{
+        .tag = .breakpoint,
+        .ty = Type.initTag(.void),
+        .src = .unneeded,
+    };
+
+    try expectEqual([_]MockRegister2{ .r0, .r1, .r2 }, try function.register_manager.allocRegs(3, .{
+        &mock_instruction,
+        &mock_instruction,
+        &mock_instruction,
+    }, &.{}));
+
+    // Exceptions
+    try expectEqual([_]MockRegister2{ .r0, .r2, .r3 }, try function.register_manager.allocRegs(3, .{ null, null, null }, &.{.r1}));
+    try expectEqualSlices(MockRegister2, &[_]MockRegister2{ .r0, .r2 }, function.spilled.items);
+}
+
+test "getReg" {
+    const allocator = std.testing.allocator;
+
+    var function = MockFunction1{
+        .allocator = allocator,
+    };
+    defer function.deinit();
+
+    var mock_instruction = ir.Inst{
+        .tag = .breakpoint,
+        .ty = Type.initTag(.void),
+        .src = .unneeded,
+    };
 
     try function.register_manager.getReg(.r3, &mock_instruction);
 
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(function.register_manager.isRegAllocated(.r3));
+    try expect(!function.register_manager.isRegAllocated(.r2));
+    try expect(function.register_manager.isRegAllocated(.r3));
+    try expect(function.register_manager.isRegFree(.r2));
+    try expect(!function.register_manager.isRegFree(.r3));
 
     // Spill r3
     try function.register_manager.getReg(.r3, &mock_instruction);
 
-    try std.testing.expect(!function.register_manager.isRegAllocated(.r2));
-    try std.testing.expect(function.register_manager.isRegAllocated(.r3));
-    try std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r3}, function.spilled.items);
+    try expect(!function.register_manager.isRegAllocated(.r2));
+    try expect(function.register_manager.isRegAllocated(.r3));
+    try expect(function.register_manager.isRegFree(.r2));
+    try expect(!function.register_manager.isRegFree(.r3));
+    try expectEqualSlices(MockRegister1, &[_]MockRegister1{.r3}, function.spilled.items);
 }