Merge pull request #662 from skandhas/pr-add-Module-class_variable_set

Add Module#class_variable_set for mruby

4 files changed, 80 insertions, 0 deletions

diff --git a/include/mruby/variable.h b/include/mruby/variable.h
index 8754e8a90..9f4a62317 100644
--- a/include/mruby/variable.h
+++ b/include/mruby/variable.h
@@ -57,6 +57,8 @@ mrb_sym mrb_class_sym(mrb_state *mrb, struct RClass *c, struct RClass *outer);
 mrb_value mrb_mod_class_variables(mrb_state*, mrb_value);
 mrb_value mrb_mod_cv_get(mrb_state *mrb, struct RClass * c, mrb_sym sym);
 mrb_value mrb_cv_get(mrb_state *mrb, mrb_value mod, mrb_sym sym);
+void mrb_mod_cv_set(mrb_state *mrb, struct RClass * c, mrb_sym sym, mrb_value v);
+void mrb_cv_set(mrb_state *mrb, mrb_value mod, mrb_sym sym, mrb_value v);
 
 /* GC functions */
 void mrb_gc_mark_gv(mrb_state*);
diff --git a/src/class.c b/src/class.c
index 19a697bc8..8d14319f2 100644
--- a/src/class.c
+++ b/src/class.c
@@ -1483,6 +1483,38 @@ mrb_mod_cvar_get(mrb_state *mrb, mrb_value mod)
   return mrb_cv_get(mrb, mod, id);
 }
 
+/* 15.2.2.4.18 */
+/*
+ *  call-seq:
+ *     obj.class_variable_set(symbol, obj)    -> obj
+ *
+ *  Sets the class variable names by <i>symbol</i> to
+ *  <i>object</i>.
+ *
+ *     class Fred
+ *       @@foo = 99
+ *       def foo
+ *         @@foo
+ *       end
+ *     end
+ *     Fred.class_variable_set(:@@foo, 101)     #=> 101
+ *     Fred.new.foo                             #=> 101
+ */
+
+static mrb_value
+mrb_mod_cvar_set(mrb_state *mrb, mrb_value mod)
+{
+  mrb_value sym, value;
+  mrb_sym id;
+  mrb_get_args(mrb, "oo", &sym, &value);
+
+  id = mrb_sym_value(mrb,sym);
+
+  check_cv_name(mrb, id);
+  mrb_cv_set(mrb, mod, id, value);
+  return value;
+}
+
 static void
 check_const_name(mrb_state *mrb, mrb_sym id)
 {
@@ -1585,6 +1617,7 @@ mrb_init_class(mrb_state *mrb)
   mrb_define_method(mrb, cls, "new", mrb_instance_new, ARGS_ANY());                  /* 15.2.3.3.3 */
   mrb_define_method(mrb, cls, "inherited", mrb_bob_init, ARGS_REQ(1));
   mrb_define_method(mrb, mod, "class_variable_get", mrb_mod_cvar_get, ARGS_REQ(1));  /* 15.2.2.4.17 */
+  mrb_define_method(mrb, mod, "class_variable_set", mrb_mod_cvar_set, ARGS_REQ(2));  /* 15.2.2.4.18 */
   mrb_define_method(mrb, mod, "extend_object", mrb_mod_extend_object, ARGS_REQ(1));  /* 15.2.2.4.25 */
   mrb_define_method(mrb, mod, "extended", mrb_bob_init, ARGS_REQ(1));                /* 15.2.2.4.26 */
   mrb_define_method(mrb, mod, "include", mrb_mod_include, ARGS_ANY());               /* 15.2.2.4.27 */
diff --git a/src/variable.c b/src/variable.c
index 147373bd4..7743c6309 100644
--- a/src/variable.c
+++ b/src/variable.c
@@ -695,6 +695,36 @@ mrb_cv_get(mrb_state *mrb, mrb_value mod, mrb_sym sym)
   return mrb_mod_cv_get(mrb, mrb_class_ptr(mod), sym);
 }
 
+void
+mrb_mod_cv_set(mrb_state *mrb, struct RClass * c, mrb_sym sym, mrb_value v)
+{
+  struct RClass * cls = c;
+
+  while (c) {
+    if (c->iv) {
+      iv_tbl *t = c->iv;
+
+      if (iv_get(mrb, t, sym, NULL)) {
+        iv_put(mrb, t, sym, v);
+        return;
+      }
+    }
+    c = c->super;
+  }
+
+  if (!cls->iv) {
+    cls->iv = iv_new(mrb);
+  }
+
+  iv_put(mrb, cls->iv, sym, v);
+}
+
+void
+mrb_cv_set(mrb_state *mrb, mrb_value mod, mrb_sym sym, mrb_value v)
+{
+  mrb_mod_cv_set(mrb, mrb_class_ptr(mod), sym, v);
+}
+
 mrb_value
 mrb_vm_cv_get(mrb_state *mrb, mrb_sym sym)
 {
diff --git a/test/t/module.rb b/test/t/module.rb
index fbc93eb08..54d4f3015 100644
--- a/test/t/module.rb
+++ b/test/t/module.rb
@@ -54,6 +54,21 @@ assert('Module#class_variable_get', '15.2.2.4.17') do
   Test4ClassVariableGet.class_variable_get(:@@cv) == 99
 end
 
+assert('Module#class_variable_set', '15.2.2.4.18') do
+  class Test4ClassVariableSet
+    @@foo = 100
+    def foo
+      @@foo
+    end
+  end
+
+  Test4ClassVariableSet.class_variable_set(:@@cv, 99)
+  Test4ClassVariableSet.class_variable_set(:@@foo, 101)
+
+  Test4ClassVariableSet.class_variables.include? :@@cv and
+  Test4ClassVariableSet.class_variable_get(:@@cv) == 99 and
+  Test4ClassVariableSet.new.foo == 101
+end
 
 assert('Module#class_variables', '15.2.2.4.19') do
   class Test4ClassVariables1