Implement callcache using declare weak references

3 files changed, 13 insertions, 1 deletions

diff --git a/gc.c b/gc.c
index ba358a656e..47e6d9a6b6 100644
--- a/gc.c
+++ b/gc.c
@@ -1209,6 +1209,17 @@ rb_gc_handle_weak_references(VALUE obj)
            rb_bug("rb_gc_handle_weak_references: unknown T_DATA");
         }
         break;
+
+      case T_IMEMO: {
+        GC_ASSERT(imemo_type(obj) == imemo_callcache);
+
+        struct rb_callcache *cc = (struct rb_callcache *)obj;
+        if (!rb_gc_handle_weak_references_alive_p(cc->klass)) {
+            vm_cc_invalidate(cc);
+        }
+
+        break;
+      }
       default:
         rb_bug("rb_gc_handle_weak_references: type not supported\n");
     }
diff --git a/imemo.c b/imemo.c
index 42f6615a5e..2d63b7cd99 100644
--- a/imemo.c
+++ b/imemo.c
@@ -372,7 +372,6 @@ rb_imemo_mark_and_move(VALUE obj, bool reference_updating)
             RUBY_ASSERT(RB_TYPE_P(cc->klass, T_CLASS) || RB_TYPE_P(cc->klass, T_ICLASS));
             RUBY_ASSERT(IMEMO_TYPE_P((VALUE)cc->cme_, imemo_ment));
 
-            rb_gc_mark_weak((VALUE *)&cc->klass);
             if ((vm_cc_super_p(cc) || vm_cc_refinement_p(cc))) {
                 rb_gc_mark_movable((VALUE)cc->cme_);
             }
diff --git a/vm_callinfo.h b/vm_callinfo.h
index 2c51bf9092..3df4b4eb0e 100644
--- a/vm_callinfo.h
+++ b/vm_callinfo.h
@@ -342,6 +342,8 @@ vm_cc_new(VALUE klass,
 {
     cc_check_class(klass);
     struct rb_callcache *cc = SHAREABLE_IMEMO_NEW(struct rb_callcache, imemo_callcache, klass);
+    rb_gc_declare_weak_references((VALUE)cc);
+
     *((struct rb_callable_method_entry_struct **)&cc->cme_) = (struct rb_callable_method_entry_struct *)cme;
     *((vm_call_handler *)&cc->call_) = call;