10 files changed, 366 insertions, 402 deletions
diff --git a/ext/objspace/objspace_dump.c b/ext/objspace/objspace_dump.c
index 83b434c3a1..80732d0282 100644
--- a/ext/objspace/objspace_dump.c
+++ b/ext/objspace/objspace_dump.c
@@ -394,9 +394,10 @@ dump_object(VALUE obj, struct dump_config *dc)
 
     dc->cur_obj = obj;
     dc->cur_obj_references = 0;
-    if (BUILTIN_TYPE(obj) == T_NODE || BUILTIN_TYPE(obj) == T_IMEMO) {
+    if (BUILTIN_TYPE(obj) == T_NODE || (BUILTIN_TYPE(obj) == T_IMEMO && !IMEMO_TYPE_P(obj, imemo_fields))) {
         dc->cur_obj_klass = 0;
-    } else {
+    }
+    else {
         dc->cur_obj_klass = RBASIC_CLASS(obj);
     }
 
@@ -414,8 +415,8 @@ dump_object(VALUE obj, struct dump_config *dc)
     dump_append(dc, obj_type(obj));
     dump_append(dc, "\"");
 
-    if (BUILTIN_TYPE(obj) != T_IMEMO) {
-        size_t shape_id = rb_obj_shape_id(obj);
+    if (BUILTIN_TYPE(obj) != T_IMEMO || IMEMO_TYPE_P(obj, imemo_fields)) {
+        size_t shape_id = rb_obj_shape_id(obj) & SHAPE_ID_OFFSET_MASK;
         dump_append(dc, ", \"shape_id\":");
         dump_append_sizet(dc, shape_id);
     }
diff --git a/gc.c b/gc.c
index f0189294bd..b0876fca5e 100644
--- a/gc.c
+++ b/gc.c
@@ -2015,27 +2015,6 @@ object_id_to_ref(void *objspace_ptr, VALUE object_id)
 static inline void
 obj_free_object_id(VALUE obj)
 {
-    if (RB_BUILTIN_TYPE(obj) == T_IMEMO) {
-        return;
-    }
-
-#if RUBY_DEBUG
-    switch (BUILTIN_TYPE(obj)) {
-      case T_CLASS:
-      case T_MODULE:
-        break;
-      default:
-        if (rb_shape_obj_has_id(obj)) {
-            VALUE id = object_id_get(obj, RBASIC_SHAPE_ID(obj)); // Crash if missing
-            if (!(FIXNUM_P(id) || RB_TYPE_P(id, T_BIGNUM))) {
-                rb_p(obj);
-                rb_bug("Corrupted object_id");
-            }
-        }
-        break;
-    }
-#endif
-
     VALUE obj_id = 0;
     if (RB_UNLIKELY(id2ref_tbl)) {
         switch (BUILTIN_TYPE(obj)) {
@@ -2043,21 +2022,32 @@ obj_free_object_id(VALUE obj)
           case T_MODULE:
             obj_id = RCLASS(obj)->object_id;
             break;
-          default: {
+          case T_IMEMO:
+            if (!IMEMO_TYPE_P(obj, imemo_fields)) {
+                return;
+            }
+            // fallthrough
+          case T_OBJECT:
+            {
             shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
             if (rb_shape_has_object_id(shape_id)) {
                 obj_id = object_id_get(obj, shape_id);
             }
             break;
           }
+          default:
+            // For generic_fields, the T_IMEMO/fields is responsible for freeing the id.
+            return;
         }
 
         if (RB_UNLIKELY(obj_id)) {
             RUBY_ASSERT(FIXNUM_P(obj_id) || RB_TYPE_P(obj_id, T_BIGNUM));
 
             if (!st_delete(id2ref_tbl, (st_data_t *)&obj_id, NULL)) {
-                // If we're currently building the table then it's not a bug
-                if (id2ref_tbl_built) {
+                // If we're currently building the table then it's not a bug.
+                // The the object is a T_IMEMO/fields, then it's possible the actual object
+                // has been garbage collected already.
+                if (id2ref_tbl_built && !RB_TYPE_P(obj, T_IMEMO)) {
                     rb_bug("Object ID seen, but not in _id2ref table: object_id=%llu object=%s", NUM2ULL(obj_id), rb_obj_info(obj));
                 }
             }
@@ -2071,7 +2061,7 @@ rb_gc_obj_free_vm_weak_references(VALUE obj)
     obj_free_object_id(obj);
 
     if (rb_obj_exivar_p(obj)) {
-        rb_free_generic_ivar((VALUE)obj);
+        rb_free_generic_ivar(obj);
     }
 
     switch (BUILTIN_TYPE(obj)) {
@@ -2316,10 +2306,6 @@ rb_obj_memsize_of(VALUE obj)
         return 0;
     }
 
-    if (rb_obj_exivar_p(obj)) {
-        size += rb_generic_ivar_memsize(obj);
-    }
-
     switch (BUILTIN_TYPE(obj)) {
       case T_OBJECT:
         if (rb_shape_obj_too_complex_p(obj)) {
@@ -3935,38 +3921,6 @@ vm_weak_table_foreach_update_weak_value(st_data_t *key, st_data_t *value, st_dat
     return iter_data->update_callback((VALUE *)value, iter_data->data);
 }
 
-static void
-free_gen_fields_tbl(VALUE obj, struct gen_fields_tbl *fields_tbl)
-{
-    if (UNLIKELY(rb_shape_obj_too_complex_p(obj))) {
-        st_free_table(fields_tbl->as.complex.table);
-    }
-
-    xfree(fields_tbl);
-}
-
-static int
-vm_weak_table_gen_fields_foreach_too_complex_i(st_data_t _key, st_data_t value, st_data_t data, int error)
-{
-    struct global_vm_table_foreach_data *iter_data = (struct global_vm_table_foreach_data *)data;
-
-    GC_ASSERT(!iter_data->weak_only);
-
-    if (SPECIAL_CONST_P((VALUE)value)) return ST_CONTINUE;
-
-    return iter_data->callback((VALUE)value, iter_data->data);
-}
-
-static int
-vm_weak_table_gen_fields_foreach_too_complex_replace_i(st_data_t *_key, st_data_t *value, st_data_t data, int existing)
-{
-    struct global_vm_table_foreach_data *iter_data = (struct global_vm_table_foreach_data *)data;
-
-    GC_ASSERT(!iter_data->weak_only);
-
-    return iter_data->update_callback((VALUE *)value, iter_data->data);
-}
-
 struct st_table *rb_generic_fields_tbl_get(void);
 
 static int
@@ -4003,60 +3957,50 @@ vm_weak_table_gen_fields_foreach(st_data_t key, st_data_t value, st_data_t data)
 
     int ret = iter_data->callback((VALUE)key, iter_data->data);
 
+    VALUE new_value = (VALUE)value;
+    VALUE new_key = (VALUE)key;
+
     switch (ret) {
       case ST_CONTINUE:
         break;
 
       case ST_DELETE:
-        free_gen_fields_tbl((VALUE)key, (struct gen_fields_tbl *)value);
         RBASIC_SET_SHAPE_ID((VALUE)key, ROOT_SHAPE_ID);
         return ST_DELETE;
 
       case ST_REPLACE: {
-        VALUE new_key = (VALUE)key;
         ret = iter_data->update_callback(&new_key, iter_data->data);
-        if (key != new_key) ret = ST_DELETE;
-        DURING_GC_COULD_MALLOC_REGION_START();
-        {
-            st_insert(rb_generic_fields_tbl_get(), (st_data_t)new_key, value);
+        if (key != new_key) {
+            ret = ST_DELETE;
         }
-        DURING_GC_COULD_MALLOC_REGION_END();
-        key = (st_data_t)new_key;
         break;
       }
 
       default:
-        return ret;
+        rb_bug("vm_weak_table_gen_fields_foreach: return value %d not supported", ret);
     }
 
     if (!iter_data->weak_only) {
-        struct gen_fields_tbl *fields_tbl = (struct gen_fields_tbl *)value;
+        int ivar_ret = iter_data->callback(new_value, iter_data->data);
+        switch (ivar_ret) {
+          case ST_CONTINUE:
+            break;
 
-        if (rb_shape_obj_too_complex_p((VALUE)key)) {
-            st_foreach_with_replace(
-                fields_tbl->as.complex.table,
-                vm_weak_table_gen_fields_foreach_too_complex_i,
-                vm_weak_table_gen_fields_foreach_too_complex_replace_i,
-                data
-            );
+          case ST_REPLACE:
+            iter_data->update_callback(&new_value, iter_data->data);
+            break;
+
+          default:
+            rb_bug("vm_weak_table_gen_fields_foreach: return value %d not supported", ivar_ret);
         }
-        else {
-            uint32_t fields_count = RSHAPE_LEN(RBASIC_SHAPE_ID((VALUE)key));
-            for (uint32_t i = 0; i < fields_count; i++) {
-                if (SPECIAL_CONST_P(fields_tbl->as.shape.fields[i])) continue;
+    }
 
-                int ivar_ret = iter_data->callback(fields_tbl->as.shape.fields[i], iter_data->data);
-                switch (ivar_ret) {
-                  case ST_CONTINUE:
-                    break;
-                  case ST_REPLACE:
-                    iter_data->update_callback(&fields_tbl->as.shape.fields[i], iter_data->data);
-                    break;
-                  default:
-                    rb_bug("vm_weak_table_gen_fields_foreach: return value %d not supported", ivar_ret);
-                }
-            }
+    if (key != new_key || value != new_value) {
+        DURING_GC_COULD_MALLOC_REGION_START();
+        {
+            st_insert(rb_generic_fields_tbl_get(), (st_data_t)new_key, new_value);
         }
+        DURING_GC_COULD_MALLOC_REGION_END();
     }
 
     return ret;
diff --git a/imemo.c b/imemo.c
index f465c0098b..a4393ffe79 100644
--- a/imemo.c
+++ b/imemo.c
@@ -147,6 +147,23 @@ rb_imemo_fields_new_complex(VALUE klass, size_t capa)
     return imemo_fields_new_complex(klass, capa);
 }
 
+static int
+imemo_fields_trigger_wb_i(st_data_t key, st_data_t value, st_data_t arg)
+{
+    VALUE field_obj = (VALUE)arg;
+    RB_OBJ_WRITTEN(field_obj, Qundef, (VALUE)value);
+    return ST_CONTINUE;
+}
+
+VALUE
+rb_imemo_fields_new_complex_tbl(VALUE klass, st_table *tbl)
+{
+    VALUE fields = imemo_fields_new(klass, sizeof(struct rb_fields));
+    IMEMO_OBJ_FIELDS(fields)->as.complex.table = tbl;
+    st_foreach(tbl, imemo_fields_trigger_wb_i, (st_data_t)fields);
+    return fields;
+}
+
 VALUE
 rb_imemo_fields_clone(VALUE fields_obj)
 {
@@ -168,6 +185,19 @@ rb_imemo_fields_clone(VALUE fields_obj)
     return clone;
 }
 
+void
+rb_imemo_fields_clear(VALUE fields_obj)
+{
+    // When replacing an imemo/fields by another one, we must clear
+    // its shape so that gc.c:obj_free_object_id won't be called.
+    if (rb_shape_obj_too_complex_p(fields_obj)) {
+        RBASIC_SET_SHAPE_ID(fields_obj, ROOT_TOO_COMPLEX_SHAPE_ID);
+    }
+    else {
+        RBASIC_SET_SHAPE_ID(fields_obj, ROOT_SHAPE_ID);
+    }
+}
+
 /* =========================================================================
  * memsize
  * ========================================================================= */
diff --git a/internal/imemo.h b/internal/imemo.h
index 849748f92f..44b41d1b1c 100644
--- a/internal/imemo.h
+++ b/internal/imemo.h
@@ -280,7 +280,9 @@ struct rb_fields {
 
 VALUE rb_imemo_fields_new(VALUE klass, size_t capa);
 VALUE rb_imemo_fields_new_complex(VALUE klass, size_t capa);
+VALUE rb_imemo_fields_new_complex_tbl(VALUE klass, st_table *tbl);
 VALUE rb_imemo_fields_clone(VALUE fields_obj);
+void rb_imemo_fields_clear(VALUE fields_obj);
 
 static inline VALUE *
 rb_imemo_fields_ptr(VALUE obj_fields)
diff --git a/internal/variable.h b/internal/variable.h
index 8da6c678a5..92017d6184 100644
--- a/internal/variable.h
+++ b/internal/variable.h
@@ -18,7 +18,6 @@
 /* variable.c */
 void rb_gc_mark_global_tbl(void);
 void rb_gc_update_global_tbl(void);
-size_t rb_generic_ivar_memsize(VALUE);
 VALUE rb_search_class_path(VALUE);
 VALUE rb_attr_delete(VALUE, ID);
 void rb_autoload_str(VALUE mod, ID id, VALUE file);
@@ -47,8 +46,7 @@ void rb_gvar_namespace_ready(const char *name);
  */
 VALUE rb_mod_set_temporary_name(VALUE, VALUE);
 
-struct gen_fields_tbl;
-int rb_gen_fields_tbl_get(VALUE obj, ID id, struct gen_fields_tbl **fields_tbl);
+int rb_gen_fields_tbl_get(VALUE obj, ID id, VALUE *fields_obj);
 void rb_obj_copy_ivs_to_hash_table(VALUE obj, st_table *table);
 void rb_obj_init_too_complex(VALUE obj, st_table *table);
 void rb_evict_ivars_to_hash(VALUE obj);
diff --git a/ractor.c b/ractor.c
index cce376c543..a4a746b495 100644
--- a/ractor.c
+++ b/ractor.c
@@ -1657,8 +1657,8 @@ obj_traverse_replace_i(VALUE obj, struct obj_traverse_replace_data *data)
 } while (0)
 
     if (UNLIKELY(rb_obj_exivar_p(obj))) {
-        struct gen_fields_tbl *fields_tbl;
-        rb_ivar_generic_fields_tbl_lookup(obj, &fields_tbl);
+        VALUE fields_obj;
+        rb_ivar_generic_fields_tbl_lookup(obj, &fields_obj);
 
         if (UNLIKELY(rb_shape_obj_too_complex_p(obj))) {
             struct obj_traverse_replace_callback_data d = {
@@ -1667,7 +1667,7 @@ obj_traverse_replace_i(VALUE obj, struct obj_traverse_replace_data *data)
                 .src = obj,
             };
             rb_st_foreach_with_replace(
-                fields_tbl->as.complex.table,
+                rb_imemo_fields_complex_tbl(fields_obj),
                 obj_iv_hash_traverse_replace_foreach_i,
                 obj_iv_hash_traverse_replace_i,
                 (st_data_t)&d
@@ -1676,8 +1676,9 @@ obj_traverse_replace_i(VALUE obj, struct obj_traverse_replace_data *data)
         }
         else {
             uint32_t fields_count = RSHAPE_LEN(RBASIC_SHAPE_ID(obj));
+            VALUE *fields = rb_imemo_fields_ptr(fields_obj);
             for (uint32_t i = 0; i < fields_count; i++) {
-                CHECK_AND_REPLACE(fields_tbl->as.shape.fields[i]);
+                CHECK_AND_REPLACE(fields[i]);
             }
         }
     }
diff --git a/test/ruby/test_encoding.rb b/test/ruby/test_encoding.rb
index ee37199be0..0ab357f53a 100644
--- a/test/ruby/test_encoding.rb
+++ b/test/ruby/test_encoding.rb
@@ -33,7 +33,7 @@ class TestEncoding < Test::Unit::TestCase
     encodings.each do |e|
       assert_raise(TypeError) { e.dup }
       assert_raise(TypeError) { e.clone }
-      assert_equal(e.object_id, Marshal.load(Marshal.dump(e)).object_id)
+      assert_same(e, Marshal.load(Marshal.dump(e)))
     end
   end
 
diff --git a/variable.c b/variable.c
index 697b6c2c28..b71c3981f3 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,33 +1239,17 @@ 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)
-{
-    return rb_gen_fields_tbl_get(obj, 0, fields_tbl);
-}
-
-static size_t
-gen_fields_tbl_bytes(size_t n)
+rb_ivar_generic_fields_tbl_lookup(VALUE obj, VALUE *fields_obj)
 {
-    return offsetof(struct gen_fields_tbl, as.shape.fields) + n * sizeof(VALUE);
+    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 {
-            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]);
-            }
-        }
+    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);
     }
 }
 
@@ -1252,47 +1259,9 @@ rb_free_generic_ivar(VALUE obj)
     if (rb_obj_exivar_p(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);
-            }
+            st_delete(generic_fields_tbl_no_ractor_check(obj), &key, &value);
         }
-        RBASIC_SET_SHAPE_ID(obj, ROOT_SHAPE_ID);
-    }
-}
-
-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(RSHAPE_CAPACITY(RBASIC_SHAPE_ID(obj)));
-        }
-    }
-    return 0;
-}
-
-static size_t
-gen_fields_tbl_count(VALUE obj, const struct gen_fields_tbl *fields_tbl)
-{
-    if (rb_shape_obj_too_complex_p(obj)) {
-        return st_table_size(fields_tbl->as.complex.table);
-    }
-    else {
-        return RSHAPE_LEN(RBASIC_SHAPE_ID(obj));
     }
 }
 
