Fix stack on procedure call return

3 files changed, 65 insertions, 24 deletions

diff --git a/Makefile b/Makefile
index 581dc1a..cdb8c93 100644
--- a/Makefile
+++ b/Makefile
@@ -17,6 +17,7 @@ BIN    := $(BUILD_DIR)/$(TARGET)
 CC             := cc
 CFLAGS         := -Wall -Wextra -pedantic -DBIN_NAME=\"$(TARGET)\"
 CFLAGS         += $(INC_FLAGS)
+NASM_FLAGS     ?= -felf64
 LDFLAGS        :=
 LDLIBS         :=
 RELEASE_CFLAGS := -DNDEBUG -O2 -static
@@ -30,8 +31,10 @@ DEBUG_CFLAGS   := -DDEBUG -O0 -g
 DEBUG ?= 0
 ifeq ($(DEBUG), 1)
     CFLAGS += $(DEBUG_CFLAGS)
+    NASM_FLAGS += -g -F dwarf
 else ifeq ($(DEBUG), 2)
     CFLAGS += $(DEBUG_CFLAGS) -fsanitize=address
+    NASM_FLAGS += -g -F dwarf
 else
     CFLAGS += $(RELEASE_CFLAGS)
 endif
@@ -54,7 +57,7 @@ tests: $(BIN)
 
 run: $(BIN)
         $(BIN) example.bdl > build/example.asm
-        nasm -felf64 build/example.asm -o build/example.o
+        nasm $(NASM_FLAGS) build/example.asm -o build/example.o
         ld build/example.o -o build/example
         @./build/example
 
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
@@ -339,6 +341,30 @@ compile_cdr(Object *obj) {
 }
 
 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
     // initialized at the start of the compilation procedure.
@@ -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
diff --git a/src/parser.c b/src/parser.c
index 64bd22d..0968ccd 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -456,7 +456,9 @@ symbol_in_env(Environment *env, Object *symbol) {
 
 void
 insert_local(Environment *env, Object *symbol, Object *value) {
-    if (find_var_index(env->locals, symbol) != -1) {
+    ssize_t idx = find_var_index(env->locals, symbol);
+    if (idx != -1) {
+        env->local_values[idx] = value;
         return;
     }
     array_push(env->locals, symbol);