1 files changed, 523 insertions, 481 deletions
diff --git a/variable.c b/variable.c
index 6d0e9832e7..e535aefe27 100644
--- a/variable.c
+++ b/variable.c
@@ -1197,8 +1197,31 @@ rb_generic_fields_tbl_get(void)
     return generic_fields_tbl_;
 }
 
+static inline VALUE
+generic_fields_lookup(VALUE obj, ID id, bool force_check_ractor)
+{
+    VALUE fields_obj = Qfalse;
+    RB_VM_LOCKING() {
+        st_table *generic_tbl = generic_fields_tbl(obj, id, false);
+        st_lookup(generic_tbl, obj, (st_data_t *)&fields_obj);
+    }
+    return fields_obj;
+}
+
+static inline void
+generic_fields_insert(VALUE obj, VALUE fields_obj)
+{
+    RUBY_ASSERT(IMEMO_TYPE_P(fields_obj, imemo_fields));
+
+    RB_VM_LOCKING() {
+        st_table *generic_tbl = generic_fields_tbl_no_ractor_check(obj);
+        st_insert(generic_tbl, obj, fields_obj);
+    }
+    RB_OBJ_WRITTEN(obj, Qundef, fields_obj);
+}
+
 int
-rb_gen_fields_tbl_get(VALUE obj, ID id, struct gen_fields_tbl **fields_tbl)
+rb_gen_fields_tbl_get(VALUE obj, ID id, VALUE *fields_obj)
 {
     RUBY_ASSERT(!RB_TYPE_P(obj, T_ICLASS));
 
@@ -1207,7 +1230,7 @@ rb_gen_fields_tbl_get(VALUE obj, ID id, struct gen_fields_tbl **fields_tbl)
 
     RB_VM_LOCKING() {
         if (st_lookup(generic_fields_tbl(obj, id, false), (st_data_t)obj, &data)) {
-            *fields_tbl = (struct gen_fields_tbl *)data;
+            *fields_obj = (VALUE)data;
             r = 1;
         }
     }
@@ -1216,133 +1239,71 @@ rb_gen_fields_tbl_get(VALUE obj, ID id, struct gen_fields_tbl **fields_tbl)
 }
 
 int
-rb_ivar_generic_fields_tbl_lookup(VALUE obj, struct gen_fields_tbl **fields_tbl)
+rb_ivar_generic_fields_tbl_lookup(VALUE obj, VALUE *fields_obj)
 {
-    return rb_gen_fields_tbl_get(obj, 0, fields_tbl);
-}
-
-static size_t
-gen_fields_tbl_bytes(size_t n)
-{
-    return offsetof(struct gen_fields_tbl, as.shape.fields) + n * sizeof(VALUE);
-}
-
-static struct gen_fields_tbl *
-gen_fields_tbl_resize(struct gen_fields_tbl *old, uint32_t n)
-{
-    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;
+    return rb_gen_fields_tbl_get(obj, 0, fields_obj);
 }
 
 void
 rb_mark_generic_ivar(VALUE obj)
 {
-    st_data_t data;
-    if (st_lookup(generic_fields_tbl_no_ractor_check(obj), (st_data_t)obj, &data)) {
-        struct gen_fields_tbl *fields_tbl = (struct gen_fields_tbl *)data;
-        if (rb_shape_obj_too_complex_p(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++) {
-                rb_gc_mark_movable(fields_tbl->as.shape.fields[i]);
-            }
-        }
+    VALUE data;
+    if (st_lookup(generic_fields_tbl_no_ractor_check(obj), (st_data_t)obj, (st_data_t *)&data)) {
+        rb_gc_mark_movable(data);
     }
 }
 
 void
 rb_free_generic_ivar(VALUE obj)
 {
-    st_data_t key = (st_data_t)obj, value;
-
-    bool too_complex = rb_shape_obj_too_complex_p(obj);
-
-    RB_VM_LOCKING() {
-        if (st_delete(generic_fields_tbl_no_ractor_check(obj), &key, &value)) {
-            struct gen_fields_tbl *fields_tbl = (struct gen_fields_tbl *)value;
-
-            if (UNLIKELY(too_complex)) {
-                st_free_table(fields_tbl->as.complex.table);
-            }
-
-            xfree(fields_tbl);
-        }
-    }
-}
+    if (rb_obj_exivar_p(obj)) {
+        st_data_t key = (st_data_t)obj, value;
 
-size_t
-rb_generic_ivar_memsize(VALUE obj)
-{
-    struct gen_fields_tbl *fields_tbl;
-
-    if (rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
-        if (rb_shape_obj_too_complex_p(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 0;
-}
-
-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);
-    }
-    else {
-        for (i = 0; i < fields_tbl->as.shape.fields_count; i++) {
-            if (!UNDEF_P(fields_tbl->as.shape.fields[i])) {
-                n++;
-            }
+        RB_VM_LOCKING() {
+            st_delete(generic_fields_tbl_no_ractor_check(obj), &key, &value);
         }
     }
-
-    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);
             break;
+          case T_IMEMO:
+            RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_fields));
+            fields_hash = rb_imemo_fields_complex_tbl(obj);
+            break;
           default:
-            RUBY_ASSERT(FL_TEST_RAW(obj, FL_EXIVAR));
-            struct gen_fields_tbl *fields_tbl = NULL;
-            rb_ivar_generic_fields_tbl_lookup(obj, &fields_tbl);
-            RUBY_ASSERT(fields_tbl);
-            fields_hash = fields_tbl->as.complex.table;
+            RUBY_ASSERT(rb_obj_exivar_p(obj));
+            VALUE fields_obj = 0;
+            rb_ivar_generic_fields_tbl_lookup(obj, &fields_obj);
+            RUBY_ASSERT(fields_obj);
+            fields_hash = rb_imemo_fields_complex_tbl(fields_obj);
             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,23 +1315,26 @@ 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);
         break;
