Allow passing a Rust closure to rb_iseq_callback (#6575)

8 files changed, 33 insertions, 16 deletions

diff --git a/cont.c b/cont.c
index 1ce60811d2..1d88088461 100644
--- a/cont.c
+++ b/cont.c
@@ -1253,7 +1253,7 @@ jit_cont_free(struct rb_jit_cont *cont)
 
 // Call a given callback against all on-stack ISEQs.
 void
-rb_jit_cont_each_iseq(rb_iseq_callback callback)
+rb_jit_cont_each_iseq(rb_iseq_callback callback, void *data)
 {
     struct rb_jit_cont *cont;
     for (cont = first_jit_cont; cont != NULL; cont = cont->next) {
@@ -1264,7 +1264,7 @@ rb_jit_cont_each_iseq(rb_iseq_callback callback)
         for (cfp = RUBY_VM_END_CONTROL_FRAME(cont->ec) - 1; ; cfp = RUBY_VM_NEXT_CONTROL_FRAME(cfp)) {
             const rb_iseq_t *iseq;
             if (cfp->pc && (iseq = cfp->iseq) != NULL && imemo_type((VALUE)iseq) == imemo_iseq) {
-                callback(iseq);
+                callback(iseq, data);
             }
 
             if (cfp == cont->ec->cfp)
diff --git a/internal/cont.h b/internal/cont.h
index 0b669f0ad5..38fab4f8ac 100644
--- a/internal/cont.h
+++ b/internal/cont.h
@@ -19,7 +19,7 @@ struct rb_execution_context_struct; /* in vm_core.c */
 void rb_fiber_reset_root_local_storage(struct rb_thread_struct *);
 void ruby_register_rollback_func_for_ensure(VALUE (*ensure_func)(VALUE), VALUE (*rollback_func)(VALUE));
 void rb_fiber_init_jit_cont(struct rb_fiber_struct *fiber);
-void rb_jit_cont_each_iseq(rb_iseq_callback callback);
+void rb_jit_cont_each_iseq(rb_iseq_callback callback, void *data);
 void rb_jit_cont_finish(void);
 
 VALUE rb_fiberptr_self(struct rb_fiber_struct *fiber);
diff --git a/iseq.h b/iseq.h
index d771874c2f..e5ab4870e8 100644
--- a/iseq.h
+++ b/iseq.h
@@ -31,7 +31,7 @@ RUBY_EXTERN const int ruby_api_version[];
 typedef struct rb_iseq_struct rb_iseq_t;
 #define rb_iseq_t rb_iseq_t
 #endif
-typedef void (*rb_iseq_callback)(const rb_iseq_t *);
+typedef void (*rb_iseq_callback)(const rb_iseq_t *, void *);
 
 extern const ID rb_iseq_shared_exc_local_tbl[];
 
diff --git a/mjit.c b/mjit.c
index 22da5fd45f..075c30fb67 100644
--- a/mjit.c
+++ b/mjit.c
@@ -951,7 +951,7 @@ mjit_capture_cc_entries(const struct rb_iseq_constant_body *compiled_iseq, const
 
 // Set up field `used_code_p` for unit iseqs whose iseq on the stack of ec.
 static void
-mark_iseq_units(const rb_iseq_t *iseq)
+mark_iseq_units(const rb_iseq_t *iseq, void *data)
 {
     if (ISEQ_BODY(iseq)->jit_unit != NULL) {
         ISEQ_BODY(iseq)->jit_unit->used_code_p = true;
@@ -982,7 +982,7 @@ unload_units(void)
     }
     // All threads have a root_fiber which has a mjit_cont. Other normal fibers also
     // have a mjit_cont. Thus we can check ISeqs in use by scanning ec of mjit_conts.
-    rb_jit_cont_each_iseq(mark_iseq_units);
+    rb_jit_cont_each_iseq(mark_iseq_units, NULL);
     // TODO: check stale_units and unload unused ones! (note that the unit is not associated to ISeq anymore)
 
     // Unload units whose total_calls is smaller than any total_calls in unit_queue.
diff --git a/yjit.c b/yjit.c
index 2a523f9787..f6e64aad65 100644
--- a/yjit.c
+++ b/yjit.c
@@ -894,11 +894,17 @@ rb_assert_cme_handle(VALUE handle)
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_ment));
 }
 
+// Used for passing a callback and other data over rb_objspace_each_objects
+struct iseq_callback_data {
+    rb_iseq_callback callback;
+    void *data;
+};
+
 // Heap-walking callback for rb_yjit_for_each_iseq().
 static int
 for_each_iseq_i(void *vstart, void *vend, size_t stride, void *data)
 {
-    const rb_iseq_callback callback = (rb_iseq_callback)data;
+    const struct iseq_callback_data *callback_data = (struct iseq_callback_data *)data;
     VALUE v = (VALUE)vstart;
     for (; v != (VALUE)vend; v += stride) {
         void *ptr = asan_poisoned_object_p(v);
@@ -906,7 +912,7 @@ for_each_iseq_i(void *vstart, void *vend, size_t stride, void *data)
 
         if (rb_obj_is_iseq(v)) {
             rb_iseq_t *iseq = (rb_iseq_t *)v;
-            callback(iseq);
+            callback_data->callback(iseq, callback_data->data);
         }
 
         asan_poison_object_if(ptr, v);
@@ -917,9 +923,10 @@ for_each_iseq_i(void *vstart, void *vend, size_t stride, void *data)
 // Iterate through the whole GC heap and invoke a callback for each iseq.
 // Used for global code invalidation.
 void
-rb_yjit_for_each_iseq(rb_iseq_callback callback)
+rb_yjit_for_each_iseq(rb_iseq_callback callback, void *data)
 {
-    rb_objspace_each_objects(for_each_iseq_i, (void *)callback);
+    struct iseq_callback_data callback_data = { .callback = callback, .data = data };
+    rb_objspace_each_objects(for_each_iseq_i, (void *)&callback_data);
 }
 
 // For running write barriers from Rust. Required when we add a new edge in the
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 53cb31beb1..ea7eb56405 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -527,6 +527,16 @@ fn get_or_create_iseq_payload(iseq: IseqPtr) -> &'static mut IseqPayload {
     unsafe { payload_non_null.as_mut() }.unwrap()
 }
 
+/// Iterate over all existing ISEQs
+pub fn for_each_iseq<F: FnMut(IseqPtr)>(mut callback: F) {
+    unsafe extern "C" fn callback_wrapper(iseq: IseqPtr, data: *mut c_void) {
+        let callback: &mut &mut dyn FnMut(IseqPtr) -> bool = unsafe { std::mem::transmute(data) };
+        callback(iseq);
+    };
+    let mut data: &mut dyn FnMut(IseqPtr) = &mut callback;
+    unsafe { rb_yjit_for_each_iseq(Some(callback_wrapper), (&mut data) as *mut _ as *mut c_void) };
+}
+
 /// Free the per-iseq payload
 #[no_mangle]
 pub extern "C" fn rb_yjit_iseq_free(payload: *mut c_void) {
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index db124c9303..00bade6b70 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1244,7 +1244,9 @@ pub const YARVINSN_trace_putobject_INT2FIX_0_: ruby_vminsn_type = 200;
 pub const YARVINSN_trace_putobject_INT2FIX_1_: ruby_vminsn_type = 201;
 pub const VM_INSTRUCTION_SIZE: ruby_vminsn_type = 202;
 pub type ruby_vminsn_type = u32;
-pub type rb_iseq_callback = ::std::option::Option<unsafe extern "C" fn(arg1: *const rb_iseq_t)>;
+pub type rb_iseq_callback = ::std::option::Option<
+    unsafe extern "C" fn(arg1: *const rb_iseq_t, arg2: *mut ::std::os::raw::c_void),
+>;
 extern "C" {
     pub fn rb_vm_insn_addr2opcode(addr: *const ::std::os::raw::c_void) -> ::std::os::raw::c_int;
 }
@@ -1540,7 +1542,7 @@ extern "C" {
     pub fn rb_assert_cme_handle(handle: VALUE);
 }
 extern "C" {
-    pub fn rb_yjit_for_each_iseq(callback: rb_iseq_callback);
+    pub fn rb_yjit_for_each_iseq(callback: rb_iseq_callback, data: *mut ::std::os::raw::c_void);
 }
 extern "C" {
     pub fn rb_yjit_obj_written(
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index 07de3374c8..3ca57b4943 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -532,9 +532,7 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
     // Stop other ractors since we are going to patch machine code.
     with_vm_lock(src_loc!(), || {
         // Make it so all live block versions are no longer valid branch targets
-        unsafe { rb_yjit_for_each_iseq(Some(invalidate_all_blocks_for_tracing)) };
-
-        extern "C" fn invalidate_all_blocks_for_tracing(iseq: IseqPtr) {
+        for_each_iseq(|iseq| {
             if let Some(payload) = unsafe { get_iseq_payload(iseq) } {
                 // C comment:
                 //   Leaking the blocks for now since we might have situations where
@@ -554,7 +552,7 @@ pub extern "C" fn rb_yjit_tracing_invalidate_all() {
 
             // Reset output code entry point
             unsafe { rb_iseq_reset_jit_func(iseq) };
-        }
+        });
 
         let cb = CodegenGlobals::get_inline_cb();