Now supporting local variables! :^)

6 files changed, 92 insertions, 29 deletions

diff --git a/cup/ast.c b/cup/ast.c
index f9b24d1..473af90 100644
--- a/cup/ast.c
+++ b/cup/ast.c
@@ -15,8 +15,7 @@ char *data_type_to_str(DataType type)
 void print_type_to_file(FILE *out, Type type)
 {
     fprintf(out, "%s", data_type_to_str(type.type));
-    for (int i = 0; i < type.indirection; i++)
-    {
+    for (int i = 0; i < type.indirection; i++) {
         fprintf(out, "*");
     }
 }
@@ -99,7 +98,7 @@ static void do_print_ast(Node *node, int depth)
     } else if (node->type == AST_FUNC) {
         dump_func(node, depth);
     } else if (node->type == AST_LITERAL) {
-        printf("(literal %d)\n", node->literal.as_int);
+        printf("%d\n", node->literal.as_int);
     } else if (node->type == AST_RETURN) {
         printf("return\n");
         do_print_ast(node->unary_expr, depth + 1);
@@ -110,13 +109,16 @@ static void do_print_ast(Node *node, int depth)
         printf("%s\n", node_type_to_str(node->type));
         do_print_ast(node->binary.left, depth + 1);
         do_print_ast(node->binary.right, depth + 1);
+    } else if (node->type == AST_VAR) {
+        assert(node->variable && node->variable->name);
+        printf("%s\n", node->variable->name);
     } else if (node->type == AST_VARDECL) {
-        printf("var %s (", node->var.name);
-        print_type_to_file(stdout, node->var.type);
+        printf("var %s (", node->var_decl.var.name);
+        print_type_to_file(stdout, node->var_decl.var.type);
         printf(")");
-        if (node->var.value != NULL) {
+        if (node->var_decl.value != NULL) {
             printf(" = \n");
-            do_print_ast(node->var.value, depth + 1);
+            do_print_ast(node->var_decl.value, depth + 1);
         } else {
             printf("\n");
         }
@@ -136,8 +138,8 @@ void dump_func(Node *node, int depth)
     if (node->func.num_locals > 0) {
         printf("\n locals: \n");
         for (int i = 0; i < node->func.num_locals; i++) {
-            printf(" - `%s`, offset: %lld (", node->func.locals[i].name, node->func.locals[i].offset);
-            print_type_to_file(stdout, node->func.locals[i].type);
+            printf(" - `%s`, offset: %lld (", node->func.locals[i]->name, node->func.locals[i]->offset);
+            print_type_to_file(stdout, node->func.locals[i]->type);
             printf(")\n");
         }
     }
diff --git a/cup/ast.h b/cup/ast.h
index c4e9836..4b18c82 100644
--- a/cup/ast.h
+++ b/cup/ast.h
@@ -25,6 +25,7 @@
   F(OP_ASSIGN, "=")                                                            \
   F(AST_LITERAL, "literal")                                                    \
   F(AST_VARDECL, "variable decl")                                              \
+  F(AST_VAR, "variable")                                                       \
   F(AST_RETURN, "return")                                                      \
   F(AST_FUNC, "func")                                                          \
   F(AST_PROGRAM, "program")                                                    \
@@ -81,7 +82,7 @@ typedef struct ast_node {
             Type return_type;
             Node *body;
 
-            Variable *locals;
+            Variable **locals;
             int num_locals;
 
             int cur_stack_offset;
@@ -101,11 +102,13 @@ typedef struct ast_node {
             };
         } literal;
 
+        // FIXME: Different struct for assignment?
         struct {
-            char *name;
-            Type type;
+            Variable var;
             Node *value;
-        } var;
+        } var_decl;
+
+        Variable *variable;
     };
 } Node;
 
diff --git a/cup/common.h b/cup/common.h
new file mode 100644
index 0000000..522483f
--- /dev/null
+++ b/cup/common.h
@@ -0,0 +1,6 @@
+#pragma once
+
+#include <stdbool.h>
+#include <stdint.h>
+
+typedef int64_t i64;
diff --git a/cup/generator.c b/cup/generator.c
index c0889cc..1575b57 100644
--- a/cup/generator.c
+++ b/cup/generator.c
@@ -8,6 +8,7 @@
 #include <assert.h>
 
 static int label_counter = 0;
+static Node *current_function = NULL;
 
 // The evaluated expression is stored into `rax`
 void generate_expr_into_rax(Node *expr, FILE *out)
@@ -18,7 +19,11 @@ void generate_expr_into_rax(Node *expr, FILE *out)
         assert(expr->literal.type.type == TYPE_INT);
         fprintf(out, "    mov rax, %d\n", expr->literal.as_int);
 