+      case T_IMEMO:
+        RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_fields));
+        fields = rb_imemo_fields_ptr(obj);
+        break;
       default:
-        RUBY_ASSERT(FL_TEST_RAW(obj, FL_EXIVAR));
-        struct gen_fields_tbl *fields_tbl = NULL;
-        rb_ivar_generic_fields_tbl_lookup(obj, &fields_tbl);
-        RUBY_ASSERT(fields_tbl);
-        fields = fields_tbl->as.shape.fields;
+        RUBY_ASSERT(rb_obj_exivar_p(obj));
+        VALUE fields_obj = 0;
+        rb_ivar_generic_fields_tbl_lookup(obj, &fields_obj);
+        RUBY_ASSERT(fields_obj);
+        fields = rb_imemo_fields_ptr(fields_obj);
         break;
     }
     return fields[attr_index];
@@ -1382,43 +1346,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 +1367,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_fields));
+            shape_id = RBASIC_SHAPE_ID(obj);
+
+            if (rb_shape_too_complex_p(shape_id)) {
+                st_table *iv_table = rb_imemo_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_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,20 +1407,23 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
             break;
         }
       default:
-        if (FL_TEST_RAW(obj, FL_EXIVAR)) {
-            struct gen_fields_tbl *fields_tbl;
-            rb_gen_fields_tbl_get(obj, id, &fields_tbl);
+        shape_id = RBASIC_SHAPE_ID(obj);
+        if (rb_obj_exivar_p(obj)) {
+            VALUE fields_obj = 0;
+            rb_gen_fields_tbl_get(obj, id, &fields_obj);
 
-            if (rb_shape_obj_too_complex_p(obj)) {
+            RUBY_ASSERT(fields_obj);
+
+            if (rb_shape_obj_too_complex_p(fields_obj)) {
                 VALUE val;
-                if (rb_st_lookup(fields_tbl->as.complex.table, (st_data_t)id, (st_data_t *)&val)) {
+                if (rb_st_lookup(rb_imemo_fields_complex_tbl(fields_obj), (st_data_t)id, (st_data_t *)&val)) {
                     return val;
                 }
                 else {
                     return undef;
                 }
             }
-            ivar_list = fields_tbl->as.shape.fields;
+            ivar_list = rb_imemo_fields_ptr(fields_obj);
         }
         else {
             return undef;
@@ -1493,13 +1458,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_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 +1485,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,34 +1493,43 @@ 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_fields));
+        fields = rb_imemo_fields_ptr(obj);
         break;
       case T_OBJECT:
         fields = ROBJECT_FIELDS(obj);
         break;
       default: {
-        struct gen_fields_tbl *fields_tbl;
-        rb_gen_fields_tbl_get(obj, id, &fields_tbl);
-        fields = fields_tbl->as.shape.fields;
+        VALUE fields_obj;
+        rb_gen_fields_tbl_get(obj, id, &fields_obj);
+        fields = rb_imemo_fields_ptr(fields_obj);
         break;
       }
     }
 
     RUBY_ASSERT(removed_shape_id != INVALID_SHAPE_ID);
 
-    attr_index_t new_fields_count = RSHAPE(next_shape_id)->next_field_index;
-
-    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;
 
-    MEMMOVE(&fields[removed_index], &fields[removed_index + 1], VALUE, trailing_fields);
+    attr_index_t new_fields_count = RSHAPE_LEN(next_shape_id);
+    if (new_fields_count) {
+        size_t trailing_fields = new_fields_count - removed_index;
+
+        MEMMOVE(&fields[removed_index], &fields[removed_index + 1], VALUE, trailing_fields);
+    }
+    else {
+        rb_free_generic_ivar(obj);
+    }
 
     if (RB_TYPE_P(obj, T_OBJECT) &&
         !RB_FL_TEST_RAW(obj, ROBJECT_EMBED) &&
@@ -1563,10 +1543,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 +1551,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_fields));
+            table = rb_imemo_fields_complex_tbl(obj);
             break;
 
           case T_OBJECT:
@@ -1583,9 +1564,9 @@ too_complex:
             break;
 
           default: {
-            struct gen_fields_tbl *fields_tbl;
-            if (rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
-                table = fields_tbl->as.complex.table;
+            VALUE fields_obj;
+            if (rb_gen_fields_tbl_get(obj, 0, &fields_obj)) {
+                table = rb_imemo_fields_complex_tbl(fields_obj);
             }
             break;
           }
@@ -1598,10 +1579,6 @@ too_complex:
         }
     }
 
-    if (locked) {
-        RB_VM_LOCK_LEAVE_LEV(&lev);
-    }
-
     return val;
 }
 
@@ -1611,54 +1588,49 @@ rb_attr_delete(VALUE obj, ID id)
     return rb_ivar_delete(obj, id, Qnil);
 }
 
+static inline void generic_update_fields_obj(VALUE obj, VALUE fields_obj, const VALUE original_fields_obj);
+
 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);
 
-    VALUE *old_fields = NULL;
-
     switch (BUILTIN_TYPE(obj)) {
       case T_OBJECT:
-        if (!(RBASIC(obj)->flags & ROBJECT_EMBED)) {
-            old_fields = ROBJECT_FIELDS(obj);
+        {
+            VALUE *old_fields = NULL;
+            if (!(RBASIC(obj)->flags & ROBJECT_EMBED)) {
+                old_fields = ROBJECT_FIELDS(obj);
+            }
+            RBASIC_SET_SHAPE_ID(obj, shape_id);
+            ROBJECT_SET_FIELDS_HASH(obj, table);
+            if (old_fields) {
+                xfree(old_fields);
+            }
         }
-        RBASIC_SET_SHAPE_ID(obj, shape_id);
-        ROBJECT_SET_FIELDS_HASH(obj, 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() {
-            struct st_table *gen_ivs = generic_fields_tbl_no_ractor_check(obj);
-
-            struct gen_fields_tbl *old_fields_tbl = NULL;
-            st_lookup(gen_ivs, (st_data_t)obj, (st_data_t *)&old_fields_tbl);
-
-            if (old_fields_tbl) {
-                /* We need to modify old_fields_tbl to have the too complex shape
-                 * and hold the table because the xmalloc could trigger a GC
-                 * compaction. We want the table to be updated rather than
-                 * the original fields. */
-                rb_obj_set_shape_id(obj, shape_id);
-                old_fields_tbl->as.complex.table = table;
-                old_fields = (VALUE *)old_fields_tbl;
-            }
-
-            struct gen_fields_tbl *fields_tbl = xmalloc(sizeof(struct gen_fields_tbl));
-            fields_tbl->as.complex.table = table;
-            st_insert(gen_ivs, (st_data_t)obj, (st_data_t)fields_tbl);
+        {
+            VALUE fields_obj = rb_imemo_fields_new_complex_tbl(rb_obj_class(obj), table);
+            RBASIC_SET_SHAPE_ID(fields_obj, shape_id);
 
+            RB_VM_LOCKING() {
+                const VALUE original_fields_obj = generic_fields_lookup(obj, 0, false);
+                generic_update_fields_obj(obj, fields_obj, original_fields_obj);
+            }
             RBASIC_SET_SHAPE_ID(obj, shape_id);
         }
     }
 
-    xfree(old_fields);
     return shape_id;
 }
 
@@ -1673,12 +1645,12 @@ rb_obj_init_too_complex(VALUE obj, st_table *table)
     obj_transition_too_complex(obj, table);
 }
 
+void rb_obj_copy_fields_to_hash_table(VALUE obj, st_table *table);
+
 // Copy all object fields, including ivars and internal object_id, etc
 shape_id_t
 rb_evict_fields_to_hash(VALUE obj)
 {
-    void rb_obj_copy_fields_to_hash_table(VALUE obj, st_table *table);
-
     RUBY_ASSERT(!rb_shape_obj_too_complex_p(obj));
 
     st_table *table = st_init_numtable_with_size(RSHAPE_LEN(RBASIC_SHAPE_ID(obj)));
@@ -1809,138 +1781,174 @@ general_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val, void *data,
     }
 }
 
