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.

1 files changed, 16 insertions, 1 deletions

diff --git a/shape.c b/shape.c
index 8fae6fe7ab..ef08030ad4 100644
--- a/shape.c
+++ b/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)));
 }