-    } else if (expr->type == OP_NEG) {
+    } else if (expr->type == AST_VAR) {
+        i64 offset = expr->variable->offset;    
+        fprintf(out, "    mov rax, [rbp-%lld]\n", offset);
+
+    }  else if (expr->type == OP_NEG) {
         generate_expr_into_rax(expr->unary_expr, out);
         fprintf(out, "    neg rax\n");
 
@@ -167,7 +172,16 @@ void generate_statement(Node *stmt, FILE *out)
 {
     if (stmt->type == AST_RETURN) {
         generate_expr_into_rax(stmt->unary_expr, out);
+        // TODO: Only do this if we have local variables
+        fprintf(out, "    mov rsp, rbp\n");
+        fprintf(out, "    pop rbp\n");
         fprintf(out, "    ret\n");
+    } else if (stmt->type == AST_VARDECL) {
+        if (stmt->var_decl.value) {
+            generate_expr_into_rax(stmt->var_decl.value, out);
+            i64 offset = stmt->var_decl.var.offset;
+            fprintf(out, "    mov [rbp-%lld], rax\n", offset);
+        }
     } else {
         fprintf(stderr, "Unsupported statement type in generate_statement: %s\n", node_type_to_str(stmt->type));
         exit(1);
@@ -181,14 +195,24 @@ void generate_block(Node *block, FILE *out)
         generate_statement(block->block.children[i], out);
 }
 
-void generate_function(Node *func, FILE *out)
+void generate_function_header(Node *func, FILE *out)
 {
     assert(func->type == AST_FUNC);
     fprintf(out, "global %s\n", func->func.name);
     fprintf(out, "%s:\n", func->func.name);
-    generate_block(func->func.body, out);
+    // TODO: Only do this if we have local variables
+    fprintf(out, "    push rbp\n");
+    fprintf(out, "    mov rbp, rsp\n");
+    fprintf(out, "    sub rsp, %d\n", func->func.cur_stack_offset);
 }
 
+void generate_function(Node *func, FILE *out)
+{
+    assert(func->type == AST_FUNC);
+    current_function = func;
+    generate_function_header(func, out);
+    generate_block(func->func.body, out);
+}
 
 void generate_asm(Node *root, FILE *out)
 {
diff --git a/cup/parser.c b/cup/parser.c
index 8c0c29c..ecd12f9 100644
--- a/cup/parser.c
+++ b/cup/parser.c
@@ -1,6 +1,7 @@
 #include "parser.h"
 #include "utils.h"
 #include <stdlib.h>
+#include <string.h>
 #include <assert.h>
 
 static Node *current_function = NULL;
@@ -47,7 +48,16 @@ NodeType binary_token_to_op(TokenType type)
     }
 }
 
-
+Variable *find_local_variable(Token *token)
+{
+    assert_token(*token, TOKEN_IDENTIFIER);
+    for (int i = 0; i < current_function->func.num_locals; i++) {
+        if (strcmp(current_function->func.locals[i]->name, token->value.as_string) == 0) {
+            return current_function->func.locals[i];
+        }
+    }
+    return NULL;
+}
 
 void Node_add_child(Node *parent, Node *child)
 {
@@ -97,26 +107,29 @@ Node *parse_expression(Lexer *);
 Node *parse_var_declaration(Lexer *lexer)
 {
     Token token = assert_token(Lexer_next(lexer), TOKEN_LET);
+    // TODO: Reuse this for globals? Or maybe just make a new function?
+    if (!current_function || current_function->type != AST_FUNC)
+        die_location(token.loc, "Variable declaration outside of function");
+    
+    token = assert_token(Lexer_next(lexer), TOKEN_IDENTIFIER);
+    if (find_local_variable(&token) != NULL)
+        die_location(token.loc, "Variable `%s` already declared", token.value.as_string);
+    
     Node *node = Node_new(AST_VARDECL);
-    node->var.name = assert_token(Lexer_next(lexer), TOKEN_IDENTIFIER).value.as_string;
+    node->var_decl.var.name = token.value.as_string;
     assert_token(Lexer_next(lexer), TOKEN_COLON);
-    node->var.type = parse_type(lexer);
+    node->var_decl.var.type = parse_type(lexer);
     assert_token(Lexer_next(lexer), TOKEN_ASSIGN);
-    node->var.value = parse_expression(lexer);
+    node->var_decl.value = parse_expression(lexer);
     assert_token(Lexer_next(lexer), TOKEN_SEMICOLON);
 
-    // Add variable to current function
-    if (!current_function || current_function->type != AST_FUNC)
-        die_location(token.loc, "Variable declaration outside of function");
+    // Set offset for variable
+    node->var_decl.var.offset = current_function->func.cur_stack_offset;
 
     int new_len = (current_function->func.num_locals + 1);
     int var_size = 8; // TODO: Compute sizes based on different types
-    current_function->func.locals = realloc(current_function->func.locals, sizeof(Variable) * new_len);
-    current_function->func.locals[current_function->func.num_locals] = (Variable) {
-        .name = node->var.name,
-        .type = node->var.type,
-        .offset = current_function->func.cur_stack_offset,
-    };
+    current_function->func.locals = realloc(current_function->func.locals, sizeof(Variable *) * new_len);
+    current_function->func.locals[current_function->func.num_locals] = &node->var_decl.var;
     current_function->func.num_locals++;
     current_function->func.cur_stack_offset += var_size;
 
@@ -146,6 +159,13 @@ Node *parse_factor(Lexer *lexer)
         assert_token(Lexer_next(lexer), TOKEN_CLOSE_PAREN);
     } else if (token.type == TOKEN_INTLIT) {
         expr = parse_literal(lexer);
+    } else if (token.type == TOKEN_IDENTIFIER) {
+        Lexer_next(lexer);
+        Variable *var = find_local_variable(&token); 
+        if (var == NULL)
+            die_location(token.loc, "Could not find variable `%s`", token.value.as_string);
+        expr = Node_new(AST_VAR);
+        expr->variable = var;
     } else {
         die_location(token.loc, ": Expected token found in parse_factor: `%s`", token_type_to_str(token.type));
         exit(1);
diff --git a/examples/variables.cup b/examples/variables.cup
new file mode 100644
index 0000000..3ee8ed2
--- /dev/null
+++ b/examples/variables.cup
@@ -0,0 +1,8 @@
+
+
+fn main() {
+    let x: int = 5 * 10;
+    let y: int = x + x;
+    let z: int = y - x / 2;
+    return z;
+}