Add workaround for `void` return in function calls

1 files changed, 11 insertions, 3 deletions

diff --git a/src/compiler.h b/src/compiler.h
index e91798a..e198127 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -358,8 +358,6 @@ compile_call(Object *obj) {
 
     // 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");
 }
@@ -489,9 +487,19 @@ compile_lambda(Object *obj) {
     context_printf("    mov     rbp, rsp\n");
 
     // Procedure body.
-    for (size_t i = 0; i < array_size(obj->body); i++) {
+    // In case the last expression of a function doesn't return anything (e.g.
+    // a `def` or `display` primitive), we store a sentinel `nil` value at the
+    // end of the stack.
+    //
+    // NOTE: This is probably better handled by a type system that
+    // allows functions to return void, but right now the caller and function
+    // creation expect that all functions return values. Failure to comply with
+    // this convention will result in a corrupted stack.
+    for (size_t i = 0; i < array_size(obj->body) - 1; i++) {
         compile_object(obj->body[i]);
     }
+    compile_nil();
+    compile_object(obj->body[array_size(obj->body) - 1]);
 
     // Return is stored in the `rax`.
     context_printf("    pop     rax\n");