1 files changed, 262 insertions, 207 deletions
diff --git a/variable.c b/variable.c
index 6d0e9832e7..93ae6bb8b2 100644
--- a/variable.c
+++ b/variable.c
@@ -1228,19 +1228,10 @@ gen_fields_tbl_bytes(size_t n)
 }
 
 static struct gen_fields_tbl *
-gen_fields_tbl_resize(struct gen_fields_tbl *old, uint32_t n)
+gen_fields_tbl_resize(struct gen_fields_tbl *old, uint32_t new_capa)
 {
-    RUBY_ASSERT(n > 0);
-
-    uint32_t len = old ? old->as.shape.fields_count : 0;
-    struct gen_fields_tbl *fields_tbl = xrealloc(old, gen_fields_tbl_bytes(n));
-
-    fields_tbl->as.shape.fields_count = n;
-    for (; len < n; len++) {
-        fields_tbl->as.shape.fields[len] = Qundef;
-    }
-
-    return fields_tbl;
+    RUBY_ASSERT(new_capa > 0);
+    return xrealloc(old, gen_fields_tbl_bytes(new_capa));
 }
 
 void
@@ -1253,7 +1244,8 @@ rb_mark_generic_ivar(VALUE obj)
             rb_mark_tbl_no_pin(fields_tbl->as.complex.table);
         }
         else {
-            for (uint32_t i = 0; i < fields_tbl->as.shape.fields_count; i++) {
+            uint32_t fields_count = RSHAPE_LEN(RBASIC_SHAPE_ID(obj));
+            for (uint32_t i = 0; i < fields_count; i++) {
                 rb_gc_mark_movable(fields_tbl->as.shape.fields[i]);
             }
         }
@@ -1290,7 +1282,7 @@ rb_generic_ivar_memsize(VALUE obj)
             return sizeof(struct gen_fields_tbl) + st_memsize(fields_tbl->as.complex.table);
         }
         else {
-            return gen_fields_tbl_bytes(fields_tbl->as.shape.fields_count);
+            return gen_fields_tbl_bytes(RSHAPE_CAPACITY(RBASIC_SHAPE_ID(obj)));
         }
     }
     return 0;
@@ -1299,36 +1291,35 @@ rb_generic_ivar_memsize(VALUE obj)
 static size_t
 gen_fields_tbl_count(VALUE obj, const struct gen_fields_tbl *fields_tbl)
 {
-    uint32_t i;
-    size_t n = 0;
-
     if (rb_shape_obj_too_complex_p(obj)) {
-        n = st_table_size(fields_tbl->as.complex.table);
+        return st_table_size(fields_tbl->as.complex.table);
     }
     else {
-        for (i = 0; i < fields_tbl->as.shape.fields_count; i++) {
-            if (!UNDEF_P(fields_tbl->as.shape.fields[i])) {
-                n++;
-            }
-        }
+        return RSHAPE_LEN(RBASIC_SHAPE_ID(obj));
     }
-
-    return n;
 }
 
 VALUE
 rb_obj_field_get(VALUE obj, shape_id_t target_shape_id)
 {
     RUBY_ASSERT(!SPECIAL_CONST_P(obj));
-    RUBY_ASSERT(RSHAPE(target_shape_id)->type == SHAPE_IVAR || RSHAPE(target_shape_id)->type == SHAPE_OBJ_ID);
+    RUBY_ASSERT(RSHAPE_TYPE_P(target_shape_id, SHAPE_IVAR) || RSHAPE_TYPE_P(target_shape_id, SHAPE_OBJ_ID));
+
+    if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
+        ASSERT_vm_locking();
+        VALUE field_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+        if (field_obj) {
+            return rb_obj_field_get(field_obj, target_shape_id);
+        }
+        return Qundef;
+    }
 
     if (rb_shape_too_complex_p(target_shape_id)) {
         st_table *fields_hash;
         switch (BUILTIN_TYPE(obj)) {
           case T_CLASS:
           case T_MODULE:
-            ASSERT_vm_locking();
-            fields_hash = RCLASS_FIELDS_HASH(obj);
+            rb_bug("Unreachable");
             break;
           case T_OBJECT:
             fields_hash = ROBJECT_FIELDS_HASH(obj);
@@ -1342,7 +1333,7 @@ rb_obj_field_get(VALUE obj, shape_id_t target_shape_id)
             break;
         }
         VALUE value = Qundef;
-        st_lookup(fields_hash, RSHAPE(target_shape_id)->edge_name, &value);
+        st_lookup(fields_hash, RSHAPE_EDGE_NAME(target_shape_id), &value);
 
 #if RUBY_DEBUG
         if (UNDEF_P(value)) {
@@ -1354,13 +1345,12 @@ rb_obj_field_get(VALUE obj, shape_id_t target_shape_id)
         return value;
     }
 
-    attr_index_t attr_index = RSHAPE(target_shape_id)->next_field_index - 1;
+    attr_index_t attr_index = RSHAPE_INDEX(target_shape_id);
     VALUE *fields;
     switch (BUILTIN_TYPE(obj)) {
       case T_CLASS:
       case T_MODULE:
-        ASSERT_vm_locking();
-        fields = RCLASS_PRIME_FIELDS(obj);
+        rb_bug("Unreachable");
         break;
       case T_OBJECT:
         fields = ROBJECT_FIELDS(obj);
@@ -1382,43 +1372,19 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
     if (SPECIAL_CONST_P(obj)) return undef;
 
     shape_id_t shape_id;
-    VALUE * ivar_list;
-    shape_id = RBASIC_SHAPE_ID(obj);
+    VALUE *ivar_list;
 
     switch (BUILTIN_TYPE(obj)) {
       case T_CLASS:
       case T_MODULE:
         {
-            bool found = false;
-            VALUE val;
-
-            RB_VM_LOCKING() {
-                if (rb_shape_too_complex_p(shape_id)) {
-                    st_table * iv_table = RCLASS_FIELDS_HASH(obj);
-                    if (rb_st_lookup(iv_table, (st_data_t)id, (st_data_t *)&val)) {
-                        found = true;
-                    }
-                    else {
-                        val = undef;
-                    }
-                }
-                else {
-                    attr_index_t index = 0;
-                    found = rb_shape_get_iv_index(shape_id, id, &index);
-
-                    if (found) {
-                        ivar_list = RCLASS_PRIME_FIELDS(obj);
-                        RUBY_ASSERT(ivar_list);
-
-                        val = ivar_list[index];
-                    }
-                    else {
-                        val = undef;
-                    }
-                }
+            VALUE val = undef;
+            VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+            if (fields_obj) {
+                val = rb_ivar_lookup(fields_obj, id, undef);
             }
 
-            if (found &&
+            if (val != undef &&
                     rb_is_instance_id(id) &&
                     UNLIKELY(!rb_ractor_main_p()) &&
                     !rb_ractor_shareable_p(val)) {
@@ -1427,10 +1393,32 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
             }
             return val;
         }
+      case T_IMEMO:
+        // Handled like T_OBJECT
+        {
+            RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_class_fields));
+            shape_id = RBASIC_SHAPE_ID(obj);
+
+            if (rb_shape_too_complex_p(shape_id)) {
+                st_table *iv_table = rb_imemo_class_fields_complex_tbl(obj);
+                VALUE val;
+                if (rb_st_lookup(iv_table, (st_data_t)id, (st_data_t *)&val)) {
+                    return val;
+                }
+                else {
+                    return undef;
+                }
+            }
+
+            RUBY_ASSERT(!rb_shape_obj_too_complex_p(obj));
+            ivar_list = rb_imemo_class_fields_ptr(obj);
+            break;
+        }
       case T_OBJECT:
         {
+            shape_id = RBASIC_SHAPE_ID(obj);
             if (rb_shape_too_complex_p(shape_id)) {
-                st_table * iv_table = ROBJECT_FIELDS_HASH(obj);
+                st_table *iv_table = ROBJECT_FIELDS_HASH(obj);
                 VALUE val;
                 if (rb_st_lookup(iv_table, (st_data_t)id, (st_data_t *)&val)) {
                     return val;
@@ -1445,6 +1433,7 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
             break;
         }
       default:
+        shape_id = RBASIC_SHAPE_ID(obj);
         if (FL_TEST_RAW(obj, FL_EXIVAR)) {
             struct gen_fields_tbl *fields_tbl;
             rb_gen_fields_tbl_get(obj, id, &fields_tbl);
@@ -1493,13 +1482,22 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
 {
     rb_check_frozen(obj);
 
-    bool locked = false;
-    unsigned int lev = 0;
     VALUE val = undef;
     if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
         IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(id);
-        RB_VM_LOCK_ENTER_LEV(&lev);
-        locked = true;
+
+        VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+        if (fields_obj) {
+            if (rb_multi_ractor_p()) {
+                fields_obj = rb_imemo_class_fields_clone(fields_obj);
+                val = rb_ivar_delete(fields_obj, id, undef);
+                RCLASS_WRITABLE_SET_FIELDS_OBJ(obj, fields_obj);
+            }
+            else {
+                val = rb_ivar_delete(fields_obj, id, undef);
+            }
+        }
+        return val;
     }
 
     shape_id_t old_shape_id = rb_obj_shape_id(obj);
@@ -1511,9 +1509,6 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
     shape_id_t next_shape_id = rb_shape_transition_remove_ivar(obj, id, &removed_shape_id);
 
     if (next_shape_id == old_shape_id) {
-        if (locked) {
-            RB_VM_LOCK_LEAVE_LEV(&lev);
-        }
         return undef;
     }
 
@@ -1522,13 +1517,17 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
         goto too_complex;
     }
 
-    RUBY_ASSERT(RSHAPE(next_shape_id)->next_field_index == RSHAPE(old_shape_id)->next_field_index - 1);
+    RUBY_ASSERT(RSHAPE_LEN(next_shape_id) == RSHAPE_LEN(old_shape_id) - 1);
 
     VALUE *fields;
     switch(BUILTIN_TYPE(obj)) {
       case T_CLASS:
       case T_MODULE:
-        fields = RCLASS_PRIME_FIELDS(obj);
+        rb_bug("Unreachable");
+        break;
+      case T_IMEMO:
+        RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_class_fields));
+        fields = rb_imemo_class_fields_ptr(obj);
         break;
       case T_OBJECT:
         fields = ROBJECT_FIELDS(obj);
@@ -1543,9 +1542,9 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
 
     RUBY_ASSERT(removed_shape_id != INVALID_SHAPE_ID);
 
-    attr_index_t new_fields_count = RSHAPE(next_shape_id)->next_field_index;
+    attr_index_t new_fields_count = RSHAPE_LEN(next_shape_id);
 
-    attr_index_t removed_index = RSHAPE(removed_shape_id)->next_field_index - 1;
+    attr_index_t removed_index = RSHAPE_INDEX(removed_shape_id);
     val = fields[removed_index];
     size_t trailing_fields = new_fields_count - removed_index;
 
@@ -1563,10 +1562,6 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
     }
     rb_obj_set_shape_id(obj, next_shape_id);
 
-    if (locked) {
-        RB_VM_LOCK_LEAVE_LEV(&lev);
-    }
-
     return val;
 
 too_complex:
@@ -1575,7 +1570,12 @@ too_complex:
         switch (BUILTIN_TYPE(obj)) {
           case T_CLASS:
           case T_MODULE:
-            table = RCLASS_WRITABLE_FIELDS_HASH(obj);
+            rb_bug("Unreachable");
+            break;
+
+          case T_IMEMO:
+            RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_class_fields));
+            table = rb_imemo_class_fields_complex_tbl(obj);
             break;
 
           case T_OBJECT:
@@ -1598,10 +1598,6 @@ too_complex:
         }
     }
 
