Do not emit shape transition warnings when YJIT is compiling

[Bug #20522]

If `Warning.warn` is redefined in Ruby, emitting a warning would invoke
Ruby code, which can't safely be done when YJIT is compiling.

Author: byroot

bootstraptest/test_yjit.rb

@@ -5013,3 +5013,37 @@ def test_body(args) = Array(1, *args)
     test_body(array)
   end
 }
+
+# compiling code shouldn't emit warnings as it may call into more Ruby code
+assert_normal_exit %q{
+  # [Bug #20522]
+  Warning[:performance] = true
+
+  module StrictWarnings
+    def warn(msg, category: nil, **)
+      raise msg
+    end
+  end
+  Warning.singleton_class.prepend(StrictWarnings)
+
+  class A
+    def compiled_method(is_private)
+      @some_ivar = is_private
+    end
+  end
+
+  100.times do |i|
+    klass = Class.new(A)
+    7.times do |j|
+      obj = klass.new
+      obj.instance_variable_set("@base_#{i}", 42)
+      obj.instance_variable_set("@ivar_#{j}", 42)
+    end
+    obj = klass.new
+    obj.instance_variable_set("@base_#{i}", 42)
+    begin
+      obj.compiled_method(true)
+    rescue
+    end
+  end
+}

shape.c

@@ -696,8 +696,8 @@ rb_shape_get_next_iv_shape(rb_shape_t* shape, ID id)
     return get_next_shape_internal(shape, id, SHAPE_IVAR, &dont_care, true);
 }
 
-rb_shape_t *
-rb_shape_get_next(rb_shape_t *shape, VALUE obj, ID id)
+static inline rb_shape_t *
+shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
 {
     RUBY_ASSERT(!is_instance_id(id) || RTEST(rb_sym2str(ID2SYM(id))));
     if (UNLIKELY(shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
@@ -730,7 +730,7 @@ rb_shape_get_next(rb_shape_t *shape, VALUE obj, ID id)
 
         if (variation_created) {
             RCLASS_EXT(klass)->variation_count++;
-            if (rb_warning_category_enabled_p(RB_WARN_CATEGORY_PERFORMANCE)) {
+            if (emit_warnings && rb_warning_category_enabled_p(RB_WARN_CATEGORY_PERFORMANCE)) {
                 if (RCLASS_EXT(klass)->variation_count >= SHAPE_MAX_VARIATIONS) {
                     rb_category_warn(
                         RB_WARN_CATEGORY_PERFORMANCE,
@@ -747,6 +747,18 @@ rb_shape_get_next(rb_shape_t *shape, VALUE obj, ID id)
     return new_shape;
 }
 
+rb_shape_t *
+rb_shape_get_next(rb_shape_t *shape, VALUE obj, ID id)
+{
+    return shape_get_next(shape, obj, id, true);
+}
+
+rb_shape_t *
+rb_shape_get_next_no_warnings(rb_shape_t *shape, VALUE obj, ID id)
+{
+    return shape_get_next(shape, obj, id, false);
+}
+
 // Same as rb_shape_get_iv_index, but uses a provided valid shape id and index
 // to return a result faster if branches of the shape tree are closely related.
 bool

shape.h

@@ -164,6 +164,7 @@ int rb_shape_frozen_shape_p(rb_shape_t* shape);
 rb_shape_t* rb_shape_transition_shape_frozen(VALUE obj);
 bool rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape, VALUE * removed);
 rb_shape_t* rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id);
+rb_shape_t* rb_shape_get_next_no_warnings(rb_shape_t* shape, VALUE obj, ID id);
 
 rb_shape_t * rb_shape_rebuild_shape(rb_shape_t * initial_shape, rb_shape_t * dest_shape);
 

yjit/bindgen/src/main.rs

@@ -100,7 +100,7 @@ fn main() {
         .allowlist_function("rb_shape_get_shape_by_id")
         .allowlist_function("rb_shape_id_offset")
         .allowlist_function("rb_shape_get_iv_index")
-        .allowlist_function("rb_shape_get_next")
+        .allowlist_function("rb_shape_get_next_no_warnings")
         .allowlist_function("rb_shape_id")
         .allowlist_function("rb_shape_obj_too_complex")
         .allowlist_var("SHAPE_ID_NUM_BITS")

yjit/src/codegen.rs

@@ -2970,7 +2970,7 @@ fn gen_set_ivar(
     // The current shape doesn't contain this iv, we need to transition to another shape.
     let new_shape = if !shape_too_complex && receiver_t_object && ivar_index.is_none() {
         let current_shape = comptime_receiver.shape_of();
-        let next_shape = unsafe { rb_shape_get_next(current_shape, comptime_receiver, ivar_name) };
+        let next_shape = unsafe { rb_shape_get_next_no_warnings(current_shape, comptime_receiver, ivar_name) };
         let next_shape_id = unsafe { rb_shape_id(next_shape) };
 
         // If the VM ran out of shapes, or this class generated too many leaf,

yjit/src/cruby_bindings.inc.rs

@@ -1063,7 +1063,11 @@ extern "C" {
     pub fn rb_shape_get_shape_id(obj: VALUE) -> shape_id_t;
     pub fn rb_shape_get_iv_index(shape: *mut rb_shape_t, id: ID, value: *mut attr_index_t) -> bool;
     pub fn rb_shape_obj_too_complex(obj: VALUE) -> bool;
-    pub fn rb_shape_get_next(shape: *mut rb_shape_t, obj: VALUE, id: ID) -> *mut rb_shape_t;
+    pub fn rb_shape_get_next_no_warnings(
+        shape: *mut rb_shape_t,
+        obj: VALUE,
+        id: ID,
+    ) -> *mut rb_shape_t;
     pub fn rb_shape_id(shape: *mut rb_shape_t) -> shape_id_t;
     pub fn rb_gvar_get(arg1: ID) -> VALUE;
     pub fn rb_gvar_set(arg1: ID, arg2: VALUE) -> VALUE;