ZJIT: Forget about dead ISEQs in Invariants

Without this, we crash during reference update.

Author: XrXr

iseq.c

@@ -174,6 +174,9 @@ rb_iseq_free(const rb_iseq_t *iseq)
             RUBY_ASSERT(rb_yjit_live_iseq_count > 0);
             rb_yjit_live_iseq_count--;
         }
+#endif
+#if USE_ZJIT
+        rb_zjit_iseq_free(iseq);
 #endif
         ruby_xfree((void *)body->iseq_encoded);
         ruby_xfree((void *)body->insns_info.body);

zjit.h

@@ -22,6 +22,7 @@ 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);
 void rb_zjit_iseq_update_references(void *payload);
+void rb_zjit_iseq_free(const rb_iseq_t *iseq);
 void rb_zjit_before_ractor_spawn(void);
 void rb_zjit_tracing_invalidate_all(void);
 #else

zjit/src/gc.rs

@@ -128,6 +128,17 @@ pub extern "C" fn rb_zjit_iseq_update_references(payload: *mut c_void) {
     with_time_stat(gc_time_ns, || iseq_update_references(payload));
 }
 
+/// GC callback for finalizing an ISEQ
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_iseq_free(iseq: IseqPtr) {
+    if !ZJITState::has_instance() {
+        return;
+    }
+    // TODO(Shopify/ruby#682): Free `IseqPayload`
+    let invariants = ZJITState::get_invariants();
+    invariants.forget_iseq(iseq);
+}
+
 /// 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

@@ -68,6 +68,17 @@ impl Invariants {
         self.update_no_ep_escape_iseq_patch_points();
     }
 
+    /// Forget an ISEQ when freeing it. We need to because a) if the address is reused, we'd be
+    /// tracking the wrong object b) dead VALUEs in the table can means we risk passing invalid
+    /// VALUEs to `rb_gc_location()`.
+    pub fn forget_iseq(&mut self, iseq: IseqPtr) {
+        // Why not patch the patch points? If the ISEQ is dead then the GC also proved that all
+        // generated code referencing the ISEQ are unreachable. We mark the ISEQs baked into
+        // generated code.
+        self.ep_escape_iseqs.remove(&iseq);
+        self.no_ep_escape_iseq_patch_points.remove(&iseq);
+    }
+
     /// 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)