Commit 228a1ce3e8

@@ -929,7 +929,7 @@ fn Function(comptime arch: std.Target.Cpu.Arch) type {
             return MCValue{ .stack_offset = stack_offset };
         }
 
-        pub fn spillInstruction(self: *Self, src: usize, reg: Register, inst: *ir.Inst) !void {
+        pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: Register, inst: *ir.Inst) !void {
             const stack_mcv = try self.allocRegOrMem(inst, false);
             const reg_mcv = self.getResolvedInstValue(inst);
             assert(reg == toCanonicalReg(reg_mcv.register));

@@ -4,6 +4,9 @@ const assert = std.debug.assert;
 const Allocator = std.mem.Allocator;
 const ir = @import("ir.zig");
 const Type = @import("type.zig").Type;
+const Module = @import("Module.zig");
+const LazySrcLoc = Module.LazySrcLoc;
+
 const log = std.log.scoped(.register_manager);
 
 pub fn RegisterManager(
@@ -22,6 +25,7 @@ pub fn RegisterManager(
 
         /// An integer whose bits represent all the registers and whether they are free.
         const FreeRegInt = std.meta.Int(.unsigned, callee_preserved_regs.len);
+        const ShiftInt = math.Log2Int(FreeRegInt);
 
         fn getFunction(self: *Self) *Function {
             return @fieldParentPtr(Function, "register_manager", self);
@@ -34,7 +38,6 @@ pub fn RegisterManager(
         fn markRegUsed(self: *Self, reg: Register) void {
             if (FreeRegInt == u0) return;
             const index = reg.allocIndex() orelse return;
-            const ShiftInt = math.Log2Int(FreeRegInt);
             const shift = @intCast(ShiftInt, index);
             const mask = @as(FreeRegInt, 1) << shift;
             self.free_registers &= ~mask;
@@ -44,7 +47,6 @@ pub fn RegisterManager(
         fn markRegFree(self: *Self, reg: Register) void {
             if (FreeRegInt == u0) return;
             const index = reg.allocIndex() orelse return;
-            const ShiftInt = math.Log2Int(FreeRegInt);
             const shift = @intCast(ShiftInt, index);
             self.free_registers |= @as(FreeRegInt, 1) << shift;
         }
@@ -54,9 +56,8 @@ pub fn RegisterManager(
         pub fn isRegAllocated(self: Self, reg: Register) bool {
             if (FreeRegInt == u0) return false;
             const index = reg.allocIndex() orelse return false;
-            const ShiftInt = math.Log2Int(FreeRegInt);
             const shift = @intCast(ShiftInt, index);
-            return self.free_registers & @as(FreeRegInt, 1) << shift != 0;
+            return self.allocated_registers & @as(FreeRegInt, 1) << shift != 0;
         }
 
         /// Before calling, must ensureCapacity + 1 on self.registers.
@@ -66,9 +67,16 @@ pub fn RegisterManager(
             if (free_index >= callee_preserved_regs.len) {
                 return null;
             }
-            const mask = @as(FreeRegInt, 1) << free_index;
+
+            // This is necessary because the return type of @ctz is 1
+            // bit longer than ShiftInt if callee_preserved_regs.len
+            // is a power of two. This int cast is always safe because
+            // free_index < callee_preserved_regs.len
+            const shift = @intCast(ShiftInt, free_index);
+            const mask = @as(FreeRegInt, 1) << shift;
             self.free_registers &= ~mask;
             self.allocated_registers |= mask;
+
             const reg = callee_preserved_regs[free_index];
             self.registers.putAssumeCapacityNoClobber(reg, inst);
             log.debug("alloc {} => {*}", .{ reg, inst });
@@ -125,3 +133,96 @@ pub fn RegisterManager(
         }
     };
 }
+
+const MockRegister = enum(u2) {
+    r0, r1, r2, r3,
+
+    pub fn allocIndex(self: MockRegister) ?u2 {
+        inline for (mock_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 Self = @This();
+
+    pub fn deinit(self: *Self) void {
+        self.register_manager.deinit(self.allocator);
+        self.spilled.deinit(self.allocator);
+    }
+    
+    pub fn spillInstruction(self: *Self, src: LazySrcLoc, reg: MockRegister, inst: *ir.Inst) !void {
+        try self.spilled.append(self.allocator, reg);
+    }
+};
+
+test "tryAllocReg: no spilling" {
+    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);
+    std.testing.expectEqual(@as(?MockRegister, .r2), function.register_manager.tryAllocReg(&mock_instruction));
+    std.testing.expectEqual(@as(?MockRegister, .r3), function.register_manager.tryAllocReg(&mock_instruction));
+    std.testing.expectEqual(@as(?MockRegister, null), function.register_manager.tryAllocReg(&mock_instruction));
+
+    std.testing.expect(function.register_manager.isRegAllocated(.r2));
+    std.testing.expect(function.register_manager.isRegAllocated(.r3));
+
+    function.register_manager.freeReg(.r2);
+    function.register_manager.freeReg(.r3);
+
+    std.testing.expect(function.register_manager.isRegAllocated(.r2));
+    std.testing.expect(function.register_manager.isRegAllocated(.r3));
+}
+
+test "allocReg: spilling" {
+    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);
+    std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction));
+    std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction));
+
+    // Spill a register
+    std.testing.expectEqual(@as(?MockRegister, .r2), try function.register_manager.allocReg(&mock_instruction));
+    std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items);
+
+    // No spilling necessary
+    function.register_manager.freeReg(.r3);
+    std.testing.expectEqual(@as(?MockRegister, .r3), try function.register_manager.allocReg(&mock_instruction));
+    std.testing.expectEqualSlices(MockRegister, &[_]MockRegister{.r2}, function.spilled.items);
+}