compiler: implement labeled switch/continue

4 files changed, 175 insertions, 32 deletions

diff --git a/lib/std/zig/Ast.zig b/lib/std/zig/Ast.zig
index 1f734cef63..b6f4ad68ee 100644
--- a/lib/std/zig/Ast.zig
+++ b/lib/std/zig/Ast.zig
@@ -1184,14 +1184,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) {
@@ -1895,6 +1888,15 @@ 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];
+    return tree.fullSwitchComponents(.{
+        .switch_token = tree.nodes.items(.main_token)[node],
+        .condition = data.lhs,
+        .sub_range = data.rhs,
+    });
+}
+
 pub fn switchCaseOne(tree: Ast, node: Node.Index) full.SwitchCase {
     const data = &tree.nodes.items(.data)[node];
     const values: *[1]Node.Index = &data.lhs;
@@ -2206,6 +2208,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);
@@ -2477,6 +2494,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),
@@ -2829,6 +2853,17 @@ pub const full = struct {
         };
     };
 
+    pub const Switch = struct {
+        ast: Components,
+        label_token: ?TokenIndex,
+
+        pub const Components = struct {
+            switch_token: TokenIndex,
+            condition: Node.Index,
+            sub_range: Node.Index,
+        };
+    };
+
     pub const SwitchCase = struct {
         inline_token: ?TokenIndex,
         /// Points to the first token after the `|`. Will either be an identifier or
@@ -3287,7 +3322,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.
diff --git a/lib/std/zig/AstGen.zig b/lib/std/zig/AstGen.zig
index 9c27d1e036..92da810b30 100644
--- a/lib/std/zig/AstGen.zig
+++ b/lib/std/zig/AstGen.zig
@@ -1144,7 +1144,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),
@@ -2160,6 +2160,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;
                             }
@@ -2234,6 +2239,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;
@@ -2258,6 +2268,15 @@ 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;
                         }
@@ -2265,8 +2284,35 @@ fn continueExpr(parent_gz: *GenZir, parent_scope: *Scope, node: Ast.Node.Index)
                     // 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
@@ -2284,12 +2330,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,
         }
@@ -2881,6 +2922,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;
@@ -7546,7 +7588,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;
@@ -7555,14 +7598,14 @@ 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 operand_node = node_datas[node].lhs;
+    const extra = tree.extraData(node_datas[node].rhs, Ast.Node.SubRange);
     const case_nodes = tree.extra_data[extra.start..extra.end];
 
-    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,
         },
@@ -7573,11 +7616,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;
@@ -7595,6 +7643,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) {
@@ -7614,7 +7671,7 @@ fn switchExpr(
                 );
             } else if (underscore_src) |some_underscore| {
                 return astgen.failNodeNotes(
-                    switch_node,
+                    node,
                     "else and '_' prong in switch expression",
                     .{},
                     &[_]u32{
@@ -7655,7 +7712,7 @@ fn switchExpr(
                 );
             } else if (else_src) |some_else| {
                 return astgen.failNodeNotes(
-                    switch_node,
+                    node,
                     "else and '_' prong in switch expression",
                     .{},
                     &[_]u32{
@@ -7704,6 +7761,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
@@ -7725,7 +7788,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);
@@ -7946,6 +8024,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),
         },
     });
@@ -7982,7 +8062,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();
     }
@@ -11824,6 +11904,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,
diff --git a/lib/std/zig/Parse.zig b/lib/std/zig/Parse.zig
index 6f557a0f55..20e69845cb 100644
--- a/lib/std/zig/Parse.zig
+++ b/lib/std/zig/Parse.zig
@@ -924,7 +924,6 @@ fn expectContainerField(p: *Parse) !Node.Index {
 ///      / KEYWORD_errdefer Payload? BlockExprStatement
 ///      / IfStatement
 ///      / LabeledStatement
-///      / SwitchExpr
 ///      / VarDeclExprStatement
 fn expectStatement(p: *Parse, allow_defer_var: bool) Error!Node.Index {
     if (p.eatToken(.keyword_comptime)) |comptime_token| {
@@ -995,7 +994,6 @@ fn expectStatement(p: *Parse, allow_defer_var: bool) Error!Node.Index {
                 .rhs = try p.expectBlockExprStatement(),
             },
         }),
-        .keyword_switch => return p.expectSwitchExpr(),
         .keyword_if => return p.expectIfStatement(),
         .keyword_enum, .keyword_struct, .keyword_union => {
             const identifier = p.tok_i + 1;
@@ -1238,7 +1236,7 @@ fn expectIfStatement(p: *Parse) !Node.Index {
     });
 }
 
