ZJIT: Centralize shape setting for ZJIT GC fastpath

This commit centralizes shape setting for objects in the ZJIT GC fastpath
into the rb_zjit_new_obj_shape function to avoid each type having to deal
with shapes individually.

array.c

@@ -2944,11 +2944,8 @@ rb_zjit_array_dup_can_fastpath(VALUE ary, size_t *alloc_size_out, VALUE *flags_o
 
     if (len > embed_capa) return false;
 
-    size_t size = sizeof(struct RArray);
-    shape_id_t shape_id = rb_shape_transition_slot_size(ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_OTHER,
-                                                        rb_gc_size_slot_size(size));
-    *alloc_size_out = size;
-    *flags_out = T_ARRAY | RARRAY_EMBED_FLAG | ((VALUE)len << RARRAY_EMBED_LEN_SHIFT) | ((VALUE)shape_id << SHAPE_FLAG_SHIFT);
+    *alloc_size_out = sizeof(struct RArray);
+    *flags_out = T_ARRAY | RARRAY_EMBED_FLAG | ((VALUE)len << RARRAY_EMBED_LEN_SHIFT);
     *len_out = len;
     return true;
 }

gc.c

@@ -1252,7 +1252,7 @@ rb_class_allocate_instance(VALUE klass)
 
 #if USE_ZJIT
 bool
-rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, shape_id_t *shape_id_out)
+rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, VALUE *flags_out)
 {
     uint32_t index_tbl_num_entries = RCLASS_MAX_IV_COUNT(klass);
 
@@ -1261,10 +1261,9 @@ rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, shape_id
         return false;
     }
 
-    size_t size = robject_embedded_size(index_tbl_num_entries);
-    *size_out = size;
-    *shape_id_out = rb_shape_transition_slot_size(rb_shape_transition_robject(0),
-                                                  rb_gc_size_slot_size(size));
+    *size_out = robject_embedded_size(index_tbl_num_entries);
+    *flags_out = T_OBJECT | rb_shape_transition_robject(0);
+
     return true;
 }
 #endif

hash.c

@@ -1467,12 +1467,8 @@ hash_alloc(VALUE klass)
 size_t
 rb_zjit_hash_new_size(VALUE *flags_out)
 {
-    size_t size = hash_slot_size(sizeof(st_table) > sizeof(ar_table));
-    // mimic rb_newobj()
-    shape_id_t shape_id = rb_shape_transition_slot_size(ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_OTHER,
-                                                        rb_gc_size_slot_size(size));
-    *flags_out = T_HASH | ((VALUE)shape_id << SHAPE_FLAG_SHIFT);
-    return size;
+    *flags_out = T_HASH;
+    return hash_slot_size(sizeof(st_table) > sizeof(ar_table));
 }
 #endif
 

range.c

@@ -87,20 +87,13 @@ void
 rb_zjit_range_new_fastpath(bool exclude_end, size_t *alloc_size_out, VALUE *flags_out)
 {
     const long len = 2;
-    size_t size = offsetof(struct RStruct, as.ary) + (sizeof(VALUE) * len);
+    *alloc_size_out = offsetof(struct RStruct, as.ary) + (sizeof(VALUE) * len);
     if (RCLASS_MAX_IV_COUNT(rb_cRange) > 0) {
-        size += sizeof(VALUE);
+        *alloc_size_out += sizeof(VALUE);
     }
 
-    VALUE flags = T_STRUCT | (len << RSTRUCT_EMBED_LEN_SHIFT) | RANGE_FL_INIT | FL_FREEZE;
-    if (exclude_end) flags |= RANGE_FL_EXCL;
-
-    shape_id_t shape_id = rb_shape_transition_slot_size(ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_EXTENDED,
-                                                         rb_gc_size_slot_size(size));
-    shape_id = rb_shape_transition_frozen(shape_id);
-
-    *alloc_size_out = size;
-    *flags_out = flags | ((VALUE)shape_id << SHAPE_FLAG_SHIFT);
+    *flags_out = T_STRUCT | (len << RSTRUCT_EMBED_LEN_SHIFT) | RANGE_FL_INIT | FL_FREEZE;
+    if (exclude_end) *flags_out |= RANGE_FL_EXCL;
 }
 #endif
 

string.c

@@ -2070,11 +2070,8 @@ rb_zjit_str_resurrect_fastpath(VALUE str, bool chilled, size_t *size_out,
     flags |= T_STRING;
     if (chilled) flags |= STR_CHILLED;
 
-    shape_id_t shape_id = rb_shape_transition_slot_size(ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_OTHER,
-                                                        rb_gc_size_slot_size(size));
-
     *size_out = size;
-    *flags_out = flags | ((VALUE)shape_id << SHAPE_FLAG_SHIFT);
+    *flags_out = flags;
     *len_out = len;
     *byte_size_out = (size_t)(len + termlen);
     return true;

zjit.c

@@ -24,6 +24,7 @@
 #include "ruby/debug.h"
 #include "internal/cont.h"
 #include "ractor_core.h"
+#include "shape.h"
 
 // This build config impacts the pointer tagging scheme and we only want to
 // support one scheme for simplicity.
@@ -184,6 +185,37 @@ rb_zjit_offset_ractor_objspace(void)
     return offsetof(rb_ractor_t, objspace);
 }
 
