Add support for lexically scoped local variables

5 files changed, 104 insertions, 41 deletions

diff --git a/src/bytecode/chunk.c b/src/bytecode/chunk.c
index 71562fa..8ff6acf 100644
--- a/src/bytecode/chunk.c
+++ b/src/bytecode/chunk.c
@@ -9,8 +9,8 @@ chunk_init(StringView name) {
     array_init(chunk->lines, 0);
     array_init(chunk->name, name.n);
     array_insert(chunk->name, name.start, name.n);
-    array_init(chunk->params, 0);
-    array_init(chunk->locals, 0);
+    chunk->n_params = 0;
+    chunk->n_locals = 0;
     return chunk;
 }
 
@@ -24,8 +24,6 @@ chunk_free(Chunk *chunk) {
     array_free(chunk->constants);
     array_free(chunk->lines);
     array_free(chunk->name);
-    array_free(chunk->params);
-    array_free(chunk->locals);
     free(chunk);
 }
 
diff --git a/src/bytecode/chunk.h b/src/bytecode/chunk.h
index c584d4a..9457fa9 100755
--- a/src/bytecode/chunk.h
+++ b/src/bytecode/chunk.h
@@ -20,10 +20,10 @@ typedef struct Chunk {
     LineInfo *lines;
     // Chunk name.
     char *name;
-    // Parameters
-    StringView *params;
-    // Locals.
-    StringView *locals;
+
+    // Number of locals and parameters.
+    size_t n_params;
+    size_t n_locals;
 } Chunk;
 
 #define NEW_CHUNK(NAME) chunk_init((StringView){(NAME), sizeof(NAME) - 1})
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;
 }
 
diff --git a/src/bytecode/debug.h b/src/bytecode/debug.h
index 54d2cdb..e2e3756 100755
--- a/src/bytecode/debug.h
+++ b/src/bytecode/debug.h
@@ -48,11 +48,7 @@ static const char* ops_str[] = {
 
 void
 disassemble_chunk(Chunk *chunk) {
-    if (array_size(chunk->name) < 1) {
-        printf("=====  main  =====\n");
-    } else {
-        printf("=====  %.*s  =====\n", (int)array_size(chunk->name), chunk->name);
-    }
+    printf("=====  %.*s  =====\n", (int)array_size(chunk->name), chunk->name);
     printf("code:\n");
     size_t offset = 0;
     while (offset < array_size(chunk->code)) {
diff --git a/src/bytecode/vm.h b/src/bytecode/vm.h
index 84d2432..287c83c 100755
--- a/src/bytecode/vm.h
+++ b/src/bytecode/vm.h
@@ -188,6 +188,13 @@ vm_interpret(VM *vm) {
                 ssize_t idx = AS_FIXNUM(array_pop(vm->stack));
                 vm->stack[frame->stack_offset + idx] = value;
             } break;
+            case OP_SET_LOCAL: {
+                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;
+            } break;
             case OP_DEF: {
                 Object value = array_pop(vm->stack);
                 Object name = array_pop(vm->stack);
@@ -280,13 +287,13 @@ vm_interpret(VM *vm) {
                 // Check the number of arguments is correct.
                 // NOTE: This is probably better handled at compilation, but for
                 // now this is simpler to implement.
-                ssize_t n_params = array_size(proc.chunk->params);
+                ssize_t n_params = proc.chunk->n_params;
+                ssize_t n_locals = proc.chunk->n_locals;
                 if (n_args < n_params) {
                     RUNTIME_ERROR(ERR_NOT_ENOUGH_ARGS);
                 } else if (n_args > n_params) {
                     RUNTIME_ERROR(ERR_TOO_MANY_ARGS);
                 }
-                ssize_t n_locals = array_size(proc.chunk->locals) - n_params;
 
 #ifdef DEBUG
                 disassemble_chunk(proc.chunk);
@@ -298,7 +305,7 @@ vm_interpret(VM *vm) {
                     CallFrame new_frame = (CallFrame){
                         .chunk = proc.chunk,
                         .rp = vm->pc,
-                        .stack_offset = array_size(vm->stack) - array_size(proc.chunk->params),
+                        .stack_offset = array_size(vm->stack) - n_params,
                     };
                     array_push(vm->frames, new_frame);
                     frame = &vm->frames[array_size(vm->frames) - 1];
@@ -315,7 +322,8 @@ vm_interpret(VM *vm) {
                     }
 
                     // Reset stack size.
-                    array_head(vm->stack)->size = frame->stack_offset + n_params + n_locals;
+                    size_t offset = frame->stack_offset + n_params + n_locals;
+                    array_head(vm->stack)->size = offset;
                 }
                 vm->pc = frame->chunk->code;
             } break;