Fix remove_class_variable for too complex classes

2 files changed, 56 insertions, 59 deletions

diff --git a/test/ruby/test_shapes.rb b/test/ruby/test_shapes.rb
index 5e77dc049e..11eb66e7a3 100644
--- a/test/ruby/test_shapes.rb
+++ b/test/ruby/test_shapes.rb
@@ -252,7 +252,7 @@ class TestShapes < Test::Unit::TestCase
     end;
   end
 
-  def test_run_out_of_shape_for_class
+  def test_run_out_of_shape_for_class_ivar
     assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
     begin;
       i = 0
@@ -275,6 +275,28 @@ class TestShapes < Test::Unit::TestCase
     end;
   end
 
+  def test_run_out_of_shape_for_class_cvar
+    assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
+    begin;
+      c = Class.new
+      i = 0
+      while RubyVM::Shape.shapes_available > 0
+        c.class_variable_set(:"@@i#{i}", 1)
+        i += 1
+      end
+
+      c.class_variable_set(:@@a, 1)
+      assert_equal(1, c.class_variable_get(:@@a))
+
+      c.class_eval { remove_class_variable(:@@a) }
+      assert_false(c.class_variable_defined?(:@@a))
+
+      assert_raise(NameError) do
+        c.class_eval { remove_class_variable(:@@a) }
+      end
+    end;
+  end
+
   def test_run_out_of_shape_generic_ivar_set
     assert_separately([], "#{<<~"begin;"}\n#{<<~'end;'}")
     begin;
diff --git a/variable.c b/variable.c
index 1710bf3355..ec3ee60b2d 100644
--- a/variable.c
+++ b/variable.c
@@ -1372,25 +1372,42 @@ rb_ivar_delete(VALUE obj, ID id, VALUE undef)
     rb_check_frozen(obj);
 
     VALUE val = undef;
-    rb_shape_t * shape = rb_shape_get_shape(obj);
+    rb_shape_t *shape = rb_shape_get_shape(obj);
 
-    switch (BUILTIN_TYPE(obj)) {
-      case T_CLASS:
-      case T_MODULE:
+    if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
         IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(id);
+    }
 
-        RB_VM_LOCK_ENTER();
-        {
-            rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
+    if (!rb_shape_transition_shape_remove_ivar(obj, id, shape, &val)) {
+        if (!rb_shape_obj_too_complex(obj)) {
+            rb_evict_ivars_to_hash(obj, shape);
         }
-        RB_VM_LOCK_LEAVE();
 
-        break;
-      default: {
-        rb_shape_transition_shape_remove_ivar(obj, id, shape, &val);
+        st_table *table = NULL;
+        switch (BUILTIN_TYPE(obj)) {
+          case T_CLASS:
+          case T_MODULE:
+            table = RCLASS_IV_HASH(obj);
+            break;
 
-        break;
-      }
+          case T_OBJECT:
+            table = ROBJECT_IV_HASH(obj);
+            break;
+
+          default: {
+            struct gen_ivtbl *ivtbl;
+            if (rb_gen_ivtbl_get(obj, 0, &ivtbl)) {
+                table = ivtbl->as.complex.table;
+            }
+            break;
+          }
+        }
+
+        if (table) {
+            if (!st_delete(table, (st_data_t *)&id, (st_data_t *)&val)) {
+                val = undef;
+            }
+        }
     }
 
     return val;
@@ -2181,60 +2198,18 @@ check_id_type(VALUE obj, VALUE *pname,
 VALUE
 rb_obj_remove_instance_variable(VALUE obj, VALUE name)
 {
-    VALUE val = Qundef;
     const ID id = id_for_var(obj, name, an, instance);
 
     // Frozen check comes here because it's expected that we raise a
     // NameError (from the id_for_var check) before we raise a FrozenError
     rb_check_frozen(obj);
 
-    if (!id) {
-        goto not_defined;
-    }
-
-    rb_shape_t * shape = rb_shape_get_shape(obj);
+    if (id) {
+        VALUE val = rb_ivar_delete(obj, id, Qundef);
 
-    if (BUILTIN_TYPE(obj) == T_CLASS || BUILTIN_TYPE(obj) == T_MODULE) {
-        IVAR_ACCESSOR_SHOULD_BE_MAIN_RACTOR(id);
+        if (val != Qundef) return val;
     }
 
-    if (!rb_shape_transition_shape_remove_ivar(obj, id, shape, &val)) {
-        if (!rb_shape_obj_too_complex(obj)) {
-            rb_evict_ivars_to_hash(obj, shape);
-        }
-
-        st_table *table = NULL;
-        switch (BUILTIN_TYPE(obj)) {
-          case T_CLASS:
-          case T_MODULE:
-            table = RCLASS_IV_HASH(obj);
-            break;
-
-          case T_OBJECT:
-            table = ROBJECT_IV_HASH(obj);
-            break;
-
-          default: {
-            struct gen_ivtbl *ivtbl;
-            if (rb_gen_ivtbl_get(obj, 0, &ivtbl)) {
-                table = ivtbl->as.complex.table;
-            }
-            break;
-          }
-        }
-
-        if (table) {
-            if (!st_delete(table, (st_data_t *)&id, (st_data_t *)&val)) {
-                val = Qundef;
-            }
-        }
-    }
-
-    if (val != Qundef) {
-        return val;
-    }
-
-  not_defined:
     rb_name_err_raise("instance variable %1$s not defined",
                       obj, name);
     UNREACHABLE_RETURN(Qnil);