✨ closures.

features/fun.sk

@@ -43,3 +43,26 @@ end
 
 assert f(fun(x) x * 2 end, 2) eq 8
 
+var g = fun (init)
+    var count = init
+    
+    var inner = fun ()
+        count = count + 1
+        count - 1
+    end
+
+    return inner
+end
+
+var h = g(5)
+var i = g(0)
+
+assert h() eq 5
+assert i() eq 0
+assert h() eq 6
+assert i() eq 1
+assert h() eq 7
+assert i() eq 2
+assert h() eq 8
+assert i() eq 3
+

lib/include/exec.h

@@ -4,6 +4,7 @@
 #include "commons.h"
 #include "state.h"
 #include "prog.h"
+struct env;
 
 struct exec
 {
@@ -17,4 +18,8 @@ void exec_execute(struct exec* self,
                   struct state* state,
                   struct prog* prog);
 
+void exec_capture_env(struct exec* self, 
+                      struct state* state,
+                      struct fun* fun,
+                      struct env* env);
 #endif

lib/include/fun.h

@@ -3,13 +3,28 @@
 
 #include "commons.h"
 
+struct closure
+{
+    int id;
+    struct value* value;
+};
+
 struct fun
 {
     struct prog* prog;
+    struct sym* sym;
+    struct vec closures;
 };
 
+void closure_init(struct closure* self, int id, struct value* value);
+struct closure* closure_new_clone(struct closure* self);
+void closure_free(struct closure* self);
+
 void fun_init(struct fun* self, struct prog* new_prog);
 void fun_free(struct fun* self);
 struct fun* fun_new_clone(struct fun* self);
 
+void fun_capture(struct fun* self, int id, struct value* value);
+struct closure* fun_try_get_closure(struct fun* self, int id);
+
 #endif

lib/include/prog.h

@@ -13,7 +13,8 @@ G(OP_USUB), G(OP_ASSERT_EQ), \
 G(OP_NOT), G(OP_AND), G(OP_OR), \
 G(OP_BR), G(OP_BRF), G(OP_LT), G(OP_GT), \
 G(OP_EQUAL), G(OP_LOCAL_STORE), G(OP_LOCAL_LOAD), \
-G(OP_CALL), G(OP_RET), G(OP_MAKE_REF)
+G(OP_CALL), G(OP_RET), G(OP_MAKE_REF), \
+G(OP_CLOSURE_LOAD), G(OP_CLOSURE_STORE)
 
 SK_ENUM_H(Opcode, OPCODE);
 

lib/include/state.h

@@ -24,6 +24,7 @@ struct frame
     struct vec stack_values;
     struct vec locals;
     struct vec stack;
+    struct fun* fun;
 };
 
 struct state
@@ -55,6 +56,7 @@ void state_pop_frame(struct state* self);
 bool state_has_top(struct state* self);
 SK state_top(struct state* self);
 struct value* state_try_get_value(struct state* self, SK value);
+struct local* state_try_get_local_by_id(struct state* self, int id);
 
 SK state_pop(struct state* self);
 
@@ -75,9 +77,15 @@ struct local* state_try_get_local(struct state* self, int id);
 SK state_local_store(struct state* self,
                      int id);
 
+void state_closure_store(struct state* self,
+                         int id);
+
 void state_local_load(struct state* self,
                       int id);
 
+void state_closure_load(struct state* self,
+                        int id);
+
 TypeKind state_common_num_type(struct state* self, SK lhs, SK rhs);
 TypeKind state_type(struct state* self, SK value);
 double state_as_real(struct state* self, SK lhs);
@@ -99,6 +107,6 @@ SK state_lt(struct state* self);
 SK state_gt(struct state* self);
 SK state_eq(struct state* self);
 
-void state_call(struct state* self);
+void state_call(struct state* self, struct fun* fun);
 void state_ret(struct state* self);
 #endif

lib/include/sym.h

@@ -2,10 +2,12 @@
 #define SK_SYM_H
 
 #include "commons.h"
