YJIT: Call YJIT hooks before enabling YJIT (#14032)

inits.c

@@ -90,7 +90,6 @@ rb_call_builtin_inits(void)
 #define BUILTIN(n) CALL(builtin_##n)
     BUILTIN(kernel);
     BUILTIN(yjit);
-    // BUILTIN(yjit_hook) is called after rb_yjit_init()
     BUILTIN(gc);
     BUILTIN(ractor);
     BUILTIN(numeric);
@@ -109,6 +108,7 @@ rb_call_builtin_inits(void)
     BUILTIN(nilclass);
     BUILTIN(marshal);
     BUILTIN(zjit);
+    BUILTIN(yjit_hook);
     Init_builtin_prelude();
 }
 #undef CALL

ruby.c

@@ -1833,12 +1833,6 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
     }
 #endif
 
-#if USE_YJIT
-    // Call yjit_hook.rb after rb_yjit_init() to use `RubyVM::YJIT.enabled?`
-    void Init_builtin_yjit_hook();
-    Init_builtin_yjit_hook();
-#endif
-
     rb_namespace_init_done();
     ruby_init_prelude();
     ruby_set_script_name(opt->script_name);

yjit.rb

@@ -64,7 +64,6 @@ module RubyVM::YJIT
     end
 
     at_exit { print_and_dump_stats } if stats
-    call_yjit_hooks
     Primitive.rb_yjit_enable(stats, stats != :quiet, log, log != :quiet, mem_size, call_threshold)
   end
 

yjit/bindgen/src/main.rs

@@ -105,6 +105,9 @@ fn main() {
         .allowlist_var("SHAPE_ID_NUM_BITS")
         .allowlist_var("SHAPE_ID_HAS_IVAR_MASK")
 
+        // From ruby/internal/eval.h
+        .allowlist_function("rb_funcall")
+
         // From ruby/internal/intern/object.h
         .allowlist_function("rb_obj_is_kind_of")
         .allowlist_function("rb_obj_frozen_p")
@@ -269,6 +272,7 @@ fn main() {
         .allowlist_function("rb_float_new")
 
         // From vm_core.h
+        .allowlist_var("rb_cRubyVM")
         .allowlist_var("rb_mRubyVMFrozenCore")
         .allowlist_var("VM_BLOCK_HANDLER_NONE")
         .allowlist_type("vm_frame_env_flags")
@@ -383,6 +387,7 @@ fn main() {
         .allowlist_function("rb_ivar_defined")
         .allowlist_function("rb_ivar_get")
         .allowlist_function("rb_mod_name")
+        .allowlist_function("rb_const_get")
 
         // From internal/vm.h
         .allowlist_var("rb_vm_insns_count")

yjit/src/cruby.rs

@@ -601,9 +601,15 @@ pub fn rust_str_to_ruby(str: &str) -> VALUE {
 
 /// Produce a Ruby symbol from a Rust string slice
 pub fn rust_str_to_sym(str: &str) -> VALUE {
+    let id = rust_str_to_id(str);
+    unsafe { rb_id2sym(id) }
+}
+
+/// Produce an ID from a Rust string slice
+pub fn rust_str_to_id(str: &str) -> ID {
     let c_str = CString::new(str).unwrap();
     let c_ptr: *const c_char = c_str.as_ptr();
-    unsafe { rb_id2sym(rb_intern(c_ptr)) }
+    unsafe { rb_intern(c_ptr) }
 }
 
 /// Produce an owned Rust String from a C char pointer

yjit/src/cruby_bindings.inc.rs

@@ -1019,6 +1019,7 @@ extern "C" {
     pub fn rb_gc_location(obj: VALUE) -> VALUE;
     pub fn rb_gc_writebarrier(old: VALUE, young: VALUE);
     pub fn rb_class_get_superclass(klass: VALUE) -> VALUE;
+    pub fn rb_funcall(recv: VALUE, mid: ID, n: ::std::os::raw::c_int, ...) -> VALUE;
     pub static mut rb_mKernel: VALUE;
     pub static mut rb_cBasicObject: VALUE;
     pub static mut rb_cArray: VALUE;
@@ -1080,6 +1081,7 @@ extern "C" {
     pub fn rb_ivar_get(obj: VALUE, name: ID) -> VALUE;
     pub fn rb_ivar_defined(obj: VALUE, name: ID) -> VALUE;
     pub fn rb_attr_get(obj: VALUE, name: ID) -> VALUE;
+    pub fn rb_const_get(space: VALUE, name: ID) -> VALUE;
     pub fn rb_obj_info_dump(obj: VALUE);
     pub fn rb_class_allocate_instance(klass: VALUE) -> VALUE;
     pub fn rb_obj_equal(obj1: VALUE, obj2: VALUE) -> VALUE;
@@ -1102,6 +1104,7 @@ extern "C" {
         klass: VALUE,
         id: ID,
     ) -> *const rb_callable_method_entry_t;
+    pub static mut rb_cRubyVM: VALUE;
     pub static mut rb_mRubyVMFrozenCore: VALUE;
     pub static mut rb_block_param_proxy: VALUE;
     pub fn rb_vm_ep_local_ep(ep: *const VALUE) -> *const VALUE;

yjit/src/yjit.rs

@@ -54,6 +54,12 @@ fn yjit_init() {
     // TODO: need to make sure that command-line options have been
     // initialized by CRuby
 
+    // Call YJIT hooks before enabling YJIT to avoid compiling the hooks themselves
+    unsafe {
+        let yjit = rb_const_get(rb_cRubyVM, rust_str_to_id("YJIT"));
+        rb_funcall(yjit, rust_str_to_id("call_yjit_hooks"), 0);
+    }
+
     // Catch panics to avoid UB for unwinding into C frames.
     // See https://doc.rust-lang.org/nomicon/exception-safety.html
     let result = std::panic::catch_unwind(|| {

yjit_hook.rb

@@ -1,8 +1,3 @@
-# If YJIT is enabled, load the YJIT-only version of builtin methods
-if defined?(RubyVM::YJIT) && RubyVM::YJIT.enabled?
-  RubyVM::YJIT.send(:call_yjit_hooks)
-end
-
 # Remove the helper defined in kernel.rb
 class Module
   undef :with_yjit