Add locals for lambda procedures

1 files changed, 9 insertions, 3 deletions

diff --git a/src/compiler.h b/src/compiler.h
index 023a039..14c6954 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -461,6 +461,8 @@ compile_lambda(Object *obj) {
 
     // Create a new compilation context.
     char *prev_context = current_context;
+    Environment *prev_env = current_env;
+    current_env = obj->env;
     current_context = NULL;
     array_init(current_context, 0);
 
@@ -470,6 +472,7 @@ compile_lambda(Object *obj) {
 
     // Initialize function call frame.
     context_printf("    push    rbp\n");
+    context_printf("    sub     rsp, %zu\n", 8 * array_size(current_env->locals));
     context_printf("    mov     rbp, rsp\n");
 
     // Procedure body.
@@ -482,6 +485,7 @@ compile_lambda(Object *obj) {
 
     // Restore the previous call frame.
     context_printf("    mov     rsp, rbp\n");
+    context_printf("    add     rsp, %zu\n", 8 * array_size(current_env->locals));
     context_printf("    pop     rbp\n");
     context_printf("    ret\n");
     context_printf("\n");
@@ -489,6 +493,7 @@ compile_lambda(Object *obj) {
     // Restore previous compilation context.
     array_push(procedures, current_context);
     current_context = prev_context;
+    current_env = prev_env;
 
     // Create tagged pointer with this lambda procedure.
     context_printf("    mov     rax, %s\n", name);
@@ -512,7 +517,8 @@ compile_def(Object *obj) {
     compile_object(obj->var_expr);
     ssize_t idx = find_local_index(current_env->locals, obj->var_name);
     context_printf("    pop     rax\n");
-    context_printf("    mov     [rbp + %ld], rax\n", 8 * (idx + 1));
+    context_printf("    mov     [rbp + %ld], rax\n", 8 * idx);
+    compile_nil();
     context_printf("    ;; <-- compile_def\n");
 }
 
@@ -520,7 +526,7 @@ void
 compile_symbol(Object *obj) {
     context_printf("    ;; --> compile_symbol\n");
     ssize_t idx = find_local_index(current_env->locals, obj);
-    context_printf("    mov     rax, [rbp + %ld]\n", 8 * (idx + 1));
+    context_printf("    mov     rax, [rbp + %ld]\n", 8 * idx);
     context_printf("    push    rax\n");
     context_printf("    ;; <-- compile_symbol\n");
 }
@@ -632,7 +638,7 @@ compile(Program program) {
     printf("    mov     r15, bdl_heap\n");
 
     // Initialize main locals.
-    printf("    sub     rsp, %zu\n", 8 * array_size(program.env->locals));
+    printf("    sub     rsp, %zu\n", 8 * array_size(current_env->locals));
     printf("    mov     rbp, rsp\n");
 
     // Keep the bottom stack value set as NIL_VAL for a default return value.