2 files changed, 43 insertions, 0 deletions

diff --git a/tests/codegen_test.c b/tests/codegen_test.c
index 1da711c..c92dfc2 100644
--- a/tests/codegen_test.c
+++ b/tests/codegen_test.c
@@ -100,6 +100,27 @@ static void test_generate_if_else(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_generate_while(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer(
+		"proc main\n{\nwhile rcx > 0\n{\nrbx += rcx\nrcx -= 1\n}\n}\n");
+	struct Program program;
+	check(context, parse_program(&lexer, &program));
+
+	char buffer[1024];
+	generate_to_buffer(&program, buffer, sizeof(buffer));
+
+	check(context, strstr(buffer, ".while_0:") != NULL);
+	check(context, strstr(buffer, "cmp rcx, 0") != NULL);
+	check(context, strstr(buffer, "jle .while_end_0") != NULL);
+	check(context, strstr(buffer, "add rbx, rcx") != NULL);
+	check(context, strstr(buffer, "jmp .while_0") != NULL);
+	check(context, strstr(buffer, ".while_end_0:") != NULL);
+	check(context, strstr(buffer, "; TODO") == NULL);
+
+	free_program(&program);
+}
+
 static void test_generate_negative(struct TestContext* context)
 {
 	struct Lexer lexer = create_lexer(
@@ -485,6 +506,7 @@ void run_codegen_tests(struct TestContext* context)
 	test_generate_text(context);
 	test_generate_if(context);
 	test_generate_if_else(context);
+	test_generate_while(context);
 	test_generate_negative(context);
 	test_generate_call(context);
 	test_generate_param_substitution(context);
diff --git a/tests/parser_test.c b/tests/parser_test.c
index b51093e..88e32a1 100644
--- a/tests/parser_test.c
+++ b/tests/parser_test.c
@@ -225,6 +225,26 @@ static void test_parse_else_if(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_parse_while(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer(
+		"proc main\n{\nwhile rcx > 0\n{\nrbx += rcx\nrcx -= 1\n}\n}\n");
+	struct Program program;
+
+	check(context, parse_program(&lexer, &program));
+	check(context, program.procs[0].body_count == 1);
+
+	struct Statement loop = program.procs[0].body[0];
+	check(context, loop.kind == STATEMENT_WHILE);
+	check(context, primary_is(loop.loop.left, "rcx"));
+	check(context, text_is(loop.loop.comparison, ">"));
+	check(context, primary_is(loop.loop.right, "0"));
+	check(context, loop.loop.body_count == 2);
+	check(context, text_is(loop.loop.body[1].assign.target, "rcx"));
+
+	free_program(&program);
+}
+
 static void test_parse_call(struct TestContext* context)
 {
 	struct Lexer lexer = create_lexer("proc main\n{\nprint_number(r12)\nf(a, b)\n}\n");
@@ -396,6 +416,7 @@ void run_parser_tests(struct TestContext* context)
 	test_parse_if(context);
 	test_parse_if_block(context);
 	test_parse_else_if(context);
+	test_parse_while(context);
 	test_parse_call(context);
 	test_parse_stack(context);
 	test_parse_sized_deref(context);