Allow classes and modules to become too complex

This makes the behavior of classes and modules when there are too many instance variables match the behavior of objects with too many instance variables.

5 files changed, 117 insertions, 30 deletions

diff --git a/gc.c b/gc.c
index b8c8ae099d..bbcdc4cedc 100644
--- a/gc.c
+++ b/gc.c
@@ -3558,7 +3558,11 @@ obj_free(rb_objspace_t *objspace, VALUE obj)
       case T_CLASS:
         rb_id_table_free(RCLASS_M_TBL(obj));
         cc_table_free(objspace, obj, FALSE);
-        if (RCLASS_IVPTR(obj)) {
+        if (rb_shape_obj_too_complex(obj)) {
+            RB_DEBUG_COUNTER_INC(obj_obj_too_complex);
+            rb_id_table_free(RCLASS_TABLE_IVPTR(obj));
+        }
+        else if (RCLASS_IVPTR(obj)) {
             xfree(RCLASS_IVPTR(obj));
         }
         if (RCLASS_CONST_TBL(obj)) {
diff --git a/internal/class.h b/internal/class.h
index 9e47a339c1..e12fa5c355 100644
--- a/internal/class.h
+++ b/internal/class.h
@@ -96,6 +96,7 @@ STATIC_ASSERT(sizeof_rb_classext_t, sizeof(struct RClass) + sizeof(rb_classext_t
 #define RCLASS_CONST_TBL(c) (RCLASS_EXT(c)->const_tbl)
 #define RCLASS_M_TBL(c) (RCLASS(c)->m_tbl)
 #define RCLASS_IVPTR(c) (RCLASS_EXT(c)->iv_ptr)
+#define RCLASS_TABLE_IVPTR(c) (struct rb_id_table *)(RCLASS_EXT(c)->iv_ptr)
 #define RCLASS_CALLABLE_M_TBL(c) (RCLASS_EXT(c)->callable_m_tbl)
 #define RCLASS_CC_TBL(c) (RCLASS_EXT(c)->cc_tbl)
 #define RCLASS_CVC_TBL(c) (RCLASS_EXT(c)->cvc_tbl)
diff --git a/shape.c b/shape.c
index 01ec0f0604..e6ba20a644 100644
--- a/shape.c
+++ b/shape.c
@@ -239,7 +239,7 @@ remove_shape_recursive(VALUE obj, ID id, rb_shape_t * shape, VALUE * removed)
             if (new_parent) {
                 bool dont_care;
                 enum ruby_value_type type = BUILTIN_TYPE(obj);
-                bool new_shape_necessary = type != T_OBJECT;
+                bool new_shape_necessary = type != T_OBJECT && type != T_CLASS && type != T_MODULE;
                 rb_shape_t * new_child = get_next_shape_internal(new_parent, shape->edge_name, shape->type, &dont_care, true, new_shape_necessary);
                 new_child->capacity = shape->capacity;
                 if (new_child->type == SHAPE_IVAR) {
@@ -316,12 +316,10 @@ rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id)
     }
 
     bool variation_created = false;
-    // For non T_OBJECTS, force a new shape
-    bool new_shape_necessary = BUILTIN_TYPE(obj) != T_OBJECT;
+    bool new_shape_necessary = BUILTIN_TYPE(obj) != T_OBJECT && BUILTIN_TYPE(obj) != T_CLASS && BUILTIN_TYPE(obj) != T_MODULE;
     rb_shape_t * new_shape = get_next_shape_internal(shape, id, SHAPE_IVAR, &variation_created, allow_new_shape, new_shape_necessary);
 
     if (!new_shape) {
-        RUBY_ASSERT(BUILTIN_TYPE(obj) == T_OBJECT);
         new_shape = rb_shape_get_shape_by_id(OBJ_TOO_COMPLEX_SHAPE_ID);
     }
 
@@ -336,6 +334,15 @@ rb_shape_get_next(rb_shape_t* shape, VALUE obj, ID id)
             RCLASS_EXT(klass)->variation_count++;
         }
     }
+    else if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
+        if (new_shape->next_iv_index > RCLASS_EXT(obj)->max_iv_count) {
+            RCLASS_EXT(obj)->max_iv_count = new_shape->next_iv_index;
+        }
+
+        if (variation_created) {
+            RCLASS_EXT(obj)->variation_count++;
+        }
+    }
 
     return new_shape;
 }