-struct gen_fields_lookup_ensure_size {
-    VALUE obj;
-    ID id;
-    struct gen_fields_tbl *fields_tbl;
-    shape_id_t shape_id;
-    bool resize;
-};
+static inline void
+generic_update_fields_obj(VALUE obj, VALUE fields_obj, const VALUE original_fields_obj)
+{
+    if (fields_obj != original_fields_obj) {
+        if (original_fields_obj) {
+            // Clear root shape to avoid triggering cleanup such as free_object_id.
+            rb_imemo_fields_clear(original_fields_obj);
+        }
 
-static int
-generic_fields_lookup_ensure_size(st_data_t *k, st_data_t *v, st_data_t u, int existing)
+        generic_fields_insert(obj, fields_obj);
+    }
+}
+
+static void
+generic_ivar_set(VALUE obj, ID id, VALUE val)
 {
-    ASSERT_vm_locking();
+    bool existing = true;
 
-    struct gen_fields_lookup_ensure_size *fields_lookup = (struct gen_fields_lookup_ensure_size *)u;
-    struct gen_fields_tbl *fields_tbl = existing ? (struct gen_fields_tbl *)*v : NULL;
+    VALUE fields_obj = generic_fields_lookup(obj, id, false);
 
-    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));
+    const VALUE original_fields_obj = fields_obj;
+    if (!fields_obj) {
+        fields_obj = rb_imemo_fields_new(rb_obj_class(obj), 1);
+    }
+    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == RBASIC_SHAPE_ID(fields_obj));
+
+    shape_id_t current_shape_id = RBASIC_SHAPE_ID(fields_obj);
+    shape_id_t next_shape_id = current_shape_id;
+
+    if (UNLIKELY(rb_shape_too_complex_p(current_shape_id))) {
+        goto too_complex;
+    }
+
+    attr_index_t index;
+    if (!rb_shape_get_iv_index(current_shape_id, id, &index)) {
+        existing = false;
+
+        index = RSHAPE_LEN(current_shape_id);
+        if (index >= SHAPE_MAX_FIELDS) {
+            rb_raise(rb_eArgError, "too many instance variables");
         }
-        else {
-            FL_SET_RAW((VALUE)*k, FL_EXIVAR);
+
+        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_fields_new_complex(rb_obj_class(obj), current_len + 1);
+            if (current_len) {
+                rb_obj_copy_fields_to_hash_table(original_fields_obj, rb_imemo_fields_complex_tbl(fields_obj));
+            }
+            RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+            goto too_complex;
         }
 
-        fields_tbl = gen_fields_tbl_resize(fields_tbl, RSHAPE_CAPACITY(fields_lookup->shape_id));
-        *v = (st_data_t)fields_tbl;
-    }
+        attr_index_t next_capacity = RSHAPE_CAPACITY(next_shape_id);
+        attr_index_t current_capacity = RSHAPE_CAPACITY(current_shape_id);
 
-    RUBY_ASSERT(FL_TEST((VALUE)*k, FL_EXIVAR));
+        if (next_capacity != current_capacity) {
+            RUBY_ASSERT(next_capacity > current_capacity);
 
-    fields_lookup->fields_tbl = fields_tbl;
-    if (fields_lookup->shape_id) {
-        rb_obj_set_shape_id(fields_lookup->obj, fields_lookup->shape_id);
+            fields_obj = rb_imemo_fields_new(rb_obj_class(obj), next_capacity);
+            if (original_fields_obj) {
+                attr_index_t fields_count = RSHAPE_LEN(current_shape_id);
+                VALUE *fields = rb_imemo_fields_ptr(fields_obj);
+                MEMCPY(fields, rb_imemo_fields_ptr(original_fields_obj), VALUE, fields_count);
+                for (attr_index_t i = 0; i < fields_count; i++) {
+                    RB_OBJ_WRITTEN(fields_obj, Qundef, fields[i]);
+                }
+            }
+        }
+
+        RUBY_ASSERT(RSHAPE(next_shape_id)->type == SHAPE_IVAR);
+        RUBY_ASSERT(index == (RSHAPE_LEN(next_shape_id) - 1));
     }
 
