Refactor `rb_shape_traverse_from_new_root` to not expose `rb_shape_t`

3 files changed, 21 insertions, 9 deletions

diff --git a/gc.c b/gc.c
index 0ca08d53e0..9552db5fe6 100644
--- a/gc.c
+++ b/gc.c
@@ -384,13 +384,14 @@ rb_gc_rebuild_shape(VALUE obj, size_t heap_id)
         return (uint32_t)orig_shape_id;
     }
 
-    rb_shape_t *orig_shape = rb_shape_get_shape_by_id(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);
+    shape_id_t initial_shape_id = (shape_id_t)(heap_id + FIRST_T_OBJECT_SHAPE_ID);
+    shape_id_t new_shape_id = rb_shape_traverse_from_new_root(initial_shape_id, orig_shape_id);
 
-    if (!new_shape) return 0;
+    if (new_shape_id == INVALID_SHAPE_ID) {
+         return 0;
+     }
 
-    return (uint32_t)rb_shape_id(new_shape);
+    return (uint32_t)new_shape_id;
 }
 
 void rb_vm_update_references(void *ptr);
diff --git a/shape.c b/shape.c
index ab32533a9d..ccf2a7a898 100644
--- a/shape.c
+++ b/shape.c
@@ -322,6 +322,9 @@ rb_shape_get_root_shape(void)
 shape_id_t
 rb_shape_id(rb_shape_t *shape)
 {
+    if (shape == NULL) {
+        return INVALID_SHAPE_ID;
+    }
     return (shape_id_t)(shape - GET_SHAPE_TREE()->shape_list);
 }
 
@@ -999,14 +1002,14 @@ rb_shape_id_offset(void)
     return sizeof(uintptr_t) - SHAPE_ID_NUM_BITS / sizeof(uintptr_t);
 }
 
-rb_shape_t *
-rb_shape_traverse_from_new_root(rb_shape_t *initial_shape, rb_shape_t *dest_shape)
+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_T_OBJECT);
     rb_shape_t *next_shape = initial_shape;
 
     if (dest_shape->type != initial_shape->type) {
-        next_shape = rb_shape_traverse_from_new_root(initial_shape, rb_shape_get_parent(dest_shape));
+        next_shape = shape_traverse_from_new_root(initial_shape, rb_shape_get_parent(dest_shape));
         if (!next_shape) {
             return NULL;
         }
@@ -1050,6 +1053,14 @@ rb_shape_traverse_from_new_root(rb_shape_t *initial_shape, rb_shape_t *dest_shap
     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 = rb_shape_get_shape_by_id(initial_shape_id);
+    rb_shape_t *dest_shape = rb_shape_get_shape_by_id(dest_shape_id);
+    return rb_shape_id(shape_traverse_from_new_root(initial_shape, dest_shape));
+}
+
 // 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_FROZEN or SHAPE_OBJ_ID.
diff --git a/shape.h b/shape.h
index dd52e20543..8228916003 100644
--- a/shape.h
+++ b/shape.h
@@ -232,7 +232,7 @@ RBASIC_FIELDS_COUNT(VALUE obj)
     return rb_shape_get_shape_by_id(rb_shape_get_shape_id(obj))->next_field_index;
 }
 
-rb_shape_t *rb_shape_traverse_from_new_root(rb_shape_t *initial_shape, rb_shape_t *orig_shape);
+shape_id_t rb_shape_traverse_from_new_root(shape_id_t initial_shape_id, shape_id_t orig_shape_id);
 
 bool rb_shape_set_shape_id(VALUE obj, shape_id_t shape_id);