ZJIT: Prepare for sharing JIT hooks with ZJIT (#14044)

k0kubun · web-flow · commit 2cd10de33097 · 2025-07-30T10:11:10.000-07:00
diff --git a/array.rb b/array.rb
@@ -212,7 +212,7 @@ def fetch_values(*indexes, &block)
     indexes
   end
 
-  with_yjit do
+  with_jit do
     if Primitive.rb_builtin_basic_definition_p(:each)
       undef :each
 
diff --git a/common.mk b/common.mk
@@ -1236,8 +1236,9 @@ BUILTIN_RB_SRCS = \
 		$(srcdir)/nilclass.rb \
 		$(srcdir)/prelude.rb \
 		$(srcdir)/gem_prelude.rb \
+		$(srcdir)/jit_hook.rb \
+		$(srcdir)/jit_undef.rb \
 		$(srcdir)/yjit.rb \
-		$(srcdir)/yjit_hook.rb \
 		$(srcdir)/zjit.rb \
 		$(empty)
 BUILTIN_RB_INCS = $(BUILTIN_RB_SRCS:.rb=.rbinc)
diff --git a/depend b/depend
@@ -9196,6 +9196,8 @@ miniinit.$(OBJEXT): {$(VPATH)}internal/warning_push.h
 miniinit.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
 miniinit.$(OBJEXT): {$(VPATH)}io.rb
 miniinit.$(OBJEXT): {$(VPATH)}iseq.h
+miniinit.$(OBJEXT): {$(VPATH)}jit_hook.rb
+miniinit.$(OBJEXT): {$(VPATH)}jit_undef.rb
 miniinit.$(OBJEXT): {$(VPATH)}kernel.rb
 miniinit.$(OBJEXT): {$(VPATH)}marshal.rb
 miniinit.$(OBJEXT): {$(VPATH)}method.h
@@ -9232,7 +9234,6 @@ miniinit.$(OBJEXT): {$(VPATH)}vm_core.h
 miniinit.$(OBJEXT): {$(VPATH)}vm_opts.h
 miniinit.$(OBJEXT): {$(VPATH)}warning.rb
 miniinit.$(OBJEXT): {$(VPATH)}yjit.rb
-miniinit.$(OBJEXT): {$(VPATH)}yjit_hook.rb
 miniinit.$(OBJEXT): {$(VPATH)}zjit.rb
 namespace.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 namespace.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
@@ -18755,6 +18756,8 @@ vm.$(OBJEXT): {$(VPATH)}internal/variable.h
 vm.$(OBJEXT): {$(VPATH)}internal/warning_push.h
 vm.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
 vm.$(OBJEXT): {$(VPATH)}iseq.h
+vm.$(OBJEXT): {$(VPATH)}jit_hook.rbinc
+vm.$(OBJEXT): {$(VPATH)}jit_undef.rbinc
 vm.$(OBJEXT): {$(VPATH)}method.h
 vm.$(OBJEXT): {$(VPATH)}missing.h
 vm.$(OBJEXT): {$(VPATH)}node.h
@@ -18797,7 +18800,6 @@ vm.$(OBJEXT): {$(VPATH)}vm_opts.h
 vm.$(OBJEXT): {$(VPATH)}vm_sync.h
 vm.$(OBJEXT): {$(VPATH)}vmtc.inc
 vm.$(OBJEXT): {$(VPATH)}yjit.h
-vm.$(OBJEXT): {$(VPATH)}yjit_hook.rbinc
 vm.$(OBJEXT): {$(VPATH)}zjit.h
 vm_backtrace.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 vm_backtrace.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
diff --git a/inits.c b/inits.c
@@ -88,8 +88,10 @@ void
 rb_call_builtin_inits(void)
 {
 #define BUILTIN(n) CALL(builtin_##n)
-    BUILTIN(kernel);
+    BUILTIN(jit_hook);
     BUILTIN(yjit);
+    BUILTIN(zjit);
+    BUILTIN(kernel);
     BUILTIN(gc);
     BUILTIN(ractor);
     BUILTIN(numeric);
@@ -107,8 +109,7 @@ rb_call_builtin_inits(void)
     BUILTIN(thread_sync);
     BUILTIN(nilclass);
     BUILTIN(marshal);
-    BUILTIN(zjit);
-    BUILTIN(yjit_hook);
+    BUILTIN(jit_undef);
     Init_builtin_prelude();
 }
 #undef CALL
diff --git a/internal/cmdlineopt.h b/internal/cmdlineopt.h
@@ -23,9 +23,6 @@ typedef struct ruby_cmdline_options {
     ruby_features_t warn;
     unsigned int dump;
     long backtrace_length_limit;
-#if USE_ZJIT
-    void *zjit;
-#endif
 
     const char *crash_report;
 
@@ -42,6 +39,9 @@ typedef struct ruby_cmdline_options {
 #if USE_YJIT
     unsigned int yjit: 1;
 #endif
+#if USE_ZJIT
+    unsigned int zjit: 1;
+#endif
 } ruby_cmdline_options_t;
 
 struct ruby_opt_message {
diff --git a/jit_hook.rb b/jit_hook.rb
@@ -0,0 +1,13 @@
+class Module
+  # Internal helper for built-in initializations to define methods only when JIT is enabled.
+  # This method is removed in jit_undef.rb.
+  private def with_jit(&block) # :nodoc:
+    # ZJIT currently doesn't compile Array#each properly, so it's disabled for now.
+    if defined?(RubyVM::ZJIT) && Primitive.rb_zjit_option_enabled_p && false # TODO: remove `&& false` (Shopify/ruby#667)
+      # We don't support lazily enabling ZJIT yet, so we can call the block right away.
+      block.call
+    elsif defined?(RubyVM::YJIT)
+      RubyVM::YJIT.send(:add_jit_hook, block)
+    end
+  end
+end
diff --git a/jit_undef.rb b/jit_undef.rb
@@ -0,0 +1,4 @@
+# Remove the helper defined in jit_hook.rb
+class Module
+  undef :with_jit
+end
diff --git a/kernel.rb b/kernel.rb
@@ -291,13 +291,3 @@ def Integer(arg, base = 0, exception: true)
     end
   end
 end
-
-class Module
-  # Internal helper for built-in initializations to define methods only when YJIT is enabled.
-  # This method is removed in yjit_hook.rb.
-  private def with_yjit(&block) # :nodoc:
-    if defined?(RubyVM::YJIT)
-      RubyVM::YJIT.send(:add_yjit_hook, block)
-    end
-  end
-end
diff --git a/numeric.rb b/numeric.rb
@@ -322,7 +322,7 @@ def denominator
     1
   end
 
-  with_yjit do
+  with_jit do
     if Primitive.rb_builtin_basic_definition_p(:downto)
       undef :downto
 
diff --git a/ruby.c b/ruby.c
@@ -1196,14 +1196,12 @@ setup_yjit_options(const char *s)
 
 #if USE_ZJIT
 static void
-setup_zjit_options(ruby_cmdline_options_t *opt, const char *s)
+setup_zjit_options(const char *s)
 {
     // The option parsing is done in zjit/src/options.rs
-    extern void *rb_zjit_init_options(void);
-    extern bool rb_zjit_parse_option(void *options, const char *s);
+    extern bool rb_zjit_parse_option(const char *s);
 
-    if (!opt->zjit) opt->zjit = rb_zjit_init_options();
-    if (!rb_zjit_parse_option(opt->zjit, s)) {
+    if (!rb_zjit_parse_option(s)) {
         rb_raise(rb_eRuntimeError, "invalid ZJIT option '%s' (--help will show valid zjit options)", s);
     }
 }
@@ -1481,7 +1479,7 @@ proc_long_options(ruby_cmdline_options_t *opt, const char *s, long argc, char **
     else if (is_option_with_optarg("zjit", '-', true, false, false)) {
 #if USE_ZJIT
         FEATURE_SET(opt->features, FEATURE_BIT(zjit));
-        setup_zjit_options(opt, s);
+        setup_zjit_options(s);
 #else
         rb_warn("Ruby was built without ZJIT support."
                 " You may need to install rustc to build Ruby with ZJIT.");
@@ -1828,8 +1826,8 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
 #endif
 #if USE_ZJIT
     if (opt->zjit) {
-        extern void rb_zjit_init(void *options);
-        rb_zjit_init(opt->zjit);
+        extern void rb_zjit_init(void);
+        rb_zjit_init();
     }
 #endif
 
@@ -2370,8 +2368,9 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt)
 #endif
 #if USE_ZJIT
     if (FEATURE_SET_P(opt->features, zjit) && !opt->zjit) {
-        extern void *rb_zjit_init_options(void);
-        opt->zjit = rb_zjit_init_options();
+        extern void rb_zjit_prepare_options(void);
+        rb_zjit_prepare_options();
+        opt->zjit = true;
     }
 #endif
 
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -1084,6 +1084,13 @@ def test_require_rubygems_with_auto_compact
     }, call_threshold: 2
   end
 
+  def test_zjit_option_uses_array_each_in_ruby
+    omit 'ZJIT wrongly compiles Array#each, so it is disabled for now'
+    assert_runs '"<internal:array>"', %q{
+      Array.instance_method(:each).source_location&.first
+    }
+  end
+
   def test_profile_under_nested_jit_call
     assert_compiles '[nil, nil, 3]', %q{
       def profile
diff --git a/vm.c b/vm.c
@@ -4509,14 +4509,21 @@ Init_vm_objects(void)
     vm->cc_refinement_table = rb_set_init_numtable();
 }
 
+#if USE_ZJIT
+extern VALUE rb_zjit_option_enabled_p(rb_execution_context_t *ec, VALUE self);
+#else
+static VALUE rb_zjit_option_enabled_p(rb_execution_context_t *ec, VALUE self) { return Qfalse; }
+#endif
+
+// Whether JIT is enabled or not, we need to load/undef `#with_jit` for other builtins.
+#include "jit_hook.rbinc"
+#include "jit_undef.rbinc"
+
 // Stub for builtin function when not building YJIT units
 #if !USE_YJIT
 void Init_builtin_yjit(void) {}
 #endif
 
-// Whether YJIT is enabled or not, we load yjit_hook.rb to remove Kernel#with_yjit.
-#include "yjit_hook.rbinc"
-
 // Stub for builtin function when not building ZJIT units
 #if !USE_ZJIT
 void Init_builtin_zjit(void) {}
diff --git a/vm_method.c b/vm_method.c
@@ -775,7 +775,7 @@ rb_method_definition_set(const rb_method_entry_t *me, rb_method_definition_t *de
                 /* setup iseq first (before invoking GC) */
                 RB_OBJ_WRITE(me, &def->body.iseq.iseqptr, iseq);
 
-                // Methods defined in `with_yjit` should be considered METHOD_ENTRY_BASIC
+                // Methods defined in `with_jit` should be considered METHOD_ENTRY_BASIC
                 if (rb_iseq_attr_p(iseq, BUILTIN_ATTR_C_TRACE)) {
                     METHOD_ENTRY_BASIC_SET((rb_method_entry_t *)me, TRUE);
                 }
diff --git a/yjit.rb b/yjit.rb
@@ -264,23 +264,23 @@ def self.simulate_oom! # :nodoc:
   end
 
   # Blocks that are called when YJIT is enabled
-  @yjit_hooks = []
+  @jit_hooks = []
 
   class << self
     # :stopdoc:
     private
 
     # Register a block to be called when YJIT is enabled
-    def add_yjit_hook(hook)
-      @yjit_hooks << hook
+    def add_jit_hook(hook)
+      @jit_hooks << hook
     end
 
-    # Run YJIT hooks registered by RubyVM::YJIT.with_yjit
-    def call_yjit_hooks
+    # Run YJIT hooks registered by `#with_jit`
+    def call_jit_hooks
       # Skip using builtin methods in Ruby if --yjit-c-builtin is given
       return if Primitive.yjit_c_builtin_p
-      @yjit_hooks.each(&:call)
-      @yjit_hooks.clear
+      @jit_hooks.each(&:call)
+      @jit_hooks.clear
     end
 
     # Print stats and dump exit locations
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
@@ -46,7 +46,7 @@ pub struct Options {
     // The number of registers allocated for stack temps
     pub num_temp_regs: usize,
 
-    // Disable Ruby builtin methods defined by `with_yjit` hooks, e.g. Array#each in Ruby
+    // Disable Ruby builtin methods defined by `with_jit` hooks, e.g. Array#each in Ruby
     pub c_builtin: bool,
 
     // Capture stats
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
@@ -57,7 +57,7 @@ fn yjit_init() {
     // 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);
+        rb_funcall(yjit, rust_str_to_id("call_jit_hooks"), 0);
     }
 
     // Catch panics to avoid UB for unwinding into C frames.
diff --git a/yjit_hook.rb b/yjit_hook.rb
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
@@ -957,7 +957,7 @@ pub use manual_defs::*;
 pub mod test_utils {
     use std::{ptr::null, sync::Once};
 
-    use crate::{options::init_options, state::rb_zjit_enabled_p, state::ZJITState};
+    use crate::{options::rb_zjit_prepare_options, state::rb_zjit_enabled_p, state::ZJITState};
 
     use super::*;
 
@@ -979,6 +979,7 @@ pub mod test_utils {
             // <https://github.com/Shopify/zjit/pull/37>, though
             let mut var: VALUE = Qnil;
             ruby_init_stack(&mut var as *mut VALUE as *mut _);
+            rb_zjit_prepare_options(); // enable `#with_jit` on builtins
             ruby_init();
 
             // Pass command line options so the VM loads core library methods defined in
@@ -994,7 +995,7 @@ pub mod test_utils {
         }
 
         // Set up globals for convenience
-        ZJITState::init(init_options());
+        ZJITState::init();
 
         // Enable zjit_* instructions
         unsafe { rb_zjit_enabled_p = true; }
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
diff --git a/zjit/src/state.rs b/zjit/src/state.rs

Original file line number	Diff line number	Diff line change
`@@ -322,7 +322,7 @@ def denominator`
`322`	`322`	`1`
`323`	`323`	`end`
`324`	`324`
`325`		`- with_yjit do`
	`325`	`+ with_jit do`
`326`	`326`	`if Primitive.rb_builtin_basic_definition_p(:downto)`
`327`	`327`	`undef :downto`
`328`	`328`