[PATCH olang v1 2/2] parser: add var definition and reference support

[Johnny Richard <johnny@johnnyrichard.com>]

In order to get the var definition parsed we also introduce the
reference node to get a functional ast test.

This first implementation assumes that all variables will be of type u32
and we don't allow variable definitions without assignment, this
assignment issue will be addressed in the future.

Signed-off-by: Johnny Richard <johnny@johnnyrichard.com>
---
 src/ast.c                                     | 28 ++++++++++
 src/ast.h                                     | 22 ++++++++
 src/parser.c                                  | 52 +++++++++++++++++++
 src/pretty_print_ast.c                        | 25 +++++++++
 .../integration/tests/0024_var_definition.ol  | 14 ++++-
 5 files changed, 139 insertions(+), 2 deletions(-)

diff --git a/src/ast.c b/src/ast.c
index 4b252ae..d2fd2fa 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -52,6 +52,22 @@ ast_new_node_fn_def(arena_t *arena, string_view_t identifier, type_t return_type
     return node_fn_def;
 }
 
+ast_node_t *
+ast_new_node_var_def(arena_t *arena, string_view_t identifier, type_t type, ast_node_t *value)
+{
+    ast_node_t *node_var_def = (ast_node_t *)arena_alloc(arena, sizeof(ast_node_t));
+    assert(node_var_def);
+
+    node_var_def->kind = AST_NODE_VAR_DEF;
+    ast_var_definition_t *var_def = &node_var_def->as_var_def;
+
+    var_def->identifier = identifier;
+    var_def->type = type;
+    var_def->value = value;
+
+    return node_var_def;
+}
+
 ast_node_t *
 ast_new_node_bin_op(arena_t *arena, ast_binary_op_kind_t kind, ast_node_t *lhs, ast_node_t *rhs)
 {
@@ -79,6 +95,18 @@ ast_new_node_literal_u32(arena_t *arena, uint32_t value)
     return node_literal;
 }
 
+ast_node_t *
+ast_new_node_ref(arena_t *arena, string_view_t identifier)
+{
+    ast_node_t *node_ref = (ast_node_t *)arena_alloc(arena, sizeof(ast_node_t));
+    assert(node_ref);
+
+    node_ref->kind = AST_NODE_REF;
+    node_ref->as_ref.identifier = identifier;
+
+    return node_ref;
+}
+
 ast_node_t *
 ast_new_node_return_stmt(arena_t *arena)
 {
diff --git a/src/ast.h b/src/ast.h
index 9bf3723..94ee6b4 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -30,10 +30,12 @@ typedef enum
     AST_NODE_PROGRAM,
     AST_NODE_BLOCK,
     AST_NODE_FN_DEF,
+    AST_NODE_VAR_DEF,
     AST_NODE_BINARY_OP,
     AST_NODE_RETURN_STMT,
     AST_NODE_IF_STMT,
     AST_NODE_LITERAL,
+    AST_NODE_REF,
     AST_NODE_UNKNOWN
 } ast_node_kind_t;
 
@@ -59,6 +61,13 @@ typedef struct ast_fn_definition
     ast_node_t *block;
 } ast_fn_definition_t;
 
+typedef struct ast_var_definition
+{
+    string_view_t identifier;
+    type_t type;
+    ast_node_t *value;
+} ast_var_definition_t;
+
 typedef enum
 {
     AST_LITERAL_U32
@@ -73,6 +82,11 @@ typedef struct ast_literal
     };
 } ast_literal_t;
 
