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

Author: k0kubun

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 @@ def self.enable(stats: false, log: false, mem_size: nil, call_threshold: nil)
     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

@@ -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/src/yjit.rs

@@ -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