[Bug #19896]

fix memory leak in vm_method

This introduces a unified reference_count to clarify who is referencing a method.
This also allows us to treat the refinement method as the def owner since it counts itself as a reference

Co-authored-by: Peter Zhu <peter@peterzhu.ca>

1 files changed, 43 insertions, 62 deletions

diff --git a/vm_method.c b/vm_method.c
index cf4cfd2e56..e2a3907c33 100644
--- a/vm_method.c
+++ b/vm_method.c
@@ -166,7 +166,6 @@ invalidate_negative_cache(ID mid)
     }
 }
 
-static rb_method_entry_t *rb_method_entry_alloc(ID called_id, VALUE owner, VALUE defined_class, const rb_method_definition_t *def);
 const rb_method_entry_t * rb_method_entry_clone(const rb_method_entry_t *src_me);
 static const rb_callable_method_entry_t *complemented_callable_method_entry(VALUE klass, ID id);
 static const rb_callable_method_entry_t *lookup_overloaded_cme(const rb_callable_method_entry_t *cme);
@@ -403,33 +402,25 @@ rb_add_method_optimized(VALUE klass, ID mid, enum method_optimized_type opt_type
 }
 
 static void
-rb_method_definition_release(rb_method_definition_t *def, int complemented)
+rb_method_definition_release(rb_method_definition_t *def)
 {
     if (def != NULL) {
-        const int alias_count = def->alias_count;
-        const int complemented_count = def->complemented_count;
-        VM_ASSERT(alias_count >= 0);
-        VM_ASSERT(complemented_count >= 0);
-
-        if (alias_count + complemented_count == 0) {
-            if (METHOD_DEBUG) fprintf(stderr, "-%p-%s:%d,%d (remove)\n", (void *)def,
-                                      rb_id2name(def->original_id), alias_count, complemented_count);
+        const int reference_count = def->reference_count;
+        def->reference_count--;
+
+        VM_ASSERT(reference_count >= 0);
+
+        if (def->reference_count == 0) {
+            if (METHOD_DEBUG) fprintf(stderr, "-%p-%s:%d (remove)\n", (void *)def,
+                                      rb_id2name(def->original_id), def->reference_count);
             if (def->type == VM_METHOD_TYPE_BMETHOD && def->body.bmethod.hooks) {
                 xfree(def->body.bmethod.hooks);
             }
             xfree(def);
         }
         else {
-            if (complemented) {
-                VM_ASSERT(def->complemented_count > 0);
-                def->complemented_count--;
-            }
-            else if (def->alias_count > 0) {
-                def->alias_count--;
-            }
-
-            if (METHOD_DEBUG) fprintf(stderr, "-%p-%s:%d->%d,%d->%d (dec)\n", (void *)def, rb_id2name(def->original_id),
-                                      alias_count, def->alias_count, complemented_count, def->complemented_count);
+            if (METHOD_DEBUG) fprintf(stderr, "-%p-%s:%d->%d (dec)\n", (void *)def, rb_id2name(def->original_id),
+                                      reference_count, def->reference_count);
         }
     }
 }
@@ -442,7 +433,7 @@ rb_free_method_entry(const rb_method_entry_t *me)
     if (me->def && me->def->iseq_overload) {
         delete_overloaded_cme((const rb_callable_method_entry_t *)me);
     }
-    rb_method_definition_release(me->def, METHOD_ENTRY_COMPLEMENTED(me));
+    rb_method_definition_release(me->def);
 }
 
 static inline rb_method_entry_t *search_method(VALUE klass, ID id, VALUE *defined_class_ptr);
@@ -509,10 +500,20 @@ setup_method_cfunc_struct(rb_method_cfunc_t *cfunc, VALUE (*func)(ANYARGS), int
     cfunc->invoker = call_cfunc_invoker_func(argc);
 }
 
