YJIT: Take care of GC references in ISEQ invariants

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

bootstraptest/test_yjit.rb

@@ -2317,6 +2317,19 @@ def foo(obj)
   foo(Foo)
 }
 
+# Test EP == BP invalidation with moving ISEQs
+assert_equal 'ok', %q{
+  def entry
+    ok = proc { :ok } # set #entry as an EP-escaping ISEQ
+    [nil].reverse_each do # avoid exiting the JIT frame on the constant
+      GC.compact # move #entry ISEQ
+    end
+    ok # should be read off of escaped EP
+  end
+
+  entry.call
+}
+
 # invokesuper edge case
 assert_equal '[:A, [:A, :B]]', %q{
   class B

iseq.c

@@ -167,7 +167,7 @@ rb_iseq_free(const rb_iseq_t *iseq)
         struct rb_iseq_constant_body *const body = ISEQ_BODY(iseq);
         rb_rjit_free_iseq(iseq); /* Notify RJIT */
 #if USE_YJIT
-        rb_yjit_iseq_free(body->yjit_payload);
+        rb_yjit_iseq_free(iseq);
         if (FL_TEST_RAW((VALUE)iseq, ISEQ_TRANSLATED)) {
             RUBY_ASSERT(rb_yjit_live_iseq_count > 0);
             rb_yjit_live_iseq_count--;
@@ -377,7 +377,7 @@ rb_iseq_mark_and_move(rb_iseq_t *iseq, bool reference_updating)
             rb_rjit_iseq_update_references(body);
 #endif
 #if USE_YJIT
-            rb_yjit_iseq_update_references(body->yjit_payload);
+            rb_yjit_iseq_update_references(iseq);
 #endif
         }
         else {

yjit.h

@@ -40,8 +40,8 @@ void rb_yjit_init(bool yjit_enabled);
 void rb_yjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop);
 void rb_yjit_constant_state_changed(ID id);
 void rb_yjit_iseq_mark(void *payload);
-void rb_yjit_iseq_update_references(void *payload);
-void rb_yjit_iseq_free(void *payload);
+void rb_yjit_iseq_update_references(const rb_iseq_t *iseq);
+void rb_yjit_iseq_free(const rb_iseq_t *iseq);
 void rb_yjit_before_ractor_spawn(void);
 void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx);
 void rb_yjit_tracing_invalidate_all(void);
@@ -65,8 +65,8 @@ static inline void rb_yjit_init(bool yjit_enabled) {}
 static inline void rb_yjit_bop_redefined(int redefined_flag, enum ruby_basic_operators bop) {}
 static inline void rb_yjit_constant_state_changed(ID id) {}
 static inline void rb_yjit_iseq_mark(void *payload) {}
-static inline void rb_yjit_iseq_update_references(void *payload) {}
-static inline void rb_yjit_iseq_free(void *payload) {}
+static inline void rb_yjit_iseq_update_references(const rb_iseq_t *iseq) {}
+static inline void rb_yjit_iseq_free(const rb_iseq_t *iseq) {}
 static inline void rb_yjit_before_ractor_spawn(void) {}
 static inline void rb_yjit_constant_ic_update(const rb_iseq_t *const iseq, IC ic, unsigned insn_idx) {}
 static inline void rb_yjit_tracing_invalidate_all(void) {}

yjit/src/core.rs

@@ -1138,8 +1138,12 @@ pub fn for_each_off_stack_iseq_payload<F: FnMut(&mut IseqPayload)>(mut callback:
 
 /// Free the per-iseq payload
 #[no_mangle]
-pub extern "C" fn rb_yjit_iseq_free(payload: *mut c_void) {
+pub extern "C" fn rb_yjit_iseq_free(iseq: IseqPtr) {
+    // Free invariants for the ISEQ
+    iseq_free_invariants(iseq);
+
     let payload = {
+        let payload = unsafe { rb_iseq_get_yjit_payload(iseq) };
         if payload.is_null() {
             // Nothing to free.
             return;
@@ -1266,7 +1270,11 @@ pub extern "C" fn rb_yjit_iseq_mark(payload: *mut c_void) {
 /// GC callback for updating GC objects in the per-iseq payload.
 /// This is a mirror of [rb_yjit_iseq_mark].
 #[no_mangle]
-pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
+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.
         return;

yjit/src/invariants.rs

@@ -177,6 +177,22 @@ 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) {
+    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) {
+    Invariants::get_instance()
+        .no_ep_escape_iseqs
+        .remove(&iseq);
+}
+
 // Checks rb_method_basic_definition_p and registers the current block for invalidation if method
 // lookup changes.
 // A "basic method" is one defined during VM boot, so we can use this to check assumptions based on