riscv: implement switch_dispatch & loop_switch_br

src/arch/riscv64/CodeGen.zig

@@ -108,7 +108,7 @@ frame_allocs: std.MultiArrayList(FrameAlloc) = .{},
 free_frame_indices: std.AutoArrayHashMapUnmanaged(FrameIndex, void) = .{},
 frame_locs: std.MultiArrayList(Mir.FrameLoc) = .{},
 
-loop_repeat_info: std.AutoHashMapUnmanaged(Air.Inst.Index, struct {
+loops: std.AutoHashMapUnmanaged(Air.Inst.Index, struct {
     /// The state to restore before branching.
     state: State,
     /// The branch target.
@@ -232,11 +232,12 @@ const MCValue = union(enum) {
             .register,
             .register_pair,
             .register_offset,
-            .load_frame,
             .load_symbol,
             .load_tlv,
             .indirect,
             => true,
+
+            .load_frame => |frame_addr| !frame_addr.index.isNamed(),
         };
     }
 
@@ -804,7 +805,7 @@ pub fn generate(
         function.frame_allocs.deinit(gpa);
         function.free_frame_indices.deinit(gpa);
         function.frame_locs.deinit(gpa);
-        function.loop_repeat_info.deinit(gpa);
+        function.loops.deinit(gpa);
         var block_it = function.blocks.valueIterator();
         while (block_it.next()) |block| block.deinit(gpa);
         function.blocks.deinit(gpa);
@@ -1588,7 +1589,7 @@ fn genBody(func: *Func, body: []const Air.Inst.Index) InnerError!void {
             .block           => try func.airBlock(inst),
             .br              => try func.airBr(inst),
             .repeat          => try func.airRepeat(inst),
-            .switch_dispatch => return func.fail("TODO implement `switch_dispatch`", .{}),
+            .switch_dispatch => try func.airSwitchDispatch(inst),
             .trap            => try func.airTrap(),
             .breakpoint      => try func.airBreakpoint(),
             .ret_addr        => try func.airRetAddr(inst),
@@ -1678,7 +1679,7 @@ fn genBody(func: *Func, body: []const Air.Inst.Index) InnerError!void {
             .field_parent_ptr => try func.airFieldParentPtr(inst),
 
             .switch_br       => try func.airSwitchBr(inst),
-            .loop_switch_br  => return func.fail("TODO implement `loop_switch_br`", .{}),
+            .loop_switch_br  => try func.airLoopSwitchBr(inst),
 
             .ptr_slice_len_ptr => try func.airPtrSliceLenPtr(inst),
             .ptr_slice_ptr_ptr => try func.airPtrSlicePtrPtr(inst),
@@ -5610,11 +5611,11 @@ fn airLoop(func: *Func, inst: Air.Inst.Index) !void {
     func.scope_generation += 1;
     const state = try func.saveState();
 
-    try func.loop_repeat_info.putNoClobber(func.gpa, inst, .{
+    try func.loops.putNoClobber(func.gpa, inst, .{
         .state = state,
         .jmp_target = @intCast(func.mir_instructions.len),
     });
-    defer assert(func.loop_repeat_info.remove(inst));
+    defer assert(func.loops.remove(inst));
 
     try func.genBody(body);
 
@@ -5671,12 +5672,7 @@ fn lowerBlock(func: *Func, inst: Air.Inst.Index, body: []const Air.Inst.Index) !
 
 fn airSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
     const switch_br = func.air.unwrapSwitch(inst);
-
-    const liveness = try func.liveness.getSwitchBr(func.gpa, inst, switch_br.cases_len + 1);
-    defer func.gpa.free(liveness.deaths);
-
     const condition = try func.resolveInst(switch_br.operand);
-    const condition_ty = func.typeOf(switch_br.operand);
 
     // If the condition dies here in this switch instruction, process
     // that death now instead of later as this has an effect on
@@ -5685,6 +5681,22 @@ fn airSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
         if (switch_br.operand.toIndex()) |op_inst| try func.processDeath(op_inst);
     }
 
+    try func.lowerSwitchBr(inst, switch_br, condition);
+
+    // We already took care of pl_op.operand earlier, so there's nothing left to do
+    func.finishAirBookkeeping();
+}
+
+fn lowerSwitchBr(
+    func: *Func,
+    inst: Air.Inst.Index,
+    switch_br: Air.UnwrappedSwitch,
+    condition: MCValue,
+) !void {
+    const condition_ty = func.typeOf(switch_br.operand);
+    const liveness = try func.liveness.getSwitchBr(func.gpa, inst, switch_br.cases_len + 1);
+    defer func.gpa.free(liveness.deaths);
+
     func.scope_generation += 1;
     const state = try func.saveState();
 
@@ -5785,8 +5797,92 @@ fn airSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
             .close_scope = true,
         });
     }
+}
+
+fn airLoopSwitchBr(func: *Func, inst: Air.Inst.Index) !void {
+    const switch_br = func.air.unwrapSwitch(inst);
+    const condition = try func.resolveInst(switch_br.operand);
+
+    const mat_cond = if (condition.isMutable() and
+        func.reuseOperand(inst, switch_br.operand, 0, condition))
+        condition
+    else mat_cond: {
+        const ty = func.typeOf(switch_br.operand);
+        const mat_cond = try func.allocRegOrMem(ty, inst, true);
+        try func.genCopy(ty, mat_cond, condition);
+        break :mat_cond mat_cond;
+    };
+    func.inst_tracking.putAssumeCapacityNoClobber(inst, InstTracking.init(mat_cond));
+
+    // If the condition dies here in this switch instruction, process
+    // that death now instead of later as this has an effect on
+    // whether it needs to be spilled in the branches
+    if (func.liveness.operandDies(inst, 0)) {
+        if (switch_br.operand.toIndex()) |op_inst| try func.processDeath(op_inst);
+    }
+
+    func.scope_generation += 1;
+    const state = try func.saveState();
+
+    try func.loops.putNoClobber(func.gpa, inst, .{
+        .state = state,
+        .jmp_target = @intCast(func.mir_instructions.len),
+    });
+    defer assert(func.loops.remove(inst));
+
+    // Stop tracking block result without forgetting tracking info
+    try func.freeValue(mat_cond);
+
+    try func.lowerSwitchBr(inst, switch_br, mat_cond);
+
+    try func.processDeath(inst);
+    func.finishAirBookkeeping();
+}
+
+fn airSwitchDispatch(func: *Func, inst: Air.Inst.Index) !void {
+    const br = func.air.instructions.items(.data)[@intFromEnum(inst)].br;
+
+    const block_ty = func.typeOfIndex(br.block_inst);
+    const block_tracking = func.inst_tracking.getPtr(br.block_inst).?;
+    const loop_data = func.loops.getPtr(br.block_inst).?;
+    done: {
+        try func.getValue(block_tracking.short, null);
+        const src_mcv = try func.resolveInst(br.operand);
+
+        if (func.reuseOperandAdvanced(inst, br.operand, 0, src_mcv, br.block_inst)) {
+            try func.getValue(block_tracking.short, br.block_inst);
+            // .long = .none to avoid merging operand and block result stack frames.
+            const current_tracking: InstTracking = .{ .long = .none, .short = src_mcv };
+            try current_tracking.materializeUnsafe(func, br.block_inst, block_tracking.*);
+            for (current_tracking.getRegs()) |src_reg| func.register_manager.freeReg(src_reg);
+            break :done;
+        }
+
+        try func.getValue(block_tracking.short, br.block_inst);
+        const dst_mcv = block_tracking.short;
+        try func.genCopy(block_ty, dst_mcv, try func.resolveInst(br.operand));
+        break :done;
+    }
+
+    // Process operand death so that it is properly accounted for in the State below.
+    if (func.liveness.operandDies(inst, 0)) {
+        if (br.operand.toIndex()) |op_inst| try func.processDeath(op_inst);
+    }
+
+    try func.restoreState(loop_data.state, &.{}, .{
+        .emit_instructions = true,
+        .update_tracking = false,
+        .resurrect = false,
+        .close_scope = false,
+    });
+
+    // Emit a jump with a relocation. It will be patched up after the block ends.
+    // Leave the jump offset undefined
+    _ = try func.jump(loop_data.jmp_target);
+
+    // Stop tracking block result without forgetting tracking info
+    try func.freeValue(block_tracking.short);
 
-    // We already took care of pl_op.operand earlier, so there's nothing left to do
     func.finishAirBookkeeping();
 }
 
@@ -5867,7 +5963,7 @@ fn airBr(func: *Func, inst: Air.Inst.Index) !void {
 
 fn airRepeat(func: *Func, inst: Air.Inst.Index) !void {
     const loop_inst = func.air.instructions.items(.data)[@intFromEnum(inst)].repeat.loop_inst;
-    const repeat_info = func.loop_repeat_info.get(loop_inst).?;
+    const repeat_info = func.loops.get(loop_inst).?;
     try func.restoreState(repeat_info.state, &.{}, .{
         .emit_instructions = true,
         .update_tracking = false,
@@ -8298,7 +8394,10 @@ fn typeOf(func: *Func, inst: Air.Inst.Ref) Type {
 
 fn typeOfIndex(func: *Func, inst: Air.Inst.Index) Type {
     const zcu = func.pt.zcu;
-    return func.air.typeOfIndex(inst, &zcu.intern_pool);
+    return switch (func.air.instructions.items(.tag)[@intFromEnum(inst)]) {
+        .loop_switch_br => func.typeOf(func.air.unwrapSwitch(inst).operand),
+        else => func.air.typeOfIndex(inst, &zcu.intern_pool),
+    };
 }
 
 fn hasFeature(func: *Func, feature: Target.riscv.Feature) bool {

test/behavior/switch_loop.zig

@@ -80,6 +80,7 @@ test "switch loop on tagged union" {
     if (builtin.zig_backend == .stage2_arm) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_sparc64) return error.SkipZigTest; // TODO
     if (builtin.zig_backend == .stage2_spirv64) return error.SkipZigTest; // TODO
+    if (builtin.zig_backend == .stage2_riscv64) return error.SkipZigTest;
 
     const S = struct {
         const U = union(enum) {