Prepare for closure capture

1 files changed, 56 insertions, 16 deletions

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);