Get rid of TOO_COMPLEX shape type

Instead it's now a `shape_id` flag.

This allows to check if an object is complex without having
to chase the `rb_shape_t` pointer.

Author: byroot

ext/objspace/objspace_dump.c

@@ -821,9 +821,6 @@ shape_id_i(shape_id_t shape_id, void *data)
       case SHAPE_T_OBJECT:
         dump_append(dc, ", \"shape_type\":\"T_OBJECT\"");
         break;
-      case SHAPE_OBJ_TOO_COMPLEX:
-        dump_append(dc, ", \"shape_type\":\"OBJ_TOO_COMPLEX\"");
-        break;
       case SHAPE_OBJ_ID:
         dump_append(dc, ", \"shape_type\":\"OBJ_ID\"");
         break;

shape.c

@@ -379,12 +379,6 @@ shape_frozen_p(shape_id_t shape_id)
 }
 #endif
 
-static inline bool
-shape_too_complex_p(rb_shape_t *shape)
-{
-    return shape->flags & SHAPE_FL_TOO_COMPLEX;
-}
-
 void
 rb_shape_each_shape_id(each_shape_callback callback, void *data)
 {
@@ -529,7 +523,6 @@ rb_shape_alloc_new_child(ID id, rb_shape_t *shape, enum shape_type shape_type)
             redblack_cache_ancestors(new_shape);
         }
         break;
-      case SHAPE_OBJ_TOO_COMPLEX:
       case SHAPE_ROOT:
       case SHAPE_T_OBJECT:
         rb_bug("Unreachable");
