Use shape bits to store capacity of objectHEADmaster

This commit uses 7 shape bits to store the capacity of the object instead
of the heap ID. This allows for slot sizes with up to capacity of 127.
This removes the abstraction leak of slot sizes from the GC, which allows
more flexibility in the GC.

This implementation should currently make no difference for the default
GC as it will end up creating the same root shapes. We can see that there
is basically no change in benchmarks:

--------------  ------------  ------------  ------------  ------------  --------------  -------------  -----------------
bench            master (ms)     RSS (MiB)   branch (ms)     RSS (MiB)  branch 1st itr  master/branch  RSS master/branch
activerecord     99.6 ± 3.9%   75.5 ± 0.3%   98.9 ± 2.4%   74.0 ± 1.3%           0.919          1.007              1.021
chunky-png      328.2 ± 0.6%   83.0 ± 3.9%  336.4 ± 0.7%   87.0 ± 2.2%           0.979          0.976              0.954
erubi-rails     440.4 ± 1.7%  138.1 ± 0.3%  431.0 ± 0.5%  135.2 ± 0.0%           0.992          1.022              1.022
hexapdf         855.6 ± 0.7%  577.0 ± 0.1%  856.7 ± 0.9%  649.0 ± 2.6%           0.993          0.999              0.889
liquid-c         21.6 ± 8.2%  69.6 ± 14.7%   23.2 ± 8.4%  64.6 ± 15.4%           0.993          0.930              1.077
liquid-compile   19.9 ± 7.9%   47.0 ± 5.7%   20.2 ± 8.7%   46.9 ± 5.1%           0.925          0.987              1.001
liquid-render    54.0 ± 2.7%   55.3 ± 8.6%   55.2 ± 2.7%   55.9 ± 9.0%           0.977          0.978              0.990
lobsters        351.6 ± 0.7%  336.8 ± 0.2%  357.0 ± 0.6%  335.4 ± 0.1%           0.997          0.985              1.004
mail             47.5 ± 3.1%   75.6 ± 2.7%   46.8 ± 1.3%   72.3 ± 0.5%           0.974          1.016              1.045
psych-load      829.1 ± 1.2%   55.0 ± 0.4%  832.9 ± 0.5%   55.3 ± 0.7%           1.006          0.995              0.995
railsbench      713.8 ± 0.4%  134.6 ± 1.1%  717.9 ± 1.0%  138.2 ± 1.2%           0.992          0.994              0.974
rubocop          73.1 ± 6.0%  106.7 ± 1.7%   73.2 ± 6.1%  106.4 ± 1.4%           0.997          0.998              1.003
ruby-lsp         69.0 ± 1.2%   78.3 ± 0.2%   67.8 ± 2.3%   74.2 ± 0.0%           1.043          1.017              1.055
sequel           21.7 ± 6.8%   56.1 ± 1.6%   21.0 ± 4.4%   55.9 ± 1.1%           0.873          1.032              1.003
shipit          608.9 ± 1.4%  160.2 ± 0.5%  619.2 ± 1.4%  157.3 ± 0.2%           1.012          0.983              1.018
--------------  ------------  ------------  ------------  ------------  --------------  -------------  -----------------

13 files changed, 139 insertions, 191 deletions

diff --git a/gc.c b/gc.c
index 1baea0cdbe..bbc4a6d8e3 100644
--- a/gc.c
+++ b/gc.c
@@ -375,11 +375,11 @@ rb_gc_shutdown_call_finalizer_p(VALUE obj)
 }
 
 void
-rb_gc_obj_changed_pool(VALUE obj, size_t heap_id)
+rb_gc_obj_changed_slot_size(VALUE obj, size_t slot_size)
 {
     RUBY_ASSERT(RB_TYPE_P(obj, T_OBJECT));
 
-    RBASIC_SET_SHAPE_ID(obj, rb_obj_shape_transition_heap(obj, heap_id));
+    RBASIC_SET_SHAPE_ID_WITH_CAPACITY(obj, rb_obj_shape_transition_capacity(obj, rb_shape_capacity_for_slot_size(slot_size)));
 }
 
 void rb_vm_update_references(void *ptr);