+/* Sets all of the required shape flags for the object including the layout type,
+ * the frozen status, and the slot size.
+ */
+VALUE
+rb_zjit_new_obj_shape(VALUE flags, size_t alloc_size)
+{
+    shape_id_t shape_id;
+    switch (flags & T_MASK) {
+      case T_OBJECT:
+        shape_id = ROOT_SHAPE_ID;
+        break;
+      case T_STRUCT:
+        shape_id = ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_EXTENDED;
+        break;
+      case T_DATA:
+        shape_id = ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_RDATA;
+        break;
+      default:
+        shape_id = ROOT_SHAPE_ID | SHAPE_ID_LAYOUT_OTHER;
+        break;
+    }
+
+    if (flags & FL_FREEZE) {
+        shape_id = rb_shape_transition_frozen(shape_id);
+    }
+
+    shape_id = rb_shape_transition_slot_size(shape_id, rb_gc_size_slot_size(alloc_size));
+
+    return (flags & SHAPE_FLAG_MASK) | ((VALUE)shape_id << SHAPE_FLAG_SHIFT);
+}
+
 VALUE
 rb_zjit_defined_ivar(VALUE obj, ID id, VALUE pushval)
 {

zjit.h

@@ -94,7 +94,8 @@ void rb_zjit_jit_frame_update_references(zjit_jit_frame_t *jit_frame);
 void rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp);
 void rb_zjit_materialize_frames_for_longjmp(const rb_execution_context_t *ec, rb_control_frame_t *cfp);
 size_t rb_zjit_hash_new_size(VALUE *flags_out);
-bool rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, shape_id_t *shape_id_out);
+VALUE rb_zjit_new_obj_shape(VALUE flags, size_t alloc_size);
+bool rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, VALUE *flags_out);
 bool rb_zjit_str_resurrect_fastpath(VALUE str, bool chilled, size_t *size_out, VALUE *flags_out, long *len_out, size_t *byte_size_out);
 bool rb_zjit_array_dup_can_fastpath(VALUE ary, size_t *alloc_size_out, VALUE *flags_out, long *len_out);
 void rb_zjit_range_new_fastpath(bool exclude_end, size_t *alloc_size_out, VALUE *flags_out);

zjit/bindgen/src/main.rs

@@ -172,6 +172,7 @@ fn main() {
         .allowlist_function("rb_gc_writebarrier")
         .allowlist_function("rb_gc_writebarrier_remember")
         .allowlist_function("rb_gc_register_mark_object")
+        .allowlist_function("rb_zjit_new_obj_shape")
 
         // VALUE variables for Ruby class objects
         .allowlist_var("rb_cBasicObject")

zjit/src/codegen.rs

@@ -2616,13 +2616,12 @@ fn gen_object_alloc_class(jit: &mut JITState, asm: &mut Assembler, class: VALUE,
     gen_prepare_leaf_call_with_gc(asm, state);
     if unsafe { rb_zjit_class_has_default_allocator(class) } {
         let mut alloc_size: usize = 0;
-        let mut shape_id: shape_id_t = 0;
+        let mut flags = VALUE(0);
         let has_fastpath = unsafe {
-            rb_zjit_class_allocate_instance_fastpath(class, &mut alloc_size, &mut shape_id)
+            rb_zjit_class_allocate_instance_fastpath(class, &mut alloc_size, &mut flags)
         };
         if has_fastpath {
-            let flags = (RUBY_T_OBJECT as u64) | ((shape_id as u64) << RB_SHAPE_FLAG_SHIFT as u64);
-            gc_fastpath::gc_fastpath_new_obj(jit, asm, alloc_size, flags, class, |_asm, _obj| {}, |asm| {
+            gc_fastpath::gc_fastpath_new_obj(jit, asm, alloc_size, flags.as_u64(), class, |_asm, _obj| {}, |asm| {
                 asm_ccall!(asm, rb_class_allocate_instance, class.into())
             })
         } else {

zjit/src/codegen/gc_fastpath.rs

@@ -5,7 +5,7 @@ use crate::cruby::{
     RB_GC_ZJIT_FASTPATH_DEFAULT, RB_GC_ZJIT_FASTPATH_MMTK,
     RUBY_OFFSET_EC_THREAD_PTR, RUBY_OFFSET_RBASIC_FLAGS, RUBY_OFFSET_RBASIC_KLASS,
     RUBY_OFFSET_THREAD_RACTOR, VALUE, VALUE_BITS, rb_zjit_offset_ractor_newobj_cache,
-    rb_zjit_offset_ractor_objspace,
+    rb_zjit_offset_ractor_objspace, rb_zjit_new_obj_shape,
 };
 use super::JITState;
 
@@ -79,6 +79,8 @@ pub(super) fn gc_fastpath_new_obj(
     init: impl Fn(&mut Assembler, Opnd),
     slow_path: impl Fn(&mut Assembler) -> lir::Opnd,
 ) -> lir::Opnd {
+    let flags = unsafe { rb_zjit_new_obj_shape(VALUE(flags as usize), alloc_size) }.as_u64();
+
     let Some(fastpath) = prepare_new_obj_fastpath(alloc_size, flags, klass) else {
         return slow_path(asm);
     };

zjit/src/cruby_bindings.inc.rs

@@ -2274,10 +2274,11 @@ unsafe extern "C" {
     pub fn rb_iseq_defined_string(type_: defined_type) -> VALUE;
     pub fn rb_zjit_profile_enable(iseq: *const rb_iseq_t);
     pub fn rb_zjit_hash_new_size(flags_out: *mut VALUE) -> usize;
+    pub fn rb_zjit_new_obj_shape(flags: VALUE, alloc_size: usize) -> VALUE;
     pub fn rb_zjit_class_allocate_instance_fastpath(
         klass: VALUE,
         size_out: *mut usize,
-        shape_id_out: *mut shape_id_t,
+        flags_out: *mut VALUE,
     ) -> bool;
     pub fn rb_zjit_str_resurrect_fastpath(
         str_: VALUE,