1 files changed, 54 insertions, 19 deletions
diff --git a/vm_method.c b/vm_method.c
index ac1f997545..d86cadc6c7 100644
--- a/vm_method.c
+++ b/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
@@ -526,9 +528,42 @@ rb_vm_ci_free(const struct rb_callinfo *ci)
 }
 
 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();
 }