ruby · maximecb · Nov 4, 2024 · Oct 21, 2024 · Oct 28, 2024 · Oct 28, 2024
diff --git a/array.c b/array.c
@@ -2604,6 +2604,39 @@ ary_fetch_next(VALUE self, VALUE *index, VALUE *value)
     return Qtrue;
 }
 
+/*
+ *  call-seq:
+ *    each {|element| ... } -> self
+ *    each -> new_enumerator
+ *
+ *  With a block given, iterates over the elements of +self+,
+ *  passing each element to the block;
+ *  returns +self+:
+ *
+ *    a = [:foo, 'bar', 2]
+ *    a.each {|element|  puts "#{element.class} #{element}" }
+ *
+ *  Output:
+ *
+ *    Symbol foo
+ *    String bar
+ *    Integer 2
+ *
+ *  Allows the array to be modified during iteration:
+ *
+ *    a = [:foo, 'bar', 2]
+ *    a.each {|element| puts element; a.clear if element.to_s.start_with?('b') }
+ *
+ *  Output:
+ *
+ *    foo
+ *    bar
+ *
+ *  With no block given, returns a new Enumerator.
+ *
+ *  Related: see {Methods for Iterating}[rdoc-ref:Array@Methods+for+Iterating].
+ */
+
 VALUE
 rb_ary_each(VALUE ary)
 {
@@ -8631,6 +8664,7 @@ Init_Array(void)
     rb_define_method(rb_cArray, "unshift", rb_ary_unshift_m, -1);
     rb_define_alias(rb_cArray,  "prepend", "unshift");
     rb_define_method(rb_cArray, "insert", rb_ary_insert, -1);
+    rb_define_method(rb_cArray, "each", rb_ary_each, 0);
     rb_define_method(rb_cArray, "each_index", rb_ary_each_index, 0);
     rb_define_method(rb_cArray, "reverse_each", rb_ary_reverse_each, 0);
     rb_define_method(rb_cArray, "length", rb_ary_length, 0);

diff --git a/array.rb b/array.rb
@@ -1,49 +1,4 @@
 class Array
-  # call-seq:
-  #   each {|element| ... } -> self
-  #   each -> new_enumerator
-  #
-  # With a block given, iterates over the elements of +self+,
-  # passing each element to the block;
-  # returns +self+:
-  #
-  #   a = [:foo, 'bar', 2]
-  #   a.each {|element|  puts "#{element.class} #{element}" }
-  #
-  # Output:
-  #
-  #   Symbol foo
-  #   String bar
-  #   Integer 2
-  #
-  # Allows the array to be modified during iteration:
-  #
-  #   a = [:foo, 'bar', 2]
-  #   a.each {|element| puts element; a.clear if element.to_s.start_with?('b') }
-  #
-  # Output:
-  #
-  #   foo
-  #   bar
-  #
-  # With no block given, returns a new Enumerator.
-  #
-  # Related: see {Methods for Iterating}[rdoc-ref:Array@Methods+for+Iterating].
-
-  def each
-    Primitive.attr! :inline_block
-
-    unless defined?(yield)
-      return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
-    end
-    _i = 0
-    value = nil
-    while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
-      yield value
-    end
-    self
-  end
-
   # call-seq:
   #   shuffle!(random: Random) -> self
   #
@@ -258,4 +213,24 @@ def fetch_values(*indexes, &block)
     indexes.map! { |i| fetch(i, &block) }
     indexes
   end
+
+  with_yjit do
+    if Array.instance_method(:each).source_location == nil # preserve monkey patches
+      undef :each
+
+      def each # :nodoc:
+        Primitive.attr! :inline_block, :c_trace
+
+        unless defined?(yield)
+          return Primitive.cexpr! 'SIZED_ENUMERATOR(self, 0, 0, ary_enum_length)'
+        end
+        _i = 0
+        value = nil
+        while Primitive.cexpr!(%q{ ary_fetch_next(self, LOCAL_PTR(_i), LOCAL_PTR(value)) })
+          yield value
+        end
+        self
+      end
+    end
+  end
 end
diff --git a/common.mk b/common.mk
@@ -1212,6 +1212,7 @@ BUILTIN_RB_SRCS = \
 		$(srcdir)/prelude.rb \
 		$(srcdir)/gem_prelude.rb \
 		$(srcdir)/yjit.rb \
+		$(srcdir)/yjit_hook.rb \
 		$(empty)
 BUILTIN_RB_INCS = $(BUILTIN_RB_SRCS:.rb=.rbinc)
 
@@ -10642,6 +10643,7 @@ 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
 node.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 node.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 node.$(OBJEXT): $(CCAN_DIR)/list/list.h
@@ -20002,6 +20004,7 @@ 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_backtrace.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 vm_backtrace.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 vm_backtrace.$(OBJEXT): $(CCAN_DIR)/list/list.h

diff --git a/compile.c b/compile.c
@@ -8960,6 +8960,9 @@ compile_builtin_attr(rb_iseq_t *iseq, const NODE *node)
         else if (strcmp(RSTRING_PTR(string), "use_block") == 0) {
             iseq_set_use_block(iseq);
         }
+        else if (strcmp(RSTRING_PTR(string), "c_trace") == 0) {
+            ISEQ_BODY(iseq)->builtin_attrs |= BUILTIN_ATTR_C_TRACE;
+        }
         else {
             goto unknown_arg;
         }

diff --git a/inits.c b/inits.c
@@ -84,6 +84,9 @@ void
 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);
@@ -95,11 +98,9 @@ rb_call_builtin_inits(void)
     BUILTIN(warning);
     BUILTIN(array);
     BUILTIN(hash);
-    BUILTIN(kernel);
     BUILTIN(symbol);
     BUILTIN(timev);
     BUILTIN(thread_sync);
-    BUILTIN(yjit);
     BUILTIN(nilclass);
     BUILTIN(marshal);
     BUILTIN(rjit_c);

diff --git a/kernel.rb b/kernel.rb
@@ -290,4 +290,12 @@ def Integer(arg, base = 0, exception: true)
       Primitive.rb_f_integer(arg, base, exception);
     end
   end
+
+  # Internal helper for builtin inits to define methods only when YJIT is enabled.
+  # This method is removed in yjit_hook.rb.
+  def with_yjit(&block) # :nodoc:
+    if defined?(RubyVM::YJIT)
+      RubyVM::YJIT.send(:add_yjit_hook, block)
+    end
+  end
 end
diff --git a/prism_compile.c b/prism_compile.c
@@ -3387,6 +3387,9 @@ pm_compile_builtin_attr(rb_iseq_t *iseq, const pm_scope_node_t *scope_node, cons
         else if (strcmp(RSTRING_PTR(string), "use_block") == 0) {
             iseq_set_use_block(iseq);
         }
+        else if (strcmp(RSTRING_PTR(string), "c_trace") == 0) {
+            ISEQ_BODY(iseq)->builtin_attrs |= BUILTIN_ATTR_C_TRACE;
+        }
         else {
             COMPILE_ERROR(iseq, node_location->line, "unknown argument to attr!: %s", RSTRING_PTR(string));
             return COMPILE_NG;

diff --git a/ruby.c b/ruby.c
@@ -1816,6 +1816,10 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
     rb_yjit_init(opt->yjit);
 #endif
 
+    // Call yjit_hook.rb after rb_yjit_init() to use `RubyVM::YJIT.enabled?`
+    void Init_builtin_yjit_hook();
+    Init_builtin_yjit_hook();
+
     ruby_set_script_name(opt->script_name);
     require_libraries(&opt->req_list);
 }

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
@@ -1677,6 +1677,71 @@ def test = :ok
     RUBY
   end
 
+  def test_yjit_option_uses_array_each_in_ruby
+    assert_separately(["--yjit"], <<~'RUBY')
+      # Array#each should be implemented in Ruby for YJIT
+      assert_equal "<internal:array>", Array.instance_method(:each).source_location.first
+
+      # The backtrace, however, should not be `from <internal:array>:XX:in 'Array#each'`
+      begin
+        [nil].each { raise }
+      rescue => e
+        assert_equal "-:11:in 'Array#each'", e.backtrace[1]
+      end
+    RUBY
+  end
+
+  def test_yjit_enable_replaces_array_each
+    assert_separately([*("--disable=yjit" if RubyVM::YJIT.enabled?)], <<~'RUBY')
+      # Array#each should be implemented in C for the interpreter
+      assert_nil Array.instance_method(:each).source_location
+
+      # The backtrace should not be `from <internal:array>:XX:in 'Array#each'`
+      begin
+        [nil].each { raise }
+      rescue => e
+        assert_equal "-:11:in 'Array#each'", e.backtrace[1]
+      end
+
+      RubyVM::YJIT.enable
+
+      # Array#each should be implemented in Ruby for YJIT
+      assert_equal "<internal:array>", Array.instance_method(:each).source_location.first
+
+      # However, the backtrace should still not be `from <internal:array>:XX:in 'Array#each'`
+      begin
+        [nil].each { raise }
+      rescue => e
+        assert_equal "-:23:in 'Array#each'", e.backtrace[1]
+      end
+    RUBY
+  end
+
+  def test_yjit_enable_preserves_array_each_monkey_patch
+    assert_separately([*("--disable=yjit" if RubyVM::YJIT.enabled?)], <<~'RUBY')
+      # Array#each should be implemented in C initially
+      assert_nil Array.instance_method(:each).source_location
+
+      # Override Array#each
+      $called = false
+      Array.prepend(Module.new {
+        def each
+          $called = true
+          super
+        end
+      })
+
+      RubyVM::YJIT.enable
+
+      # The monkey-patch should still be alive
+      [].each {}
+      assert_true $called
+
+      # YJIT should not replace Array#each with the "<internal:array>" one
+      assert_equal "-", Array.instance_method(:each).source_location.first
+    RUBY
+  end
+
   private
 
   def code_gc_helpers

diff --git a/tool/mk_builtin_loader.rb b/tool/mk_builtin_loader.rb
@@ -6,7 +6,7 @@
 
 SUBLIBS = {}
 REQUIRED = {}
-BUILTIN_ATTRS = %w[leaf inline_block use_block]
+BUILTIN_ATTRS = %w[leaf inline_block use_block c_trace]
 
 module CompileWarning
   @@warnings = 0

diff --git a/vm.c b/vm.c
@@ -4425,6 +4425,9 @@ Init_vm_objects(void)
 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 RJIT units
 #if !USE_RJIT
 void Init_builtin_rjit(void) {}

diff --git a/vm_backtrace.c b/vm_backtrace.c
@@ -448,7 +448,9 @@ location_to_str(rb_backtrace_location_t *loc)
     VALUE file, owner = Qnil, name;
     int lineno;
 
-    if (loc->cme && loc->cme->def->type == VM_METHOD_TYPE_CFUNC) {
+    if (loc->cme && (loc->cme->def->type == VM_METHOD_TYPE_CFUNC || (loc->cme->def->type == VM_METHOD_TYPE_ISEQ
+                    // Ruby methods with `Primitive.attr! :c_trace` should behave like C methods
+                    && (loc->cme->def->body.iseq.iseqptr->body->builtin_attrs & BUILTIN_ATTR_C_TRACE) != 0))) {
         if (loc->iseq && loc->pc) {
             file = rb_iseq_path(loc->iseq);
             lineno = calc_lineno(loc->iseq, loc->pc);
@@ -684,13 +686,21 @@ rb_ec_partial_backtrace_object(const rb_execution_context_t *ec, long start_fram
                         const VALUE *pc = cfp->pc;
                         loc = &bt->backtrace[bt->backtrace_size++];
                         RB_OBJ_WRITE(btobj, &loc->cme, rb_vm_frame_method_entry(cfp));
-                        RB_OBJ_WRITE(btobj, &loc->iseq, iseq);
-                        loc->pc = pc;
-                        bt_update_cfunc_loc(cfunc_counter, loc-1, iseq, pc);
-                        if (do_yield) {
-                            bt_yield_loc(loc - cfunc_counter, cfunc_counter+1, btobj);
+                        // Ruby methods with `Primitive.attr! :c_trace` should behave like C methods
+                        if ((cfp->iseq->body->builtin_attrs & BUILTIN_ATTR_C_TRACE) != 0) {
+                            loc->iseq = NULL;
+                            loc->pc = NULL;
+                            cfunc_counter++;
+                        }
+                        else {
+                            RB_OBJ_WRITE(btobj, &loc->iseq, iseq);
+                            loc->pc = pc;
+                            bt_update_cfunc_loc(cfunc_counter, loc-1, iseq, pc);
+                            if (do_yield) {
+                                bt_yield_loc(loc - cfunc_counter, cfunc_counter+1, btobj);
+                            }
+                            cfunc_counter = 0;
                         }
-                        cfunc_counter = 0;
                     }
                     skip_next_frame = is_rescue_or_ensure_frame(cfp);
                 }

diff --git a/vm_core.h b/vm_core.h
@@ -372,6 +372,8 @@ enum rb_builtin_attr {
     BUILTIN_ATTR_SINGLE_NOARG_LEAF = 0x02,
     // This attribute signals JIT to duplicate the iseq for each block iseq so that its `yield` will be monomorphic.
     BUILTIN_ATTR_INLINE_BLOCK = 0x04,
+    // The iseq acts like a C method in backtraces.
+    BUILTIN_ATTR_C_TRACE = 0x08,
 };
 
 typedef VALUE (*rb_jit_func_t)(struct rb_execution_context_struct *, struct rb_control_frame_struct *);

diff --git a/yjit.rb b/yjit.rb
@@ -37,7 +37,7 @@ def self.reset_stats!
   # whether to enable \YJIT compilation logging or not.
   #
   # `stats`:
-  # * `false`: Disable stats.
+  # * `false`: Don't enable stats.
   # * `true`: Enable stats. Print stats at exit.
   # * `:quiet`: Enable stats. Do not print stats at exit.
   #
@@ -48,6 +48,7 @@ def self.reset_stats!
   def self.enable(stats: false, log: false)
     return false if enabled?
     at_exit { print_and_dump_stats } if stats
+    call_yjit_hooks
     Primitive.rb_yjit_enable(stats, stats != :quiet, log, log != :quiet)
   end
 
@@ -247,10 +248,24 @@ def self.simulate_oom! # :nodoc:
     at_exit { print_and_dump_stats }
   end
 
+  # Blocks that are called when YJIT is enabled
+  @yjit_hooks = []
+
   class << self
     # :stopdoc:
     private
 
+    # Lazily call a given block when YJIT is enabled
+    def add_yjit_hook(hook)
+      @yjit_hooks << hook
+    end
+
+    # Run YJIT hooks registered by RubyVM::YJIT.with_yjit
+    def call_yjit_hooks
+      @yjit_hooks.each(&:call)
+      @yjit_hooks.clear
+    end
+
     # Print stats and dump exit locations
     def print_and_dump_stats # :nodoc:
       if Primitive.rb_yjit_print_stats_p