From 62984fb580b9e355cf9eacf9de53bff8c895c0b5 Mon Sep 17 00:00:00 2001
From: Bad Diode <bd@badd10de.dev>
Date: Mon, 15 Nov 2021 22:33:53 +0100
Subject: Add initial boilerplate for closure capture

---
 src/compiler.h | 20 ++++++++++++++++----
 src/parser.c   | 32 +++++++++++++++++++++++++++++++-
 src/parser.h   |  1 +
 3 files changed, 48 insertions(+), 5 deletions(-)

diff --git a/src/compiler.h b/src/compiler.h
index 1312d99..e91798a 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -508,18 +508,30 @@ compile_lambda(Object *obj) {
     current_context = prev_context;
     current_env = prev_env;
 
-    // Create tagged pointer with this lambda procedure.
+    // 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));
+        // TODO: What about capturing captured variables or parameters?
+        assert(idx != -1 && "unexpected index");
+    }
+
+    // Create tagged pointer with this lambda procedure.
     context_printf("    mov     rax, r15\n");
     context_printf("    or      rax, %zu\n", LAMBDA_TAG);
 
     // Push compiled object to the stack.
     context_printf("    push    rax\n");
 
-    // TODO: When we add closed variables they will be stored here, for now just
-    // incrementing by a 64 bit pointer size, as that is what we are storing.
-    context_printf("    add     r15, 8\n");
+    // Adjust the heap pointer depending on the number of variables captured.
+    context_printf("    add     r15, %ld\n", 8 * (array_size(obj->env->captured) + 2));
 
     context_printf("    ;; <-- compile_lambda\n");
 }
diff --git a/src/parser.c b/src/parser.c
index 0968ccd..e72675f 100644
--- a/src/parser.c
+++ b/src/parser.c
@@ -473,6 +473,14 @@ insert_params(Environment *env, Object *symbol) {
     array_push(env->params, symbol);
 }
 
+void
+insert_captured(Environment *env, Object *symbol) {
+    if (find_var_index(env->captured, symbol) != -1) {
+        return;
+    }
+    array_push(env->captured, symbol);
+}
+
 void
 semantic_analysis(Environment *env, Object *obj, Errors *errors) {
     if (obj == NULL || obj->visited) {
@@ -481,7 +489,27 @@ semantic_analysis(Environment *env, Object *obj, Errors *errors) {
     obj->visited = true;
     switch (obj->type) {
         case OBJ_TYPE_SYMBOL: {
-            Object *found = symbol_in_env(env, obj);
+            Object *found = NULL;
+
+            Environment *cur_env = env;
+            while (cur_env != NULL) {
+                ssize_t idx = find_var_index(cur_env->locals, obj);
+                if (idx != -1) {
+                    found = cur_env->local_values[idx];
+                    if (cur_env != env && cur_env->parent != NULL) {
+                        insert_captured(env, obj);
+                    }
+                    break;
+                }
+                idx = find_var_index(cur_env->params, obj);
+                if (idx != -1) {
+                    found = cur_env->params[idx];
+                    break;
+                }
+                cur_env = cur_env->parent;
+            }
+
+            // Check if symbol is in other environments.
             if (found == NULL) {
                 error_push(errors, (Error){
                     .type = ERR_TYPE_PARSER,
@@ -645,6 +673,7 @@ env_alloc(Environment *parent) {
     array_init(env->locals, 0);
     array_init(env->local_values, 0);
     array_init(env->params, 0);
+    array_init(env->captured, 0);
     env->parent = parent;
     array_push(environments, env);
     return env;
@@ -687,6 +716,7 @@ free_objects(void) {
             array_free(env->locals);
             array_free(env->local_values);
             array_free(env->params);
+            array_free(env->captured);
             free(env);
         }
         array_free(environments);
diff --git a/src/parser.h b/src/parser.h
index ca7260b..60a307c 100644
--- a/src/parser.h
+++ b/src/parser.h
@@ -7,6 +7,7 @@ typedef struct Environment {
     struct Object **locals;
     struct Object **local_values;
     struct Object **params;
+    struct Object **captured;
     struct Environment *parent;
 } Environment;
 
-- 
cgit v1.2.1