Optimize callcache invalidation for refinements

Fixes [Bug #21201]

This change addresses a performance regression where defining methods
inside `refine` blocks caused severe slowdowns. The issue was due to
`rb_clear_all_refinement_method_cache()` triggering a full object
space scan via `rb_objspace_each_objects` to find and invalidate
affected callcaches, which is very inefficient.

To fix this, I introduce `vm->cc_refinement_table` to track
callcaches related to refinements. This allows us to invalidate
only the necessary callcaches without scanning the entire heap,
resulting in significant performance improvement.

Author: alpaca-tc

gc.c

@@ -2095,6 +2095,15 @@ rb_gc_obj_free_vm_weak_references(VALUE obj)
         break;
       case T_IMEMO:
         switch (imemo_type(obj)) {
+          case imemo_callcache: {
+            const struct rb_callcache *cc = (const struct rb_callcache *)obj;
+
+            if (vm_cc_refinement_p(cc)) {
+                rb_vm_delete_cc_refinement(cc);
+            }
+
+            break;
+          }
           case imemo_callinfo:
             rb_vm_ci_free((const struct rb_callinfo *)obj);
             break;
@@ -3929,6 +3938,23 @@ vm_weak_table_foreach_update_weak_key(st_data_t *key, st_data_t *value, st_data_
     return ret;
 }
 
+static int
+vm_weak_table_cc_refinement_foreach(st_data_t key, st_data_t data, int error)
+{
+    struct global_vm_table_foreach_data *iter_data = (struct global_vm_table_foreach_data *)data;
+
+    return iter_data->callback((VALUE)key, iter_data->data);
+}
+
+static int
+vm_weak_table_cc_refinement_foreach_update_update(st_data_t *key, st_data_t data, int existing)
+{
+    struct global_vm_table_foreach_data *iter_data = (struct global_vm_table_foreach_data *)data;
+
+    return iter_data->update_callback((VALUE *)key, iter_data->data);
+}
+
+
 static int
 vm_weak_table_str_sym_foreach(st_data_t key, st_data_t value, st_data_t data, int error)
 {
@@ -4178,8 +4204,21 @@ rb_gc_vm_weak_table_foreach(vm_table_foreach_callback_func callback,
         );
         break;
       }
+      case RB_GC_VM_CC_REFINEMENT_TABLE: {
+        if (vm->cc_refinement_table) {
+            set_foreach_with_replace(
+              vm->cc_refinement_table,
+              vm_weak_table_cc_refinement_foreach,
+              vm_weak_table_cc_refinement_foreach_update_update,
+              (st_data_t)&foreach_data
+            );
+        }
+        break;
+      }
       case RB_GC_VM_WEAK_TABLE_COUNT:
         rb_bug("Unreacheable");
+      default:
+        rb_bug("rb_gc_vm_weak_table_foreach: unknown table %d", table);
     }
 }
 

gc/gc.h

@@ -31,6 +31,7 @@ enum rb_gc_vm_weak_tables {
     RB_GC_VM_ID2REF_TABLE,
     RB_GC_VM_GENERIC_FIELDS_TABLE,
     RB_GC_VM_FROZEN_STRINGS_TABLE,
+    RB_GC_VM_CC_REFINEMENT_TABLE,
     RB_GC_VM_WEAK_TABLE_COUNT
 };
 

internal/set_table.h

@@ -37,6 +37,8 @@ size_t rb_set_table_size(const struct set_table *tbl);
 set_table *rb_set_init_table_with_size(set_table *tab, const struct st_hash_type *, st_index_t);
 #define set_init_numtable rb_set_init_numtable
 set_table *rb_set_init_numtable(void);
+#define set_init_numtable_with_size rb_set_init_numtable_with_size
+set_table *rb_set_init_numtable_with_size(st_index_t size);
 #define set_delete rb_set_delete
 int rb_set_delete(set_table *, st_data_t *); /* returns 0:notfound 1:deleted */
 #define set_insert rb_set_insert

method.h

@@ -254,6 +254,9 @@ void rb_scope_visibility_set(rb_method_visibility_t);
 
 VALUE rb_unnamed_parameters(int arity);
 
+void rb_vm_insert_cc_refinement(const struct rb_callcache *cc);
+void rb_vm_delete_cc_refinement(const struct rb_callcache *cc);
+
 void rb_clear_method_cache(VALUE klass_or_module, ID mid);
 void rb_clear_all_refinement_method_cache(void);
 void rb_invalidate_method_caches(struct rb_id_table *cm_tbl, struct rb_id_table *cc_tbl);

st.c

@@ -2465,6 +2465,12 @@ set_init_numtable(void)
     return set_init_table_with_size(NULL, &type_numhash, 0);
 }
 
+set_table *
+set_init_numtable_with_size(st_index_t size)
+{
+    return set_init_table_with_size(NULL, &type_numhash, size);
+}
+
 size_t
 set_table_size(const struct set_table *tbl)
 {

vm.c

@@ -3194,6 +3194,10 @@ ruby_vm_destruct(rb_vm_t *vm)
             st_free_table(vm->ci_table);
             vm->ci_table = NULL;
         }
