Add proper function calls with closed variables

1 files changed, 43 insertions, 12 deletions

diff --git a/src/compiler.h b/src/compiler.h
index e198127..3d61833 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -344,20 +344,46 @@ void
 compile_call(Object *obj) {
     context_printf("    ;; --> compile_call\n");
     compile_object(obj->head);
-    size_t n_args = 0;
+    context_printf("    pop     rax\n");
+    context_printf("    mov     rcx, PTR_MASK\n");
+    context_printf("    and     rcx, rax\n");
+    context_printf("    mov     rax, [rcx]\n");
+    context_printf("    push    rax\n");
+
+    // Get the number of parameters/captured variables for this function.
+    Object *fun = obj->head;
+    // FIXME: this is horrible and WILL BREAK.
+    while (!IS_LAMBDA(fun)) {
+        if (IS_SYMBOL(fun)) {
+            fun = symbol_in_env(current_env, fun);
+            continue;
+        }
+        if (IS_PAIR(fun)) {
+            fun = fun->head;
+        }
+    }
+
+    size_t n_args = array_size(fun->env->params);
+    size_t n_cap = array_size(fun->env->captured);
+    size_t offset = n_args + n_cap;
+
+    // Push captured variables.
+    for (size_t i = 0; i < n_cap; i++) {
+        context_printf("    mov     rax, [rcx + 8 * %zu]\n", i + 1);
+        context_printf("    push    rax\n");
+    }
+
+    // Compile arguments.
     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("    mov     rax, [rsp + %zu]\n", 8 * offset);
     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));
+    context_printf("    add     rsp, %zu\n", 8 * (offset + 1));
     context_printf("    push    rax\n");
     context_printf("    ;; <-- compile_call\n");
 }
@@ -519,14 +545,13 @@ compile_lambda(Object *obj) {
     // Add function address.
     context_printf("    mov     rax, %s\n", name);
     context_printf("    mov     [r15], rax\n");
-    context_printf("    mov     rax, %ld\n", array_size(obj->env->captured));
-    context_printf("    mov     [r15 + 8], rax\n");
 
     // Add captured variables to the heap.
     for (size_t i = 0; i < array_size(obj->env->captured); i++) {
         ssize_t idx = find_var_index(current_env->locals, obj->env->captured[i]);
-        context_printf("    mov     rax, [rbp + %ld]\n", 8 * idx);
-        context_printf("    mov     [r15 + %ld], rax\n", 8 * (i + 2));
+        context_printf("    mov     rax, rbp\n");
+        context_printf("    add     rax, %ld\n", 8 * idx);
+        context_printf("    mov     [r15 + %ld], rax\n", 8 * (i + 1));
         // TODO: What about capturing captured variables or parameters?
         assert(idx != -1 && "unexpected index");
     }
@@ -569,7 +594,13 @@ compile_symbol(Object *obj) {
     //      (display a))
     idx = find_var_index(current_env->captured, obj);
     if (idx != -1) {
-        // TODO: not implemented.
+        size_t n_locals = array_size(current_env->locals);
+        size_t n_params = array_size(current_env->params);
+        size_t n_cap = array_size(current_env->captured);
+        size_t offset = 8 * (n_locals + n_params + n_cap - idx + 1);
+        context_printf("    mov     rcx, [rbp + %ld]\n", offset);
+        context_printf("    mov     rax, [rcx]\n");
+        context_printf("    push    rax\n");
         context_printf("    ;; <-- compile_symbol\n");
         return;
     }