Prepare for closure capture

10 files changed, 139 insertions, 63 deletions

diff --git a/Makefile b/Makefile
index 30784b2..05656f9 100755
--- a/Makefile
+++ b/Makefile
@@ -20,7 +20,7 @@ CFLAGS         += $(INC_FLAGS)
 LDFLAGS        :=
 LDLIBS         :=
 RELEASE_CFLAGS := -DNDEBUG -O2 -static
-DEBUG_CFLAGS   := -DDEBUG -O0 -g -fsanitize=address
+DEBUG_CFLAGS   := -DDEBUG -O0 -g
 
 .PHONY: build tests run clean
 
@@ -30,6 +30,8 @@ DEBUG_CFLAGS   := -DDEBUG -O0 -g -fsanitize=address
 DEBUG ?= 0
 ifeq ($(DEBUG), 1)
     CFLAGS += $(DEBUG_CFLAGS)
+else ifeq ($(DEBUG), 2)
+    CFLAGS += $(DEBUG_CFLAGS) -fsanitize=address
 else
     CFLAGS += $(RELEASE_CFLAGS)
 endif
diff --git a/src/bytecode/chunk.c b/src/bytecode/chunk.c
index 8ff6acf..af4a3a2 100644
--- a/src/bytecode/chunk.c
+++ b/src/bytecode/chunk.c
@@ -19,7 +19,7 @@ chunk_free(Chunk *chunk) {
     array_free(chunk->code);
     for (size_t i = 0; i < array_size(chunk->constants); i++) {
         Object obj = chunk->constants[i];
-        object_free(obj);
+        object_free(&obj);
     }
     array_free(chunk->constants);
     array_free(chunk->lines);
diff --git a/src/bytecode/compiler.h b/src/bytecode/compiler.h
index 124c345..09e49a1 100755
--- a/src/bytecode/compiler.h
+++ b/src/bytecode/compiler.h
@@ -9,6 +9,7 @@
 typedef struct Scope {
     StringView *params;
     StringView *locals;
+    StringView *captured;
 } Scope;
 
 typedef struct Compiler {
@@ -33,6 +34,7 @@ enter_scope(Compiler *compiler) {
     Scope *scope = &compiler->scopes[compiler->scope_depth++];
     array_init(scope->params, 0);
     array_init(scope->locals, 0);
+    array_init(scope->captured, 0);
 }
 
 void
@@ -40,6 +42,7 @@ exit_scope(Compiler *compiler) {
     Scope *scope = &compiler->scopes[--compiler->scope_depth];
     array_free(scope->params);
     array_free(scope->locals);
+    array_free(scope->captured);
 }
 
 Scope *
@@ -74,6 +77,16 @@ find_local_index(Scope *scope, Token tok) {
     return -1;
 }
 
+ssize_t
+find_captued_index(Scope *scope, Token tok) {
+    for (size_t i = 0; i < array_size(scope->captured); i++) {
+        if (sv_equal(&tok.value, &scope->captured[i])) {
+            return i;
+        }
+    }
+    return -1;
+}
+
 void
 emit_constant(Chunk *chunk, Token tok, Object obj) {
     size_t prev_size = array_size(chunk->constants);
@@ -84,7 +97,7 @@ emit_constant(Chunk *chunk, Token tok, Object obj) {
     // If the non value constant was already present we need to properly free
     // the memory from the object given to this function.
     if (prev_size == array_size(chunk->constants)) {
-        object_free(obj);
+        object_free(&obj);
     }
 }
 
@@ -329,7 +342,8 @@ compile_declare_op(Chunk *chunk, Compiler *compiler, Token start, Ops op) {
 void
 compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
     enter_scope(compiler);
-    Object fun = make_lambda(name);
+    Chunk *proc_chunk = chunk_init(name);
+    Object fun = make_lambda(proc_chunk);
 
     // Prepeare parameters.
     Token tok = next_token(compiler);
@@ -371,8 +385,8 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
             }
             array_push(scope->params, tok.value);
             array_push(scope->locals, tok.value);
-            fun.chunk->n_params++;
-            fun.chunk->n_locals++;
+            fun.closure->chunk->n_params++;
+            fun.closure->chunk->n_locals++;
         }
     } else if (tok.type != TOKEN_NIL) {
         error_push((Error){
@@ -400,9 +414,9 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
             next_token(compiler);
             break;
         }
-        parse_tree(fun.chunk, compiler);
+        parse_tree(fun.closure->chunk, compiler);
     }
-    add_code(fun.chunk, OP_RETURN, start.line, start.column);
+    add_code(fun.closure->chunk, OP_RETURN, start.line, start.column);
     emit_constant(chunk, start, fun);
     exit_scope(compiler);
 }
@@ -609,23 +623,49 @@ parse_tree(Chunk *chunk, Compiler *compiler) {
         } break;
         case TOKEN_SYMBOL: {
             if (compiler->scope_depth > 1) {
-                size_t depth = compiler->scope_depth - 1;
+                Scope *current_scope = get_current_scope(compiler);
                 ssize_t idx = -1;
-                do {
-                    Scope *scope = &compiler->scopes[depth];
-                    idx = find_local_index(scope, tok);
-                    if (idx >= 0) {
-                        break;
-                    }
-                    depth--;
-                } while (depth > 0);
+                // Check if the variable was already captured.
+                idx = find_captued_index(current_scope, tok);
+                if (idx >= 0) {
+                    emit_constant(chunk, tok, FIXNUM_VAL(idx));
+                    add_code(chunk, OP_CAPTURED, tok.line, tok.column);
+                    return;
+                }
 
+                // Check current scope locals. If we find it, emit OP_LOCAL.
+                idx = find_local_index(current_scope, tok);
                 if (idx >= 0) {
-                    emit_constant(chunk, tok, FIXNUM_VAL(depth));
                     emit_constant(chunk, tok, FIXNUM_VAL(idx));
                     add_code(chunk, OP_LOCAL, tok.line, tok.column);
                     return;
                 }
+
+                // Descend scopes, if we find the symbol at a different depth,
+                // we need to capture the variable.
+                size_t depth = compiler->scope_depth - 2;
+                while (depth > 0) {
+                    Scope *scope = &compiler->scopes[depth];
+                    idx = find_local_index(scope, tok);
+                    if (idx >= 0) {
+                        // Capture this local.
+                        emit_constant(chunk, tok, FIXNUM_VAL(idx));
+                        emit_constant(chunk, tok, FIXNUM_VAL(depth));
+                        add_code(chunk, OP_CAPTURE_LOCAL, tok.line, tok.column);
+
+                        // Put captured variable on stack.
+                        ssize_t captured_idx = array_size(current_scope->captured);
+                        emit_constant(chunk, tok, FIXNUM_VAL(captured_idx));
+                        add_code(chunk, OP_CAPTURED, tok.line, tok.column);
+                        array_push(current_scope->captured, tok.value);
+                        return;
+                    }
+                    depth--;
+                }
+
+                // TODO: Capture globals?
+
+                // NOTE: set! must know how to deal also with captured vars.
             }
 
             Object obj = make_symbol(tok.value);
diff --git a/src/bytecode/debug.h b/src/bytecode/debug.h
index e34b65f..7078744 100755
--- a/src/bytecode/debug.h
+++ b/src/bytecode/debug.h
@@ -11,6 +11,8 @@ static const char* ops_str[] = {
     // Load/store ops.
     [OP_CONSTANT] = "OP_CONSTANT",
     [OP_LOCAL] = "OP_LOCAL",
+    [OP_CAPTURED] = "OP_CAPTURED",
+    [OP_CAPTURE_LOCAL] = "OP_CAPTURE_LOCAL",
     [OP_DEF_LOCAL] = "OP_DEF_LOCAL",
     [OP_SET_LOCAL] = "OP_SET_LOCAL",
     [OP_DEF] = "OP_DEF",
diff --git a/src/bytecode/hashtable.h b/src/bytecode/hashtable.h
index 1f55048..7c0c380 100644
--- a/src/bytecode/hashtable.h
+++ b/src/bytecode/hashtable.h
@@ -151,7 +151,7 @@ _ht_maybe_grow(HashTable *table) {
     // Free old arrays.
     array_free(old_pairs);
     for (size_t i = 0; i < array_size(old_keys); i++) {
-        object_free(old_keys[i]);
+        object_free(&old_keys[i]);
     }
     array_free(old_keys);
     array_free(old_values);
@@ -198,7 +198,7 @@ ht_free(HashTable *table) {
     }
     array_free(table->pairs);
     for (size_t i = 0; i < array_size(table->keys); i++) {
-        object_free(table->keys[i]);
+        object_free(&table->keys[i]);
     }
     array_free(table->keys);
     array_free(table->values);
diff --git a/src/bytecode/main.c b/src/bytecode/main.c
index 7f2042e..4b80f71 100755
--- a/src/bytecode/main.c
+++ b/src/bytecode/main.c
@@ -57,8 +57,9 @@ process_source(const StringView *source) {
 #endif
 
     // Interpret chunk.
+    Object main_proc = make_lambda(main);
     CallFrame frame = (CallFrame){
-        .chunk = main,
+        .closure = main_proc.closure,
     };
     array_push(vm.frames, frame);
     vm_interpret(&vm);
diff --git a/src/bytecode/objects.c b/src/bytecode/objects.c
index 3b4a2eb..c2fb989 100644
--- a/src/bytecode/objects.c
+++ b/src/bytecode/objects.c
@@ -23,11 +23,13 @@ make_symbol(StringView sv) {
 }
 
 Object
-make_lambda(StringView name) {
+make_lambda(Chunk *chunk) {
     Object obj = {
         .type = OBJ_TYPE_LAMBDA,
-        .chunk = chunk_init(name),
     };
+    obj.closure = malloc(sizeof(Closure));
+    obj.closure->chunk = chunk;
+    array_init(obj.closure->values, 0);
     return obj;
 }
 
@@ -58,8 +60,9 @@ object_display(Object obj) {
             // printf(")");
         } break;
         case OBJ_TYPE_LAMBDA: {
+            Chunk *chunk = obj.closure->chunk;
             printf("#{procedure:%.*s}",
-                    (int)array_size(obj.chunk->name), obj.chunk->name);
+                    (int)array_size(chunk->name), chunk->name);
         } break;
         case OBJ_TYPE_ERR: {
             printf("#{error}");
@@ -69,13 +72,21 @@ object_display(Object obj) {
 }
 
 void
-object_free(Object obj) {
-    if (IS_STRING(obj) || IS_SYMBOL(obj)) {
-        array_free(obj.text);
+object_free(Object *obj) {
+    if (IS_STRING(*obj) || IS_SYMBOL(*obj)) {
+        array_free(obj->text);
         return;
     }
-    if (IS_LAMBDA(obj)) {
-        chunk_free(obj.chunk);
+    if (IS_LAMBDA(*obj)) {
+        Closure *closure = obj->closure;
+        for (size_t i = 0; i < array_size(closure->values); i++) {
+            object_free(&closure->values[i]);
+        }
+        array_free(closure->values);
+        // NOTE: we are leaking memory without this, we'll need a GC
+        // soon...
+        // chunk_free(closure->chunk);
+        free(closure);
     }
 }
 
@@ -104,7 +115,7 @@ object_equal(Object a, Object b) {
             }
         } break;
         case OBJ_TYPE_LAMBDA: {
-            return a.chunk == b.chunk;
+            return a.closure == b.closure;
         } break;
         default: {
             return false;
diff --git a/src/bytecode/objects.h b/src/bytecode/objects.h
index 6c286b5..a9a7d0f 100755
--- a/src/bytecode/objects.h
+++ b/src/bytecode/objects.h
@@ -17,6 +17,20 @@ typedef enum ObjectType {
     OBJ_TYPE_ERR,
 } ObjectType;
 
+typedef struct Object Object;
+
+typedef struct Closure {
+    // Non-owning reference to a chunk.
+    Chunk *chunk;
+    // Captured values for this closure.
+    Object *values;
+} Closure;
+
+// typdef struct ConsCell {
+//     struct Object *car;
+//     struct Object *cdr;
+// } ConsCell;
+
 typedef struct Object {
     ObjectType type;
     bool marked;
@@ -29,26 +43,18 @@ typedef struct Object {
         char *text;
 
         // OBJ_TYPE_PAIR
-        // struct {
-        //     struct Object *car;
-        //     struct Object *cdr;
-        // };
+        // ConsCell *cons_cell;
 
         // OBJ_TYPE_LAMBDA
-        Chunk *chunk;
-        // struct {
-        //     struct Object *params;
-        //     struct Object *body;
-        //     struct Environment *env;
-        // };
+        Closure *closure;
     };
 } Object;
 
 Object make_string(StringView sv);
 Object make_symbol(StringView sv);
-Object make_lambda(StringView name);
+Object make_lambda(Chunk *chunk);
 void object_display(Object obj);
-void object_free(Object obj);
+void object_free(Object *obj);
 bool object_equal(Object a, Object b);
 Object object_copy(Object src);
 
diff --git a/src/bytecode/ops.h b/src/bytecode/ops.h
index 52c774a..d45c27e 100755
--- a/src/bytecode/ops.h
+++ b/src/bytecode/ops.h
@@ -5,6 +5,9 @@ typedef enum Ops {
     // Load/store ops.
     OP_CONSTANT,
     OP_LOCAL,
+    OP_CAPTURED,
+    OP_CAPTURE_LOCAL,
+    // OP_SET_CAPTURED,
     OP_DEF_LOCAL,
     OP_SET_LOCAL,
     OP_DEF,
diff --git a/src/bytecode/vm.h b/src/bytecode/vm.h
index 3fba3d7..a32d0a1 100755
--- a/src/bytecode/vm.h
+++ b/src/bytecode/vm.h
@@ -13,7 +13,7 @@
 
 typedef struct CallFrame {
     // Current code being run.
-    Chunk *chunk;
+    Closure *closure;
     // Current program counter for this call.
     // Ref to stack. Is this really needed?
     // Object *locals;
@@ -65,7 +65,7 @@ vm_reset(VM *vm) {
     fprintf(stderr, "stack trace:\n"); \
     for (ssize_t i = array_size(vm->frames) - 1; i >= 0 ; i--) { \
         CallFrame frame = vm->frames[i]; \
-        Chunk *chunk = frame.chunk; \
+        Chunk *chunk = frame.closure->chunk; \
         size_t instruction = vm->pc - chunk->code - 1; \
         fprintf(stderr, "\t%-4ld  ->  ", i); \
         fprintf(stderr, "%.*s",(int)array_size(chunk->name), chunk->name); \
@@ -80,8 +80,8 @@ vm_reset(VM *vm) {
     error_push((Error){ \
         .type = ERR_TYPE_RUNTIME, \
         .value = (ERR), \
-        .line = frame->chunk->lines[vm->pc - frame->chunk->code - 1].line, \
-        .col = frame->chunk->lines[vm->pc - frame->chunk->code - 1].col, \
+        .line = chunk->lines[vm->pc - chunk->code - 1].line, \
+        .col = chunk->lines[vm->pc - chunk->code - 1].col, \
     }); \
     STACK_TRACE() \
     return
@@ -149,10 +149,11 @@ vm_reset(VM *vm) {
 void
 vm_interpret(VM *vm) {
     CallFrame *frame = &vm->frames[0];
-    vm->pc = frame->chunk->code;
+    Chunk *chunk = frame->closure->chunk;
+    vm->pc = chunk->code;
     frame->rp = NULL;
 
-    if (frame->chunk->code == NULL || array_size(frame->chunk->code) == 0) {
+    if (chunk->code == NULL || array_size(chunk->code) == 0) {
         error_push((Error){
             .type = ERR_TYPE_RUNTIME,
             .value = ERR_EMPTY_CHUNK,
@@ -170,22 +171,31 @@ vm_interpret(VM *vm) {
             }
         }
         printf(" ]\nop: ");
-        disassemble_instruction(frame->chunk, (vm->pc - frame->chunk->code));
+        disassemble_instruction(chunk, (vm->pc - chunk->code));
 #endif
         u8 instruction = *vm->pc++;
         switch (instruction) {
             case OP_LOCAL: {
                 ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
-                ssize_t depth = AS_FIXNUM(array_pop(vm->stack));
-                depth = array_size(vm->frames) - depth;
-                CallFrame frame = vm->frames[depth];
+                CallFrame frame = vm->frames[array_size(vm->frames) - 1];
                 array_push(vm->stack, vm->stack[frame.stack_offset + idx]);
             } break;
             case OP_CONSTANT: {
                 u8 constant = *vm->pc++;
-                Object obj = frame->chunk->constants[constant];
+                Object obj = chunk->constants[constant];
                 array_push(vm->stack, obj);
             } break;
+            case OP_CAPTURED: {
+                assert(false && "not implemented");
+            } break;
+            case OP_CAPTURE_LOCAL: {
+                assert(false && "not implemented");
+                // Object value = array_pop(vm->stack);
+                // ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
+                // ssize_t depth = AS_FIXNUM(array_pop(vm->stack));
+                // CallFrame frame = vm->frames[depth];
+                // vm->stack[frame.stack_offset + idx] = value;
+            } break;
             case OP_DEF_LOCAL: {
                 Object value = array_pop(vm->stack);
                 ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
@@ -194,8 +204,7 @@ vm_interpret(VM *vm) {
             case OP_SET_LOCAL: {
                 Object value = array_pop(vm->stack);
                 ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
-                ssize_t depth = AS_FIXNUM(array_pop(vm->stack));
-                CallFrame frame = vm->frames[depth];
+                CallFrame frame = vm->frames[array_size(vm->frames) - 1];
                 vm->stack[frame.stack_offset + idx] = value;
             } break;
             case OP_DEF: {
@@ -290,8 +299,8 @@ vm_interpret(VM *vm) {
                 // Check the number of arguments is correct.
                 // NOTE: This is probably better handled at compilation, but for
                 // now this is simpler to implement.
-                ssize_t n_params = proc.chunk->n_params;
-                ssize_t n_locals = proc.chunk->n_locals - n_params;
+                ssize_t n_params = proc.closure->chunk->n_params;
+                ssize_t n_locals = proc.closure->chunk->n_locals - n_params;
                 if (n_args < n_params) {
                     RUNTIME_ERROR(ERR_NOT_ENOUGH_ARGS);
                 } else if (n_args > n_params) {
@@ -299,20 +308,21 @@ vm_interpret(VM *vm) {
                 }
 
 #ifdef DEBUG
-                disassemble_chunk(proc.chunk);
+                disassemble_chunk(proc.closure->chunk);
                 printf("n_locals: %ld\n", n_locals);
                 printf("n_params: %ld\n", n_params);
 #endif
                 // Tail-call optimization.
-                if (proc.chunk != frame->chunk || *vm->pc != OP_RETURN) {
+                if (proc.closure->chunk != chunk || *vm->pc != OP_RETURN) {
                     // Prepare new call frame.
                     CallFrame new_frame = (CallFrame){
-                        .chunk = proc.chunk,
+                        .closure = proc.closure,
                         .rp = vm->pc,
                         .stack_offset = array_size(vm->stack) - n_params,
                     };
                     array_push(vm->frames, new_frame);
                     frame = &vm->frames[array_size(vm->frames) - 1];
+                    chunk = frame->closure->chunk;
 
                     // Prepare local slots.
                     for (ssize_t i = 0; i < n_locals; i++) {
@@ -330,7 +340,7 @@ vm_interpret(VM *vm) {
                     size_t offset = frame->stack_offset + n_params + n_locals;
                     array_head(vm->stack)->size = offset;
                 }
-                vm->pc = frame->chunk->code;
+                vm->pc = chunk->code;
             } break;
             case OP_RETURN: {
                 if (frame->rp == NULL) {
@@ -347,6 +357,7 @@ vm_interpret(VM *vm) {
                 array_head(vm->stack)->size = frame->stack_offset;
                 array_push(vm->stack, ret);
                 frame = &vm->frames[array_size(vm->frames) - 1];
+                chunk = frame->closure->chunk;
             } break;
             case OP_DROP: {
                 array_head(vm->stack)->size = 0;
@@ -360,8 +371,8 @@ vm_interpret(VM *vm) {
     error_push((Error){
         .type = ERR_TYPE_RUNTIME,
         .value = ERR_PC_OOB,
-        .line = frame->chunk->lines[0].line,
-        .col = frame->chunk->lines[0].col,
+        .line = chunk->lines[0].line,
+        .col = chunk->lines[0].col,
     });
 }