@@ -1321,12 +1290,16 @@ rb_obj_field_get(VALUE obj, shape_id_t target_shape_id)
           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(rb_obj_exivar_p(obj));
-            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;
+            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;
@@ -1352,12 +1325,16 @@ rb_obj_field_get(VALUE obj, shape_id_t target_shape_id)
       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(rb_obj_exivar_p(obj));
-        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;
+        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];
@@ -1432,19 +1409,21 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
       default:
         shape_id = RBASIC_SHAPE_ID(obj);
         if (rb_obj_exivar_p(obj)) {
-            struct gen_fields_tbl *fields_tbl;
-            rb_gen_fields_tbl_get(obj, id, &fields_tbl);
+            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;
@@ -1530,9 +1509,9 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
         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;
       }
     }
@@ -1585,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;
           }
@@ -1609,6 +1588,8 @@ 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)
 {
@@ -1619,46 +1600,37 @@ obj_transition_too_complex(VALUE obj, st_table *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:
         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,135 +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;
-    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 VALUE *
-generic_ivar_set_shape_fields(VALUE obj, void *data)
+        generic_fields_insert(obj, fields_obj);
+    }
+}
+
+static void
+generic_ivar_set(VALUE obj, ID id, VALUE val)
 {
-    RUBY_ASSERT(!rb_shape_obj_too_complex_p(obj));
+    bool existing = true;
 
-    struct gen_fields_lookup_ensure_size *fields_lookup = data;
-    struct gen_fields_tbl *fields_tbl = NULL;
+    VALUE fields_obj = generic_fields_lookup(obj, id, false);
 
-    // We can't use st_update, since when resizing the fields table GC can
-    // happen, which will modify the st_table and may rebuild it
-    RB_VM_LOCKING() {
-        st_table *tbl = generic_fields_tbl(obj, fields_lookup->id, false);
-        int existing = st_lookup(tbl, (st_data_t)obj, (st_data_t *)&fields_tbl);
+    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));
 