-/// LabeledStatement <- BlockLabel? (Block / LoopStatement)
+/// LabeledStatement <- BlockLabel? (Block / LoopStatement / SwitchExpr)
 fn parseLabeledStatement(p: *Parse) !Node.Index {
     const label_token = p.parseBlockLabel();
     const block = try p.parseBlock();
@@ -1247,6 +1245,9 @@ fn parseLabeledStatement(p: *Parse) !Node.Index {
     const loop_stmt = try p.parseLoopStatement();
     if (loop_stmt != 0) return loop_stmt;
 
+    const switch_expr = try p.parseSwitchExpr();
+    if (switch_expr != 0) return switch_expr;
+
     if (label_token != 0) {
         const after_colon = p.tok_i;
         const node = try p.parseTypeExpr();
@@ -2072,7 +2073,7 @@ fn expectTypeExpr(p: *Parse) Error!Node.Index {
 ///      / KEYWORD_break BreakLabel? Expr?
 ///      / KEYWORD_comptime Expr
 ///      / KEYWORD_nosuspend Expr
-///      / KEYWORD_continue BreakLabel?
+///      / KEYWORD_continue BreakLabel? Expr?
 ///      / KEYWORD_resume Expr
 ///      / KEYWORD_return Expr?
 ///      / BlockLabel? LoopExpr
@@ -2098,7 +2099,7 @@ fn parsePrimaryExpr(p: *Parse) !Node.Index {
                 .main_token = p.nextToken(),
                 .data = .{
                     .lhs = try p.parseBreakLabel(),
-                    .rhs = undefined,
+                    .rhs = try p.parseExpr(),
                 },
             });
         },
@@ -2627,7 +2628,6 @@ fn parseSuffixExpr(p: *Parse) !Node.Index {
 ///      / KEYWORD_anyframe
 ///      / KEYWORD_unreachable
 ///      / STRINGLITERAL
-///      / SwitchExpr
 ///
 /// ContainerDecl <- (KEYWORD_extern / KEYWORD_packed)? ContainerDeclAuto
 ///
@@ -2647,6 +2647,7 @@ fn parseSuffixExpr(p: *Parse) !Node.Index {
 /// LabeledTypeExpr
 ///     <- BlockLabel Block
 ///      / BlockLabel? LoopTypeExpr
+///      / BlockLabel? SwitchExpr
 ///
 /// LoopTypeExpr <- KEYWORD_inline? (ForTypeExpr / WhileTypeExpr)
 fn parsePrimaryTypeExpr(p: *Parse) !Node.Index {
@@ -2753,6 +2754,10 @@ fn parsePrimaryTypeExpr(p: *Parse) !Node.Index {
                     p.tok_i += 2;
                     return p.parseWhileTypeExpr();
                 },
+                .keyword_switch => {
+                    p.tok_i += 2;
+                    return p.expectSwitchExpr();
+                },
                 .l_brace => {
                     p.tok_i += 2;
                     return p.parseBlock();
@@ -3029,8 +3034,17 @@ fn parseWhileTypeExpr(p: *Parse) !Node.Index {
 }
 
 /// SwitchExpr <- KEYWORD_switch LPAREN Expr RPAREN LBRACE SwitchProngList RBRACE
+fn parseSwitchExpr(p: *Parse) !Node.Index {
+    const switch_token = p.eatToken(.keyword_switch) orelse return null_node;
+    return p.expectSwitchSuffix(switch_token);
+}
+
 fn expectSwitchExpr(p: *Parse) !Node.Index {
     const switch_token = p.assertToken(.keyword_switch);
+    return p.expectSwitchSuffix(switch_token);
+}
+
+fn expectSwitchSuffix(p: *Parse, switch_token: TokenIndex) !Node.Index {
     _ = try p.expectToken(.l_paren);
     const expr_node = try p.expectExpr();
     _ = try p.expectToken(.r_paren);
diff --git a/lib/std/zig/Zir.zig b/lib/std/zig/Zir.zig
index af4ddaad6a..0186b45f74 100644
--- a/lib/std/zig/Zir.zig
+++ b/lib/std/zig/Zir.zig
@@ -314,6 +314,9 @@ pub const Inst = struct {
         /// break instruction in a block, and the target block is the parent.
         /// Uses the `break` union field.
         break_inline,
+        /// Branch from within a switch case to the case specified by the operand.
+        /// Uses the `break` union field. `block_inst` refers to a `switch_block` or `switch_block_ref`.
+        switch_continue,
         /// Checks that comptime control flow does not happen inside a runtime block.
         /// Uses the `un_node` union field.
         check_comptime_control_flow,
@@ -1273,6 +1276,7 @@ pub const Inst = struct {
                 .panic,
                 .trap,
                 .check_comptime_control_flow,
+                .switch_continue,
                 => true,
             };
         }
@@ -1512,6 +1516,7 @@ pub const Inst = struct {
                 .break_inline,
                 .condbr,
                 .condbr_inline,
+                .switch_continue,
                 .compile_error,
                 .ret_node,
                 .ret_load,
@@ -1597,6 +1602,7 @@ pub const Inst = struct {
                 .bool_br_or = .pl_node,
                 .@"break" = .@"break",
                 .break_inline = .@"break",
+                .switch_continue = .@"break",
                 .check_comptime_control_flow = .un_node,
                 .for_len = .pl_node,
                 .call = .pl_node,
@@ -2288,6 +2294,7 @@ pub const Inst = struct {
         },
         @"break": struct {
             operand: Ref,
+            /// Index of a `Break` payload.
             payload_index: u32,
         },
         dbg_stmt: LineColumn,
@@ -2945,9 +2952,13 @@ pub const Inst = struct {
             has_under: bool,
             /// 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
+            /// be comptime-known via `inline`.
+            any_non_inline_capture: bool,
+            has_continue: bool,
             scalar_cases_len: ScalarCasesLen,
 
-            pub const ScalarCasesLen = u28;
+            pub const ScalarCasesLen = u26;
 
             pub fn specialProng(bits: Bits) SpecialProng {
                 const has_else: u2 = @intFromBool(bits.has_else);
@@ -3750,6 +3761,7 @@ fn findDeclsInner(
         .bool_br_or,
         .@"break",
         .break_inline,
+        .switch_continue,
         .check_comptime_control_flow,
         .builtin_call,
         .cmp_lt,