[PATCH olang 2/2] ast: permit multi declarations on translation unit

[Carlos Maniero <carlos@maniero.me>]

At this point the parser still parsing only a single function, but the
ast is ready to support multiple declarations.

Signed-off-by: Carlos Maniero <carlos@maniero.me>
---
 src/ast.c                   |  7 +++++--
 src/ast.h                   |  4 ++--
 src/checker.c               |  7 ++++++-
 src/codegen_linux_aarch64.c | 22 +++++++++++++++++++---
 src/codegen_linux_x86_64.c  | 22 +++++++++++++++++++---
 src/parser.c                |  6 +++++-
 src/pretty_print_ast.c      | 13 +++++++++++--
 tests/unit/parser_test.c    | 11 ++++++++---
 8 files changed, 75 insertions(+), 17 deletions(-)

diff --git a/src/ast.c b/src/ast.c
index f5fa483..bb74679 100644
--- a/src/ast.c
+++ b/src/ast.c
@@ -23,7 +23,7 @@
 #include "string_view.h"
 
 ast_node_t *
-ast_new_translation_unit(arena_t *arena, ast_node_t *fn_def)
+ast_new_translation_unit(arena_t *arena)
 {
     ast_node_t *node = (ast_node_t *)arena_alloc(arena, sizeof(ast_node_t));
     assert(node);
@@ -31,7 +31,10 @@ ast_new_translation_unit(arena_t *arena, ast_node_t *fn_def)
     node->kind = AST_NODE_TRANSLATION_UNIT;
     ast_translation_unit_t *translation_unit = &node->as_translation_unit;
 
-    translation_unit->fn = fn_def;
+    translation_unit->decls = (list_t *)arena_alloc(arena, sizeof(list_t));
+    assert(translation_unit->decls);
+
+    list_init(translation_unit->decls, arena);
 
     return node;
 }
diff --git a/src/ast.h b/src/ast.h
index 718c80f..6cfbfc0 100644
--- a/src/ast.h
+++ b/src/ast.h
@@ -48,7 +48,7 @@ typedef struct ast_block
 
 typedef struct ast_translation_unit
 {
-    ast_node_t *fn;
+    list_t *decls;
 } ast_translation_unit_t;
 
 typedef struct ast_fn_definition
@@ -146,7 +146,7 @@ typedef struct ast_node
 } ast_node_t;
 
 ast_node_t *
-ast_new_translation_unit(arena_t *arena, ast_node_t *fn_def);
+ast_new_translation_unit(arena_t *arena);
 
 ast_node_t *
 ast_new_node_fn_def(arena_t *arena, string_view_t id, string_view_t return_type, ast_node_t *block);
diff --git a/src/checker.c b/src/checker.c
index e9bfacb..814d052 100644
--- a/src/checker.c
+++ b/src/checker.c
@@ -54,7 +54,12 @@ populate_scope(checker_t *checker, scope_t *scope, ast_node_t *ast)
 {
     switch (ast->kind) {
         case AST_NODE_TRANSLATION_UNIT: {
-            populate_scope(checker, scope, ast->as_translation_unit.fn);
+            list_item_t *item = list_head(ast->as_translation_unit.decls);
+
+            while (item != NULL) {
+                populate_scope(checker, scope, (ast_node_t *)item->value);
+                item = list_next(item);
+            }
             return;
         }
 
diff --git a/src/codegen_linux_aarch64.c b/src/codegen_linux_aarch64.c
index e8ae729..d8187ab 100644
--- a/src/codegen_linux_aarch64.c
+++ b/src/codegen_linux_aarch64.c
@@ -49,10 +49,26 @@ codegen_linux_aarch64_emit_translation_unit(FILE *out, ast_node_t *node)
     assert(node->kind == AST_NODE_TRANSLATION_UNIT);
     ast_translation_unit_t translation_unit = node->as_translation_unit;
 
-    ast_fn_definition_t fn = translation_unit.fn->as_fn_def;
+    list_item_t *item = list_head(translation_unit.decls);
 
-    assert(string_view_eq_to_cstr(fn.id, "main"));
-    codegen_linux_aarch64_emit_function(out, &fn);
+    bool main_found = false;
+
+    while (item != NULL) {
+        ast_node_t *decl = (ast_node_t *)item->value;
+
+        if (decl->kind == AST_NODE_FN_DEF) {
+            ast_fn_definition_t fn = decl->as_fn_def;
+            codegen_linux_aarch64_emit_function(out, &fn);
+
+            main_found = main_found || string_view_eq_to_cstr(fn.id, "main");
+        } else {
+            assert(0 && "translation unit only supports function declarations");
+        }
+
+        item = list_next(item);
+    }
+
+    assert(main_found && "main function is required.");
 }
 
 static void
diff --git a/src/codegen_linux_x86_64.c b/src/codegen_linux_x86_64.c
index 37d4575..0173443 100644
--- a/src/codegen_linux_x86_64.c
+++ b/src/codegen_linux_x86_64.c
@@ -63,10 +63,26 @@ codegen_linux_x86_64_emit_translation_unit(codegen_x86_64_t *codegen, ast_node_t
     assert(node->kind == AST_NODE_TRANSLATION_UNIT);
     ast_translation_unit_t translation_unit = node->as_translation_unit;
 
-    ast_fn_definition_t fn = translation_unit.fn->as_fn_def;
+    list_item_t *item = list_head(translation_unit.decls);
 
-    assert(string_view_eq_to_cstr(fn.id, "main"));
-    codegen_linux_x86_64_emit_function(codegen, &fn);
+    bool main_found = false;
+
+    while (item != NULL) {
+        ast_node_t *decl = (ast_node_t *)item->value;
+
+        if (decl->kind == AST_NODE_FN_DEF) {
+            ast_fn_definition_t fn = decl->as_fn_def;
+            codegen_linux_x86_64_emit_function(codegen, &fn);
+
+            main_found = main_found || string_view_eq_to_cstr(fn.id, "main");
+        } else {
+            assert(0 && "translation unit only supports function declarations");
+        }
+
+        item = list_next(item);
+    }
+
+    assert(main_found && "main function is required.");
 }
 
 static void
diff --git a/src/parser.c b/src/parser.c
index 9332f6e..c79f3bd 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -80,7 +80,11 @@ parser_parse_translation_unit(parser_t *parser)
         return NULL;
     }
 