-    return ST_CONTINUE;
-}
+    VALUE *fields = rb_imemo_fields_ptr(fields_obj);
+    RB_OBJ_WRITE(fields_obj, &fields[index], val);
 
-static VALUE *
-generic_ivar_set_shape_fields(VALUE obj, void *data)
-{
-    RUBY_ASSERT(!rb_shape_obj_too_complex_p(obj));
+    if (!existing) {
+        RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+    }
 
-    struct gen_fields_lookup_ensure_size *fields_lookup = data;
+    generic_update_fields_obj(obj, fields_obj, original_fields_obj);
 
-    RB_VM_LOCKING() {
-        st_update(generic_fields_tbl(obj, fields_lookup->id, false), (st_data_t)obj, generic_fields_lookup_ensure_size, (st_data_t)fields_lookup);
+    if (!existing) {
+        RBASIC_SET_SHAPE_ID(obj, next_shape_id);
     }
 
-    FL_SET_RAW(obj, FL_EXIVAR);
+    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == RBASIC_SHAPE_ID(fields_obj));
 
-    return fields_lookup->fields_tbl->as.shape.fields;
-}
+    return;
 
-static void
-generic_ivar_set_shape_resize_fields(VALUE obj, attr_index_t _old_capa, attr_index_t new_capa, void *data)
-{
-    struct gen_fields_lookup_ensure_size *fields_lookup = data;
+too_complex:
+    {
+        st_table *table = rb_imemo_fields_complex_tbl(fields_obj);
+        existing = st_insert(table, (st_data_t)id, (st_data_t)val);
+        RB_OBJ_WRITTEN(fields_obj, Qundef, val);
 
-    fields_lookup->resize = true;
-}
+        generic_update_fields_obj(obj, fields_obj, original_fields_obj);
 
-static void
-generic_ivar_set_set_shape_id(VALUE obj, shape_id_t shape_id, void *data)
-{
-    struct gen_fields_lookup_ensure_size *fields_lookup = data;
+        if (!existing) {
+            RBASIC_SET_SHAPE_ID(obj, next_shape_id);
+        }
+    }
 
-    fields_lookup->shape_id = shape_id;
-}
+    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == RBASIC_SHAPE_ID(fields_obj));
 
-static shape_id_t
-generic_ivar_set_transition_too_complex(VALUE obj, void *_data)
-{
-    shape_id_t new_shape_id = rb_evict_fields_to_hash(obj);
-    FL_SET_RAW(obj, FL_EXIVAR);
-    return new_shape_id;
+    return;
 }
 
