ruby · casperisfine · Jun 6, 2025 · Jun 6, 2025 · Jun 7, 2025
diff --git a/internal/variable.h b/internal/variable.h
@@ -52,7 +52,7 @@ int rb_gen_fields_tbl_get(VALUE obj, ID id, struct gen_fields_tbl **fields_tbl);
 void rb_obj_copy_ivs_to_hash_table(VALUE obj, st_table *table);
 void rb_obj_init_too_complex(VALUE obj, st_table *table);
 void rb_evict_ivars_to_hash(VALUE obj);
-void rb_evict_fields_to_hash(VALUE obj);
+shape_id_t rb_evict_fields_to_hash(VALUE obj);
 VALUE rb_obj_field_get(VALUE obj, shape_id_t target_shape_id);
 void rb_ivar_set_internal(VALUE obj, ID id, VALUE val);
 void rb_obj_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val);

diff --git a/object.c b/object.c
@@ -339,17 +339,15 @@ rb_obj_copy_ivar(VALUE dest, VALUE obj)
     shape_id_t dest_shape_id = src_shape_id;
     shape_id_t initial_shape_id = RBASIC_SHAPE_ID(dest);
 
-    if (RSHAPE(initial_shape_id)->heap_index != RSHAPE(src_shape_id)->heap_index || !rb_shape_canonical_p(src_shape_id)) {
-        RUBY_ASSERT(RSHAPE(initial_shape_id)->type == SHAPE_T_OBJECT);
+    RUBY_ASSERT(RSHAPE(initial_shape_id)->type == SHAPE_T_OBJECT);
 
-        dest_shape_id = rb_shape_rebuild(initial_shape_id, src_shape_id);
-        if (UNLIKELY(rb_shape_too_complex_p(dest_shape_id))) {
-            st_table *table = rb_st_init_numtable_with_size(src_num_ivs);
-            rb_obj_copy_ivs_to_hash_table(obj, table);
-            rb_obj_init_too_complex(dest, table);
+    dest_shape_id = rb_shape_rebuild(initial_shape_id, src_shape_id);
+    if (UNLIKELY(rb_shape_too_complex_p(dest_shape_id))) {
+        st_table *table = rb_st_init_numtable_with_size(src_num_ivs);
+        rb_obj_copy_ivs_to_hash_table(obj, table);
+        rb_obj_init_too_complex(dest, table);
 
-            return;
-        }
+        return;
     }
 
     VALUE *src_buf = ROBJECT_FIELDS(obj);

diff --git a/shape.c b/shape.c
@@ -723,15 +723,67 @@ remove_shape_recursive(rb_shape_t *shape, ID id, rb_shape_t **removed_shape)
     }
 }
 
+
+static shape_id_t
+shape_transition_object_id(shape_id_t original_shape_id)
+{
+    RUBY_ASSERT(!rb_shape_has_object_id(original_shape_id));
+
+    bool dont_care;
+    rb_shape_t *shape = get_next_shape_internal(RSHAPE(original_shape_id), ruby_internal_object_id, SHAPE_OBJ_ID, &dont_care, true);
+
+    RUBY_ASSERT(shape);
+
+    return shape_id(shape, original_shape_id) | SHAPE_ID_FL_HAS_OBJECT_ID;
+}
+
+shape_id_t
+rb_shape_transition_object_id(VALUE obj)
+{
+    return shape_transition_object_id(RBASIC_SHAPE_ID(obj));
+}
+
+shape_id_t
+rb_shape_object_id(shape_id_t original_shape_id)
+{
+    RUBY_ASSERT(rb_shape_has_object_id(original_shape_id));
+
+    rb_shape_t *shape = RSHAPE(original_shape_id);
+    while (shape->type != SHAPE_OBJ_ID) {
+        if (UNLIKELY(shape->parent_id == INVALID_SHAPE_ID)) {
+            rb_bug("Missing object_id in shape tree");
+        }
+        shape = RSHAPE(shape->parent_id);
+    }
+
+    return shape_id(shape, original_shape_id) | SHAPE_ID_FL_HAS_OBJECT_ID;
+}
+
 static inline shape_id_t
 transition_complex(shape_id_t shape_id)
 {
-        if (rb_shape_has_object_id(shape_id)) {
-        return ROOT_TOO_COMPLEX_WITH_OBJ_ID | (shape_id & SHAPE_ID_FLAGS_MASK);
+    uint8_t heap_index = RSHAPE(shape_id)->heap_index;
+    shape_id_t next_shape_id;
+
+    if (heap_index) {
+        next_shape_id = rb_shape_root(heap_index - 1) | SHAPE_ID_FL_TOO_COMPLEX;
+        if (rb_shape_has_object_id(shape_id)) {
+            next_shape_id = shape_transition_object_id(next_shape_id);
+        }
+    }
+    else {
+        if (rb_shape_has_object_id(shape_id)) {
+            next_shape_id = ROOT_TOO_COMPLEX_WITH_OBJ_ID | (shape_id & SHAPE_ID_FLAGS_MASK);
+        }
+        else {
+            next_shape_id = ROOT_TOO_COMPLEX_SHAPE_ID | (shape_id & SHAPE_ID_FLAGS_MASK);
+        }
     }
-    return ROOT_TOO_COMPLEX_SHAPE_ID | (shape_id & SHAPE_ID_FLAGS_MASK);
-}
 
+    RUBY_ASSERT(rb_shape_has_object_id(shape_id) == rb_shape_has_object_id(next_shape_id));
+
+    return next_shape_id;
+}
 
 shape_id_t
 rb_shape_transition_remove_ivar(VALUE obj, ID id, shape_id_t *removed_shape_id)
