6 files changed, 170 insertions(+), 9 deletions(-)

A examples/weak-tables.janet
M src/core/gc.c
M src/core/gc.h
M src/core/state.h
M src/core/table.c
M src/core/vm.c

@@ 0,0 1,20 @@
+(def weak-k (table/new 10 :k))
+(def weak-v (table/new 10 :v))
+(def weak-kv (table/new 10 :kv))
+
+(put weak-kv (gensym) 10)
+(put weak-kv :hello :world)
+(put weak-k :abc123zz77asda :stuff)
+(put weak-k true :abc123zz77asda)
+(put weak-k :zyzzyz false)
+(put weak-v (gensym) 10)
+(put weak-v 20 (gensym))
+(print "before gc")
+(tracev weak-k)
+(tracev weak-v)
+(tracev weak-kv)
+(gccollect)
+(print "after gc")
+(tracev weak-k)
+(tracev weak-v)
+(tracev weak-kv)

@@ 133,6 133,24 @@ static void janet_mark_many(const Janet *values, int32_t n) {
 }
 
 /* Mark a bunch of key values items in memory */
+static void janet_mark_keys(const JanetKV *kvs, int32_t n) {
+    const JanetKV *end = kvs + n;
+    while (kvs < end) {
+        janet_mark(kvs->key);
+        kvs++;
+    }
+}
+
+/* Mark a bunch of key values items in memory */
+static void janet_mark_values(const JanetKV *kvs, int32_t n) {
+    const JanetKV *end = kvs + n;
+    while (kvs < end) {
+        janet_mark(kvs->value);
+        kvs++;
+    }
+}
+
+/* Mark a bunch of key values items in memory */
 static void janet_mark_kvs(const JanetKV *kvs, int32_t n) {
     const JanetKV *end = kvs + n;
     while (kvs < end) {


@@ 154,7 172,15 @@ recur: /* Manual tail recursion */
     if (janet_gc_reachable(table))
         return;
     janet_gc_mark(table);
-    janet_mark_kvs(table->data, table->capacity);
+    enum JanetMemoryType memtype = janet_gc_type(table);
+    if (memtype == JANET_MEMORY_TABLE_WEAKK) {
+        janet_mark_values(table->data, table->capacity);
+    } else if (memtype == JANET_MEMORY_TABLE_WEAKV) {
+        janet_mark_keys(table->data, table->capacity);
+    } else if (memtype == JANET_MEMORY_TABLE) {
+        janet_mark_kvs(table->data, table->capacity);
+    }
+    /* do nothing for JANET_MEMORY_TABLE_WEAKKV */
     if (table->proto) {
         table = table->proto;
         goto recur;


@@ 329,12 355,89 @@ static void janet_deinit_block(JanetGCObject *mem) {
     }
 }
 
+/* Check that a value x has been visited in the mark phase */
+static int janet_check_liveref(Janet x) {
+    switch (janet_type(x)) {
+        default:
+            return 1;
+        case JANET_ARRAY:
+        case JANET_TABLE:
+        case JANET_FUNCTION:
+        case JANET_BUFFER:
+        case JANET_FIBER:
+            return janet_gc_reachable(janet_unwrap_pointer(x));
+        case JANET_STRING:
+        case JANET_SYMBOL:
+        case JANET_KEYWORD:
+            return janet_gc_reachable(janet_string_head(janet_unwrap_string(x)));
+        case JANET_ABSTRACT:
+            return janet_gc_reachable(janet_abstract_head(janet_unwrap_abstract(x)));
+        case JANET_TUPLE:
+            return janet_gc_reachable(janet_tuple_head(janet_unwrap_tuple(x)));
+        case JANET_STRUCT:
+            return janet_gc_reachable(janet_struct_head(janet_unwrap_struct(x)));
+    }
+}
+
 /* Iterate over all allocated memory, and free memory that is not
  * marked as reachable. Flip the gc color flag for next sweep. */
 void janet_sweep() {
     JanetGCObject *previous = NULL;
-    JanetGCObject *current = janet_vm.blocks;
+    JanetGCObject *current = janet_vm.weak_blocks;
     JanetGCObject *next;
+
+    /* Sweep weak heap to drop weak refs */
+    while (NULL != current) {
+        next = current->data.next;
+        if (current->flags & JANET_MEM_REACHABLE) {
+            /* Check for dead references */
+            enum JanetMemoryType type = janet_gc_type(current);
+            JanetTable *table = (JanetTable *) current;
+            int check_values = (type == JANET_MEMORY_TABLE_WEAKV) || (type == JANET_MEMORY_TABLE_WEAKKV);
+            int check_keys = (type == JANET_MEMORY_TABLE_WEAKK) || (type == JANET_MEMORY_TABLE_WEAKKV);
+            JanetKV *end = table->data + table->capacity;
+            JanetKV *kvs = table->data;
+            while (kvs < end) {
+                int drop = 0;
+                if (check_keys && !janet_check_liveref(kvs->key)) drop = 1;
+                if (check_values && !janet_check_liveref(kvs->value)) drop = 1;
+                if (drop) {
+                    /* Inlined from janet_table_remove without search */
+                    table->count--;
+                    table->deleted++;
+                    kvs->key = janet_wrap_nil();
+                    kvs->value = janet_wrap_false();
+                }
+                kvs++;
+            }
+        }
+        current = next;
+    }
+
+    /* Sweep weak heap to free blocks */
+    previous = NULL;
+    current = janet_vm.weak_blocks;
+    while (NULL != current) {
+        next = current->data.next;
+        if (current->flags & (JANET_MEM_REACHABLE | JANET_MEM_DISABLED)) {
+            previous = current;
+            current->flags &= ~JANET_MEM_REACHABLE;
+        } else {
+            janet_vm.block_count--;
+            janet_deinit_block(current);
+            if (NULL != previous) {
+                previous->data.next = next;
+            } else {
+                janet_vm.weak_blocks = next;
+            }
+            janet_free(current);
+        }
+        current = next;
+    }
+
+    /* Sweep main heap to free blocks */
+    previous = NULL;
+    current = janet_vm.blocks;
     while (NULL != current) {
         next = current->data.next;
         if (current->flags & (JANET_MEM_REACHABLE | JANET_MEM_DISABLED)) {


@@ 352,6 455,7 @@ void janet_sweep() {
         }
         current = next;
     }
+
 #ifdef JANET_EV
     /* Sweep threaded abstract types for references to decrement */
     JanetKV *items = janet_vm.threaded_abstracts.data;


@@ 409,8 513,15 @@ void *janet_gcalloc(enum JanetMemoryType type, size_t size) {
 
     /* Prepend block to heap list */
     janet_vm.next_collection += size;
-    mem->data.next = janet_vm.blocks;
-    janet_vm.blocks = mem;
+    if (type < JANET_MEMORY_TABLE_WEAKK) {
+        /* normal heap */
+        mem->data.next = janet_vm.blocks;
+        janet_vm.blocks = mem;
+    } else {
+        /* weak heap */
+        mem->data.next = janet_vm.weak_blocks;
+        janet_vm.weak_blocks = mem;
+    }
     janet_vm.block_count++;
 
     return (void *)mem;


@@ 442,7 553,7 @@ void janet_collect(void) {
     if (janet_vm.gc_suspend) return;
     depth = JANET_RECURSION_GUARD;
     janet_vm.gc_mark_phase = 1;
-    /* Try and prevent many major collections back to back.
+    /* Try to prevent many major collections back to back.
      * A full collection will take O(janet_vm.block_count) time.
      * If we have a large heap, make sure our interval is not too
      * small so we won't make many collections over it. This is just a

@@ 57,6 57,9 @@ enum JanetMemoryType {
     JANET_MEMORY_FUNCENV,
     JANET_MEMORY_FUNCDEF,
     JANET_MEMORY_THREADED_ABSTRACT,
+    JANET_MEMORY_TABLE_WEAKK,
+    JANET_MEMORY_TABLE_WEAKV,
+    JANET_MEMORY_TABLE_WEAKKV
 };
 
 /* To allocate collectable memory, one must call janet_alloc, initialize the memory,

@@ 121,6 121,7 @@ struct JanetVM {
 
     /* Garbage collection */
     void *blocks;
+    void *weak_blocks;
     size_t gc_interval;
     size_t next_collection;
     size_t block_count;

@@ 87,11 87,27 @@ void janet_table_deinit(JanetTable *table) {
 }
 
 /* Create a new table */
+
 JanetTable *janet_table(int32_t capacity) {
     JanetTable *table = janet_gcalloc(JANET_MEMORY_TABLE, sizeof(JanetTable));
     return janet_table_init_impl(table, capacity, 0);
 }
 
+JanetTable *janet_table_weakk(int32_t capacity) {
+    JanetTable *table = janet_gcalloc(JANET_MEMORY_TABLE_WEAKK, sizeof(JanetTable));
+    return janet_table_init_impl(table, capacity, 0);
+}
+
+JanetTable *janet_table_weakv(int32_t capacity) {
+    JanetTable *table = janet_gcalloc(JANET_MEMORY_TABLE_WEAKV, sizeof(JanetTable));
+    return janet_table_init_impl(table, capacity, 0);
+}
+
+JanetTable *janet_table_weakkv(int32_t capacity) {
+    JanetTable *table = janet_gcalloc(JANET_MEMORY_TABLE_WEAKKV, sizeof(JanetTable));
+    return janet_table_init_impl(table, capacity, 0);
+}
+
 /* Find the bucket that contains the given key. Will also return
  * bucket where key should go if not in the table. */
 JanetKV *janet_table_find(JanetTable *t, Janet key) {


@@ 293,14 309,23 @@ JanetTable *janet_table_proto_flatten(JanetTable *t) {
 /* C Functions */
 
 JANET_CORE_FN(cfun_table_new,
-              "(table/new capacity)",
+              "(table/new capacity &opt weak)",
               "Creates a new empty table with pre-allocated memory "
               "for `capacity` entries. This means that if one knows the number of "
               "entries going into a table on creation, extra memory allocation "
-              "can be avoided. Returns the new table.") {
-    janet_fixarity(argc, 1);
+              "can be avoided. Optionally provide a keyword flags `:kv` to create a table with "
+              "weak referenecs to keys, values, or both. "
+              "Returns the new table.") {
+    janet_arity(argc, 1, 2);
     int32_t cap = janet_getnat(argv, 0);
-    return janet_wrap_table(janet_table(cap));
+    if (argc == 1) {
+        return janet_wrap_table(janet_table(cap));
+    }
+    uint32_t flags = janet_getflags(argv, 1, "kv");
+    if (flags == 0) return janet_wrap_table(janet_table(cap));
+    if (flags == 1) return janet_wrap_table(janet_table_weakk(cap));
+    if (flags == 2) return janet_wrap_table(janet_table_weakv(cap));
+    return janet_wrap_table(janet_table_weakkv(cap));
 }
 
 JANET_CORE_FN(cfun_table_getproto,

@@ 1585,6 1585,7 @@ int janet_init(void) {
 
     /* Garbage collection */
     janet_vm.blocks = NULL;
+    janet_vm.weak_blocks = NULL;
     janet_vm.next_collection = 0;
     janet_vm.gc_interval = 0x400000;
     janet_vm.block_count = 0;