feat: calculate stack size instead of hard-coding

3 files changed, 14 insertions, 4 deletions

diff --git a/ast.c b/ast.c
index 8a586c2..6aa0e7e 100644
--- a/ast.c
+++ b/ast.c
@@ -218,6 +218,7 @@ typedef enum AstNodeKind AstNodeKind;
 #define node_idx __i1
 #define node_op __i1
 #define node_stack_offset __i1
+#define node_stack_size __i1
 
 struct AstNode {
     AstNodeKind kind;
diff --git a/codegen.c b/codegen.c
index 7b9e88c..a1167d4 100644
--- a/codegen.c
+++ b/codegen.c
@@ -7,6 +7,7 @@ typedef enum GenMode GenMode;
 struct CodeGen {
     int next_label;
     int* loop_labels;
+    AstNode* current_func;
 };
 typedef struct CodeGen CodeGen;
 
@@ -50,7 +51,7 @@ void codegen_func_prologue(CodeGen* g, AstNode* ast) {
     for (int i = 0; i < ast->node_params->node_len; ++i) {
         printf("  push %s\n", param_reg(i));
     }
-    printf("  sub rsp, %d\n", 8 * LVAR_MAX);
+    printf("  sub rsp, %d\n", ast->node_stack_size);
 }
 
 void codegen_func_epilogue(CodeGen* g, AstNode* ast) {
@@ -298,22 +299,23 @@ void codegen_func_call(CodeGen* g, AstNode* ast) {
     const char* func_name = ast->name;
 
     if (strcmp(func_name, "__ducc_va_start") == 0) {
+        int stack_size = g->current_func->node_stack_size;
         printf("  # __ducc_va_start BEGIN\n");
         for (int i = 0; i < 6; ++i) {
             printf("  mov rax, %s\n", param_reg(i));
-            printf("  mov [rbp-%d], rax\n", 8 + (LVAR_MAX - 4 - i) * 8);
+            printf("  mov [rbp-%d], rax\n", 8 + (stack_size - 4 - i) * 8);
         }
         AstNode* va_list_args = ast->node_args->node_items;
         codegen_expr(g, va_list_args, GenMode_lval);
         printf("  pop rdi\n");
         printf("  mov rax, rbp\n");
-        printf("  sub rax, %d\n", 8 + (LVAR_MAX - 1) * 8);
+        printf("  sub rax, %d\n", 8 + (stack_size - 1) * 8);
         printf("  mov [rdi], rax\n");
         printf("  mov DWORD PTR [rax], 8\n");
         printf("  mov DWORD PTR [rax+4], 0\n");
         printf("  mov QWORD PTR [rax+8], 0\n");
         printf("  mov rdi, rbp\n");
-        printf("  sub rdi, %d\n", 8 + (LVAR_MAX - 4) * 8);
+        printf("  sub rdi, %d\n", 8 + (stack_size - 4) * 8);
         printf("  mov QWORD PTR [rax+16], rdi\n");
         printf("  # __ducc_va_start END\n");
         return;
@@ -536,6 +538,7 @@ void codegen_stmt(CodeGen* g, AstNode* ast) {
 }
 
 void codegen_func(CodeGen* g, AstNode* ast) {
+    g->current_func = ast;
     printf("%s:\n", ast->name);
 
     codegen_func_prologue(g, ast);
@@ -547,6 +550,7 @@ void codegen_func(CodeGen* g, AstNode* ast) {
     codegen_func_epilogue(g, ast);
 
     printf("\n");
+    g->current_func = NULL;
 }
 
 void codegen(Program* prog) {
diff --git a/parse.c b/parse.c
index 7262c6c..1ee9ce7 100644
--- a/parse.c
+++ b/parse.c
@@ -1102,6 +1102,11 @@ AstNode* parse_func_decl_or_def(Parser* p) {
     func->name = name;
     func->node_params = params;
     func->node_body = body;
+    if (p->n_lvars == 0) {
+        func->node_stack_size = 0;
+    } else {
+        func->node_stack_size = p->lvars[p->n_lvars - 1].stack_offset + type_sizeof(p->lvars[p->n_lvars - 1].ty);
+    }
     return func;
 }