ZJIT: Remove dead CMEs from Invariants

Author: XrXr

vm_method.c

@@ -859,6 +859,17 @@ rb_free_method_entry_vm_weak_references(const rb_method_entry_t *me)
 void
 rb_free_method_entry(const rb_method_entry_t *me)
 {
+#if USE_ZJIT
+    if (METHOD_ENTRY_CACHED(me)) {
+        rb_zjit_cme_free((const rb_callable_method_entry_t *)me);
+    }
+#endif
+
+#if USE_YJIT
+    // YJIT rb_yjit_root_mark() roots CMEs in `Invariants`,
+    // to remove from `Invariants` here.
+#endif
+
     rb_method_definition_release(me->def);
 }
 

zjit.h

@@ -18,6 +18,7 @@ void rb_zjit_profile_insn(uint32_t insn, rb_execution_context_t *ec);
 void rb_zjit_profile_enable(const rb_iseq_t *iseq);
 void rb_zjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
 void rb_zjit_cme_invalidate(const rb_callable_method_entry_t *cme);
+void rb_zjit_cme_free(const rb_callable_method_entry_t *cme);
 void rb_zjit_invalidate_no_ep_escape(const rb_iseq_t *iseq);
 void rb_zjit_constant_state_changed(ID id);
 void rb_zjit_iseq_mark(void *payload);

zjit/src/gc.rs

@@ -139,6 +139,16 @@ pub extern "C" fn rb_zjit_iseq_free(iseq: IseqPtr) {
     invariants.forget_iseq(iseq);
 }
 
+/// GC callback for finalizing a CME
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_cme_free(cme: *const rb_callable_method_entry_struct) {
+    if !ZJITState::has_instance() {
+        return;
+    }
+    let invariants = ZJITState::get_invariants();
+    invariants.forget_cme(cme);
+}
+
 /// GC callback for updating object references after all object moves
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_root_update_references() {

zjit/src/invariants.rs

@@ -80,6 +80,11 @@ impl Invariants {
         self.no_ep_escape_iseq_patch_points.remove(&iseq);
     }
 
+    /// Forget a CME when freeing it. See [Self::forget_iseq] for reasoning.
+    pub fn forget_cme(&mut self, cme: *const rb_callable_method_entry_t) {
+        self.cme_patch_points.remove(&cme);
+    }
+
     /// Update ISEQ references in Invariants::ep_escape_iseqs
     fn update_ep_escape_iseqs(&mut self) {
         let updated = std::mem::take(&mut self.ep_escape_iseqs)