@@ -523,7 +530,6 @@ rb_shape_obj_too_complex(VALUE obj)
 void
 rb_shape_set_too_complex(VALUE obj)
 {
-    RUBY_ASSERT(BUILTIN_TYPE(obj) == T_OBJECT);
     RUBY_ASSERT(!rb_shape_obj_too_complex(obj));
     rb_shape_set_shape_id(obj, OBJ_TOO_COMPLEX_SHAPE_ID);
 }
diff --git a/test/ruby/test_shapes.rb b/test/ruby/test_shapes.rb
index b1a2ba2f1b..e34256415b 100644
--- a/test/ruby/test_shapes.rb
+++ b/test/ruby/test_shapes.rb
@@ -105,7 +105,27 @@ class TestShapes < Test::Unit::TestCase
       obj.instance_variable_set(:"@a#{_1}", 1)
     end
 
-    assert_false RubyVM::Shape.of(obj).too_complex?
+    assert_predicate RubyVM::Shape.of(obj), :too_complex?
+  end
+
+  def test_too_many_ivs_on_module
+    obj = Module.new
+
+    (RubyVM::Shape::SHAPE_MAX_NUM_IVS + 1).times do
+      obj.instance_variable_set(:"@a#{_1}", 1)
+    end
+
+    assert_predicate RubyVM::Shape.of(obj), :too_complex?
+  end
+
+  def test_too_many_ivs_on_builtin
+    obj = "string"
+
+    (RubyVM::Shape::SHAPE_MAX_NUM_IVS + 1).times do
+      obj.instance_variable_set(:"@a#{_1}", 1)
+    end
+
+    refute RubyVM::Shape.of(obj).too_complex?
   end
 
   def test_removing_when_too_many_ivs_on_class
diff --git a/variable.c b/variable.c
index c0b4625e2e..64279d881b 100644
--- a/variable.c
+++ b/variable.c
@@ -1136,6 +1136,16 @@ rb_ivar_lookup(VALUE obj, ID id, VALUE undef)
                 shape_id = RCLASS_SHAPE_ID(obj);
 #endif
 
+                if (rb_shape_obj_too_complex(obj)) {
+                    struct rb_id_table * iv_table = RCLASS_TABLE_IVPTR(obj);
+                    if (rb_id_table_lookup(iv_table, id, &val)) {
+                        return val;
+                    }
+                    else {
+                        return undef;
+                    }
+                }
+
                 attr_index_t index = 0;
                 shape = rb_shape_get_shape_by_id(shape_id);
                 found = rb_shape_get_iv_index(shape, id, &index);
@@ -1268,6 +1278,8 @@ generic_ivar_set(VALUE obj, ID id, VALUE val)
     attr_index_t index;
     // The returned shape will have `id` in its iv_table
     rb_shape_t *shape = rb_shape_get_shape(obj);
