Refactor rb_shape_transition_shape_capa out

Right now the `rb_shape_get_next` shape caller need to
first check if there is capacity left, and if not call
`rb_shape_transition_shape_capa` before it can call `rb_shape_get_next`.

And on each of these it needs to checks if we got a TOO_COMPLEX
back.

All this logic is duplicated in the interpreter, YJIT and RJIT.

Instead we can have `rb_shape_get_next` do the capacity transition
when needed. The caller can compare the old and new shapes capacity
to know if resizing is needed. It also can check for TOO_COMPLEX
only once.

internal/variable.h

@@ -64,7 +64,6 @@ VALUE rb_gvar_get(ID);
 VALUE rb_gvar_set(ID, VALUE);
 VALUE rb_gvar_defined(ID);
 void rb_const_warn_if_deprecated(const rb_const_entry_t *, VALUE, ID);
-rb_shape_t * rb_grow_iv_list(VALUE obj);
 void rb_ensure_iv_list_size(VALUE obj, uint32_t len, uint32_t newsize);
 attr_index_t rb_obj_ivar_set(VALUE obj, ID id, VALUE val);
 

lib/ruby_vm/rjit/insn_compiler.rb

@@ -502,41 +502,26 @@ def setinstancevariable(jit, ctx, asm)
           shape = C.rb_shape_get_shape_by_id(shape_id)
 
           current_capacity = shape.capacity
-          # If the object doesn't have the capacity to store the IV,
-          # then we'll need to allocate it.
-          needs_extension = shape.next_iv_index >= current_capacity
-
-          # We can write to the object, but we need to transition the shape
-          ivar_index = shape.next_iv_index
-
-          capa_shape =
-            if needs_extension
-              # We need to add an extended table to the object
-              # First, create an outgoing transition that increases the capacity
-              C.rb_shape_transition_shape_capa(shape)
-            else
-              nil
-            end
-
-          dest_shape =
-            if capa_shape
-              C.rb_shape_get_next(capa_shape, comptime_receiver, ivar_name)
-            else
-              C.rb_shape_get_next(shape, comptime_receiver, ivar_name)
-            end
+          dest_shape = C.rb_shape_get_next(shape, comptime_receiver, ivar_name)
           new_shape_id = C.rb_shape_id(dest_shape)
 
           if new_shape_id == C::OBJ_TOO_COMPLEX_SHAPE_ID
             asm.incr_counter(:setivar_too_complex)
             return CantCompile
           end
 
+          ivar_index = shape.next_iv_index
+
+          # If the new shape has a different capacity, we need to
+          # reallocate the object.
+          needs_extension = dest_shape.capacity != shape.capacity
+
           if needs_extension
             # Generate the C call so that runtime code will increase
             # the capacity and set the buffer.
             asm.mov(C_ARGS[0], :rax)
             asm.mov(C_ARGS[1], current_capacity)
-            asm.mov(C_ARGS[2], capa_shape.capacity)
+            asm.mov(C_ARGS[2], dest_shape.capacity)
             asm.call(C.rb_ensure_iv_list_size)
 
             # Load the receiver again after the function call

rjit_c.rb

@@ -171,12 +171,6 @@ def rb_method_entry_at(klass, mid)
       me_addr == 0 ? nil : rb_method_entry_t.new(me_addr)
     end
 
-    def rb_shape_transition_shape_capa(shape)
-      _shape = shape.to_i
-      shape_addr = Primitive.cexpr! 'SIZET2NUM((size_t)rb_shape_transition_shape_capa((rb_shape_t *)NUM2SIZET(_shape)))'
-      rb_shape_t.new(shape_addr)
-    end
-
     def rb_shape_get_next(shape, obj, id)
       _shape = shape.to_i
       shape_addr = Primitive.cexpr! 'SIZET2NUM((size_t)rb_shape_get_next((rb_shape_t *)NUM2SIZET(_shape), obj, (ID)NUM2SIZET(id)))'

shape.c

@@ -44,6 +44,8 @@ static ID id_frozen;
 static ID id_t_object;
 static ID size_pool_edge_names[SIZE_POOL_COUNT];
 
