stage2 AArch64: Fix issue in binOp and add regression test

2 files changed, 61 insertions, 19 deletions

diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig
index 8bd5324ed1..8b5503f293 100644
--- a/src/arch/aarch64/CodeGen.zig
+++ b/src/arch/aarch64/CodeGen.zig
@@ -981,13 +981,19 @@ fn binOpRegister(
     if (lhs_is_register) self.register_manager.freezeRegs(&.{lhs.register});
     if (rhs_is_register) self.register_manager.freezeRegs(&.{rhs.register});
 
+    const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
+
     const lhs_reg = if (lhs_is_register) lhs.register else blk: {
         const track_inst: ?Air.Inst.Index = if (maybe_inst) |inst| inst: {
             const bin_op = self.air.instructions.items(.data)[inst].bin_op;
             break :inst Air.refToIndex(bin_op.lhs).?;
         } else null;
+
         const reg = try self.register_manager.allocReg(track_inst);
         self.register_manager.freezeRegs(&.{reg});
+
+        if (track_inst) |inst| branch.inst_table.putAssumeCapacity(inst, .{ .register = reg });
+
         break :blk reg;
     };
     defer self.register_manager.unfreezeRegs(&.{lhs_reg});
@@ -997,8 +1003,12 @@ fn binOpRegister(
             const bin_op = self.air.instructions.items(.data)[inst].bin_op;
             break :inst Air.refToIndex(bin_op.rhs).?;
         } else null;
+
         const reg = try self.register_manager.allocReg(track_inst);
         self.register_manager.freezeRegs(&.{reg});
+
+        if (track_inst) |inst| branch.inst_table.putAssumeCapacity(inst, .{ .register = reg });
+
         break :blk reg;
     };
     defer self.register_manager.unfreezeRegs(&.{rhs_reg});
@@ -1077,6 +1087,8 @@ fn binOpImmediate(
 
     if (lhs_is_register) self.register_manager.freezeRegs(&.{lhs.register});
 
+    const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
+
     const lhs_reg = if (lhs_is_register) lhs.register else blk: {
         const track_inst: ?Air.Inst.Index = if (maybe_inst) |inst| inst: {
             const bin_op = self.air.instructions.items(.data)[inst].bin_op;
@@ -1084,8 +1096,12 @@ fn binOpImmediate(
                 if (lhs_and_rhs_swapped) bin_op.rhs else bin_op.lhs,
             ).?;
         } else null;
+
         const reg = try self.register_manager.allocReg(track_inst);
         self.register_manager.freezeRegs(&.{reg});
+
+        if (track_inst) |inst| branch.inst_table.putAssumeCapacity(inst, .{ .register = reg });
+
         break :blk reg;
     };
     defer self.register_manager.unfreezeRegs(&.{lhs_reg});
@@ -3141,7 +3157,7 @@ fn genSetReg(self: *Self, ty: Type, reg: Register, mcv: MCValue) InnerError!void
                 .tag = .cset,
                 .data = .{ .r_cond = .{
                     .rd = reg,
-                    .cond = condition,
+                    .cond = condition.negate(),
                 } },
             });
         },
diff --git a/test/stage2/aarch64.zig b/test/stage2/aarch64.zig
index 580a375aef..b16a29f56f 100644
--- a/test/stage2/aarch64.zig
+++ b/test/stage2/aarch64.zig
@@ -17,15 +17,8 @@ pub fn addCases(ctx: *TestContext) !void {
         var case = ctx.exe("linux_aarch64 hello world", linux_aarch64);
         // Regular old hello world
         case.addCompareOutput(
-            \\pub export fn _start() noreturn {
+            \\pub fn main() void {
             \\    print();
-            \\    exit();
-            \\}
-            \\
-            \\fn doNothing() void {}
-            \\
-            \\fn answer() u64 {
-            \\    return 0x1234abcd1234abcd;
             \\}
             \\
             \\fn print() void {
@@ -38,16 +31,6 @@ pub fn addCases(ctx: *TestContext) !void {
             \\        : "memory", "cc"
             \\    );
             \\}
-            \\
-            \\fn exit() noreturn {
-            \\    asm volatile ("svc #0"
-            \\        :
-            \\        : [number] "{x8}" (93),
-            \\          [arg1] "{x0}" (0)
-            \\        : "memory", "cc"
-            \\    );
-            \\    unreachable;
-            \\}
         ,
             "Hello, World!\n",
         );
@@ -129,6 +112,49 @@ pub fn addCases(ctx: *TestContext) !void {
         );
     }
 
+    {
+        var case = ctx.exe("large add function", linux_aarch64);
+
+        case.addCompareOutput(
+            \\pub fn main() void {
+            \\    assert(add(3, 4) == 791);
+            \\}
+            \\
+            \\fn add(a: u32, b: u32) u32 {
+            \\    const x: u32 = blk: {
+            \\        const c = a + b; // 7
+            \\        const d = a + c; // 10
+            \\        const e = d + b; // 14
+            \\        const f = d + e; // 24
+            \\        const g = e + f; // 38
+            \\        const h = f + g; // 62
+            \\        const i = g + h; // 100
+            \\        const j = i + d; // 110
+            \\        const k = i + j; // 210
+            \\        const l = k + c; // 217
+            \\        const m = l + d; // 227
+            \\        const n = m + e; // 241
+            \\        const o = n + f; // 265
+            \\        const p = o + g; // 303
+            \\        const q = p + h; // 365
+            \\        const r = q + i; // 465
+            \\        const s = r + j; // 575
+            \\        const t = s + k; // 785
+            \\        break :blk t;
+            \\    };
+            \\    const y = x + a; // 788
+            \\    const z = y + a; // 791
+            \\    return z;
+            \\}
+            \\
+            \\fn assert(ok: bool) void {
+            \\    if (!ok) unreachable;
+            \\}
+        ,
+            "",
+        );
+    }
+
     // macOS tests
     {
         var case = ctx.exe("hello world with updates", macos_aarch64);