Backport CallCache for refiment to ruby_3_3

gc.c

@@ -3801,6 +3801,12 @@ obj_free(rb_objspace_t *objspace, VALUE obj)
             }
           case imemo_callcache:
             RB_DEBUG_COUNTER_INC(obj_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_constcache:
             RB_DEBUG_COUNTER_INC(obj_imemo_constcache);
@@ -14468,3 +14474,21 @@ ruby_xrealloc2(void *ptr, size_t n, size_t new_size)
 #endif
     return ruby_xrealloc2_body(ptr, n, new_size);
 }
+
+int
+gc_update_references_weak_table_i(VALUE obj)
+{
+    int ret;
+    asan_unpoisoning_object(obj) {
+        ret = BUILTIN_TYPE(obj) == T_MOVED ? ST_REPLACE : ST_CONTINUE;
+    }
+    return ret;
+}
+
+int
+gc_update_references_weak_table_replace_i(VALUE *obj)
+{
+    *obj = rb_gc_location(*obj);
+
+    return ST_CONTINUE;
+}

internal/gc.h

@@ -362,3 +362,6 @@ ruby_sized_realloc_n(void *ptr, size_t new_count, size_t element_size, size_t ol
 #define ruby_sized_xrealloc2 ruby_sized_xrealloc2_inlined
 #define ruby_sized_xfree ruby_sized_xfree_inlined
 #endif /* INTERNAL_GC_H */
+
+int gc_update_references_weak_table_i(VALUE obj);
+int gc_update_references_weak_table_replace_i(VALUE *obj);

method.h

@@ -248,6 +248,9 @@ void rb_scope_visibility_set(rb_method_visibility_t);
 
 VALUE rb_unnamed_parameters(int arity);
 
+void rb_vm_insert_cc_refinement(st_table *cc_refinement_table, 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);
 

vm.c

@@ -2812,6 +2812,61 @@ rb_vm_call_cfunc(VALUE recv, VALUE (*func)(VALUE), VALUE arg,
 
 /* vm */
 
+struct global_vm_cc_refinement_foreach_data {
+    st_table *new_tbl;
+};
+
+static int
+vm_weak_table_cc_refinement_foreach_i(st_data_t key, st_data_t value, st_data_t data)
+{
+    struct global_vm_cc_refinement_foreach_data *cc_refinement_foreach_data = (struct global_vm_cc_refinement_foreach_data *)data;
+    struct st_table *new_tbl = cc_refinement_foreach_data->new_tbl;
+
+    int ret = gc_update_references_weak_table_i((VALUE)key);
+
+    const struct rb_callcache *cc = (struct rb_callcache *)key;
+    switch (ret) {
+        case ST_CONTINUE:
+            break;
+        case ST_DELETE:
+            return ret;
+        case ST_REPLACE: {
+            VALUE new_key = (VALUE)key;
+            gc_update_references_weak_table_replace_i(&new_key);
+            cc = (struct rb_callcache *)new_key;
+            break;
+        }
+        default:
+          rb_bug("vm_weak_table_cc_refinement_foreach_i: return value %d not supported", ret);
+    }
+
+    DURING_GC_COULD_MALLOC_REGION_START();
+    {
+        rb_vm_insert_cc_refinement(new_tbl, cc);
+    }
+    DURING_GC_COULD_MALLOC_REGION_END();
+
+    return ret;
+}
+
+static st_table *
+vm_weak_table_cc_refinement_rebuild(st_table *cc_refinement_table)
+{
+    st_table *new_tbl = NULL;
+    DURING_GC_COULD_MALLOC_REGION_START();
+    {
+        new_tbl = st_init_numtable_with_size(st_table_size(cc_refinement_table));
+    }
+    DURING_GC_COULD_MALLOC_REGION_END();
+
+    struct global_vm_cc_refinement_foreach_data cc_refinement_foreach_data = {
+        .new_tbl = new_tbl,
+    };
+
+    st_foreach(cc_refinement_table, vm_weak_table_cc_refinement_foreach_i, (st_data_t)&cc_refinement_foreach_data);
+    return new_tbl;
+}
+
 void
 rb_vm_update_references(void *ptr)
 {
@@ -2844,6 +2899,14 @@ rb_vm_update_references(void *ptr)
             vm->coverages = rb_gc_location(vm->coverages);
             vm->me2counter = rb_gc_location(vm->me2counter);
         }
+
+        if (vm->cc_refinement_table && st_table_size(vm->cc_refinement_table) > 0) {
+            st_table *old_tbl = vm->cc_refinement_table;
+            st_table *new_tbl = vm_weak_table_cc_refinement_rebuild(vm->cc_refinement_table);
+
+            vm->cc_refinement_table = new_tbl;
+            st_free_table(old_tbl);
+        }
     }
 }
 
