1 files changed, 197 insertions, 11 deletions

diff --git a/id_table.c b/id_table.c
index 6bb067d09a..299b7d3ae5 100644
--- a/id_table.c
+++ b/id_table.c
@@ -38,16 +38,9 @@ typedef struct rb_id_item {
     VALUE    val;
 } item_t;
 
-struct rb_id_table {
-    int capa;
-    int num;
-    int used;
-    item_t *items;
-};
-
 #if SIZEOF_VALUE == 8
 #define ITEM_GET_KEY(tbl, i) ((tbl)->items[i].key)
-#define ITEM_KEY_ISSET(tbl, i) ((tbl)->items[i].key)
+#define ITEM_KEY_ISSET(tbl, i) ((tbl)->items && (tbl)->items[i].key)
 #define ITEM_COLLIDED(tbl, i) ((tbl)->items[i].collision)
 #define ITEM_SET_COLLIDED(tbl, i) ((tbl)->items[i].collision = 1)
 static inline void
@@ -80,9 +73,10 @@ round_capa(int capa)
     return (capa + 1) << 2;
 }
 
-static struct rb_id_table *
-rb_id_table_init(struct rb_id_table *tbl, int capa)
+struct rb_id_table *
+rb_id_table_init(struct rb_id_table *tbl, size_t s_capa)
 {
+    int capa = (int)s_capa;
     MEMZERO(tbl, struct rb_id_table, 1);
     if (capa > 0) {
         capa = round_capa(capa);
@@ -96,7 +90,13 @@ struct rb_id_table *
 rb_id_table_create(size_t capa)
 {
     struct rb_id_table *tbl = ALLOC(struct rb_id_table);
-    return rb_id_table_init(tbl, (int)capa);
+    return rb_id_table_init(tbl, capa);
+}
+
+void
+rb_id_table_free_items(struct rb_id_table *tbl)
+{
+    xfree(tbl->items);
 }
 
 void
@@ -291,6 +291,10 @@ rb_id_table_foreach_values(struct rb_id_table *tbl, rb_id_table_foreach_values_f
 {
     int i, capa = tbl->capa;
 
+    if (!tbl->items) {
+        return;
+    }
+
     for (i=0; i<capa; i++) {
         if (ITEM_KEY_ISSET(tbl, i)) {
             enum rb_id_table_iterator_result ret = (*func)(tbl->items[i].val, data);
@@ -324,3 +328,185 @@ rb_id_table_foreach_values_with_replace(struct rb_id_table *tbl, rb_id_table_for
     }
 }
 
+static void
+managed_id_table_free(void *data)
+{
+    struct rb_id_table *tbl = (struct rb_id_table *)data;
+    rb_id_table_free_items(tbl);
+}
+
+static size_t
+managed_id_table_memsize(const void *data)
+{
+    const struct rb_id_table *tbl = (const struct rb_id_table *)data;
+    return rb_id_table_memsize(tbl) - sizeof(struct rb_id_table);
+}
+
+const rb_data_type_t rb_managed_id_table_type = {
+    .wrap_struct_name = "VM/managed_id_table",
+    .function = {
+        .dmark = NULL, // Nothing to mark
+        .dfree = managed_id_table_free,
+        .dsize = managed_id_table_memsize,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE,
+};
+
+static inline struct rb_id_table *
+managed_id_table_ptr(VALUE obj)
+{
+    RUBY_ASSERT(RB_TYPE_P(obj, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(obj), &rb_managed_id_table_type));
+
+    return RTYPEDDATA_GET_DATA(obj);
+}
+
+VALUE
+rb_managed_id_table_create(const rb_data_type_t *type, size_t capa)
+{
+    struct rb_id_table *tbl;
+    VALUE obj = TypedData_Make_Struct(0, struct rb_id_table, type, tbl);
+    RB_OBJ_SET_SHAREABLE(obj);
+    rb_id_table_init(tbl, capa); // NOTE: this can cause GC, so dmark and dsize need to check tbl->items
+    return obj;
+}
+
+VALUE
+rb_managed_id_table_new(size_t capa)
+{
+    return rb_managed_id_table_create(&rb_managed_id_table_type, capa);
+}
+
+static enum rb_id_table_iterator_result
+managed_id_table_dup_i(ID id, VALUE val, void *data)
+{
+    struct rb_id_table *new_tbl = (struct rb_id_table *)data;
+    rb_id_table_insert(new_tbl, id, val);
+    return ID_TABLE_CONTINUE;
+}
+
+VALUE
+rb_managed_id_table_dup(VALUE old_table)
+{
+    struct rb_id_table *new_tbl;
+    VALUE obj = TypedData_Make_Struct(0, struct rb_id_table, RTYPEDDATA_TYPE(old_table), new_tbl);
+    struct rb_id_table *old_tbl = managed_id_table_ptr(old_table);
+    rb_id_table_init(new_tbl, old_tbl->num + 1);
+    rb_id_table_foreach(old_tbl, managed_id_table_dup_i, new_tbl);
+    return obj;
+}
+
+int
+rb_managed_id_table_lookup(VALUE table, ID id, VALUE *valp)
+{
+    return rb_id_table_lookup(managed_id_table_ptr(table), id, valp);
+}
+
+int
+rb_managed_id_table_insert(VALUE table, ID id, VALUE val)
+{
+    return rb_id_table_insert(managed_id_table_ptr(table), id, val);
+}
+
+size_t
+rb_managed_id_table_size(VALUE table)
+{
+    return rb_id_table_size(managed_id_table_ptr(table));
+}
+
+void
+rb_managed_id_table_foreach(VALUE table, rb_id_table_foreach_func_t *func, void *data)
+{
+    rb_id_table_foreach(managed_id_table_ptr(table), func, data);
+}
+
+void
+rb_managed_id_table_foreach_values(VALUE table, rb_id_table_foreach_values_func_t *func, void *data)
+{
+    rb_id_table_foreach_values(managed_id_table_ptr(table), func, data);
+}
+
+int
+rb_managed_id_table_delete(VALUE table, ID id)
+{
+    return rb_id_table_delete(managed_id_table_ptr(table), id);
+}
+
+static enum rb_id_table_iterator_result
+marked_id_table_mark_i(VALUE val, void *data)
+{
+    rb_gc_mark_movable(val);
+    return ID_TABLE_CONTINUE;
+}
+
+static void
+marked_id_table_mark(void *ptr)
+{
+    struct rb_id_table *tbl = (struct rb_id_table *)ptr;
+    rb_id_table_foreach_values(tbl, marked_id_table_mark_i, NULL);
+}
+
+static enum rb_id_table_iterator_result
+marked_id_table_compact_check_i(VALUE value, void *data)
+{
+    if (rb_gc_location(value) != value) {
+        return ID_TABLE_REPLACE;
+    }
+    return ID_TABLE_CONTINUE;
+}
+
+static enum rb_id_table_iterator_result
+marked_id_table_compact_replace_i(VALUE *value, void *data, int existing)
+{
+    *value = rb_gc_location(*value);
+    return ID_TABLE_CONTINUE;
+}
+
+static void
+marked_id_table_compact(void *ptr)
+{
+    struct rb_id_table *tbl = (struct rb_id_table *)ptr;
+    rb_id_table_foreach_values_with_replace(tbl, marked_id_table_compact_check_i, marked_id_table_compact_replace_i, NULL);
+}
+
+const rb_data_type_t rb_marked_id_table_type = {
+    .wrap_struct_name = "VM/marked_id_table",
+    .function = {
+        .dmark = marked_id_table_mark,
+        .dfree = managed_id_table_free,
+        .dsize = managed_id_table_memsize,
+        .dcompact = marked_id_table_compact,
+    },
+    .parent = &rb_managed_id_table_type,
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE,
+};
+
+VALUE
+rb_marked_id_table_new(size_t capa)
+{
+    return rb_managed_id_table_create(&rb_marked_id_table_type, capa);
+}
+
+int
+rb_marked_id_table_insert(VALUE table, ID id, VALUE val)
+{
+    int result = rb_managed_id_table_insert(table, id, val);
+    RB_OBJ_WRITTEN(table, Qundef, val);
+    return result;
+}
+
+static enum rb_id_table_iterator_result
+marked_id_table_dup_i(VALUE val, void *data)
+{
+    VALUE new_table = (VALUE)data;
+    RB_OBJ_WRITTEN(new_table, Qundef, val);
+    return ID_TABLE_CONTINUE;
+}
+
+VALUE
+rb_marked_id_table_dup(VALUE old_table)
+{
+    VALUE new_table = rb_managed_id_table_dup(old_table);
+    rb_managed_id_table_foreach_values(new_table, marked_id_table_dup_i, (void *)new_table);
+    return new_table;
+}