YJIT: Replace Array#each only when YJIT is enabled (#11955)

* YJIT: Replace Array#each only when YJIT is enabled

* Add comments about BUILTIN_ATTR_C_TRACE

* Make Ruby Array#each available with --yjit as well

* Fix all paths that expect a C location

* Use method_basic_definition_p to detect patches

* Copy a comment about C_TRACE flag to compilers

* Rephrase a comment about add_yjit_hook

* Give METHOD_ENTRY_BASIC flag to Array#each

* Add --yjit-c-builtin option

* Allow inconsistent source_location in test-spec

* Refactor a check of BUILTIN_ATTR_C_TRACE

* Set METHOD_ENTRY_BASIC without touching vm->running

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)
 {
@@ -8634,6 +8667,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);

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 @@ class Array
     indexes.map! { |i| fetch(i, &block) }
     indexes
   end
+
+  with_yjit do
+    if Primitive.rb_builtin_basic_definition_p(:each)
+      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

builtin.h

@@ -106,6 +106,12 @@ rb_vm_lvar(rb_execution_context_t *ec, int index)
 #endif
 }
 
+static inline VALUE
+rb_builtin_basic_definition_p(rb_execution_context_t *ec, VALUE klass, VALUE id_sym)
+{
+    return rb_method_basic_definition_p(klass, rb_sym2id(id_sym)) ? Qtrue : Qfalse;
+}
+
 #define LOCAL_PTR(local) local ## __ptr
 
 // dump/load

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)
 
@@ -10674,6 +10675,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
@@ -20036,6 +20038,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

compile.c

@@ -8960,6 +8960,10 @@ 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) {
+            // Let the iseq act like a C method in backtraces
+            ISEQ_BODY(iseq)->builtin_attrs |= BUILTIN_ATTR_C_TRACE;
+        }
         else {
             goto unknown_arg;
         }

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);

kernel.rb

@@ -290,4 +290,12 @@ module Kernel
       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

prism_compile.c

@@ -3387,6 +3387,10 @@ 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) {
+            // Let the iseq act like a C method in backtraces
+            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;

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);
 }

spec/ruby/optional/capi/debug_spec.rb

@@ -37,9 +37,12 @@ describe "C-API Debug function" do
 
     it "matches the locations in rb_debug_inspector_backtrace_locations" do
       frames = @o.rb_debug_inspector_open(42)
-      frames.each do |_s, _klass, binding, _iseq, backtrace_location|
+      frames.each do |_s, klass, binding, iseq, backtrace_location|
         if binding
-          binding.source_location.should == [backtrace_location.path, backtrace_location.lineno]
+          # YJIT modifies Array#each backtraces but leaves its source_location as is
+          unless defined?(RubyVM::YJIT) && klass == Array && iseq.label == "each"
+            binding.source_location.should == [backtrace_location.path, backtrace_location.lineno]
+          end
           method_name = binding.eval('__method__')
           if method_name
             method_name.should == backtrace_location.base_label.to_sym

test/ruby/test_yjit.rb

@@ -1677,6 +1677,71 @@ class TestYJIT < Test::Unit::TestCase
     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

tool/mk_builtin_loader.rb

@@ -6,7 +6,7 @@ require_relative 'ruby_vm/helpers/c_escape'
 
 SUBLIBS = {}
 REQUIRED = {}
-BUILTIN_ATTRS = %w[leaf inline_block use_block]
+BUILTIN_ATTRS = %w[leaf inline_block use_block c_trace]
 
 module CompileWarning
   @@warnings = 0

vm.c

@@ -4437,6 +4437,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) {}

vm_backtrace.c

@@ -265,10 +265,26 @@ retry:
     }
 }
 
