Refactor OBJ_TOO_COMPLEX_SHAPE_ID to not be referenced outside shape.h

Also refactor checks for `->type == SHAPE_OBJ_TOO_COMPLEX`.
ruby · byroot · May 8, 2025 · Apr 30, 2025 · May 5, 2025 · Apr 21, 2025
commit e6d3ed3554fba9dd47fd4cc4490213dfb31f046a
diff --git a/gc.c b/gc.c
@@ -379,9 +379,10 @@ rb_gc_set_shape(VALUE obj, uint32_t shape_id)
 uint32_t
 rb_gc_rebuild_shape(VALUE obj, size_t heap_id)
 {
-    rb_shape_t *orig_shape = rb_shape_get_shape(obj);
+    shape_id_t orig_shape_id = rb_shape_get_shape_id(obj);
+    rb_shape_t *orig_shape = rb_shape_get_shape_by_id(orig_shape_id);
 
-    if (rb_shape_obj_too_complex(obj)) return (uint32_t)OBJ_TOO_COMPLEX_SHAPE_ID;
+    if (rb_shape_too_complex_p(orig_shape)) return orig_shape_id;
 
     rb_shape_t *initial_shape = rb_shape_get_shape_by_id((shape_id_t)(heap_id + FIRST_T_OBJECT_SHAPE_ID));
     rb_shape_t *new_shape = rb_shape_traverse_from_new_root(initial_shape, orig_shape);

diff --git a/object.c b/object.c
@@ -362,7 +362,7 @@ rb_obj_copy_ivar(VALUE dest, VALUE obj)
         RUBY_ASSERT(initial_shape->type == SHAPE_T_OBJECT);
 
         shape_to_set_on_dest = rb_shape_rebuild_shape(initial_shape, src_shape);
-        if (UNLIKELY(rb_shape_id(shape_to_set_on_dest) == OBJ_TOO_COMPLEX_SHAPE_ID)) {
+        if (UNLIKELY(rb_shape_too_complex_p(shape_to_set_on_dest))) {
             st_table * table = rb_st_init_numtable_with_size(src_num_ivs);
             rb_obj_copy_ivs_to_hash_table(obj, table);
             rb_obj_convert_to_too_complex(dest, table);
@@ -371,7 +371,7 @@ rb_obj_copy_ivar(VALUE dest, VALUE obj)
         }
     }
 
-    RUBY_ASSERT(src_num_ivs <= shape_to_set_on_dest->capacity || rb_shape_id(shape_to_set_on_dest) == OBJ_TOO_COMPLEX_SHAPE_ID);
+    RUBY_ASSERT(src_num_ivs <= shape_to_set_on_dest->capacity || rb_shape_too_complex_p(shape_to_set_on_dest));
     if (initial_shape->capacity < shape_to_set_on_dest->capacity) {
         rb_ensure_iv_list_size(dest, initial_shape->capacity, shape_to_set_on_dest->capacity);
         dest_buf = ROBJECT_FIELDS(dest);
@@ -507,7 +507,7 @@ rb_obj_clone_setup(VALUE obj, VALUE clone, VALUE kwfreeze)
 
         if (RB_OBJ_FROZEN(obj)) {
             rb_shape_t *next_shape = rb_shape_transition_shape_frozen(clone);
-            if (!rb_shape_obj_too_complex(clone) && next_shape->type == SHAPE_OBJ_TOO_COMPLEX) {
+            if (!rb_shape_obj_too_complex(clone) && rb_shape_too_complex_p(next_shape)) {
                 rb_evict_ivars_to_hash(clone);
             }
             else {
@@ -531,7 +531,7 @@ rb_obj_clone_setup(VALUE obj, VALUE clone, VALUE kwfreeze)
         rb_shape_t *next_shape = rb_shape_transition_shape_frozen(clone);
         // If we're out of shapes, but we want to freeze, then we need to
         // evacuate this clone to a hash
-        if (!rb_shape_obj_too_complex(clone) && next_shape->type == SHAPE_OBJ_TOO_COMPLEX) {
+        if (!rb_shape_obj_too_complex(clone) && rb_shape_too_complex_p(next_shape)) {
             rb_evict_ivars_to_hash(clone);
         }
         else {

diff --git a/shape.c b/shape.c
@@ -495,7 +495,7 @@ get_next_shape_internal(rb_shape_t *shape, ID id, enum shape_type shape_type, bo
     rb_shape_t *res = NULL;
 
     // There should never be outgoing edges from "too complex"
-    RUBY_ASSERT(rb_shape_id(shape) != OBJ_TOO_COMPLEX_SHAPE_ID);
+    RUBY_ASSERT(!rb_shape_too_complex_p(shape));
 
     *variation_created = false;
 
@@ -573,7 +573,7 @@ get_next_shape_internal(rb_shape_t *shape, ID id, enum shape_type shape_type, bo
     return res;
 }
 
-int
+bool
 rb_shape_frozen_shape_p(rb_shape_t *shape)
 {
     return SHAPE_FROZEN == (enum shape_type)shape->type;
@@ -703,6 +703,11 @@ rb_shape_transition_shape_frozen(VALUE obj)
     return next_shape;
 }
 
+rb_shape_t *
+rb_shape_transition_shape_too_complex(VALUE obj)
+{
+    return rb_shape_get_shape_by_id(OBJ_TOO_COMPLEX_SHAPE_ID);
+}
 /*
  * This function is used for assertions where we don't want to increment
  * max_iv_count
@@ -1012,6 +1017,18 @@ rb_shape_obj_too_complex(VALUE obj)
     return rb_shape_get_shape_id(obj) == OBJ_TOO_COMPLEX_SHAPE_ID;
 }
 
+bool
+rb_shape_too_complex_p(rb_shape_t *shape)
+{
+    return rb_shape_id(shape) == OBJ_TOO_COMPLEX_SHAPE_ID;
+}
+
+bool
+rb_shape_id_too_complex_p(shape_id_t shape_id)
+{
+    return shape_id == OBJ_TOO_COMPLEX_SHAPE_ID;
+}
+
 size_t
 rb_shape_edges_count(rb_shape_t *shape)
 {

diff --git a/shape.h b/shape.h
@@ -157,11 +157,14 @@ rb_shape_t *rb_shape_get_next_iv_shape(rb_shape_t *shape, ID id);
 bool rb_shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value);
 bool rb_shape_get_iv_index_with_hint(shape_id_t shape_id, ID id, attr_index_t *value, shape_id_t *shape_id_hint);
 RUBY_FUNC_EXPORTED bool rb_shape_obj_too_complex(VALUE obj);
+bool rb_shape_too_complex_p(rb_shape_t *shape);
+bool rb_shape_id_too_complex_p(shape_id_t shape_id);
 
 void rb_shape_set_shape(VALUE obj, rb_shape_t *shape);
 rb_shape_t *rb_shape_get_shape(VALUE obj);
-int rb_shape_frozen_shape_p(rb_shape_t *shape);
+bool rb_shape_frozen_shape_p(rb_shape_t *shape);
 rb_shape_t *rb_shape_transition_shape_frozen(VALUE obj);
+rb_shape_t *rb_shape_transition_shape_too_complex(VALUE obj);
 bool rb_shape_transition_shape_remove_ivar(VALUE obj, ID id, rb_shape_t *shape, VALUE *removed);
 rb_shape_t *rb_shape_get_next(rb_shape_t *shape, VALUE obj, ID id);
 rb_shape_t *rb_shape_get_next_no_warnings(rb_shape_t *shape, VALUE obj, ID id);

diff --git a/variable.c b/variable.c
@@ -1482,6 +1482,7 @@ void
 rb_obj_convert_to_too_complex(VALUE obj, st_table *table)
 {
     RUBY_ASSERT(!rb_shape_obj_too_complex(obj));
+    rb_shape_t *too_complex_shape = rb_shape_transition_shape_too_complex(obj);
 
     VALUE *old_fields = NULL;
 
@@ -1490,13 +1491,13 @@ rb_obj_convert_to_too_complex(VALUE obj, st_table *table)
         if (!(RBASIC(obj)->flags & ROBJECT_EMBED)) {
             old_fields = ROBJECT_FIELDS(obj);
         }
-        rb_shape_set_shape_id(obj, OBJ_TOO_COMPLEX_SHAPE_ID);
+        rb_shape_set_shape(obj, too_complex_shape);
         ROBJECT_SET_FIELDS_HASH(obj, table);
         break;
       case T_CLASS:
       case T_MODULE:
         old_fields = RCLASS_FIELDS(obj);
-        rb_shape_set_shape_id(obj, OBJ_TOO_COMPLEX_SHAPE_ID);
+        rb_shape_set_shape(obj, too_complex_shape);
         RCLASS_SET_FIELDS_HASH(obj, table);
         break;
       default:
@@ -1513,9 +1514,9 @@ rb_obj_convert_to_too_complex(VALUE obj, st_table *table)
                  * compaction. We want the table to be updated rather than
                  * the original fields. */
 #if SHAPE_IN_BASIC_FLAGS
-                rb_shape_set_shape_id(obj, OBJ_TOO_COMPLEX_SHAPE_ID);
+                rb_shape_set_shape(obj, too_complex_shape);
 #else
-                old_fields_tbl->shape_id = OBJ_TOO_COMPLEX_SHAPE_ID;
+                old_fields_tbl->shape_id = rb_shape_id(too_complex_shape);
 #endif
                 old_fields_tbl->as.complex.table = table;
                 old_fields = (VALUE *)old_fields_tbl;
@@ -1524,10 +1525,11 @@ rb_obj_convert_to_too_complex(VALUE obj, st_table *table)
             struct gen_fields_tbl *fields_tbl = xmalloc(sizeof(struct gen_fields_tbl));
             fields_tbl->as.complex.table = table;
             st_insert(gen_ivs, (st_data_t)obj, (st_data_t)fields_tbl);
+
 #if SHAPE_IN_BASIC_FLAGS
-            rb_shape_set_shape_id(obj, OBJ_TOO_COMPLEX_SHAPE_ID);
+            rb_shape_set_shape(obj, too_complex_shape);
 #else
-            fields_tbl->shape_id = OBJ_TOO_COMPLEX_SHAPE_ID;
+            fields_tbl->shape_id = rb_shape_id(too_complex_shape);
 #endif
         }
         RB_VM_LOCK_LEAVE();
@@ -1570,7 +1572,7 @@ general_ivar_set(VALUE obj, ID id, VALUE val, void *data,
 
     rb_shape_t *current_shape = rb_shape_get_shape(obj);
 
-    if (UNLIKELY(current_shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
+    if (UNLIKELY(rb_shape_too_complex_p(current_shape))) {
         goto too_complex;
     }
 
@@ -1584,7 +1586,7 @@ general_ivar_set(VALUE obj, ID id, VALUE val, void *data,
         }
 
         rb_shape_t *next_shape = rb_shape_get_next(current_shape, obj, id);
-        if (UNLIKELY(next_shape->type == SHAPE_OBJ_TOO_COMPLEX)) {
+        if (UNLIKELY(rb_shape_too_complex_p(next_shape))) {
             transition_too_complex_func(obj, data);
             goto too_complex;
         }
@@ -1709,7 +1711,7 @@ generic_ivar_set_too_complex_table(VALUE obj, void *data)
     if (!rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
         fields_tbl = xmalloc(sizeof(struct gen_fields_tbl));
 #if !SHAPE_IN_BASIC_FLAGS
-        fields_tbl->shape_id = OBJ_TOO_COMPLEX_SHAPE_ID;
+        fields_tbl->shape_id = rb_shape_id(rb_shape_transition_shape_too_complex(obj));
 #endif
         fields_tbl->as.complex.table = st_init_numtable_with_size(1);
 
@@ -1886,7 +1888,7 @@ void rb_obj_freeze_inline(VALUE x)
 
         // If we're transitioning from "not complex" to "too complex"
         // then evict ivars.  This can happen if we run out of shapes
-        if (!rb_shape_obj_too_complex(x) && next_shape->type == SHAPE_OBJ_TOO_COMPLEX) {
+        if (!rb_shape_obj_too_complex(x) && rb_shape_too_complex_p(next_shape)) {
             rb_evict_ivars_to_hash(x);
         }
         rb_shape_set_shape(x, next_shape);
@@ -2029,7 +2031,6 @@ iterate_over_shapes_with_callback(rb_shape_t *shape, rb_ivar_foreach_callback_fu
       case SHAPE_FROZEN:
         return iterate_over_shapes_with_callback(rb_shape_get_parent(shape), callback, itr_data);
       case SHAPE_OBJ_TOO_COMPLEX:
-      default:
         rb_bug("Unreachable");
     }
 }
@@ -2117,7 +2118,7 @@ rb_copy_generic_ivar(VALUE clone, VALUE obj)
         if (rb_shape_obj_too_complex(obj)) {
             new_fields_tbl = xmalloc(sizeof(struct gen_fields_tbl));
 #if !SHAPE_IN_BASIC_FLAGS
-            new_fields_tbl->shape_id = OBJ_TOO_COMPLEX_SHAPE_ID;
+            new_fields_tbl->shape_id = old_fields_tbl->shape_id;
 #endif
             new_fields_tbl->as.complex.table = st_copy(obj_fields_tbl->as.complex.table);
         }
@@ -2140,7 +2141,7 @@ rb_copy_generic_ivar(VALUE clone, VALUE obj)
         }
         RB_VM_LOCK_LEAVE();
 
-        rb_shape_t * obj_shape = rb_shape_get_shape(obj);
+        rb_shape_t *obj_shape = rb_shape_get_shape(obj);
         if (rb_shape_frozen_shape_p(obj_shape)) {
             rb_shape_set_shape_id(clone, obj_shape->parent_id);
         }

diff --git a/vm_insnhelper.c b/vm_insnhelper.c
@@ -1289,7 +1289,7 @@ vm_getivar(VALUE obj, ID id, const rb_iseq_t *iseq, IVC ic, const struct rb_call
     }
 
     if (LIKELY(cached_id == shape_id)) {
-        RUBY_ASSERT(cached_id != OBJ_TOO_COMPLEX_SHAPE_ID);
+        RUBY_ASSERT(!rb_shape_id_too_complex_p(cached_id));
 
         if (index == ATTR_INDEX_NOT_SET) {
             return default_value;
@@ -1330,7 +1330,7 @@ vm_getivar(VALUE obj, ID id, const rb_iseq_t *iseq, IVC ic, const struct rb_call
         }
 #endif
 
-        if (shape_id == OBJ_TOO_COMPLEX_SHAPE_ID) {
+        if (rb_shape_id_too_complex_p(shape_id)) {
             st_table *table = NULL;
             switch (BUILTIN_TYPE(obj)) {
               case T_CLASS:
@@ -1408,7 +1408,7 @@ vm_getivar(VALUE obj, ID id, const rb_iseq_t *iseq, IVC ic, const struct rb_call
 static void
 populate_cache(attr_index_t index, shape_id_t next_shape_id, ID id, const rb_iseq_t *iseq, IVC ic, const struct rb_callcache *cc, bool is_attr)
 {
-    RUBY_ASSERT(next_shape_id != OBJ_TOO_COMPLEX_SHAPE_ID);
+    RUBY_ASSERT(!rb_shape_id_too_complex_p(next_shape_id));
 
     // Cache population code
     if (is_attr) {
@@ -1436,7 +1436,7 @@ vm_setivar_slowpath(VALUE obj, ID id, VALUE val, const rb_iseq_t *iseq, IVC ic,
 
         shape_id_t next_shape_id = ROBJECT_SHAPE_ID(obj);
 
-        if (next_shape_id != OBJ_TOO_COMPLEX_SHAPE_ID) {
+        if (!rb_shape_id_too_complex_p(next_shape_id)) {
             populate_cache(index, next_shape_id, id, iseq, ic, cc, is_attr);
         }
 
@@ -1517,7 +1517,7 @@ vm_setivar(VALUE obj, ID id, VALUE val, shape_id_t dest_shape_id, attr_index_t i
             VM_ASSERT(!rb_ractor_shareable_p(obj) || rb_obj_frozen_p(obj));
 
             shape_id_t shape_id = ROBJECT_SHAPE_ID(obj);
-            RUBY_ASSERT(dest_shape_id != OBJ_TOO_COMPLEX_SHAPE_ID);
+            RUBY_ASSERT(!rb_shape_id_too_complex_p(dest_shape_id));
 
             if (LIKELY(shape_id == dest_shape_id)) {
                 RUBY_ASSERT(dest_shape_id != INVALID_SHAPE_ID && shape_id != INVALID_SHAPE_ID);

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
@@ -101,8 +101,8 @@ fn main() {
         .allowlist_function("rb_shape_get_next_no_warnings")
         .allowlist_function("rb_shape_id")
         .allowlist_function("rb_shape_obj_too_complex")
+        .allowlist_function("rb_shape_too_complex_p")
         .allowlist_var("SHAPE_ID_NUM_BITS")
-        .allowlist_var("OBJ_TOO_COMPLEX_SHAPE_ID")
 
         // From ruby/internal/intern/object.h
         .allowlist_function("rb_obj_is_kind_of")

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -3109,14 +3109,16 @@ fn gen_set_ivar(
     };
 
     // The current shape doesn't contain this iv, we need to transition to another shape.
+    let mut new_shape_too_complex = false;
     let new_shape = if !shape_too_complex && receiver_t_object && ivar_index.is_none() {
         let current_shape = comptime_receiver.shape_of();
         let next_shape = unsafe { rb_shape_get_next_no_warnings(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 {
+        new_shape_too_complex = unsafe { rb_shape_too_complex_p(next_shape) };
+        if new_shape_too_complex {
             Some((next_shape_id, None, 0_usize))
         } else {
             let current_capacity = unsafe { (*current_shape).capacity };
@@ -3138,7 +3140,6 @@ fn gen_set_ivar(
     } else {
         None
     };
-    let new_shape_too_complex = matches!(new_shape, Some((OBJ_TOO_COMPLEX_SHAPE_ID, _, _)));
 
     // If the receiver isn't a T_OBJECT, or uses a custom allocator,
     // then just write out the IV write as a function call.

diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
@@ -115,7 +115,6 @@ fn main() {
         .allowlist_function("rb_shape_id")
         .allowlist_function("rb_shape_obj_too_complex")
         .allowlist_var("SHAPE_ID_NUM_BITS")
-        .allowlist_var("OBJ_TOO_COMPLEX_SHAPE_ID")
 
         // From ruby/internal/intern/object.h
         .allowlist_function("rb_obj_is_kind_of")

diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs