ZJIT: Register builtin CMEs before prelude to avoid prepend crash

Split rb_zjit_init into rb_zjit_init_builtin_cmes (called before
ruby_init_prelude) and rb_zjit_init (called after). The prelude may
load bundler via BUNDLER_SETUP which can call Kernel.prepend, moving
core methods to an origin iclass. Registering method annotations
before the prelude ensures we capture CMEs while core classes are
pristine.

Author: rwstauner

ruby.c

@@ -1824,6 +1824,10 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
 #if USE_YJIT
     rb_yjit_init_builtin_cmes();
 #endif
+#if USE_ZJIT
+    extern void rb_zjit_init_builtin_cmes(void);
+    rb_zjit_init_builtin_cmes();
+#endif
 
     ruby_init_prelude();
 

test/ruby/test_zjit.rb

@@ -112,6 +112,17 @@ def test_zjit_disable
     RUBY
   end
 
+  def test_zjit_prelude_kernel_prepend
+    # Simulate what bundler/setup can do: prepend a module to Kernel during
+    # the prelude via the BUNDLER_SETUP mechanism in rubygems.rb:
+    #   require ENV["BUNDLER_SETUP"] if ENV["BUNDLER_SETUP"] && !defined?(Bundler)
+    Tempfile.create(["kernel_prepend", ".rb"]) do |f|
+      f.write("Kernel.prepend(Module.new)\n")
+      f.flush
+      assert_separately([{ "BUNDLER_SETUP" => f.path }, "--enable=gems", "--zjit"], "", ignore_stderr: true)
+    end
+  end
+
   def test_zjit_enable_respects_existing_options
     assert_separately(['--zjit-disable', '--zjit-stats-quiet'], <<~RUBY)
       refute_predicate RubyVM::ZJIT, :enabled?

zjit/src/state.rs

@@ -292,25 +292,28 @@ impl ZJITState {
     }
 }
 
-/// Initialize ZJIT at boot. This is called even if ZJIT is disabled.
+/// Initialize IDs and annotate builtin C method entries.
+/// Must be called at boot before ruby_init_prelude() since the prelude
+/// could redefine core methods (e.g. Kernel.prepend via bundler).
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_init(zjit_enabled: bool) {
+pub extern "C" fn rb_zjit_init_builtin_cmes() {
     use InitializationState::*;
 
     debug_assert!(
         matches!(unsafe { &ZJIT_STATE }, Uninitialized),
-        "rb_zjit_init should only be called once during boot",
+        "rb_zjit_init_builtin_cmes should only be called once during boot",
     );
 
-    // Initialize IDs and method annotations.
-    // cruby_methods::init() must be called at boot,
-    // as cmes could have been re-defined after boot.
     cruby::ids::init();
 
     let method_annotations = cruby_methods::init();
 
     unsafe { ZJIT_STATE = Initialized(method_annotations); }
+}
 
+/// Initialize ZJIT at boot. This is called even if ZJIT is disabled.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_init(zjit_enabled: bool) {
     // If --zjit, enable ZJIT immediately
     if zjit_enabled {
         zjit_enable();