From 0a0050695254411047122a760ca22ad31c6bd9f8 Mon Sep 17 00:00:00 2001
From: Mustafa Quraish <mustafaq9@gmail.com>
Date: Sat, 5 Feb 2022 21:26:56 -0500
Subject: [compiler.cup] Support for+while loops

---
 compiler/codegen.cup | 26 ++++++++++++++++++++++++++
 compiler/parser.cup  | 53 ++++++++++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/compiler/codegen.cup b/compiler/codegen.cup
index 975844d..ffe7c1f 100644
--- a/compiler/codegen.cup
+++ b/compiler/codegen.cup
@@ -290,6 +290,32 @@ fn generate_statement(node: Node*) {
             generate_block(node.d.conditional.els);
         emit_asm(".if"); emit_num(label); emit_asm(":\n");
 
+    } else if (node.typ == AST_WHILE) {
+        let label = ++gen_label_counter;
+        emit_asm(".loop_s"); emit_num(label); emit_asm(":\n");
+        generate_expr_into_rax(node.d.looop.cond);
+        emit_asm("    cmp rax, 0\n");
+        emit_asm("    je .loop_e"); emit_num(label); emit_asm("\n");
+        generate_statement(node.d.looop.body);
+        emit_asm("    jmp .loop_s"); emit_num(label); emit_asm("\n");
+        emit_asm(".loop_e"); emit_num(label); emit_asm(":\n");
+
+    } else if (node.typ == AST_FOR) {
+        let label = ++gen_label_counter;
+        if (node.d.looop.init)
+            generate_statement(node.d.looop.init);
+        emit_asm(".loop_s"); emit_num(label); emit_asm(":\n");
+        if (node.d.looop.cond) {
+            generate_expr_into_rax(node.d.looop.cond);
+            emit_asm("    cmp rax, 0\n");
+            emit_asm("    je .loop_e"); emit_num(label); emit_asm("\n");
+        }
+        generate_statement(node.d.looop.body);
+        if (node.d.looop.step)
+            generate_statement(node.d.looop.step);
+        emit_asm("    jmp .loop_s"); emit_num(label); emit_asm("\n");
+        emit_asm(".loop_e"); emit_num(label); emit_asm(":\n");
+
     } else {
         // Default to a simple expression statement
         generate_expr_into_rax(node);
diff --git a/compiler/parser.cup b/compiler/parser.cup
index 1335fba..914c4e6 100644
--- a/compiler/parser.cup
+++ b/compiler/parser.cup
@@ -413,6 +413,7 @@ fn add_variable_to_current_block(var: Variable*) {
 }
 
 fn parse_var_declaration(lexer: Lexer*): Node* {
+
     let token: Token;
     lexer_next_assert(lexer, &token, TOKEN_LET);
     lexer_next_assert(lexer, &token, TOKEN_IDENTIFIER);
@@ -486,7 +487,48 @@ fn parse_function_params(lexer: Lexer*, func: Node*) {
 }
 
 fn parse_block(lexer: Lexer*): Node*;
+fn parse_statement(lexer: Lexer*): Node*;
+
+fn parse_for_loop(lexer: Lexer*): Node* {
+    let token: Token;
+    lexer_next_assert(lexer, &token, TOKEN_FOR);
+
+    let looop = node_new(AST_FOR);
+    lexer_next_assert(lexer, &token, TOKEN_OPEN_PAREN);
+
+    // NOTE: We're going to put the for loop in it's own block
+    //       so that any declarations in the init of the loop
+    //       can only be referenced within the loop.
+    let node = node_new(AST_BLOCK);
+    node.d.block.children = vector_new_sized(1);
+    node.d.block.locals = vector_new_sized(1);
+    block_add_child(node, looop);
+    block_stack_push(node);
 
+    // All of the expressions in the for loop are optional
+
+    lexer_peek(lexer, &token);
+    if (token.typ == TOKEN_LET)
+        looop.d.looop.init = parse_var_declaration(lexer);
+    else if (token.typ != TOKEN_SEMICOLON)
+        looop.d.looop.init = parse_expression(lexer);
+    lexer_next_assert(lexer, &token, TOKEN_SEMICOLON);
+
+    lexer_peek(lexer, &token);
+    if (token.typ != TOKEN_SEMICOLON)
+        looop.d.looop.cond = parse_expression(lexer);
+    lexer_next_assert(lexer, &token, TOKEN_SEMICOLON);
+
+    lexer_peek(lexer, &token);
+    if (token.typ != TOKEN_CLOSE_PAREN)
+        looop.d.looop.step = parse_expression(lexer);
+    lexer_next_assert(lexer, &token, TOKEN_CLOSE_PAREN);
+
+    looop.d.looop.body = parse_statement(lexer);
+    block_stack_pop();
+
+    return node;
+}
 
 fn parse_statement(lexer: Lexer*): Node* {
     let node: Node*;
@@ -524,9 +566,16 @@ fn parse_statement(lexer: Lexer*): Node* {
             node.d.conditional.els = parse_statement(lexer);
         }
     } else if (token.typ == TOKEN_WHILE) {
-        die("while is not implemented yet");
+        lexer_next(lexer, &token);
+        node = node_new(AST_WHILE);
+        lexer_next_assert(lexer, &token, TOKEN_OPEN_PAREN);
+        node.d.looop.cond = parse_expression(lexer);
+        lexer_next_assert(lexer, &token, TOKEN_CLOSE_PAREN);
+        node.d.looop.body = parse_statement(lexer);
+
     } else if (token.typ == TOKEN_FOR) {
-        die("for is not implemented yet");
+        node = parse_for_loop(lexer);
+
     } else if (token.typ == TOKEN_DEFER) {
         die("defer is not implemented yet");
     } else if (token.typ == TOKEN_LET) {
-- 
cgit v1.2.3