@@ -3060,6 +3123,10 @@ ruby_vm_destruct(rb_vm_t *vm)
             st_free_table(vm->ci_table);
             vm->ci_table = NULL;
         }
+        if (vm->cc_refinement_table) {
+            st_free_table(vm->cc_refinement_table);
+            vm->cc_refinement_table = NULL;
+        }
         if (vm->frozen_strings) {
             st_free_table(vm->frozen_strings);
             vm->frozen_strings = 0;
@@ -3155,6 +3222,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_st_memsize(vm->cc_refinement_table) +
         vm_memsize_constant_cache() +
         GET_SHAPE_TREE()->cache_size * sizeof(redblack_node_t)
     );
@@ -4234,6 +4302,7 @@ Init_vm_objects(void)
     vm->mark_object_ary = rb_ary_hidden_new(128);
     vm->loading_table = st_init_strtable();
     vm->ci_table = st_init_table(&vm_ci_hashtype);
+    vm->cc_refinement_table = st_init_numtable();
     vm->frozen_strings = st_init_table_with_size(&rb_fstring_hash_type, 10000);
 }
 

vm_callinfo.h

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

vm_core.h

@@ -753,6 +753,7 @@ typedef struct rb_vm_struct {
     st_table *ci_table;
     struct rb_id_table *negative_cme_table;
     st_table *overloaded_cme_table; // cme -> overloaded_cme
+    st_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

@@ -314,27 +314,25 @@ rb_clear_method_cache(VALUE klass_or_module, ID mid)
 void rb_cc_table_free(VALUE klass);
 
 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 = asan_poisoned_object_p(v);
-        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 value, st_data_t data)
+{
+    VALUE v = (VALUE)key;
+    void *ptr = asan_poisoned_object_p(v);
+    asan_unpoison_object(v, false);
 
-        if (ptr) {
-            asan_poison_object(v);
+    if (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) {
+        asan_poison_object(v);
+    }
+
+    return ST_CONTINUE;
 }
 
 static st_index_t
@@ -451,10 +449,48 @@ rb_vm_ci_free(const struct rb_callinfo *ci)
     RB_VM_LOCK_LEAVE();
 }
 
+void
+rb_vm_insert_cc_refinement(st_table *cc_refinement_table, const struct rb_callcache *cc)
+{
+    st_data_t key = (st_data_t)cc;
+
+    if (cc_refinement_table) {
+        st_insert(cc_refinement_table, key, 1);
+    } else {
+        rb_vm_t *vm = GET_VM();
+        RB_VM_LOCK_ENTER();
+        {
+            st_insert(vm->cc_refinement_table, key, 1);
+        }
+        RB_VM_LOCK_LEAVE();
+    }
+}
+
+void rb_st_compact_table(st_table *tab);
+
+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;
+
+    st_delete(vm->cc_refinement_table, &key, NULL);
+}
+
 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();
+    {
+        st_foreach(vm->cc_refinement_table, invalidate_cc_refinement, (st_data_t)NULL);
+        st_clear(vm->cc_refinement_table);
+        rb_st_compact_table(vm->cc_refinement_table);
+    }
+    RB_VM_LOCK_LEAVE();
     rb_yjit_invalidate_all_method_lookup_assumptions();
 }