Permit explicit tags with '_' switch prong

Mainly affects ZIR representation of switch_block[_ref]
and special prong (detection) logic for switch.
Adds a new SpecialProng tag 'absorbing_under' that allows
specifying additional explicit tags in a '_' prong which
are respected when checking that every value is handled
during semantic analysis but are not transformed into AIR
and instead 'absorbed' by the '_' branch.

Author: Justus2308

lib/std/zig/Ast.zig

@@ -2824,6 +2824,24 @@ pub const full = struct {
             arrow_token: TokenIndex,
             target_expr: Node.Index,
         };
+
+        /// Returns:
+        ///   `null` if case is not special
+        ///   `.none` if case is else prong
+        ///   Index of underscore otherwise
+        pub fn isSpecial(case: *const SwitchCase, tree: *const Ast) ?Node.OptionalIndex {
+            if (case.ast.values.len == 0) {
+                return .none;
+            }
+            for (case.ast.values) |val| {
+                if (tree.nodeTag(val) == .identifier and
+                    mem.eql(u8, tree.tokenSlice(tree.nodeMainToken(val)), "_"))
+                {
+                    return val.toOptional();
+                }
+            }
+            return null;
+        }
     };
 
     pub const Asm = struct {

lib/std/zig/AstGen.zig

@@ -7656,6 +7656,7 @@ fn switchExpr(
     var special_node: Ast.Node.OptionalIndex = .none;
     var else_src: ?Ast.TokenIndex = null;
     var underscore_src: ?Ast.TokenIndex = null;
+    var underscore_node: Ast.Node.OptionalIndex = .none;
     for (case_nodes) |case_node| {
         const case = tree.fullSwitchCase(case_node).?;
         if (case.payload_token) |payload_token| {
@@ -7676,7 +7677,7 @@ fn switchExpr(
                 any_non_inline_capture = true;
             }
         }
-        // Check for else/`_` prong.
+        // Check for else prong.
         if (case.ast.values.len == 0) {
             const case_src = case.ast.arrow_token - 1;
             if (else_src) |src| {
@@ -7715,56 +7716,60 @@ fn switchExpr(
             special_prong = .@"else";
             else_src = case_src;
             continue;
-        } else if (case.ast.values.len == 1 and
-            tree.nodeTag(case.ast.values[0]) == .identifier and
-            mem.eql(u8, tree.tokenSlice(tree.nodeMainToken(case.ast.values[0])), "_"))
-        {
-            const case_src = case.ast.arrow_token - 1;
-            if (underscore_src) |src| {
-                return astgen.failTokNotes(
-                    case_src,
-                    "multiple '_' prongs in switch expression",
-                    .{},
-                    &[_]u32{
-                        try astgen.errNoteTok(
-                            src,
-                            "previous '_' prong here",
-                            .{},
-                        ),
-                    },
-                );
-            } else if (else_src) |some_else| {
-                return astgen.failNodeNotes(
-                    node,
-                    "else and '_' prong in switch expression",
-                    .{},
-                    &[_]u32{
-                        try astgen.errNoteTok(
-                            some_else,
-                            "else prong here",
-                            .{},
-                        ),
-                        try astgen.errNoteTok(
-                            case_src,
-                            "'_' prong here",
-                            .{},
-                        ),
-                    },
-                );
-            }
-            if (case.inline_token != null) {
-                return astgen.failTok(case_src, "cannot inline '_' prong", .{});
-            }
-            special_node = case_node.toOptional();
-            special_prong = .under;
-            underscore_src = case_src;
-            continue;
         }
 
+        // Check for '_' prong.
+        var found_underscore = false;
         for (case.ast.values) |val| {
-            if (tree.nodeTag(val) == .string_literal)
-                return astgen.failNode(val, "cannot switch on strings", .{});
+            switch (tree.nodeTag(val)) {
+                .identifier => if (mem.eql(u8, tree.tokenSlice(tree.nodeMainToken(val)), "_")) {
+                    const case_src = case.ast.arrow_token - 1;
+                    if (underscore_src) |src| {
+                        return astgen.failTokNotes(
+                            case_src,
+                            "multiple '_' prongs in switch expression",
+                            .{},
+                            &[_]u32{
+                                try astgen.errNoteTok(
+                                    src,
+                                    "previous '_' prong here",
+                                    .{},
+                                ),
+                            },
+                        );
+                    } else if (else_src) |some_else| {
+                        return astgen.failNodeNotes(
+                            node,
+                            "else and '_' prong in switch expression",
+                            .{},
+                            &[_]u32{
+                                try astgen.errNoteTok(
+                                    some_else,
+                                    "else prong here",
+                                    .{},
+                                ),
+                                try astgen.errNoteTok(
+                                    case_src,
+                                    "'_' prong here",
+                                    .{},
+                                ),
+                            },
+                        );
+                    }
+                    if (case.inline_token != null) {
+                        return astgen.failTok(case_src, "cannot inline '_' prong", .{});
+                    }
+                    special_node = case_node.toOptional();
+                    special_prong = if (case.ast.values.len == 1) .under else .absorbing_under;
+                    underscore_src = case_src;
+                    underscore_node = val.toOptional();
+                    found_underscore = true;
+                },
+                .string_literal => return astgen.failNode(val, "cannot switch on strings", .{}),
+                else => {},
+            }
         }
+        if (found_underscore) continue;
 
         if (case.ast.values.len == 1 and tree.nodeTag(case.ast.values[0]) != .switch_range) {
             scalar_cases_len += 1;
@@ -7928,14 +7933,23 @@ fn switchExpr(
 
         const header_index: u32 = @intCast(payloads.items.len);
         const body_len_index = if (is_multi_case) blk: {
-            payloads.items[multi_case_table + multi_case_index] = header_index;
-            multi_case_index += 1;
+            if (case_node.toOptional() == special_node) {
+                assert(special_prong == .absorbing_under);
+                payloads.items[case_table_start] = header_index;
+            } else {
+                payloads.items[multi_case_table + multi_case_index] = header_index;
+                multi_case_index += 1;
+            }
             try payloads.resize(gpa, header_index + 3); // items_len, ranges_len, body_len
 
             // items
             var items_len: u32 = 0;
             for (case.ast.values) |item_node| {
-                if (tree.nodeTag(item_node) == .switch_range) continue;
+                if (item_node.toOptional() == underscore_node or
+                    tree.nodeTag(item_node) == .switch_range)
+                {
+                    continue;
+                }
                 items_len += 1;
 
                 const item_inst = try comptimeExpr(parent_gz, scope, item_ri, item_node, .switch_item);
@@ -7945,7 +7959,9 @@ fn switchExpr(
             // ranges
             var ranges_len: u32 = 0;
             for (case.ast.values) |range| {
-                if (tree.nodeTag(range) != .switch_range) continue;
+                if (tree.nodeTag(range) != .switch_range) {
+                    continue;
+                }
                 ranges_len += 1;
 
                 const first_node, const last_node = tree.nodeData(range).node_and_node;
@@ -7960,6 +7976,7 @@ fn switchExpr(
             payloads.items[header_index + 1] = ranges_len;
             break :blk header_index + 2;
         } else if (case_node.toOptional() == special_node) blk: {
+            assert(special_prong != .absorbing_under);
             payloads.items[case_table_start] = header_index;
             try payloads.resize(gpa, header_index + 1); // body_len
             break :blk header_index;
@@ -8015,15 +8032,13 @@ fn switchExpr(
     try astgen.extra.ensureUnusedCapacity(gpa, @typeInfo(Zir.Inst.SwitchBlock).@"struct".fields.len +
         @intFromBool(multi_cases_len != 0) +
         @intFromBool(any_has_tag_capture) +
-        payloads.items.len - case_table_end +
-        (case_table_end - case_table_start) * @typeInfo(Zir.Inst.As).@"struct".fields.len);
+        payloads.items.len - scratch_top);
 
     const payload_index = astgen.addExtraAssumeCapacity(Zir.Inst.SwitchBlock{
         .operand = raw_operand,
         .bits = Zir.Inst.SwitchBlock.Bits{
             .has_multi_cases = multi_cases_len != 0,
-            .has_else = special_prong == .@"else",
-            .has_under = special_prong == .under,
+            .special_prong = special_prong,
             .any_has_tag_capture = any_has_tag_capture,
             .any_non_inline_capture = any_non_inline_capture,
             .has_continue = switch_full.label_token != null and block_scope.label.?.used_for_continue,
@@ -8042,13 +8057,30 @@ fn switchExpr(
     const zir_datas = astgen.instructions.items(.data);
     zir_datas[@intFromEnum(switch_block)].pl_node.payload_index = payload_index;
 
-    for (payloads.items[case_table_start..case_table_end], 0..) |start_index, i| {
+    var normal_case_table_start = case_table_start;
+    if (special_prong != .none) {
+        normal_case_table_start += 1;
+
+        const start_index = payloads.items[case_table_start];
         var body_len_index = start_index;
         var end_index = start_index;
-        const table_index = case_table_start + i;
-        if (table_index < scalar_case_table) {
+        if (special_prong == .absorbing_under) {
+            body_len_index += 2;
+            const items_len = payloads.items[start_index];
+            const ranges_len = payloads.items[start_index + 1];
+            end_index += 3 + items_len + 2 * ranges_len;
+        } else {
             end_index += 1;
-        } else if (table_index < multi_case_table) {
+        }
+        const prong_info: Zir.Inst.SwitchBlock.ProngInfo = @bitCast(payloads.items[body_len_index]);
+        end_index += prong_info.body_len;
+        astgen.extra.appendSliceAssumeCapacity(payloads.items[start_index..end_index]);
+    }
+    for (payloads.items[normal_case_table_start..case_table_end], 0..) |start_index, i| {
+        var body_len_index = start_index;
+        var end_index = start_index;
+        const table_index = normal_case_table_start + i;
+        if (table_index < multi_case_table) {
             body_len_index += 1;
             end_index += 2;
         } else {

lib/std/zig/Zir.zig

@@ -3225,8 +3225,15 @@ pub const Inst = struct {
 
     /// 0. multi_cases_len: u32 // If has_multi_cases is set.
     /// 1. tag_capture_inst: u32 // If any_has_tag_capture is set. Index of instruction prongs use to refer to the inline tag capture.
-    /// 2. else_body { // If has_else or has_under is set.
+    /// 2. else_body { // If special_prong != .none
+    ///        items_len: u32, // If special_prong == .absorbing_under
+    ///        ranges_len: u32, // If special_prong == .absorbing_under
     ///        info: ProngInfo,
+    ///        item: Ref, // for every items_len
+    ///        ranges: { // for every ranges_len
+    ///            item_first: Ref,
+    ///            item_last: Ref,
+    ///        }
     ///        body member Index for every info.body_len
     ///     }
     /// 3. scalar_cases: { // for every scalar_cases_len
@@ -3238,7 +3245,7 @@ pub const Inst = struct {
     ///        items_len: u32,
     ///        ranges_len: u32,
     ///        info: ProngInfo,
-    ///        item: Ref // for every items_len
+    ///        item: Ref, // for every items_len
     ///        ranges: { // for every ranges_len
     ///            item_first: Ref,
     ///            item_last: Ref,
@@ -3274,10 +3281,8 @@ pub const Inst = struct {
         pub const Bits = packed struct(u32) {
             /// If true, one or more prongs have multiple items.
             has_multi_cases: bool,
-            /// If true, there is an else prong. This is mutually exclusive with `has_under`.
-            has_else: bool,
-            /// If true, there is an underscore prong. This is mutually exclusive with `has_else`.
-            has_under: bool,
+            /// Information about the special prong.
+            special_prong: SpecialProng,
             /// If true, at least one prong has an inline tag capture.
             any_has_tag_capture: bool,
             /// If true, at least one prong has a capture which may not
@@ -3287,17 +3292,6 @@ pub const Inst = struct {
             scalar_cases_len: ScalarCasesLen,
 
             pub const ScalarCasesLen = u26;
-
-            pub fn specialProng(bits: Bits) SpecialProng {
-                const has_else: u2 = @intFromBool(bits.has_else);
-                const has_under: u2 = @intFromBool(bits.has_under);
-                return switch ((has_else << 1) | has_under) {
-                    0b00 => .none,
-                    0b01 => .under,
-                    0b10 => .@"else",
-                    0b11 => unreachable,
-                };
-            }
         };
 
         pub const MultiProng = struct {
@@ -3872,7 +3866,18 @@ pub const Inst = struct {
     };
 };
 
-pub const SpecialProng = enum { none, @"else", under };
+pub const SpecialProng = enum(u2) {
+    none,
+    /// Simple else prong.
+    /// `else => {}`
+    @"else",
+    /// Simple '_' prong.
+    /// `_ => {}`
+    under,
+    /// '_' prong with additional items.
+    /// `a, _, b => {}`
+    absorbing_under,
+};
 
 pub const DeclIterator = struct {
     extra_index: u32,
@@ -4716,7 +4721,7 @@ fn findTrackableSwitch(
     }
 
     const has_special = switch (kind) {
-        .normal => extra.data.bits.specialProng() != .none,
+        .normal => extra.data.bits.special_prong != .none,
         .err_union => has_special: {
             // Handle `non_err_body` first.
             const prong_info: Inst.SwitchBlock.ProngInfo = @bitCast(zir.extra[extra_index]);
@@ -4731,6 +4736,23 @@ fn findTrackableSwitch(
     };
 
     if (has_special) {
+        if (kind == .normal) {
+            if (extra.data.bits.special_prong == .absorbing_under) {
+                const items_len = zir.extra[extra_index];
+                extra_index += 1;
+                const ranges_len = zir.extra[extra_index];
+                extra_index += 1;
+                const prong_info: Zir.Inst.SwitchBlock.ProngInfo = @bitCast(zir.extra[extra_index]);
+                extra_index += 1;
+
+                extra_index += items_len + ranges_len * 2;
+
+                const body = zir.bodySlice(extra_index, prong_info.body_len);
+                extra_index += body.len;
+
+                try zir.findTrackableBody(gpa, contents, defers, body);
+            }
+        }
         const prong_info: Inst.SwitchBlock.ProngInfo = @bitCast(zir.extra[extra_index]);
         extra_index += 1;
         const body = zir.bodySlice(extra_index, prong_info.body_len);