Refactor rb_shape_get_iv_index to take a shape_id_t

Further reduce exposure of `rb_shape_t`.

Author: byroot

shape.c

@@ -731,6 +731,35 @@ rb_shape_get_next_iv_shape(shape_id_t shape_id, ID id)
     return rb_shape_id(next_shape);
 }
 
+static bool
+shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
+{
+    while (shape->parent_id != INVALID_SHAPE_ID) {
+        if (shape->edge_name == id) {
+            enum shape_type shape_type;
+            shape_type = (enum shape_type)shape->type;
+
+            switch (shape_type) {
+              case SHAPE_IVAR:
+                RUBY_ASSERT(shape->next_field_index > 0);
+                *value = shape->next_field_index - 1;
+                return true;
+              case SHAPE_ROOT:
+              case SHAPE_T_OBJECT:
+                return false;
+              case SHAPE_OBJ_TOO_COMPLEX:
+              case SHAPE_OBJ_ID:
+              case SHAPE_FROZEN:
+                rb_bug("Ivar should not exist on transition");
+            }
+        }
+
+        shape = RSHAPE(shape->parent_id);
+    }
+
+    return false;
+}
+
 static inline rb_shape_t *
 shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
 {
@@ -741,7 +770,7 @@ shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
 
 #if RUBY_DEBUG
     attr_index_t index;
-    if (rb_shape_get_iv_index(shape, id, &index)) {
+    if (shape_get_iv_index(shape, id, &index)) {
         rb_bug("rb_shape_get_next: trying to create ivar that already exists at index %u", index);
     }
 #endif
@@ -803,14 +832,14 @@ bool
 rb_shape_get_iv_index_with_hint(shape_id_t shape_id, ID id, attr_index_t *value, shape_id_t *shape_id_hint)
 {
     attr_index_t index_hint = *value;
-    rb_shape_t *shape = RSHAPE(shape_id);
-    rb_shape_t *initial_shape = shape;
 
     if (*shape_id_hint == INVALID_SHAPE_ID) {
         *shape_id_hint = shape_id;
-        return rb_shape_get_iv_index(shape, id, value);
+        return rb_shape_get_iv_index(shape_id, id, value);
     }
 
+    rb_shape_t *shape = RSHAPE(shape_id);
+    rb_shape_t *initial_shape = shape;
     rb_shape_t *shape_hint = RSHAPE(*shape_id_hint);
 
     // We assume it's likely shape_id_hint and shape_id have a close common
@@ -850,36 +879,7 @@ rb_shape_get_iv_index_with_hint(shape_id_t shape_id, ID id, attr_index_t *value,
         shape = initial_shape;
     }
     *shape_id_hint = shape_id;
-    return rb_shape_get_iv_index(shape, id, value);
-}
-
-static bool
-shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
-{
-    while (shape->parent_id != INVALID_SHAPE_ID) {
-        if (shape->edge_name == id) {
-            enum shape_type shape_type;
-            shape_type = (enum shape_type)shape->type;
-
-            switch (shape_type) {
-              case SHAPE_IVAR:
-                RUBY_ASSERT(shape->next_field_index > 0);
-                *value = shape->next_field_index - 1;
-                return true;
-              case SHAPE_ROOT:
-              case SHAPE_T_OBJECT:
-                return false;
-              case SHAPE_OBJ_TOO_COMPLEX:
-              case SHAPE_OBJ_ID:
-              case SHAPE_FROZEN:
-                rb_bug("Ivar should not exist on transition");
-            }
-        }
-
-        shape = RSHAPE(shape->parent_id);
-    }
-
-    return false;
+    return shape_get_iv_index(shape, id, value);
 }
 
 static bool
@@ -909,11 +909,13 @@ shape_cache_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
 }
 
 bool
-rb_shape_get_iv_index(rb_shape_t *shape, ID id, attr_index_t *value)
+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(rb_shape_id(shape) != ROOT_TOO_COMPLEX_SHAPE_ID);
+    RUBY_ASSERT(!rb_shape_too_complex_p(shape));
 
     if (!shape_cache_get_iv_index(shape, id, value)) {
         // If it wasn't in the ancestor cache, then don't do a linear search

shape.h

@@ -122,7 +122,7 @@ int32_t rb_shape_id_offset(void);
 RUBY_FUNC_EXPORTED rb_shape_t *rb_shape_lookup(shape_id_t shape_id);
 RUBY_FUNC_EXPORTED shape_id_t rb_obj_shape_id(VALUE obj);
 shape_id_t rb_shape_get_next_iv_shape(shape_id_t shape_id, ID id);
-bool rb_shape_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);
 bool rb_shape_get_iv_index_with_hint(shape_id_t shape_id, ID id, attr_index_t *value, shape_id_t *shape_id_hint);
 RUBY_FUNC_EXPORTED bool rb_shape_obj_too_complex_p(VALUE obj);
 bool rb_shape_too_complex_p(rb_shape_t *shape);

variable.c

@@ -1376,8 +1376,6 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
 
     shape_id_t shape_id;
     VALUE * ivar_list;
-    rb_shape_t * shape;
-
     shape_id = RBASIC_SHAPE_ID(obj);
 
     switch (BUILTIN_TYPE(obj)) {
@@ -1399,8 +1397,7 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
                 }
                 else {
                     attr_index_t index = 0;
-                    shape = RSHAPE(shape_id);
-                    found = rb_shape_get_iv_index(shape, id, &index);
+                    found = rb_shape_get_iv_index(shape_id, id, &index);
 
                     if (found) {
                         ivar_list = RCLASS_PRIME_FIELDS(obj);
@@ -1463,8 +1460,7 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
     }
 
     attr_index_t index = 0;
-    shape = RSHAPE(shape_id);
-    if (rb_shape_get_iv_index(shape, id, &index)) {
+    if (rb_shape_get_iv_index(shape_id, id, &index)) {
         return ivar_list[index];
     }
 
@@ -1717,15 +1713,16 @@ general_ivar_set(VALUE obj, ID id, VALUE val, void *data,
         .existing = true
     };
 
-    rb_shape_t *current_shape = rb_obj_shape(obj);
+    shape_id_t current_shape_id = RBASIC_SHAPE_ID(obj);
 
-    if (UNLIKELY(rb_shape_too_complex_p(current_shape))) {
+    if (UNLIKELY(rb_shape_id_too_complex_p(current_shape_id))) {
         goto too_complex;
     }
 
     attr_index_t index;
-    if (!rb_shape_get_iv_index(current_shape, id, &index)) {
+    if (!rb_shape_get_iv_index(current_shape_id, id, &index)) {
         result.existing = false;
+        rb_shape_t *current_shape = RSHAPE(current_shape_id);
 
         index = current_shape->next_field_index;
         if (index >= SHAPE_MAX_FIELDS) {
@@ -2172,7 +2169,7 @@ rb_ivar_defined(VALUE obj, ID id)
         return Qtrue;
     }
     else {
-        return RBOOL(rb_shape_get_iv_index(rb_obj_shape(obj), id, &index));
+        return RBOOL(rb_shape_get_iv_index(RBASIC_SHAPE_ID(obj), id, &index));
     }
 }
 

yjit/src/codegen.rs

@@ -2904,9 +2904,8 @@ fn gen_get_ivar(
 
     let ivar_index = unsafe {
         let shape_id = comptime_receiver.shape_id_of();
-        let shape = rb_shape_lookup(shape_id);
         let mut ivar_index: u32 = 0;
-        if rb_shape_get_iv_index(shape, ivar_name, &mut ivar_index) {
+        if rb_shape_get_iv_index(shape_id, ivar_name, &mut ivar_index) {
             Some(ivar_index as usize)
         } else {
             None
@@ -3107,9 +3106,8 @@ fn gen_set_ivar(
     let shape_too_complex = comptime_receiver.shape_too_complex();
     let ivar_index = if !shape_too_complex {
         let shape_id = comptime_receiver.shape_id_of();
-        let shape = unsafe { rb_shape_lookup(shape_id) };
         let mut ivar_index: u32 = 0;
-        if unsafe { rb_shape_get_iv_index(shape, ivar_name, &mut ivar_index) } {
+        if unsafe { rb_shape_get_iv_index(shape_id, ivar_name, &mut ivar_index) } {
             Some(ivar_index as usize)
         } else {
             None
@@ -3397,9 +3395,8 @@ fn gen_definedivar(
 
     let shape_id = comptime_receiver.shape_id_of();
     let ivar_exists = unsafe {
-        let shape = rb_shape_lookup(shape_id);
         let mut ivar_index: u32 = 0;
-        rb_shape_get_iv_index(shape, ivar_name, &mut ivar_index)
+        rb_shape_get_iv_index(shape_id, ivar_name, &mut ivar_index)
     };
 
     // Guard heap object (recv_opnd must be used before stack_pop)