From a640723d31262904b4de14be55357fb426873d7f Mon Sep 17 00:00:00 2001
From: Jean Boussier <jean.boussier@gmail.com>
Date: Sat, 7 Jun 2025 16:48:26 +0200
Subject: [PATCH] Simplify `rb_gc_rebuild_shape`

Now that there no longer multiple shape roots, all we need to do
when moving an object from one slot to the other is to update the
`heap_index` part of the shape_id.

Since this never need to create a shape transition, it will always
work and never result in a complex shape.
---
 gc.c    | 14 ++-----------
 shape.c | 63 ++++++---------------------------------------------------
 shape.h |  3 +--
 3 files changed, 9 insertions(+), 71 deletions(-)

diff --git a/gc.c b/gc.c
index 9d314c7416..3e7d88209f 100644
--- a/gc.c
+++ b/gc.c
@@ -381,19 +381,9 @@ rb_gc_set_shape(VALUE obj, uint32_t shape_id)
 uint32_t
 rb_gc_rebuild_shape(VALUE obj, size_t heap_id)
 {
-    shape_id_t orig_shape_id = rb_obj_shape_id(obj);
-    if (rb_shape_too_complex_p(orig_shape_id)) {
-        return (uint32_t)orig_shape_id;
-    }
+    RUBY_ASSERT(RB_TYPE_P(obj, T_OBJECT));
 
-    shape_id_t initial_shape_id = rb_shape_root(heap_id);
-    shape_id_t new_shape_id = rb_shape_traverse_from_new_root(initial_shape_id, orig_shape_id);
-
-    if (new_shape_id == INVALID_SHAPE_ID) {
-         return 0;
-     }
-
-    return (uint32_t)new_shape_id;
+    return (uint32_t)rb_shape_transition_heap(obj, heap_id);
 }
 
 void rb_vm_update_references(void *ptr);
diff --git a/shape.c b/shape.c
index 6cebec7cd7..a0295a4405 100644
--- a/shape.c
+++ b/shape.c
@@ -823,6 +823,12 @@ rb_shape_transition_complex(VALUE obj)
     return transition_complex(RBASIC_SHAPE_ID(obj));
 }
 
+shape_id_t
+rb_shape_transition_heap(VALUE obj, size_t heap_index)
+{
+     return (RBASIC_SHAPE_ID(obj) & (~SHAPE_ID_HEAP_INDEX_MASK)) | rb_shape_root(heap_index);
+}
+
 /*
  * This function is used for assertions where we don't want to increment
  * max_iv_count
@@ -1065,63 +1071,6 @@ rb_shape_id_offset(void)
     return sizeof(uintptr_t) - SHAPE_ID_NUM_BITS / sizeof(uintptr_t);
 }
 
-static rb_shape_t *
-shape_traverse_from_new_root(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
-{
-    RUBY_ASSERT(initial_shape->type == SHAPE_ROOT);
-    rb_shape_t *next_shape = initial_shape;
-
-    if (dest_shape->type != initial_shape->type) {
-        next_shape = shape_traverse_from_new_root(initial_shape, RSHAPE(dest_shape->parent_id));
-        if (!next_shape) {
-            return NULL;
-        }
-    }
-
-    switch ((enum shape_type)dest_shape->type) {
-      case SHAPE_IVAR:
-      case SHAPE_OBJ_ID:
-        if (!next_shape->edges) {
-            return NULL;
-        }
-
-        VALUE lookup_result;
-        if (SINGLE_CHILD_P(next_shape->edges)) {
-            rb_shape_t *child = SINGLE_CHILD(next_shape->edges);
-            if (child->edge_name == dest_shape->edge_name) {
-                return child;
-            }
-            else {
-                return NULL;
-            }
-        }
-        else {
-            if (rb_managed_id_table_lookup(next_shape->edges, dest_shape->edge_name, &lookup_result)) {
-                next_shape = (rb_shape_t *)lookup_result;
-            }
-            else {
-                return NULL;
-            }
-        }
-        break;
-      case SHAPE_ROOT:
-        break;
-    }
-
-    return next_shape;
-}
-
-shape_id_t
-rb_shape_traverse_from_new_root(shape_id_t initial_shape_id, shape_id_t dest_shape_id)
-{
-    rb_shape_t *initial_shape = RSHAPE(initial_shape_id);
-    rb_shape_t *dest_shape = RSHAPE(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
 // a different SHAPE_T_OBJECT, and don't cary over non-canonical transitions
 // such as SHAPE_OBJ_ID.
diff --git a/shape.h b/shape.h
index abfd27c075..7fec93af9f 100644
--- a/shape.h
+++ b/shape.h
@@ -164,6 +164,7 @@ shape_id_t rb_shape_transition_remove_ivar(VALUE obj, ID id, shape_id_t *removed
 shape_id_t rb_shape_transition_add_ivar(VALUE obj, ID id);
 shape_id_t rb_shape_transition_add_ivar_no_warnings(VALUE obj, ID id);
 shape_id_t rb_shape_transition_object_id(VALUE obj);
+shape_id_t rb_shape_transition_heap(VALUE obj, size_t heap_index);
 shape_id_t rb_shape_object_id(shape_id_t original_shape_id);
 
 void rb_shape_free_all(void);
@@ -302,8 +303,6 @@ RBASIC_FIELDS_COUNT(VALUE obj)
     return RSHAPE(rb_obj_shape_id(obj))->next_field_index;
 }
 
-shape_id_t rb_shape_traverse_from_new_root(shape_id_t initial_shape_id, shape_id_t orig_shape_id);
-
 bool rb_obj_set_shape_id(VALUE obj, shape_id_t shape_id);
 
 static inline bool