From f1670733249fb30d755bad1f88c0e54b26bdf49e Mon Sep 17 00:00:00 2001
From: Max Bernstein <ruby@bernsteinbear.com>
Date: Wed, 3 Dec 2025 16:06:41 -0500
Subject: Move imemo fields check out of shape_get_next

Not every caller (for example, YJIT) actually needs to pass the object.
YJIT (and, in the future, ZJIT) only need to pass the class.
---
 shape.c | 50 +++++++++++++++++++++++++++++---------------------
 1 file changed, 29 insertions(+), 21 deletions(-)

(limited to 'shape.c')

diff --git a/shape.c b/shape.c
index c6f5af2519..13d1edf36c 100644
--- a/shape.c
+++ b/shape.c
@@ -801,7 +801,7 @@ shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
 }
 
 static inline rb_shape_t *
-shape_get_next(rb_shape_t *shape, enum shape_type shape_type, VALUE obj, ID id, bool emit_warnings)
+shape_get_next(rb_shape_t *shape, enum shape_type shape_type, VALUE klass, ID id, bool emit_warnings)
 {
     RUBY_ASSERT(!is_instance_id(id) || RTEST(rb_sym2str(ID2SYM(id))));
 
@@ -812,23 +812,6 @@ shape_get_next(rb_shape_t *shape, enum shape_type shape_type, VALUE obj, ID id,
     }
 #endif
 
-    VALUE klass;
-    if (IMEMO_TYPE_P(obj, imemo_fields)) {
-        VALUE owner = rb_imemo_fields_owner(obj);
-        switch (BUILTIN_TYPE(owner)) {
-          case T_CLASS:
-          case T_MODULE:
-            klass = rb_singleton_class(owner);
-            break;
-          default:
-            klass = rb_obj_class(owner);
-            break;
-        }
-    }
-    else {
-        klass = rb_obj_class(obj);
-    }
-
     bool allow_new_shape = RCLASS_VARIATION_COUNT(klass) < SHAPE_MAX_VARIATIONS;
     bool variation_created = false;
     rb_shape_t *new_shape = get_next_shape_internal(shape, id, shape_type, &variation_created, allow_new_shape);
@@ -862,6 +845,28 @@ shape_get_next(rb_shape_t *shape, enum shape_type shape_type, VALUE obj, ID id,
     return new_shape;
 }
 
+static VALUE
+obj_get_owner_class(VALUE obj)
+{
+    VALUE klass;
+    if (IMEMO_TYPE_P(obj, imemo_fields)) {
+        VALUE owner = rb_imemo_fields_owner(obj);
+        switch (BUILTIN_TYPE(owner)) {
+          case T_CLASS:
+          case T_MODULE:
+            klass = rb_singleton_class(owner);
+            break;
+          default:
+            klass = rb_obj_class(owner);
+            break;
+        }
+    }
+    else {
+        klass = rb_obj_class(obj);
+    }
+    return klass;
+}
+
 static rb_shape_t *
 remove_shape_recursive(VALUE obj, rb_shape_t *shape, ID id, rb_shape_t **removed_shape)
 {
@@ -884,7 +889,8 @@ remove_shape_recursive(VALUE obj, rb_shape_t *shape, ID id, rb_shape_t **removed
             // We found a new parent.  Create a child of the new parent that
             // has the same attributes as this shape.
             if (new_parent) {
-                rb_shape_t *new_child = shape_get_next(new_parent, shape->type, obj, shape->edge_name, true);
+                VALUE klass = obj_get_owner_class(obj);
+                rb_shape_t *new_child = shape_get_next(new_parent, shape->type, klass, shape->edge_name, true);
                 RUBY_ASSERT(!new_child || new_child->capacity <= shape->capacity);
                 return new_child;
             }
@@ -933,7 +939,8 @@ 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));
 
-    rb_shape_t *next_shape = shape_get_next(RSHAPE(original_shape_id), SHAPE_IVAR, obj, id, true);
+    VALUE klass = obj_get_owner_class(obj);
+    rb_shape_t *next_shape = shape_get_next(RSHAPE(original_shape_id), SHAPE_IVAR, klass, id, true);
     if (next_shape) {
         return shape_id(next_shape, original_shape_id);
     }
@@ -948,7 +955,8 @@ 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));
 
-    rb_shape_t *next_shape = shape_get_next(RSHAPE(original_shape_id), SHAPE_IVAR, obj, id, false);
+    VALUE klass = obj_get_owner_class(obj);
+    rb_shape_t *next_shape = shape_get_next(RSHAPE(original_shape_id), SHAPE_IVAR, klass, id, false);
     if (next_shape) {
         return shape_id(next_shape, original_shape_id);
     }
-- 
cgit v1.2.3