+// Return true if a given location is a C method or supposed to behave like one.
+static inline bool
+location_cfunc_p(rb_backtrace_location_t *loc)
+{
+    if (!loc->cme) return false;
+
+    switch (loc->cme->def->type) {
+      case VM_METHOD_TYPE_CFUNC:
+        return true;
+      case VM_METHOD_TYPE_ISEQ:
+        return rb_iseq_attr_p(loc->cme->def->body.iseq.iseqptr, BUILTIN_ATTR_C_TRACE);
+      default:
+        return false;
+    }
+}
+
 static VALUE
 location_label(rb_backtrace_location_t *loc)
 {
-    if (loc->cme && loc->cme->def->type == VM_METHOD_TYPE_CFUNC) {
+    if (location_cfunc_p(loc)) {
         return rb_gen_method_name(loc->cme->owner, rb_id2str(loc->cme->def->original_id));
     }
     else {
@@ -314,7 +330,7 @@ location_label_m(VALUE self)
 static VALUE
 location_base_label(rb_backtrace_location_t *loc)
 {
-    if (loc->cme && loc->cme->def->type == VM_METHOD_TYPE_CFUNC) {
+    if (location_cfunc_p(loc)) {
         return rb_id2str(loc->cme->def->original_id);
     }
 
@@ -448,7 +464,7 @@ 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 (location_cfunc_p(loc)) {
         if (loc->iseq && loc->pc) {
             file = rb_iseq_path(loc->iseq);
             lineno = calc_lineno(loc->iseq, loc->pc);
@@ -684,13 +700,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);
+                        // Ruby methods with `Primitive.attr! :c_trace` should behave like C methods
-                        loc->pc = pc;
+                        if (rb_iseq_attr_p(cfp->iseq, BUILTIN_ATTR_C_TRACE)) {
-                        bt_update_cfunc_loc(cfunc_counter, loc-1, iseq, pc);
+                            loc->iseq = NULL;
-                        if (do_yield) {
+                            loc->pc = NULL;
-                            bt_yield_loc(loc - cfunc_counter, cfunc_counter+1, btobj);
+                            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);
                 }

vm_core.h

@@ -395,6 +395,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 *);
@@ -604,6 +606,12 @@ rb_iseq_check(const rb_iseq_t *iseq)
     return iseq;
 }
 
+static inline bool
+rb_iseq_attr_p(const rb_iseq_t *iseq, enum rb_builtin_attr attr)
+{
+    return (ISEQ_BODY(iseq)->builtin_attrs & attr) == attr;
+}
+
 static inline const rb_iseq_t *
 def_iseq_ptr(rb_method_definition_t *def)
 {

vm_method.c

@@ -641,6 +641,11 @@ 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
+                if (rb_iseq_attr_p(iseq, BUILTIN_ATTR_C_TRACE)) {
+                    METHOD_ENTRY_BASIC_SET((rb_method_entry_t *)me, TRUE);
+                }
+
                 if (ISEQ_BODY(iseq)->mandatory_only_iseq) def->iseq_overload = 1;
 
                 if (0) vm_cref_dump("rb_method_definition_create", cref);

yjit.c

@@ -1244,6 +1244,14 @@ VALUE rb_yjit_code_gc(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_simulate_oom_bang(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_get_exit_locations(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_enable(rb_execution_context_t *ec, VALUE self, VALUE gen_stats, VALUE print_stats, VALUE gen_compilation_log, VALUE print_compilation_log);
+VALUE rb_yjit_c_builtin_p(rb_execution_context_t *ec, VALUE self);
+
+// Allow YJIT_C_BUILTIN macro to force --yjit-c-builtin
+#ifdef YJIT_C_BUILTIN
+static VALUE yjit_c_builtin_p(rb_execution_context_t *ec, VALUE self) { return Qtrue; }
+#else
+#define yjit_c_builtin_p rb_yjit_c_builtin_p
+#endif
 
 // Preprocessed yjit.rb generated during build
 #include "yjit.rbinc"

yjit.rb

@@ -37,7 +37,7 @@ module RubyVM::YJIT
   # 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 @@ module RubyVM::YJIT
   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,26 @@ module RubyVM::YJIT
     at_exit { print_and_dump_stats }
   end
 
+  # Blocks that are called when YJIT is enabled
+  @yjit_hooks = []
+
   class << self
     # :stopdoc:
     private
 
+    # Register a block to be called 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
+      # 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
+    end
+
     # Print stats and dump exit locations
     def print_and_dump_stats # :nodoc:
       if Primitive.rb_yjit_print_stats_p

yjit/src/cruby_bindings.inc.rs

@@ -492,6 +492,7 @@ pub type rb_iseq_type = u32;
 pub const BUILTIN_ATTR_LEAF: rb_builtin_attr = 1;
 pub const BUILTIN_ATTR_SINGLE_NOARG_LEAF: rb_builtin_attr = 2;
 pub const BUILTIN_ATTR_INLINE_BLOCK: rb_builtin_attr = 4;
+pub const BUILTIN_ATTR_C_TRACE: rb_builtin_attr = 8;
 pub type rb_builtin_attr = u32;
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]

yjit/src/options.rs

@@ -1,5 +1,5 @@
 use std::{ffi::{CStr, CString}, ptr::null, fs::File};
-use crate::{backend::current::TEMP_REGS, stats::Counter};
+use crate::{backend::current::TEMP_REGS, cruby::*, stats::Counter};
 use std::os::raw::{c_char, c_int, c_uint};
 
 // Call threshold for small deployments and command-line apps
@@ -46,6 +46,9 @@ 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
+    pub c_builtin: bool,
+
     // Capture stats
     pub gen_stats: bool,
 
@@ -94,6 +97,7 @@ pub static mut OPTIONS: Options = Options {
     no_type_prop: false,
     max_versions: 4,
     num_temp_regs: 5,
+    c_builtin: false,
     gen_stats: false,
     trace_exits: None,
     print_stats: true,
@@ -270,6 +274,10 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
             }
         },
 
+        ("c-builtin", _) => unsafe {
+            OPTIONS.c_builtin = true;
+        },
+
         ("code-gc", _) => unsafe {
             OPTIONS.code_gc = true;
         },
@@ -413,3 +421,13 @@ pub extern "C" fn rb_yjit_show_usage(help: c_int, highlight: c_int, width: c_uin
         unsafe { ruby_show_usage_line(name.as_ptr(), null(), description.as_ptr(), help, highlight, width, columns) }
     }
 }
+
+/// Return true if --yjit-c-builtin is given
+#[no_mangle]
+pub extern "C" fn rb_yjit_c_builtin_p(_ec: EcPtr, _self: VALUE) -> VALUE {
+    if get_option!(c_builtin) {
+        Qtrue
+    } else {
+        Qfalse
+    }
+}

yjit_hook.rb

@@ -0,0 +1,9 @@
+# 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
+module Kernel
+  undef :with_yjit
+end