Merge pull request #10129 from joachimschmidt557/stage2-aarch64

stage2 AArch64: Implement conditional branches

4 files changed, 673 insertions, 96 deletions

diff --git a/src/arch/aarch64/CodeGen.zig b/src/arch/aarch64/CodeGen.zig
index d3c72fe1b1..994b9a9939 100644
--- a/src/arch/aarch64/CodeGen.zig
+++ b/src/arch/aarch64/CodeGen.zig
@@ -177,7 +177,7 @@ const StackAllocation = struct {
 };
 
 const BlockData = struct {
-    relocs: std.ArrayListUnmanaged(Reloc),
+    relocs: std.ArrayListUnmanaged(Mir.Inst.Index),
     /// The first break instruction encounters `null` here and chooses a
     /// machine code value for the block result, populating this field.
     /// Following break instructions encounter that value and use it for
@@ -185,18 +185,6 @@ const BlockData = struct {
     mcv: MCValue,
 };
 
-const Reloc = union(enum) {
-    /// The value is an offset into the `Function` `code` from the beginning.
-    /// To perform the reloc, write 32-bit signed little-endian integer
-    /// which is a relative jump, based on the address following the reloc.
-    rel32: usize,
-    /// A branch in the ARM instruction set
-    arm_branch: struct {
-        pos: usize,
-        cond: @import("../arm/bits.zig").Condition,
-    },
-};
-
 const BigTomb = struct {
     function: *Self,
     inst: Air.Inst.Index,
@@ -426,6 +414,12 @@ fn gen(self: *Self) !void {
             });
         }
 
+        // add sp, sp, #stack_size
+        _ = try self.addInst(.{
+            .tag = .add_immediate,
+            .data = .{ .rr_imm12_sh = .{ .rd = .xzr, .rn = .xzr, .imm12 = @intCast(u12, aligned_stack_end) } },
+        });
+
         // ldp fp, lr, [sp], #16
         _ = try self.addInst(.{
             .tag = .ldp,
@@ -437,12 +431,6 @@ fn gen(self: *Self) !void {
             } },
         });
 
