YJIT: Support calling bmethods in Ractors

Co-authored-by: Luke Gruber <luke.gru@gmail.com>

vm_core.h

@@ -2074,8 +2074,9 @@ rb_ec_ractor_ptr(const rb_execution_context_t *ec)
 static inline rb_serial_t
 rb_ec_ractor_id(const rb_execution_context_t *ec)
 {
-    VM_ASSERT(ec->ractor_id == rb_ractor_id(rb_ec_ractor_ptr(ec)));
-    return ec->ractor_id;
+    rb_serial_t ractor_id = ec->ractor_id;
+    RUBY_ASSERT(ractor_id);
+    return ractor_id;
 }
 
 static inline rb_vm_t *

yjit.c

@@ -473,6 +473,12 @@ rb_yjit_invokeblock_sp_pops(const struct rb_callinfo *ci)
     return 1 - sp_inc_of_invokeblock(ci); // + 1 to ignore return value push
 }
 
+rb_serial_t
+rb_yjit_cme_ractor_serial(const rb_callable_method_entry_t *cme)
+{
+    return cme->def->body.bmethod.defined_ractor_id;
+}
+
 // Setup jit_return to avoid returning a non-Qundef value on a non-FINISH frame.
 // See [jit_compile_exception] for details.
 void

yjit/bindgen/src/main.rs

@@ -272,6 +272,7 @@ fn main() {
         .allowlist_function("rb_optimized_call")
         .allowlist_function("rb_yjit_sendish_sp_pops")
         .allowlist_function("rb_yjit_invokeblock_sp_pops")
+        .allowlist_function("rb_yjit_cme_ractor_serial")
         .allowlist_function("rb_yjit_set_exception_return")
         .allowlist_function("rb_jit_str_concat_codepoint")
         .allowlist_type("rstring_offsets")

yjit/src/codegen.rs

@@ -7396,11 +7396,12 @@ fn gen_send_bmethod(
     let capture = unsafe { proc_block.as_.captured.as_ref() };
     let iseq = unsafe { *capture.code.iseq.as_ref() };
 
-    // Optimize for single ractor mode and avoid runtime check for
-    // "defined with an un-shareable Proc in a different Ractor"
-    if !assume_single_ractor_mode(jit, asm) {
-        gen_counter_incr(jit, asm, Counter::send_bmethod_ractor);
-        return None;
+    if !procv.shareable_p() {
+        let ractor_serial = unsafe { rb_yjit_cme_ractor_serial(cme) };
+        asm_comment!(asm, "guard current ractor == {}", ractor_serial);
+        let current_ractor_serial = asm.load(Opnd::mem(64, EC, RUBY_OFFSET_EC_RACTOR_ID));
+        asm.cmp(current_ractor_serial, Opnd::UImm(ractor_serial));
+        asm.jne(Target::side_exit(Counter::send_bmethod_ractor));
     }
 
     // Passing a block to a block needs logic different from passing

yjit/src/cruby.rs

@@ -361,6 +361,11 @@ impl VALUE {
         !self.special_const_p()
     }
 
+    /// Shareability between ractors. `RB_OBJ_SHAREABLE_P()`.
+    pub fn shareable_p(self) -> bool {
+        (self.builtin_flags() & RUBY_FL_SHAREABLE as usize) != 0
+    }
+
     /// Return true if the value is a Ruby Fixnum (immediate-size integer)
     pub fn fixnum_p(self) -> bool {
         let VALUE(cval) = self;
@@ -772,6 +777,7 @@ mod manual_defs {
     pub const RUBY_OFFSET_EC_INTERRUPT_FLAG: i32 = 32; // rb_atomic_t (u32)
     pub const RUBY_OFFSET_EC_INTERRUPT_MASK: i32 = 36; // rb_atomic_t (u32)
     pub const RUBY_OFFSET_EC_THREAD_PTR: i32 = 48;
+    pub const RUBY_OFFSET_EC_RACTOR_ID: i32 = 64;
 
     // Constants from rb_thread_t in vm_core.h
     pub const RUBY_OFFSET_THREAD_SELF: i32 = 16;

yjit/src/cruby_bindings.inc.rs

@@ -1198,6 +1198,7 @@ extern "C" {
     pub fn rb_yjit_shape_index(shape_id: shape_id_t) -> attr_index_t;
     pub fn rb_yjit_sendish_sp_pops(ci: *const rb_callinfo) -> usize;
     pub fn rb_yjit_invokeblock_sp_pops(ci: *const rb_callinfo) -> usize;
+    pub fn rb_yjit_cme_ractor_serial(cme: *const rb_callable_method_entry_t) -> rb_serial_t;
     pub fn rb_yjit_set_exception_return(
         cfp: *mut rb_control_frame_t,
         leave_exit: *mut ::std::os::raw::c_void,