+
+    RUBY_ASSERT(!rb_shape_obj_too_complex(obj));
     bool found = rb_shape_get_iv_index(shape, id, &index);
     if (!found) {
         index = shape->next_iv_index;
@@ -1655,6 +1667,7 @@ iterate_over_shapes_with_callback(rb_shape_t *shape, rb_ivar_foreach_callback_fu
             break;
           case T_CLASS:
           case T_MODULE:
+            RUBY_ASSERT(!rb_shape_obj_too_complex(itr_data->obj));
             iv_list = RCLASS_IVPTR(itr_data->obj);
             break;
           default:
@@ -1724,7 +1737,13 @@ class_ivar_each(VALUE obj, rb_ivar_foreach_callback_func *func, st_data_t arg)
     struct iv_itr_data itr_data;
     itr_data.obj = obj;
     itr_data.arg = arg;
-    iterate_over_shapes_with_callback(shape, func, &itr_data);
+    itr_data.func = func;
+    if (rb_shape_obj_too_complex(obj)) {
+        rb_id_table_foreach(RCLASS_TABLE_IVPTR(obj), each_hash_iv, &itr_data);
+    }
+    else {
+        iterate_over_shapes_with_callback(shape, func, &itr_data);
+    }
 }
 
 void
@@ -1984,7 +2003,14 @@ rb_obj_remove_instance_variable(VALUE obj, VALUE name)
       case T_CLASS:
       case T_MODULE:
         IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(id);
-        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
+        if (rb_shape_obj_too_complex(obj)) {
+            if (rb_id_table_lookup(RCLASS_TABLE_IVPTR(obj), id, &val)) {
+                rb_id_table_delete(RCLASS_TABLE_IVPTR(obj), id);
+            }
+        } else {
+            rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
+        }
+
         break;
       case T_OBJECT: {
         if (rb_shape_obj_too_complex(obj)) {
@@ -3961,35 +3987,65 @@ rb_class_ivar_set(VALUE obj, ID key, VALUE value)
 
     RB_VM_LOCK_ENTER();
     {
-        rb_shape_t * shape = rb_shape_get_shape(obj);
-        attr_index_t idx;
-        found = rb_shape_get_iv_index(shape, key, &idx);
-
-        if (found) {
-            // Changing an existing instance variable
-            RUBY_ASSERT(RCLASS_IVPTR(obj));
-
-            RCLASS_IVPTR(obj)[idx] = value;
+        if (rb_shape_obj_too_complex(obj)) {
+            struct rb_id_table * iv_table = RCLASS_TABLE_IVPTR(obj);
+            rb_id_table_insert(iv_table, key, value);
             RB_OBJ_WRITTEN(obj, Qundef, value);
+            found = 0;
         }
         else {
-            // Creating and setting a new instance variable
+            rb_shape_t * shape = rb_shape_get_shape(obj);
+            attr_index_t idx;
+            found = rb_shape_get_iv_index(shape, key, &idx);
 
-            // Move to a shape which fits the new ivar
-            idx = shape->next_iv_index;
-            shape = rb_shape_get_next(shape, obj, key);
+            if (found) {
+                // Changing an existing instance variable
+                RUBY_ASSERT(RCLASS_IVPTR(obj));
 
-            // We always allocate a power of two sized IV array. This way we
-            // only need to realloc when we expand into a new power of two size
-            if ((idx & (idx - 1)) == 0) {
-                size_t newsize = idx ? idx * 2 : 1;
-                REALLOC_N(RCLASS_IVPTR(obj), VALUE, newsize);
+                RCLASS_IVPTR(obj)[idx] = value;
+                RB_OBJ_WRITTEN(obj, Qundef, value);
             }
+            else {
+                // Creating and setting a new instance variable
+
+                // Move to a shape which fits the new ivar
+                idx = shape->next_iv_index;
+                shape = rb_shape_get_next(shape, obj, key);
+
+                // stop using shapes if we are now too complex
+                if (shape->type == SHAPE_OBJ_TOO_COMPLEX) {
+                    struct rb_id_table * table = rb_id_table_create(shape->next_iv_index);
+
+                    // Evacuate all previous values from shape into id_table
+                    rb_ivar_foreach(obj, rb_obj_evacuate_ivs_to_hash_table, (st_data_t)table);
+
+                    // insert the new value
+                    rb_id_table_insert(table, key, value);
+                    RB_OBJ_WRITTEN(obj, Qundef, value);
 
-            RUBY_ASSERT(RCLASS_IVPTR(obj));
+                    rb_shape_set_too_complex(obj);
+                    RUBY_ASSERT(rb_shape_obj_too_complex(obj));
 
-            RB_OBJ_WRITE(obj, &RCLASS_IVPTR(obj)[idx], value);
-            rb_shape_set_shape(obj, shape);
+                    if (RCLASS_IVPTR(obj)) {
+                        xfree(RCLASS_IVPTR(obj));
+                    }
+
+                    RCLASS_EXT(obj)->iv_ptr = (VALUE *)table;
+                }
+                else {
+                    // We always allocate a power of two sized IV array. This way we
+                    // only need to realloc when we expand into a new power of two size
+                    if ((idx & (idx - 1)) == 0) {
+                        size_t newsize = idx ? idx * 2 : 1;
+                        REALLOC_N(RCLASS_IVPTR(obj), VALUE, newsize);
+                    }
+
+                    RUBY_ASSERT(RCLASS_IVPTR(obj));
+
+                    RB_OBJ_WRITE(obj, &RCLASS_IVPTR(obj)[idx], value);
+                    rb_shape_set_shape(obj, shape);
+                }
+            }
         }
     }
     RB_VM_LOCK_LEAVE();