7 files changed, 21 insertions, 24 deletions
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 96dfa55c69..12f9a6c505 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -502,8 +502,6 @@ module RubyVM::RJIT
           shape = C.rb_shape_get_shape_by_id(shape_id)
 
           current_capacity = shape.capacity
-          new_capacity = current_capacity * 2
-
           # 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
@@ -515,7 +513,7 @@ module RubyVM::RJIT
             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, new_capacity)
+              C.rb_shape_transition_shape_capa(shape)
             else
               nil
             end
@@ -538,7 +536,7 @@ module RubyVM::RJIT
             # the capacity and set the buffer.
             asm.mov(C_ARGS[0], :rax)
             asm.mov(C_ARGS[1], current_capacity)
-            asm.mov(C_ARGS[2], new_capacity)
+            asm.mov(C_ARGS[2], capa_shape.capacity)
             asm.call(C.rb_ensure_iv_list_size)
 
             # Load the receiver again after the function call
diff --git a/rjit_c.rb b/rjit_c.rb
index ef99df211a..72e8221cbc 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -171,9 +171,9 @@ module RubyVM::RJIT # :nodoc: all
       me_addr == 0 ? nil : rb_method_entry_t.new(me_addr)
     end
 
-    def rb_shape_transition_shape_capa(shape, new_capacity)
+    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), NUM2UINT(new_capacity)))'
+      shape_addr = Primitive.cexpr! 'SIZET2NUM((size_t)rb_shape_transition_shape_capa((rb_shape_t *)NUM2SIZET(_shape)))'
       rb_shape_t.new(shape_addr)
     end
 
diff --git a/shape.c b/shape.c
index dcdb9d28aa..89a2c3bd0b 100644
--- a/shape.c
+++ b/shape.c
@@ -417,8 +417,8 @@ rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id)
     return new_shape;
 }
 
-rb_shape_t *
-rb_shape_transition_shape_capa(rb_shape_t* shape, uint32_t new_capacity)
+static inline rb_shape_t *
+rb_shape_transition_shape_capa_create(rb_shape_t* shape, uint32_t new_capacity)
 {
     ID edge_name = rb_make_temporary_id(new_capacity);
     bool dont_care;
@@ -427,6 +427,12 @@ rb_shape_transition_shape_capa(rb_shape_t* shape, uint32_t new_capacity)
     return new_shape;
 }
 
+rb_shape_t *
+rb_shape_transition_shape_capa(rb_shape_t* shape)
+{
+    return rb_shape_transition_shape_capa_create(shape, shape->capacity * 2);
+}
+
 bool
 rb_shape_get_iv_index(rb_shape_t * shape, ID id, attr_index_t *value)
 {
@@ -541,7 +547,7 @@ rb_shape_rebuild_shape(rb_shape_t * initial_shape, rb_shape_t * dest_shape)
       case SHAPE_IVAR:
         if (midway_shape->capacity <= midway_shape->next_iv_index) {
             // There isn't enough room to write this IV, so we need to increase the capacity
-            midway_shape = rb_shape_transition_shape_capa(midway_shape, midway_shape->capacity * 2);
+            midway_shape = rb_shape_transition_shape_capa(midway_shape);
         }
 
         midway_shape = rb_shape_get_next_iv_shape(midway_shape, dest_shape->edge_name);
@@ -828,7 +834,7 @@ Init_default_shapes(void)
     // Shapes by size pool
     for (int i = 1; i < SIZE_POOL_COUNT; i++) {
         uint32_t capa = (uint32_t)((rb_size_pool_slot_size(i) - offsetof(struct RObject, as.ary)) / sizeof(VALUE));
-        rb_shape_t * new_shape = rb_shape_transition_shape_capa(root, capa);
+        rb_shape_t * new_shape = rb_shape_transition_shape_capa_create(root, capa);
         new_shape->type = SHAPE_INITIAL_CAPACITY;
         new_shape->size_pool_index = i;
         RUBY_ASSERT(rb_shape_id(new_shape) == (shape_id_t)i);
diff --git a/shape.h b/shape.h
index 025dfb4dd8..5689b48067 100644
--- a/shape.h
+++ b/shape.h
@@ -152,7 +152,7 @@ rb_shape_t* rb_shape_get_shape(VALUE obj);
 int rb_shape_frozen_shape_p(rb_shape_t* shape);
 void rb_shape_transition_shape_frozen(VALUE obj);
 void 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, uint32_t new_capacity);
+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);
diff --git a/variable.c b/variable.c
index c30971bd08..15dc4831d8 100644
--- a/variable.c
+++ b/variable.c
@@ -1444,13 +1444,10 @@ rb_shape_t *
 rb_grow_iv_list(VALUE obj)
 {
     rb_shape_t * initial_shape = rb_shape_get_shape(obj);
-    uint32_t len = initial_shape->capacity;
-    RUBY_ASSERT(len > 0);
-    uint32_t newsize = (uint32_t)(len * 2);
+    RUBY_ASSERT(initial_shape->capacity > 0);
+    rb_shape_t * res = rb_shape_transition_shape_capa(initial_shape);
 
-    rb_shape_t * res = rb_shape_transition_shape_capa(initial_shape, newsize);
-
-    rb_ensure_iv_list_size(obj, len, newsize);
+    rb_ensure_iv_list_size(obj, initial_shape->capacity, res->capacity);
 
     rb_shape_set_shape(obj, res);
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index e4fa1a3665..5dff45ba8d 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -2403,7 +2403,6 @@ fn gen_setinstancevariable(
         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.
@@ -2416,7 +2415,7 @@ fn gen_setinstancevariable(
             // 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, new_capacity) })
+            Some(unsafe { rb_shape_transition_shape_capa(shape) })
         } else {
             None
         };
@@ -2429,7 +2428,7 @@ fn gen_setinstancevariable(
 
         let new_shape_id = unsafe { rb_shape_id(dest_shape) };
         let needs_extension = if needs_extension {
-            Some((current_capacity, new_capacity))
+            Some((current_capacity, unsafe { (*dest_shape).capacity }))
         } else {
             None
         };
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index bf554a9f85..fde68a80b5 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1179,10 +1179,7 @@ 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,
-        new_capacity: u32,
-    ) -> *mut rb_shape_t;
+    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;