-static st_table *
-generic_ivar_set_too_complex_table(VALUE obj, void *data)
+static void
+generic_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val)
 {
-    struct gen_fields_lookup_ensure_size *fields_lookup = data;
+    bool existing = true;
 
-    struct gen_fields_tbl *fields_tbl;
-    if (!rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
-        fields_tbl = xmalloc(sizeof(struct gen_fields_tbl));
-        fields_tbl->as.complex.table = st_init_numtable_with_size(1);
+    VALUE fields_obj = generic_fields_lookup(obj, RSHAPE_EDGE_NAME(target_shape_id), false);
+    const VALUE original_fields_obj = fields_obj;
 
-        RB_VM_LOCKING() {
-            st_insert(generic_fields_tbl(obj, fields_lookup->id, false), (st_data_t)obj, (st_data_t)fields_tbl);
+    shape_id_t current_shape_id = fields_obj ? RBASIC_SHAPE_ID(fields_obj) : ROOT_SHAPE_ID;
+
+    if (UNLIKELY(rb_shape_too_complex_p(target_shape_id))) {
+        if (UNLIKELY(!rb_shape_too_complex_p(current_shape_id))) {
+            attr_index_t current_len = RSHAPE_LEN(current_shape_id);
+            fields_obj = rb_imemo_fields_new_complex(rb_obj_class(obj), current_len + 1);
+            if (current_len) {
+                rb_obj_copy_fields_to_hash_table(original_fields_obj, rb_imemo_fields_complex_tbl(fields_obj));
+            }
+
+            current_shape_id = target_shape_id;
         }
 
-        FL_SET_RAW(obj, FL_EXIVAR);
-    }
+        existing = false;
+        st_table *table = rb_imemo_fields_complex_tbl(fields_obj);
 
-    RUBY_ASSERT(rb_shape_obj_too_complex_p(obj));
+        RUBY_ASSERT(RSHAPE_EDGE_NAME(target_shape_id));
+        st_insert(table, (st_data_t)RSHAPE_EDGE_NAME(target_shape_id), (st_data_t)val);
+        RB_OBJ_WRITTEN(fields_obj, Qundef, val);
+        RBASIC_SET_SHAPE_ID(fields_obj, target_shape_id);
+    }
+    else {
+        attr_index_t index = RSHAPE_INDEX(target_shape_id);
+        if (index >= RSHAPE_CAPACITY(current_shape_id)) {
+            fields_obj = rb_imemo_fields_new(rb_obj_class(obj), index);
+            if (original_fields_obj) {
+                attr_index_t fields_count = RSHAPE_LEN(current_shape_id);
+                VALUE *fields = rb_imemo_fields_ptr(fields_obj);
+                MEMCPY(fields, rb_imemo_fields_ptr(original_fields_obj), VALUE, fields_count);
+                for (attr_index_t i = 0; i < fields_count; i++) {
+                    RB_OBJ_WRITTEN(fields_obj, Qundef, fields[i]);
+                }
+            }
+        }
 
-    return fields_tbl->as.complex.table;
-}
+        VALUE *table = rb_imemo_fields_ptr(fields_obj);
+        RB_OBJ_WRITE(fields_obj, &table[index], val);
 
-static void
-generic_ivar_set(VALUE obj, ID id, VALUE val)
-{
-    struct gen_fields_lookup_ensure_size fields_lookup = {
-        .obj = obj,
-        .id = id,
-        .resize = false,
-    };
+        if (RSHAPE_LEN(target_shape_id) > RSHAPE_LEN(current_shape_id)) {
+            existing = false;
+            RBASIC_SET_SHAPE_ID(fields_obj, target_shape_id);
+        }
+    }
 
-    general_ivar_set(obj, id, val, &fields_lookup,
-                     generic_ivar_set_shape_fields,
-                     generic_ivar_set_shape_resize_fields,
-                     generic_ivar_set_set_shape_id,
-                     generic_ivar_set_transition_too_complex,
-                     generic_ivar_set_too_complex_table);
-}
+    generic_update_fields_obj(obj, fields_obj, original_fields_obj);
 
-static void
-generic_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val)
-{
-    struct gen_fields_lookup_ensure_size fields_lookup = {
-        .obj = obj,
-        .resize = false,
-    };
+    if (!existing) {
+        RBASIC_SET_SHAPE_ID(obj, target_shape_id);
+    }
 
-    general_field_set(obj, target_shape_id, val, &fields_lookup,
-                      generic_ivar_set_shape_fields,
-                      generic_ivar_set_shape_resize_fields,
-                      generic_ivar_set_set_shape_id,
-                      generic_ivar_set_transition_too_complex,
-                      generic_ivar_set_too_complex_table);
+    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == RBASIC_SHAPE_ID(fields_obj));
 }
 
 void
@@ -2052,11 +2060,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 +2134,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 +2143,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 +2163,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_fields));
+            table = rb_imemo_fields_complex_tbl(obj);
             break;
 
           case T_OBJECT:
@@ -2156,11 +2176,10 @@ ivar_defined0(VALUE obj, ID id)
             break;
 
           default: {
-            struct gen_fields_tbl *fields_tbl;
-            if (rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
-                table = fields_tbl->as.complex.table;
+            VALUE fields_obj;
+            if (rb_gen_fields_tbl_get(obj, 0, &fields_obj)) {
+                table = rb_imemo_fields_complex_tbl(fields_obj);
             }
-            break;
           }
         }
 
@@ -2180,12 +2199,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 +2222,46 @@ 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 *fields;
+    switch (BUILTIN_TYPE(itr_data->obj)) {
+      case T_OBJECT:
+        RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
+        fields = ROBJECT_FIELDS(itr_data->obj);
+        break;
+      case T_IMEMO:
+        RUBY_ASSERT(IMEMO_TYPE_P(itr_data->obj, imemo_fields));
+        RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
+
+        fields = rb_imemo_fields_ptr(itr_data->obj);
+        break;
       default:
-        UNREACHABLE_RETURN(false);
+        rb_bug("Unreachable");
     }
