stage2: break and continue out of loops

2 files changed, 109 insertions, 44 deletions

diff --git a/src/Module.zig b/src/Module.zig
index 20cb7bf195..3806c0dea8 100644
--- a/src/Module.zig
+++ b/src/Module.zig
@@ -796,6 +796,10 @@ pub const Scope = struct {
         /// The first N instructions in a function body ZIR are arg instructions.
         instructions: std.ArrayListUnmanaged(*zir.Inst) = .{},
         label: ?Label = null,
+        break_block: ?*zir.Inst.Block = null,
+        continue_block: ?*zir.Inst.Block = null,
+        /// only valid if label != null or (continue_block and break_block) != null
+        break_result_loc: astgen.ResultLoc = undefined,
 
         pub const Label = struct {
             token: ast.TokenIndex,
diff --git a/src/astgen.zig b/src/astgen.zig
index 622a18c443..c4ac5b5fc9 100644
--- a/src/astgen.zig
+++ b/src/astgen.zig
@@ -261,6 +261,7 @@ pub fn expr(mod: *Module, scope: *Scope, rl: ResultLoc, node: *ast.Node) InnerEr
         .Block => return rlWrapVoid(mod, scope, rl, node, try blockExpr(mod, scope, node.castTag(.Block).?)),
         .LabeledBlock => return labeledBlockExpr(mod, scope, rl, node.castTag(.LabeledBlock).?, .block),
         .Break => return rlWrap(mod, scope, rl, try breakExpr(mod, scope, node.castTag(.Break).?)),
+        .Continue => return rlWrap(mod, scope, rl, try continueExpr(mod, scope, node.castTag(.Continue).?)),
         .PtrType => return rlWrap(mod, scope, rl, try ptrType(mod, scope, node.castTag(.PtrType).?)),
         .GroupedExpression => return expr(mod, scope, rl, node.castTag(.GroupedExpression).?.expr),
         .ArrayType => return rlWrap(mod, scope, rl, try arrayType(mod, scope, node.castTag(.ArrayType).?)),
@@ -291,7 +292,6 @@ pub fn expr(mod: *Module, scope: *Scope, rl: ResultLoc, node: *ast.Node) InnerEr
         .StructInitializer => return mod.failNode(scope, node, "TODO implement astgen.expr for .StructInitializer", .{}),
         .StructInitializerDot => return mod.failNode(scope, node, "TODO implement astgen.expr for .StructInitializerDot", .{}),
         .Suspend => return mod.failNode(scope, node, "TODO implement astgen.expr for .Suspend", .{}),
-        .Continue => return mod.failNode(scope, node, "TODO implement astgen.expr for .Continue", .{}),
         .AnyType => return mod.failNode(scope, node, "TODO implement astgen.expr for .AnyType", .{}),
         .FnProto => return mod.failNode(scope, node, "TODO implement astgen.expr for .FnProto", .{}),
         .ContainerDecl => return mod.failNode(scope, node, "TODO implement astgen.expr for .ContainerDecl", .{}),
@@ -339,48 +339,97 @@ fn breakExpr(mod: *Module, parent_scope: *Scope, node: *ast.Node.ControlFlowExpr
     const tree = parent_scope.tree();
     const src = tree.token_locs[node.ltoken].start;
 
-    if (node.getLabel()) |break_label| {
-        // Look for the label in the scope.
-        var scope = parent_scope;
-        while (true) {
-            switch (scope.tag) {
-                .gen_zir => {
-                    const gen_zir = scope.cast(Scope.GenZIR).?;
+    // Look for the label in the scope.
+    var scope = parent_scope;
+    while (true) {
+        switch (scope.tag) {
+            .gen_zir => {
+                const gen_zir = scope.cast(Scope.GenZIR).?;
+
+                const block_inst = blk: {
+                    if (node.getLabel()) |break_label| {
+                        if (gen_zir.label) |label| {
+                            if (try tokenIdentEql(mod, parent_scope, label.token, break_label)) {
+                                break :blk label.block_inst;
+                            }
+                        }
+                    } else if (gen_zir.break_block) |inst| {
+                        break :blk inst;
+                    }
+                    scope = gen_zir.parent;
+                    continue;
+                };
+
+                if (node.getRHS()) |rhs| {
+                    // Most result location types can be forwarded directly; however
+                    // if we need to write to a pointer which has an inferred type,
+                    // proper type inference requires peer type resolution on the block's
+                    // break operand expressions.
+                    const branch_rl: ResultLoc = switch (gen_zir.break_result_loc) {
+                        .discard, .none, .ty, .ptr, .ref => gen_zir.break_result_loc,
+                        .inferred_ptr, .bitcasted_ptr, .block_ptr => .{ .block_ptr = block_inst },
+                    };
+                    const operand = try expr(mod, parent_scope, branch_rl, rhs);
+                    return try addZIRInst(mod, parent_scope, src, zir.Inst.Break, .{
+                        .block = block_inst,
+                        .operand = operand,
+                    }, .{});
+                } else {
+                    return try addZIRInst(mod, parent_scope, src, zir.Inst.BreakVoid, .{
+                        .block = block_inst,
+                    }, .{});
+                }
+            },
+            .local_val => scope = scope.cast(Scope.LocalVal).?.parent,
+            .local_ptr => scope = scope.cast(Scope.LocalPtr).?.parent,
+            else => if (node.getLabel()) |break_label| {
+                const label_name = try identifierTokenString(mod, parent_scope, break_label);
+                return mod.failTok(parent_scope, break_label, "label not found: '{}'", .{label_name});
+            } else {
+                return mod.failTok(parent_scope, src, "break expression outside loop", .{});
+            },
+        }
+    }
+}
+
+fn continueExpr(mod: *Module, parent_scope: *Scope, node: *ast.Node.ControlFlowExpression) InnerError!*zir.Inst {
+    const tree = parent_scope.tree();
+    const src = tree.token_locs[node.ltoken].start;
+
+    // Look for the label in the scope.
+    var scope = parent_scope;
+    while (true) {
+        switch (scope.tag) {
+            .gen_zir => {
+                const gen_zir = scope.cast(Scope.GenZIR).?;
+                const continue_block = gen_zir.continue_block orelse {
+                    scope = gen_zir.parent;
+                    continue;
+                };
+                if (node.getLabel()) |break_label| blk: {
                     if (gen_zir.label) |label| {
                         if (try tokenIdentEql(mod, parent_scope, label.token, break_label)) {
-                            if (node.getRHS()) |rhs| {
-                                // Most result location types can be forwarded directly; however
-                                // if we need to write to a pointer which has an inferred type,
-                                // proper type inference requires peer type resolution on the block's
-                                // break operand expressions.
-                                const branch_rl: ResultLoc = switch (label.result_loc) {
-                                    .discard, .none, .ty, .ptr, .ref => label.result_loc,
-                                    .inferred_ptr, .bitcasted_ptr, .block_ptr => .{ .block_ptr = label.block_inst },
-                                };
-                                const operand = try expr(mod, parent_scope, branch_rl, rhs);
-                                return try addZIRInst(mod, scope, src, zir.Inst.Break, .{
-                                    .block = label.block_inst,
-                                    .operand = operand,
-                                }, .{});
-                            } else {
-                                return try addZIRInst(mod, scope, src, zir.Inst.BreakVoid, .{
-                                    .block = label.block_inst,
-                                }, .{});
-                            }
+                            break :blk;
                         }
                     }
+                    // found continue but either it has a different label, or no label
                     scope = gen_zir.parent;
-                },
-                .local_val => scope = scope.cast(Scope.LocalVal).?.parent,
-                .local_ptr => scope = scope.cast(Scope.LocalPtr).?.parent,
-                else => {
-                    const label_name = try identifierTokenString(mod, parent_scope, break_label);
-                    return mod.failTok(parent_scope, break_label, "label not found: '{}'", .{label_name});
-                },
-            }
+                    continue;
+                }
+
+                return addZIRInst(mod, parent_scope, src, zir.Inst.BreakVoid, .{
+                    .block = continue_block,
+                }, .{});
+            },
+            .local_val => scope = scope.cast(Scope.LocalVal).?.parent,
+            .local_ptr => scope = scope.cast(Scope.LocalPtr).?.parent,
+            else => if (node.getLabel()) |break_label| {
+                const label_name = try identifierTokenString(mod, parent_scope, break_label);
+                return mod.failTok(parent_scope, break_label, "label not found: '{}'", .{label_name});
+            } else {
+                return mod.failTok(parent_scope, src, "continue expression outside loop", .{});
+            },
         }
-    } else {
-        return mod.failNode(parent_scope, &node.base, "TODO implement break from loop", .{});
     }
 }
 
@@ -426,11 +475,11 @@ fn labeledBlockExpr(
         .decl = parent_scope.decl().?,
         .arena = gen_zir.arena,
         .instructions = .{},
+        .break_result_loc = rl,
         // TODO @as here is working around a stage1 miscompilation bug :(
         .label = @as(?Scope.GenZIR.Label, Scope.GenZIR.Label{
             .token = block_node.label,
             .block_inst = block_inst,
-            .result_loc = rl,
         }),
     };
     defer block_scope.instructions.deinit(mod.gpa);
@@ -1289,9 +1338,6 @@ fn whileExpr(mod: *Module, scope: *Scope, rl: ResultLoc, while_node: *ast.Node.W
         }
     }
 
-    if (while_node.label) |tok|
-        return mod.failTok(scope, tok, "TODO labeled while", .{});
-
     if (while_node.inline_token) |tok|
         return mod.failTok(scope, tok, "TODO inline while", .{});
 
@@ -1308,6 +1354,7 @@ fn whileExpr(mod: *Module, scope: *Scope, rl: ResultLoc, while_node: *ast.Node.W
         .decl = expr_scope.decl,
         .arena = expr_scope.arena,
         .instructions = .{},
+        .break_result_loc = rl,
     };
     defer loop_scope.instructions.deinit(mod.gpa);
 
@@ -1348,6 +1395,14 @@ fn whileExpr(mod: *Module, scope: *Scope, rl: ResultLoc, while_node: *ast.Node.W
     const while_block = try addZIRInstBlock(mod, scope, while_src, .block, .{
         .instructions = try expr_scope.arena.dupe(*zir.Inst, expr_scope.instructions.items),
     });
+    loop_scope.break_block = while_block;
+    loop_scope.continue_block = cond_block;
+    if (while_node.label) |some| {
+        loop_scope.label =  @as(?Scope.GenZIR.Label, Scope.GenZIR.Label{
+            .token = some,
+            .block_inst = while_block,
+        });
+    }
 
     const then_src = tree.token_locs[while_node.body.lastToken()].start;
     var then_scope: Scope.GenZIR = .{
@@ -1414,9 +1469,6 @@ fn whileExpr(mod: *Module, scope: *Scope, rl: ResultLoc, while_node: *ast.Node.W
 }
 
 fn forExpr(mod: *Module, scope: *Scope, rl: ResultLoc, for_node: *ast.Node.For) InnerError!*zir.Inst {
-    if (for_node.label) |tok|
-        return mod.failTok(scope, tok, "TODO labeled for", .{});
-
     if (for_node.inline_token) |tok|
         return mod.failTok(scope, tok, "TODO inline for", .{});
 
@@ -1458,6 +1510,7 @@ fn forExpr(mod: *Module, scope: *Scope, rl: ResultLoc, for_node: *ast.Node.For)
         .decl = for_scope.decl,
         .arena = for_scope.arena,
         .instructions = .{},
+        .break_result_loc = rl,
     };
     defer loop_scope.instructions.deinit(mod.gpa);
 
@@ -1499,6 +1552,14 @@ fn forExpr(mod: *Module, scope: *Scope, rl: ResultLoc, for_node: *ast.Node.For)
     const for_block = try addZIRInstBlock(mod, scope, for_src, .block, .{
         .instructions = try for_scope.arena.dupe(*zir.Inst, for_scope.instructions.items),
     });
+    loop_scope.break_block = for_block;
+    loop_scope.continue_block = cond_block;
+    if (for_node.label) |some| {
+        loop_scope.label =  @as(?Scope.GenZIR.Label, Scope.GenZIR.Label{
+            .token = some,
+            .block_inst = for_block,
+        });
+    }
 
     // while body
     const then_src = tree.token_locs[for_node.body.lastToken()].start;