More precisely iterate over Object instance variables

Shapes provides us with an (almost) exact count of instance variables.
We only need to check for Qundef when an IV has been "undefined"
Prefer to use ROBJECT_IV_COUNT when iterating IVs

gc.c

@@ -7270,7 +7270,7 @@ gc_mark_children(rb_objspace_t *objspace, VALUE obj)
         {
             const VALUE * const ptr = ROBJECT_IVPTR(obj);
 
-            uint32_t i, len = ROBJECT_NUMIV(obj);
+            uint32_t i, len = ROBJECT_IV_COUNT(obj);
             for (i  = 0; i < len; i++) {
                 gc_mark(objspace, ptr[i]);
             }
@@ -10019,7 +10019,7 @@ gc_ref_update_object(rb_objspace_t *objspace, VALUE v)
     }
 #endif
 
-    for (uint32_t i = 0; i < numiv; i++) {
+    for (uint32_t i = 0; i < ROBJECT_IV_COUNT(v); i++) {
         UPDATE_IF_MOVED(objspace, ptr[i]);
     }
 }

ractor.c

@@ -2312,7 +2312,7 @@ obj_traverse_i(VALUE obj, struct obj_traverse_data *data)
 
       case T_OBJECT:
         {
-            uint32_t len = ROBJECT_NUMIV(obj);
+            uint32_t len = ROBJECT_IV_COUNT(obj);
             VALUE *ptr = ROBJECT_IVPTR(obj);
 
             for (uint32_t i=0; i<len; i++) {
@@ -2766,7 +2766,7 @@ obj_traverse_replace_i(VALUE obj, struct obj_traverse_replace_data *data)
             if (data->move) rb_obj_transient_heap_evacuate(obj, TRUE);
 #endif
 
-            uint32_t len = ROBJECT_NUMIV(obj);
+            uint32_t len = ROBJECT_IV_COUNT(obj);
             VALUE *ptr = ROBJECT_IVPTR(obj);
 
             for (uint32_t i=0; i<len; i++) {

shape.h

@@ -125,6 +125,15 @@ bool rb_shape_get_iv_index(rb_shape_t * shape, ID id, attr_index_t * value);
 shape_id_t rb_shape_id(rb_shape_t * shape);
 MJIT_SYMBOL_EXPORT_END
 
+static inline uint32_t
+ROBJECT_IV_COUNT(VALUE obj)
+{
+    RBIMPL_ASSERT_TYPE(obj, RUBY_T_OBJECT);
+    uint32_t ivc = rb_shape_get_shape_by_id(ROBJECT_SHAPE_ID(obj))->iv_count;
+    RUBY_ASSERT(ivc <= ROBJECT_NUMIV(obj));
+    return ivc;
+}
+
 rb_shape_t * rb_shape_alloc(ID edge_name, rb_shape_t * parent);
 rb_shape_t * rb_shape_alloc_with_parent_id(ID edge_name, shape_id_t parent_id);
 

variable.c

@@ -1443,43 +1443,32 @@ rb_init_iv_list(VALUE obj)
     rb_ensure_iv_list_size(obj, len, newsize < len ? len : newsize);
 }
 
-// Return the instance variable index for a given name and T_OBJECT object. The
-// mapping between name and index lives on `rb_obj_class(obj)` and is created
-// if not already present.
-//
-// @note May raise when there are too many instance variables.
-// @note YJIT uses this function at compile time to simplify the work needed to
-//       access the variable at runtime.
-static uint32_t
-rb_obj_ensure_iv_index_mapping(VALUE obj, ID id)
+static VALUE
+obj_ivar_set(VALUE obj, ID id, VALUE val)
 {
-    RUBY_ASSERT(RB_TYPE_P(obj, T_OBJECT));
     attr_index_t index;
 
-    // Ensure there is a transition for IVAR +id+
-    rb_shape_transition_shape(obj, id, rb_shape_get_shape_by_id(ROBJECT_SHAPE_ID(obj)));
-
     // Get the current shape
     rb_shape_t * shape = rb_shape_get_shape_by_id(ROBJECT_SHAPE_ID(obj));
 
     if (!rb_shape_get_iv_index(shape, id, &index)) {
-        rb_bug("unreachable.  Shape was not found for id: %s", rb_id2name(id));
+        shape = rb_shape_get_next(shape, obj, id);
+        index = shape->iv_count - 1;
     }
 
     uint32_t len = ROBJECT_NUMIV(obj);
+
+    // Reallocating can kick off GC.  We can't set the new shape
+    // on this object until the buffer has been allocated, otherwise
+    // GC could read off the end of the buffer.
     if (len <= index) {
-        uint32_t newsize = (uint32_t)((shape->iv_count + 1) * 1.25);
+        uint32_t newsize = (uint32_t)((len + 1) * 1.25);
         rb_ensure_iv_list_size(obj, len, newsize);
     }
-    RUBY_ASSERT(index <= ROBJECT_NUMIV(obj));
-    return index;
-}
 
-static VALUE
-obj_ivar_set(VALUE obj, ID id, VALUE val)
-{
-    attr_index_t index = rb_obj_ensure_iv_index_mapping(obj, id);
     RB_OBJ_WRITE(obj, &ROBJECT_IVPTR(obj)[index], val);
+    rb_shape_set_shape(obj, shape);
+
     return val;
 }
 
@@ -1768,7 +1757,7 @@ rb_ivar_count(VALUE obj)
     switch (BUILTIN_TYPE(obj)) {
       case T_OBJECT:
         if (rb_shape_get_shape(obj)->iv_count > 0) {
-            st_index_t i, count, num = ROBJECT_NUMIV(obj);
+            st_index_t i, count, num = ROBJECT_IV_COUNT(obj);
             const VALUE *const ivptr = ROBJECT_IVPTR(obj);
             for (i = count = 0; i < num; ++i) {
                 if (ivptr[i] != Qundef) {

vm_insnhelper.c

@@ -1268,8 +1268,7 @@ vm_setivar_slowpath(VALUE obj, ID id, VALUE val, const rb_iseq_t *iseq, IVC ic,
 
             if (shape != next_shape) {
                 RUBY_ASSERT(next_shape->parent_id == rb_shape_id(shape));
-                rb_shape_set_shape(obj, next_shape);
-                next_shape_id = ROBJECT_SHAPE_ID(obj);
+                next_shape_id = rb_shape_id(next_shape);
             }
 
             if (rb_shape_get_iv_index(next_shape, id, &index)) { // based off the hash stored in the transition tree
@@ -1289,6 +1288,9 @@ vm_setivar_slowpath(VALUE obj, ID id, VALUE val, const rb_iseq_t *iseq, IVC ic,
                 rb_init_iv_list(obj);
             }
 
+            if (shape != next_shape) {
+                rb_shape_set_shape(obj, next_shape);
+            }
             VALUE *ptr = ROBJECT_IVPTR(obj);
             RB_OBJ_WRITE(obj, &ptr[index], val);
             RB_DEBUG_COUNTER_INC(ivar_set_ic_miss_iv_hit);