-    if (locked) {
-        RB_VM_LOCK_LEAVE_LEV(&lev);
-    }
-
     return val;
 }
 
@@ -1614,6 +1610,10 @@ rb_attr_delete(VALUE obj, ID id)
 static shape_id_t
 obj_transition_too_complex(VALUE obj, st_table *table)
 {
+    if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
+        return obj_transition_too_complex(RCLASS_WRITABLE_ENSURE_FIELDS_OBJ(obj), table);
+    }
+
     RUBY_ASSERT(!rb_shape_obj_too_complex_p(obj));
     shape_id_t shape_id = rb_shape_transition_complex(obj);
 
@@ -1629,9 +1629,7 @@ obj_transition_too_complex(VALUE obj, st_table *table)
         break;
       case T_CLASS:
       case T_MODULE:
-        old_fields = RCLASS_PRIME_FIELDS(obj);
-        RBASIC_SET_SHAPE_ID(obj, shape_id);
-        RCLASS_SET_FIELDS_HASH(obj, table);
+        rb_bug("Unreachable");
         break;
       default:
         RB_VM_LOCKING() {
@@ -1827,8 +1825,8 @@ generic_fields_lookup_ensure_size(st_data_t *k, st_data_t *v, st_data_t u, int e
 
     if (!existing || fields_lookup->resize) {
         if (existing) {
-            RUBY_ASSERT(RSHAPE(fields_lookup->shape_id)->type == SHAPE_IVAR || RSHAPE(fields_lookup->shape_id)->type == SHAPE_OBJ_ID);
-            RUBY_ASSERT(RSHAPE_CAPACITY(RSHAPE(fields_lookup->shape_id)->parent_id) < RSHAPE_CAPACITY(fields_lookup->shape_id));
+            RUBY_ASSERT(RSHAPE_TYPE_P(fields_lookup->shape_id, SHAPE_IVAR) || RSHAPE_TYPE_P(fields_lookup->shape_id, SHAPE_OBJ_ID));
+            RUBY_ASSERT(RSHAPE_CAPACITY(RSHAPE_PARENT(fields_lookup->shape_id)) < RSHAPE_CAPACITY(fields_lookup->shape_id));
         }
         else {
             FL_SET_RAW((VALUE)*k, FL_EXIVAR);
@@ -2052,11 +2050,20 @@ rb_vm_set_ivar_id(VALUE obj, ID id, VALUE val)
 bool
 rb_obj_set_shape_id(VALUE obj, shape_id_t shape_id)
 {
-    if (rb_obj_shape_id(obj) == shape_id) {
+    shape_id_t old_shape_id = rb_obj_shape_id(obj);
+    if (old_shape_id == shape_id) {
         return false;
     }
 
+    if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
+        // Avoid creating the fields_obj just to freeze the class
+        if (!(shape_id == SPECIAL_CONST_SHAPE_ID && old_shape_id == ROOT_SHAPE_ID)) {
+            RBASIC_SET_SHAPE_ID(RCLASS_WRITABLE_ENSURE_FIELDS_OBJ(obj), shape_id);
+        }
+    }
+    // FIXME: How to do multi-shape?
     RBASIC_SET_SHAPE_ID(obj, shape_id);
+
     return true;
 }
 
@@ -2117,8 +2124,6 @@ rb_ivar_set_internal(VALUE obj, ID id, VALUE val)
     ivar_set(obj, id, val);
 }
 
-static void class_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val);
-
 void
 rb_obj_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val)
 {
@@ -2128,8 +2133,8 @@ rb_obj_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val)
         break;
       case T_CLASS:
       case T_MODULE:
-        ASSERT_vm_locking();
-        class_field_set(obj, target_shape_id, val);
+        // The only field is object_id and T_CLASS handle it differently.
+        rb_bug("Unreachable");
         break;
       default:
         generic_field_set(obj, target_shape_id, val);
@@ -2148,7 +2153,12 @@ ivar_defined0(VALUE obj, ID id)
         switch (BUILTIN_TYPE(obj)) {
           case T_CLASS:
           case T_MODULE:
-            table = (st_table *)RCLASS_FIELDS_HASH(obj);
+            rb_bug("Unreachable");
+            break;
+
+          case T_IMEMO:
+            RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_class_fields));
+            table = rb_imemo_class_fields_complex_tbl(obj);
             break;
 
           case T_OBJECT:
@@ -2180,12 +2190,15 @@ rb_ivar_defined(VALUE obj, ID id)
 {
     if (SPECIAL_CONST_P(obj)) return Qfalse;
 
-    VALUE defined;
+    VALUE defined = Qfalse;
     switch (BUILTIN_TYPE(obj)) {
       case T_CLASS:
       case T_MODULE:
-        RB_VM_LOCKING() {
-            defined = ivar_defined0(obj, id);
+        {
+            VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+            if (fields_obj) {
+                defined = ivar_defined0(fields_obj, id);
+            }
         }
         break;
       default:
@@ -2200,60 +2213,50 @@ struct iv_itr_data {
     struct gen_fields_tbl *fields_tbl;
     st_data_t arg;
     rb_ivar_foreach_callback_func *func;
+    VALUE *fields;
     bool ivar_only;
 };
 
-/*
- * Returns a flag to stop iterating depending on the result of +callback+.
- */
-static bool
-iterate_over_shapes_with_callback(rb_shape_t *shape, rb_ivar_foreach_callback_func *callback, struct iv_itr_data *itr_data)
+static int
+iterate_over_shapes_callback(shape_id_t shape_id, void *data)
 {
-    switch ((enum shape_type)shape->type) {
-      case SHAPE_ROOT:
-        return false;
-      case SHAPE_OBJ_ID:
-        if (itr_data->ivar_only) {
-            return iterate_over_shapes_with_callback(RSHAPE(shape->parent_id), callback, itr_data);
-        }
-        // fallthrough
-      case SHAPE_IVAR:
-        ASSUME(callback);
-        if (iterate_over_shapes_with_callback(RSHAPE(shape->parent_id), callback, itr_data)) {
-            return true;
-        }
+    struct iv_itr_data *itr_data = data;
 
-        VALUE * iv_list;
-        switch (BUILTIN_TYPE(itr_data->obj)) {
-          case T_OBJECT:
-            RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
-            iv_list = ROBJECT_FIELDS(itr_data->obj);
-            break;
-          case T_CLASS:
-          case T_MODULE:
-            RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
-            iv_list = RCLASS_PRIME_FIELDS(itr_data->obj);
-            break;
-          default:
-            iv_list = itr_data->fields_tbl->as.shape.fields;
-            break;
-        }
-        VALUE val = iv_list[shape->next_field_index - 1];
-        if (!UNDEF_P(val)) {
-            switch (callback(shape->edge_name, val, itr_data->arg)) {
-              case ST_CHECK:
-              case ST_CONTINUE:
-                break;
-              case ST_STOP:
-                return true;
-              default:
-                rb_bug("unreachable");
-            }
-        }
-        return false;
+    if (itr_data->ivar_only && !RSHAPE_TYPE_P(shape_id, SHAPE_IVAR)) {
+        return ST_CONTINUE;
+    }
+
+    VALUE *iv_list;
+    switch (BUILTIN_TYPE(itr_data->obj)) {
+      case T_OBJECT:
+        RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
+        iv_list = ROBJECT_FIELDS(itr_data->obj);
+        break;
+      case T_CLASS:
+      case T_MODULE:
+        rb_bug("Unreachable");
+      case T_IMEMO:
+        RUBY_ASSERT(IMEMO_TYPE_P(itr_data->obj, imemo_class_fields));
+        RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
+
+        iv_list = rb_imemo_class_fields_ptr(itr_data->obj);
+        break;
       default:
-        UNREACHABLE_RETURN(false);
+        iv_list = itr_data->fields_tbl->as.shape.fields;
+        break;
     }
+
+    VALUE val = iv_list[RSHAPE_INDEX(shape_id)];
+    return itr_data->func(RSHAPE_EDGE_NAME(shape_id), val, itr_data->arg);
+}
+
+/*
+ * Returns a flag to stop iterating depending on the result of +callback+.
+ */
+static void
+iterate_over_shapes(shape_id_t shape_id, rb_ivar_foreach_callback_func *callback, struct iv_itr_data *itr_data)
+{
+    rb_shape_foreach_field(shape_id, iterate_over_shapes_callback, itr_data);
 }
 
 static int
@@ -2279,7 +2282,8 @@ obj_fields_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg, b
         rb_st_foreach(ROBJECT_FIELDS_HASH(obj), each_hash_iv, (st_data_t)&itr_data);
     }
     else {
-        iterate_over_shapes_with_callback(RSHAPE(shape_id), func, &itr_data);
+        itr_data.fields = ROBJECT_FIELDS(obj);
+        iterate_over_shapes(shape_id, func, &itr_data);
     }
 }
 
@@ -2302,28 +2306,30 @@ gen_fields_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg, b
         rb_st_foreach(fields_tbl->as.complex.table, each_hash_iv, (st_data_t)&itr_data);
     }
     else {
-        iterate_over_shapes_with_callback(RSHAPE(shape_id), func, &itr_data);
+        itr_data.fields = fields_tbl->as.shape.fields;
+        iterate_over_shapes(shape_id, func, &itr_data);
     }
 }
 
 static void