-        if (!existing || fields_lookup->resize) {
-            uint32_t new_capa = RSHAPE_CAPACITY(fields_lookup->shape_id);
-            uint32_t old_capa = RSHAPE_CAPACITY(RSHAPE_PARENT(fields_lookup->shape_id));
+    shape_id_t current_shape_id = RBASIC_SHAPE_ID(fields_obj);
+    shape_id_t next_shape_id = current_shape_id;
 
-            if (existing) {
-                RUBY_ASSERT(RSHAPE_TYPE_P(fields_lookup->shape_id, SHAPE_IVAR) || RSHAPE_TYPE_P(fields_lookup->shape_id, SHAPE_OBJ_ID));
-                RUBY_ASSERT(old_capa < new_capa);
-                RUBY_ASSERT(fields_tbl);
-            }
-            else {
-                RUBY_ASSERT(!fields_tbl);
-                RUBY_ASSERT(old_capa == 0);
-            }
-            RUBY_ASSERT(new_capa > 0);
+    if (UNLIKELY(rb_shape_too_complex_p(current_shape_id))) {
+        goto too_complex;
+    }
 
-            struct gen_fields_tbl *old_fields_tbl = fields_tbl;
-            fields_tbl = xmalloc(gen_fields_tbl_bytes(new_capa));
-            if (old_fields_tbl) {
-                memcpy(fields_tbl, old_fields_tbl, gen_fields_tbl_bytes(old_capa));
-            }
-            st_insert(tbl, (st_data_t)obj, (st_data_t)fields_tbl);
-            if (old_fields_tbl) {
-                xfree(old_fields_tbl);
+    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");
+        }
+
+        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;
         }
 
