4 files changed, 70 insertions, 0 deletions

diff --git a/tests/codegen_test.c b/tests/codegen_test.c
index c56ca7a..61e441f 100644
--- a/tests/codegen_test.c
+++ b/tests/codegen_test.c
@@ -399,8 +399,36 @@ static void test_generate_load(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_generate_enum_struct(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer(
+		"enum Color\n{\nRed,\nGreen,\nBlue\n}\n"
+		"struct Point\n{\nx: qword\ny: qword\nflag: byte\n}\n"
+		"proc main\n{\nrax = Color.Blue\nrbx = Point.y\nrcx = Point.flag\nrdx = Point.size\n}\n");
+	struct Program program;
+	check(context, parse_program(&lexer, &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, "mov rax, 2") != NULL);   // Color.Blue -> 2
+	check(context, strstr(buffer, "mov rbx, 8") != NULL);   // Point.y -> 8
+	check(context, strstr(buffer, "mov rcx, 16") != NULL);  // Point.flag -> 16
+	check(context, strstr(buffer, "mov rdx, 17") != NULL);  // Point.size -> 17
+
+	free_program(&program);
+}
+
 void run_codegen_tests(struct TestContext* context)
 {
+	test_generate_enum_struct(context);
 	test_generate_load(context);
 	test_generate_logical_registers(context);
 	test_generate_logical_disabled(context);
diff --git a/tests/lexer_test.c b/tests/lexer_test.c
index 7260d6c..38644d9 100644
--- a/tests/lexer_test.c
+++ b/tests/lexer_test.c
@@ -62,6 +62,13 @@ static void test_size_keywords(struct TestContext* context)
 	check(context, scan_token(&lexer).type == TOKEN_QWORD);
 }
 
+static void test_enum_struct_keywords(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer("enum struct");
+	check(context, scan_token(&lexer).type == TOKEN_ENUM);
+	check(context, scan_token(&lexer).type == TOKEN_STRUCT);
+}
+
 static void test_line_counting(struct TestContext* context)
 {
 	struct Lexer lexer = create_lexer("a\nb\nc");
@@ -89,6 +96,7 @@ void run_lexer_tests(struct TestContext* context)
 	test_literals(context);
 	test_keywords(context);
 	test_size_keywords(context);
+	test_enum_struct_keywords(context);
 	test_line_counting(context);
 	test_comments(context);
 }
diff --git a/tests/parser_test.c b/tests/parser_test.c
index 5a45e9e..71e1433 100644
--- a/tests/parser_test.c
+++ b/tests/parser_test.c
@@ -267,6 +267,27 @@ static void test_parse_register_size_suffix(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_parse_enum_struct(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer(
+		"enum Color\n{\nRed,\nGreen,\nBlue\n}\nstruct Point\n{\nx\ny: byte\n}\n");
+	struct Program program;
+
+	check(context, parse_program(&lexer, &program));
+	check(context, program.enum_count == 1);
+	check(context, text_is(program.enums[0].name, "Color"));
+	check(context, program.enums[0].member_count == 3);
+	check(context, text_is(program.enums[0].members[1], "Green"));
+
+	check(context, program.struct_count == 1);
+	check(context, text_is(program.structs[0].name, "Point"));
+	check(context, program.structs[0].field_count == 2);
+	check(context, program.structs[0].fields[0].size == STORE_SIZE_QWORD);
+	check(context, program.structs[0].fields[1].size == STORE_SIZE_BYTE);
+
+	free_program(&program);
+}
+
 static void test_parse_errors(struct TestContext* context)
 {
 	struct Program program;
@@ -300,5 +321,6 @@ void run_parser_tests(struct TestContext* context)
 	test_parse_directives(context);
 	test_parse_bad_directive(context);
 	test_parse_register_size_suffix(context);
+	test_parse_enum_struct(context);
 	test_parse_errors(context);
 }
diff --git a/tests/sema_test.c b/tests/sema_test.c
index 2d6e187..e89f608 100644
--- a/tests/sema_test.c
+++ b/tests/sema_test.c
@@ -97,6 +97,17 @@ static void test_deref_needs_register(struct TestContext* context)
 	check(context, analyze_source("proc main\n{\nrax = ^rsi\n}\n"));
 }
 
+static void test_enum_struct_members(struct TestContext* context)
+{
+	check(context, analyze_source(
+		"enum Color\n{\nRed,\nGreen\n}\nstruct Point\n{\nx\ny\n}\n"
+		"proc main\n{\nrax = Color.Green\nrbx = Point.y\nrcx = Point.size\n}\n"));
+	check(context, !analyze_source(
+		"enum Color\n{\nRed\n}\nproc main\n{\nrax = Color.Nope\n}\n"));
+	check(context, !analyze_source(
+		"struct Point\n{\nx\n}\nproc main\n{\nrax = Point.z\n}\n"));
+}
+
 static void test_references_resolve(struct TestContext* context)
 {
 	// registers, params, consts, data (+ .len), stack buffers, labels, calls
@@ -133,5 +144,6 @@ void run_sema_tests(struct TestContext* context)
 	test_const_expr_rejects_register(context);
 	test_const_expr_rejects_data(context);
 	test_deref_needs_register(context);
+	test_enum_struct_members(context);
 	test_references_resolve(context);
 }