ZJIT: Forget about dead ISEQs in `Invariants`

Without this, we crash during reference update.

4 files changed, 26 insertions, 0 deletions

diff --git a/iseq.c b/iseq.c
index d8891353f1..082ea28c4a 100644
--- a/iseq.c
+++ b/iseq.c
@@ -175,6 +175,9 @@ rb_iseq_free(const rb_iseq_t *iseq)
             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);
         ruby_xfree((void *)body->insns_info.positions);
diff --git a/zjit.h b/zjit.h
index 9735cab6d4..24da476fd9 100644
--- a/zjit.h
+++ b/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
diff --git a/zjit/src/gc.rs b/zjit/src/gc.rs
index bf07f1f340..f53536e261 100644
--- a/zjit/src/gc.rs
+++ b/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() {
diff --git a/zjit/src/invariants.rs b/zjit/src/invariants.rs
index 2e67c33a6d..cea97bc8ee 100644
--- a/zjit/src/invariants.rs
+++ b/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)