3 files changed, 95 insertions, 5 deletions

diff --git a/tests/codegen_test.c b/tests/codegen_test.c
index 64a422d..5236268 100644
--- a/tests/codegen_test.c
+++ b/tests/codegen_test.c
@@ -26,7 +26,6 @@ static void test_generate_consts_and_data(struct TestContext* context)
 	check(context, strstr(buffer, "msg: db `hi`") != NULL);
 	check(context, strstr(buffer, ".len equ $ - msg") != NULL);
 	check(context, strstr(buffer, "section .text") != NULL);
-	check(context, strstr(buffer, "global _start") != NULL);
 
 	free_program(&program);
 }
@@ -48,7 +47,7 @@ static void test_generate_text(struct TestContext* context)
 	buffer[read] = '\0';
 	fclose(out);
 
-	check(context, strstr(buffer, "_start:") != NULL);
+	check(context, strstr(buffer, "main:") != NULL);
 	check(context, strstr(buffer, "mov rax, SYS_WRITE") != NULL);
 	check(context, strstr(buffer, "mov rsi, message") != NULL);
 	check(context, strstr(buffer, "mov rdx, message.len") != NULL);
@@ -228,8 +227,58 @@ static void test_generate_sized_store(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_generate_entry_and_bits(struct TestContext* context)
+{
+	struct Lexer with_entry = create_lexer("[bits: 64]\n[entry: main]\nproc main\n{\nsyscall\n}\nproc helper\n{\nsyscall\n}\n");
+	struct Program program;
+	check(context, parse_program(&with_entry, &program));
+
+	FILE* out = tmpfile();
+	generate_nasm(&program, out);
+	fflush(out);
+	rewind(out);
+
+	char buffer[1024];
+	size_t read = fread(buffer, 1, sizeof(buffer) - 1, out);
+	buffer[read] = '\0';
+	fclose(out);
+
+	check(context, strstr(buffer, "bits 64") != NULL);
+	check(context, strstr(buffer, "global main") != NULL);
+	check(context, strstr(buffer, "main:") != NULL);
+	check(context, strstr(buffer, "_start") == NULL);
+	// the entry proc does not return; the non-entry helper does
+	check(context, strstr(buffer, "helper:\n\tsyscall\n\tret") != NULL);
+
+	free_program(&program);
+}
+
+static void test_generate_no_entry(struct TestContext* context)
+{
+	struct Lexer no_entry = create_lexer("proc main\n{\nsyscall\n}\n");
+	struct Program program;
+	check(context, parse_program(&no_entry, &program));
+
+	FILE* out = tmpfile();
+	generate_nasm(&program, out);
+	fflush(out);
+	rewind(out);
+
+	char buffer[1024];
+	size_t read = fread(buffer, 1, sizeof(buffer) - 1, out);
+	buffer[read] = '\0';
+	fclose(out);
+
+	check(context, strstr(buffer, "global") == NULL);
+	check(context, strstr(buffer, "main:") != NULL);
+
+	free_program(&program);
+}
+
 void run_codegen_tests(struct TestContext* context)
 {
+	test_generate_entry_and_bits(context);
+	test_generate_no_entry(context);
 	test_generate_consts_and_data(context);
 	test_generate_text(context);
 	test_generate_if(context);
diff --git a/tests/parser_test.c b/tests/parser_test.c
index 61f8c37..3b60261 100644
--- a/tests/parser_test.c
+++ b/tests/parser_test.c
@@ -208,6 +208,32 @@ static void test_parse_sized_deref(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_parse_directives(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer("[bits: 32]\n[entry: kmain]\n[enable: logical_registers]\nproc kmain\n{\nsyscall\n}\n");
+	struct Program program;
+
+	check(context, parse_program(&lexer, &program));
+	check(context, program.config.bits == 32);
+	check(context, program.config.has_entry);
+	check(context, text_is(program.config.entry, "kmain"));
+	check(context, program.config.logical_registers);
+
+	free_program(&program);
+}
+
+static void test_parse_bad_directive(struct TestContext* context)
+{
+	struct Lexer bad_bits = create_lexer("[bits: 16]\nproc main\n{\nsyscall\n}\n");
+	struct Program program;
+	check(context, !parse_program(&bad_bits, &program));
+	free_program(&program);
+
+	struct Lexer bad_key = create_lexer("[target: nasm]\nproc main\n{\nsyscall\n}\n");
+	check(context, !parse_program(&bad_key, &program));
+	free_program(&program);
+}
+
 static void test_parse_errors(struct TestContext* context)
 {
 	struct Program program;
@@ -237,5 +263,7 @@ void run_parser_tests(struct TestContext* context)
 	test_parse_call(context);
 	test_parse_stack(context);
 	test_parse_sized_deref(context);
+	test_parse_directives(context);
+	test_parse_bad_directive(context);
 	test_parse_errors(context);
 }
diff --git a/tests/sema_test.c b/tests/sema_test.c
index a1946d9..b7d3805 100644
--- a/tests/sema_test.c
+++ b/tests/sema_test.c
@@ -24,9 +24,20 @@ static void test_valid_program(struct TestContext* context)
 	check(context, analyze_source("const N = 1\ndata msg = \"hi\"\nproc main\n{\nsyscall\n}\n"));
 }
 
-static void test_missing_main(struct TestContext* context)
+static void test_no_entry_is_ok(struct TestContext* context)
 {
-	check(context, !analyze_source("proc helper\n{\nsyscall\n}\n"));
+	// without an [entry: ...] directive there is no required entry point
+	check(context, analyze_source("proc helper\n{\nsyscall\n}\n"));
+}
+
+static void test_defined_entry(struct TestContext* context)
+{
+	check(context, analyze_source("[entry: start]\nproc start\n{\nsyscall\n}\n"));
+}
+
+static void test_undefined_entry(struct TestContext* context)
+{
+	check(context, !analyze_source("[entry: main]\nproc helper\n{\nsyscall\n}\n"));
 }
 
 static void test_duplicate_const(struct TestContext* context)
@@ -42,7 +53,9 @@ static void test_duplicate_across_kinds(struct TestContext* context)
 void run_sema_tests(struct TestContext* context)
 {
 	test_valid_program(context);
-	test_missing_main(context);
+	test_no_entry_is_ok(context);
+	test_defined_entry(context);
+	test_undefined_entry(context);
 	test_duplicate_const(context);
 	test_duplicate_across_kinds(context);
 }