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.

2 files changed, 17 insertions, 15 deletions

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 0302ecd536..cd6e649aa0 100644
--- a/yjit/src/core.rs
+++ b/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.
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index a9432f8745..4814e7a1c7 100644
--- a/yjit/src/invariants.rs
+++ b/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.