@@ -604,7 +604,6 @@ typedef struct gc_function_map {
     void *(*ractor_cache_alloc)(void *objspace_ptr, void *ractor);
     void (*set_params)(void *objspace_ptr);
     void (*init)(void);
-    size_t *(*heap_sizes)(void *objspace_ptr);
     // Shutdown
     void (*shutdown_free_objects)(void *objspace_ptr);
     void (*objspace_free)(void *objspace_ptr);
@@ -622,9 +621,9 @@ typedef struct gc_function_map {
     VALUE (*stress_get)(void *objspace_ptr);
     struct rb_gc_vm_context *(*get_vm_context)(void *objspace_ptr);
     // Object allocation
-    VALUE (*new_obj)(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size);
+    VALUE (*new_obj)(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size, size_t *actual_alloc_size);
     size_t (*obj_slot_size)(VALUE obj);
-    size_t (*heap_id_for_size)(void *objspace_ptr, size_t size);
+    size_t (*size_slot_size)(void *objspace_ptr, size_t size);
     bool (*size_allocatable_p)(size_t size);
     // Malloc
     void *(*malloc)(void *objspace_ptr, size_t size, bool gc_allowed);
@@ -783,7 +782,6 @@ ruby_modular_gc_init(void)
     load_modular_gc_func(ractor_cache_alloc);
     load_modular_gc_func(set_params);
     load_modular_gc_func(init);
-    load_modular_gc_func(heap_sizes);
     // Shutdown
     load_modular_gc_func(shutdown_free_objects);
     load_modular_gc_func(objspace_free);
@@ -803,7 +801,7 @@ ruby_modular_gc_init(void)
     // Object allocation
     load_modular_gc_func(new_obj);
     load_modular_gc_func(obj_slot_size);
-    load_modular_gc_func(heap_id_for_size);
+    load_modular_gc_func(size_slot_size);
     load_modular_gc_func(size_allocatable_p);
     // Malloc
     load_modular_gc_func(malloc);
@@ -871,7 +869,6 @@ ruby_modular_gc_init(void)
 # define rb_gc_impl_ractor_cache_alloc rb_gc_functions.ractor_cache_alloc
 # define rb_gc_impl_set_params rb_gc_functions.set_params
 # define rb_gc_impl_init rb_gc_functions.init
-# define rb_gc_impl_heap_sizes rb_gc_functions.heap_sizes
 // Shutdown
 # define rb_gc_impl_shutdown_free_objects rb_gc_functions.shutdown_free_objects
 # define rb_gc_impl_objspace_free rb_gc_functions.objspace_free
@@ -891,7 +888,7 @@ ruby_modular_gc_init(void)
 // Object allocation
 # define rb_gc_impl_new_obj rb_gc_functions.new_obj
 # define rb_gc_impl_obj_slot_size rb_gc_functions.obj_slot_size
-# define rb_gc_impl_heap_id_for_size rb_gc_functions.heap_id_for_size
+# define rb_gc_impl_size_slot_size rb_gc_functions.size_slot_size
 # define rb_gc_impl_size_allocatable_p rb_gc_functions.size_allocatable_p
 // Malloc
 # define rb_gc_impl_malloc rb_gc_functions.malloc
@@ -1033,7 +1030,11 @@ rb_newobj(rb_execution_context_t *ec, VALUE klass, VALUE flags, shape_id_t shape
 {
     GC_ASSERT((flags & FL_WB_PROTECTED) == 0);
     rb_ractor_t *cr = rb_ec_ractor_ptr(ec);
-    VALUE obj = rb_gc_impl_new_obj(rb_gc_get_objspace(), cr->newobj_cache, klass, flags, wb_protected, size);
+    size_t actual_alloc_size;
+    VALUE obj = rb_gc_impl_new_obj(rb_gc_get_objspace(), cr->newobj_cache, klass, flags, wb_protected, size, &actual_alloc_size);
+
+    GC_ASSERT(actual_alloc_size >= size);
+    shape_id = (shape_id & ~SHAPE_ID_CAPACITY_MASK) | rb_shape_id_with_capacity(rb_shape_capacity_for_slot_size(actual_alloc_size));
 
 #if RACTOR_CHECK_MODE
     void rb_ractor_setup_belonging(VALUE obj);
@@ -1092,7 +1093,7 @@ rb_newobj_of(VALUE klass, VALUE flags, size_t size)
 static
 VALUE class_allocate_complex_instance(VALUE klass, uint32_t capacity)
 {
-    shape_id_t initial_shape_id = rb_shape_id_with_robject_layout(rb_shape_root(rb_gc_heap_id_for_size(sizeof(struct RObject))));
+    shape_id_t initial_shape_id = rb_shape_id_with_robject_layout(0);
     VALUE obj = rb_newobj_of_with_shape(klass, T_OBJECT, initial_shape_id, sizeof(struct RObject));
     rb_obj_init_complex(obj, rb_st_init_numtable_with_size(capacity));
     return obj;
@@ -1105,8 +1106,8 @@ rb_class_allocate_instance(VALUE klass)
     VALUE obj;
 
     // Directly start as COMPLEX if we know we're over the limit.
-    RUBY_ASSERT(rb_shape_tree.max_capacity > 0);
-    if (RB_UNLIKELY(index_tbl_num_entries > rb_shape_tree.max_capacity)) {
+    RUBY_ASSERT(SHAPE_MAX_CAPACITY > 0);
+    if (RB_UNLIKELY(index_tbl_num_entries > SHAPE_MAX_CAPACITY)) {
         obj = class_allocate_complex_instance(klass, index_tbl_num_entries);
     }
     else {
@@ -1117,7 +1118,7 @@ rb_class_allocate_instance(VALUE klass)
 
         // There might be a NEWOBJ tracepoint callback, and it may set fields.
         // So the shape must be passed to `NEWOBJ_OF`.
-        obj = rb_newobj_of_with_shape(klass, T_OBJECT, rb_shape_id_with_robject_layout(rb_shape_root(rb_gc_heap_id_for_size(size))), size);
+        obj = rb_newobj_of_with_shape(klass, T_OBJECT, rb_shape_id_with_robject_layout(0), size);
 
         #if RUBY_DEBUG
             VALUE *ptr = ROBJECT_FIELDS(obj);
@@ -2187,7 +2188,6 @@ object_id0(VALUE obj)
     id = generate_next_object_id();
     rb_obj_field_set(obj, object_id_shape_id, 0, id);
 
-    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == object_id_shape_id);
     RUBY_ASSERT(rb_obj_shape_has_id(obj));
 
     if (RB_UNLIKELY(id2ref_tbl)) {
@@ -3989,9 +3989,9 @@ rb_gc_prepare_heap(void)
 }
 
 size_t
-rb_gc_heap_id_for_size(size_t size)
+rb_gc_size_slot_size(size_t size)
 {
-    return rb_gc_impl_heap_id_for_size(rb_gc_get_objspace(), size);
+    return rb_gc_impl_size_slot_size(rb_gc_get_objspace(), size);
 }
 
 bool
@@ -4704,12 +4704,6 @@ rb_gc_initial_stress_set(VALUE flag)
     initial_stress = flag;
 }
 
-size_t *
-rb_gc_heap_sizes(void)
-{
-    return rb_gc_impl_heap_sizes(rb_gc_get_objspace());
-}
-
 VALUE
 rb_gc_enable(void)
 {
diff --git a/gc/default/default.c b/gc/default/default.c
index 4d4bd1af25..b39a1bd855 100644
--- a/gc/default/default.c
+++ b/gc/default/default.c
@@ -2610,24 +2610,9 @@ heap_idx_for_size(size_t size)
 }
 
 size_t
-rb_gc_impl_heap_id_for_size(void *objspace_ptr, size_t size)
+rb_gc_impl_size_slot_size(void *objspace_ptr, size_t size)
 {
-    return heap_idx_for_size(size);
-}
-
-
-static size_t heap_sizes[HEAP_COUNT + 1] = { 0 };
-
-size_t *
-rb_gc_impl_heap_sizes(void *objspace_ptr)
-{
-    if (heap_sizes[0] == 0) {
-        for (unsigned char i = 0; i < HEAP_COUNT; i++) {
-            heap_sizes[i] = heap_slot_size(i);
-        }
-    }
-
-    return heap_sizes;
+    return heap_slot_size((unsigned char)heap_idx_for_size(size));
 }
 
 NOINLINE(static VALUE newobj_cache_miss(rb_objspace_t *objspace, rb_ractor_newobj_cache_t *cache, size_t heap_idx, bool vm_locked));
@@ -2740,7 +2725,7 @@ newobj_slowpath_wb_unprotected(VALUE klass, VALUE flags, rb_objspace_t *objspace
 }
 
 VALUE
-rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size)
+rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size, size_t *actual_alloc_size)
 {
     VALUE obj;
     rb_objspace_t *objspace = objspace_ptr;
@@ -2755,6 +2740,7 @@ rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags
     }
 
     size_t heap_idx = heap_idx_for_size(alloc_size);
+    *actual_alloc_size = heap_slot_size((unsigned char)heap_idx);
 
     rb_ractor_newobj_cache_t *cache = (rb_ractor_newobj_cache_t *)cache_ptr;
 
@@ -7483,7 +7469,7 @@ gc_move(rb_objspace_t *objspace, VALUE src, VALUE dest, struct heap_page *src_pa
     memcpy((void *)dest, (void *)src, MIN(src_slot_size, slot_size));
 
     if (src_slot_size != slot_size && RB_TYPE_P(src, T_OBJECT)) {
-        rb_gc_obj_changed_pool(dest, dest_page->heap - heaps);
+        rb_gc_obj_changed_slot_size(dest, slot_size - RVALUE_OVERHEAD);
     }
 
     if (RVALUE_OVERHEAD > 0) {
diff --git a/gc/gc.h b/gc/gc.h
index d147a3122a..2a20d65c56 100644
--- a/gc/gc.h
+++ b/gc/gc.h
@@ -110,7 +110,7 @@ MODULAR_GC_FN void rb_gc_mark_roots(void *objspace, const char **categoryp);
 MODULAR_GC_FN void rb_gc_ractor_newobj_cache_foreach(void (*func)(void *cache, void *data), void *data);
 MODULAR_GC_FN bool rb_gc_multi_ractor_p(void);
 MODULAR_GC_FN bool rb_gc_shutdown_call_finalizer_p(VALUE obj);
-MODULAR_GC_FN void rb_gc_obj_changed_pool(VALUE obj, size_t heap_id);
+MODULAR_GC_FN void rb_gc_obj_changed_slot_size(VALUE obj, size_t slot_size);
 MODULAR_GC_FN void rb_gc_prepare_heap_process_object(VALUE obj);
 MODULAR_GC_FN bool rb_memerror_reentered(void);
 MODULAR_GC_FN bool rb_obj_id_p(VALUE);
diff --git a/gc/gc_impl.h b/gc/gc_impl.h
index b8417fd03a..892831c8c1 100644
--- a/gc/gc_impl.h
+++ b/gc/gc_impl.h
@@ -38,7 +38,6 @@ GC_IMPL_FN void rb_gc_impl_objspace_init(void *objspace_ptr);
 GC_IMPL_FN void *rb_gc_impl_ractor_cache_alloc(void *objspace_ptr, void *ractor);
 GC_IMPL_FN void rb_gc_impl_set_params(void *objspace_ptr);
 GC_IMPL_FN void rb_gc_impl_init(void);
-GC_IMPL_FN size_t *rb_gc_impl_heap_sizes(void *objspace_ptr);
 // Shutdown
 GC_IMPL_FN void rb_gc_impl_shutdown_free_objects(void *objspace_ptr);
 GC_IMPL_FN void rb_gc_impl_objspace_free(void *objspace_ptr);
@@ -56,9 +55,9 @@ GC_IMPL_FN VALUE rb_gc_impl_config_get(void *objspace_ptr);
 GC_IMPL_FN void rb_gc_impl_config_set(void *objspace_ptr, VALUE hash);
 GC_IMPL_FN struct rb_gc_vm_context *rb_gc_impl_get_vm_context(void *objspace_ptr);
 // Object allocation
-GC_IMPL_FN VALUE rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size);
+GC_IMPL_FN VALUE rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size, size_t *actual_alloc_size);
 GC_IMPL_FN size_t rb_gc_impl_obj_slot_size(VALUE obj);
-GC_IMPL_FN size_t rb_gc_impl_heap_id_for_size(void *objspace_ptr, size_t size);
+GC_IMPL_FN size_t rb_gc_impl_size_slot_size(void *objspace_ptr, size_t size);
 GC_IMPL_FN bool rb_gc_impl_size_allocatable_p(size_t size);
 // Malloc
 /*
diff --git a/gc/mmtk/mmtk.c b/gc/mmtk/mmtk.c
index d996b92352..3600a3d2e3 100644
--- a/gc/mmtk/mmtk.c
+++ b/gc/mmtk/mmtk.c
@@ -687,12 +687,6 @@ rb_gc_impl_init(void)
     rb_define_singleton_method(rb_mGC, "verify_compaction_references", rb_f_notimplement, -1);
 }
 
-size_t *
-rb_gc_impl_heap_sizes(void *objspace_ptr)
-{
-    return heap_sizes;
-}
-
 int
 rb_mmtk_obj_free_iter_wrapper(VALUE obj, void *data)
 {
@@ -902,7 +896,7 @@ mmtk_post_alloc_fast_immix(struct objspace *objspace, struct MMTk_ractor_cache *
 }
 
 VALUE
-rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size)
+rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size, size_t *actual_alloc_size)
 {
 #define MMTK_ALLOCATION_SEMANTICS_DEFAULT 0
     struct objspace *objspace = objspace_ptr;
@@ -916,6 +910,7 @@ rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags
             break;
         }
     }
+    *actual_alloc_size = alloc_size;
 
     if (objspace->gc_stress) {
         mmtk_handle_user_collection_request(ractor_cache, false, false);
@@ -962,10 +957,10 @@ rb_gc_impl_obj_slot_size(VALUE obj)
 }
 
 size_t
-rb_gc_impl_heap_id_for_size(void *objspace_ptr, size_t size)
+rb_gc_impl_size_slot_size(void *objspace_ptr, size_t size)
 {
     for (int i = 0; i < MMTK_HEAP_COUNT; i++) {
-        if (size <= heap_sizes[i]) return i;
+        if (size <= heap_sizes[i]) return heap_sizes[i];
     }
 
     rb_bug("size too big");
diff --git a/gc/wbcheck/wbcheck.c b/gc/wbcheck/wbcheck.c
index 08b644f4bd..5da87f5762 100644
--- a/gc/wbcheck/wbcheck.c
+++ b/gc/wbcheck/wbcheck.c
@@ -536,12 +536,6 @@ rb_gc_impl_init(void)
     // Stub implementation
 }
 
-size_t *
-rb_gc_impl_heap_sizes(void *objspace_ptr)
-{
-    return heap_sizes;
-}
-
 // Shutdown
 void
 rb_gc_impl_shutdown_free_objects(void *objspace_ptr)
@@ -1075,7 +1069,7 @@ lock_and_maybe_gc(void *objspace_ptr)
 }
 
 VALUE
-rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size)
+rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags, bool wb_protected, size_t alloc_size, size_t *actual_alloc_size)
 {
     unsigned int lev = RB_GC_VM_LOCK();
     rb_gc_vm_barrier();
@@ -1084,12 +1078,14 @@ rb_gc_impl_new_obj(void *objspace_ptr, void *cache_ptr, VALUE klass, VALUE flags
     maybe_gc(objspace_ptr);
 
     // Ensure minimum allocation size of BASE_SLOT_SIZE
-    alloc_size = heap_sizes[rb_gc_impl_heap_id_for_size(objspace_ptr, alloc_size)];
+    alloc_size = rb_gc_impl_size_slot_size(objspace_ptr, alloc_size);
 
     // Allocate memory for the object
     VALUE *mem = malloc(alloc_size);
     if (!mem) rb_bug("FIXME: malloc failed");
 
+    *actual_alloc_size = alloc_size;
+
     // Initialize the object
     VALUE obj = (VALUE)mem;
     RBASIC(obj)->flags = flags;
@@ -1119,10 +1115,10 @@ rb_gc_impl_obj_slot_size(VALUE obj)
 }
 
 size_t
-rb_gc_impl_heap_id_for_size(void *objspace_ptr, size_t size)
+rb_gc_impl_size_slot_size(void *objspace_ptr, size_t size)
 {
     for (int i = 0; i < HEAP_COUNT; i++) {
-        if (size <= heap_sizes[i]) return i;
+        if (size <= heap_sizes[i]) return heap_sizes[i];
     }
     rb_bug("size too big");
 }
@@ -1952,4 +1948,3 @@ rb_gc_impl_copy_attributes(void *objspace_ptr, VALUE dest, VALUE obj)
     }
     rb_gc_impl_copy_finalizer(objspace_ptr, dest, obj);
 }
-
diff --git a/imemo.c b/imemo.c
index 3814fbd453..a823b74b5b 100644
--- a/imemo.c
+++ b/imemo.c
@@ -133,7 +133,7 @@ rb_imemo_cdhash_new(size_t size, const struct st_hash_type *type)
 VALUE
 rb_imemo_fields_new(VALUE owner, shape_id_t shape_id, bool shareable)
 {
-    size_t capa = RSHAPE_CAPACITY(shape_id);
+    size_t capa = RSHAPE(shape_id)->capacity;
     size_t embedded_size = offsetof(struct rb_fields, as.embed) + capa * sizeof(VALUE);
     RUBY_ASSERT(rb_gc_size_allocatable_p(embedded_size));
     VALUE fields = rb_imemo_new(imemo_fields, owner, embedded_size, shareable);
diff --git a/internal/gc.h b/internal/gc.h
index 8e85381f09..6d1df9363f 100644
--- a/internal/gc.h
+++ b/internal/gc.h
@@ -205,8 +205,7 @@ void *rb_gc_ractor_cache_alloc(rb_ractor_t *ractor);
 void rb_gc_ractor_cache_free(void *cache);
 
 bool rb_gc_size_allocatable_p(size_t size);
-size_t *rb_gc_heap_sizes(void);
-size_t rb_gc_heap_id_for_size(size_t size);
+size_t rb_gc_size_slot_size(size_t size);
 
 void rb_gc_mark_and_move(VALUE *ptr);
 
diff --git a/shape.c b/shape.c
index e6f4726a1d..875b62099b 100644
--- a/shape.c
+++ b/shape.c
@@ -411,12 +411,13 @@ rb_obj_shape_id(VALUE obj)
         VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
         shape_id_t base = ROOT_SHAPE_ID;
         if (fields_obj) {
-            // Remove the layout from the fields object.  We want to
-            // combine the shape of the fields object with the layout of the
-            // class / module object.
-            base = RBASIC_SHAPE_ID(fields_obj) & ~SHAPE_ID_LAYOUT_MASK;
+            // Remove the layout and capacity from the fields object. We want to
+            // combine the shape tree state of the fields object with the layout
+            // and object slot capacity of the class / module object.
+            base = RBASIC_SHAPE_ID(fields_obj) & ~(SHAPE_ID_LAYOUT_MASK | SHAPE_ID_CAPACITY_MASK);
         }
-        return rb_shape_layout(RBASIC_SHAPE_ID(obj)) | base;
+        shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
+        return rb_shape_layout(shape_id) | (shape_id & SHAPE_ID_CAPACITY_MASK) | base;
     }
     return RBASIC_SHAPE_ID(obj);
 }
@@ -513,17 +514,14 @@ redblack_cache_ancestors(rb_shape_t *shape)
 static attr_index_t
 shape_grow_capa(attr_index_t current_capa)
 {
-    const attr_index_t *capacities = rb_shape_tree.capacities;
-    size_t heaps_count = rb_shape_tree.heaps_count;
-
-    // First try to use the next size that will be embeddable in a larger object slot.
-    for (size_t i = 0; i < heaps_count; i++) {
-        attr_index_t capa = capacities[i];
-        if (capa > current_capa) {
-            return capa;
-        }
+    size_t next_size = rb_obj_embedded_size(current_capa + 1);
+    if (UNLIKELY(!rb_gc_size_allocatable_p(next_size))) {
+        return SHAPE_MAX_CAPACITY;
     }
-    return capacities[rb_shape_tree.heaps_count - 1];
+
+    attr_index_t next_capa = rb_shape_capacity_for_slot_size(rb_gc_size_slot_size(next_size));
+    RUBY_ASSERT(next_capa > current_capa);
+    return next_capa;
 }
 
 static rb_shape_t *
@@ -702,7 +700,7 @@ rb_shape_transition_object_id(shape_id_t original_shape_id)
 
     bool dont_care;
     rb_shape_t *shape = NULL;
-    if (LIKELY(original_shape->next_field_index < rb_shape_tree.max_capacity)) {
+    if (LIKELY(original_shape->next_field_index < SHAPE_MAX_CAPACITY)) {
         shape = get_next_shape_internal(original_shape, id_object_id, SHAPE_OBJ_ID, &dont_care, true);
     }
     if (!shape) {
@@ -767,8 +765,8 @@ shape_get_next(rb_shape_t *shape, enum shape_type shape_type, VALUE klass, ID id
     }
 #endif
 
-    RUBY_ASSERT(rb_shape_tree.max_capacity > 0);
-    if (UNLIKELY(shape->next_field_index >= rb_shape_tree.max_capacity)) {
+    RUBY_ASSERT(SHAPE_MAX_CAPACITY > 0);
+    if (UNLIKELY(shape->next_field_index >= SHAPE_MAX_CAPACITY)) {
         return NULL;
     }
 
@@ -1306,19 +1304,14 @@ rb_shape_verify_consistency(VALUE obj, shape_id_t shape_id)
         }
     }
 
-    uint8_t flags_heap_index = rb_shape_heap_index(shape_id);
+    attr_index_t shape_id_capacity = rb_shape_capacity(shape_id);
     if (RB_TYPE_P(obj, T_OBJECT)) {
-        RUBY_ASSERT(flags_heap_index > 0);
-        size_t shape_id_slot_size = rb_shape_tree.capacities[flags_heap_index - 1] * sizeof(VALUE) + sizeof(struct RBasic);
+        RUBY_ASSERT(shape_id_capacity > 0);
+        size_t shape_id_slot_size = shape_id_capacity * sizeof(VALUE) + sizeof(struct RBasic);
         size_t actual_slot_size = rb_gc_obj_slot_size(obj);
 
         if (shape_id_slot_size != actual_slot_size) {
-            rb_bug("shape_id heap_index flags mismatch: shape_id_slot_size=%zu, gc_slot_size=%zu\n", shape_id_slot_size, actual_slot_size);
-        }
-    }
-    else {
-        if (flags_heap_index) {
-            rb_bug("shape_id indicate heap_index > 0 but object is not T_OBJECT: %s", rb_obj_info(obj));
+            rb_bug("shape_id capacity flags mismatch: shape_id_slot_size=%zu, gc_slot_size=%zu\n", shape_id_slot_size, actual_slot_size);
         }
     }
 
@@ -1395,7 +1388,7 @@ shape_id_t_to_rb_cShape(shape_id_t shape_id)
             INT2NUM(shape->parent_offset),
             rb_shape_edge_name(shape),
             INT2NUM(shape->next_field_index),
-            INT2NUM(rb_shape_heap_index(shape_id)),
+            INT2NUM(rb_shape_capacity(shape_id)),
             INT2NUM(shape->type),
             INT2NUM(RSHAPE_CAPACITY(shape_id)));
     rb_obj_freeze(obj);
@@ -1565,30 +1558,6 @@ rb_shape_find_by_id(VALUE mod, VALUE id)
 void
 Init_default_shapes(void)
 {
-    size_t *heap_sizes = rb_gc_heap_sizes();
-    size_t heaps_count = 0;
-    while (heap_sizes[heaps_count]) {
-        heaps_count++;
-    }
-
-    if (heaps_count > SHAPE_ID_HEAP_INDEX_MAX) {
-        rb_bug("Init_default_shapes initialized with %zu heaps, only up to %u are supported", heaps_count, SHAPE_ID_HEAP_INDEX_MAX);
-    }
-
-    size_t index;
-    for (index = 0; index < heaps_count; index++) {
-        if (heap_sizes[index] > sizeof(struct RBasic)) {
-            size_t capa = (heap_sizes[index] - sizeof(struct RBasic)) / sizeof(VALUE);
-            RUBY_ASSERT(capa < ATTR_INDEX_NOT_SET);
-            rb_shape_tree.capacities[index] = (attr_index_t)capa;
-        }
-        else {
-            rb_shape_tree.capacities[index] = 0;
-        }
-    }
-    rb_shape_tree.heaps_count = heaps_count;
-    rb_shape_tree.max_capacity = rb_shape_tree.capacities[heaps_count - 1];
-
 #ifdef HAVE_MMAP
     size_t shape_list_mmap_size = rb_size_mul_or_raise(SHAPE_BUFFER_SIZE, sizeof(rb_shape_t), rb_eRuntimeError);
     rb_shape_tree.shape_list = (rb_shape_t *)mmap(NULL, shape_list_mmap_size,
@@ -1664,7 +1633,7 @@ Init_shape(void)
             "parent_offset",
             "edge_name",
             "next_field_index",
-            "heap_index",
+            "embedded_capacity",
             "type",
             "capacity",
             NULL);
@@ -1682,7 +1651,7 @@ Init_shape(void)
     rb_define_const(rb_cShape, "SHAPE_ID_NUM_BITS", INT2NUM(SHAPE_ID_NUM_BITS));
     rb_define_const(rb_cShape, "SHAPE_FLAG_SHIFT", INT2NUM(SHAPE_FLAG_SHIFT));
     rb_define_const(rb_cShape, "SHAPE_MAX_VARIATIONS", INT2NUM(SHAPE_MAX_VARIATIONS));
-    rb_define_const(rb_cShape, "SHAPE_MAX_FIELDS", INT2NUM(rb_shape_tree.max_capacity));
+    rb_define_const(rb_cShape, "SHAPE_MAX_FIELDS", INT2NUM(SHAPE_MAX_CAPACITY));
     rb_define_const(rb_cShape, "SIZEOF_RB_SHAPE_T", INT2NUM(sizeof(rb_shape_t)));
     rb_define_const(rb_cShape, "SIZEOF_REDBLACK_NODE_T", INT2NUM(sizeof(redblack_node_t)));
     rb_define_const(rb_cShape, "SHAPE_BUFFER_SIZE", INT2NUM(sizeof(rb_shape_t) * SHAPE_BUFFER_SIZE));
diff --git a/shape.h b/shape.h
index 7df66f045f..ab126b5a89 100644
--- a/shape.h
+++ b/shape.h
@@ -14,32 +14,31 @@ STATIC_ASSERT(shape_id_num_bits, SHAPE_ID_NUM_BITS == sizeof(shape_id_t) * CHAR_
 #define SHAPE_BUFFER_SIZE (1 << SHAPE_ID_OFFSET_NUM_BITS)
 #define SHAPE_ID_OFFSET_MASK (SHAPE_BUFFER_SIZE - 1)
 
-#define SHAPE_ID_HEAP_INDEX_BITS 4
-#define SHAPE_ID_HEAP_INDEX_MAX (((attr_index_t)1 << SHAPE_ID_HEAP_INDEX_BITS) - 1)
+#define SHAPE_ID_CAPACITY_BITS 7
+#define SHAPE_ID_CAPACITY_MAX ((1U << SHAPE_ID_CAPACITY_BITS) - 1)
 
-#define SHAPE_ID_HEAP_INDEX_OFFSET SHAPE_ID_OFFSET_NUM_BITS
-#define SHAPE_ID_FL_USHIFT (SHAPE_ID_OFFSET_NUM_BITS + SHAPE_ID_HEAP_INDEX_BITS)
+#define SHAPE_ID_CAPACITY_OFFSET SHAPE_ID_OFFSET_NUM_BITS
+#define SHAPE_ID_FL_USHIFT (SHAPE_ID_OFFSET_NUM_BITS + SHAPE_ID_CAPACITY_BITS)
 
 // shape_id_t bits:
 //      0-18 SHAPE_ID_OFFSET_MASK
 //              index in rb_shape_tree.shape_list. Allow to access `rb_shape_t *`.
 //              This is the part that describe how fields are laid out in memory.
-//      19-22 SHAPE_ID_HEAP_INDEX_MASK
-//              index in rb_shape_tree.capacities. Allow to access slot size.
-//              Currently always 0 except for T_OBJECT.
-//      23 SHAPE_ID_FL_COMPLEX
+//      19-25 SHAPE_ID_CAPACITY_MASK
+//              Embedded field capacity for T_OBJECT objects.
+//      26 SHAPE_ID_FL_COMPLEX
 //              The object is backed by a `st_table`.
-//      24 SHAPE_ID_FL_FROZEN
+//      27 SHAPE_ID_FL_FROZEN
 //              Whether the object is frozen or not.
-//      25 SHAPE_ID_FL_HAS_OBJECT_ID
+//      28 SHAPE_ID_FL_HAS_OBJECT_ID
 //              Whether the object has an `SHAPE_OBJ_ID` transition.
-//      26-27 SHAPE_ID_LAYOUT_MASK
+//      29-30 SHAPE_ID_LAYOUT_MASK
 //              The object's physical field layout.
 
 enum shape_id_fl_type {
 #define RBIMPL_SHAPE_ID_FL(n) (1<<(SHAPE_ID_FL_USHIFT+n))
 
-    SHAPE_ID_HEAP_INDEX_MASK = ((1 << SHAPE_ID_HEAP_INDEX_BITS) - 1) << SHAPE_ID_HEAP_INDEX_OFFSET,
+    SHAPE_ID_CAPACITY_MASK = ((1 << SHAPE_ID_CAPACITY_BITS) - 1) << SHAPE_ID_CAPACITY_OFFSET,
 
     SHAPE_ID_FL_COMPLEX = RBIMPL_SHAPE_ID_FL(0),
     SHAPE_ID_FL_FROZEN = RBIMPL_SHAPE_ID_FL(1),
@@ -64,7 +63,7 @@ enum shape_id_fl_type {
     SHAPE_ID_LAYOUT_MASK = SHAPE_ID_LAYOUT_OTHER,
 
     SHAPE_ID_FL_NON_CANONICAL_MASK = SHAPE_ID_FL_FROZEN | SHAPE_ID_FL_HAS_OBJECT_ID,
-    SHAPE_ID_FLAGS_MASK = SHAPE_ID_HEAP_INDEX_MASK | SHAPE_ID_FL_NON_CANONICAL_MASK | SHAPE_ID_FL_COMPLEX | SHAPE_ID_LAYOUT_MASK,
+    SHAPE_ID_FLAGS_MASK = SHAPE_ID_CAPACITY_MASK | SHAPE_ID_FL_NON_CANONICAL_MASK | SHAPE_ID_FL_COMPLEX | SHAPE_ID_LAYOUT_MASK,
 
 #undef RBIMPL_SHAPE_ID_FL
 };
@@ -75,13 +74,13 @@ enum shape_id_mask {
     SHAPE_ID_HAS_IVAR_MASK = SHAPE_ID_FL_COMPLEX | (SHAPE_ID_OFFSET_MASK - 1),
 };
 
-// The interpreter doesn't care about frozen status, slot size, or object id, and
+// The interpreter doesn't care about frozen status, embedded capacity, or object id, and
 // has its own checks for physical field layout when reading ivars.
 // So we normalize shape_id by clearing these bits to improve cache hits.
 // JITs however might care about some of it.
-#define SHAPE_ID_READ_ONLY_MASK (~(SHAPE_ID_FL_FROZEN | SHAPE_ID_HEAP_INDEX_MASK | SHAPE_ID_FL_HAS_OBJECT_ID | SHAPE_ID_LAYOUT_MASK))
+#define SHAPE_ID_READ_ONLY_MASK (~(SHAPE_ID_FL_FROZEN | SHAPE_ID_CAPACITY_MASK | SHAPE_ID_FL_HAS_OBJECT_ID | SHAPE_ID_LAYOUT_MASK))
 // For write it's the same idea, but here we do care about frozen status.
-#define SHAPE_ID_WRITE_MASK (~(SHAPE_ID_HEAP_INDEX_MASK | SHAPE_ID_FL_HAS_OBJECT_ID | SHAPE_ID_LAYOUT_MASK))
+#define SHAPE_ID_WRITE_MASK (~(SHAPE_ID_CAPACITY_MASK | SHAPE_ID_FL_HAS_OBJECT_ID | SHAPE_ID_LAYOUT_MASK))
 
 typedef uint32_t redblack_id_t;
 
@@ -90,9 +89,18 @@ typedef uint32_t redblack_id_t;
 
 #define SHAPE_MAX_VARIATIONS 8
 
+// TODO: Fix this so that the max capacity does not depend on the CPU architecture
+#if RBASIC_SHAPE_ID_FIELD
+# define SHAPE_MAX_CAPACITY 125
+#else
+# define SHAPE_MAX_CAPACITY 126
+#endif
+
 #define INVALID_SHAPE_ID (SHAPE_BUFFER_SIZE - 1)
 #define ATTR_INDEX_NOT_SET ((attr_index_t)-1)
 
+STATIC_ASSERT(shape_max_capacity, SHAPE_MAX_CAPACITY <= SHAPE_ID_CAPACITY_MAX);
+
 #define ROOT_SHAPE_ID                   0x0
 #define ROOT_SHAPE_WITH_OBJ_ID          0x1
 #define ROOT_COMPLEX_SHAPE_ID       (ROOT_SHAPE_ID | SHAPE_ID_FL_COMPLEX)
@@ -126,9 +134,6 @@ enum shape_flags {
 
 typedef struct {
     rb_shape_t *shape_list;
-    attr_index_t max_capacity;
-    attr_index_t heaps_count;
-    attr_index_t capacities[SHAPE_ID_HEAP_INDEX_MAX];
 } rb_shape_tree_t;
 
 RUBY_SYMBOL_EXPORT_BEGIN
@@ -181,7 +186,7 @@ rb_shape_layout(shape_id_t shape_id)
 }
 
 static inline void
-RBASIC_SET_SHAPE_ID(VALUE obj, shape_id_t shape_id)
+RBASIC_SET_SHAPE_ID_WITH_CAPACITY(VALUE obj, shape_id_t shape_id)
 {
     RUBY_ASSERT(!RB_SPECIAL_CONST_P(obj));
     RUBY_ASSERT(!RB_TYPE_P(obj, T_IMEMO) || IMEMO_TYPE_P(obj, imemo_fields));
@@ -192,6 +197,15 @@ RBASIC_SET_SHAPE_ID(VALUE obj, shape_id_t shape_id)
     RUBY_ASSERT(rb_shape_verify_consistency(obj, shape_id));
 }
 
+static inline void
+RBASIC_SET_SHAPE_ID(VALUE obj, shape_id_t shape_id)
+{
+    RUBY_ASSERT(!RB_SPECIAL_CONST_P(obj));
+
+    shape_id = (shape_id & ~SHAPE_ID_CAPACITY_MASK) | (RBASIC_SHAPE_ID(obj) & SHAPE_ID_CAPACITY_MASK);
+    RBASIC_SET_SHAPE_ID_WITH_CAPACITY(obj, shape_id);
+}
+
 static inline shape_id_t
 RSHAPE_FLAGS(shape_id_t shape_id)
 {
@@ -265,22 +279,30 @@ rb_shape_id_with_robject_layout(shape_id_t shape_id)
     return (shape_id & ~SHAPE_ID_LAYOUT_MASK) | SHAPE_ID_LAYOUT_ROBJECT;
 }
 
-static inline uint8_t
-rb_shape_heap_index(shape_id_t shape_id)
+static inline attr_index_t
+rb_shape_capacity(shape_id_t shape_id)
 {
-    return (uint8_t)((shape_id & SHAPE_ID_HEAP_INDEX_MASK) >> SHAPE_ID_HEAP_INDEX_OFFSET);
+    return (attr_index_t)((shape_id & SHAPE_ID_CAPACITY_MASK) >> SHAPE_ID_CAPACITY_OFFSET);
 }
 
 static inline shape_id_t
-rb_shape_root(size_t heap_id)
+rb_shape_id_with_capacity(size_t capacity)
 {
-    shape_id_t heap_index = (shape_id_t)(heap_id + 1);
-    shape_id_t heap_flags = heap_index << SHAPE_ID_HEAP_INDEX_OFFSET;
+    shape_id_t capacity_flags = (shape_id_t)capacity << SHAPE_ID_CAPACITY_OFFSET;
+
+    RUBY_ASSERT(capacity <= SHAPE_ID_CAPACITY_MAX);
+    RUBY_ASSERT((capacity_flags & SHAPE_ID_CAPACITY_MASK) == capacity_flags);
+    RUBY_ASSERT(rb_shape_capacity(capacity_flags) == capacity);
 
-    RUBY_ASSERT((heap_flags & SHAPE_ID_HEAP_INDEX_MASK) == heap_flags);
-    RUBY_ASSERT(rb_shape_heap_index(heap_flags) == heap_index);
+    return ROOT_SHAPE_ID | capacity_flags;
+}
 
-    return ROOT_SHAPE_ID | heap_flags;
+static inline attr_index_t
+rb_shape_capacity_for_slot_size(size_t slot_size)
+{
+    size_t capacity = (slot_size - sizeof(struct RBasic)) / sizeof(VALUE);
+    RUBY_ASSERT(capacity <= SHAPE_ID_CAPACITY_MAX);
+    return (attr_index_t)capacity;
 }
 
 static inline shape_id_t
@@ -308,19 +330,9 @@ RSHAPE_TYPE_P(shape_id_t shape_id, enum shape_type type)
 }
 
 static inline attr_index_t
-RSHAPE_EMBEDDED_CAPACITY(shape_id_t shape_id)
-{
-    uint8_t heap_index = rb_shape_heap_index(shape_id);
-    if (heap_index) {
-        return rb_shape_tree.capacities[heap_index - 1];
-    }
-    return 0;
-}
-
-static inline attr_index_t
 RSHAPE_CAPACITY(shape_id_t shape_id)
 {
-    attr_index_t embedded_capacity = RSHAPE_EMBEDDED_CAPACITY(shape_id);
+    attr_index_t embedded_capacity = rb_shape_capacity(shape_id);
 
     if (embedded_capacity > RSHAPE(shape_id)->capacity) {
         return embedded_capacity;
@@ -489,10 +501,7 @@ rb_shape_transition_complex(shape_id_t shape_id)
         next_shape_id = rb_shape_layout(shape_id) | ROOT_COMPLEX_WITH_OBJ_ID;
     }
 
-    uint8_t heap_index = rb_shape_heap_index(shape_id);
-    if (heap_index) {
-        next_shape_id |= rb_shape_root(heap_index - 1);
-    }
+    next_shape_id |= shape_id & SHAPE_ID_CAPACITY_MASK;
 
     RUBY_ASSERT(rb_shape_has_object_id(shape_id) == rb_shape_has_object_id(next_shape_id));
 
@@ -508,9 +517,9 @@ rb_shape_transition_offset(shape_id_t shape_id, shape_id_t offset)
 }
 
 static inline shape_id_t
-rb_shape_transition_heap(shape_id_t shape_id, size_t heap_index)
+rb_shape_transition_capacity(shape_id_t shape_id, size_t capacity)
 {
-    return (shape_id & (~SHAPE_ID_HEAP_INDEX_MASK)) | rb_shape_root(heap_index);
+    return (shape_id & (~SHAPE_ID_CAPACITY_MASK)) | rb_shape_id_with_capacity(capacity);
 }
 
 shape_id_t rb_shape_transition_object_id(shape_id_t shape_id);
@@ -529,9 +538,9 @@ rb_obj_shape_transition_complex(VALUE obj)
 }
 
 static inline shape_id_t
-rb_obj_shape_transition_heap(VALUE obj, size_t heap_index)
+rb_obj_shape_transition_capacity(VALUE obj, size_t capacity)
 {
-    return rb_shape_transition_heap(RBASIC_SHAPE_ID(obj), heap_index);
+    return rb_shape_transition_capacity(RBASIC_SHAPE_ID(obj), capacity);
 }
 
 static inline shape_id_t
@@ -626,10 +635,9 @@ rb_setivar_cache_pack(shape_id_t shape_offset, shape_id_t dest_shape_offset, att
     return packed_cache;
 }
 
-
-ALWAYS_INLINE(static shape_id_t rb_setivar_cache_revalidate(shape_id_t shape_id, rb_setivar_cache cache));
+ALWAYS_INLINE(static shape_id_t rb_setivar_cache_revalidate(shape_id_t shape_id, shape_id_t fields_shape_id, rb_setivar_cache cache));
 static shape_id_t
-rb_setivar_cache_revalidate(shape_id_t shape_id, rb_setivar_cache cache)
+rb_setivar_cache_revalidate(shape_id_t shape_id, shape_id_t fields_shape_id, rb_setivar_cache cache)
 {
     RUBY_ASSERT(shape_id != INVALID_SHAPE_ID);
     RUBY_ASSERT(cache.dest_shape_offset == INVALID_SHAPE_ID || cache.dest_shape_offset == RSHAPE_OFFSET(cache.dest_shape_offset));
@@ -639,7 +647,7 @@ rb_setivar_cache_revalidate(shape_id_t shape_id, rb_setivar_cache cache)
         return INVALID_SHAPE_ID;
     }
 
-    if (UNLIKELY(cache.index >= RSHAPE_CAPACITY(shape_id))) {
+    if (UNLIKELY(cache.index >= RSHAPE_CAPACITY(fields_shape_id))) {
         // That's still a hit in term of layout, but the object will need to be resized,
         // so unfortunately we'll have to go through the slow path regardless...
         return INVALID_SHAPE_ID;
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index dceb14413a..8abbfb1509 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -1404,7 +1404,7 @@ vm_setivar_class(VALUE obj, VALUE val, rb_setivar_cache cache)
     }
 
     shape_id_t shape_id = RBASIC_SHAPE_ID(fields_obj);
-    shape_id_t dest_shape_id = rb_setivar_cache_revalidate(shape_id, cache);
+    shape_id_t dest_shape_id = rb_setivar_cache_revalidate(shape_id, RBASIC_SHAPE_ID(fields_obj), cache);
     if (UNLIKELY(dest_shape_id == INVALID_SHAPE_ID)) {
         return Qundef;
     }
@@ -1426,14 +1426,17 @@ NOINLINE(static VALUE vm_setivar_default(VALUE obj, ID id, VALUE val, rb_setivar
 static VALUE
 vm_setivar_default(VALUE obj, ID id, VALUE val, rb_setivar_cache cache)
 {
+    VALUE fields_obj = rb_obj_fields(obj, id);
+    if (UNLIKELY(!fields_obj)) {
+        return Qundef;
+    }
+
     shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
-    shape_id_t dest_shape_id = rb_setivar_cache_revalidate(shape_id, cache);
+    shape_id_t dest_shape_id = rb_setivar_cache_revalidate(shape_id, RBASIC_SHAPE_ID(fields_obj), cache);
     if (UNLIKELY(dest_shape_id == INVALID_SHAPE_ID)) {
         return Qundef;
     }
 
-    VALUE fields_obj = rb_obj_fields(obj, id);
-    RUBY_ASSERT(fields_obj);
     RB_OBJ_WRITE(fields_obj, &rb_imemo_fields_ptr(fields_obj)[cache.index], val);
 
     if (shape_id != dest_shape_id) {
@@ -1458,7 +1461,7 @@ vm_setivar(VALUE obj, VALUE val, rb_setivar_cache cache)
             VM_ASSERT(!rb_ractor_shareable_p(obj) || rb_obj_frozen_p(obj));
 
             shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
-            shape_id_t dest_shape_id = rb_setivar_cache_revalidate(shape_id, cache);
+            shape_id_t dest_shape_id = rb_setivar_cache_revalidate(shape_id, shape_id, cache);
             if (UNLIKELY(dest_shape_id == INVALID_SHAPE_ID)) {
                 break;
             }
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 272c10fde9..2b31f5cddd 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -632,7 +632,7 @@ pub const VM_ENV_FLAG_ISOLATED: vm_frame_env_flags = 16;
 pub type vm_frame_env_flags = u32;
 pub type attr_index_t = u8;
 pub type shape_id_t = u32;
-pub const SHAPE_ID_HAS_IVAR_MASK: shape_id_mask = 8912894;
+pub const SHAPE_ID_HAS_IVAR_MASK: shape_id_mask = 67633150;
 pub type shape_id_mask = u32;
 #[repr(C)]
 pub struct rb_cvar_class_tbl_entry {
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index d64876308f..cd945e0b2b 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -1488,17 +1488,17 @@ pub const VM_ENV_FLAG_ISOLATED: vm_frame_env_flags = 16;
 pub type vm_frame_env_flags = u32;
 pub type attr_index_t = u8;
 pub type shape_id_t = u32;
-pub const SHAPE_ID_HEAP_INDEX_MASK: shape_id_fl_type = 7864320;
-pub const SHAPE_ID_FL_COMPLEX: shape_id_fl_type = 8388608;
-pub const SHAPE_ID_FL_FROZEN: shape_id_fl_type = 16777216;
-pub const SHAPE_ID_FL_HAS_OBJECT_ID: shape_id_fl_type = 33554432;
+pub const SHAPE_ID_CAPACITY_MASK: shape_id_fl_type = 66584576;
+pub const SHAPE_ID_FL_COMPLEX: shape_id_fl_type = 67108864;
+pub const SHAPE_ID_FL_FROZEN: shape_id_fl_type = 134217728;
+pub const SHAPE_ID_FL_HAS_OBJECT_ID: shape_id_fl_type = 268435456;
 pub const SHAPE_ID_LAYOUT_ROBJECT: shape_id_fl_type = 0;
-pub const SHAPE_ID_LAYOUT_RCLASS: shape_id_fl_type = 67108864;
-pub const SHAPE_ID_LAYOUT_RDATA: shape_id_fl_type = 134217728;
-pub const SHAPE_ID_LAYOUT_OTHER: shape_id_fl_type = 201326592;
-pub const SHAPE_ID_LAYOUT_MASK: shape_id_fl_type = 201326592;
-pub const SHAPE_ID_FL_NON_CANONICAL_MASK: shape_id_fl_type = 50331648;
-pub const SHAPE_ID_FLAGS_MASK: shape_id_fl_type = 267911168;
+pub const SHAPE_ID_LAYOUT_RCLASS: shape_id_fl_type = 536870912;
+pub const SHAPE_ID_LAYOUT_RDATA: shape_id_fl_type = 1073741824;
+pub const SHAPE_ID_LAYOUT_OTHER: shape_id_fl_type = 1610612736;
+pub const SHAPE_ID_LAYOUT_MASK: shape_id_fl_type = 1610612736;
+pub const SHAPE_ID_FL_NON_CANONICAL_MASK: shape_id_fl_type = 402653184;
+pub const SHAPE_ID_FLAGS_MASK: shape_id_fl_type = 2146959360;
 pub type shape_id_fl_type = u32;
 pub const CONST_DEPRECATED: rb_const_flag_t = 256;
 pub const CONST_VISIBILITY_MASK: rb_const_flag_t = 255;