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.

3 files changed, 22 insertions, 36 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 359d5ccb9c..88b7f6b9c5 100644
--- a/yjit/bindgen/src/main.rs
+++ b/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")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 16f29ead8b..be6d38a466 100644
--- a/yjit/src/codegen.rs
+++ b/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
     };
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 83024e0b7e..431314ff86 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/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;