-class_fields_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg, bool ivar_only)
+class_fields_each(VALUE fields_obj, rb_ivar_foreach_callback_func *func, st_data_t arg, bool ivar_only)
 {
-    RUBY_ASSERT(RB_TYPE_P(obj, T_CLASS) || RB_TYPE_P(obj, T_MODULE));
+    IMEMO_TYPE_P(fields_obj, imemo_class_fields);
 
     struct iv_itr_data itr_data = {
-        .obj = obj,
+        .obj = fields_obj,
         .arg = arg,
         .func = func,
         .ivar_only = ivar_only,
     };
 
-    shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
+    shape_id_t shape_id = RBASIC_SHAPE_ID(fields_obj);
     if (rb_shape_too_complex_p(shape_id)) {
-        rb_st_foreach(RCLASS_WRITABLE_FIELDS_HASH(obj), each_hash_iv, (st_data_t)&itr_data);
+        rb_st_foreach(rb_imemo_class_fields_complex_tbl(fields_obj), each_hash_iv, (st_data_t)&itr_data);
     }
     else {
-        iterate_over_shapes_with_callback(RSHAPE(shape_id), func, &itr_data);
+        itr_data.fields = rb_imemo_class_fields_ptr(fields_obj);
+        iterate_over_shapes(shape_id, func, &itr_data);
     }
 }
 
