Transition complex objects to "too complex" shape

When an object becomes "too complex" (in other words it has too many
variations in the shape tree), we transition it to use a "too complex"
shape and use a hash for storing instance variables.

Without this patch, there were rare cases where shape tree growth could
"explode" and cause performance degradation on what would otherwise have
been cached fast paths.

This patch puts a limit on shape tree growth, and gracefully degrades in
the rare case where there could be a factorial growth in the shape tree.

For example:

```ruby
class NG; end

HUGE_NUMBER.times do
  NG.new.instance_variable_set(:"@unique_ivar_#{_1}", 1)
end
```

We consider objects to be "too complex" when the object's class has more
than SHAPE_MAX_VARIATIONS (currently 8) leaf nodes in the shape tree and
the object introduces a new variation (a new leaf node) associated with
that class.

For example, new variations on instances of the following class would be
considered "too complex" because those instances create more than 8
leaves in the shape tree:

```ruby
class Foo; end
9.times { Foo.new.instance_variable_set(":@uniq_#{_1}", 1) }
```

However, the following class is *not* too complex because it only has
one leaf in the shape tree:

```ruby
class Foo
  def initialize
    @a = @b = @c = @d = @e = @f = @g = @h = @i = nil
  end
end
9.times { Foo.new }
``

This case is rare, so we don't expect this change to impact performance
of most applications, but it needs to be handled.

Co-Authored-By: Aaron Patterson <tenderlove@ruby-lang.org>

4 files changed, 42 insertions, 16 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index e33b0d5dba..996d7fa1c0 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -93,7 +93,9 @@ fn main() {
         .allowlist_function("rb_shape_get_next")
         .allowlist_function("rb_shape_id")
         .allowlist_function("rb_shape_transition_shape_capa")
+        .allowlist_function("rb_shape_obj_too_complex")
         .allowlist_var("SHAPE_ID_NUM_BITS")
+        .allowlist_var("OBJ_TOO_COMPLEX_SHAPE_ID")
 
         // From ruby/internal/intern/object.h
         .allowlist_function("rb_obj_is_kind_of")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index ab87038c67..cec5c4671b 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1963,6 +1963,11 @@ fn gen_get_ivar(
     recv_opnd: YARVOpnd,
     side_exit: CodePtr,
 ) -> CodegenStatus {
+    // If the object has a too complex shape, we exit
+    if comptime_receiver.shape_too_complex() {
+        return CantCompile;
+    }
+
     let comptime_val_klass = comptime_receiver.class_of();
     let starting_context = ctx.clone(); // make a copy for use with jit_chain_guard
 
@@ -2192,7 +2197,8 @@ fn gen_setinstancevariable(
 
     // If the comptime receiver is frozen, writing an IV will raise an exception
     // and we don't want to JIT code to deal with that situation.
-    if comptime_receiver.is_frozen() {
+    // If the object has a too complex shape, we will also exit
+    if comptime_receiver.is_frozen() || comptime_receiver.shape_too_complex() {
         return CantCompile;
     }
 
@@ -2281,39 +2287,53 @@ fn gen_setinstancevariable(
             megamorphic_side_exit,
         );
 
-        let write_val = ctx.stack_pop(1);
+        let write_val;
 
         match ivar_index {
             // If we don't have an instance variable index, then we need to
             // transition out of the current shape.
             None => {
-                let mut shape = comptime_receiver.shape_of();
+                let shape = comptime_receiver.shape_of();
+
+                let current_capacity = unsafe { (*shape).capacity };
+                let new_capacity = current_capacity * 2;
 
                 // If the object doesn't have the capacity to store the IV,
                 // then we'll need to allocate it.
-                let needs_extension = unsafe { (*shape).next_iv_index >= (*shape).capacity };
+                let needs_extension = unsafe { (*shape).next_iv_index >= current_capacity };
 
                 // We can write to the object, but we need to transition the shape
                 let ivar_index = unsafe { (*shape).next_iv_index } as usize;
 
-                if needs_extension {
-                    let current_capacity = unsafe { (*shape).capacity };
-                    let newsize = current_capacity * 2;
-
+                let capa_shape = if needs_extension {
                     // We need to add an extended table to the object
                     // First, create an outgoing transition that increases the
                     // capacity
-                    shape = unsafe {
-                        rb_shape_transition_shape_capa(shape, newsize)
-                    };
+                    Some(unsafe { rb_shape_transition_shape_capa(shape, new_capacity) })
+                } else {
+                    None
+                };
+
+                let dest_shape = if capa_shape.is_none() {
+                    unsafe { rb_shape_get_next(shape, comptime_receiver, ivar_name) }
+                } else {
+                    unsafe { rb_shape_get_next(capa_shape.unwrap(), comptime_receiver, ivar_name) }
+                };
+
+                let new_shape_id = unsafe { rb_shape_id(dest_shape) };
 
+                if new_shape_id == OBJ_TOO_COMPLEX_SHAPE_ID {
+                    return CantCompile;
+                }
+
+                if needs_extension {
                     // Generate the C call so that runtime code will increase
                     // the capacity and set the buffer.
                     asm.ccall(rb_ensure_iv_list_size as *const u8,
                               vec![
                                   recv,
                                   Opnd::UImm(current_capacity.into()),
-                                  Opnd::UImm(newsize.into())
+                                  Opnd::UImm(new_capacity.into())
                               ]
                     );
 
@@ -2321,10 +2341,7 @@ fn gen_setinstancevariable(
                     recv = asm.load(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF))
                 }
 
-                let new_shape_id = unsafe {
-                    rb_shape_id(rb_shape_get_next(shape, comptime_receiver, ivar_name))
-                };
-
+                write_val = ctx.stack_pop(1);
                 gen_write_iv(asm, comptime_receiver, recv, ivar_index, write_val, needs_extension);
 
                 asm.comment("write shape");
@@ -2342,6 +2359,7 @@ fn gen_setinstancevariable(
                 // the iv index by searching up the shape tree.  If we've
                 // made the transition already, then there's no reason to
                 // update the shape on the object.  Just set the IV.
+                write_val = ctx.stack_pop(1);
                 gen_write_iv(asm, comptime_receiver, recv, ivar_index, write_val, false);
             },
         }
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index b6228fe64b..ba09d4119f 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -398,6 +398,10 @@ impl VALUE {
         unsafe { rb_obj_frozen_p(self) != VALUE(0) }
     }
 
+    pub fn shape_too_complex(self) -> bool {
+        unsafe { rb_shape_obj_too_complex(self) }
+    }
+
     pub fn shape_id_of(self) -> u32 {
         unsafe { rb_shape_get_shape_id(self) }
     }
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 3e00aa3689..af77747861 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -124,6 +124,7 @@ impl<T> ::std::cmp::PartialEq for __BindgenUnionField<T> {
 }
 impl<T> ::std::cmp::Eq for __BindgenUnionField<T> {}
 pub const SHAPE_ID_NUM_BITS: u32 = 32;
+pub const OBJ_TOO_COMPLEX_SHAPE_ID: u32 = 11;
 pub const INTEGER_REDEFINED_OP_FLAG: u32 = 1;
 pub const FLOAT_REDEFINED_OP_FLAG: u32 = 2;
 pub const STRING_REDEFINED_OP_FLAG: u32 = 4;
@@ -1112,6 +1113,7 @@ extern "C" {
     pub fn rb_shape_get_next(shape: *mut rb_shape_t, obj: VALUE, id: ID) -> *mut rb_shape_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_id(shape: *mut rb_shape_t) -> shape_id_t;
+    pub fn rb_shape_obj_too_complex(obj: VALUE) -> bool;
     pub fn rb_ary_tmp_new_from_values(
         arg1: VALUE,
         arg2: ::std::os::raw::c_long,