Add OP_SET_CAPTURED operation

4 files changed, 37 insertions, 23 deletions

diff --git a/src/bytecode/compiler.h b/src/bytecode/compiler.h
index d35c7c6..bc39d47 100755
--- a/src/bytecode/compiler.h
+++ b/src/bytecode/compiler.h
@@ -171,8 +171,6 @@ parse_symbol(Chunk *chunk, Compiler *compiler, Token tok) {
         }
 
         // TODO: Capture globals?
-
-        // NOTE: set! must know how to deal also with captured vars.
     }
 
     Object obj = make_symbol(tok.value);
@@ -326,28 +324,38 @@ compile_declare_op(Chunk *chunk, Compiler *compiler, Token start, Ops op) {
             }
             emit_constant(chunk, name, FIXNUM_VAL(idx));
         } else if (op == OP_SET) {
-            size_t depth = compiler->scope_depth - 1;
+            // FIXME: This is fucking ugly.
+            Scope *current_scope = get_current_scope(compiler);
             ssize_t idx = -1;
-            // Check if name is local in this or any previous scope.
-            do {
-                Scope *scope = &compiler->scopes[depth];
-                idx = find_local_index(scope, name);
-                if (idx >= 0) {
-                    break;
-                }
-                depth--;
-            } while (depth > 0);
-
+            // Check if the variable was already captured.
+            idx = find_captued_index(current_scope, name);
             if (idx >= 0) {
-                // If the value is found emit OP_SET_LOCAL with tree parameters:
-                // The new value, the scope depth, and the scope index.
-                op = OP_SET_LOCAL;
-                emit_constant(chunk, name, FIXNUM_VAL(depth));
                 emit_constant(chunk, name, FIXNUM_VAL(idx));
+                op = OP_SET_CAPTURED;
             } else {
-                // If not found at all, emit set for the global scope.
-                Object obj = make_symbol(name.value);
-                emit_constant(chunk, name, obj);
+                idx = find_local_index(current_scope, name);
+                if (idx >= 0) {
+                    emit_constant(chunk, name, FIXNUM_VAL(idx));
+                    op = OP_SET_LOCAL;
+                } else {
+                    size_t depth = compiler->scope_depth - 2;
+                    while (depth > 0) {
+                        Scope *scope = &compiler->scopes[depth];
+                        idx = find_local_index(scope, name);
+                        if (idx >= 0) {
+                            op = OP_SET_CAPTURED;
+                            ssize_t captured_idx = array_size(current_scope->captured);
+                            emit_constant(chunk, name, FIXNUM_VAL(captured_idx));
+                            array_push(current_scope->captured, name);
+                            break;
+                        }
+                        depth--;
+                    }
+                    if (idx < 0) {
+                        Object obj = make_symbol(name.value);
+                        emit_constant(chunk, name, obj);
+                    }
+                }
             }
         }
     }
@@ -471,7 +479,7 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
     size_t n_captured = array_size(scope->captured);
     if (n_captured > 0) {
         compiler->scope_depth--;
-        for (ssize_t i = 0; i < n_captured; i++) {
+        for (size_t i = 0; i < n_captured; i++) {
             Token tok = scope->captured[i];
             print_token(tok);
             parse_symbol(chunk, compiler, tok);
diff --git a/src/bytecode/debug.h b/src/bytecode/debug.h
index 7078744..b21d8e6 100755
--- a/src/bytecode/debug.h
+++ b/src/bytecode/debug.h
@@ -13,6 +13,7 @@ static const char* ops_str[] = {
     [OP_LOCAL] = "OP_LOCAL",
     [OP_CAPTURED] = "OP_CAPTURED",
     [OP_CAPTURE_LOCAL] = "OP_CAPTURE_LOCAL",
+    [OP_SET_CAPTURED] = "OP_SET_CAPTURED",
     [OP_DEF_LOCAL] = "OP_DEF_LOCAL",
     [OP_SET_LOCAL] = "OP_SET_LOCAL",
     [OP_DEF] = "OP_DEF",
diff --git a/src/bytecode/ops.h b/src/bytecode/ops.h
index d45c27e..a43aed6 100755
--- a/src/bytecode/ops.h
+++ b/src/bytecode/ops.h
@@ -7,7 +7,7 @@ typedef enum Ops {
     OP_LOCAL,
     OP_CAPTURED,
     OP_CAPTURE_LOCAL,
-    // OP_SET_CAPTURED,
+    OP_SET_CAPTURED,
     OP_DEF_LOCAL,
     OP_SET_LOCAL,
     OP_DEF,
diff --git a/src/bytecode/vm.h b/src/bytecode/vm.h
index f1c525d..4ff9743 100755
--- a/src/bytecode/vm.h
+++ b/src/bytecode/vm.h
@@ -191,12 +191,17 @@ vm_interpret(VM *vm) {
             case OP_CAPTURE_LOCAL: {
                 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++) {
+                for (ssize_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_SET_CAPTURED: {
+                Object value = array_pop(vm->stack);
+                ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
+                frame->closure->values[idx] = value;
+            } break;
             case OP_DEF_LOCAL: {
                 Object value = array_pop(vm->stack);
                 ssize_t idx = AS_FIXNUM(array_pop(vm->stack));