+typedef struct ast_ref
+{
+    string_view_t identifier;
+} ast_ref_t;
+
 typedef enum ast_binary_op_kind
 {
     AST_BINOP_ADDITION,
@@ -121,8 +135,10 @@ typedef struct ast_node
     {
         ast_program_t as_program;
         ast_fn_definition_t as_fn_def;
+        ast_var_definition_t as_var_def;
         ast_binary_op_t as_bin_op;
         ast_literal_t as_literal;
+        ast_ref_t as_ref;
         ast_block_t as_block;
         ast_return_stmt_t as_return_stmt;
         ast_if_stmt_t as_if_stmt;
@@ -135,12 +151,18 @@ ast_new_program(arena_t *arena, ast_node_t *fn_def);
 ast_node_t *
 ast_new_node_fn_def(arena_t *arena, string_view_t identifier, type_t return_type, ast_node_t *block);
 
+ast_node_t *
+ast_new_node_var_def(arena_t *arena, string_view_t identifier, type_t type, ast_node_t *value);
+
 ast_node_t *
 ast_new_node_bin_op(arena_t *arena, ast_binary_op_kind_t kind, ast_node_t *lhs, ast_node_t *rhs);
 
 ast_node_t *
 ast_new_node_literal_u32(arena_t *arena, uint32_t value);
 
+ast_node_t *
+ast_new_node_ref(arena_t *arena, string_view_t identifier);
+
 ast_node_t *
 ast_new_node_return_stmt(arena_t *arena);
 
diff --git a/src/parser.c b/src/parser.c
index 8f49f5e..d71500f 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -45,6 +45,9 @@ parser_parse_return_stmt(parser_t *parser);
 static ast_node_t *
 parser_parse_if_stmt(parser_t *parser);
 
+static ast_node_t *
+parser_parse_var_def(parser_t *parser);
+
 ast_node_t *
 parser_parse_fn_definition(parser_t *parser);
 
@@ -233,6 +236,9 @@ parser_parse_factor(parser_t *parser)
         case TOKEN_NUMBER:
             return ast_new_node_literal_u32(parser->arena, string_view_to_u32(token.value));
 
+        case TOKEN_IDENTIFIER:
+            return ast_new_node_ref(parser->arena, token.value);
+
         case TOKEN_OPAREN: {
             ast_node_t *expr = parser_parse_expr(parser);
             if (expr == NULL) {
@@ -348,7 +354,15 @@ StartLoop:
             node = parser_parse_if_stmt(parser);
             break;
         }
+        case TOKEN_VAR: {
+            node = parser_parse_var_def(parser);
+            break;
+        }
+        case TOKEN_CCURLY: {
+            goto EndLoop;
+        }
         default: {
+            // FIXME: write a better error message
             goto EndLoop;
         }
     }
@@ -443,6 +457,44 @@ parser_parse_if_stmt(parser_t *parser)
     return node_if_stmt;
 }
 
+static ast_node_t *
+parser_parse_var_def(parser_t *parser)
+{
+    if (!skip_expected_token(parser, TOKEN_VAR)) {
+        return NULL;
+    }
+
+    token_t identifier_token;
+    if (!expected_next_token(parser, &identifier_token, TOKEN_IDENTIFIER)) {
+        return NULL;
+    }
+
+    if (!skip_expected_token(parser, TOKEN_COLON)) {
+        return NULL;
+    }
+
+    type_t var_type;
+    if (!parser_parse_type(parser, &var_type)) {
+        return NULL;
+    }
+
+    // FIXME: assignment must be optional according to the spec
+    if (!skip_expected_token(parser, TOKEN_EQ)) {
+        return NULL;
+    }
+
+    ast_node_t *expr = parser_parse_expr(parser);
+    if (expr == NULL) {
+        return NULL;
+    }
+
+    ast_node_t *var_node = ast_new_node_var_def(parser->arena, identifier_token.value, var_type, expr);
+
+    skip_line_feeds(parser->lexer);
+
+    return var_node;
+}
+
 static bool
 skip_expected_token(parser_t *parser, token_kind_t expected_kind)
 {
diff --git a/src/pretty_print_ast.c b/src/pretty_print_ast.c
index 3c12716..d6eef67 100644
--- a/src/pretty_print_ast.c
+++ b/src/pretty_print_ast.c
@@ -198,6 +198,31 @@ ast_node_to_pretty_print_node(ast_node_t *ast, arena_t *arena)
 
             return node;
         }
+        case AST_NODE_VAR_DEF: {
+            pretty_print_node_t *node = pretty_print_node_new(arena);
+            ast_var_definition_t var = ast->as_var_def;
+
+            char name[256];
+            sprintf(name, "Var_Definition <name:" SV_FMT "> <kind:u32>", SV_ARG(var.identifier));
+            node->name = (char *)arena_alloc(arena, sizeof(char) * (strlen(name) + 1));
+            strcpy(node->name, name);
+
+            pretty_print_node_t *child = ast_node_to_pretty_print_node(var.value, arena);
+            list_append(node->children, child);
+
+            return node;
+        }
+        case AST_NODE_REF: {
+            pretty_print_node_t *node = pretty_print_node_new(arena);
+            ast_ref_t ref = ast->as_ref;
+
+            char name[256];
+            sprintf(name, "Reference <name:" SV_FMT ">", SV_ARG(ref.identifier));
+            node->name = (char *)arena_alloc(arena, sizeof(char) * (strlen(name) + 1));
+            strcpy(node->name, name);
+
+            return node;
+        }
         case AST_NODE_BINARY_OP: {
             pretty_print_node_t *node = pretty_print_node_new(arena);
             ast_binary_op_t binop = ast->as_bin_op;
diff --git a/tests/integration/tests/0024_var_definition.ol b/tests/integration/tests/0024_var_definition.ol
index d0b68da..241fc1a 100644
--- a/tests/integration/tests/0024_var_definition.ol
+++ b/tests/integration/tests/0024_var_definition.ol
@@ -14,10 +14,20 @@
 # along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
 fn main(): u32 {
-  var code: u32 = 0;
-  return code;
+  var code: u32 = 0
+  return code
 }
 
 # TEST test_contains_tokens WITH
 # ./tests/0024_var_definition.ol:17:3: <var>
 # END
+
+# TEST test_ast WITH
+# Translation_Unit
+# `-Function_Definition <name:main> <return:0>
+#   `-Block
+#     |-Var_Definition <name:code> <kind:u32>
+#     | `-Literal <kind:u32> <value:0>
+#     `-Return_Statement
+#       `-Reference <name:code>
+# END
-- 
2.46.0