compiler: implement labeled switch/continue

lib/std/zig/Ast.zig

@@ -1197,14 +1197,7 @@ pub fn lastToken(tree: Ast, node: Node.Index) TokenIndex {
             n = extra.sentinel;
         },
 
-        .@"continue" => {
-            if (datas[n].lhs != 0) {
-                return datas[n].lhs + end_offset;
-            } else {
-                return main_tokens[n] + end_offset;
-            }
-        },
-        .@"break" => {
+        .@"continue", .@"break" => {
             if (datas[n].rhs != 0) {
                 n = datas[n].rhs;
             } else if (datas[n].lhs != 0) {
@@ -1908,6 +1901,25 @@ pub fn taggedUnionEnumTag(tree: Ast, node: Node.Index) full.ContainerDecl {
     });
 }
 
+pub fn switchFull(tree: Ast, node: Node.Index) full.Switch {
+    const data = &tree.nodes.items(.data)[node];
+    const main_token = tree.nodes.items(.main_token)[node];
+    const switch_token: TokenIndex, const label_token: ?TokenIndex = switch (tree.tokens.items(.tag)[main_token]) {
+        .identifier => .{ main_token + 2, main_token },
+        .keyword_switch => .{ main_token, null },
+        else => unreachable,
+    };
+    const extra = tree.extraData(data.rhs, Ast.Node.SubRange);
+    return .{
+        .ast = .{
+            .switch_token = switch_token,
+            .condition = data.lhs,
+            .cases = tree.extra_data[extra.start..extra.end],
+        },
+        .label_token = label_token,
+    };
+}
+
 pub fn switchCaseOne(tree: Ast, node: Node.Index) full.SwitchCase {
     const data = &tree.nodes.items(.data)[node];
     const values: *[1]Node.Index = &data.lhs;
@@ -2217,6 +2229,21 @@ fn fullContainerDeclComponents(tree: Ast, info: full.ContainerDecl.Components) f
     return result;
 }
 
+fn fullSwitchComponents(tree: Ast, info: full.Switch.Components) full.Switch {
+    const token_tags = tree.tokens.items(.tag);
+    const tok_i = info.switch_token -| 1;
+    var result: full.Switch = .{
+        .ast = info,
+        .label_token = null,
+    };
+    if (token_tags[tok_i] == .colon and
+        token_tags[tok_i -| 1] == .identifier)
+    {
+        result.label_token = tok_i - 1;
+    }
+    return result;
+}
+
 fn fullSwitchCaseComponents(tree: Ast, info: full.SwitchCase.Components, node: Node.Index) full.SwitchCase {
     const token_tags = tree.tokens.items(.tag);
     const node_tags = tree.nodes.items(.tag);
@@ -2488,6 +2515,13 @@ pub fn fullContainerDecl(tree: Ast, buffer: *[2]Ast.Node.Index, node: Node.Index
     };
 }
 
+pub fn fullSwitch(tree: Ast, node: Node.Index) ?full.Switch {
+    return switch (tree.nodes.items(.tag)[node]) {
+        .@"switch", .switch_comma => tree.switchFull(node),
+        else => null,
+    };
+}
+
 pub fn fullSwitchCase(tree: Ast, node: Node.Index) ?full.SwitchCase {
     return switch (tree.nodes.items(.tag)[node]) {
         .switch_case_one, .switch_case_inline_one => tree.switchCaseOne(node),
@@ -2840,6 +2874,17 @@ pub const full = struct {
         };
     };
 
+    pub const Switch = struct {
+        ast: Components,
+        label_token: ?TokenIndex,
+
+        pub const Components = struct {
+            switch_token: TokenIndex,
+            condition: Node.Index,
+            cases: []const Node.Index,
+        };
+    };
+
     pub const SwitchCase = struct {
         inline_token: ?TokenIndex,
         /// Points to the first token after the `|`. Will either be an identifier or
@@ -3250,6 +3295,7 @@ pub const Node = struct {
         /// main_token is the `(`.
         async_call_comma,
         /// `switch(lhs) {}`. `SubRange[rhs]`.
+        /// `main_token` is the identifier of a preceding label, if any; otherwise `switch`.
         @"switch",
         /// Same as switch except there is known to be a trailing comma
         /// before the final rbrace
@@ -3294,7 +3340,8 @@ pub const Node = struct {
         @"suspend",
         /// `resume lhs`. rhs is unused.
         @"resume",
-        /// `continue`. lhs is token index of label if any. rhs is unused.
+        /// `continue :lhs rhs`
+        /// both lhs and rhs may be omitted.
         @"continue",
         /// `break :lhs rhs`
         /// both lhs and rhs may be omitted.

lib/std/zig/AstGen.zig

@@ -1140,7 +1140,7 @@ fn expr(gz: *GenZir, scope: *Scope, ri: ResultInfo, node: Ast.Node.Index) InnerE
         .error_set_decl => return errorSetDecl(gz, ri, node),
         .array_access => return arrayAccess(gz, scope, ri, node),
         .@"comptime" => return comptimeExprAst(gz, scope, ri, node),
-        .@"switch", .switch_comma => return switchExpr(gz, scope, ri.br(), node),
+        .@"switch", .switch_comma => return switchExpr(gz, scope, ri.br(), node, tree.fullSwitch(node).?),
 
         .@"nosuspend" => return nosuspendExpr(gz, scope, ri, node),
         .@"suspend" => return suspendExpr(gz, scope, node),
@@ -2154,6 +2154,11 @@ fn breakExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index) Inn
                     if (break_label != 0) {
                         if (block_gz.label) |*label| {
                             if (try astgen.tokenIdentEql(label.token, break_label)) {
+                                const maybe_switch_tag = astgen.instructions.items(.tag)[@intFromEnum(label.block_inst)];
+                                switch (maybe_switch_tag) {
+                                    .switch_block, .switch_block_ref => return astgen.failNode(node, "cannot break from switch", .{}),
+                                    else => {},
+                                }
                                 label.used = true;
                                 break :blk label.block_inst;
                             }
@@ -2228,6 +2233,11 @@ fn continueExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index)
     const tree = astgen.tree;
     const node_datas = tree.nodes.items(.data);
     const break_label = node_datas[node].lhs;
+    const rhs = node_datas[node].rhs;
+
+    if (break_label == 0 and rhs != 0) {
+        return astgen.failNode(node, "cannot continue with operand without label", .{});
+    }
 
     // Look for the label in the scope.
     var scope = parent_scope;
@@ -2252,15 +2262,51 @@ fn continueExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index)
                 if (break_label != 0) blk: {
                     if (gen_zir.label) |*label| {
                         if (try astgen.tokenIdentEql(label.token, break_label)) {
+                            const maybe_switch_tag = astgen.instructions.items(.tag)[@intFromEnum(label.block_inst)];
+                            if (rhs != 0) switch (maybe_switch_tag) {
+                                .switch_block, .switch_block_ref => {},
+                                else => return astgen.failNode(node, "cannot continue loop with operand", .{}),
+                            } else switch (maybe_switch_tag) {
+                                .switch_block, .switch_block_ref => return astgen.failNode(node, "cannot continue switch without operand", .{}),
+                                else => {},
+                            }
+
                             label.used = true;
                             break :blk;
                         }
                     }
                     // found continue but either it has a different label, or no label
                     scope = gen_zir.parent;
                     continue;
+                } else if (gen_zir.label) |label| {
+                    // This `continue` is unlabeled. If the gz we've found corresponds to a labeled
+                    // `switch`, ignore it and continue to parent scopes.
+                    switch (astgen.instructions.items(.tag)[@intFromEnum(label.block_inst)]) {
+                        .switch_block, .switch_block_ref => {
+                            scope = gen_zir.parent;
+                            continue;
+                        },
+                        else => {},
+                    }
+                }
+
+                if (rhs != 0) {
+                    // We need to figure out the result info to use.
+                    // The type should match
+                    const operand = try reachableExpr(parent_gz, parent_scope, gen_zir.continue_result_info, rhs, node);
+
+                    try genDefers(parent_gz, scope, parent_scope, .normal_only);
+
+                    // As our last action before the continue, "pop" the error trace if needed
+                    if (!gen_zir.is_comptime)
+                        _ = try parent_gz.addRestoreErrRetIndex(.{ .block = continue_block }, .always, node);
+
+                    _ = try parent_gz.addBreakWithSrcNode(.switch_continue, continue_block, operand, rhs);
+                    return Zir.Inst.Ref.unreachable_value;
                 }
 
+                try genDefers(parent_gz, scope, parent_scope, .normal_only);
+
                 const break_tag: Zir.Inst.Tag = if (gen_zir.is_inline)
                     .break_inline
                 else
@@ -2278,12 +2324,7 @@ fn continueExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index)
             },
             .local_val => scope = scope.cast(Scope.LocalVal).?.parent,
             .local_ptr => scope = scope.cast(Scope.LocalPtr).?.parent,
-            .defer_normal => {
-                const defer_scope = scope.cast(Scope.Defer).?;
-                scope = defer_scope.parent;
-                try parent_gz.addDefer(defer_scope.index, defer_scope.len);
-            },
-            .defer_error => scope = scope.cast(Scope.Defer).?.parent,
+            .defer_normal, .defer_error => scope = scope.cast(Scope.Defer).?.parent,
             .namespace => break,
             .top => unreachable,
         }
@@ -2843,6 +2884,7 @@ fn addEnsureResult(gz: *GenZir, maybe_unused_result: Zir.Inst.Ref, statement: As
             .panic,
             .trap,
             .check_comptime_control_flow,
+            .switch_continue,
             => {
                 noreturn_src_node = statement;
                 break :b true;
@@ -7569,7 +7611,8 @@ fn switchExpr(
     parent_gz: *GenZir,
     scope: *Scope,
     ri: ResultInfo,
-    switch_node: Ast.Node.Index,
+    node: Ast.Node.Index,
+    switch_full: Ast.full.Switch,
 ) InnerError!Zir.Inst.Ref {
     const astgen = parent_gz.astgen;
     const gpa = astgen.gpa;
@@ -7578,14 +7621,13 @@ fn switchExpr(
     const node_tags = tree.nodes.items(.tag);
     const main_tokens = tree.nodes.items(.main_token);
     const token_tags = tree.tokens.items(.tag);
-    const operand_node = node_datas[switch_node].lhs;
-    const extra = tree.extraData(node_datas[switch_node].rhs, Ast.Node.SubRange);
-    const case_nodes = tree.extra_data[extra.start..extra.end];
+    const operand_node = switch_full.ast.condition;
+    const case_nodes = switch_full.ast.cases;
 
-    const need_rl = astgen.nodes_need_rl.contains(switch_node);
+    const need_rl = astgen.nodes_need_rl.contains(node);
     const block_ri: ResultInfo = if (need_rl) ri else .{
         .rl = switch (ri.rl) {
-            .ptr => .{ .ty = (try ri.rl.resultType(parent_gz, switch_node)).? },
+            .ptr => .{ .ty = (try ri.rl.resultType(parent_gz, node)).? },
             .inferred_ptr => .none,
             else => ri.rl,
         },
@@ -7596,11 +7638,16 @@ fn switchExpr(
     const LocTag = @typeInfo(ResultInfo.Loc).Union.tag_type.?;
     const need_result_rvalue = @as(LocTag, block_ri.rl) != @as(LocTag, ri.rl);
 
+    if (switch_full.label_token) |label_token| {
+        try astgen.checkLabelRedefinition(scope, label_token);
+    }
+
     // We perform two passes over the AST. This first pass is to collect information
     // for the following variables, make note of the special prong AST node index,
     // and bail out with a compile error if there are multiple special prongs present.
     var any_payload_is_ref = false;
     var any_has_tag_capture = false;
+    var any_non_inline_capture = false;
     var scalar_cases_len: u32 = 0;
     var multi_cases_len: u32 = 0;
     var inline_cases_len: u32 = 0;
@@ -7618,6 +7665,15 @@ fn switchExpr(
             if (token_tags[ident + 1] == .comma) {
                 any_has_tag_capture = true;
             }
+
+            // If the first capture is ignored, then there is no runtime-known
+            // capture, as the tag capture must be for an inline prong.
+            // This check isn't perfect, because for things like enums, the
+            // first prong *is* comptime-known for inline prongs! But such
+            // knowledge requires semantic analysis.
+            if (!mem.eql(u8, tree.tokenSlice(ident), "_")) {
+                any_non_inline_capture = true;
+            }
         }
         // Check for else/`_` prong.
         if (case.ast.values.len == 0) {
@@ -7637,7 +7693,7 @@ fn switchExpr(
                 );
             } else if (underscore_src) |some_underscore| {
                 return astgen.failNodeNotes(
-                    switch_node,
+                    node,
                     "else and '_' prong in switch expression",
                     .{},
                     &[_]u32{
@@ -7678,7 +7734,7 @@ fn switchExpr(
                 );
             } else if (else_src) |some_else| {
                 return astgen.failNodeNotes(
-                    switch_node,
+                    node,
                     "else and '_' prong in switch expression",
                     .{},
                     &[_]u32{
@@ -7727,6 +7783,12 @@ fn switchExpr(
     const raw_operand = try expr(parent_gz, scope, operand_ri, operand_node);
     const item_ri: ResultInfo = .{ .rl = .none };
 
+    // If this switch is labeled, it will have `continue`s targeting it, and thus we need the operand type
+    // to provide a result type.
+    const raw_operand_ty_ref = if (switch_full.label_token != null) t: {
+        break :t try parent_gz.addUnNode(.typeof, raw_operand, operand_node);
+    } else undefined;
+
     // This contains the data that goes into the `extra` array for the SwitchBlock/SwitchBlockMulti,
     // except the first cases_nodes.len slots are a table that indexes payloads later in the array, with
     // the special case index coming first, then scalar_case_len indexes, then multi_cases_len indexes
@@ -7748,7 +7810,22 @@ fn switchExpr(
     try emitDbgStmtForceCurrentIndex(parent_gz, operand_lc);
     // This gets added to the parent block later, after the item expressions.
     const switch_tag: Zir.Inst.Tag = if (any_payload_is_ref) .switch_block_ref else .switch_block;
-    const switch_block = try parent_gz.makeBlockInst(switch_tag, switch_node);
+    const switch_block = try parent_gz.makeBlockInst(switch_tag, node);
+
+    if (switch_full.label_token) |label_token| {
+        block_scope.continue_block = switch_block.toOptional();
+        block_scope.continue_result_info = .{
+            .rl = if (any_payload_is_ref)
+                .{ .ref_coerced_ty = raw_operand_ty_ref }
+            else
+                .{ .coerced_ty = raw_operand_ty_ref },
+        };
+
+        block_scope.label = .{
+            .token = label_token,
+            .block_inst = switch_block,
+        };
+    }
 
     // We re-use this same scope for all cases, including the special prong, if any.
     var case_scope = parent_gz.makeSubBlock(&block_scope.base);
@@ -7969,6 +8046,8 @@ fn switchExpr(
             .has_else = special_prong == .@"else",
             .has_under = special_prong == .under,
             .any_has_tag_capture = any_has_tag_capture,
+            .any_non_inline_capture = any_non_inline_capture,
+            .has_continue = switch_full.label_token != null,
             .scalar_cases_len = @intCast(scalar_cases_len),
         },
     });
@@ -8005,7 +8084,7 @@ fn switchExpr(
     }
 
     if (need_result_rvalue) {
-        return rvalue(parent_gz, ri, switch_block.toRef(), switch_node);
+        return rvalue(parent_gz, ri, switch_block.toRef(), node);
     } else {
         return switch_block.toRef();
     }
@@ -11894,6 +11973,7 @@ const GenZir = struct {
     continue_block: Zir.Inst.OptionalIndex = .none,
     /// Only valid when setBreakResultInfo is called.
     break_result_info: AstGen.ResultInfo = undefined,
+    continue_result_info: AstGen.ResultInfo = undefined,
 
     suspend_node: Ast.Node.Index = 0,
     nosuspend_node: Ast.Node.Index = 0,