From 603dc6f53a2a5d21b98597891979ffe8e84cc2df Mon Sep 17 00:00:00 2001
From: bog <bog@gozmail.net>
Date: Fri, 5 Apr 2024 21:51:19 +0200
Subject: [PATCH] :bug: closure capture and local variables collisions.

---
 features/fun.sk    | 14 +++++++++++
 lib/src/compiler.c | 60 +++++++++++++++++++++++++++++++++++++---------
 lib/src/exec.c     | 59 ++++++++++++++++++++++++++++++++++-----------
 lib/src/state.c    |  1 +
 lib/src/sym.c      |  1 +
 5 files changed, 110 insertions(+), 25 deletions(-)

diff --git a/features/fun.sk b/features/fun.sk
index 9acfa35..7168ebf 100644
--- a/features/fun.sk
+++ b/features/fun.sk
@@ -66,3 +66,17 @@ assert i() eq 2
 assert h() eq 8
 assert i() eq 3
 
+var j = fun (count)
+    return fun ()
+        count = count + 1
+        count - 1
+    end
+end
+
+var k = j(37)
+assert k() eq 37
+assert k() eq 38
+assert k() eq 39
+assert k() eq 40
+assert k() eq 41
+
diff --git a/lib/src/compiler.c b/lib/src/compiler.c
index 10da719..2859978 100644
--- a/lib/src/compiler.c
+++ b/lib/src/compiler.c
@@ -149,6 +149,7 @@ void compiler_compile(struct compiler* self,
             }
 
             compiler_compile(self, expr, prog, sym);
+
             if (symbol->closure_id != SK_NO_CLOSURE)
             {
                 prog_add_instr(prog, OP_CLOSURE_STORE, symbol->id);
@@ -167,7 +168,27 @@ void compiler_compile(struct compiler* self,
                 break;
             }
 
-            if (symbol->closure_id != SK_NO_CLOSURE)
+            struct env* itr = sym->env->parent;
+            struct symbol* context = NULL;
+
+            while (itr)
+            {
+                struct symbol* s = env_try_get(itr, symbol->name);
+
+                if (s)
+                {
+                    context = s;
+                    break;
+                }
+
+                itr = itr->parent;
+            }
+
+            if (0 && context)
+            {
+                prog_add_instr(prog, OP_LOCAL_LOAD, context->id);
+            }
+            else if (symbol->closure_id != SK_NO_CLOSURE)
             {
                 prog_add_instr(prog, OP_CLOSURE_LOAD, symbol->id);
             }
@@ -189,14 +210,7 @@ void compiler_compile(struct compiler* self,
             prog_init(p);
 
             fun_init(fun, p);
-
-            for (size_t i=0; i<sym->env->symbols.size; i++)
-            {
-                struct symbol* s = sym->env->symbols.data[i];
-                sym_decl_closure(fun->sym, s->id, s->name);
-            }
-
-            struct node* args = node->children.data[0];
+            struct node* params = node->children.data[0];
             struct node* body = node->children.data[1];
 
             char* var_name =
@@ -204,9 +218,9 @@ void compiler_compile(struct compiler* self,
                 self->var_names.data[self->var_names.size - 1]
                 : NULL;
 
-            for (size_t i=0; i<args->children.size; i++)
+            for (size_t i=0; i<params->children.size; i++)
             {
-                struct node* child = args->children.data[i];
+                struct node* child = params->children.data[i];
                 char* name = child->token->value;
                 sym_decl_var(fun->sym, name);
             }
@@ -215,6 +229,30 @@ void compiler_compile(struct compiler* self,
             {
                 sym_decl_var(fun->sym, var_name);
             }
+
+            // Closure
+            {
+                for (size_t i=0; i<sym->env->symbols.size; i++)
+                {
+                    struct symbol* s = sym->env->symbols.data[i];
+                    sym_decl_closure(fun->sym, s->id, s->name);
+                }
+
+                if (sym->env->parent)
+                {
+                    for (
+                    size_t i=0;
+                    i<sym->env->parent->symbols.size;
+                    i++
+                )
+                    {
+                        struct symbol* s =
+                            sym->env->parent->symbols.data[i];
+
+                        sym_decl_closure(fun->sym, s->id, s->name);
+                    }
+                }
+            }
             compiler_compile(self, body, p, fun->sym);
 
             union val val;
diff --git a/lib/src/exec.c b/lib/src/exec.c
index 518c138..188e540 100644
--- a/lib/src/exec.c
+++ b/lib/src/exec.c
@@ -308,24 +308,55 @@ void exec_capture_env(struct exec* self,
                       struct env* env)
 {
     (void) self;
-    for (size_t i=0; i<env->symbols.size; i++)
+
+    if (env->parent)
     {
-        struct symbol* symbol = env->symbols.data[i];
-        if (symbol->closure_id != SK_NO_CLOSURE)
+        for (size_t i=0; i<env->parent->symbols.size; i++)
         {
-            struct local* loc =
-                state_try_get_local_by_id(
-                    state,
-                    symbol->closure_id
-                );
+            struct symbol* symbol = env->parent->symbols.data[i];
+            if (symbol->closure_id != SK_NO_CLOSURE)
+            {
+                struct local* loc =
+                    state_try_get_local_by_id(
+                        state,
+                        symbol->closure_id
+                    );
 
-            struct value* val =
-                state_try_get_value(
-                    state,
-                    loc->addr
-                );
+                if(!loc) { continue; }
 
-            fun_capture(fun, symbol->id, val);
+                struct value* val =
+                    state_try_get_value(
+                        state,
+                        loc->addr
+                    );
+
+                fun_capture(fun, symbol->id, val);
+            }
+        }
+    }
+    if (env)
+    {
+        for (size_t i=0; i<env->symbols.size; i++)
+        {
+            struct symbol* symbol = env->symbols.data[i];
+            if (symbol->closure_id != SK_NO_CLOSURE)
+            {
+                struct local* loc =
+                    state_try_get_local_by_id(
+                        state,
+                        symbol->closure_id
+                    );
+
+                if(!loc) { continue; }
+
+                struct value* val =
+                    state_try_get_value(
+                        state,
+                        loc->addr
+                    );
+
+                fun_capture(fun, symbol->id, val);
+            }
         }
     }
 }
diff --git a/lib/src/state.c b/lib/src/state.c
index 81bd65c..f905c24 100644
--- a/lib/src/state.c
+++ b/lib/src/state.c
@@ -314,6 +314,7 @@ void state_closure_load(struct state* self,
                         int id)
 {
     assert(self);
+
     struct frame* frame = state_frame(self);
     assert(frame->fun);
 
diff --git a/lib/src/sym.c b/lib/src/sym.c
index 0c075e0..3d2907c 100644
--- a/lib/src/sym.c
+++ b/lib/src/sym.c
@@ -142,6 +142,7 @@ int sym_decl_closure(struct sym* self, int id, char const* name)
 {
     assert(self);
     assert(name);
+
     struct symbol* symbol = malloc(sizeof(struct symbol));
     symbol_init(symbol, self->closure_counter, name, self->env);
     symbol->is_const = false;