-    return ast_new_translation_unit(parser->arena, fn);
+    ast_node_t *translation_unit_node = ast_new_translation_unit(parser->arena);
+
+    list_append(translation_unit_node->as_translation_unit.decls, fn);
+
+    return translation_unit_node;
 }
 
 static ast_binary_op_kind_t
diff --git a/src/pretty_print_ast.c b/src/pretty_print_ast.c
index db646c5..8116e60 100644
--- a/src/pretty_print_ast.c
+++ b/src/pretty_print_ast.c
@@ -116,8 +116,17 @@ ast_node_to_pretty_print_node(ast_node_t *ast, arena_t *arena)
             pretty_print_node_t *node = pretty_print_node_new(arena);
             node->name = "Translation_Unit";
 
-            pretty_print_node_t *fn_node = ast_node_to_pretty_print_node(ast->as_translation_unit.fn, arena);
-            list_append(node->children, fn_node);
+            list_item_t *item = list_head(ast->as_translation_unit.decls);
+
+            while (item != NULL) {
+                ast_node_t *decl = (ast_node_t *)item->value;
+
+                pretty_print_node_t *fn_node = ast_node_to_pretty_print_node(decl, arena);
+                list_append(node->children, fn_node);
+
+                item = list_next(item);
+            }
+
             return node;
         }
         case AST_NODE_FN_DEF: {
diff --git a/tests/unit/parser_test.c b/tests/unit/parser_test.c
index 4e229be..a7c60d1 100644
--- a/tests/unit/parser_test.c
+++ b/tests/unit/parser_test.c
@@ -46,10 +46,15 @@ parse_translation_unit_test(const MunitParameter params[], void *user_data_or_fi
     assert_uint(translation_unit_node->kind, ==, AST_NODE_TRANSLATION_UNIT);
 
     ast_translation_unit_t translation_unit = translation_unit_node->as_translation_unit;
-    assert_not_null(translation_unit.fn);
-    assert_uint(translation_unit.fn->kind, ==, AST_NODE_FN_DEF);
 
-    ast_fn_definition_t fn = translation_unit.fn->as_fn_def;
+    assert_uint(list_size(translation_unit.decls), ==, 1);
+
+    ast_node_t *fn_node = (ast_node_t *)list_head(translation_unit.decls)->value;
+
+    assert_not_null(fn_node);
+    assert_uint(fn_node->kind, ==, AST_NODE_FN_DEF);
+
+    ast_fn_definition_t fn = fn_node->as_fn_def;
     assert_memory_equal(fn.id.size, fn.id.chars, "main");
     assert_memory_equal(fn.return_type.size, fn.return_type.chars, "u32");
 
-- 
2.34.1