loop support

7 files changed, 179 insertions, 19 deletions

diff --git a/ast.c b/ast.c
index 69173bd..7c7432b 100644
--- a/ast.c
+++ b/ast.c
@@ -116,12 +116,38 @@ static void expr_destroy(struct expr_node* node) {
 static void if_destroy(struct if_node* node) {
     expr_destroy(node->cond);
     free(node->cond);
+
     stmt_destroy(node->true_branch);
     free(node->true_branch);
+
     if (node->false_branch != NULL) {
         stmt_destroy(node->false_branch);
         free(node->false_branch);
     }
+
+    scope_destroy(node->scope);
+    free(node->scope);
+}
+
+static void loop_destroy(struct loop_node* node) {
+    if (node->init != NULL) {
+        expr_destroy(node->init);
+        free(node->init);
+    }
+    if (node->cond != NULL) {
+        expr_destroy(node->cond);
+        free(node->cond);
+    }
+    if (node->incr != NULL) {
+        expr_destroy(node->incr);
+        free(node->incr);
+    }
+
+    stmt_destroy(node->body);
+    free(node->body);
+
+    scope_destroy(node->scope);
+    free(node->scope);
 }
 
 static void stmt_destroy(struct stmt_node* node) {
@@ -142,6 +168,9 @@ static void stmt_destroy(struct stmt_node* node) {
         case STMT_IF:
             if_destroy(&node->inner.if_);
             break;
+        case STMT_LOOP:
+            loop_destroy(&node->inner.loop);
+            break;
     }
 }
 
diff --git a/ast.h b/ast.h
index 6ff95df..71c400c 100644
--- a/ast.h
+++ b/ast.h
@@ -139,6 +139,14 @@ struct if_node {
     struct scope* scope;
 };
 
+struct loop_node {
+    struct expr_node* init; /* null if absent */
+    struct expr_node* cond; /* null if absent */
+    struct expr_node* incr; /* null if absent */
+    struct stmt_node* body;
+    struct scope* scope;
+};
+
 struct stmt_node {
     enum {
         STMT_EMPTY,
@@ -147,6 +155,7 @@ struct stmt_node {
         STMT_RETURN,
         STMT_GROUP,
         STMT_IF,
+        STMT_LOOP,
     } type;
     union {
         struct expr_node expr;
@@ -154,6 +163,7 @@ struct stmt_node {
         struct return_node return_;
         struct group_node group;
         struct if_node if_;
+        struct loop_node loop;
     } inner;
 
     struct stmt_node* next;
diff --git a/codegen.c b/codegen.c
index 79deb46..f9dbc15 100644
--- a/codegen.c
+++ b/codegen.c
@@ -35,20 +35,26 @@ static const struct storage_location MULDIV_OVERFLOW_LOC = {
 static struct scope* scope;
 static const struct fn_decl_node* active_fn;
 static unsigned long long branch_counter = 0;
+static unsigned long long loop_counter = 0;
 
-static void enter_scope(struct scope* child_scope) {
+static void enter_scope(
+    struct scope* child_scope,
+    unsigned long long bp_offset
+) {
     if (child_scope == NULL ||
             child_scope->next_out != scope)
         CGEN_PANIC("enter_scope: scopes are misaligned");
 
     scope = child_scope;
+    scope->bp_offset = bp_offset;
 }
 
-static void exit_scope(struct scope* child_scope) {
+static void exit_scope(struct scope* child_scope, bool save_bp_offset) {
     if (child_scope != scope || child_scope->next_out == NULL)
         CGEN_PANIC("exit_scope: scopes are misaligned");
 
     scope = child_scope->next_out;
+    if (save_bp_offset) scope->bp_offset = child_scope->bp_offset;
 }
 
 static struct lval_def allocate_register(const struct type_def* type) {
@@ -305,6 +311,11 @@ static struct var_def* emit_var_decl(
     const struct var_decl_node* node
 ) {
     struct lval_def var_dst = make_stack_lval(outfile, node->type.def_ref);
+
+    fprintf(outfile, "\t; '%s' lives in: ", node->def_ref->name);
+    emit_storage_loc(outfile, &var_dst.loc, var_dst.type->sz);
+    fprintf(outfile, "\n");
+
     struct var_def* var_def = node->def_ref;
     var_def->loc = var_dst.loc;
     return var_def;
@@ -562,39 +573,63 @@ static void emit_group_contents(FILE* outfile, const struct group_node* node) {
 }
 
 static void emit_group(FILE* outfile, const struct group_node* node) {
-    enter_scope(node->scope);
-    scope->bp_offset = scope->next_out->bp_offset; /* don't reset bp */
+    enter_scope(node->scope, scope->bp_offset);
 
     emit_group_contents(outfile, node);
 
     /* don't reset sp because alloca needs to work */
     scope->next_out->bp_offset = scope->bp_offset;
-    exit_scope(node->scope);
+    exit_scope(node->scope, true);
 }
 
 static void emit_if(FILE* outfile, const struct if_node* node) {
-    enter_scope(node->scope);
+    enter_scope(node->scope, scope->bp_offset);
 
     struct lval_def cond_result =
         allocate_temporary(outfile, node->cond->resolved_type);
     emit_expr(outfile, node->cond, &cond_result);
     emit_cmp_zero(outfile, &cond_result);
 
-    unsigned long long false_branch_num = branch_counter++;
-    fprintf(outfile, "\tjz branch%lld\n", false_branch_num);
+    unsigned long long branch_num = ++branch_counter;
+    fprintf(outfile, "\tjz branch_false@%lld\n", branch_num);
     emit_stmt(outfile, node->true_branch);
 
     if (node->false_branch == NULL) {
-        fprintf(outfile, "branch%lld:\n", false_branch_num);
+        fprintf(outfile, "branch_false@%lld:\n", branch_num);
     } else {
-        unsigned long long end_branch_num = branch_counter++;
-        fprintf(outfile, "\tjmp branch%lld\n", end_branch_num);
-        fprintf(outfile, "branch%lld:\n", false_branch_num);
+        fprintf(outfile, "\tjmp branch_done@%lld\n", branch_num);
+        fprintf(outfile, "branch_false@%lld:\n", branch_num);
         emit_stmt(outfile, node->false_branch);
-        fprintf(outfile, "branch%lld:\n", end_branch_num);
+        fprintf(outfile, "branch_done@%lld:\n", branch_num);
+    }
+
+    exit_scope(node->scope, true);
+}
+
+static void emit_loop(FILE* outfile, const struct loop_node* node) {
+    enter_scope(node->scope, scope->bp_offset);
+
+    if (node->init != NULL) emit_expr(outfile, node->init, NULL);
+
+    unsigned long long loop_num = ++loop_counter;
+    if (node->cond != NULL) {
+        struct lval_def cond_dst =
+            allocate_temporary(outfile, node->cond->resolved_type);
+        fprintf(outfile, "loop_head@%lld:\n", loop_num);
+        emit_expr(outfile, node->cond, &cond_dst);
+        emit_cmp_zero(outfile, &cond_dst);
+        fprintf(outfile, "\tjz loop_done@%lld\n", loop_num);
+    } else {
+        fprintf(outfile, "loop_head@%lld:\n", loop_num);
     }
 
-    exit_scope(node->scope);
+    emit_stmt(outfile, node->body);
+    if (node->incr != NULL) emit_expr(outfile, node->incr, NULL);
+
+    fprintf(outfile, "\tjmp loop_head@%lld\n", loop_num);
+    fprintf(outfile, "loop_done@%lld:\n", loop_num);
+
+    exit_scope(node->scope, true);
 }
 
 static void emit_stmt(FILE* outfile, const struct stmt_node* node) {
@@ -616,11 +651,14 @@ static void emit_stmt(FILE* outfile, const struct stmt_node* node) {
         case STMT_IF:
             emit_if(outfile, &node->inner.if_);
             break;
+        case STMT_LOOP:
+            emit_loop(outfile, &node->inner.loop);
+            break;
     }
 }
 
 static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
-    enter_scope(node->scope);
+    enter_scope(node->scope, 0);
 
     if (active_fn != NULL)
         CGEN_PANIC(
@@ -660,9 +698,9 @@ static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
         emit_mov(outfile, &arg_dst, &arg_src);
     }
 
-    enter_scope(node->body.scope);
+    enter_scope(node->body.scope, scope->bp_offset);
     emit_group_contents(outfile, &node->body);
-    exit_scope(node->body.scope);
+    exit_scope(node->body.scope, true);
 
     fprintf(outfile, RETURN_LABEL_FMT ":\n", node->name);
     fprintf(outfile, "\tmov rsp, rbp\n");
@@ -670,7 +708,7 @@ static void emit_fn_decl(FILE* outfile, const struct fn_decl_node* node) {
     fprintf(outfile, "\tret\n");
 
     active_fn = NULL;
-    exit_scope(node->scope);
+    exit_scope(node->scope, false);
 }
 
 static void emit_root_node(FILE* outfile, const struct root_node* node) {
diff --git a/parser.c b/parser.c
index 4779e73..b094780 100644
--- a/parser.c
+++ b/parser.c
@@ -128,7 +128,7 @@ static void parse_expr_assign(struct expr_node* p_node) {
                 .type = LVAL_VAR_REF,
                 .inner.var_ref = p_node->inner.var_ref,
             };
-            return;
+            break;
         default:
             PARSER_PANIC("expression is not assignable");
     }
@@ -346,6 +346,55 @@ static void parse_if(struct if_node* p_node) {
     scope_pop(&scope);
 }
 
+static void parse_while(struct loop_node* p_node) {
+    expect_kw("while");
+
+    scope_push(&scope);
+    p_node->scope = scope;
+
+    expect(TK_LPAREN);
+    p_node->cond = protected_alloc(sizeof(struct expr_node));
+    parse_expr(p_node->cond);
+    expect(TK_RPAREN);
+
+    p_node->body = protected_alloc(sizeof(struct stmt_node));
+    parse_stmt(p_node->body);
+
+    scope_pop(&scope);
+}
+
+static void parse_for(struct loop_node* p_node) {
+    expect_kw("for");
+
+    scope_push(&scope);
+    p_node->scope = scope;
+
+    expect(TK_LPAREN);
+    peek_or_panic();
+    if (tok.type != TK_SEMI) {
+        p_node->init = protected_alloc(sizeof(struct expr_node));
+        parse_expr(p_node->init);
+    }
+    expect(TK_SEMI);
+    peek_or_panic();
+    if (tok.type != TK_SEMI) {
+        p_node->cond = protected_alloc(sizeof(struct expr_node));
+        parse_expr(p_node->cond);
+    }
+    expect(TK_SEMI);
+    peek_or_panic();
+    if (tok.type != TK_RPAREN) {
+        p_node->incr = protected_alloc(sizeof(struct expr_node));
+        parse_expr(p_node->incr);
+    }
+    expect(TK_RPAREN);
+
+    p_node->body = protected_alloc(sizeof(struct stmt_node));
+    parse_stmt(p_node->body);
+
+    scope_pop(&scope);
+}
+
 static void parse_stmt_assign(struct stmt_node* p_node) {
     peek_or_panic();
     if (tok.type != TK_ASSIGN) return;
@@ -385,6 +434,14 @@ static void parse_stmt(struct stmt_node* p_node) {
                 p_node->type = STMT_IF;
                 parse_if(&p_node->inner.if_);
                 return;
+            } else if (strcmp(tok.data.ident, "while") == 0) {
+                p_node->type = STMT_LOOP;
+                parse_while(&p_node->inner.loop);
+                return;
+            } else if (strcmp(tok.data.ident, "for") == 0) {
+                p_node->type = STMT_LOOP;
+                parse_for(&p_node->inner.loop);
+                return;
             } else if (strcmp(tok.data.ident, "return") == 0) {
                 p_node->type = STMT_RETURN;
                 parse_return(&p_node->inner.return_);
diff --git a/test/forever.c b/test/forever.c
new file mode 100644
index 0000000..acb73e1
--- /dev/null
+++ b/test/forever.c
@@ -0,0 +1,5 @@
+int main (int argc, char** argv) {
+    for(;;);
+    while (1);
+    for(;1;);
+}
diff --git a/test/loop.c b/test/loop.c
new file mode 100644
index 0000000..a17444f
--- /dev/null
+++ b/test/loop.c
@@ -0,0 +1,7 @@
+int main (int argc, char** argv) {
+    int i;
+    for (i = 0; i; i = i + 1) {
+        i = i / 2;
+    }
+    return i;
+}
diff --git a/type_checker.c b/type_checker.c
index aabd7c9..d8534e2 100644
--- a/type_checker.c
+++ b/type_checker.c
@@ -184,6 +184,17 @@ static void type_check_if(struct if_node* node) {
     exit_scope(node->scope);
 }
 
+static void type_check_loop(struct loop_node* node) {
+    enter_scope(node->scope);
+
+    if (node->init != NULL) type_check_expr(node->init);
+    if (node->cond != NULL) type_check_expr(node->cond);
+    if (node->incr != NULL) type_check_expr(node->incr);
+    type_check_stmt(node->body);
+
+    exit_scope(node->scope);
+}
+
 static void type_check_stmt(struct stmt_node* node) {
     switch (node->type) {
         case STMT_EMPTY:
@@ -203,6 +214,9 @@ static void type_check_stmt(struct stmt_node* node) {
         case STMT_IF:
             type_check_if(&node->inner.if_);
             break;
+        case STMT_LOOP:
+            type_check_loop(&node->inner.loop);
+            break;
     }
 }