-        if (fields_lookup->shape_id) {
-            rb_obj_set_shape_id(fields_lookup->obj, fields_lookup->shape_id);
+        attr_index_t next_capacity = RSHAPE_CAPACITY(next_shape_id);
+        attr_index_t current_capacity = RSHAPE_CAPACITY(current_shape_id);
+
+        if (next_capacity != current_capacity) {
+            RUBY_ASSERT(next_capacity > current_capacity);
+
+            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 fields_tbl->as.shape.fields;
-}
+    VALUE *fields = rb_imemo_fields_ptr(fields_obj);
+    RB_OBJ_WRITE(fields_obj, &fields[index], val);
 
-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;
+    if (!existing) {
+        RBASIC_SET_SHAPE_ID(fields_obj, next_shape_id);
+    }
 
-    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);
-    return new_shape_id;
-}
+    return;
 
-static st_table *
-generic_ivar_set_too_complex_table(VALUE obj, 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);
 
-    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);
+        generic_update_fields_obj(obj, fields_obj, original_fields_obj);
 
-        RB_VM_LOCKING() {
-            st_insert(generic_fields_tbl(obj, fields_lookup->id, false), (st_data_t)obj, (st_data_t)fields_tbl);
+        if (!existing) {
+            RBASIC_SET_SHAPE_ID(obj, next_shape_id);
         }
     }
 
