Fix up some edge cases with refcounting

3 files changed, 8 insertions, 6 deletions

diff --git a/builtins/datatypes.h b/builtins/datatypes.h
index 5ac84be5..e73e3c16 100644
--- a/builtins/datatypes.h
+++ b/builtins/datatypes.h
@@ -34,6 +34,8 @@ typedef struct {
     uint32_t next_bucket;
 } bucket_t;
 
+#define TABLE_MAX_DATA_REFCOUNT 3
+
 typedef struct {
     uint32_t count:31, last_free:31;
     uint8_t data_refcount:2;
diff --git a/builtins/table.c b/builtins/table.c
index f9623342..e0c94593 100644
--- a/builtins/table.c
+++ b/builtins/table.c
@@ -97,10 +97,10 @@ static inline void hshow(const table_t *t)
 
 static void maybe_copy_on_write(table_t *t, const TypeInfo *type)
 {
-    if (t->entries.data_refcount)
+    if (t->entries.data_refcount != 0)
         Array$compact(&t->entries, entry_size(type));
 
-    if (t->bucket_info && t->bucket_info->data_refcount) {
+    if (t->bucket_info && t->bucket_info->data_refcount != 0) {
         int64_t size = sizeof(bucket_info_t) + sizeof(bucket_t[t->bucket_info->count]);
         t->bucket_info = memcpy(GC_MALLOC(size), t->bucket_info, size);
         t->bucket_info->data_refcount = 0;
@@ -109,8 +109,8 @@ static void maybe_copy_on_write(table_t *t, const TypeInfo *type)
 
 public void Table$mark_copy_on_write(table_t *t)
 {
-    t->entries.data_refcount = 3;
-    if (t->bucket_info) t->bucket_info->data_refcount = 3;
+    ARRAY_INCREF(t->entries);
+    if (t->bucket_info) t->bucket_info->data_refcount = TABLE_MAX_DATA_REFCOUNT;
 }
 
 // Return address of value or NULL
diff --git a/compile.c b/compile.c
index 8157a79a..1bad1238 100644
--- a/compile.c
+++ b/compile.c
@@ -2134,7 +2134,7 @@ CORD compile(env_t *env, ast_t *ast)
         case TableType: {
             if (streq(f->field, "keys")) {
                 return CORD_all("({ table_t *t = ", compile_to_pointer_depth(env, f->fielded, 1, false), ";\n"
-                                "t->entries.data_refcount = 3;\n"
+                                "ARRAY_INCREF(t->entries);\n"
                                 "t->entries; })");
             } else if (streq(f->field, "values")) {
                 auto table = Match(value_t, TableType);
@@ -2143,7 +2143,7 @@ CORD compile(env_t *env, ast_t *ast)
                 if (align > 1 && offset % align > 0)
                     offset += align - (offset % align);
                 return CORD_all("({ table_t *t = ", compile_to_pointer_depth(env, f->fielded, 1, false), ";\n"
-                                "t->entries.data_refcount = 3;\n"
+                                "ARRAY_INCREF(t->entries.data_refcount);\n"
                                 "(array_t){.data = t->entries.data + ", CORD_asprintf("%zu", offset),
                                 ",\n .length=t->entries.length,\n .stride=t->entries.stride,\n .data_refcount=3};})");
             } else if (streq(f->field, "fallback")) {