WIP

Author: byroot

class.c

@@ -278,7 +278,7 @@ class_duplicate_iclass_classext(VALUE iclass, rb_classext_t *mod_ext, const rb_n
     else {
         RCLASSEXT_M_TBL(ext) = RCLASSEXT_M_TBL(mod_ext);
     }
-    RCLASSEXT_SET_FIELDS(ext, rb_imemo_obj_fields_new_complex(st_init_numtable()));
+
     RCLASSEXT_CONST_TBL(ext) = RCLASSEXT_CONST_TBL(mod_ext);
     RCLASSEXT_CVC_TBL(ext) = RCLASSEXT_CVC_TBL(mod_ext);
 
@@ -318,11 +318,11 @@ rb_class_duplicate_classext(rb_classext_t *orig, VALUE klass, const rb_namespace
     // TODO: consider shapes for performance
     if (RCLASSEXT_FIELDS(orig)) {
         st_table *tbl = st_copy((st_table *)RCLASSEXT_FIELDS(orig));
-        RCLASSEXT_SET_FIELDS(ext, rb_imemo_obj_fields_new_complex(tbl));
+        RCLASSEXT_SET_FIELDS(ext, rb_imemo_obj_fields_new_complex(klass, tbl));
         rb_autoload_copy_table_for_namespace((st_table *)RCLASSEXT_FIELDS(ext), ns);
     }
     else {
-        RCLASSEXT_SET_FIELDS(ext, rb_imemo_obj_fields_new_complex(st_init_numtable()));
+        RCLASSEXT_SET_FIELDS(ext, rb_imemo_obj_fields_new_complex(klass, st_init_numtable()));
     }
 
     if (RCLASSEXT_SHARED_CONST_TBL(orig)) {

imemo.c

@@ -112,15 +112,17 @@ rb_imemo_tmpbuf_parser_heap(void *buf, rb_imemo_tmpbuf_t *old_heap, size_t cnt)
 static const size_t IMEMO_BUF_OVERHEAD = sizeof(VALUE); // flags
 
 VALUE
-rb_imemo_obj_fields_new(size_t capa)
+rb_imemo_obj_fields_new(VALUE klass, size_t capa)
 {
+    klass = rb_singleton_class(klass);
+
     if (rb_gc_size_allocatable_p(capa + IMEMO_BUF_OVERHEAD)) {
-        VALUE fields = rb_imemo_new(imemo_obj_fields, 0, capa + IMEMO_BUF_OVERHEAD);
+        VALUE fields = rb_imemo_new(imemo_obj_fields, klass, capa + IMEMO_BUF_OVERHEAD);
         RUBY_ASSERT(IMEMO_TYPE_P(fields, imemo_obj_fields));
         return fields;
     }
     else {
-        VALUE fields = rb_imemo_new(imemo_obj_fields, 0, sizeof(VALUE *));
+        VALUE fields = rb_imemo_new(imemo_obj_fields, klass, sizeof(VALUE *));
         FL_SET_RAW(fields, OBJ_FIELD_EXTERNAL);
         // TODO: malloc + tag.
         rb_bug("TODO: malloc + tag.");
@@ -129,9 +131,9 @@ rb_imemo_obj_fields_new(size_t capa)
 }
 
 VALUE
-rb_imemo_obj_fields_new_complex(st_table *tbl)
+rb_imemo_obj_fields_new_complex(VALUE klass, st_table *tbl)
 {
-    VALUE fields = rb_imemo_obj_fields_new(sizeof(tbl));
+    VALUE fields = rb_imemo_obj_fields_new(klass, sizeof(tbl));
     IMEMO_OBJ_FIELDS(fields)->as.complex.table = tbl;
     return fields;
 }

internal/class.h

@@ -569,7 +569,7 @@ RCLASS_SET_FIELDS_HASH(VALUE obj, st_table *tbl)
     RUBY_ASSERT(RB_TYPE_P(obj, RUBY_T_CLASS) || RB_TYPE_P(obj, RUBY_T_MODULE));
     RUBY_ASSERT(rb_shape_obj_too_complex_p(obj));
 
-    VALUE fields = rb_imemo_obj_fields_new(sizeof(tbl));
+    VALUE fields = rb_imemo_obj_fields_new(obj, sizeof(tbl));
     IMEMO_OBJ_FIELDS(fields)->as.complex.table = tbl;
     RCLASSEXT_SET_FIELDS(RCLASS_EXT_PRIME(obj), fields);
 }
@@ -580,7 +580,7 @@ RCLASS_WRITE_FIELDS_HASH(VALUE obj, st_table *tbl)
     RUBY_ASSERT(RB_TYPE_P(obj, RUBY_T_CLASS) || RB_TYPE_P(obj, RUBY_T_MODULE));
     RUBY_ASSERT(rb_shape_obj_too_complex_p(obj));
 
-    VALUE fields = rb_imemo_obj_fields_new(sizeof(tbl));
+    VALUE fields = rb_imemo_obj_fields_new(obj, sizeof(tbl));
     IMEMO_OBJ_FIELDS(fields)->as.complex.table = tbl;
     RCLASSEXT_SET_FIELDS(RCLASS_EXT_WRITABLE(obj), fields);
 }

internal/imemo.h

@@ -260,9 +260,7 @@ MEMO_V2_SET(struct MEMO *m, VALUE v)
 
 struct rb_obj_fields {
     VALUE flags;
-    // We're padding the struct so that capacity match T_OBJECT.
-    // Ideally we'd have a different root to make use of that extra slot.
-    VALUE _unused;
+    VALUE klass;
     union {
         struct {
             VALUE fields[1];
@@ -280,8 +278,8 @@ struct rb_obj_fields {
 #define OBJ_FIELD_COMPLEX  IMEMO_FL_USER1
 #define IMEMO_OBJ_FIELDS(fields) ((struct rb_obj_fields *)fields)
 
-VALUE rb_imemo_obj_fields_new(size_t capa);
-VALUE rb_imemo_obj_fields_new_complex(st_table *tbl);
+VALUE rb_imemo_obj_fields_new(VALUE klass, size_t capa);
+VALUE rb_imemo_obj_fields_new_complex(VALUE klass, st_table *tbl);
 
 static inline VALUE *
 rb_imemo_obj_fields_ptr(VALUE obj_fields)

shape.c

@@ -810,12 +810,17 @@ shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
         klass = rb_singleton_class(obj);
         allow_new_root = true;
         break;
+      case T_IMEMO:
+        RUBY_ASSERT(IMEMO_TYPE_P(obj, imemo_obj_fields));
+        allow_new_root = true;
+        klass = Qfalse;
+        break;
       default:
         klass = rb_obj_class(obj);
         break;
     }
 
-    bool allow_new_shape = RCLASS_VARIATION_COUNT(klass) < SHAPE_MAX_VARIATIONS;
+    bool allow_new_shape = !klass || RCLASS_VARIATION_COUNT(klass) < SHAPE_MAX_VARIATIONS;
     bool variation_created = false;
     allow_new_root &= allow_new_shape;
 
@@ -833,7 +838,7 @@ shape_get_next(rb_shape_t *shape, VALUE obj, ID id, bool emit_warnings)
     rb_shape_t *new_shape = get_next_shape_internal(shape, id, SHAPE_IVAR, &variation_created, allow_new_shape);
 
     // Check if we should update max_iv_count on the object's class
-    if (obj != klass && new_shape->next_field_index > RCLASS_MAX_IV_COUNT(klass)) {
+    if (klass && new_shape->next_field_index > RCLASS_MAX_IV_COUNT(klass)) {
         RCLASS_SET_MAX_IV_COUNT(klass, new_shape->next_field_index);
     }
 

shape.h

@@ -94,15 +94,15 @@ static inline shape_id_t
 get_shape_id_from_flags(VALUE obj)
 {
     RUBY_ASSERT(!RB_SPECIAL_CONST_P(obj));
-    RUBY_ASSERT(!RB_TYPE_P(obj, T_IMEMO));
+    RUBY_ASSERT(!RB_TYPE_P(obj, T_IMEMO) || IMEMO_TYPE_P(obj, imemo_obj_fields));
     return (shape_id_t)((RBASIC(obj)->flags) >> SHAPE_FLAG_SHIFT);
 }
 
 static inline void
 set_shape_id_in_flags(VALUE obj, shape_id_t shape_id)
 {
     RUBY_ASSERT(!RB_SPECIAL_CONST_P(obj));
-    RUBY_ASSERT(!RB_TYPE_P(obj, T_IMEMO));
+    RUBY_ASSERT(!RB_TYPE_P(obj, T_IMEMO) || IMEMO_TYPE_P(obj, imemo_obj_fields));
     // Ractors are occupying the upper 32 bits of flags, but only in debug mode
     // Object shapes are occupying top bits
     RBASIC(obj)->flags &= SHAPE_FLAG_MASK;

variable.c

@@ -1415,17 +1415,17 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
       case T_MODULE:
         {
             bool found = false;
-            VALUE val;
+            VALUE val = undef;
 
-            RB_VM_LOCK_ENTER();
-            {
+            VALUE fields_obj = RCLASS_PRIME_FIELDS_VALUE(obj); // TODO: figure namespaces
+            if (fields_obj) {
 #if !SHAPE_IN_BASIC_FLAGS
-                shape_id = RCLASS_SHAPE_ID(obj);
+                shape_id = get_shape_id_from_flags(fields_obj);
 #endif
 
-                if (rb_shape_obj_too_complex_p(obj)) {
-                    st_table * iv_table = RCLASS_FIELDS_HASH(obj);
-                    if (rb_st_lookup(iv_table, (st_data_t)id, (st_data_t *)&val)) {
+                if (rb_shape_obj_too_complex_p(fields_obj)) {
+                    st_table *fields_tbl = IMEMO_OBJ_FIELDS(fields_obj)->as.complex.table;
+                    if (rb_st_lookup(fields_tbl, (st_data_t)id, (st_data_t *)&val)) {
                         found = true;
                     }
                     else {
@@ -1438,7 +1438,7 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
                     found = rb_shape_get_iv_index(shape, id, &index);
 
                     if (found) {
-                        ivar_list = RCLASS_PRIME_FIELDS(obj);
+                        ivar_list = rb_imemo_obj_fields_ptr(fields_obj);
                         RUBY_ASSERT(ivar_list);
 
                         val = ivar_list[index];
@@ -1448,7 +1448,6 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
                     }
                 }
             }
-            RB_VM_LOCK_LEAVE();
 
             if (found &&
                     rb_is_instance_id(id) &&
@@ -4738,7 +4737,7 @@ class_ivar_set_shape_fields(VALUE obj, void *_data)
 static void
 class_ivar_set_shape_resize_fields(VALUE obj, attr_index_t old_capa, attr_index_t new_capa, void *_data)
 {
-    VALUE new_fields = rb_imemo_obj_fields_new(new_capa);
+    VALUE new_fields = rb_imemo_obj_fields_new(obj, new_capa);
     if (old_capa) {
         MEMCPY(rb_imemo_obj_fields_ptr(new_fields), RCLASS_PRIME_FIELDS(obj), VALUE, old_capa);
     }
@@ -4765,27 +4764,85 @@ class_ivar_set_too_complex_table(VALUE obj, void *_data)
     return RCLASS_WRITABLE_FIELDS_HASH(obj);
 }
 
+int
+class_atomic_ivar_set(VALUE obj, ID id, VALUE val)
+{
+    bool existing = true;
+
+    VALUE original_fields_obj = RCLASS_PRIME_FIELDS_VALUE(obj); // TODO: figure out namespaces
+    VALUE fields_obj = original_fields_obj;
+
+    if (!fields_obj) {
+        fields_obj = rb_imemo_obj_fields_new(obj, 0);
+    }
+
+    rb_shape_t *current_shape = rb_obj_shape(fields_obj);
+    if (UNLIKELY(rb_shape_too_complex_p(current_shape))) {
+        goto too_complex;
+    }
+
+    attr_index_t index;
+    shape_id_t next_shape_id = 0;
+    if (!rb_shape_get_iv_index(current_shape, id, &index)) {
+        existing = false;
+
+        index = current_shape->next_field_index;
+        if (index >= SHAPE_MAX_FIELDS) {
+            rb_raise(rb_eArgError, "too many instance variables");
+        }
+
+        next_shape_id = rb_shape_transition_add_ivar(obj, id);
+        rb_shape_t *next_shape = RSHAPE(next_shape_id);
+        if (UNLIKELY(rb_shape_too_complex_p(next_shape))) {
+            rb_bug("TODO: transition_too_complex");
+            goto too_complex;
+        }
+        else if (UNLIKELY(next_shape->capacity != current_shape->capacity)) {
+            RUBY_ASSERT(next_shape->capacity > current_shape->capacity);
+            fields_obj = rb_imemo_obj_fields_new(obj, next_shape->capacity);
+        }
+
+        RUBY_ASSERT(next_shape->type == SHAPE_IVAR);
+        RUBY_ASSERT(index == (next_shape->next_field_index - 1));
+    }
+
+    VALUE *fields = rb_imemo_obj_fields_ptr(fields_obj);
+    RB_OBJ_WRITE(obj, &fields[index], val);
+    if (!existing) {
+        rb_shape_set_shape_id(fields_obj, next_shape_id);
+    }
+
+    if (fields_obj != original_fields_obj) {
+        RUBY_ATOMIC_VALUE_SET(RCLASS_PRIME_FIELDS_VALUE(obj), fields_obj);
+    }
+
+    return existing;
+
+too_complex:
+    {
+        rb_bug("TODO: handle too_complex");
+    }
+    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_LOCK_ENTER();
-    {
-        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;
+    if (rb_multi_ractor_p()) {
+        return class_atomic_ivar_set(obj, id, val);
+    }
+    else {
+        return 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;
     }
-    RB_VM_LOCK_LEAVE();
-
-    return existing;
 }
 
 static void