-        // add sp, sp, #stack_size
-        _ = try self.addInst(.{
-            .tag = .add_immediate,
-            .data = .{ .rr_imm12_sh = .{ .rd = .xzr, .rn = .xzr, .imm12 = @intCast(u12, aligned_stack_end) } },
-        });
-
         // ret lr
         _ = try self.addInst(.{
             .tag = .ret,
@@ -1358,7 +1346,9 @@ fn airArg(self: *Self, inst: Air.Inst.Index) !void {
             const stack_offset = try self.allocMem(inst, abi_size, abi_align);
             try self.genSetStack(ty, stack_offset, MCValue{ .register = reg });
 
-            break :blk MCValue{ .stack_offset = stack_offset };
+            // TODO correct loading and storing from memory
+            // break :blk MCValue{ .stack_offset = stack_offset };
+            break :blk result;
         },
         else => result,
     };
@@ -1634,8 +1624,6 @@ fn airRetLoad(self: *Self, inst: Air.Inst.Index) !void {
 }
 
 fn airCmp(self: *Self, inst: Air.Inst.Index, op: math.CompareOperator) !void {
-    _ = op;
-
     const bin_op = self.air.instructions.items(.data)[inst].bin_op;
     if (self.liveness.isUnused(inst))
         return self.finishAir(inst, .dead, .{ bin_op.lhs, bin_op.rhs, .none });
@@ -1646,10 +1634,79 @@ fn airCmp(self: *Self, inst: Air.Inst.Index, op: math.CompareOperator) !void {
 
     const lhs = try self.resolveInst(bin_op.lhs);
     const rhs = try self.resolveInst(bin_op.rhs);
-    _ = lhs;
-    _ = rhs;
+    const result: MCValue = result: {
+        const lhs_is_register = lhs == .register;
+        const rhs_is_register = rhs == .register;
+        // lhs should always be a register
+        const rhs_should_be_register = switch (rhs) {
+            .immediate => |imm| imm < 0 or imm > std.math.maxInt(u12),
+            else => true,
+        };
+
+        var lhs_mcv = lhs;
+        var rhs_mcv = rhs;
+
+        // Allocate registers
+        if (rhs_should_be_register) {
+            if (!lhs_is_register and !rhs_is_register) {
+                const regs = try self.register_manager.allocRegs(2, .{
+                    Air.refToIndex(bin_op.rhs).?, Air.refToIndex(bin_op.lhs).?,
+                }, &.{});
+                lhs_mcv = MCValue{ .register = regs[0] };
+                rhs_mcv = MCValue{ .register = regs[1] };
+            } else if (!rhs_is_register) {
+                rhs_mcv = MCValue{ .register = try self.register_manager.allocReg(Air.refToIndex(bin_op.rhs).?, &.{}) };
+            }
+        }
+        if (!lhs_is_register) {
+            lhs_mcv = MCValue{ .register = try self.register_manager.allocReg(Air.refToIndex(bin_op.lhs).?, &.{}) };
+        }
+
+        // Move the operands to the newly allocated registers
+        const branch = &self.branch_stack.items[self.branch_stack.items.len - 1];
+        if (lhs_mcv == .register and !lhs_is_register) {
+            try self.genSetReg(ty, lhs_mcv.register, lhs);
+            branch.inst_table.putAssumeCapacity(Air.refToIndex(bin_op.lhs).?, lhs);
+        }
+        if (rhs_mcv == .register and !rhs_is_register) {
+            try self.genSetReg(ty, rhs_mcv.register, rhs);
+            branch.inst_table.putAssumeCapacity(Air.refToIndex(bin_op.rhs).?, rhs);
+        }
+
+        // The destination register is not present in the cmp instruction
+        // The signedness of the integer does not matter for the cmp instruction
+        switch (rhs_mcv) {
+            .register => |reg| {
+                _ = try self.addInst(.{
+                    .tag = .cmp_shifted_register,
+                    .data = .{ .rrr_imm6_shift = .{
+                        .rd = .xzr,
+                        .rn = lhs_mcv.register,
+                        .rm = reg,
+                        .imm6 = 0,
+                        .shift = .lsl,
+                    } },
+                });
+            },
+            .immediate => |imm| {
+                _ = try self.addInst(.{
+                    .tag = .cmp_immediate,
+                    .data = .{ .rr_imm12_sh = .{
+                        .rd = .xzr,
+                        .rn = lhs_mcv.register,
+                        .imm12 = @intCast(u12, imm),
+                    } },
+                });
+            },
+            else => unreachable,
+        }
 
-    return self.fail("TODO implement cmp for {}", .{self.target.cpu.arch});
+        break :result switch (ty.isSignedInt()) {
+            true => MCValue{ .compare_flags_signed = op },
+            false => MCValue{ .compare_flags_unsigned = op },
+        };
+    };
+    return self.finishAir(inst, result, .{ bin_op.lhs, bin_op.rhs, .none });
 }
 
 fn airDbgStmt(self: *Self, inst: Air.Inst.Index) !void {
@@ -1667,9 +1724,153 @@ fn airDbgStmt(self: *Self, inst: Air.Inst.Index) !void {
 }
 
 fn airCondBr(self: *Self, inst: Air.Inst.Index) !void {
-    _ = inst;
+    const pl_op = self.air.instructions.items(.data)[inst].pl_op;
+    const cond = try self.resolveInst(pl_op.operand);
+    const extra = self.air.extraData(Air.CondBr, pl_op.payload);
+    const then_body = self.air.extra[extra.end..][0..extra.data.then_body_len];
+    const else_body = self.air.extra[extra.end + then_body.len ..][0..extra.data.else_body_len];
+    const liveness_condbr = self.liveness.getCondBr(inst);
+
+    const reloc: Mir.Inst.Index = switch (cond) {
+        .compare_flags_signed,
+        .compare_flags_unsigned,
+        => try self.addInst(.{
+            .tag = .b_cond,
+            .data = .{
+                .inst_cond = .{
+                    .inst = undefined, // populated later through performReloc
+                    .cond = switch (cond) {
+                        .compare_flags_signed => |cmp_op| blk: {
+                            // Here we map to the opposite condition because the jump is to the false branch.
+                            const condition = Instruction.Condition.fromCompareOperatorSigned(cmp_op);
+                            break :blk condition.negate();
+                        },
+                        .compare_flags_unsigned => |cmp_op| blk: {
+                            // Here we map to the opposite condition because the jump is to the false branch.
+                            const condition = Instruction.Condition.fromCompareOperatorUnsigned(cmp_op);
+                            break :blk condition.negate();
+                        },
+                        else => unreachable,
+                    },
+                },
+            },
+        }),
+        else => return self.fail("TODO implement condr when condition is {s}", .{@tagName(cond)}),
+    };
+
+    // Capture the state of register and stack allocation state so that we can revert to it.
+    const parent_next_stack_offset = self.next_stack_offset;
+    const parent_free_registers = self.register_manager.free_registers;
+    var parent_stack = try self.stack.clone(self.gpa);
+    defer parent_stack.deinit(self.gpa);
+    const parent_registers = self.register_manager.registers;
+
+    try self.branch_stack.append(.{});
+
+    try self.ensureProcessDeathCapacity(liveness_condbr.then_deaths.len);
+    for (liveness_condbr.then_deaths) |operand| {
+        self.processDeath(operand);
+    }
+    try self.genBody(then_body);
+
+    // Revert to the previous register and stack allocation state.
+
+    var saved_then_branch = self.branch_stack.pop();
+    defer saved_then_branch.deinit(self.gpa);
+
+    self.register_manager.registers = parent_registers;
+
+    self.stack.deinit(self.gpa);
+    self.stack = parent_stack;
+    parent_stack = .{};
+
+    self.next_stack_offset = parent_next_stack_offset;
+    self.register_manager.free_registers = parent_free_registers;
+
+    try self.performReloc(reloc);
+    const else_branch = self.branch_stack.addOneAssumeCapacity();
+    else_branch.* = .{};
+
+    try self.ensureProcessDeathCapacity(liveness_condbr.else_deaths.len);
+    for (liveness_condbr.else_deaths) |operand| {
+        self.processDeath(operand);
+    }
+    try self.genBody(else_body);
+
+    // At this point, each branch will possibly have conflicting values for where
+    // each instruction is stored. They agree, however, on which instructions are alive/dead.
+    // We use the first ("then") branch as canonical, and here emit
+    // instructions into the second ("else") branch to make it conform.
+    // We continue respect the data structure semantic guarantees of the else_branch so
+    // that we can use all the code emitting abstractions. This is why at the bottom we
+    // assert that parent_branch.free_registers equals the saved_then_branch.free_registers
+    // rather than assigning it.
+    const parent_branch = &self.branch_stack.items[self.branch_stack.items.len - 2];
+    try parent_branch.inst_table.ensureUnusedCapacity(self.gpa, else_branch.inst_table.count());
+
+    const else_slice = else_branch.inst_table.entries.slice();
+    const else_keys = else_slice.items(.key);
+    const else_values = else_slice.items(.value);
+    for (else_keys) |else_key, else_idx| {
+        const else_value = else_values[else_idx];
+        const canon_mcv = if (saved_then_branch.inst_table.fetchSwapRemove(else_key)) |then_entry| blk: {
+            // The instruction's MCValue is overridden in both branches.
+            parent_branch.inst_table.putAssumeCapacity(else_key, then_entry.value);
+            if (else_value == .dead) {
+                assert(then_entry.value == .dead);
+                continue;
+            }
+            break :blk then_entry.value;
+        } else blk: {
+            if (else_value == .dead)
+                continue;
+            // The instruction is only overridden in the else branch.
+            var i: usize = self.branch_stack.items.len - 2;
+            while (true) {
+                i -= 1; // If this overflows, the question is: why wasn't the instruction marked dead?
+                if (self.branch_stack.items[i].inst_table.get(else_key)) |mcv| {
+                    assert(mcv != .dead);
+                    break :blk mcv;
+                }
+            }
+        };
+        log.debug("consolidating else_entry {d} {}=>{}", .{ else_key, else_value, canon_mcv });
+        // TODO make sure the destination stack offset / register does not already have something
+        // going on there.
+        try self.setRegOrMem(self.air.typeOfIndex(else_key), canon_mcv, else_value);
+        // TODO track the new register / stack allocation
+    }
+    try parent_branch.inst_table.ensureUnusedCapacity(self.gpa, saved_then_branch.inst_table.count());
+    const then_slice = saved_then_branch.inst_table.entries.slice();
+    const then_keys = then_slice.items(.key);
+    const then_values = then_slice.items(.value);
+    for (then_keys) |then_key, then_idx| {
+        const then_value = then_values[then_idx];
+        // We already deleted the items from this table that matched the else_branch.
+        // So these are all instructions that are only overridden in the then branch.
+        parent_branch.inst_table.putAssumeCapacity(then_key, then_value);
+        if (then_value == .dead)
+            continue;
+        const parent_mcv = blk: {
+            var i: usize = self.branch_stack.items.len - 2;
+            while (true) {
+                i -= 1;
+                if (self.branch_stack.items[i].inst_table.get(then_key)) |mcv| {
+                    assert(mcv != .dead);
+                    break :blk mcv;
+                }
+            }
+        };
+        log.debug("consolidating then_entry {d} {}=>{}", .{ then_key, parent_mcv, then_value });
+        // TODO make sure the destination stack offset / register does not already have something
+        // going on there.
+        try self.setRegOrMem(self.air.typeOfIndex(then_key), parent_mcv, then_value);
+        // TODO track the new register / stack allocation
+    }
 
-    return self.fail("TODO implement condbr {}", .{self.target.cpu.arch});
+    self.branch_stack.pop().deinit(self.gpa);
+
+    return self.finishAir(inst, .unreach, .{ pl_op.operand, .none, .none });
 }
 
 fn isNull(self: *Self, operand: MCValue) !MCValue {
@@ -1860,10 +2061,12 @@ fn airSwitch(self: *Self, inst: Air.Inst.Index) !void {
     return self.fail("TODO airSwitch for {}", .{self.target.cpu.arch});
 }
 
-fn performReloc(self: *Self, reloc: Reloc) !void {
-    switch (reloc) {
-        .rel32 => return self.fail("TODO reloc.rel32 for {}", .{self.target.cpu.arch}),
-        .arm_branch => return self.fail("TODO reloc.arm_branch for {}", .{self.target.cpu.arch}),
+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),
+        else => unreachable,
     }
 }
 
@@ -1903,7 +2106,10 @@ fn brVoid(self: *Self, block: Air.Inst.Index) !void {
     // Emit a jump with a relocation. It will be patched up after the block ends.
     try block_data.relocs.ensureUnusedCapacity(self.gpa, 1);
 
-    return self.fail("TODO implement brvoid for {}", .{self.target.cpu.arch});
+    block_data.relocs.appendAssumeCapacity(try self.addInst(.{
+        .tag = .b,
+        .data = .{ .inst = undefined }, // populated later through performReloc
+    }));
 }
 
 fn airAsm(self: *Self, inst: Air.Inst.Index) !void {
@@ -2050,8 +2256,6 @@ fn genSetStack(self: *Self, ty: Type, stack_offset: u32, mcv: MCValue) InnerErro
             return self.fail("TODO implement set stack variable from embedded_in_code", .{});
         },
         .register => |reg| {
-            _ = reg;
-
             const abi_size = ty.abiSize(self.target.*);
             const adj_off = stack_offset + abi_size;
 
@@ -2115,6 +2319,25 @@ fn genSetReg(self: *Self, ty: Type, reg: Register, mcv: MCValue) InnerError!void
                 else => unreachable, // unexpected register size
             }
         },
+        .compare_flags_unsigned,
+        .compare_flags_signed,
+        => |op| {
+            const condition = switch (mcv) {
+                .compare_flags_unsigned => Instruction.Condition.fromCompareOperatorUnsigned(op),
+                .compare_flags_signed => Instruction.Condition.fromCompareOperatorSigned(op),
+                else => unreachable,
+            };
+
+            _ = try self.addInst(.{
+                .tag = .cset,
+                .data = .{ .rrr_cond = .{
+                    .rd = reg,
+                    .rn = .xzr,
+                    .rm = .xzr,
+                    .cond = condition,
+                } },
+            });
+        },
         .immediate => |x| {
             _ = try self.addInst(.{
                 .tag = .movz,
diff --git a/src/arch/aarch64/Emit.zig b/src/arch/aarch64/Emit.zig
index 665d529245..ff619cd71c 100644
--- a/src/arch/aarch64/Emit.zig
+++ b/src/arch/aarch64/Emit.zig
@@ -14,6 +14,7 @@ const DW = std.dwarf;
 const leb128 = std.leb;
 const Instruction = bits.Instruction;
 const Register = bits.Register;
+const log = std.log.scoped(.aarch64_emit);
 const DebugInfoOutput = @import("../../codegen.zig").DebugInfoOutput;
 
 mir: Mir,
@@ -47,9 +48,16 @@ const InnerError = error{
 };
 
 const BranchType = enum {
+    b_cond,
     unconditional_branch_immediate,
 
-    const default = BranchType.unconditional_branch_immediate;
+    fn default(tag: Mir.Inst.Tag) BranchType {
+        return switch (tag) {
+            .b, .bl => .unconditional_branch_immediate,
+            .b_cond => .b_cond,
+            else => unreachable,
+        };
+    }
 };
 
 pub fn emitMir(
@@ -65,8 +73,11 @@ pub fn emitMir(
         const inst = @intCast(u32, index);
         switch (tag) {
             .add_immediate => try emit.mirAddSubtractImmediate(inst),
+            .cmp_immediate => try emit.mirAddSubtractImmediate(inst),
             .sub_immediate => try emit.mirAddSubtractImmediate(inst),
 
+            .b_cond => try emit.mirConditionalBranchImmediate(inst),
+
             .b => try emit.mirBranch(inst),
             .bl => try emit.mirBranch(inst),
 
@@ -78,6 +89,10 @@ pub fn emitMir(
 
             .call_extern => try emit.mirCallExtern(inst),
 
+            .cmp_shifted_register => try emit.mirAddSubtractShiftedRegister(inst),
+
+            .cset => try emit.mirConditionalSelect(inst),
+
             .dbg_line => try emit.mirDbgLine(inst),
 
             .dbg_prologue_end => try emit.mirDebugPrologueEnd(),
@@ -107,29 +122,50 @@ pub fn emitMir(
 }
 
 pub fn deinit(emit: *Emit) void {
+    var iter = emit.branch_forward_origins.valueIterator();
+    while (iter.next()) |origin_list| {
+        origin_list.deinit(emit.bin_file.allocator);
+    }
+
     emit.branch_types.deinit(emit.bin_file.allocator);
     emit.branch_forward_origins.deinit(emit.bin_file.allocator);
     emit.code_offset_mapping.deinit(emit.bin_file.allocator);
     emit.* = undefined;
 }
 
-fn optimalBranchType(emit: *Emit, offset: i64) !BranchType {
+fn optimalBranchType(emit: *Emit, tag: Mir.Inst.Tag, offset: i64) !BranchType {
     assert(offset & 0b11 == 0);
 
-    // TODO handle conditional branches
-    if (std.math.cast(i26, offset >> 2)) |_| {
-        return BranchType.unconditional_branch_immediate;
-    } else |_| {
-        return emit.fail("TODO support branches larger than +-128 MiB", .{});
+    switch (tag) {
+        .b, .bl => {
+            if (std.math.cast(i26, offset >> 2)) |_| {
+                return BranchType.unconditional_branch_immediate;
+            } else |_| {
+                return emit.fail("TODO support branches larger than +-128 MiB", .{});
+            }
+        },
+        .b_cond => {
+            if (std.math.cast(i19, offset >> 2)) |_| {
+                return BranchType.b_cond;
+            } else |_| {
+                return emit.fail("TODO support conditional branches larger than +-1 MiB", .{});
+            }
+        },
+        else => unreachable,
     }
 }
 
 fn instructionSize(emit: *Emit, inst: Mir.Inst.Index) usize {
     const tag = emit.mir.instructions.items(.tag)[inst];
-    switch (tag) {
-        .b, .bl => switch (emit.branch_types.get(inst).?) {
+
+    if (isBranch(tag)) {
+        switch (emit.branch_types.get(inst).?) {
             .unconditional_branch_immediate => return 4,
-        },
+            .b_cond => return 4,
+        }
+    }
+
+    switch (tag) {
         .load_memory => {
             if (emit.bin_file.options.pie) {
                 // adrp, ldr
@@ -146,10 +182,32 @@ fn instructionSize(emit: *Emit, inst: Mir.Inst.Index) usize {
                 return 5 * 4;
             }
         },
+        .call_extern => return 4,
+        .dbg_line,
+        .dbg_epilogue_begin,
+        .dbg_prologue_end,
+        => return 0,
         else => return 4,
     }
 }
 
+fn isBranch(tag: Mir.Inst.Tag) bool {
+    return switch (tag) {
+        .b, .bl, .b_cond => true,
+        else => false,
+    };
+}
+
+fn branchTarget(emit: *Emit, inst: Mir.Inst.Index) Mir.Inst.Index {
+    const tag = emit.mir.instructions.items(.tag)[inst];
+
+    switch (tag) {
+        .b, .bl => return emit.mir.instructions.items(.data)[inst].inst,
+        .b_cond => return emit.mir.instructions.items(.data)[inst].inst_cond.inst,
+        else => unreachable,
+    }
+}
+
 fn lowerBranches(emit: *Emit) !void {
     const mir_tags = emit.mir.instructions.items(.tag);
     const allocator = emit.bin_file.allocator;
@@ -162,41 +220,38 @@ fn lowerBranches(emit: *Emit) !void {
     // generating MIR
     for (mir_tags) |tag, index| {
         const inst = @intCast(u32, index);
-        switch (tag) {
-            .b, .bl => {
-                const target_inst = emit.mir.instructions.items(.data)[inst].inst;
-
-                // Remember this branch instruction
-                try emit.branch_types.put(allocator, inst, BranchType.default);
-
-                // Forward branches require some extra stuff: We only
-                // know their offset once we arrive at the target
-                // instruction. Therefore, we need to be able to
-                // access the branch instruction when we visit the
-                // target instruction in order to manipulate its type
-                // etc.
-                if (target_inst > inst) {
-                    // Remember the branch instruction index
-                    try emit.code_offset_mapping.put(allocator, inst, 0);
-
-                    if (emit.branch_forward_origins.getPtr(target_inst)) |origin_list| {
-                        try origin_list.append(allocator, inst);
-                    } else {
-                        var origin_list: std.ArrayListUnmanaged(Mir.Inst.Index) = .{};
-                        try origin_list.append(allocator, inst);
-                        try emit.branch_forward_origins.put(allocator, target_inst, origin_list);
-                    }
+        if (isBranch(tag)) {
+            const target_inst = emit.branchTarget(inst);
+
+            // Remember this branch instruction
+            try emit.branch_types.put(allocator, inst, BranchType.default(tag));
+
+            // Forward branches require some extra stuff: We only
+            // know their offset once we arrive at the target
+            // instruction. Therefore, we need to be able to
+            // access the branch instruction when we visit the
+            // target instruction in order to manipulate its type
+            // etc.
+            if (target_inst > inst) {
+                // Remember the branch instruction index
+                try emit.code_offset_mapping.put(allocator, inst, 0);
+
+                if (emit.branch_forward_origins.getPtr(target_inst)) |origin_list| {
+                    try origin_list.append(allocator, inst);
+                } else {
+                    var origin_list: std.ArrayListUnmanaged(Mir.Inst.Index) = .{};
+                    try origin_list.append(allocator, inst);
+                    try emit.branch_forward_origins.put(allocator, target_inst, origin_list);
                 }
+            }
 
-                // Remember the target instruction index so that we
-                // update the real code offset in all future passes
-                //
-                // putNoClobber may not be used as the put operation
-                // may clobber the entry when multiple branches branch
-                // to the same target instruction
-                try emit.code_offset_mapping.put(allocator, target_inst, 0);
-            },
-            else => {}, // not a branch
+            // Remember the target instruction index so that we
+            // update the real code offset in all future passes
+            //
+            // putNoClobber may not be used as the put operation
+            // may clobber the entry when multiple branches branch
+            // to the same target instruction
+            try emit.code_offset_mapping.put(allocator, target_inst, 0);
         }
     }
 
@@ -220,21 +275,20 @@ fn lowerBranches(emit: *Emit) !void {
 
             // If this instruction is a backward branch, calculate the
             // offset, which may potentially update the branch type
-            switch (tag) {
-                .b, .bl => {
-                    const target_inst = emit.mir.instructions.items(.data)[inst].inst;
-                    if (target_inst < inst) {
-                        const target_offset = emit.code_offset_mapping.get(target_inst).?;
-                        const offset = @intCast(i64, target_offset) - @intCast(i64, current_code_offset + 8);
-                        const branch_type = emit.branch_types.getPtr(inst).?;
-                        const optimal_branch_type = try emit.optimalBranchType(offset);
-                        if (branch_type.* != optimal_branch_type) {
-                            branch_type.* = optimal_branch_type;
-                            all_branches_lowered = false;
-                        }
+            if (isBranch(tag)) {
+                const target_inst = emit.branchTarget(inst);
+                if (target_inst < inst) {
+                    const target_offset = emit.code_offset_mapping.get(target_inst).?;
+                    const offset = @intCast(i64, target_offset) - @intCast(i64, current_code_offset);
+                    const branch_type = emit.branch_types.getPtr(inst).?;
+                    const optimal_branch_type = try emit.optimalBranchType(tag, offset);
+                    if (branch_type.* != optimal_branch_type) {
+                        branch_type.* = optimal_branch_type;
+                        all_branches_lowered = false;
                     }
-                },
-                else => {},
+
+                    log.debug("lowerBranches: branch {} has offset {}", .{ inst, offset });
+                }
             }
 
             // If this instruction is the target of one or more
@@ -242,14 +296,17 @@ fn lowerBranches(emit: *Emit) !void {
             // potentially update the branch type
             if (emit.branch_forward_origins.get(inst)) |origin_list| {
                 for (origin_list.items) |forward_branch_inst| {
+                    const branch_tag = emit.mir.instructions.items(.tag)[forward_branch_inst];
                     const forward_branch_inst_offset = emit.code_offset_mapping.get(forward_branch_inst).?;
-                    const offset = @intCast(i64, forward_branch_inst_offset) - @intCast(i64, current_code_offset + 8);
+                    const offset = @intCast(i64, current_code_offset) - @intCast(i64, forward_branch_inst_offset);
                     const branch_type = emit.branch_types.getPtr(forward_branch_inst).?;
-                    const optimal_branch_type = try emit.optimalBranchType(offset);
+                    const optimal_branch_type = try emit.optimalBranchType(branch_tag, offset);
                     if (branch_type.* != optimal_branch_type) {
                         branch_type.* = optimal_branch_type;
                         all_branches_lowered = false;
                     }
+
+                    log.debug("lowerBranches: branch {} has offset {}", .{ forward_branch_inst, offset });
                 }
             }
 
@@ -347,6 +404,12 @@ fn mirAddSubtractImmediate(emit: *Emit, inst: Mir.Inst.Index) !void {
             rr_imm12_sh.imm12,
             rr_imm12_sh.sh == 1,
         )),
+        .cmp_immediate => try emit.writeInstruction(Instruction.subs(
+            rr_imm12_sh.rd,
+            rr_imm12_sh.rn,
+            rr_imm12_sh.imm12,
+            rr_imm12_sh.sh == 1,
+        )),
         .sub_immediate => try emit.writeInstruction(Instruction.sub(
             rr_imm12_sh.rd,
             rr_imm12_sh.rn,
@@ -357,12 +420,37 @@ fn mirAddSubtractImmediate(emit: *Emit, inst: Mir.Inst.Index) !void {
     }
 }
 
+fn mirConditionalBranchImmediate(emit: *Emit, inst: Mir.Inst.Index) !void {
+    const tag = emit.mir.instructions.items(.tag)[inst];
+    const inst_cond = emit.mir.instructions.items(.data)[inst].inst_cond;
+
+    const offset = @intCast(i64, emit.code_offset_mapping.get(inst_cond.inst).?) - @intCast(i64, emit.code.items.len);
+    const branch_type = emit.branch_types.get(inst).?;
+    log.debug("mirConditionalBranchImmediate: {} offset={}", .{ inst, offset });
+
+    switch (branch_type) {
+        .b_cond => switch (tag) {
+            .b_cond => try emit.writeInstruction(Instruction.bCond(inst_cond.cond, @intCast(i21, offset))),
+            else => unreachable,
+        },
+        else => unreachable,
+    }
+}
+
 fn mirBranch(emit: *Emit, inst: Mir.Inst.Index) !void {
     const tag = emit.mir.instructions.items(.tag)[inst];
     const target_inst = emit.mir.instructions.items(.data)[inst].inst;
 
-    const offset = @intCast(i64, emit.code_offset_mapping.get(target_inst).?) - @intCast(i64, emit.code.items.len + 8);
+    log.debug("branch {}(tag: {}) -> {}(tag: {})", .{
+        inst,
+        tag,
+        target_inst,
+        emit.mir.instructions.items(.tag)[target_inst],
+    });
+
+    const offset = @intCast(i64, emit.code_offset_mapping.get(target_inst).?) - @intCast(i64, emit.code.items.len);
     const branch_type = emit.branch_types.get(inst).?;
+    log.debug("mirBranch: {} offset={}", .{ inst, offset });
 
     switch (branch_type) {
         .unconditional_branch_immediate => switch (tag) {
@@ -370,6 +458,7 @@ fn mirBranch(emit: *Emit, inst: Mir.Inst.Index) !void {
             .bl => try emit.writeInstruction(Instruction.bl(@intCast(i28, offset))),
             else => unreachable,
         },
+        else => unreachable,
     }
 }
 
@@ -453,6 +542,37 @@ fn mirCallExtern(emit: *Emit, inst: Mir.Inst.Index) !void {
     }
 }
 
+fn mirAddSubtractShiftedRegister(emit: *Emit, inst: Mir.Inst.Index) !void {
+    const tag = emit.mir.instructions.items(.tag)[inst];
+    const rrr_imm6_shift = emit.mir.instructions.items(.data)[inst].rrr_imm6_shift;
+
+    switch (tag) {
+        .cmp_shifted_register => try emit.writeInstruction(Instruction.subsShiftedRegister(
+            rrr_imm6_shift.rd,
+            rrr_imm6_shift.rn,
+            rrr_imm6_shift.rm,
+            rrr_imm6_shift.shift,
+            rrr_imm6_shift.imm6,
+        )),
+        else => unreachable,
+    }
+}
+
+fn mirConditionalSelect(emit: *Emit, inst: Mir.Inst.Index) !void {
+    const tag = emit.mir.instructions.items(.tag)[inst];
+    const rrr_cond = emit.mir.instructions.items(.data)[inst].rrr_cond;
+
+    switch (tag) {
+        .cset => try emit.writeInstruction(Instruction.csinc(
+            rrr_cond.rd,
+            rrr_cond.rn,
+            rrr_cond.rm,
+            rrr_cond.cond,
+        )),
+        else => unreachable,
+    }
+}
+
 fn mirLoadMemory(emit: *Emit, inst: Mir.Inst.Index) !void {
     assert(emit.mir.instructions.items(.tag)[inst] == .load_memory);
     const payload = emit.mir.instructions.items(.data)[inst].payload;
diff --git a/src/arch/aarch64/Mir.zig b/src/arch/aarch64/Mir.zig
index 43e7c7f1ed..de8415c1e8 100644
--- a/src/arch/aarch64/Mir.zig
+++ b/src/arch/aarch64/Mir.zig
@@ -26,6 +26,8 @@ pub const Inst = struct {
     pub const Tag = enum(u16) {
         /// Add (immediate)
         add_immediate,
+        /// Branch conditionally
+        b_cond,
         /// Branch
         b,
         /// Branch with Link
@@ -36,13 +38,19 @@ pub const Inst = struct {
         brk,
         /// Pseudo-instruction: Call extern
         call_extern,
+        /// Compare (immediate)
+        cmp_immediate,
+        /// Compare (shifted register)
+        cmp_shifted_register,
+        /// Conditional set
+        cset,
         /// Pseudo-instruction: End of prologue
         dbg_prologue_end,
         /// Pseudo-instruction: Beginning of epilogue
         dbg_epilogue_begin,
         /// Pseudo-instruction: Update debug line
         dbg_line,
-        /// Psuedo-instruction: Load memory
+        /// Pseudo-instruction: Load memory
         ///
         /// Payload is `LoadMemory`
         load_memory,
@@ -97,7 +105,7 @@ pub const Inst = struct {
         ///
         /// Used by e.g. nop
         nop: void,
-        /// Another instruction.
+        /// Another instruction
         ///
         /// Used by e.g. b
         inst: Index,
@@ -117,6 +125,13 @@ pub const Inst = struct {
         ///
         /// Used by e.g. blr
         reg: Register,
+        /// Another instruction and a condition
+        ///
+        /// Used by e.g. b_cond
+        inst_cond: struct {
+            inst: Index,
+            cond: bits.Instruction.Condition,
+        },
         /// A register, an unsigned 16-bit immediate, and an optional shift
         ///
         /// Used by e.g. movz
@@ -141,6 +156,25 @@ pub const Inst = struct {
             imm12: u12,
             sh: u1 = 0,
         },
+        /// Three registers and a shift (shift type and 6-bit amount)
+        ///
+        /// Used by e.g. cmp_shifted_register
+        rrr_imm6_shift: struct {
+            rd: Register,
+            rn: Register,
+            rm: Register,
+            imm6: u6,
+            shift: bits.Instruction.AddSubtractShiftedRegisterShift,
+        },
+        /// Three registers and a condition
+        ///
+        /// Used by e.g. cset
+        rrr_cond: struct {
+            rd: Register,
+            rn: Register,
+            rm: Register,
+            cond: bits.Instruction.Condition,
+        },
         /// Three registers and a LoadStoreOffset
         ///
         /// Used by e.g. str_register
diff --git a/src/arch/aarch64/bits.zig b/src/arch/aarch64/bits.zig
index f751170818..26974a3b6a 100644
--- a/src/arch/aarch64/bits.zig
+++ b/src/arch/aarch64/bits.zig
@@ -295,6 +295,18 @@ pub const Instruction = union(enum) {
         op: u1,
         sf: u1,
     },
+    add_subtract_shifted_register: packed struct {
+        rd: u5,
+        rn: u5,
+        imm6: u6,
+        rm: u5,
+        fixed_1: u1 = 0b0,
+        shift: u2,
+        fixed_2: u5 = 0b01011,
+        s: u1,
+        op: u1,
+        sf: u1,
+    },
     conditional_branch: struct {
         cond: u4,
         o0: u1,
@@ -309,6 +321,17 @@ pub const Instruction = union(enum) {
         fixed: u6 = 0b011010,
         sf: u1,
     },
+    conditional_select: struct {
+        rd: u5,
+        rn: u5,
+        op2: u2,
+        cond: u4,
+        rm: u5,
+        fixed: u8 = 0b11010100,
+        s: u1,
+        op: u1,
+        sf: u1,
+    },
 
     pub const Shift = struct {
         shift: Type = .lsl,
@@ -376,6 +399,57 @@ pub const Instruction = union(enum) {
         /// Integer: Always
         /// Floating point: Always
         nv,
+
+        /// Converts a std.math.CompareOperator into a condition flag,
+        /// i.e. returns the condition that is true iff the result of the
+        /// comparison is true. Assumes signed comparison
+        pub fn fromCompareOperatorSigned(op: std.math.CompareOperator) Condition {
+            return switch (op) {
+                .gte => .ge,
+                .gt => .gt,
+                .neq => .ne,
+                .lt => .lt,
+                .lte => .le,
+                .eq => .eq,
+            };
+        }
+
+        /// Converts a std.math.CompareOperator into a condition flag,
+        /// i.e. returns the condition that is true iff the result of the
+        /// comparison is true. Assumes unsigned comparison
+        pub fn fromCompareOperatorUnsigned(op: std.math.CompareOperator) Condition {
+            return switch (op) {
+                .gte => .cs,
+                .gt => .hi,
+                .neq => .ne,
+                .lt => .cc,
+                .lte => .ls,
+                .eq => .eq,
+            };
+        }
+
+        /// Returns the condition which is true iff the given condition is
+        /// false (if such a condition exists)
+        pub fn negate(cond: Condition) Condition {
+            return switch (cond) {
+                .eq => .ne,
+                .ne => .eq,
+                .cs => .cc,
+                .cc => .cs,
+                .mi => .pl,
+                .pl => .mi,
+                .vs => .vc,
+                .vc => .vs,
+                .hi => .ls,
+                .ls => .hi,
+                .ge => .lt,
+                .lt => .ge,
+                .gt => .le,
+                .le => .gt,
+                .al => unreachable,
+                .nv => unreachable,
+            };
+        }
     };
 
     pub fn toU32(self: Instruction) u32 {
@@ -391,9 +465,11 @@ pub const Instruction = union(enum) {
             .no_operation => |v| @bitCast(u32, v),
             .logical_shifted_register => |v| @bitCast(u32, v),
             .add_subtract_immediate => |v| @bitCast(u32, v),
+            .add_subtract_shifted_register => |v| @bitCast(u32, v),
             // TODO once packed structs work, this can be refactored
             .conditional_branch => |v| @as(u32, v.cond) | (@as(u32, v.o0) << 4) | (@as(u32, v.imm19) << 5) | (@as(u32, v.o1) << 24) | (@as(u32, v.fixed) << 25),
             .compare_and_branch => |v| @as(u32, v.rt) | (@as(u32, v.imm19) << 5) | (@as(u32, v.op) << 24) | (@as(u32, v.fixed) << 25) | (@as(u32, v.sf) << 31),
+            .conditional_select => |v| @as(u32, v.rd) | @as(u32, v.rn) << 5 | @as(u32, v.op2) << 10 | @as(u32, v.cond) << 12 | @as(u32, v.rm) << 16 | @as(u32, v.fixed) << 21 | @as(u32, v.s) << 29 | @as(u32, v.op) << 30 | @as(u32, v.sf) << 31,
         };
     }
 
@@ -804,6 +880,35 @@ pub const Instruction = union(enum) {
         };
     }
 
+    pub const AddSubtractShiftedRegisterShift = enum(u2) { lsl, lsr, asr, _ };
+
+    fn addSubtractShiftedRegister(
+        op: u1,
+        s: u1,
+        shift: AddSubtractShiftedRegisterShift,
+        rd: Register,
+        rn: Register,
+        rm: Register,
+        imm6: u6,
+    ) Instruction {
+        return Instruction{
+            .add_subtract_shifted_register = .{
+                .rd = rd.id(),
+                .rn = rn.id(),
+                .imm6 = imm6,
+                .rm = rm.id(),
+                .shift = @enumToInt(shift),
+                .s = s,
+                .op = op,
+                .sf = switch (rd.size()) {
+                    32 => 0b0,
+                    64 => 0b1,
+                    else => unreachable, // unexpected register size
+                },
+            },
+        };
+    }
+
     fn conditionalBranch(
         o0: u1,
         o1: u1,
@@ -841,6 +946,33 @@ pub const Instruction = union(enum) {
         };
     }
 
+    fn conditionalSelect(
+        op2: u2,
+        op: u1,
+        s: u1,
+        rd: Register,
+        rn: Register,
+        rm: Register,
+        cond: Condition,
+    ) Instruction {
+        return Instruction{
+            .conditional_select = .{
+                .rd = rd.id(),
+                .rn = rn.id(),
+                .op2 = op2,
+                .cond = @enumToInt(cond),
+                .rm = rm.id(),
+                .s = s,
+                .op = op,
+                .sf = switch (rd.size()) {
+                    32 => 0b0,
+                    64 => 0b1,
+                    else => unreachable, // unexpected register size
+                },
+            },
+        };
+    }
+
     // Helper functions for assembly syntax functions
 
     // Move wide (immediate)
@@ -1055,6 +1187,48 @@ pub const Instruction = union(enum) {
         return addSubtractImmediate(0b1, 0b1, rd, rn, imm, shift);
     }
 
+    // Add/subtract (shifted register)
+
+    pub fn addShiftedRegister(
+        rd: Register,
+        rn: Register,
+        rm: Register,
+        shift: AddSubtractShiftedRegisterShift,
+        imm6: u6,
+    ) Instruction {
+        return addSubtractShiftedRegister(0b0, 0b0, shift, rd, rn, rm, imm6);
+    }
+
+    pub fn addsShiftedRegister(
+        rd: Register,
+        rn: Register,
+        rm: Register,
+        shift: AddSubtractShiftedRegisterShift,
+        imm6: u6,
+    ) Instruction {
+        return addSubtractShiftedRegister(0b0, 0b1, shift, rd, rn, rm, imm6);
+    }
+
+    pub fn subShiftedRegister(
+        rd: Register,
+        rn: Register,
+        rm: Register,
+        shift: AddSubtractShiftedRegisterShift,
+        imm6: u6,
+    ) Instruction {
+        return addSubtractShiftedRegister(0b1, 0b0, shift, rd, rn, rm, imm6);
+    }
+
+    pub fn subsShiftedRegister(
+        rd: Register,
+        rn: Register,
+        rm: Register,
+        shift: AddSubtractShiftedRegisterShift,
+        imm6: u6,
+    ) Instruction {
+        return addSubtractShiftedRegister(0b1, 0b1, shift, rd, rn, rm, imm6);
+    }
+
     // Conditional branch
 
     pub fn bCond(cond: Condition, offset: i21) Instruction {
@@ -1070,6 +1244,24 @@ pub const Instruction = union(enum) {
     pub fn cbnz(rt: Register, offset: i21) Instruction {
         return compareAndBranch(0b1, rt, offset);
     }
+
+    // Conditional select
+
+    pub fn csel(rd: Register, rn: Register, rm: Register, cond: Condition) Instruction {
+        return conditionalSelect(0b00, 0b0, 0b0, rd, rn, rm, cond);
+    }
+
+    pub fn csinc(rd: Register, rn: Register, rm: Register, cond: Condition) Instruction {
+        return conditionalSelect(0b01, 0b0, 0b0, rd, rn, rm, cond);
+    }
+
+    pub fn csinv(rd: Register, rn: Register, rm: Register, cond: Condition) Instruction {
+        return conditionalSelect(0b00, 0b1, 0b0, rd, rn, rm, cond);
+    }
+
+    pub fn csneg(rd: Register, rn: Register, rm: Register, cond: Condition) Instruction {
+        return conditionalSelect(0b01, 0b1, 0b0, rd, rn, rm, cond);
+    }
 };
 
 test {
@@ -1231,6 +1423,14 @@ test "serialize instructions" {
             .inst = Instruction.cbz(.x10, 40),
             .expected = 0b1_011010_0_0000000000000001010_01010,
         },
+        .{ // add x0, x1, x2, lsl #5
+            .inst = Instruction.addShiftedRegister(.x0, .x1, .x2, .lsl, 5),
+            .expected = 0b1_0_0_01011_00_0_00010_000101_00001_00000,
+        },
+        .{ // csinc x1, x2, x4, eq
+            .inst = Instruction.csinc(.x1, .x2, .x4, .eq),
+            .expected = 0b1_0_0_11010100_00100_0000_0_1_00010_00001,
+        },
     };
 
     for (testcases) |case| {