+        if (vm->cc_refinement_table) {
+            rb_set_free_table(vm->cc_refinement_table);
+            vm->cc_refinement_table = NULL;
+        }
         RB_ALTSTACK_FREE(vm->main_altstack);
 
         struct global_object_list *next;
@@ -3294,6 +3298,7 @@ vm_memsize(const void *ptr)
         vm_memsize_builtin_function_table(vm->builtin_function_table) +
         rb_id_table_memsize(vm->negative_cme_table) +
         rb_st_memsize(vm->overloaded_cme_table) +
+        rb_set_memsize(vm->cc_refinement_table) +
         vm_memsize_constant_cache()
     );
 
@@ -4503,6 +4508,7 @@ Init_vm_objects(void)
     vm->mark_object_ary = pin_array_list_new(Qnil);
     vm->loading_table = st_init_strtable();
     vm->ci_table = st_init_table(&vm_ci_hashtype);
+    vm->cc_refinement_table = rb_set_init_numtable();
 }
 
 // Stub for builtin function when not building YJIT units

vm_callinfo.h

@@ -345,6 +345,7 @@ vm_cc_new(VALUE klass,
         break;
       case cc_type_refinement:
         *(VALUE *)&cc->flags |= VM_CALLCACHE_REFINEMENT;
+        rb_vm_insert_cc_refinement(cc);
         break;
     }
 

vm_core.h

@@ -803,6 +803,7 @@ typedef struct rb_vm_struct {
     struct rb_id_table *negative_cme_table;
     st_table *overloaded_cme_table; // cme -> overloaded_cme
     set_table *unused_block_warning_table;
+    set_table *cc_refinement_table;
 
     // This id table contains a mapping from ID to ICs. It does this with ID
     // keys and nested st_tables as values. The nested tables have ICs as keys

vm_method.c

@@ -393,27 +393,29 @@ rb_invalidate_method_caches(struct rb_id_table *cm_tbl, struct rb_id_table *cc_t
 }
 
 static int
-invalidate_all_refinement_cc(void *vstart, void *vend, size_t stride, void *data)
-{
-    VALUE v = (VALUE)vstart;
-    for (; v != (VALUE)vend; v += stride) {
-        void *ptr = rb_asan_poisoned_object_p(v);
-        rb_asan_unpoison_object(v, false);
-
-        if (RBASIC(v)->flags) { // liveness check
-            if (imemo_type_p(v, imemo_callcache)) {
-                const struct rb_callcache *cc = (const struct rb_callcache *)v;
-                if (vm_cc_refinement_p(cc) && cc->klass) {
-                    vm_cc_invalidate(cc);
-                }
-            }
-        }
+invalidate_cc_refinement(st_data_t key, st_data_t data)
+{
+    VALUE v = (VALUE)key;
+    void *ptr = rb_asan_poisoned_object_p(v);
+    rb_asan_unpoison_object(v, false);
 
-        if (ptr) {
-            rb_asan_poison_object(v);
+    if (rb_gc_pointer_to_heap_p(v) &&
+        !rb_objspace_garbage_object_p(v) &&
+        RBASIC(v)->flags) { // liveness check
+        const struct rb_callcache *cc = (const struct rb_callcache *)v;
+
+        VM_ASSERT(vm_cc_refinement_p(cc));
+
+        if (cc->klass) {
+            vm_cc_invalidate(cc);
         }
     }
-    return 0; // continue to iteration
+
+    if (ptr) {
+        rb_asan_poison_object(v);
+    }
+
+    return ST_CONTINUE;
 }
 
 static st_index_t
@@ -525,10 +527,43 @@ rb_vm_ci_free(const struct rb_callinfo *ci)
     st_delete(vm->ci_table, &key, NULL);
 }
 
+void
+rb_vm_insert_cc_refinement(const struct rb_callcache *cc)
+{
+    st_data_t key = (st_data_t)cc;
+
+    rb_vm_t *vm = GET_VM();
+    RB_VM_LOCK_ENTER();
+    {
+        rb_set_insert(vm->cc_refinement_table, key);
+    }
+    RB_VM_LOCK_LEAVE();
+}
+
+void
+rb_vm_delete_cc_refinement(const struct rb_callcache *cc)
+{
+    ASSERT_vm_locking();
+
+    rb_vm_t *vm = GET_VM();
+    st_data_t key = (st_data_t)cc;
+
+    rb_set_delete(vm->cc_refinement_table, &key);
+}
+
 void
 rb_clear_all_refinement_method_cache(void)
 {
-    rb_objspace_each_objects(invalidate_all_refinement_cc, NULL);
+    rb_vm_t *vm = GET_VM();
+
+    RB_VM_LOCK_ENTER();
+    {
+        rb_set_foreach(vm->cc_refinement_table, invalidate_cc_refinement, (st_data_t)NULL);
+        rb_set_clear(vm->cc_refinement_table);
+        rb_set_compact_table(vm->cc_refinement_table);
+    }
+    RB_VM_LOCK_LEAVE();
+
     rb_yjit_invalidate_all_method_lookup_assumptions();
 }