-    RUBY_ASSERT(rb_shape_obj_too_complex_p(obj));
+    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == RBASIC_SHAPE_ID(fields_obj));
 
-    return fields_tbl->as.complex.table;
+    return;
 }
 
 static void
-generic_ivar_set(VALUE obj, ID id, VALUE val)
+generic_field_set(VALUE obj, shape_id_t target_shape_id, VALUE val)
 {
-    struct gen_fields_lookup_ensure_size fields_lookup = {
-        .obj = obj,
-        .id = id,
-        .resize = false,
-    };
+    bool existing = true;
 
-    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);
-}
+    VALUE fields_obj = generic_fields_lookup(obj, RSHAPE_EDGE_NAME(target_shape_id), false);
+    const VALUE original_fields_obj = 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,
-    };
+    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;
+        }
+
+        existing = false;
+        st_table *table = rb_imemo_fields_complex_tbl(fields_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]);
+                }
+            }
+        }
+
+        VALUE *table = rb_imemo_fields_ptr(fields_obj);
+        RB_OBJ_WRITE(fields_obj, &table[index], val);
 
-    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);
+        if (RSHAPE_LEN(target_shape_id) > RSHAPE_LEN(current_shape_id)) {
+            existing = false;
+            RBASIC_SET_SHAPE_ID(fields_obj, target_shape_id);
+        }
+    }
+
+    generic_update_fields_obj(obj, fields_obj, original_fields_obj);
+
+    if (!existing) {
+        RBASIC_SET_SHAPE_ID(obj, target_shape_id);
+    }
+
+    RUBY_ASSERT(RBASIC_SHAPE_ID(obj) == RBASIC_SHAPE_ID(fields_obj));
 }
 
 void
