Merge pull request #24381 from Justus2308/switch-better-underscore

Enhance switch on non-exhaustive enums

1 files changed, 66 insertions, 37 deletions

diff --git a/src/Zcu.zig b/src/Zcu.zig
index 158badfcaa..53da90a9c2 100644
--- a/src/Zcu.zig
+++ b/src/Zcu.zig
@@ -1679,20 +1679,37 @@ pub const SrcLoc = struct {
                 return tree.nodeToSpan(condition);
             },
 
-            .node_offset_switch_special_prong => |node_off| {
+            .node_offset_switch_else_prong => |node_off| {
                 const tree = try src_loc.file_scope.getTree(zcu);
                 const switch_node = node_off.toAbsolute(src_loc.base_node);
                 _, const extra_index = tree.nodeData(switch_node).node_and_extra;
                 const case_nodes = tree.extraDataSlice(tree.extraData(extra_index, Ast.Node.SubRange), Ast.Node.Index);
                 for (case_nodes) |case_node| {
                     const case = tree.fullSwitchCase(case_node).?;
-                    const is_special = (case.ast.values.len == 0) or
-                        (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])), "_"));
-                    if (!is_special) continue;
+                    if (case.ast.values.len == 0) {
+                        return tree.nodeToSpan(case_node);
+                    }
+                } else unreachable;
+            },
 
-                    return tree.nodeToSpan(case_node);
+            .node_offset_switch_under_prong => |node_off| {
+                const tree = try src_loc.file_scope.getTree(zcu);
+                const switch_node = node_off.toAbsolute(src_loc.base_node);
+                _, const extra_index = tree.nodeData(switch_node).node_and_extra;
+                const case_nodes = tree.extraDataSlice(tree.extraData(extra_index, Ast.Node.SubRange), Ast.Node.Index);
+                for (case_nodes) |case_node| {
+                    const case = tree.fullSwitchCase(case_node).?;
+                    for (case.ast.values) |val| {
+                        if (tree.nodeTag(val) == .identifier and
+                            mem.eql(u8, tree.tokenSlice(tree.nodeMainToken(val)), "_"))
+                        {
+                            return tree.tokensToSpan(
+                                tree.firstToken(case_node),
+                                tree.lastToken(case_node),
+                                tree.nodeMainToken(val),
+                            );
+                        }
+                    }
                 } else unreachable;
             },
 
