2 files changed, 64 insertions, 0 deletions

diff --git a/src/codegen/nasm.c b/src/codegen/nasm.c
index 49c13c4..50c16e2 100644
--- a/src/codegen/nasm.c
+++ b/src/codegen/nasm.c
@@ -239,12 +239,39 @@ static void emit_statement(struct Emitter* emitter, struct Statement* statement)
 		case STATEMENT_CALL:
 			emit_call(emitter, &statement->call);
 			break;
+		case STATEMENT_STACK:
+			break;
 		default:
 			fprintf(out, "\t; TODO: unsupported statement\n");
 			break;
 	}
 }
 
+static uint64_t token_to_u64(struct Token token)
+{
+	uint64_t value = 0;
+	for (size_t i = 0; i < token.length; i += 1)
+		value = value * 10 + (uint64_t)(token.start[i] - '0');
+
+	return value;
+}
+
+static uint64_t proc_stack_size(struct ProcDecl* proc)
+{
+	uint64_t total = 0;
+	for (size_t i = 0; i < proc->body_count; i += 1)
+	{
+		struct Statement* statement = &proc->body[i];
+		if (statement->kind == STATEMENT_STACK)
+			total += token_to_u64(statement->stack.size);
+	}
+
+	if (total % 16 != 0)
+		total += 16 - (total % 16);
+
+	return total;
+}
+
 static void emit_proc(struct Program* program, struct ProcDecl* proc, FILE* out)
 {
 	struct Emitter emitter;
@@ -259,11 +286,23 @@ static void emit_proc(struct Program* program, struct ProcDecl* proc, FILE* out)
 	else
 		fprintf(out, "%.*s:\n", (int)proc->name.length, proc->name.start);
 
+	uint64_t stack_size = proc_stack_size(proc);
+	if (stack_size > 0)
+	{
+		fprintf(out, "\tpush rbp\n");
+		fprintf(out, "\tmov rbp, rsp\n");
+		fprintf(out, "\tsub rsp, %llu\n", (unsigned long long)stack_size);
+	}
+
 	for (size_t i = 0; i < proc->body_count; i += 1)
 		emit_statement(&emitter, &proc->body[i]);
 
 	if (!is_entry)
+	{
+		if (stack_size > 0)
+			fprintf(out, "\tleave\n");
 		fprintf(out, "\tret\n");
+	}
 }
 
 void generate_nasm(struct Program* program, FILE* out)
diff --git a/tests/codegen_test.c b/tests/codegen_test.c
index f4df522..4f39855 100644
--- a/tests/codegen_test.c
+++ b/tests/codegen_test.c
@@ -152,6 +152,30 @@ static void test_generate_divide(struct TestContext* context)
 	free_program(&program);
 }
 
+static void test_generate_stack_frame(struct TestContext* context)
+{
+	struct Lexer lexer = create_lexer("proc work\n{\nstack buffer[32]\nsyscall\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, "push rbp") != NULL);
+	check(context, strstr(buffer, "mov rbp, rsp") != NULL);
+	check(context, strstr(buffer, "sub rsp, 32") != NULL);
+	check(context, strstr(buffer, "leave") != NULL);
+
+	free_program(&program);
+}
+
 void run_codegen_tests(struct TestContext* context)
 {
 	test_generate_consts_and_data(context);
@@ -160,4 +184,5 @@ void run_codegen_tests(struct TestContext* context)
 	test_generate_call(context);
 	test_generate_param_substitution(context);
 	test_generate_divide(context);
+	test_generate_stack_frame(context);
 }