YJIT: Fix reference update for Invariants::no_ep_escape_iseqs

yjit.c

@@ -1047,7 +1047,6 @@ rb_yjit_multi_ractor_p(void)
 void
 rb_assert_iseq_handle(VALUE handle)
 {
-    RUBY_ASSERT_ALWAYS(rb_objspace_markable_object_p(handle));
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_iseq));
 }
 
@@ -1164,20 +1163,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() } {
@@ -360,7 +348,7 @@ pub extern "C" fn rb_yjit_constant_state_changed(id: ID) {
 /// Callback for marking GC objects inside [Invariants].
 /// See `struct yjijt_root_struct` in C.
 #[no_mangle]
-pub extern "C" fn rb_yjit_root_mark() {
+pub extern "C" fn rb_yjit_root_mark(_: *mut c_void) {
     // Call rb_gc_mark on exit location's raw_samples to
     // wrap frames in a GC allocated object. This needs to be called
     // at the same time as root mark.
@@ -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.