@@ -2165,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;
           }
         }
 
@@ -2225,27 +2235,23 @@ iterate_over_shapes_callback(shape_id_t shape_id, void *data)
         return ST_CONTINUE;
     }
 
-    VALUE *iv_list;
+    VALUE *fields;
     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);
+        fields = 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_fields));
         RUBY_ASSERT(!rb_shape_obj_too_complex_p(itr_data->obj));
 
-        iv_list = rb_imemo_fields_ptr(itr_data->obj);
+        fields = rb_imemo_fields_ptr(itr_data->obj);
         break;
       default:
-        iv_list = itr_data->fields_tbl->as.shape.fields;
-        break;
+        rb_bug("Unreachable");
     }
 
-    VALUE val = iv_list[RSHAPE_INDEX(shape_id)];
+    VALUE val = fields[RSHAPE_INDEX(shape_id)];
     return itr_data->func(RSHAPE_EDGE_NAME(shape_id), val, itr_data->arg);
 }
 
@@ -2287,31 +2293,7 @@ obj_fields_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg, b
 }
 
 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 {
-        itr_data.fields = fields_tbl->as.shape.fields;
-        iterate_over_shapes(shape_id, func, &itr_data);
-    }
-}
-
-static void
-class_fields_each(VALUE fields_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)
 {
     IMEMO_TYPE_P(fields_obj, imemo_fields);
 
@@ -2335,8 +2317,8 @@ class_fields_each(VALUE fields_obj, rb_ivar_foreach_callback_func *func, st_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);
 
@@ -2344,19 +2326,16 @@ rb_copy_generic_ivar(VALUE dest, VALUE obj)
         return;
     }
 
