Implement marked_id_table

Similar to `managed_id_table`, but the values are marked.
This can be used to implement RCU pattern on id tables
that contain marked objects (e.g. `class->const_tbl`).

Author: byroot

id_table.c

@@ -87,8 +87,8 @@ rb_id_table_init(struct rb_id_table *tbl, size_t s_capa)
     MEMZERO(tbl, struct rb_id_table, 1);
     if (capa > 0) {
         capa = round_capa(capa);
-        tbl->capa = (int)capa;
         tbl->items = ZALLOC_N(item_t, capa);
+        tbl->capa = (int)capa;
     }
     return tbl;
 }
@@ -341,7 +341,7 @@ managed_id_table_free(void *data)
 static size_t
 managed_id_table_memsize(const void *data)
 {
-    const struct rb_id_table *tbl = (const struct rb_id_table *)data;
+    const struct rb_id_table *tbl = (struct rb_id_table *)data;
     return rb_id_table_memsize(tbl) - sizeof(struct rb_id_table);
 }
 
@@ -427,3 +427,142 @@ rb_managed_id_table_foreach(VALUE table, rb_id_table_foreach_func_t *func, void
 
     rb_id_table_foreach(RTYPEDDATA_GET_DATA(table), func, data);
 }
+
+static enum rb_id_table_iterator_result
+marked_id_table_mark_i(ID key, VALUE value, void *ref)
+{
+    rb_gc_mark_movable(value);
+    return ID_TABLE_CONTINUE;
+}
+
+static void
+marked_id_table_mark(void *data)
+{
+    struct rb_id_table *tbl = (struct rb_id_table *)data;
+    rb_id_table_foreach(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 *data)
+{
+    struct rb_id_table *tbl = (struct rb_id_table *)data;
+    rb_id_table_foreach_values_with_replace(tbl, marked_id_table_compact_check_i, marked_id_table_compact_replace_i, NULL);
+}
+
+static const rb_data_type_t marked_id_table_type = {
+    .wrap_struct_name = "VM/marked_id_table",
+    .function = {
+        .dmark = marked_id_table_mark,
+        .dfree = (RUBY_DATA_FUNC)managed_id_table_free,
+        .dsize = managed_id_table_memsize,
+        .dcompact = marked_id_table_compact,
+    },
+    .flags = RUBY_TYPED_FREE_IMMEDIATELY | RUBY_TYPED_WB_PROTECTED | RUBY_TYPED_EMBEDDABLE,
+};
+
+VALUE
+rb_marked_id_table_new(size_t capa)
+{
+    struct rb_id_table *tbl;
+    VALUE obj = TypedData_Make_Struct(0, struct rb_id_table, &marked_id_table_type, tbl);
+    rb_id_table_init(tbl, capa);
+    return obj;
+}
+
+struct marked_id_table_dup_args {
+    struct rb_id_table *tbl;
+    VALUE owner;
+};
+
+static enum rb_id_table_iterator_result
+marked_id_table_dup_i(ID id, VALUE val, void *data)
+{
+    struct marked_id_table_dup_args *args = (struct marked_id_table_dup_args *)data;
+
+    rb_id_table_insert(args->tbl, id, val);
+    RB_OBJ_WRITTEN(args->owner, Qundef, val);
+    return ID_TABLE_CONTINUE;
+}
+
+size_t
+rb_marked_id_table_size(VALUE table)
+{
+    RUBY_ASSERT(RB_TYPE_P(table, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(table), &marked_id_table_type));
+
+    return rb_id_table_size(RTYPEDDATA_GET_DATA(table));
+}
+
+VALUE
+rb_marked_id_table_dup(VALUE old_table)
+{
+    RUBY_ASSERT(RB_TYPE_P(old_table, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(old_table), &marked_id_table_type));
+
+    struct rb_id_table *new_tbl;
+    VALUE obj = TypedData_Make_Struct(0, struct rb_id_table, &marked_id_table_type, new_tbl);
+    struct rb_id_table *old_tbl = RTYPEDDATA_GET_DATA(old_table);
+
+    struct marked_id_table_dup_args args = {
+        .tbl = new_tbl,
+        .owner = obj,
+    };
+
+    rb_id_table_init(new_tbl, old_tbl->num + 1);
+    rb_id_table_foreach(old_tbl, marked_id_table_dup_i, (void *)&args);
+    return obj;
+}
+
+int
+rb_marked_id_table_lookup(VALUE table, ID id, VALUE *valp)
+{
+    RUBY_ASSERT(RB_TYPE_P(table, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(table), &marked_id_table_type));
+
+    return rb_id_table_lookup(RTYPEDDATA_GET_DATA(table), id, valp);
+}
+
+int
+rb_marked_id_table_insert(VALUE table, ID id, VALUE val)
+{
+    RUBY_ASSERT(RB_TYPE_P(table, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(table), &marked_id_table_type));
+
+    int retval = rb_id_table_insert(RTYPEDDATA_GET_DATA(table), id, val);
+    RB_OBJ_WRITTEN(table, Qundef, val);
+    return retval;
+}
+
+int
+rb_marked_id_table_delete(VALUE table, ID id)
+{
+    RUBY_ASSERT(RB_TYPE_P(table, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(table), &marked_id_table_type));
+
+    return rb_id_table_delete(RTYPEDDATA_GET_DATA(table), id);
+}
+
+void
+rb_marked_id_table_foreach(VALUE table, rb_id_table_foreach_func_t *func, void *data)
+{
+    RUBY_ASSERT(RB_TYPE_P(table, T_DATA));
+    RUBY_ASSERT(rb_typeddata_inherited_p(RTYPEDDATA_TYPE(table), &marked_id_table_type));
+
+    rb_id_table_foreach(RTYPEDDATA_GET_DATA(table), func, data);
+}

id_table.h

@@ -42,6 +42,15 @@ int rb_managed_id_table_lookup(VALUE table, ID id, VALUE *valp);
 size_t rb_managed_id_table_size(VALUE table);
 void rb_managed_id_table_foreach(VALUE table, rb_id_table_foreach_func_t *func, void *data);
 
+VALUE rb_marked_id_table_new(size_t capa);
+VALUE rb_marked_id_table_dup(VALUE table);
+int rb_marked_id_table_insert(VALUE table, ID id, VALUE val);
+int rb_marked_id_table_lookup(VALUE table, ID id, VALUE *valp);
+int rb_marked_id_table_delete(VALUE table, ID id);
+
+size_t rb_marked_id_table_size(VALUE table);
+void rb_marked_id_table_foreach(VALUE table, rb_id_table_foreach_func_t *func, void *data);
+
 RUBY_SYMBOL_EXPORT_BEGIN
 size_t rb_id_table_size(const struct rb_id_table *tbl);
 RUBY_SYMBOL_EXPORT_END