stage2 AArch64: implement cond_br for other MCValues

3 files changed, 74 insertions, 14 deletions

diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig
index 7c207039b9..fade038eb0 100644
--- a/src/arch/aarch64/CodeGen.zig
+++ b/src/arch/aarch64/CodeGen.zig
@@ -923,12 +923,12 @@ fn airNot(self: *Self, inst: Air.Inst.Index) !void {
                 };
                 break :result r;
             },
-            else => {},
+            else => {
+                return self.fail("TODO implement NOT for {}", .{self.target.cpu.arch});
+            },
         }
-
-        return self.fail("TODO implement NOT for {}", .{self.target.cpu.arch});
     };
-    _ = result;
+    return self.finishAir(inst, result, .{ ty_op.operand, .none, .none });
 }
 
 fn airMin(self: *Self, inst: Air.Inst.Index) !void {
@@ -1411,7 +1411,7 @@ fn airSliceLen(self: *Self, inst: Air.Inst.Index) !void {
             .dead, .unreach => unreachable,
             .register => unreachable, // a slice doesn't fit in one register
             .stack_offset => |off| {
-                break :result MCValue{ .stack_offset = off + 8 };
+                break :result MCValue{ .stack_offset = off };
             },
             .memory => |addr| {
                 break :result MCValue{ .memory = addr + 8 };
@@ -2425,7 +2425,22 @@ fn airCondBr(self: *Self, inst: Air.Inst.Index) !void {
                 },
             },
         }),
-        else => return self.fail("TODO implement condbr when condition is {s}", .{@tagName(cond)}),
+        else => blk: {
+            const reg = switch (cond) {
+                .register => |r| r,
+                else => try self.copyToTmpRegister(Type.bool, cond),
+            };
+
+            break :blk try self.addInst(.{
+                .tag = .cbz,
+                .data = .{
+                    .r_inst = .{
+                        .rt = reg,
+                        .inst = undefined, // populated later through performReloc
+                    },
+                },
+            });
+        },
     };
 
     // Capture the state of register and stack allocation state so that we can revert to it.
