Add OP_SET_CAPTURED operation

1 files changed, 29 insertions, 21 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);