From 43a627927f989034ae64846abea73bc8bdb545be Mon Sep 17 00:00:00 2001
From: Jakub Konka <kubkon@jakubkonka.com>
Date: Sat, 7 May 2022 10:31:08 +0200
Subject: [PATCH] x64: fix misused register locks

---
 src/arch/x86_64/CodeGen.zig | 45 ++++++++++++++++---------------------
 src/register_manager.zig    |  8 ++++++-
 2 files changed, 26 insertions(+), 27 deletions(-)

diff --git a/src/arch/x86_64/CodeGen.zig b/src/arch/x86_64/CodeGen.zig
index 8bb7111142..7df315d7e1 100644
--- a/src/arch/x86_64/CodeGen.zig
+++ b/src/arch/x86_64/CodeGen.zig
@@ -211,6 +211,16 @@ pub const MCValue = union(enum) {
             else => false,
         };
     }
+
+    fn asRegister(mcv: MCValue) ?Register {
+        return switch (mcv) {
+            .register,
+            .register_overflow_unsigned,
+            .register_overflow_signed,
+            => |reg| reg,
+            else => null,
+        };
+    }
 };
 
 const Branch = struct {
@@ -842,20 +852,15 @@ fn finishAir(self: *Self, inst: Air.Inst.Index, result: MCValue, operands: [Live
         const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
         branch.inst_table.putAssumeCapacityNoClobber(inst, result);
 
-        switch (result) {
-            .register,
-            .register_overflow_signed,
-            .register_overflow_unsigned,
-            => |reg| {
-                // In some cases (such as bitcast), an operand
-                // may be the same MCValue as the result. If
-                // that operand died and was a register, it
-                // was freed by processDeath. We have to
-                // "re-allocate" the register.
-                if (self.register_manager.isRegFree(reg)) {
-                    self.register_manager.getRegAssumeFree(reg, inst);
-                }
-            },
+        if (result.asRegister()) |reg| {
+            // In some cases (such as bitcast), an operand
+            // may be the same MCValue as the result. If
+            // that operand died and was a register, it
+            // was freed by processDeath. We have to
+            // "re-allocate" the register.
+            if (self.register_manager.isRegFree(reg)) {
+                self.register_manager.getRegAssumeFree(reg, inst);
+            }
         }
     }
     self.finishAirBookkeeping();
@@ -4011,12 +4016,6 @@ fn airRet(self: *Self, inst: Air.Inst.Index) !void {
     const ret_ty = self.fn_type.fnReturnType();
     switch (self.ret_mcv) {
         .stack_offset => {
-            var reg_locks: [2]RegisterLock = undefined;
-            self.register_manager.freezeRegsAssumeUnused(2, .{ .rax, .rcx }, &reg_locks);
-            defer for (reg_locks) |reg| {
-                self.register_manager.unfreezeReg(reg);
-            };
-
             const reg = try self.copyToTmpRegister(Type.usize, self.ret_mcv);
             const reg_lock = self.register_manager.freezeRegAssumeUnused(reg);
             defer self.register_manager.unfreezeReg(reg_lock);
@@ -4051,12 +4050,6 @@ fn airRetLoad(self: *Self, inst: Air.Inst.Index) !void {
     const elem_ty = ptr_ty.elemType();
     switch (self.ret_mcv) {
         .stack_offset => {
-            var reg_locks: [2]RegisterLock = undefined;
-            self.register_manager.freezeRegsAssumeUnused(2, .{ .rax, .rcx }, &reg_locks);
-            defer for (reg_locks) |reg| {
-                self.register_manager.unfreezeReg(reg);
-            };
-
             const reg = try self.copyToTmpRegister(Type.usize, self.ret_mcv);
             const reg_lock = self.register_manager.freezeRegAssumeUnused(reg);
             defer self.register_manager.unfreezeReg(reg_lock);
diff --git a/src/register_manager.zig b/src/register_manager.zig
index 7e96d87af2..6e73181f2a 100644
--- a/src/register_manager.zig
+++ b/src/register_manager.zig
@@ -127,7 +127,11 @@ pub fn RegisterManager(
         /// Only the owner of the `RegisterLock` can unfreeze the
         /// register later.
         pub fn freezeReg(self: *Self, reg: Register) ?RegisterLock {
-            if (self.isRegFrozen(reg)) return null;
+            log.debug("freezing {}", .{reg});
+            if (self.isRegFrozen(reg)) {
+                log.debug("  register already locked", .{});
+                return null;
+            }
             const mask = getRegisterMask(reg) orelse return null;
             self.frozen_registers |= mask;
             return RegisterLock{ .register = reg };
@@ -136,6 +140,7 @@ pub fn RegisterManager(
         /// Like `freezeReg` but asserts the register was unused always
         /// returning a valid lock.
         pub fn freezeRegAssumeUnused(self: *Self, reg: Register) RegisterLock {
+            log.debug("freezing asserting free {}", .{reg});
             assert(!self.isRegFrozen(reg));
             const mask = getRegisterMask(reg) orelse unreachable;
             self.frozen_registers |= mask;
@@ -158,6 +163,7 @@ pub fn RegisterManager(
         /// Requires `RegisterLock` to unfreeze a register.
         /// Call `freezeReg` to obtain the lock first.
         pub fn unfreezeReg(self: *Self, lock: RegisterLock) void {
+            log.debug("unfreezing {}", .{lock.register});
             const mask = getRegisterMask(lock.register) orelse return;
             self.frozen_registers &= ~mask;
         }