+
+    VALUE val = fields[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,81 +2287,55 @@ 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);
-    }
-}
-
-static void
-gen_fields_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg, bool ivar_only)
-{
-    struct gen_fields_tbl *fields_tbl;
-    if (!rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) return;
-
-    struct iv_itr_data itr_data = {
-        .obj = obj,
-        .fields_tbl = fields_tbl,
-        .arg = arg,
-        .func = func,
-        .ivar_only = ivar_only,
-    };
-
-    shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
-    if (rb_shape_too_complex_p(shape_id)) {
-        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 = ROBJECT_FIELDS(obj);
+        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)
+imemo_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_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_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_fields_ptr(fields_obj);
+        iterate_over_shapes(shape_id, func, &itr_data);
     }
 }
 
 void
 rb_copy_generic_ivar(VALUE dest, VALUE obj)
 {
-    struct gen_fields_tbl *obj_fields_tbl;
-    struct gen_fields_tbl *new_fields_tbl;
+    VALUE fields_obj;
+    VALUE new_fields_obj;
 
     rb_check_frozen(dest);
 
-    if (!FL_TEST(obj, FL_EXIVAR)) {
-        goto clear;
-    }
-
-    unsigned long src_num_ivs = rb_ivar_count(obj);
-    if (!src_num_ivs) {
-        goto clear;
+    if (!rb_obj_exivar_p(obj)) {
+        return;
     }
 
     shape_id_t src_shape_id = rb_obj_shape_id(obj);
 
-    if (rb_gen_fields_tbl_get(obj, 0, &obj_fields_tbl)) {
-        if (gen_fields_tbl_count(obj, obj_fields_tbl) == 0)
+    if (rb_gen_fields_tbl_get(obj, 0, &fields_obj)) {
+        unsigned long src_num_ivs = rb_ivar_count(fields_obj);
+        if (!src_num_ivs) {
             goto clear;
-
-        FL_SET(dest, FL_EXIVAR);
+        }
 
         if (rb_shape_too_complex_p(src_shape_id)) {
-            rb_shape_copy_complex_ivars(dest, obj, src_shape_id, obj_fields_tbl->as.complex.table);
+            rb_shape_copy_complex_ivars(dest, obj, src_shape_id, rb_imemo_fields_complex_tbl(fields_obj));
             return;
         }
 
@@ -2361,64 +2343,49 @@ 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))) {
                 st_table *table = rb_st_init_numtable_with_size(src_num_ivs);
                 rb_obj_copy_ivs_to_hash_table(obj, table);
                 rb_obj_init_too_complex(dest, table);
-
                 return;
             }
         }
 
         if (!RSHAPE_LEN(dest_shape_id)) {
             rb_obj_set_shape_id(dest, dest_shape_id);
-            FL_UNSET(dest, FL_EXIVAR);
             return;
         }
 
-        new_fields_tbl = gen_fields_tbl_resize(0, RSHAPE_CAPACITY(dest_shape_id));
-
-        VALUE *src_buf = obj_fields_tbl->as.shape.fields;
-        VALUE *dest_buf = new_fields_tbl->as.shape.fields;
-
+        new_fields_obj = rb_imemo_fields_new(rb_obj_class(dest), RSHAPE_CAPACITY(dest_shape_id));
+        VALUE *src_buf = rb_imemo_fields_ptr(fields_obj);
+        VALUE *dest_buf = rb_imemo_fields_ptr(new_fields_obj);
         rb_shape_copy_fields(dest, dest_buf, dest_shape_id, obj, src_buf, src_shape_id);
+        RBASIC_SET_SHAPE_ID(new_fields_obj, dest_shape_id);
 
-        /*
-         * c.fields_tbl may change in gen_fields_copy due to realloc,
-         * no need to free
-         */
         RB_VM_LOCKING() {
             generic_fields_tbl_no_ractor_check(dest);
-            st_insert(generic_fields_tbl_no_ractor_check(obj), (st_data_t)dest, (st_data_t)new_fields_tbl);
+            st_insert(generic_fields_tbl_no_ractor_check(obj), (st_data_t)dest, (st_data_t)new_fields_obj);
+            RB_OBJ_WRITTEN(dest, Qundef, new_fields_obj);
         }
 
-        rb_obj_set_shape_id(dest, dest_shape_id);
+        RBASIC_SET_SHAPE_ID(dest, dest_shape_id);
     }
     return;
 
   clear:
-    if (FL_TEST(dest, FL_EXIVAR)) {
-        RBASIC_SET_SHAPE_ID(dest, ROOT_SHAPE_ID);
-        rb_free_generic_ivar(dest);
-        FL_UNSET(dest, FL_EXIVAR);
-    }
+    rb_free_generic_ivar(dest);
 }
 
 void
 rb_replace_generic_ivar(VALUE clone, VALUE obj)
 {
-    RUBY_ASSERT(FL_TEST(obj, FL_EXIVAR));
-
     RB_VM_LOCKING() {
         st_data_t fields_tbl, obj_data = (st_data_t)obj;
         if (st_delete(generic_fields_tbl_, &obj_data, &fields_tbl)) {
-            FL_UNSET_RAW(obj, FL_EXIVAR);
-
             st_insert(generic_fields_tbl_, (st_data_t)clone, fields_tbl);
-            FL_SET_RAW(clone, FL_EXIVAR);
         }
         else {
             rb_bug("unreachable");
@@ -2431,19 +2398,30 @@ 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_fields)) {
+            imemo_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) {
+                imemo_fields_each(fields_obj, func, arg, ivar_only);
+            }
         }
         break;
       default:
-        if (FL_TEST(obj, FL_EXIVAR)) {
-            gen_fields_each(obj, func, arg, ivar_only);
+        if (rb_obj_exivar_p(obj)) {
+            VALUE fields_obj = 0;
+            if (!rb_gen_fields_tbl_get(obj, 0, &fields_obj)) return;
+
+            imemo_fields_each(fields_obj, func, arg, ivar_only);
         }
         break;
     }
@@ -2465,16 +2443,46 @@ rb_ivar_count(VALUE obj)
       case T_OBJECT:
         iv_count = ROBJECT_FIELDS_COUNT(obj);
         break;
+
       case T_CLASS:
       case T_MODULE:
-        iv_count = RCLASS_FIELDS_COUNT(obj);
+        {
+            VALUE fields_obj = RCLASS_WRITABLE_FIELDS_OBJ(obj);
+            if (!fields_obj) {
+                return 0;
+            }
+            if (rb_shape_obj_too_complex_p(fields_obj)) {
+                iv_count = rb_st_table_size(rb_imemo_fields_complex_tbl(fields_obj));
+            }
+            else {
+                iv_count = RBASIC_FIELDS_COUNT(fields_obj);
+            }
+        }
         break;
+
+      case T_IMEMO:
+        RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_fields));
+
+        if (rb_shape_obj_too_complex_p(obj)) {
+            iv_count = rb_st_table_size(rb_imemo_fields_complex_tbl(obj));
+        }
+        else {
+            iv_count = RBASIC_FIELDS_COUNT(obj);
+        }
+        break;
+
       default:
-        if (FL_TEST(obj, FL_EXIVAR)) {
-            struct gen_fields_tbl *fields_tbl;
+        if (rb_obj_exivar_p(obj)) {
+
+            if (rb_shape_obj_too_complex_p(obj)) {
+                VALUE fields_obj;
 
-            if (rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
-                iv_count = gen_fields_tbl_count(obj, fields_tbl);
+                if (rb_gen_fields_tbl_get(obj, 0, &fields_obj)) {
+                    iv_count = rb_st_table_size(rb_imemo_fields_complex_tbl(fields_obj));
+                }
+            }
+            else {
+                iv_count = RBASIC_FIELDS_COUNT(obj);
             }
         }
         break;
@@ -4674,71 +4682,107 @@ 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_fields_new(rb_singleton_class(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_fields_new_complex(rb_singleton_class(klass), current_len + 1);
+            if (current_len) {
+                rb_obj_copy_fields_to_hash_table(original_fields_obj, rb_imemo_fields_complex_tbl(fields_obj));
+                RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+            }
+            goto too_complex;
+        }
 
-    return RCLASS_WRITABLE_FIELDS_HASH(obj);
+        attr_index_t next_capacity = RSHAPE_CAPACITY(next_shape_id);
+        attr_index_t current_capacity = RSHAPE_CAPACITY(current_shape_id);
+
+        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_fields_new(rb_singleton_class(klass), next_capacity);
+            if (original_fields_obj) {
+                MEMCPY(rb_imemo_fields_ptr(fields_obj), rb_imemo_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_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_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;
+
+    bool existing = class_fields_ivar_set(obj, original_fields_obj, id, val, rb_multi_ractor_p(), &new_fields_obj);
+
+    if (new_fields_obj != original_fields_obj) {
+        RCLASS_WRITABLE_SET_FIELDS_OBJ(obj, new_fields_obj);
     }
 
-    return existing;
-}
+    // 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));
 
-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);
+    return existing;
 }
 
 static int
@@ -4754,9 +4798,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);
 }