From 9c5fb457fe6063b7545515397aa45cecb7af66bf Mon Sep 17 00:00:00 2001
From: Bad Diode <bd@badd10de.dev>
Date: Wed, 17 Nov 2021 10:23:44 +0100
Subject: Fix local variables on recursive calls

---
 src/compiler.h | 34 ++++++++++++++++++++++++++--------
 1 file changed, 26 insertions(+), 8 deletions(-)

diff --git a/src/compiler.h b/src/compiler.h
index e0d071c..6ca4467 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -378,9 +378,6 @@ compile_call_body(Object *obj) {
         obj = obj->tail;
         compile_object(obj->head);
     }
-    context_printf("    mov     rax, [rsp + %zu]\n", 8 * offset);
-    context_printf("    jmp     rax\n");
-
     return offset;
 }
 
@@ -397,6 +394,10 @@ compile_call(Object *obj) {
     // Function call compilation without start/end.
     size_t offset = compile_call_body(obj);
 
+    // Call function.
+    context_printf("    mov     rdi, [rsp + %zu]\n", 8 * offset);
+    context_printf("    jmp     rdi\n");
+
     // Restore stack to previous location and store the result on top.
     context_printf("%s:\n", lab_end);
     context_printf("    add     rsp, %zu\n", 8 * (offset + 2));
@@ -528,7 +529,6 @@ compile_lambda(Object *obj) {
     size_t n_locals = array_size(current_env->locals);
     size_t n_params = array_size(current_env->params);
     size_t n_captured = array_size(current_env->captured);
-    size_t offset = 8 * (n_locals + n_params + n_captured + 1);
 
     // Initialize function call frame.
     context_printf("    sub     rsp, %zu\n", 8 * n_locals);
@@ -546,25 +546,43 @@ compile_lambda(Object *obj) {
     for (size_t i = 0; i < array_size(obj->body) - 1; i++) {
         compile_object(obj->body[i]);
     }
-    compile_nil();
     Object *last_expr = obj->body[array_size(obj->body) - 1];
 
     // Tail Call Optimization.
     // TODO: also for if statements
+    // FIXME: only pairs that are not primitives.
     if (IS_PAIR(last_expr)) {
         // Discard the previous stack frame.
         context_printf("    mov     rsp, rbp\n");
-        context_printf("    add     rsp, %zu\n", offset);
 
-        compile_call_body(last_expr);
+        size_t old_offset = n_locals + n_captured + n_params;
+        size_t new_offset = compile_call_body(last_expr);
+        context_printf("    mov     rdi, [rbp - 8]\n");
+        for (size_t i = 0; i < new_offset + 1; i++) {
+            context_printf("    mov     rax, [rbp - 8 * %zu]\n", i + 1);
+            context_printf("    mov     [rbp + 8 * %zu], rax\n", old_offset - i);
+        }
+
+        // Set the stack pointer at the end of given parameters.
+        context_printf("    mov     rsp, rbp\n");
+        ssize_t offset_diff = old_offset - new_offset;
+        if (offset_diff > 0) {
+            context_printf("    add     rsp, 8 * %zu\n", offset_diff);
+        } else {
+            context_printf("    sub     rsp, 8 * %zu\n", offset_diff);
+        }
+
+        context_printf("    jmp     rdi\n");
     } else {
+        compile_nil();
         compile_object(last_expr);
 
         // Return is stored in the `rax`.
         context_printf("    pop     rax\n");
 
         // Restore the previous call frame.
-        context_printf("    mov     rdi, [rbp + %zu]\n", offset);
+        size_t rp_offset = (n_locals + n_params + n_captured + 1);
+        context_printf("    mov     rdi, [rbp + %zu]\n", 8 * rp_offset);
         context_printf("    mov     rsp, rbp\n");
         context_printf("    add     rsp, %zu\n", 8 * n_locals);
         context_printf("    jmp     rdi\n");
-- 
cgit v1.2.1