@@ -619,9 +612,6 @@ get_next_shape_internal_atomic(rb_shape_t *shape, ID id, enum shape_type shape_t
 static rb_shape_t *
 get_next_shape_internal(rb_shape_t *shape, ID id, enum shape_type shape_type, bool *variation_created, bool new_variations_allowed)
 {
-    // There should never be outgoing edges from "too complex", except for SHAPE_OBJ_ID
-    RUBY_ASSERT(!shape_too_complex_p(shape) || shape_type == SHAPE_OBJ_ID);
-
     if (rb_multi_ractor_p()) {
         return get_next_shape_internal_atomic(shape, id, shape_type, variation_created, new_variations_allowed);
     }
@@ -691,6 +681,7 @@ remove_shape_recursive(rb_shape_t *shape, ID id, rb_shape_t **removed_shape)
     if (shape->parent_id == INVALID_SHAPE_ID) {
         // We've hit the top of the shape tree and couldn't find the
         // IV we wanted to remove, so return NULL
+        *removed_shape = NULL;
         return NULL;
     }
     else {
@@ -706,23 +697,14 @@ remove_shape_recursive(rb_shape_t *shape, ID id, rb_shape_t **removed_shape)
             // We found a new parent.  Create a child of the new parent that
             // has the same attributes as this shape.
             if (new_parent) {
-                if (UNLIKELY(shape_too_complex_p(new_parent))) {
-                    return new_parent;
-                }
-
                 bool dont_care;
                 rb_shape_t *new_child = get_next_shape_internal(new_parent, shape->edge_name, shape->type, &dont_care, true);
-                if (UNLIKELY(!new_child)) {
-                    return new_child;
-                }
-
-                RUBY_ASSERT(new_child->capacity <= shape->capacity);
-
+                RUBY_ASSERT(!new_child || new_child->capacity <= shape->capacity);
                 return new_child;
             }
             else {
                 // We went all the way to the top of the shape tree and couldn't
-                // find an IV to remove, so return NULL
+                // find an IV to remove so return NULL.
                 return NULL;
             }
         }
@@ -732,19 +714,27 @@ remove_shape_recursive(rb_shape_t *shape, ID id, rb_shape_t **removed_shape)
 shape_id_t
 rb_shape_transition_remove_ivar(VALUE obj, ID id, shape_id_t *removed_shape_id)
 {
-    shape_id_t shape_id = rb_obj_shape_id(obj);
-    rb_shape_t *shape = RSHAPE(shape_id);
+    shape_id_t original_shape_id = RBASIC_SHAPE_ID(obj);
 
-    RUBY_ASSERT(!shape_too_complex_p(shape));
-    RUBY_ASSERT(!shape_frozen_p(shape_id));
+    RUBY_ASSERT(!rb_shape_too_complex_p(original_shape_id));
+    RUBY_ASSERT(!shape_frozen_p(original_shape_id));
 
     rb_shape_t *removed_shape = NULL;
-    rb_shape_t *new_shape = remove_shape_recursive(shape, id, &removed_shape);
-    if (new_shape) {
+    rb_shape_t *new_shape = remove_shape_recursive(RSHAPE(original_shape_id), id, &removed_shape);
+
+    if (removed_shape) {
         *removed_shape_id = raw_shape_id(removed_shape);
-        return raw_shape_id(new_shape);
     }
-    return shape_id;
+
+    if (new_shape) {
+        return shape_id(new_shape, original_shape_id);
+    }
+    else if (removed_shape) {
+        // We found the shape to remove, but couldn't create a new variation.
+        // We must transition to TOO_COMPLEX.
+        return ROOT_TOO_COMPLEX_SHAPE_ID | (original_shape_id & SHAPE_ID_FLAGS_MASK);
+    }
+    return original_shape_id;
 }
 
 shape_id_t
@@ -760,7 +750,7 @@ shape_id_t
 rb_shape_transition_complex(VALUE obj)
 {
     shape_id_t original_shape_id = RBASIC_SHAPE_ID(obj);
-    return shape_id(ROOT_SHAPE_ID | SHAPE_ID_FL_TOO_COMPLEX, original_shape_id);
+    return ROOT_TOO_COMPLEX_SHAPE_ID | (original_shape_id & SHAPE_ID_FLAGS_MASK);
 }
 
 static inline bool
@@ -832,7 +822,6 @@ shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
               case SHAPE_ROOT:
               case SHAPE_T_OBJECT:
                 return false;
-              case SHAPE_OBJ_TOO_COMPLEX:
               case SHAPE_OBJ_ID:
                 rb_bug("Ivar should not exist on transition");
             }
@@ -848,9 +837,6 @@ static inline rb_shape_t *
 shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
 {
     RUBY_ASSERT(!is_instance_id(id) || RTEST(rb_sym2str(ID2SYM(id))));
-    if (UNLIKELY(shape_too_complex_p(shape))) {
-        return shape;
-    }
 
 #if RUBY_DEBUG
     attr_index_t index;
@@ -874,6 +860,11 @@ shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
     bool variation_created = false;
     rb_shape_t *new_shape = get_next_shape_internal(shape, id, SHAPE_IVAR, &variation_created, allow_new_shape);
 
+    if (!new_shape) {
+        // We could create a new variation, transitioning to TOO_COMPLEX.
+        return NULL;
+    }
+
     // Check if we should update max_iv_count on the object's class
     if (obj != klass && new_shape->next_field_index > RCLASS_MAX_IV_COUNT(klass)) {
         RCLASS_SET_MAX_IV_COUNT(klass, new_shape->next_field_index);
@@ -999,11 +990,11 @@ shape_cache_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
 bool
 rb_shape_get_iv_index(shape_id_t shape_id, ID id, attr_index_t *value)
 {
-    rb_shape_t *shape = RSHAPE(shape_id);
-
     // It doesn't make sense to ask for the index of an IV that's stored
     // on an object that is "too complex" as it uses a hash for storing IVs
-    RUBY_ASSERT(!shape_too_complex_p(shape));
+    RUBY_ASSERT(!rb_shape_too_complex_p(shape_id));
+
+    rb_shape_t *shape = RSHAPE(shape_id);
 
     if (!shape_cache_get_iv_index(shape, id, value)) {
         // If it wasn't in the ancestor cache, then don't do a linear search
@@ -1066,9 +1057,6 @@ shape_traverse_from_new_root(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
       case SHAPE_ROOT:
       case SHAPE_T_OBJECT:
         break;
-      case SHAPE_OBJ_TOO_COMPLEX:
-        rb_bug("Unreachable");
-        break;
     }
 
     return next_shape;
@@ -1085,17 +1073,17 @@ rb_shape_traverse_from_new_root(shape_id_t initial_shape_id, shape_id_t dest_sha
 // Rebuild a similar shape with the same ivars but starting from
 // a different SHAPE_T_OBJECT, and don't cary over non-canonical transitions
 // such as SHAPE_OBJ_ID.
-rb_shape_t *
-rb_shape_rebuild_shape(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
+static rb_shape_t *
+shape_rebuild(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
 {
     rb_shape_t *midway_shape;
 
     RUBY_ASSERT(initial_shape->type == SHAPE_T_OBJECT || initial_shape->type == SHAPE_ROOT);
 
     if (dest_shape->type != initial_shape->type) {
-        midway_shape = rb_shape_rebuild_shape(initial_shape, RSHAPE(dest_shape->parent_id));
-        if (UNLIKELY(raw_shape_id(midway_shape) == ROOT_TOO_COMPLEX_SHAPE_ID)) {
-            return midway_shape;
+        midway_shape = shape_rebuild(initial_shape, RSHAPE(dest_shape->parent_id));
+        if (UNLIKELY(!midway_shape)) {
+            return NULL;
         }
     }
     else {
@@ -1110,9 +1098,6 @@ rb_shape_rebuild_shape(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
       case SHAPE_ROOT:
       case SHAPE_T_OBJECT:
         break;
-      case SHAPE_OBJ_TOO_COMPLEX:
-        rb_bug("Unreachable");
-        break;
     }
 
     return midway_shape;
@@ -1124,7 +1109,7 @@ rb_shape_rebuild(shape_id_t initial_shape_id, shape_id_t dest_shape_id)
     RUBY_ASSERT(!rb_shape_too_complex_p(initial_shape_id));
     RUBY_ASSERT(!rb_shape_too_complex_p(dest_shape_id));
 
-    return raw_shape_id(rb_shape_rebuild_shape(RSHAPE(initial_shape_id), RSHAPE(dest_shape_id)));
+    return raw_shape_id(shape_rebuild(RSHAPE(initial_shape_id), RSHAPE(dest_shape_id)));
 }
 
 void
@@ -1171,7 +1156,7 @@ rb_shape_copy_complex_ivars(VALUE dest, VALUE obj, shape_id_t src_shape_id, st_t
 RUBY_FUNC_EXPORTED bool
 rb_shape_obj_too_complex_p(VALUE obj)
 {
-    return shape_too_complex_p(obj_shape(obj));
+    return !RB_SPECIAL_CONST_P(obj) && rb_shape_too_complex_p(RBASIC_SHAPE_ID(obj));
 }
 
 size_t
@@ -1210,8 +1195,7 @@ static VALUE
 shape_too_complex(VALUE self)
 {
     shape_id_t shape_id = NUM2INT(rb_struct_getmember(self, rb_intern("id")));
-    rb_shape_t *shape = RSHAPE(shape_id);
-    return RBOOL(shape_too_complex_p(shape));
+    return RBOOL(rb_shape_too_complex_p(shape_id));
 }
 
 static VALUE

shape.h

@@ -31,7 +31,7 @@ typedef uint32_t redblack_id_t;
 #define ROOT_SHAPE_ID               0x0
 #define ROOT_TOO_COMPLEX_SHAPE_ID   (ROOT_SHAPE_ID | SHAPE_ID_FL_TOO_COMPLEX)
 #define SPECIAL_CONST_SHAPE_ID      (ROOT_SHAPE_ID | SHAPE_ID_FL_FROZEN)
-#define FIRST_T_OBJECT_SHAPE_ID     0x2
+#define FIRST_T_OBJECT_SHAPE_ID     0x1
 
 extern ID ruby_internal_object_id;
 
@@ -63,7 +63,6 @@ enum shape_type {
     SHAPE_IVAR,
     SHAPE_OBJ_ID,
     SHAPE_T_OBJECT,
-    SHAPE_OBJ_TOO_COMPLEX,
 };
 
 enum shape_flags {

variable.c

@@ -2240,10 +2240,6 @@ iterate_over_shapes_with_callback(rb_shape_t *shape, rb_ivar_foreach_callback_fu
             }
         }
         return false;
-      case SHAPE_OBJ_TOO_COMPLEX:
-      default:
-        rb_bug("Unreachable");
-        UNREACHABLE_RETURN(false);
     }
 }
 

yjit.c

@@ -832,6 +832,12 @@ rb_yjit_set_exception_return(rb_control_frame_t *cfp, void *leave_exit, void *le
     }
 }
 
+bool
+rb_yjit_shape_too_complex_p(shape_id_t shape_id)
+{
+    return rb_shape_too_complex_p(shape_id);
+}
+
 // Primitives used by yjit.rb
 VALUE rb_yjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_yjit_print_stats_p(rb_execution_context_t *ec, VALUE self);

yjit/bindgen/src/main.rs

@@ -100,7 +100,7 @@ fn main() {
         .allowlist_function("rb_shape_get_iv_index")
         .allowlist_function("rb_shape_transition_add_ivar_no_warnings")
         .allowlist_function("rb_shape_obj_too_complex_p")
-        .allowlist_function("rb_shape_too_complex_p")
+        .allowlist_function("rb_yjit_shape_too_complex_p")
         .allowlist_var("SHAPE_ID_NUM_BITS")
 
         // From ruby/internal/intern/object.h

yjit/src/codegen.rs

@@ -3124,7 +3124,7 @@ fn gen_set_ivar(
 
         // 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).
-        new_shape_too_complex = unsafe { rb_shape_too_complex_p(next_shape_id) };
+        new_shape_too_complex = unsafe { rb_yjit_shape_too_complex_p(next_shape_id) };
         if new_shape_too_complex {
             Some((next_shape_id, None, 0_usize))
         } else {