From 26f1b9c35d337c0814158077fdc8f56b817e0b14 Mon Sep 17 00:00:00 2001
From: Bad Diode <bd@badd10de.dev>
Date: Mon, 15 Nov 2021 22:32:02 +0100
Subject: Fix stack on procedure call return

---
 src/compiler.h | 80 ++++++++++++++++++++++++++++++++++++++++++----------------
 1 file changed, 58 insertions(+), 22 deletions(-)

(limited to 'src/compiler.h')

diff --git a/src/compiler.h b/src/compiler.h
index 2a44220..1312d99 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -28,6 +28,8 @@ static Environment *current_env = NULL;
 // TODO: Separate c/h files
 // TODO: Create a "driver.c" file with the (display) function for external
 // linkage or assembly inlining.
+// TODO: Ensure we don't compile a function multiple times (for example with
+// a function that contains internal functions).
 
 // Immediate constants.
 #define NIL_VAL      47LU
@@ -338,6 +340,30 @@ compile_cdr(Object *obj) {
     context_printf("    ;; <-- compile_cdr\n");
 }
 
+void
+compile_call(Object *obj) {
+    context_printf("    ;; --> compile_call\n");
+    compile_object(obj->head);
+    size_t n_args = 0;
+    while (obj->tail != NULL) {
+        obj = obj->tail;
+        compile_object(obj->head);
+        n_args++;
+    }
+    context_printf("    mov     rax, [rsp + %zu]\n", 8 * n_args);
+    context_printf("    mov     rcx, PTR_MASK\n");
+    context_printf("    and     rcx, rax\n");
+    context_printf("    mov     rax, [rcx]\n");
+    context_printf("    call    rax\n");
+
+    // Restore stack to previous location and store the result on top.
+    context_printf("    add     rsp, %zu\n", 8 * (n_args + 1));
+    // TODO: Handle the case where the function doesn't return anything, for
+    // example if the last parameter is (display ...) or (def ...).
+    context_printf("    push    rax\n");
+    context_printf("    ;; <-- compile_call\n");
+}
+
 void
 compile_proc_call(Object *obj) {
     // TODO: Probably we want to use a hash table for these lookups that is
@@ -364,7 +390,6 @@ compile_proc_call(Object *obj) {
         compile_object(obj->tail->head);
         context_printf("    pop     rdi\n");
         context_printf("    call    display\n");
-        compile_nil();
     } else if (sv_equal(&obj->head->text, &STRING("not"))) {
         compile_not(obj->tail);
     } else if (sv_equal(&obj->head->text, &STRING("and"))) {
@@ -388,19 +413,7 @@ compile_proc_call(Object *obj) {
     } else if (sv_equal(&obj->head->text, &STRING("cdr"))) {
         compile_cdr(obj->tail);
     } else {
-        compile_object(obj->head);
-        size_t n_args = 0;
-        while (obj->tail != NULL) {
-            obj = obj->tail;
-            compile_object(obj->head);
-            n_args++;
-        }
-        context_printf("    mov     rax, [rsp + %zu]\n", 8 * n_args);
-        context_printf("    mov     rcx, PTR_MASK\n");
-        context_printf("    and     rcx, rax\n");
-        context_printf("    mov     rax, [rcx]\n");
-        context_printf("    call    rax\n");
-        context_printf("    push    rax\n");
+        compile_call(obj);
     }
 }
 
@@ -518,28 +531,51 @@ compile_def(Object *obj) {
     ssize_t idx = find_var_index(current_env->locals, obj->var_name);
     context_printf("    pop     rax\n");
     context_printf("    mov     [rbp + %ld], rax\n", 8 * idx);
-    compile_nil();
     context_printf("    ;; <-- compile_def\n");
 }
 
 void
 compile_symbol(Object *obj) {
     context_printf("    ;; --> compile_symbol\n");
-    ssize_t idx = find_var_index(current_env->locals, obj);
+    ssize_t idx = -1;
+
+    // TODO: Is a captured variable?
+    // FIXME: Order might be an issue, for example if the variable was initially
+    // captured but then declared as a local?
+    // (def a 40)
+    // (fun ext ()
+    //      (display a)
+    //      (def a 10)
+    //      (display a))
+    idx = find_var_index(current_env->captured, obj);
+    if (idx != -1) {
+        // TODO: not implemented.
+        context_printf("    ;; <-- compile_symbol\n");
+        return;
+    }
+
+    // Is a local variable?
+    idx = find_var_index(current_env->locals, obj);
     if (idx != -1) {
         context_printf("    mov     rax, [rbp + %ld]\n", 8 * idx);
         context_printf("    push    rax\n");
         context_printf("    ;; <-- compile_symbol\n");
         return;
     }
+
+    // Is a function parameter?
     idx = find_var_index(current_env->params, obj);
+    if (idx != -1) {
+        size_t n_locals = array_size(current_env->locals);
+        size_t n_params = array_size(current_env->params);
+        size_t offset = 8 * (n_locals + n_params - idx + 1);
+        context_printf("    mov     rax, [rbp + %ld]\n", offset);
+        context_printf("    push    rax\n");
+        context_printf("    ;; <-- compile_symbol\n");
+        return;
+    }
+
     assert(idx != -1 && "unexpected index");
-    size_t n_locals = array_size(current_env->locals);
-    size_t n_params = array_size(current_env->params);
-    size_t offset = 8 * (n_locals + n_params - idx + 1);
-    context_printf("    mov     rax, [rbp + %ld]\n", offset);
-    context_printf("    push    rax\n");
-    context_printf("    ;; <-- compile_symbol\n");
 }
 
 void
-- 
cgit v1.2.1