-    unsigned long src_num_ivs = rb_ivar_count(obj);
-    if (!src_num_ivs) {
-        goto clear;
-    }
-
     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;
+        }
 
         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;
         }
 
@@ -2371,7 +2350,6 @@ rb_copy_generic_ivar(VALUE dest, VALUE obj)
                 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;
             }
         }
@@ -2381,25 +2359,19 @@ rb_copy_generic_ivar(VALUE dest, VALUE obj)
             return;
         }
 
-        uint32_t dest_capa = RSHAPE_CAPACITY(dest_shape_id);
-        RUBY_ASSERT(dest_capa > 0);
-        new_fields_tbl = xmalloc(gen_fields_tbl_bytes(dest_capa));
-
-        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;
 
@@ -2428,7 +2400,7 @@ rb_field_foreach(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg,
     switch (BUILTIN_TYPE(obj)) {
       case T_IMEMO:
         if (IMEMO_TYPE_P(obj, imemo_fields)) {
-            class_fields_each(obj, func, arg, ivar_only);
+            imemo_fields_each(obj, func, arg, ivar_only);
         }
         break;
       case T_OBJECT:
@@ -2440,13 +2412,16 @@ rb_field_foreach(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg,
             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);
+                imemo_fields_each(fields_obj, func, arg, ivar_only);
             }
         }
         break;
       default:
         if (rb_obj_exivar_p(obj)) {
-            gen_fields_each(obj, func, arg, ivar_only);
+            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;
     }
@@ -2468,6 +2443,7 @@ rb_ivar_count(VALUE obj)
       case T_OBJECT:
         iv_count = ROBJECT_FIELDS_COUNT(obj);
         break;
