1 files changed, 85 insertions, 24 deletions

diff --git a/src/bytecode/compiler.h b/src/bytecode/compiler.h
index 7497ea7..f9e1da1 100755
--- a/src/bytecode/compiler.h
+++ b/src/bytecode/compiler.h
@@ -4,16 +4,49 @@
 #include "chunk.h"
 #include "lexer.h"
 
+#define MAX_DEPTH 1024
+
+typedef struct Scope {
+    StringView *params;
+    StringView *locals;
+} Scope;
+
 typedef struct Compiler {
     Token *tokens;
     size_t current;
+    size_t scope_depth;
+    Scope scopes[MAX_DEPTH];
 } Compiler;
 
+
 // Mimics the functionality in the Scanner functions, but for entire tokens.
 Token next_token(Compiler *compiler);
 Token peek_token(const Compiler *compiler);
 bool has_next_token(const Compiler *compiler);
 
+// Scope initialization/exit.
+void enter_scope(Compiler *compiler);
+void exit_scope(Compiler *compiler);
+
+void
+enter_scope(Compiler *compiler) {
+    Scope *scope = &compiler->scopes[compiler->scope_depth++];
+    array_init(scope->params, 0);
+    array_init(scope->locals, 0);
+}
+
+void
+exit_scope(Compiler *compiler) {
+    Scope *scope = &compiler->scopes[--compiler->scope_depth];
+    array_free(scope->params);
+    array_free(scope->locals);
+}
+
+Scope *
+get_current_scope(Compiler *compiler) {
+    return &compiler->scopes[compiler->scope_depth - 1];
+}
+
 Chunk * compile(Token *tokens);
 
 Token
@@ -32,10 +65,9 @@ has_next_token(const Compiler *compiler) {
 }
 
 ssize_t
-find_local_index(Chunk *chunk, Token tok) {
-    // NOTE: This is dumb and potentially slow.
-    for (size_t i = 0; i < array_size(chunk->locals); i++) {
-        if (sv_equal(&tok.value, &chunk->locals[i])) {
+find_local_index(Scope *scope, Token tok) {
+    for (size_t i = 0; i < array_size(scope->locals); i++) {
+        if (sv_equal(&tok.value, &scope->locals[i])) {
             return i;
         }
     }
@@ -217,29 +249,49 @@ compile_declare_op(Chunk *chunk, Compiler *compiler, Token start, Ops op) {
         });
         return;
     }
-    // TODO: If we are inside a function and we are using OP_DEF, we just
-    // declare the local variable directly in the stack. No need for symbols!
-    // TODO: If we are inside a function and we are using OP_SET, check if the
-    // local variable has been defined before, if not try to read from the
-    // globals for setting.
-    if (sv_equal(&STRING(""), &STR_ARRAY(chunk->name))) {
+
+    if (compiler->scope_depth <= 1) {
         Object obj = make_symbol(name.value);
         emit_constant(chunk, name, obj);
     } else {
-        // TODO: only do this if we are defining! not setting.
-        // Check if we already have the local
-        ssize_t idx = find_local_index(chunk, name);
-        if (idx < 0) {
-            array_push(chunk->locals, name.value);
-            idx = array_size(chunk->locals) - 1;
-        }
         if (op == OP_DEF) {
             op = OP_DEF_LOCAL;
+            // Check if the local is already registered.
+            Scope *scope = get_current_scope(compiler);
+            ssize_t idx = find_local_index(scope, name);
+            if (idx < 0) {
+                array_push(scope->locals, name.value);
+                idx = chunk->n_locals++;
+            }
+            emit_constant(chunk, name, FIXNUM_VAL(idx));
         } else if (op == OP_SET) {
-            op = OP_SET_LOCAL;
+            size_t depth = compiler->scope_depth - 1;
+            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);
+
+            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));
+            } else {
+                // If not found at all, emit set for the global scope.
+                Object obj = make_symbol(name.value);
+                emit_constant(chunk, name, obj);
+            }
         }
-        emit_constant(chunk, name, FIXNUM_VAL(idx));
     }
+    // NOTE: We can have compiler support for preemptively finding if globals
+    // exist or not.
 
     Token tok = peek_token(compiler);
     if (name.type == TOKEN_EOF || tok.type == TOKEN_EOF) {
@@ -276,6 +328,7 @@ 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);
 
     // Prepeare parameters.
@@ -305,7 +358,8 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
                 return;
             }
             // Check if parameters name already exists.
-            ssize_t idx = find_local_index(fun.chunk, tok);
+            Scope *scope = get_current_scope(compiler);
+            ssize_t idx = find_local_index(scope, tok);
             if (idx >= 0) {
                 error_push((Error){
                     .type = ERR_TYPE_COMPILER,
@@ -315,8 +369,10 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
                 });
                 return;
             }
-            array_push(fun.chunk->params, tok.value);
-            array_push(fun.chunk->locals, tok.value);
+            array_push(scope->params, tok.value);
+            array_push(scope->locals, tok.value);
+            fun.chunk->n_params++;
+            fun.chunk->n_locals++;
         }
     } else if (tok.type != TOKEN_NIL) {
         error_push((Error){
@@ -348,6 +404,7 @@ compile_lambda(Chunk *chunk, Compiler *compiler, Token start, StringView name) {
     }
     add_code(fun.chunk, OP_RETURN, start.line, start.column);
     emit_constant(chunk, start, fun);
+    exit_scope(compiler);
 }
 
 void
@@ -551,7 +608,8 @@ parse_tree(Chunk *chunk, Compiler *compiler) {
             return;
         } break;
         case TOKEN_SYMBOL: {
-            ssize_t idx = find_local_index(chunk, tok);
+            Scope *scope = get_current_scope(compiler);
+            ssize_t idx = find_local_index(scope, tok);
             if (idx < 0) {
                 Object obj = make_symbol(tok.value);
                 emit_constant(chunk, tok, obj);
@@ -583,11 +641,13 @@ parse_tree(Chunk *chunk, Compiler *compiler) {
 Chunk *
 compile(Token *tokens) {
     Chunk *chunk = NULL;
-    chunk = NEW_CHUNK("");
+    chunk = NEW_CHUNK("main");
     Compiler compiler = (Compiler){
         .tokens = tokens,
         .current = 0,
+        .scope_depth = 0,
     };
+    enter_scope(&compiler);
     Token main_start = peek_token(&compiler);
     while (has_next_token(&compiler)) {
         Token start = peek_token(&compiler);
@@ -598,6 +658,7 @@ compile(Token *tokens) {
         add_code(chunk, OP_DROP, start.line, start.column);
     }
     add_code(chunk, OP_RETURN, main_start.line, main_start.column);
+    exit_scope(&compiler);
     return chunk;
 }