+#define SK_NO_CLOSURE (-1)
 
 struct symbol
 {
     int id;
+    int closure_id;
     char* name;
     bool is_const;
     struct env* env;
@@ -20,11 +22,14 @@ struct env
 struct sym
 {
     int id_counter;
+    int closure_counter;
     struct env* env;
 };
 
 void env_init(struct env* self);
+struct env* env_new_clone(struct env* self);
 void env_free(struct env* self);
+
 struct symbol* env_try_get(struct env* self, char const* name);
 
 void symbol_init(struct symbol* self,
@@ -35,10 +40,12 @@ void symbol_init(struct symbol* self,
 void symbol_free(struct symbol* self);
 
 void sym_init(struct sym* self);
+struct sym* sym_new_clone(struct sym* self);
 void sym_free(struct sym* self);
 
 int sym_decl_var(struct sym* self, char const* name);
 int sym_decl_const(struct sym* self, char const* name);
+int sym_decl_closure(struct sym* self, int id, char const* name);
 
 void sym_open_scope(struct sym* self);
 void sym_close_scope(struct sym* self);

lib/src/compiler.c

@@ -104,7 +104,9 @@ void compiler_compile(struct compiler* self,
 
             struct node* expr = node->children.data[1];
             compiler_compile(self, expr, prog, sym);
+
             int id = sym_decl_var(sym, ident->token->value);
+
             prog_add_instr(prog, OP_LOCAL_STORE, id);
             char* name = vec_pop(&self->var_names);
             free(name);
@@ -147,14 +149,32 @@ void compiler_compile(struct compiler* self,
             }
 
             compiler_compile(self, expr, prog, sym);
-            prog_add_instr(prog, OP_LOCAL_STORE, symbol->id);
+            if (symbol->closure_id != SK_NO_CLOSURE)
+            {
+                prog_add_instr(prog, OP_CLOSURE_STORE, symbol->id);
+            }
+            else
+            {
+                prog_add_instr(prog, OP_LOCAL_STORE, symbol->id);
+            }
         } break;
 
         case NODE_IDENT: {
             struct symbol* symbol = sym_try_get(sym,
                                                 node->token->value);
-            assert(symbol);
-            prog_add_instr(prog, OP_LOCAL_LOAD, symbol->id);
+            if (!symbol)
+            {
+                break;
+            }
+
+            if (symbol->closure_id != SK_NO_CLOSURE)
+            {
+                prog_add_instr(prog, OP_CLOSURE_LOAD, symbol->id);
+            }
+            else
+            {
+                prog_add_instr(prog, OP_LOCAL_LOAD, symbol->id);
+            }
         } break;
 
         case NODE_RETURN: {
@@ -168,11 +188,17 @@ void compiler_compile(struct compiler* self,
             struct prog* p = malloc(sizeof(struct prog));
             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* body = node->children.data[1];
 
-            struct sym mysym;
-            sym_init(&mysym);
             char* var_name =
                 self->var_names.size > 0 ?
                 self->var_names.data[self->var_names.size - 1]
@@ -182,16 +208,14 @@ void compiler_compile(struct compiler* self,
             {
                 struct node* child = args->children.data[i];
                 char* name = child->token->value;
-                sym_decl_var(&mysym, name);
+                sym_decl_var(fun->sym, name);
             }
 
             if (var_name)
             {
-                sym_decl_var(&mysym, var_name);
+                sym_decl_var(fun->sym, var_name);
             }
-            compiler_compile(self, body, p, &mysym);
-            sym_free(&mysym);
-            fun_init(fun, p);
+            compiler_compile(self, body, p, fun->sym);
 
             union val val;
             val.fun = fun;

lib/src/exec.c

@@ -1,4 +1,5 @@
 #include "exec.h"
+#include "sym.h"
 #include "fun.h"
 
 void exec_init(struct exec* self)
@@ -33,11 +34,21 @@ void exec_execute(struct exec* self,
                 self->pc++;
             } break;
 
+            case OP_CLOSURE_STORE: {
+                state_closure_store(state, param);
+                self->pc++;
+            } break;
+
             case OP_LOCAL_LOAD: {
                 state_local_load(state, param);
                 self->pc++;
             } break;
 
+            case OP_CLOSURE_LOAD: {
+                state_closure_load(state, param);
+                self->pc++;
+            } break;
+
             case OP_ASSERT_EQ: {
                 SK rhs = state_pop(state);
                 SK lhs = state_pop(state);
@@ -91,7 +102,12 @@ void exec_execute(struct exec* self,
 
                 struct value* f = value_new_clone(state_try_get_value(state, function));
 
-                state_call(state);
+                struct fun* fun = state_try_deref(
+                    state,
+                    f->val.ref
+                )->val.fun;
+
+                state_call(state, fun);
                 for (size_t i=0; i<params.size; i++)
                 {
                     size_t k = params.size - 1 - i;
@@ -139,9 +155,14 @@ void exec_execute(struct exec* self,
                 switch (constant->type)
                 {
                     case TYPE_FUN: {
+
+                        struct fun* fun = fun_new_clone(constant->val.fun);
+                        struct sym* s = fun->sym;
+                        exec_capture_env(self, state, fun, s->env);
+
                         state_push_fun(
                             state,
-                            fun_new_clone(constant->val.fun),
+                            fun,
                             constant->line
                         );
                     } break;
@@ -280,3 +301,31 @@ void exec_execute(struct exec* self,
         }
     }
 }
+
+void exec_capture_env(struct exec* self,
+                      struct state* state,
+                      struct fun* fun,
+                      struct env* env)
+{
+    (void) self;
+    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
+                );
+
+            struct value* val =
+                state_try_get_value(
+                    state,
+                    loc->addr
+                );
+
+            fun_capture(fun, symbol->id, val);
+        }
+    }
+}

lib/src/fun.c

@@ -1,10 +1,38 @@
 #include "fun.h"
 #include "prog.h"
+#include "sym.h"
+
+void closure_init(struct closure* self, int id, struct value* value)
+{
+    assert(self);
+    assert(value);
+
+    self->id = id;
+    self->value = value_new_clone(value);
+}
+
+struct closure* closure_new_clone(struct closure* self)
+{
+    assert(self);
+    struct closure* clone = malloc(sizeof(struct closure));
+    closure_init(clone, self->id, self->value);
+    return clone;
+}
+
+void closure_free(struct closure* self)
+{
+    assert(self);
+    value_free(self->value);
+    free(self->value);
+}
 
 void fun_init(struct fun* self, struct prog* new_prog)
 {
     assert(self);
     self->prog = new_prog;
+    self->sym = malloc(sizeof(struct sym));
+    sym_init(self->sym);
+    vec_init(&self->closures);
 }
 
 void fun_free(struct fun* self)
@@ -13,6 +41,10 @@ void fun_free(struct fun* self)
     prog_free(self->prog);
     free(self->prog);
     self->prog = NULL;
+    sym_free(self->sym);
+    free(self->sym);
+    vec_free_elements(&self->closures, (void*) closure_free);
+    vec_free(&self->closures);
 }
 
 struct fun* fun_new_clone(struct fun* self)
@@ -20,5 +52,43 @@ struct fun* fun_new_clone(struct fun* self)
     (void) self;
     struct fun* clone = malloc(sizeof(struct fun));
     fun_init(clone, prog_new_clone(self->prog));
+    sym_free(clone->sym);
+    free(clone->sym);
+    clone->sym = sym_new_clone(self->sym);
+
+    for (size_t i=0; i<self->closures.size; i++)
+    {
+        struct closure* closure = self->closures.data[i];
+        struct closure* new_closure = closure_new_clone(closure);
+        vec_push(&clone->closures, new_closure);
+    }
     return clone;
 }
+
+void fun_capture(struct fun* self, int id, struct value* value)
+{
+    assert(self);
+    assert(value);
+
+    struct closure* c = malloc(sizeof(struct closure));
+    closure_init(c, id, value);
+
+    vec_push(&self->closures, c);
+}
+
+struct closure* fun_try_get_closure(struct fun* self, int id)
+{
+    assert(self);
+
+    for (size_t i=0; i<self->closures.size; i++)
+    {
+        struct closure* c = self->closures.data[i];
+        
+        if (c->id == id)
+        {
+            return c;
+        }
+    }
+
+    return NULL;
+}

lib/src/state.c

@@ -23,6 +23,7 @@ void frame_init(struct frame* self)
     vec_init(&self->stack_values);
     vec_init(&self->locals);
     vec_init(&self->stack);
+    self->fun = NULL;
 }
 
 void frame_free(struct frame* self)
@@ -137,6 +138,23 @@ struct value* state_try_get_value(struct state* self, SK value)
     return NULL;
 }
 
+struct local* state_try_get_local_by_id(struct state* self, int id)
+{
+    struct frame* frame = state_frame(self);
+
+    for (size_t i=0; i<frame->locals.size; i++)
+    {
+        struct local* loc = frame->locals.data[i];
+
+        if (loc->id == id)
+        {
+            return loc;
+        }
+    }
+
+    return NULL;
+}
+
 SK state_pop(struct state* self)
 {
     assert(self);
@@ -257,6 +275,21 @@ SK state_local_store(struct state* self,
     return local->addr;
 }
 
+void state_closure_store(struct state* self,
+                         int id)
+{
+    assert(self);
+    struct frame* frame = state_frame(self);
+    assert(frame->fun);
+
+    size_t addr = state_pop(self);
+    struct value* value = state_try_get_value(self, addr);
+    struct closure* closure = fun_try_get_closure(frame->fun, id);
+    value_free(closure->value);
+    free(closure->value);
+    closure->value = value_new_clone(value);
+}
+
 void state_local_load(struct state* self,
                       int id)
 {
@@ -277,6 +310,28 @@ void state_local_load(struct state* self,
     abort();
 }
 
+void state_closure_load(struct state* self,
+                        int id)
+{
+    assert(self);
+    struct frame* frame = state_frame(self);
+    assert(frame->fun);
+
+    for (size_t i=0; i<frame->fun->closures.size; i++)
+    {
+        struct closure* closure = frame->fun->closures.data[i];
+
+        if (closure->id == id)
+        {
+            struct value* value = closure->value;
+            state_push(self, value->type, value->val, value->line);
+            return;
+        }
+    }
+
+    fprintf(stderr, "cannot load closure %d\n", id);
+    abort();
+}
 TypeKind state_common_num_type(struct state* self, SK lhs, SK rhs)
 {
     assert(self);
@@ -674,10 +729,12 @@ SK state_eq(struct state* self)
     );
 }
 
-void state_call(struct state* self)
+void state_call(struct state* self, struct fun* fun)
 {
     assert(self);
     state_push_frame(self);
+    struct frame* frame = state_frame(self);
+    frame->fun = fun;
 }
 
 void state_ret(struct state* self)

lib/src/sym.c

@@ -7,6 +7,30 @@ void env_init(struct env* self)
     self->parent = NULL;
 }
 
+struct env* env_new_clone(struct env* self)
+{
+    assert(self);
+    struct env* clone = malloc(sizeof(struct env));
+    env_init(clone);
+
+    for (size_t i=0; i<self->symbols.size; i++)
+    {
+        struct symbol const* symbol = self->symbols.data[i];
+        struct symbol* new_symbol = malloc(sizeof(struct symbol));
+        new_symbol->id = symbol->id;
+        new_symbol->name = strdup(symbol->name);
+        new_symbol->is_const = symbol->is_const;
+        new_symbol->closure_id = symbol->closure_id;
+        new_symbol->env = clone;
+
+        vec_push(&clone->symbols, new_symbol);
+    }
+
+    clone->parent = self->parent ? env_new_clone(self->parent) : NULL;
+
+    return clone;
+}
+
 void env_free(struct env* self)
 {
     assert(self);
@@ -28,7 +52,7 @@ struct symbol* env_try_get(struct env* self, char const* name)
             return symbol;
         }
     }
-    
+
     if (self->parent)
     {
         return env_try_get(self->parent, name);
@@ -46,6 +70,7 @@ void symbol_init(struct symbol* self,
     self->id = id;
     self->name = strdup(name);
     self->is_const = false;
+    self->closure_id = SK_NO_CLOSURE;
     self->env = env;
 }
 
@@ -59,10 +84,26 @@ void sym_init(struct sym* self)
 {
     assert(self);
     self->id_counter = 1;
+    self->closure_counter = 1;
     self->env = malloc(sizeof(struct env));
     env_init(self->env);
 }
 
+struct sym* sym_new_clone(struct sym* self)
+{
+    assert(self);
+    struct sym* clone = malloc(sizeof(struct sym));
+    sym_init(clone);
+
+    clone->id_counter = self->id_counter;
+    clone->closure_counter = self->closure_counter;
+    env_free(clone->env);
+    free(clone->env);
+    clone->env = self->env ? env_new_clone(self->env) : NULL;
+
+    return clone;
+}
+
 void sym_free(struct sym* self)
 {
     assert(self);
@@ -97,6 +138,20 @@ int sym_decl_const(struct sym* self, char const* name)
     return self->id_counter - 1;
 }
 
+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;
+    symbol->closure_id = id;
+    vec_push(&self->env->symbols, symbol);
+
+    self->closure_counter++;
+    return self->closure_counter - 1;
+}
+
 void sym_open_scope(struct sym* self)
 {
     assert(self);