+
       case T_CLASS:
       case T_MODULE:
         {
@@ -2476,16 +2452,37 @@ rb_ivar_count(VALUE obj)
                 return 0;
             }
             if (rb_shape_obj_too_complex_p(fields_obj)) {
-                return rb_st_table_size(rb_imemo_fields_complex_tbl(fields_obj));
+                iv_count = rb_st_table_size(rb_imemo_fields_complex_tbl(fields_obj));
+            }
+            else {
+                iv_count = RBASIC_FIELDS_COUNT(fields_obj);
             }
-            return 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 (rb_obj_exivar_p(obj)) {
-            struct gen_fields_tbl *fields_tbl;
 
-            if (rb_gen_fields_tbl_get(obj, 0, &fields_tbl)) {
-                iv_count = gen_fields_tbl_count(obj, fields_tbl);
+            if (rb_shape_obj_too_complex_p(obj)) {
+                VALUE fields_obj;
+
+                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;
diff --git a/variable.h b/variable.h
index 54b7fc5461..82a79c63ce 100644
--- a/variable.h
+++ b/variable.h
@@ -12,18 +12,7 @@
 
 #include "shape.h"
 
-struct gen_fields_tbl {
-    union {
-        struct {
-            VALUE fields[1];
-        } shape;
-        struct {
-            st_table *table;
-        } complex;
-    } as;
-};
-
-int rb_ivar_generic_fields_tbl_lookup(VALUE obj, struct gen_fields_tbl **);
+int rb_ivar_generic_fields_tbl_lookup(VALUE obj, VALUE *);
 void rb_copy_complex_ivars(VALUE dest, VALUE obj, shape_id_t src_shape_id, st_table *fields_table);
 
 void rb_free_rb_global_tbl(void);
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index 689521eaae..2fe5e26928 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -1259,9 +1259,11 @@ vm_getivar(VALUE obj, ID id, const rb_iseq_t *iseq, IVC ic, const struct rb_call
         }
       default:
         if (rb_obj_exivar_p(obj)) {
-            struct gen_fields_tbl *fields_tbl;
-            rb_gen_fields_tbl_get(obj, id, &fields_tbl);
-            ivar_list = fields_tbl->as.shape.fields;
+            VALUE fields_obj = 0;
+            if (!rb_gen_fields_tbl_get(obj, id, &fields_obj)) {
+                return default_value;
+            }
+            ivar_list = rb_imemo_fields_ptr(fields_obj);
         }
         else {
             return default_value;
@@ -1333,9 +1335,9 @@ vm_getivar(VALUE obj, ID id, const rb_iseq_t *iseq, IVC ic, const struct rb_call
                 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;
               }
@@ -1456,7 +1458,7 @@ vm_setivar_default(VALUE obj, ID id, VALUE val, shape_id_t dest_shape_id, attr_i
 {
     shape_id_t shape_id = RBASIC_SHAPE_ID(obj);
 
-    struct gen_fields_tbl *fields_tbl = 0;
+    VALUE fields_obj = 0;
 
     // Cache hit case
     if (shape_id == dest_shape_id) {
@@ -1474,13 +1476,13 @@ vm_setivar_default(VALUE obj, ID id, VALUE val, shape_id_t dest_shape_id, attr_i
         return Qundef;
     }
 
-    rb_gen_fields_tbl_get(obj, 0, &fields_tbl);
+    rb_gen_fields_tbl_get(obj, 0, &fields_obj);
 
     if (shape_id != dest_shape_id) {
         RBASIC_SET_SHAPE_ID(obj, dest_shape_id);
     }
 
-    RB_OBJ_WRITE(obj, &fields_tbl->as.shape.fields[index], val);
+    RB_OBJ_WRITE(obj, &rb_imemo_fields_ptr(fields_obj)[index], val);
 
     RB_DEBUG_COUNTER_INC(ivar_set_ic_hit);