@@ -2770,8 +2785,9 @@ fn airSwitch(self: *Self, inst: Air.Inst.Index) !void {
 fn performReloc(self: *Self, inst: Mir.Inst.Index) !void {
     const tag = self.mir_instructions.items(.tag)[inst];
     switch (tag) {
-        .b_cond => self.mir_instructions.items(.data)[inst].inst_cond.inst = @intCast(Air.Inst.Index, self.mir_instructions.len),
-        .b => self.mir_instructions.items(.data)[inst].inst = @intCast(Air.Inst.Index, self.mir_instructions.len),
+        .cbz => self.mir_instructions.items(.data)[inst].r_inst.inst = @intCast(Mir.Inst.Index, self.mir_instructions.len),
+        .b_cond => self.mir_instructions.items(.data)[inst].inst_cond.inst = @intCast(Mir.Inst.Index, self.mir_instructions.len),
+        .b => self.mir_instructions.items(.data)[inst].inst = @intCast(Mir.Inst.Index, self.mir_instructions.len),
         else => unreachable,
     }
 }
diff --git a/src/arch/aarch64/Emit.zig b/src/arch/aarch64/Emit.zig
index 1b30c78562..1772f08aa8 100644
--- a/src/arch/aarch64/Emit.zig
+++ b/src/arch/aarch64/Emit.zig
@@ -50,11 +50,13 @@ const InnerError = error{
 };
 
 const BranchType = enum {
+    cbz,
     b_cond,
     unconditional_branch_immediate,
 
     fn default(tag: Mir.Inst.Tag) BranchType {
         return switch (tag) {
+            .cbz => .cbz,
             .b, .bl => .unconditional_branch_immediate,
             .b_cond => .b_cond,
             else => unreachable,
@@ -83,6 +85,8 @@ pub fn emitMir(
             .b => try emit.mirBranch(inst),
             .bl => try emit.mirBranch(inst),
 
+            .cbz => try emit.mirCompareAndBranch(inst),
+
             .blr => try emit.mirUnconditionalBranchRegister(inst),
             .ret => try emit.mirUnconditionalBranchRegister(inst),
 
@@ -160,15 +164,22 @@ fn optimalBranchType(emit: *Emit, tag: Mir.Inst.Tag, offset: i64) !BranchType {
     assert(offset & 0b11 == 0);
 
     switch (tag) {
+        .cbz => {
+            if (std.math.cast(i19, @shrExact(offset, 2))) |_| {
+                return BranchType.cbz;
+            } else |_| {
+                return emit.fail("TODO support cbz branches larger than +-1 MiB", .{});
+            }
+        },
         .b, .bl => {
-            if (std.math.cast(i26, offset >> 2)) |_| {
+            if (std.math.cast(i26, @shrExact(offset, 2))) |_| {
                 return BranchType.unconditional_branch_immediate;
             } else |_| {
-                return emit.fail("TODO support branches larger than +-128 MiB", .{});
+                return emit.fail("TODO support unconditional branches larger than +-128 MiB", .{});
             }
         },
         .b_cond => {
-            if (std.math.cast(i19, offset >> 2)) |_| {
+            if (std.math.cast(i19, @shrExact(offset, 2))) |_| {
                 return BranchType.b_cond;
             } else |_| {
                 return emit.fail("TODO support conditional branches larger than +-1 MiB", .{});
@@ -183,8 +194,10 @@ fn instructionSize(emit: *Emit, inst: Mir.Inst.Index) usize {
 
     if (isBranch(tag)) {
         switch (emit.branch_types.get(inst).?) {
-            .unconditional_branch_immediate => return 4,
-            .b_cond => return 4,
+            .cbz,
+            .unconditional_branch_immediate,
+            .b_cond,
+            => return 4,
         }
     }
 
@@ -222,7 +235,11 @@ fn instructionSize(emit: *Emit, inst: Mir.Inst.Index) usize {
 
 fn isBranch(tag: Mir.Inst.Tag) bool {
     return switch (tag) {
-        .b, .bl, .b_cond => true,
+        .cbz,
+        .b,
+        .bl,
+        .b_cond,
+        => true,
         else => false,
     };
 }
@@ -231,6 +248,7 @@ fn branchTarget(emit: *Emit, inst: Mir.Inst.Index) Mir.Inst.Index {
     const tag = emit.mir.instructions.items(.tag)[inst];
 
     switch (tag) {
+        .cbz => return emit.mir.instructions.items(.data)[inst].r_inst.inst,
         .b, .bl => return emit.mir.instructions.items(.data)[inst].inst,
         .b_cond => return emit.mir.instructions.items(.data)[inst].inst_cond.inst,
         else => unreachable,
@@ -494,6 +512,23 @@ fn mirBranch(emit: *Emit, inst: Mir.Inst.Index) !void {
     }
 }
 
+fn mirCompareAndBranch(emit: *Emit, inst: Mir.Inst.Index) !void {
+    const tag = emit.mir.instructions.items(.tag)[inst];
+    const r_inst = emit.mir.instructions.items(.data)[inst].r_inst;
+
+    const offset = @intCast(i64, emit.code_offset_mapping.get(r_inst.inst).?) - @intCast(i64, emit.code.items.len);
+    const branch_type = emit.branch_types.get(inst).?;
+    log.debug("mirCompareAndBranch: {} offset={}", .{ inst, offset });
+
+    switch (branch_type) {
+        .cbz => switch (tag) {
+            .cbz => try emit.writeInstruction(Instruction.cbz(r_inst.rt, @intCast(i21, offset))),
+            else => unreachable,
+        },
+        else => unreachable,
+    }
+}
+
 fn mirUnconditionalBranchRegister(emit: *Emit, inst: Mir.Inst.Index) !void {
     const tag = emit.mir.instructions.items(.tag)[inst];
     const reg = emit.mir.instructions.items(.data)[inst].reg;
diff --git a/src/arch/aarch64/Mir.zig b/src/arch/aarch64/Mir.zig
index d2263c6e32..65b80549a9 100644
--- a/src/arch/aarch64/Mir.zig
+++ b/src/arch/aarch64/Mir.zig
@@ -40,6 +40,8 @@ pub const Inst = struct {
         brk,
         /// Pseudo-instruction: Call extern
         call_extern,
+        /// Compare and Branch on Zero
+        cbz,
         /// Compare (immediate)
         cmp_immediate,
         /// Compare (shifted register)
@@ -184,6 +186,13 @@ pub const Inst = struct {
             rd: Register,
             cond: bits.Instruction.Condition,
         },
+        /// A register and another instruction
+        ///
+        /// Used by e.g. cbz
+        r_inst: struct {
+            rt: Register,
+            inst: Index,
+        },
         /// Two registers
         ///
         /// Used by e.g. mov_register