Box: Free rb_classext_t struct for a box when the box is GCed

4 files changed, 50 insertions, 2 deletions

diff --git a/box.c b/box.c
index 282c575638..3fbc4e9fe2 100644
--- a/box.c
+++ b/box.c
@@ -149,6 +149,7 @@ box_entry_initialize(rb_box_t *box)
     box->loading_table = st_init_strtable();
     box->ruby_dln_libmap = rb_hash_new_with_size(0);
     box->gvar_tbl = rb_hash_new_with_size(0);
+    box->classext_cow_classes = st_init_numtable();
 
     box->is_user = true;
     box->is_optional = true;
@@ -199,6 +200,9 @@ rb_box_entry_mark(void *ptr)
     }
     rb_gc_mark(box->ruby_dln_libmap);
     rb_gc_mark(box->gvar_tbl);
+    if (box->classext_cow_classes) {
+        rb_mark_tbl(box->classext_cow_classes);
+    }
 }
 
 static int
@@ -233,9 +237,36 @@ box_root_free(void *ptr)
     }
 }
 
+static int
+free_classext_for_box(st_data_t _key, st_data_t obj_value, st_data_t box_arg)
+{
+    rb_classext_t *ext;
+    VALUE obj = (VALUE)obj_value;
+    const rb_box_t *box = (const rb_box_t *)box_arg;
+
+    if (RB_TYPE_P(obj, T_CLASS) || RB_TYPE_P(obj, T_MODULE)) {
+        ext = rb_class_unlink_classext(obj, box);
+        rb_class_classext_free(obj, ext, false);
+    }
+    else if (RB_TYPE_P(obj, T_ICLASS)) {
+        ext = rb_class_unlink_classext(obj, box);
+        rb_iclass_classext_free(obj, ext, false);
+    }
+    else {
+        rb_bug("Invalid type of object in classext_cow_classes: %s", rb_type_str(BUILTIN_TYPE(obj)));
+    }
+    return ST_CONTINUE;
+}
+
 static void
 box_entry_free(void *ptr)
 {
+    const rb_box_t *box = (const rb_box_t *)ptr;
+
+    if (box->classext_cow_classes) {
+        st_foreach(box->classext_cow_classes, free_classext_for_box, (st_data_t)box);
+    }
+
     box_root_free(ptr);
     xfree(ptr);
 }
@@ -750,6 +781,7 @@ initialize_root_box(void)
 
     root->ruby_dln_libmap = rb_hash_new_with_size(0);
     root->gvar_tbl = rb_hash_new_with_size(0);
+    root->classext_cow_classes = NULL; // classext CoW never happen on the root box
 
     vm->root_box = root;
 
diff --git a/class.c b/class.c
index 2f94b80147..0cec526b74 100644
--- a/class.c
+++ b/class.c
@@ -86,6 +86,17 @@ cvar_table_free_i(VALUE value, void *ctx)
     return ID_TABLE_CONTINUE;
 }
 
+rb_classext_t *
+rb_class_unlink_classext(VALUE klass, const rb_box_t *box)
+{
+    st_data_t ext;
+    st_data_t key = (st_data_t)box->box_object;
+    VALUE obj_id = rb_obj_id(klass);
+    st_delete(box->classext_cow_classes, &obj_id, 0);
+    st_delete(RCLASS_CLASSEXT_TBL(klass), &key, &ext);
+    return (rb_classext_t *)ext;
+}
+
 void
 rb_class_classext_free(VALUE klass, rb_classext_t *ext, bool is_prime)
 {
@@ -156,7 +167,7 @@ struct rb_class_set_box_classext_args {
 };
 
 static int
-rb_class_set_box_classext_update(st_data_t *key_ptr, st_data_t *val_ptr, st_data_t a, int existing)
+set_box_classext_update(st_data_t *key_ptr, st_data_t *val_ptr, st_data_t a, int existing)
 {
     struct rb_class_set_box_classext_args *args = (struct rb_class_set_box_classext_args *)a;
 
@@ -182,7 +193,10 @@ rb_class_set_box_classext(VALUE obj, const rb_box_t *box, rb_classext_t *ext)
         .ext = ext,
     };
 
-    st_update(RCLASS_CLASSEXT_TBL(obj), (st_data_t)box->box_object, rb_class_set_box_classext_update, (st_data_t)&args);
+    VM_ASSERT(BOX_USER_P(box));
+
+    st_update(RCLASS_CLASSEXT_TBL(obj), (st_data_t)box->box_object, set_box_classext_update, (st_data_t)&args);
+    st_insert(box->classext_cow_classes, (st_data_t)rb_obj_id(obj), obj);
 
     // FIXME: This is done here because this is the first time the objects in
     // the classext are exposed via this class. It's likely that if GC
diff --git a/internal/box.h b/internal/box.h
index 41cad63482..c341f046d3 100644
--- a/internal/box.h
+++ b/internal/box.h
@@ -34,6 +34,7 @@ struct rb_box_struct {
     VALUE ruby_dln_libmap;
 
     VALUE gvar_tbl;
+    struct st_table *classext_cow_classes;
 
     bool is_user;
     bool is_optional;
diff --git a/internal/class.h b/internal/class.h
index f122d2f189..296a07ae29 100644
--- a/internal/class.h
+++ b/internal/class.h
@@ -513,6 +513,7 @@ void rb_undef_methods_from(VALUE klass, VALUE super);
 VALUE rb_class_inherited(VALUE, VALUE);
 VALUE rb_keyword_error_new(const char *, VALUE);
 
+rb_classext_t *rb_class_unlink_classext(VALUE klass, const rb_box_t *box);
 void rb_class_classext_free(VALUE klass, rb_classext_t *ext, bool is_prime);
 void rb_iclass_classext_free(VALUE klass, rb_classext_t *ext, bool is_prime);