From 32642ac9cb00b59fef97c1888e0424b0eb4db784 Mon Sep 17 00:00:00 2001
From: Andrew Kelley <superjoe30@gmail.com>
Date: Thu, 4 Feb 2016 02:49:12 -0700
Subject: [PATCH] for loop supports break and continue

See #51
---
 src/all_types.hpp    |  2 ++
 src/analyze.cpp      | 21 ++++++++++++++++++---
 src/codegen.cpp      | 13 +++++++++----
 test/self_hosted.zig | 20 ++++++++++++++++++--
 4 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/src/all_types.hpp b/src/all_types.hpp
index 6cf27163b..c6baee6a4 100644
--- a/src/all_types.hpp
+++ b/src/all_types.hpp
@@ -474,6 +474,7 @@ struct AstNodeWhileExpr {
     // populated by semantic analyzer
     bool condition_always_true;
     bool contains_break;
+    bool contains_continue;
     Expr resolved_expr;
     BlockContext *block_context;
 };
@@ -486,6 +487,7 @@ struct AstNodeForExpr {
 
     // populated by semantic analyzer
     bool contains_break;
+    bool contains_continue;
     Expr resolved_expr;
     VariableTableEntry *elem_var;
     VariableTableEntry *index_var;
diff --git a/src/analyze.cpp b/src/analyze.cpp
index 5aad26cf3..38bc656a1 100644
--- a/src/analyze.cpp
+++ b/src/analyze.cpp
@@ -3354,6 +3354,7 @@ static TypeTableEntry *analyze_for_expr(CodeGen *g, ImportTableEntry *import, Bl
     }
 
     BlockContext *child_context = new_block_context(node, context);
+    child_context->parent_loop_node = node;
 
     AstNode *elem_var_node = node->data.for_expr.elem_node;
     elem_var_node->block_context = child_context;
@@ -3385,8 +3386,13 @@ static TypeTableEntry *analyze_break_expr(CodeGen *g, ImportTableEntry *import,
 
     AstNode *loop_node = context->parent_loop_node;
     if (loop_node) {
-        assert(loop_node->type == NodeTypeWhileExpr);
-        loop_node->data.while_expr.contains_break = true;
+        if (loop_node->type == NodeTypeWhileExpr) {
+            loop_node->data.while_expr.contains_break = true;
+        } else if (loop_node->type == NodeTypeForExpr) {
+            loop_node->data.for_expr.contains_break = true;
+        } else {
+            zig_unreachable();
+        }
     } else {
         add_node_error(g, node, buf_sprintf("'break' expression outside loop"));
     }
@@ -3396,7 +3402,16 @@ static TypeTableEntry *analyze_break_expr(CodeGen *g, ImportTableEntry *import,
 static TypeTableEntry *analyze_continue_expr(CodeGen *g, ImportTableEntry *import, BlockContext *context,
         TypeTableEntry *expected_type, AstNode *node)
 {
-    if (!context->parent_loop_node) {
+    AstNode *loop_node = context->parent_loop_node;
+    if (loop_node) {
+        if (loop_node->type == NodeTypeWhileExpr) {
+            loop_node->data.while_expr.contains_continue = true;
+        } else if (loop_node->type == NodeTypeForExpr) {
+            loop_node->data.for_expr.contains_continue = true;
+        } else {
+            zig_unreachable();
+        }
+    } else {
         add_node_error(g, node, buf_sprintf("'continue' expression outside loop"));
     }
     return g->builtin_types.entry_unreachable;
diff --git a/src/codegen.cpp b/src/codegen.cpp
index 7dd91b6a6..5c07b3409 100644
--- a/src/codegen.cpp
+++ b/src/codegen.cpp
@@ -2089,6 +2089,7 @@ static LLVMValueRef gen_for_expr(CodeGen *g, AstNode *node) {
     LLVMBasicBlockRef cond_block = LLVMAppendBasicBlock(g->cur_fn->fn_value, "ForCond");
     LLVMBasicBlockRef body_block = LLVMAppendBasicBlock(g->cur_fn->fn_value, "ForBody");
     LLVMBasicBlockRef end_block = LLVMAppendBasicBlock(g->cur_fn->fn_value, "ForEnd");
+    LLVMBasicBlockRef continue_block = LLVMAppendBasicBlock(g->cur_fn->fn_value, "ForContinue");
 
     LLVMValueRef array_val = gen_array_base_ptr(g, node->data.for_expr.array_expr);
     add_debug_source_node(g, node);
@@ -2122,17 +2123,21 @@ static LLVMValueRef gen_for_expr(CodeGen *g, AstNode *node) {
     gen_assign_raw(g, node, BinOpTypeAssign, elem_var->value_ref, elem_val,
             elem_var->type, child_type);
     g->break_block_stack.append(end_block);
-    g->continue_block_stack.append(cond_block);
+    g->continue_block_stack.append(continue_block);
     gen_expr(g, node->data.for_expr.body);
     g->break_block_stack.pop();
     g->continue_block_stack.pop();
     if (get_expr_type(node->data.for_expr.body)->id != TypeTableEntryIdUnreachable) {
         add_debug_source_node(g, node);
-        LLVMValueRef new_index_val = LLVMBuildAdd(g->builder, index_val, one_const, "");
-        LLVMBuildStore(g->builder, new_index_val, index_ptr);
-        LLVMBuildBr(g->builder, cond_block);
+        LLVMBuildBr(g->builder, continue_block);
     }
 
+    LLVMPositionBuilderAtEnd(g->builder, continue_block);
+    add_debug_source_node(g, node);
+    LLVMValueRef new_index_val = LLVMBuildAdd(g->builder, index_val, one_const, "");
+    LLVMBuildStore(g->builder, new_index_val, index_ptr);
+    LLVMBuildBr(g->builder, cond_block);
+
     LLVMPositionBuilderAtEnd(g->builder, end_block);
     return nullptr;
 }
diff --git a/test/self_hosted.zig b/test/self_hosted.zig
index 763689286..82f4d553c 100644
--- a/test/self_hosted.zig
+++ b/test/self_hosted.zig
@@ -52,17 +52,33 @@ error SecondError;
 #attribute("test")
 fn constant_enum_with_payload() {
     should_be_empty(AnEnumWithPayload.Empty);
-    should_be_13(AnEnumWithPayload.Full(13));
+    should_be_not_empty(AnEnumWithPayload.Full(13));
 }
 
 fn should_be_empty(x: AnEnumWithPayload) {
     if (x != AnEnumWithPayload.Empty) unreachable{}
 }
 
-fn should_be_13(x: AnEnumWithPayload) {
+fn should_be_not_empty(x: AnEnumWithPayload) {
+    if (x == AnEnumWithPayload.Empty) unreachable{}
 }
 
 enum AnEnumWithPayload {
     Empty,
     Full: i32,
 }
+
+
+#attribute("test")
+fn continue_in_for_loop() {
+    const array = []i32 {1, 2, 3, 4, 5};
+    var sum : i32 = 0;
+    for (x, array) {
+        sum += x;
+        if (x < 3) {
+            continue;
+        }
+        break;
+    }
+    if (sum != 6) unreachable{}
+}