@@ -2361,7 +2367,7 @@ rb_copy_generic_ivar(VALUE dest, VALUE obj)
         shape_id_t initial_shape_id = rb_obj_shape_id(dest);
 
         if (!rb_shape_canonical_p(src_shape_id)) {
-            RUBY_ASSERT(RSHAPE(initial_shape_id)->type == SHAPE_ROOT);
+            RUBY_ASSERT(RSHAPE_TYPE_P(initial_shape_id, SHAPE_ROOT));
 
             dest_shape_id = rb_shape_rebuild(initial_shape_id, src_shape_id);
             if (UNLIKELY(rb_shape_too_complex_p(dest_shape_id))) {
@@ -2431,18 +2437,26 @@ rb_field_foreach(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg,
 {
     if (SPECIAL_CONST_P(obj)) return;
     switch (BUILTIN_TYPE(obj)) {
+      case T_IMEMO:
+        if (IMEMO_TYPE_P(obj, imemo_class_fields)) {
+            class_fields_each(obj, func, arg, ivar_only);
+        }
+        break;
       case T_OBJECT:
         obj_fields_each(obj, func, arg, ivar_only);
         break;
       case T_CLASS:
       case T_MODULE:
-        IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(0);
-        RB_VM_LOCKING() {
-            class_fields_each(obj, func, arg, ivar_only);
+        {
+            IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(0);
+            VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+            if (fields_obj) {
+                class_fields_each(fields_obj, func, arg, ivar_only);
+            }
         }
         break;
       default:
-        if (FL_TEST(obj, FL_EXIVAR)) {
+        if (FL_TEST_RAW(obj, FL_EXIVAR)) {
             gen_fields_each(obj, func, arg, ivar_only);
         }
         break;
@@ -2467,8 +2481,16 @@ rb_ivar_count(VALUE obj)
         break;
       case T_CLASS:
       case T_MODULE:
-        iv_count = RCLASS_FIELDS_COUNT(obj);
-        break;
+        {
+            VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+            if (!fields_obj) {
+                return 0;
+            }
+            if (rb_shape_obj_too_complex_p(fields_obj)) {
+                return rb_st_table_size(rb_imemo_class_fields_complex_tbl(fields_obj));
+            }
+            return RBASIC_FIELDS_COUNT(fields_obj);
+        }
       default:
         if (FL_TEST(obj, FL_EXIVAR)) {
             struct gen_fields_tbl *fields_tbl;
@@ -4674,71 +4696,106 @@ rb_iv_set(VALUE obj, const char *name, VALUE val)
     return rb_ivar_set(obj, id, val);
 }
 
-static VALUE *
-class_ivar_set_shape_fields(VALUE obj, void *_data)
+static bool
+class_fields_ivar_set(VALUE klass, VALUE fields_obj, ID id, VALUE val, bool concurrent, VALUE *new_fields_obj)
 {
-    RUBY_ASSERT(!rb_shape_obj_too_complex_p(obj));
+    bool existing = true;
+    const VALUE original_fields_obj = fields_obj;
+    fields_obj = original_fields_obj ? original_fields_obj : rb_imemo_class_fields_new(klass, 1);
 
-    return RCLASS_PRIME_FIELDS(obj);
-}
+    shape_id_t current_shape_id = RBASIC_SHAPE_ID(fields_obj);
+    shape_id_t next_shape_id = current_shape_id;
 
-static void
-class_ivar_set_shape_resize_fields(VALUE obj, attr_index_t _old_capa, attr_index_t new_capa, void *_data)
-{
-    REALLOC_N(RCLASS_PRIME_FIELDS(obj), VALUE, new_capa);
-}
+    if (UNLIKELY(rb_shape_too_complex_p(current_shape_id))) {
+        goto too_complex;
+    }
 
-static void
-class_ivar_set_set_shape_id(VALUE obj, shape_id_t shape_id, void *_data)
-{
-    rb_obj_set_shape_id(obj, shape_id);
-}
+    attr_index_t index;
+    if (!rb_shape_get_iv_index(current_shape_id, id, &index)) {
+        existing = false;
 
-static shape_id_t
-class_ivar_set_transition_too_complex(VALUE obj, void *_data)
-{
-    return rb_evict_fields_to_hash(obj);
-}
+        index = RSHAPE_LEN(current_shape_id);
+        if (index >= SHAPE_MAX_FIELDS) {
+            rb_raise(rb_eArgError, "too many instance variables");
+        }
 
-static st_table *
-class_ivar_set_too_complex_table(VALUE obj, void *_data)
-{
-    RUBY_ASSERT(rb_shape_obj_too_complex_p(obj));
+        next_shape_id = rb_shape_transition_add_ivar(fields_obj, id);
+        if (UNLIKELY(rb_shape_too_complex_p(next_shape_id))) {
+            attr_index_t current_len = RSHAPE_LEN(current_shape_id);
+            fields_obj = rb_imemo_class_fields_new_complex(klass, current_len + 1);
+            if (current_len) {
+                rb_obj_copy_fields_to_hash_table(original_fields_obj, rb_imemo_class_fields_complex_tbl(fields_obj));
+                RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+            }
+            goto too_complex;
+        }
+
+        attr_index_t next_capacity = RSHAPE_CAPACITY(next_shape_id);
+        attr_index_t current_capacity = RSHAPE_CAPACITY(current_shape_id);
 
-    return RCLASS_WRITABLE_FIELDS_HASH(obj);
+        if (concurrent || next_capacity != current_capacity) {
+            RUBY_ASSERT(concurrent || next_capacity > current_capacity);
+
+            // We allocate a new fields_obj even when concurrency isn't a concern
+            // so that we're embedded as long as possible.
+            fields_obj = rb_imemo_class_fields_new(klass, next_capacity);
+            if (original_fields_obj) {
+                MEMCPY(rb_imemo_class_fields_ptr(fields_obj), rb_imemo_class_fields_ptr(original_fields_obj), VALUE, RSHAPE_LEN(current_shape_id));
+            }
+        }
+
+        RUBY_ASSERT(RSHAPE(next_shape_id)->type == SHAPE_IVAR);
+        RUBY_ASSERT(index == (RSHAPE_LEN(next_shape_id) - 1));
+    }
+
+    VALUE *fields = rb_imemo_class_fields_ptr(fields_obj);
+    RB_OBJ_WRITE(fields_obj, &fields[index], val);
+
+    if (!existing) {
+        RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+    }
+
+    *new_fields_obj = fields_obj;
+    return existing;
+
+too_complex:
+    {
+        st_table *table = rb_imemo_class_fields_complex_tbl(fields_obj);
+        existing = st_insert(table, (st_data_t)id, (st_data_t)val);
+        RB_OBJ_WRITTEN(fields_obj, Qundef, val);
+
+        if (fields_obj != original_fields_obj) {
+            RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+        }
+    }
+
+    *new_fields_obj = fields_obj;
+    return existing;
 }
 
 int
 rb_class_ivar_set(VALUE obj, ID id, VALUE val)
 {
     RUBY_ASSERT(RB_TYPE_P(obj, T_CLASS) || RB_TYPE_P(obj, T_MODULE));
-    bool existing = false;
     rb_check_frozen(obj);
 
     rb_class_ensure_writable(obj);
 
-    RB_VM_LOCKING() {
-        existing = general_ivar_set(obj, id, val, NULL,
-                                    class_ivar_set_shape_fields,
-                                    class_ivar_set_shape_resize_fields,
-                                    class_ivar_set_set_shape_id,
-                                    class_ivar_set_transition_too_complex,
-                                    class_ivar_set_too_complex_table).existing;
-    }
+    const VALUE original_fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+    VALUE new_fields_obj = 0;
 
-    return existing;
-}
+    bool existing = class_fields_ivar_set(obj, original_fields_obj, id, val, rb_multi_ractor_p(), &new_fields_obj);
 
-static void
-class_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val)
-{
-    RUBY_ASSERT(RB_TYPE_P(obj, T_CLASS) || RB_TYPE_P(obj, T_MODULE));
-    general_field_set(obj, target_shape_id, val, NULL,
-                      class_ivar_set_shape_fields,
-                      class_ivar_set_shape_resize_fields,
-                      class_ivar_set_set_shape_id,
-                      class_ivar_set_transition_too_complex,
-                      class_ivar_set_too_complex_table);
+    if (new_fields_obj != original_fields_obj) {
+        RCLASS_WRITABLE_SET_FIELDS_OBJ(obj, new_fields_obj);
+
+        // TODO: What should we set as the T_CLASS shape_id?
+        // In most case we can replicate the single `fields_obj` shape
+        // but in namespaced case?
+        // Perhaps INVALID_SHAPE_ID?
+        RBASIC_SET_SHAPE_ID(obj, RBASIC_SHAPE_ID(new_fields_obj));
+    }
+    return existing;
 }
 
 static int
@@ -4754,9 +4811,7 @@ rb_fields_tbl_copy(VALUE dst, VALUE src)
 {
     RUBY_ASSERT(rb_type(dst) == rb_type(src));
     RUBY_ASSERT(RB_TYPE_P(dst, T_CLASS) || RB_TYPE_P(dst, T_MODULE));
-
     RUBY_ASSERT(RSHAPE_TYPE_P(RBASIC_SHAPE_ID(dst), SHAPE_ROOT));
-    RUBY_ASSERT(!RCLASS_PRIME_FIELDS(dst));
 
     rb_ivar_foreach(src, tbl_copy_i, dst);
 }