@@ -1703,12 +1720,6 @@ pub const SrcLoc = struct {
                 const case_nodes = tree.extraDataSlice(tree.extraData(extra_index, Ast.Node.SubRange), Ast.Node.Index);
                 for (case_nodes) |case_node| {
                     const case = tree.fullSwitchCase(case_node).?;
-                    const is_special = (case.ast.values.len == 0) or
-                        (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])), "_"));
-                    if (is_special) continue;
-
                     for (case.ast.values) |item_node| {
                         if (tree.nodeTag(item_node) == .switch_range) {
                             return tree.nodeToSpan(item_node);
@@ -2113,28 +2124,35 @@ pub const SrcLoc = struct {
 
                 var multi_i: u32 = 0;
                 var scalar_i: u32 = 0;
-                const case = for (case_nodes) |case_node| {
+                var underscore_node: Ast.Node.OptionalIndex = .none;
+                const case = case: for (case_nodes) |case_node| {
                     const case = tree.fullSwitchCase(case_node).?;
-                    const is_special = special: {
-                        if (case.ast.values.len == 0) break :special true;
-                        if (case.ast.values.len == 1 and tree.nodeTag(case.ast.values[0]) == .identifier) {
-                            break :special mem.eql(u8, tree.tokenSlice(tree.nodeMainToken(case.ast.values[0])), "_");
+                    if (case.ast.values.len == 0) {
+                        if (want_case_idx == LazySrcLoc.Offset.SwitchCaseIndex.special_else) {
+                            break :case case;
                         }
-                        break :special false;
-                    };
-                    if (is_special) {
-                        if (want_case_idx.isSpecial()) {
-                            break case;
-                        }
-                        continue;
+                        continue :case;
                     }
+                    if (underscore_node == .none) for (case.ast.values) |val_node| {
+                        if (tree.nodeTag(val_node) == .identifier and
+                            mem.eql(u8, tree.tokenSlice(tree.nodeMainToken(val_node)), "_"))
+                        {
+                            underscore_node = val_node.toOptional();
+                            if (want_case_idx == LazySrcLoc.Offset.SwitchCaseIndex.special_under) {
+                                break :case case;
+                            }
+                            continue :case;
+                        }
+                    };
 
                     const is_multi = case.ast.values.len != 1 or
                         tree.nodeTag(case.ast.values[0]) == .switch_range;
 
                     switch (want_case_idx.kind) {
-                        .scalar => if (!is_multi and want_case_idx.index == scalar_i) break case,
-                        .multi => if (is_multi and want_case_idx.index == multi_i) break case,
+                        .scalar => if (!is_multi and want_case_idx.index == scalar_i)
+                            break :case case,
+                        .multi => if (is_multi and want_case_idx.index == multi_i)
+                            break :case case,
                     }
 
                     if (is_multi) {
@@ -2148,7 +2166,10 @@ pub const SrcLoc = struct {
                     .switch_case_item,
                     .switch_case_item_range_first,
                     .switch_case_item_range_last,
-                    => |x| x.item_idx,
+                    => |x| item_idx: {
+                        assert(want_case_idx != LazySrcLoc.Offset.SwitchCaseIndex.special_else);
+                        break :item_idx x.item_idx;
+                    },
                     .switch_capture, .switch_tag_capture => {
                         const start = switch (src_loc.lazy) {
                             .switch_capture => case.payload_token.?,
@@ -2173,7 +2194,11 @@ pub const SrcLoc = struct {
                     .single => {
                         var item_i: 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;
+                            }
                             if (item_i != want_item.index) {
                                 item_i += 1;
                                 continue;
@@ -2184,7 +2209,9 @@ pub const SrcLoc = struct {
                     .range => {
                         var range_i: u32 = 0;
                         for (case.ast.values) |item_node| {
-                            if (tree.nodeTag(item_node) != .switch_range) continue;
+                            if (tree.nodeTag(item_node) != .switch_range) {
+                                continue;
+                            }
                             if (range_i != want_item.index) {
                                 range_i += 1;
                                 continue;
@@ -2363,10 +2390,14 @@ pub const LazySrcLoc = struct {
         /// by taking this AST node index offset from the containing base node,
         /// which points to a switch expression AST node. Next, navigate to the operand.
         node_offset_switch_operand: Ast.Node.Offset,
-        /// The source location points to the else/`_` prong of a switch expression, found
+        /// The source location points to the else prong of a switch expression, found
+        /// by taking this AST node index offset from the containing base node,
+        /// which points to a switch expression AST node. Next, navigate to the else prong.
+        node_offset_switch_else_prong: Ast.Node.Offset,
+        /// The source location points to the `_` prong of a switch expression, found
         /// by taking this AST node index offset from the containing base node,
-        /// which points to a switch expression AST node. Next, navigate to the else/`_` prong.
-        node_offset_switch_special_prong: Ast.Node.Offset,
+        /// which points to a switch expression AST node. Next, navigate to the `_` prong.
+        node_offset_switch_under_prong: Ast.Node.Offset,
         /// The source location points to all the ranges of a switch expression, found
         /// by taking this AST node index offset from the containing base node,
         /// which points to a switch expression AST node. Next, navigate to any of the
@@ -2562,10 +2593,8 @@ pub const LazySrcLoc = struct {
             kind: enum(u1) { scalar, multi },
             index: u31,
 
-            pub const special: SwitchCaseIndex = @bitCast(@as(u32, std.math.maxInt(u32)));
-            pub fn isSpecial(idx: SwitchCaseIndex) bool {
-                return @as(u32, @bitCast(idx)) == @as(u32, @bitCast(special));
-            }
+            pub const special_else: SwitchCaseIndex = @bitCast(@as(u32, std.math.maxInt(u32)));
+            pub const special_under: SwitchCaseIndex = @bitCast(@as(u32, std.math.maxInt(u32) - 1));
         };
 
         pub const SwitchItemIndex = packed struct(u32) {