@@ -754,7 +806,9 @@ rb_shape_transition_remove_ivar(VALUE obj, ID id, shape_id_t *removed_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 transition_complex(original_shape_id);
+        shape_id_t next_shape_id = transition_complex(original_shape_id);
+        RUBY_ASSERT(rb_shape_has_object_id(next_shape_id) == rb_shape_has_object_id(original_shape_id));
+        return next_shape_id;
     }
     return original_shape_id;
 }
@@ -774,41 +828,6 @@ rb_shape_transition_complex(VALUE obj)
     return transition_complex(RBASIC_SHAPE_ID(obj));
 }
 
-shape_id_t
-rb_shape_transition_object_id(VALUE obj)
-{
-    shape_id_t original_shape_id = RBASIC_SHAPE_ID(obj);
-
-    RUBY_ASSERT(!rb_shape_has_object_id(original_shape_id));
-
-    rb_shape_t *shape = NULL;
-    if (!rb_shape_too_complex_p(original_shape_id)) {
-        bool dont_care;
-        shape = get_next_shape_internal(RSHAPE(original_shape_id), ruby_internal_object_id, SHAPE_OBJ_ID, &dont_care, true);
-    }
-
-    if (!shape) {
-        shape = RSHAPE(ROOT_TOO_COMPLEX_WITH_OBJ_ID);
-    }
-    return shape_id(shape, original_shape_id) | SHAPE_ID_FL_HAS_OBJECT_ID;
-}
-
-shape_id_t
-rb_shape_object_id(shape_id_t original_shape_id)
-{
-    RUBY_ASSERT(rb_shape_has_object_id(original_shape_id));
-
-    rb_shape_t *shape = RSHAPE(original_shape_id);
-    while (shape->type != SHAPE_OBJ_ID) {
-        if (UNLIKELY(shape->parent_id == INVALID_SHAPE_ID)) {
-            rb_bug("Missing object_id in shape tree");
-        }
-        shape = RSHAPE(shape->parent_id);
-    }
-
-    return shape_id(shape, original_shape_id) | SHAPE_ID_FL_HAS_OBJECT_ID;
-}
-
 /*
  * This function is used for assertions where we don't want to increment
  * max_iv_count
@@ -918,7 +937,13 @@ rb_shape_transition_add_ivar(VALUE obj, ID id)
     shape_id_t original_shape_id = RBASIC_SHAPE_ID(obj);
     RUBY_ASSERT(!shape_frozen_p(original_shape_id));
 
-    return shape_id(shape_get_next(RSHAPE(original_shape_id), obj, id, true), original_shape_id);
+    rb_shape_t *next_shape = shape_get_next(RSHAPE(original_shape_id), obj, id, true);
+    if (next_shape) {
+        return shape_id(next_shape, original_shape_id);
+    }
+    else {
+        return transition_complex(original_shape_id);
+    }
 }
 
 shape_id_t
@@ -927,7 +952,13 @@ rb_shape_transition_add_ivar_no_warnings(VALUE obj, ID id)
     shape_id_t original_shape_id = RBASIC_SHAPE_ID(obj);
     RUBY_ASSERT(!shape_frozen_p(original_shape_id));
 
-    return shape_id(shape_get_next(RSHAPE(original_shape_id), obj, id, false), original_shape_id);
+    rb_shape_t *next_shape = shape_get_next(RSHAPE(original_shape_id), obj, id, false);
+    if (next_shape) {
+        return shape_id(next_shape, original_shape_id);
+    }
+    else {
+        return transition_complex(original_shape_id);
+    }
 }
 
 // Same as rb_shape_get_iv_index, but uses a provided valid shape id and index
@@ -1092,7 +1123,10 @@ rb_shape_traverse_from_new_root(shape_id_t initial_shape_id, shape_id_t dest_sha
 {
     rb_shape_t *initial_shape = RSHAPE(initial_shape_id);
     rb_shape_t *dest_shape = RSHAPE(dest_shape_id);
-    return shape_id(shape_traverse_from_new_root(initial_shape, dest_shape), dest_shape_id);
+
+    // Keep all dest_shape_id flags except for the heap_index.
+    shape_id_t dest_flags = (dest_shape_id & ~SHAPE_ID_HEAP_INDEX_MASK) | (initial_shape_id & SHAPE_ID_HEAP_INDEX_MASK);
+    return shape_id(shape_traverse_from_new_root(initial_shape, dest_shape), dest_flags);
 }
 
 // Rebuild a similar shape with the same ivars but starting from
@@ -1136,7 +1170,13 @@ 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(shape_rebuild(RSHAPE(initial_shape_id), RSHAPE(dest_shape_id)));
+    rb_shape_t *next_shape = shape_rebuild(RSHAPE(initial_shape_id), RSHAPE(dest_shape_id));
+    if (next_shape) {
+        return shape_id(next_shape, initial_shape_id);
+    }
+    else {
+        return transition_complex(initial_shape_id | (dest_shape_id & SHAPE_ID_FL_HAS_OBJECT_ID));
+    }
 }
 
 void
@@ -1214,6 +1254,10 @@ rb_shape_memsize(shape_id_t shape_id)
 bool
 rb_shape_verify_consistency(VALUE obj, shape_id_t shape_id)
 {
+    if (shape_id == INVALID_SHAPE_ID) {
+        rb_bug("Can't set INVALID_SHAPE_ID on an object");
+    }
+
     rb_shape_t *shape = RSHAPE(shape_id);
 
     bool has_object_id = false;
@@ -1238,6 +1282,11 @@ rb_shape_verify_consistency(VALUE obj, shape_id_t shape_id)
         }
     }
 
+    uint8_t flags_heap_index = rb_shape_heap_index(shape_id);
+    if (flags_heap_index != shape->heap_index) {
+        rb_bug("shape_id heap_index flags mismatch: flags=%u, transition=%u\n", flags_heap_index, shape->heap_index);
+    }
+
     return true;
 }
 #endif
@@ -1288,6 +1337,7 @@ shape_id_t_to_rb_cShape(shape_id_t shape_id)
 
     VALUE obj = rb_struct_new(rb_cShape,
             INT2NUM(shape_id),
+            INT2NUM(shape_id & SHAPE_ID_OFFSET_MASK),
             INT2NUM(shape->parent_id),
             rb_shape_edge_name(shape),
             INT2NUM(shape->next_field_index),
@@ -1525,15 +1575,23 @@ Init_default_shapes(void)
 
     // Make shapes for T_OBJECT
     size_t *sizes = rb_gc_heap_sizes();
-    for (int i = 0; sizes[i] > 0; i++) {
+    int i;
+    for (i = 0; sizes[i] > 0; i++) {
         rb_shape_t *t_object_shape = rb_shape_alloc_with_parent_id(0, INVALID_SHAPE_ID);
         t_object_shape->type = SHAPE_T_OBJECT;
-        t_object_shape->heap_index = i;
+        t_object_shape->heap_index = i + 1;
         t_object_shape->capacity = (uint32_t)((sizes[i] - offsetof(struct RObject, as.ary)) / sizeof(VALUE));
         t_object_shape->edges = rb_managed_id_table_new(256);
         t_object_shape->ancestor_index = LEAF;
         RUBY_ASSERT(t_object_shape == RSHAPE(rb_shape_root(i)));
     }
+
+    // Prebuild all ROOT + OBJ_ID shapes so that even when we run out of shape we can always transtion to
+    // COMPLEX + OBJ_ID.
+    bool dont_care;
+    for (i = 0; sizes[i] > 0; i++) {
+        get_next_shape_internal(RSHAPE(rb_shape_root(i)), ruby_internal_object_id, SHAPE_OBJ_ID, &dont_care, true);
+    }
 }
 
 void
@@ -1552,6 +1610,7 @@ Init_shape(void)
      * :nodoc: */
     VALUE rb_cShape = rb_struct_define_under(rb_cRubyVM, "Shape",
             "id",
+            "raw_id",
             "parent_id",
             "edge_name",
             "next_field_index",

diff --git a/shape.h b/shape.h
@@ -16,8 +16,18 @@ STATIC_ASSERT(shape_id_num_bits, SHAPE_ID_NUM_BITS == sizeof(shape_id_t) * CHAR_
 #define SHAPE_ID_FL_FROZEN (SHAPE_FL_FROZEN << SHAPE_ID_OFFSET_NUM_BITS)
 #
F438
define SHAPE_ID_FL_HAS_OBJECT_ID (SHAPE_FL_HAS_OBJECT_ID << SHAPE_ID_OFFSET_NUM_BITS)
 #define SHAPE_ID_FL_TOO_COMPLEX (SHAPE_FL_TOO_COMPLEX << SHAPE_ID_OFFSET_NUM_BITS)
+#define SHAPE_ID_FL_EMBEDDED (SHAPE_FL_EMBEDDED << SHAPE_ID_OFFSET_NUM_BITS)
 #define SHAPE_ID_FL_NON_CANONICAL_MASK (SHAPE_FL_NON_CANONICAL_MASK << SHAPE_ID_OFFSET_NUM_BITS)
-#define SHAPE_ID_READ_ONLY_MASK (~SHAPE_ID_FL_FROZEN)
+
+#define SHAPE_ID_HEAP_INDEX_BITS 3
+#define SHAPE_ID_HEAP_INDEX_OFFSET (SHAPE_ID_NUM_BITS - SHAPE_ID_HEAP_INDEX_BITS - 1) // FIXME: -1 to avoid crashing YJIT
+#define SHAPE_ID_HEAP_INDEX_MAX ((1 << SHAPE_ID_HEAP_INDEX_BITS) - 1)
+#define SHAPE_ID_HEAP_INDEX_MASK (SHAPE_ID_HEAP_INDEX_MAX << SHAPE_ID_HEAP_INDEX_OFFSET)
+
+// The interpreter doesn't care about embeded or frozen status when reading ivars.
+// So we normalize shape_id by clearing these bits to improve cache hits.
+// JITs however might care about it.
+#define SHAPE_ID_READ_ONLY_MASK (~(SHAPE_ID_FL_FROZEN | SHAPE_ID_FL_EMBEDDED | SHAPE_ID_HEAP_INDEX_MASK))
 
 typedef uint32_t redblack_id_t;
 
@@ -72,6 +82,7 @@ enum shape_flags {
     SHAPE_FL_FROZEN             = 1 << 0,
     SHAPE_FL_HAS_OBJECT_ID      = 1 << 1,
     SHAPE_FL_TOO_COMPLEX        = 1 << 2,
+    SHAPE_FL_EMBEDDED           = 1 << 3,
 
     SHAPE_FL_NON_CANONICAL_MASK = SHAPE_FL_FROZEN | SHAPE_FL_HAS_OBJECT_ID,
 };
@@ -189,10 +200,19 @@ rb_shape_canonical_p(shape_id_t shape_id)
     return !(shape_id & SHAPE_ID_FL_NON_CANONICAL_MASK);
 }
 
+static inline uint8_t
+rb_shape_heap_index(shape_id_t shape_id)
+{
+    return (uint8_t)((shape_id & SHAPE_ID_HEAP_INDEX_MASK) >> SHAPE_ID_HEAP_INDEX_OFFSET);
+}
+
 static inline shape_id_t
 rb_shape_root(size_t heap_id)
 {
-    return (shape_id_t)(heap_id + FIRST_T_OBJECT_SHAPE_ID);
+    shape_id_t heap_index = (shape_id_t)heap_id;
+
+    shape_id_t shape_id = (heap_index + FIRST_T_OBJECT_SHAPE_ID);
+    return shape_id | ((heap_index + 1) << SHAPE_ID_HEAP_INDEX_OFFSET) | SHAPE_ID_FL_EMBEDDED;
 }
 
 static inline bool

diff --git a/test/ruby/test_shapes.rb b/test/ruby/test_shapes.rb
@@ -976,7 +976,7 @@ def test_iv_index
     example.add_foo # makes a transition
     add_foo_shape = RubyVM::Shape.of(example)
     assert_equal([:@foo], example.instance_variables)
-    assert_equal(initial_shape.id, add_foo_shape.parent.id)
+    assert_equal(initial_shape.raw_id, add_foo_shape.parent.raw_id)
     assert_equal(1, add_foo_shape.next_field_index)
 
     example.remove_foo # makes a transition
@@ -987,7 +987,7 @@ def test_iv_index
     example.add_bar # makes a transition
     bar_shape = RubyVM::Shape.of(example)
     assert_equal([:@bar], example.instance_variables)
-    assert_equal(initial_shape.id, bar_shape.parent_id)
+    assert_equal(initial_shape.raw_id, bar_shape.parent_id)
     assert_equal(1, bar_shape.next_field_index)
   end