+rb_shape_t * rb_shape_transition_shape_capa(rb_shape_t * shape);
+
 #define LEAF 0
 #define BLACK 0x0
 #define RED 0x1
@@ -667,7 +669,9 @@ rb_shape_t *
 rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id)
 {
     RUBY_ASSERT(!is_instance_id(id) || RTEST(rb_sym2str(ID2SYM(id))));
-    RUBY_ASSERT(shape->type != SHAPE_OBJ_TOO_COMPLEX);
+    if (UNLIKELY(shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
+        return shape;
+    }
 
     bool allow_new_shape = true;
 
@@ -676,6 +680,14 @@ rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id)
         allow_new_shape = RCLASS_EXT(klass)->variation_count < SHAPE_MAX_VARIATIONS;
     }
 
+    if (UNLIKELY(shape->next_iv_index >= shape->capacity)) {
+        RUBY_ASSERT(shape->next_iv_index == shape->capacity);
+        shape = rb_shape_transition_shape_capa(shape);
+        if (UNLIKELY(shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
+            return shape;
+        }
+    }
+
     bool variation_created = false;
     rb_shape_t * new_shape = get_next_shape_internal(shape, id, SHAPE_IVAR, &variation_created, allow_new_shape);
 
@@ -721,6 +733,9 @@ rb_shape_transition_shape_capa_create(rb_shape_t* shape, size_t new_capacity)
 rb_shape_t *
 rb_shape_transition_shape_capa(rb_shape_t* shape)
 {
+    if (UNLIKELY(shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
+        return shape;
+    }
     return rb_shape_transition_shape_capa_create(shape, rb_malloc_grow_capa(shape->capacity, sizeof(VALUE)));
 }
 

shape.h

@@ -163,7 +163,6 @@ rb_shape_t* rb_shape_get_shape(VALUE obj);
 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_transition_shape_capa(rb_shape_t * shape);
 rb_shape_t* rb_shape_get_next(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);

variable.c

@@ -1442,45 +1442,34 @@ general_ivar_set(VALUE obj, ID id, VALUE val, void *data,
         .existing = true
     };
 
-    rb_shape_t *shape = rb_shape_get_shape(obj);
+    rb_shape_t *current_shape = rb_shape_get_shape(obj);
 
-    if (UNLIKELY(shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
+    if (UNLIKELY(current_shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
         goto too_complex;
     }
 
     attr_index_t index;
-    if (!rb_shape_get_iv_index(shape, id, &index)) {
+    if (!rb_shape_get_iv_index(current_shape, id, &index)) {
         result.existing = false;
 
-        index = shape->next_iv_index;
+        index = current_shape->next_iv_index;
         if (index >= MAX_IVARS) {
             rb_raise(rb_eArgError, "too many instance variables");
         }
 
-        if (UNLIKELY(shape->next_iv_index >= shape->capacity)) {
-            RUBY_ASSERT(shape->next_iv_index == shape->capacity);
-
-            rb_shape_t *next_shape = rb_shape_transition_shape_capa(shape);
-            if (next_shape->type == SHAPE_OBJ_TOO_COMPLEX) {
-                transition_too_complex_func(obj, data);
-                goto too_complex;
-            }
-
-            RUBY_ASSERT(next_shape->type == SHAPE_CAPACITY_CHANGE);
-            RUBY_ASSERT(next_shape->capacity > shape->capacity);
-            shape_resize_ivptr_func(obj, shape->capacity, next_shape->capacity, data);
-            shape = next_shape;
-        }
-
-        shape = rb_shape_get_next(shape, obj, id);
-        if (shape->type == SHAPE_OBJ_TOO_COMPLEX) {
+        rb_shape_t *next_shape = rb_shape_get_next(current_shape, obj, id);
+        if (UNLIKELY(next_shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
             transition_too_complex_func(obj, data);
             goto too_complex;
         }
+        else if (UNLIKELY(next_shape->capacity != current_shape->capacity)) {
+            RUBY_ASSERT(next_shape->capacity > current_shape->capacity);
+            shape_resize_ivptr_func(obj, current_shape->capacity, next_shape->capacity, data);
+        }
 
-        RUBY_ASSERT(shape->type == SHAPE_IVAR);
-        RUBY_ASSERT(index == (shape->next_iv_index - 1));
-        set_shape_func(obj, shape, data);
+        RUBY_ASSERT(next_shape->type == SHAPE_IVAR);
+        RUBY_ASSERT(index == (next_shape->next_iv_index - 1));
+        set_shape_func(obj, next_shape, data);
     }
 
     VALUE *table = shape_ivptr_func(obj, data);
@@ -1647,24 +1636,6 @@ rb_ensure_iv_list_size(VALUE obj, uint32_t current_capacity, uint32_t new_capaci
     }
 }
 
-// @note May raise when there are too many instance variables.
-rb_shape_t *
-rb_grow_iv_list(VALUE obj)
-{
-    rb_shape_t * initial_shape = rb_shape_get_shape(obj);
-    RUBY_ASSERT(initial_shape->capacity > 0);
-    rb_shape_t * res = rb_shape_transition_shape_capa(initial_shape);
-    if (res->type == SHAPE_OBJ_TOO_COMPLEX) { // Out of shapes
-        rb_evict_ivars_to_hash(obj, initial_shape);
-    }
-    else {
-        rb_ensure_iv_list_size(obj, initial_shape->capacity, res->capacity);
-        rb_shape_set_shape(obj, res);
-    }
-
-    return res;
-}
-
 int
 rb_obj_evacuate_ivs_to_hash_table(ID key, VALUE val, st_data_t arg)
 {

yjit/bindgen/src/main.rs

@@ -101,7 +101,6 @@ fn main() {
         .allowlist_function("rb_shape_get_iv_index")
         .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")

yjit/src/codegen.rs

@@ -2450,45 +2450,33 @@ fn gen_setinstancevariable(
         None
     };
 
-    // Get the next shape information if it needs transition
+    // 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 shape = comptime_receiver.shape_of();
-
-        let current_capacity = unsafe { (*shape).capacity };
-
-        // 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 >= current_capacity };
+        let current_shape = comptime_receiver.shape_of();
+        let next_shape = unsafe { rb_shape_get_next(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,
+        // it may be de-optimized into OBJ_TOO_COMPLEX_SHAPE (hash-table).
+        if next_shape_id == OBJ_TOO_COMPLEX_SHAPE_ID {
+            Some((next_shape_id, None, 0 as usize))
+        } else {
+            let current_capacity = unsafe { (*current_shape).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 the new shape has a different capacity, or is TOO_COMPLEX, we'll have to
+            // reallocate it.
+            let needs_extension = unsafe { (*current_shape).capacity != (*next_shape).capacity };
 
-        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
-            Some(unsafe { rb_shape_transition_shape_capa(shape) })
-        } else {
-            None
-        };
+            // We can write to the object, but we need to transition the shape
+            let ivar_index = unsafe { (*current_shape).next_iv_index } as usize;
 
-        let dest_shape = if let Some(capa_shape) = capa_shape {
-            if OBJ_TOO_COMPLEX_SHAPE_ID == unsafe { rb_shape_id(capa_shape) } {
-              capa_shape
+            let needs_extension = if needs_extension {
+                Some((current_capacity, unsafe { (*next_shape).capacity }))
             } else {
-              unsafe { rb_shape_get_next(capa_shape, comptime_receiver, ivar_name) }
-            }
-        } else {
-            unsafe { rb_shape_get_next(shape, comptime_receiver, ivar_name) }
-        };
-
-        let new_shape_id = unsafe { rb_shape_id(dest_shape) };
-        let needs_extension = if needs_extension {
-            Some((current_capacity, unsafe { (*dest_shape).capacity }))
-        } else {
-            None
-        };
-        Some((new_shape_id, needs_extension, ivar_index))
+                None
+            };
+            Some((next_shape_id, needs_extension, ivar_index))
+        }
     } else {
         None
     };

yjit/src/cruby_bindings.inc.rs

@@ -978,7 +978,6 @@ 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_transition_shape_capa(shape: *mut rb_shape_t) -> *mut rb_shape_t;
     pub fn rb_shape_get_next(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;