Modify class variable definition in singleton class; fix #3539

2 files changed, 27 insertions, 4 deletions

diff --git a/src/variable.c b/src/variable.c
index 2e823453f..f5d261dec 100644
--- a/src/variable.c
+++ b/src/variable.c
@@ -819,12 +819,22 @@ mrb_mod_cv_set(mrb_state *mrb, struct RClass *c, mrb_sym sym, mrb_value v)
     c = c->super;
   }
 
-  if (!cls->iv) {
-    cls->iv = iv_new(mrb);
+  if (cls && cls->tt == MRB_TT_SCLASS) {
+    mrb_value klass;
+    klass = mrb_obj_iv_get(mrb, (struct RObject*)cls,
+                          mrb_intern_lit(mrb, "__attached__"));
+
+    c = mrb_class_ptr(klass);
+  }else{
+    c = cls;
+  }
+
+  if (!c->iv) {
+    c->iv = iv_new(mrb);
   }
 
-  mrb_write_barrier(mrb, (struct RBasic*)cls);
-  iv_put(mrb, cls->iv, sym, v);
+  mrb_write_barrier(mrb, (struct RBasic*)c);
+  iv_put(mrb, c->iv, sym, v);
 }
 
 MRB_API void
diff --git a/test/t/class.rb b/test/t/class.rb
index 605b7ec40..54fadc0ca 100644
--- a/test/t/class.rb
+++ b/test/t/class.rb
@@ -384,6 +384,19 @@ assert('class variable and class << self style class method') do
   assert_equal("value", ClassVariableTest.class_variable)
 end
 
+assert('class variable definition in singleton_class') do
+  class ClassVariableDefinitionInSingletonTest
+    class << self
+      @@class_variable = "value"
+    end
+    def class_variable
+      @@class_variable
+    end
+  end
+
+  assert_equal("value", ClassVariableDefinitionInSingletonTest.new.class_variable)
+end
+
 assert('class variable in module and class << self style class method') do
   module ClassVariableInModuleTest
     @@class_variable = "value"