+static rb_method_definition_t *
+method_definition_addref(rb_method_definition_t *def, bool complemented)
+{
+    if (!complemented &&  def->reference_count > 0) def->aliased = true;
+    def->reference_count++;
+    if (METHOD_DEBUG) fprintf(stderr, "+%p-%s:%d\n", (void *)def, rb_id2name(def->original_id), def->reference_count);
+    return def;
+}
+
 void
 rb_method_definition_set(const rb_method_entry_t *me, rb_method_definition_t *def, void *opts)
 {
-    *(rb_method_definition_t **)&me->def = def;
+    rb_method_definition_release(me->def);
+    *(rb_method_definition_t **)&me->def = method_definition_addref(def, METHOD_ENTRY_COMPLEMENTED(me));
 
     if (opts != NULL) {
         switch (def->type) {
@@ -642,25 +643,10 @@ rb_method_definition_create(rb_method_type_t type, ID mid)
     return def;
 }
 
-static rb_method_definition_t *
-method_definition_addref(rb_method_definition_t *def)
-{
-    def->alias_count++;
-    if (METHOD_DEBUG) fprintf(stderr, "+%p-%s:%d\n", (void *)def, rb_id2name(def->original_id), def->alias_count);
-    return def;
-}
-
-static rb_method_definition_t *
-method_definition_addref_complement(rb_method_definition_t *def)
-{
-    def->complemented_count++;
-    if (METHOD_DEBUG) fprintf(stderr, "+%p-%s:%d\n", (void *)def, rb_id2name(def->original_id), def->complemented_count);
-    return def;
-}
-
 static rb_method_entry_t *
-rb_method_entry_alloc(ID called_id, VALUE owner, VALUE defined_class, const rb_method_definition_t *def)
+rb_method_entry_alloc(ID called_id, VALUE owner, VALUE defined_class, rb_method_definition_t *def, bool complement)
 {
+    if (def) method_definition_addref(def, complement);
     rb_method_entry_t *me = (rb_method_entry_t *)rb_imemo_new(imemo_ment, (VALUE)def, (VALUE)called_id, owner, defined_class);
     return me;
 }
@@ -682,9 +668,9 @@ filter_defined_class(VALUE klass)
 }
 
 rb_method_entry_t *
-rb_method_entry_create(ID called_id, VALUE klass, rb_method_visibility_t visi, const rb_method_definition_t *def)
+rb_method_entry_create(ID called_id, VALUE klass, rb_method_visibility_t visi, rb_method_definition_t *def)
 {
-    rb_method_entry_t *me = rb_method_entry_alloc(called_id, klass, filter_defined_class(klass), def);
+    rb_method_entry_t *me = rb_method_entry_alloc(called_id, klass, filter_defined_class(klass), def, false);
     METHOD_ENTRY_FLAGS_SET(me, visi, ruby_running ? FALSE : TRUE);
     if (def != NULL) method_definition_reset(me);
     return me;
@@ -693,13 +679,7 @@ rb_method_entry_create(ID called_id, VALUE klass, rb_method_visibility_t visi, c
 const rb_method_entry_t *
 rb_method_entry_clone(const rb_method_entry_t *src_me)
 {
-    rb_method_entry_t *me = rb_method_entry_alloc(src_me->called_id, src_me->owner, src_me->defined_class, src_me->def);
-    if (METHOD_ENTRY_COMPLEMENTED(src_me)) {
-        method_definition_addref_complement(src_me->def);
-    }
-    else {
-        method_definition_addref(src_me->def);
-    }
+    rb_method_entry_t *me = rb_method_entry_alloc(src_me->called_id, src_me->owner, src_me->defined_class, src_me->def, METHOD_ENTRY_COMPLEMENTED(src_me));
 
     METHOD_ENTRY_FLAGS_COPY(me, src_me);
     return me;
@@ -725,10 +705,8 @@ rb_method_entry_complement_defined_class(const rb_method_entry_t *src_me, ID cal
         refined.owner = orig_me->owner;
         def = NULL;
     }
-    else {
-        method_definition_addref_complement(def);
-    }
-    me = rb_method_entry_alloc(called_id, src_me->owner, defined_class, def);
+
+    me = rb_method_entry_alloc(called_id, src_me->owner, defined_class, def, true);
     METHOD_ENTRY_FLAGS_COPY(me, src_me);
     METHOD_ENTRY_COMPLEMENTED_SET(me);
     if (!def) {
@@ -744,7 +722,8 @@ rb_method_entry_complement_defined_class(const rb_method_entry_t *src_me, ID cal
 void
 rb_method_entry_copy(rb_method_entry_t *dst, const rb_method_entry_t *src)
 {
-    *(rb_method_definition_t **)&dst->def = method_definition_addref(src->def);
+    rb_method_definition_release(dst->def);
+    *(rb_method_definition_t **)&dst->def = method_definition_addref(src->def, METHOD_ENTRY_COMPLEMENTED(src));
     method_definition_reset(dst);
     dst->called_id = src->called_id;
     RB_OBJ_WRITE((VALUE)dst, &dst->owner, src->owner);
@@ -771,7 +750,8 @@ make_method_entry_refined(VALUE owner, rb_method_entry_t *me)
             rb_method_entry_alloc(me->called_id, me->owner,
                                   me->defined_class ?
                                   me->defined_class : owner,
-                                  method_definition_addref(me->def));
+                                  me->def,
+                                  true);
         METHOD_ENTRY_FLAGS_COPY(refined.orig_me, me);
         refined.owner = owner;
 
@@ -899,7 +879,7 @@ rb_method_entry_make(VALUE klass, ID mid, VALUE defined_class, rb_method_visibil
 
         if (RTEST(ruby_verbose) &&
             type != VM_METHOD_TYPE_UNDEF &&
-            (old_def->alias_count == 0) &&
+            (old_def->aliased == false) &&
             (!old_def->no_redef_warning) &&
             !make_refined &&
             old_def->type != VM_METHOD_TYPE_UNDEF &&
@@ -929,7 +909,9 @@ rb_method_entry_make(VALUE klass, ID mid, VALUE defined_class, rb_method_visibil
 
     /* create method entry */
     me = rb_method_entry_create(mid, defined_class, visi, NULL);
-    if (def == NULL) def = rb_method_definition_create(type, original_id);
+    if (def == NULL) {
+        def = rb_method_definition_create(type, original_id);
+    }
     rb_method_definition_set(me, def, opts);
 
     rb_clear_method_cache(klass, mid);
@@ -967,7 +949,7 @@ rb_method_entry_make(VALUE klass, ID mid, VALUE defined_class, rb_method_visibil
     return me;
 }
 
-static rb_method_entry_t *rb_method_entry_alloc(ID called_id, VALUE owner, VALUE defined_class, const rb_method_definition_t *def);
+static rb_method_entry_t *rb_method_entry_alloc(ID called_id, VALUE owner, VALUE defined_class, rb_method_definition_t *def, bool refined);
 
 static st_table *
 overloaded_cme_table(void)
@@ -1057,7 +1039,8 @@ get_overloaded_cme(const rb_callable_method_entry_t *cme)
         rb_method_entry_t *me = rb_method_entry_alloc(cme->called_id,
                                                       cme->owner,
                                                       cme->defined_class,
-                                                      def);
+                                                      def,
+                                                      false);
 
         ASSERT_vm_locking();
         st_insert(overloaded_cme_table(), (st_data_t)cme, (st_data_t)me);
@@ -1136,9 +1119,7 @@ method_entry_set(VALUE klass, ID mid, const rb_method_entry_t *me,
     if (newme == me) {
         me->def->no_redef_warning = TRUE;
     }
-    else {
-        method_definition_addref(me->def);
-    }
+
     method_added(klass, mid);
     return newme;
 }
@@ -1352,7 +1333,7 @@ negative_cme(ID mid)
         cme = (rb_callable_method_entry_t *)cme_data;
     }
     else {
-        cme = (rb_callable_method_entry_t *)rb_method_entry_alloc(mid, Qnil, Qnil, NULL);
+        cme = (rb_callable_method_entry_t *)rb_method_entry_alloc(mid, Qnil, Qnil, NULL, false);
         rb_id_table_insert(vm->negative_cme_table, mid, (VALUE)cme);
     }