Increase reference accuracy of on-stack close envs.

Using a bitset to indicate which stack values are upvalues, we
can more accurately track when a reference to a stack value
persists after the stack frame exits.

CHANGELOG.md

@@ -5,7 +5,7 @@ All notable changes to this project will be documented in this file.
 - Add `os/chmod`.
 - Add `chr` macro.
 - Allow `_` in the `match` macro to match anything without creating a binding
-  or doing unification.
+  or doing unification. Also change behavior of matching nil.
 - Add `:range-to` and `:down-to` verbs in the `loop` macro.
 - Fix `and` and `or` macros returning nil instead of false in some cases.
 - Allow matching successfully against nil values in the `match` macro.

src/core/bytecode.c

@@ -212,6 +212,7 @@ JanetFuncDef *janet_funcdef_alloc(void) {
     def->environments = NULL;
     def->constants = NULL;
     def->bytecode = NULL;
+    def->closure_bitset = NULL;
     def->flags = 0;
     def->slotcount = 0;
     def->arity = 0;

src/core/compile.c

@@ -102,6 +102,7 @@ void janetc_scope(JanetScope *s, JanetCompiler *c, int flags, const char *name)
     scope.bytecode_start = janet_v_count(c->buffer);
     scope.flags = flags;
     scope.parent = c->scope;
+    janetc_regalloc_init(&scope.ua);
     /* Inherit slots */
     if ((!(flags & JANET_SCOPE_FUNCTION)) && c->scope) {
         janetc_regalloc_clone(&scope.ra, &(c->scope->ra));
@@ -149,6 +150,7 @@ void janetc_popscope(JanetCompiler *c) {
     janet_v_free(oldscope->envs);
     janet_v_free(oldscope->defs);
     janetc_regalloc_deinit(&oldscope->ra);
+    janetc_regalloc_deinit(&oldscope->ua);
     /* Update pointer */
     if (newscope)
         newscope->child = NULL;
@@ -236,6 +238,11 @@ found:
         scope = scope->parent;
     janet_assert(scope, "invalid scopes");
     scope->flags |= JANET_SCOPE_ENV;
+
+    /* In the function scope, allocate the slot as an upvalue */
+    janetc_regalloc_touch(&scope->ua, ret.index);
+
+    /* Iterate through child scopes and make sure environment is propagated */
     scope = scope->child;
 
     /* Propagate env up to current scope */
@@ -737,6 +744,21 @@ JanetFuncDef *janetc_pop_funcdef(JanetCompiler *c) {
         def->flags |= JANET_FUNCDEF_FLAG_NEEDSENV;
     }
 
+    /* Copy upvalue bitset */
+    if (scope->ua.count) {
+        /* Number of u32s we need to create a bitmask for all slots */
+        int32_t numchunks = (def->slotcount + 31) >> 5;
+        uint32_t *chunks = malloc(sizeof(uint32_t) * numchunks);
+        if (NULL == chunks) {
+            JANET_OUT_OF_MEMORY;
+        }
+        memcpy(chunks, scope->ua.chunks, sizeof(uint32_t) * numchunks);
+        /* Register allocator preallocates some registers [240-255, high 16 bits of chunk index 7], we can ignore those. */
+        if (scope->ua.count > 7) chunks[7] &= 0xFFFFU;
+        def->closure_bitset = chunks;
+        def->flags |= JANET_FUNCDEF_FLAG_HASCLOBITSET;
+    }
+
     /* Pop the scope */
     janetc_popscope(c);
 

src/core/compile.h

@@ -127,7 +127,10 @@ struct JanetScope {
     /* Regsiter allocator */
     JanetcRegisterAllocator ra;
 
-    /* Referenced closure environents. The values at each index correspond
+    /* Upvalue allocator */
+    JanetcRegisterAllocator ua;
+
+    /* Referenced closure environments. The values at each index correspond
      * to which index to get the environment from in the parent. The environment
      * that corresponds to the direct parent's stack will always have value 0. */
     int32_t *envs;

src/core/fiber.c

@@ -218,13 +218,27 @@ int janet_fiber_funcframe(JanetFiber *fiber, JanetFunction *func) {
 static void janet_env_detach(JanetFuncEnv *env) {
     /* Check for closure environment */
     if (env) {
-        size_t s = sizeof(Janet) * (size_t) env->length;
+        int32_t len = env->length;
+        size_t s = sizeof(Janet) * (size_t) len;
         Janet *vmem = malloc(s);
         janet_vm_next_collection += (uint32_t) s;
         if (NULL == vmem) {
             JANET_OUT_OF_MEMORY;
         }
-        safe_memcpy(vmem, env->as.fiber->data + env->offset, s);
+        Janet *values = env->as.fiber->data + env->offset;
+        safe_memcpy(vmem, values, s);
+        uint32_t *bitset = janet_stack_frame(values)->func->def->closure_bitset;
+        if (bitset) {
+            /* Clear unneeded references in closure environment */
+            for (int32_t i = 0; i < len; i += 32) {
+                uint32_t mask = ~(bitset[i >> 5]);
+                int32_t maxj = i + 32 > len ? len : i + 32;
+                for (int32_t j = i; j < maxj; j++) {
+                    if (mask & 1) vmem[j] = janet_wrap_nil();
+                    mask >>= 1;
+                }
+            }
+        }
         env->offset = 0;
         env->as.values = vmem;
     }

src/core/gc.c

@@ -309,6 +309,7 @@ static void janet_deinit_block(JanetGCObject *mem) {
             free(def->constants);
             free(def->bytecode);
             free(def->sourcemap);
+            free(def->closure_bitset);
         }
         break;
     }

src/core/marsh.c

@@ -188,8 +188,14 @@ static void marshal_one_env(MarshalState *st, JanetFuncEnv *env, int flags) {
         pushint(st, 0);
         pushint(st, env->length);
         Janet *values = env->as.fiber->data + env->offset;
-        for (int32_t i = 0; i < env->length; i++)
-            marshal_one(st, values[i], flags + 1);
+        uint32_t *bitset = janet_stack_frame(values)->func->def->closure_bitset;
+        for (int32_t i = 0; i < env->length; i++) {
+            if (1 & (bitset[i >> 5] >> (i & 0x1F))) {
+                marshal_one(st, values[i], flags + 1);
+            } else {
+                pushbyte(st, LB_NIL);
+            }
+        }
     } else {
         janet_env_maybe_detach(env);
         pushint(st, env->offset);
@@ -214,6 +220,16 @@ static void janet_func_addflags(JanetFuncDef *def) {
     if (def->sourcemap) def->flags |= JANET_FUNCDEF_FLAG_HASSOURCEMAP;
 }
 
+/* Marshal a sequence of u32s */
+static void janet_marshal_u32s(MarshalState *st, const uint32_t *u32s, int32_t n) {
+    for (int32_t i = 0; i < n; i++) {
+        pushbyte(st, u32s[i] & 0xFF);
+        pushbyte(st, (u32s[i] >> 8) & 0xFF);
+        pushbyte(st, (u32s[i] >> 16) & 0xFF);
+        pushbyte(st, (u32s[i] >> 24) & 0xFF);
+    }
+}
+
 /* Marshal a function def */
 static void marshal_one_def(MarshalState *st, JanetFuncDef *def, int flags) {
     MARSH_STACKCHECK;
@@ -248,12 +264,7 @@ static void marshal_one_def(MarshalState *st, JanetFuncDef *def, int flags) {
         marshal_one(st, def->constants[i], flags);
 
     /* marshal the bytecode */
-    for (int32_t i = 0; i < def->bytecode_length; i++) {
-        pushbyte(st, def->bytecode[i] & 0xFF);
-        pushbyte(st, (def->bytecode[i] >> 8) & 0xFF);
-        pushbyte(st, (def->bytecode[i] >> 16) & 0xFF);
-        pushbyte(st, (def->bytecode[i] >> 24) & 0xFF);
-    }
+    janet_marshal_u32s(st, def->bytecode, def->bytecode_length);
 
     /* marshal the environments if needed */
     for (int32_t i = 0; i < def->environments_length; i++)
@@ -273,6 +284,11 @@ static void marshal_one_def(MarshalState *st, JanetFuncDef *def, int flags) {
             current = map.line;
         }
     }
+
+    /* Marshal closure bitset, if needed */
+    if (def->flags & JANET_FUNCDEF_FLAG_HASCLOBITSET) {
+        janet_marshal_u32s(st, def->closure_bitset, ((def->slotcount + 31) >> 5));
+    }
 }
 
 #define JANET_FIBER_FLAG_HASCHILD (1 << 29)
@@ -716,6 +732,20 @@ static const uint8_t *unmarshal_one_env(
     return data;
 }
 
+/* Unmarshal a series of u32s */
+static const uint8_t *janet_unmarshal_u32s(UnmarshalState *st, const uint8_t *data, uint32_t *into, int32_t n) {
+    for (int32_t i = 0; i < n; i++) {
+        MARSH_EOS(st, data + 3);
+        into[i] =
+            (uint32_t)(data[0]) |
+            ((uint32_t)(data[1]) << 8) |
+            ((uint32_t)(data[2]) << 16) |
+            ((uint32_t)(data[3]) << 24);
+        data += 4;
+    }
+    return data;
+}
+
 /* Unmarshal a funcdef */
 static const uint8_t *unmarshal_one_def(
     UnmarshalState *st,
@@ -739,6 +769,7 @@ static const uint8_t *unmarshal_one_def(
         def->bytecode_length = 0;
         def->name = NULL;
         def->source = NULL;
+        def->closure_bitset = NULL;
         janet_v_push(st->lookup_defs, def);
 
         /* Set default lengths to zero */
@@ -794,15 +825,7 @@ static const uint8_t *unmarshal_one_def(
         if (!def->bytecode) {
             JANET_OUT_OF_MEMORY;
         }
-        for (int32_t i = 0; i < bytecode_length; i++) {
-            MARSH_EOS(st, data + 3);
-            def->bytecode[i] =
-                (uint32_t)(data[0]) |
-                ((uint32_t)(data[1]) << 8) |
-                ((uint32_t)(data[2]) << 16) |
-                ((uint32_t)(data[3]) << 24);
-            data += 4;
-        }
+        data = janet_unmarshal_u32s(st, data, def->bytecode, bytecode_length);
         def->bytecode_length = bytecode_length;
 
         /* Unmarshal environments */
@@ -849,6 +872,15 @@ static const uint8_t *unmarshal_one_def(
             def->sourcemap = NULL;
         }
 
+        /* Unmarshal closure bitset if needed */
+        if (def->flags & JANET_FUNCDEF_FLAG_HASCLOBITSET) {
+            def->closure_bitset = malloc(sizeof(uint32_t) * def->slotcount);
+            if (NULL == def->closure_bitset) {
+                JANET_OUT_OF_MEMORY;
+            }
+            data = janet_unmarshal_u32s(st, data, def->closure_bitset, (def->slotcount + 31) >> 5);
+        }
+
         /* Validate */
         if (janet_verify(def))
             janet_panic("funcdef has invalid bytecode");

src/include/janet.h

@@ -805,6 +805,7 @@ struct JanetAbstractHead {
 #define JANET_FUNCDEF_FLAG_HASENVS 0x400000
 #define JANET_FUNCDEF_FLAG_HASSOURCEMAP 0x800000
 #define JANET_FUNCDEF_FLAG_STRUCTARG 0x1000000
+#define JANET_FUNCDEF_FLAG_HASCLOBITSET 0x2000000
 #define JANET_FUNCDEF_FLAG_TAG 0xFFFF
 
 /* Source mapping structure for a bytecode instruction */
@@ -820,6 +821,7 @@ struct JanetFuncDef {
     Janet *constants;
     JanetFuncDef **defs;
     uint32_t *bytecode;
+    uint32_t *closure_bitset; /* Bit set indicating which slots can be referenced by closures. */
 
     /* Various debug information */
     JanetSourceMapping *sourcemap;

test/suite8.janet

@@ -175,4 +175,10 @@
   (assert (= 1 (f1)) "marshal-live-closure 1")
   (assert (= 2 (f2)) "marshal-live-closure 2"))
 
+(do
+  (var a 1)
+  (defn b [x] (+ a x))
+  (def c (unmarshal (marshal b)))
+  (assert (= 2 (c 1)) "marshal-on-stack-closure 1"))
+
 (end-suite)