feat: introduce upvalues to catch local variables used inside closures to be used when the enclosing function goes from stack

Author: aymaneelmaini

src/chunk.h

@@ -16,6 +16,8 @@ typedef enum
     OP_GET_GLOBAL,
     OP_DEFINE_GLOBAL,
     OP_SET_GLOBAL,
+    OP_GET_UPVALUE,
+    OP_SET_UPVALUE,
     OP_EQUAL,
     OP_GREATER,
     OP_LESS,

src/compiler.c

@@ -55,6 +55,12 @@ typedef struct
     bool  is_immutable;
 } Local;
 
+typedef struct
+{
+    u8   index;
+    bool islocal;
+} Upvalue;
+
 typedef enum
 {
     TYPE_FUNCTION,
@@ -67,9 +73,10 @@ typedef struct Compiler
     ObjFunction*     function;
     FunctionType     type;
 
-    Local locals[UINT8_COUNT];
-    int   local_count;
-    int   scope_depth;
+    Local   locals[UINT8_COUNT];
+    int     local_count;
+    Upvalue upvalues[UINT8_COUNT];
+    int     scope_depth;
 } Compiler;
 
 static bool immutable_globals[UINT8_MAX];
@@ -303,6 +310,43 @@ static int resolve_local(Compiler* compiler, Token* name)
     return -1;
 }
 
+static int add_upvalue(Compiler* compiler, u8 index, bool islocal)
+{
+    int upvalue_count = compiler->function->upvalue_count;
+    for (int i = 0; i < upvalue_count; i++)
+    {
+        Upvalue* upvalue = &compiler->upvalues[i];
+        if (upvalue->index == index && upvalue->islocal == islocal)
+            return i;
+    }
+
+    if (upvalue_count == UINT8_COUNT)
+    {
+        error("Too many closure variables in function");
+        return 0;
+    }
+
+    compiler->upvalues[upvalue_count].islocal = islocal;
+    compiler->upvalues[upvalue_count].index = index;
+    return compiler->function->upvalue_count++;
+}
+
+static int resolve_upvalue(Compiler* compiler, Token* name)
+{
+    if (compiler->enclosing == NULL)
+        return -1;
+
+    int local = resolve_local(compiler->enclosing, name);
+    if (local != -1)
+        return add_upvalue(compiler, (u8)local, true);
+
+    int upvalue = resolve_upvalue(compiler->enclosing, name);
+    if (upvalue != -1)
+        return add_upvalue(compiler, (u8)upvalue, false);
+
+    return -1;
+}
+
 static void add_local(Token name, bool is_immutable)
 {
     if (current->local_count == UINT8_COUNT)
@@ -503,6 +547,11 @@ static void named_variable(Token name, bool can_assign)
         get_op = OP_GET_LOCAL;
         set_op = OP_SET_LOCAL;
     }
+    else if ((arg = resolve_upvalue(current, &name)) != -1)
+    {
+        get_op = OP_GET_UPVALUE;
+        set_op = OP_SET_UPVALUE;
+    }
     else
     {
         arg = identifier_constant(&name);
@@ -680,6 +729,12 @@ static void function(FunctionType type)
 
     ObjFunction* function = end_compiler();
     emit_bytes(OP_CLOSURE, make_constant(OBJ_VAL(function)));
+
+    for (int i = 0; i < function->upvalue_count; i++)
+    {
+        emit_byte(compiler.upvalues[i].islocal ? 1 : 1);
+        emit_byte(compiler.upvalues[i].index);
+    }
 }
 
 static void fun_declaration()

src/debug.c

@@ -2,6 +2,7 @@
 
 #include "chunk.h"
 #include "debug.h"
+#include "object.h"
 #include "value.h"
 
 static int simple_instruction(const char* name, int offset);
@@ -52,6 +53,10 @@ int disassemble_instruction(Chunk* chunk, int offset)
         return constant_instruction("OP_DEFINE_GLOBAL", chunk, offset);
     case OP_SET_GLOBAL:
         return constant_instruction("OP_SET_GLOBAL", chunk, offset);
+    case OP_GET_UPVALUE:
+        return byte_instruction("OP_GET_UPVALUE", chunk, offset);
+    case OP_SET_UPVALUE:
+        return byte_instruction("OP_SET_UPVALUE", chunk, offset);
     case OP_EQUAL:
         return simple_instruction("OP_EQUAL", offset);
     case OP_GREATER:
@@ -87,6 +92,15 @@ int disassemble_instruction(Chunk* chunk, int offset)
         print_value(chunk->constants.values[constant]);
         printf("\n");
 
+        ObjFunction* function = AS_FUNCTION(chunk->constants.values[constant]);
+        for (int j = 0; j < function->upvalue_count; j++)
+        {
+            int islocal = chunk->code[offset++];
+            int index = chunk->code[offset++];
+            printf("%04     |              %s %d\n", offset - 2,
+                   islocal ? "local" : "upvalue", index);
+        }
+
         return offset;
     }
     case OP_RETURN:

src/memory.c

@@ -48,6 +48,11 @@ static void free_object(Obj* object)
         FREE(ObjString, object);
         break;
     }
+    case OBJ_UPVALUE:
+    {
+        FREE(ObjUpvalue, object);
+        break;
+    }
     }
 }
 void free_objects()

src/object.c

@@ -35,6 +35,7 @@ ObjFunction* new_function()
     ObjFunction* function = ALLOCATE_OBJ(ObjFunction, OBJ_FUNCTION);
 
     function->arity = 0;
+    function->upvalue_count = 0;
     function->name = NULL;
     init_chunk(&function->chunk);
     return function;
@@ -136,5 +137,10 @@ void print_object(Value value)
         printf("%s", AS_CSTRING(value));
         break;
     }
+    case OBJ_UPVALUE:
+    {
+        printf("upvalue");
+        break;
+    }
     }
 }

src/object.h

@@ -26,6 +26,7 @@ typedef enum
     OBJ_FUNCTION,
     OBJ_NATIVE,
     OBJ_STRING,
+    OBJ_UPVALUE,
 } ObjType;
 
 struct Obj
@@ -38,6 +39,7 @@ typedef struct
 {
     Obj        obj;
     int        arity;
+    int        upvalue_count;
     Chunk      chunk;
     ObjString* name;
 } ObjFunction;