Add tentative implementation of captured variables

2 files changed, 84 insertions, 61 deletions

diff --git a/src/bytecode/compiler.h b/src/bytecode/compiler.h
index 09e49a1..d35c7c6 100755
--- a/src/bytecode/compiler.h
+++ b/src/bytecode/compiler.h
@@ -9,7 +9,7 @@
 typedef struct Scope {
     StringView *params;
     StringView *locals;
-    StringView *captured;
+    Token *captured;
 } Scope;
 
 typedef struct Compiler {
@@ -80,7 +80,7 @@ find_local_index(Scope *scope, Token tok) {
 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])) {
+        if (sv_equal(&tok.value, &scope->captured[i].value)) {
             return i;
         }
     }
@@ -132,6 +132,54 @@ parse_fixnum(Chunk *chunk, Token tok) {
     emit_constant(chunk, tok, FIXNUM_VAL(num * sign));
 }
 
+void
+parse_symbol(Chunk *chunk, Compiler *compiler, Token tok) {
+    if (compiler->scope_depth > 1) {
+        Scope *current_scope = get_current_scope(compiler);
+        ssize_t idx = -1;
+        // 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(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) {
+                // 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);
+                return;
+            }
+            depth--;
+        }
+
+        // TODO: Capture globals?
+
+        // NOTE: set! must know how to deal also with captured vars.
+    }
+
+    Object obj = make_symbol(tok.value);
+    emit_constant(chunk, tok, obj);
+    add_code(chunk, OP_GET, tok.line, tok.column);
+}
+
 void parse_tree(Chunk *chunk, Compiler *compiler);
 
 void
@@ -417,7 +465,29 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
         parse_tree(fun.closure->chunk, compiler);
     }
     add_code(fun.closure->chunk, OP_RETURN, start.line, start.column);
-    emit_constant(chunk, start, fun);
+
+    // Prepare closure value capture.
+    Scope *scope = get_current_scope(compiler);
+    size_t n_captured = array_size(scope->captured);
+    if (n_captured > 0) {
+        compiler->scope_depth--;
+        for (ssize_t i = 0; i < n_captured; i++) {
+            Token tok = scope->captured[i];
+            print_token(tok);
+            parse_symbol(chunk, compiler, tok);
+        }
+        // Number of captured values.
+        emit_constant(chunk, tok, FIXNUM_VAL(n_captured));
+
+        // Push created lambda to stack.
+        emit_constant(chunk, start, fun);
+
+        // Emit capture local instruction.
+        add_code(chunk, OP_CAPTURE_LOCAL, tok.line, tok.column);
+        compiler->scope_depth++;
+    } else {
+        emit_constant(chunk, start, fun);
+    }
     exit_scope(compiler);
 }
 
@@ -622,55 +692,7 @@ parse_tree(Chunk *chunk, Compiler *compiler) {
             return;
         } break;
         case TOKEN_SYMBOL: {
-            if (compiler->scope_depth > 1) {
-                Scope *current_scope = get_current_scope(compiler);
-                ssize_t idx = -1;
-                // 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(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);
-            emit_constant(chunk, tok, obj);
-            add_code(chunk, OP_GET, tok.line, tok.column);
+            parse_symbol(chunk, compiler, tok);
             return;
         } break;
         case TOKEN_NIL: {
diff --git a/src/bytecode/vm.h b/src/bytecode/vm.h
index a32d0a1..f1c525d 100755
--- a/src/bytecode/vm.h
+++ b/src/bytecode/vm.h
@@ -177,8 +177,7 @@ vm_interpret(VM *vm) {
         switch (instruction) {
             case OP_LOCAL: {
                 ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
-                CallFrame frame = vm->frames[array_size(vm->frames) - 1];
-                array_push(vm->stack, vm->stack[frame.stack_offset + idx]);
+                array_push(vm->stack, vm->stack[frame->stack_offset + idx]);
             } break;
             case OP_CONSTANT: {
                 u8 constant = *vm->pc++;
@@ -186,15 +185,17 @@ vm_interpret(VM *vm) {
                 array_push(vm->stack, obj);
             } break;
             case OP_CAPTURED: {
-                assert(false && "not implemented");
+                ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
+                array_push(vm->stack, frame->closure->values[idx]);
             } 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;
+                Object proc = array_pop(vm->stack);
+                ssize_t n_captured = AS_FIXNUM(array_pop(vm->stack));
+                for (size_t i = 0; i < n_captured; i++) {
+                    Object value = array_pop(vm->stack);
+                    array_push(proc.closure->values, value);
+                }
+                array_push(vm->stack, proc);
             } break;
             case OP_DEF_LOCAL: {
                 Object value = array_pop(vm->stack);