YJIT: Fix reference update for `Invariants::no_ep_escape_iseqs`

Previously, the update was done in the ISEQ callback. That effectively
never updated anything because the callback itself is given an intact
reference, so it could update its content, and `rb_gc_location(iseq)`
never returned a new address. Update the whole table once in the YJIT
root instead.

yjit.c

@@ -1164,20 +1164,14 @@ yjit_root_memsize(const void *ptr)
     return 0; // TODO: more accurate accounting
 }
 
-// GC callback during compaction
-static void
-yjit_root_update_references(void *ptr)
-{
-    // Do nothing since we use rb_gc_mark(), which pins.
-}
-
 void rb_yjit_root_mark(void *ptr); // in Rust
+void rb_yjit_root_update_references(void *ptr); // in Rust
 
 // Custom type for interacting with the GC
 // TODO: make this write barrier protected
 static const rb_data_type_t yjit_root_type = {
     "yjit_root",
-    {rb_yjit_root_mark, yjit_root_free, yjit_root_memsize, yjit_root_update_references},
+    {rb_yjit_root_mark, yjit_root_free, yjit_root_memsize, rb_yjit_root_update_references},
     0, 0, RUBY_TYPED_FREE_IMMEDIATELY
 };
 

yjit/src/core.rs

@@ -1271,9 +1271,6 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
 /// This is a mirror of [rb_yjit_iseq_mark].
 #[no_mangle]
 pub extern "C" fn rb_yjit_iseq_update_references(iseq: IseqPtr) {
-    // Update ISEQ references in invariants
-    iseq_update_references_in_invariants(iseq);
-
     let payload = unsafe { rb_iseq_get_yjit_payload(iseq) };
     let payload = if payload.is_null() {
         // Nothing to update.

yjit/src/invariants.rs

@@ -177,18 +177,6 @@ pub fn iseq_escapes_ep(iseq: IseqPtr) -> bool {
         .map_or(false, |blocks| blocks.is_empty())
 }
 
-/// Update ISEQ references in invariants on GC compaction
-pub fn iseq_update_references_in_invariants(iseq: IseqPtr) {
-    if unsafe { INVARIANTS.is_none() } {
-        return;
-    }
-    let no_ep_escape_iseqs = &mut Invariants::get_instance().no_ep_escape_iseqs;
-    if let Some(blocks) = no_ep_escape_iseqs.remove(&iseq) {
-        let new_iseq = unsafe { rb_gc_location(iseq.into()) }.as_iseq();
-        no_ep_escape_iseqs.insert(new_iseq, blocks);
-    }
-}
-
 /// Forget an ISEQ remembered in invariants
 pub fn iseq_free_invariants(iseq: IseqPtr) {
     if unsafe { INVARIANTS.is_none() } {
@@ -388,6 +376,23 @@ pub extern "C" fn rb_yjit_root_mark() {
     }
 }
 
+#[no_mangle]
+pub extern "C" fn rb_yjit_root_update_references(_: *mut c_void) {
+    if unsafe { INVARIANTS.is_none() } {
+        return;
+    }
+    let no_ep_escape_iseqs = &mut Invariants::get_instance().no_ep_escape_iseqs;
+
+    // Make a copy of the table with updated ISEQ keys
+    let mut updated_copy = HashMap::with_capacity(no_ep_escape_iseqs.len());
+    for (iseq, blocks) in mem::take(no_ep_escape_iseqs) {
+        let new_iseq = unsafe { rb_gc_location(iseq.into()) }.as_iseq();
+        updated_copy.insert(new_iseq, blocks);
+    }
+
+    *no_ep_escape_iseqs = updated_copy;
+}
+
 /// Remove all invariant assumptions made by the block by removing the block as
 /// as a key in all of the relevant tables.
 /// For safety, the block has to be initialized and the vm lock must be held.