Added support for Random as an argument to shuffle and shuffle!. Refactored...

Added support for Random as an argument to shuffle and shuffle!. Refactored random gem to use DATA instance type and hold mt_state inside the DATA_PTR instead of in an instance variable.

mrbgems/mruby-random/src/random.c

@@ -6,6 +6,7 @@
 
 #include "mruby.h"
 #include "mruby/variable.h"
+#include "mruby/class.h"
 #include "mruby/data.h"
 #include "mruby/array.h"
 #include "mt19937ar.h"
@@ -19,23 +20,11 @@
 #define INSTANCE_RAND_SEED_KEY_CSTR_LEN 16
 
 #define MT_STATE_KEY          "$mrb_i_mt_state"
-#define MT_STATE_KEY_CSTR_LEN 15
 
 static const struct mrb_data_type mt_state_type = {
   MT_STATE_KEY, mrb_free,
 };
 
-static mt_state *mrb_mt_get_context(mrb_state *mrb,  mrb_value self)
-{
-  mt_state *t;
-  mrb_value context;
-
-  context = mrb_iv_get(mrb, self, mrb_intern2(mrb, MT_STATE_KEY, MT_STATE_KEY_CSTR_LEN));
-  t = DATA_GET_PTR(mrb, context, &mt_state_type, mt_state);
-
-  return t;
-}
-
 static void mt_g_srand(unsigned long seed)
 {
   init_genrand(seed);
@@ -51,7 +40,8 @@ static double mt_g_rand_real()
   return genrand_real1();
 }
 
-static mrb_value mrb_random_mt_g_srand(mrb_state *mrb, mrb_value seed)
+static mrb_value 
+mrb_random_mt_g_srand(mrb_state *mrb, mrb_value seed)
 {
   if (mrb_nil_p(seed)) {
     seed = mrb_fixnum_value(time(NULL) + mt_g_rand());
@@ -65,7 +55,8 @@ static mrb_value mrb_random_mt_g_srand(mrb_state *mrb, mrb_value seed)
   return seed;
 }
 
-static mrb_value mrb_random_mt_g_rand(mrb_state *mrb, mrb_value max)
+static mrb_value 
+mrb_random_mt_g_rand(mrb_state *mrb, mrb_value max)
 {
   mrb_value value;
 
@@ -78,22 +69,26 @@ static mrb_value mrb_random_mt_g_rand(mrb_state *mrb, mrb_value max)
   return value;
 }
  
-static void mt_srand(mt_state *t, unsigned long seed)
+static void 
+mt_srand(mt_state *t, unsigned long seed)
 {
   mrb_random_init_genrand(t, seed);
 }  
 
-static unsigned long mt_rand(mt_state *t)
+static unsigned long 
+mt_rand(mt_state *t)
 {
   return mrb_random_genrand_int32(t);
 }  
 
-static double mt_rand_real(mt_state *t)
+static double 
+mt_rand_real(mt_state *t)
 {
   return mrb_random_genrand_real1(t);
 }  
 
-static mrb_value mrb_random_mt_srand(mrb_state *mrb, mt_state *t, mrb_value seed)
+static mrb_value 
+mrb_random_mt_srand(mrb_state *mrb, mt_state *t, mrb_value seed)
 { 
   if (mrb_nil_p(seed)) {
     seed = mrb_fixnum_value(time(NULL) + mt_rand(t));
@@ -107,7 +102,8 @@ static mrb_value mrb_random_mt_srand(mrb_state *mrb, mt_state *t, mrb_value seed
   return seed;
 }
 
-static mrb_value mrb_random_mt_rand(mrb_state *mrb, mt_state *t, mrb_value max)
+static mrb_value 
+mrb_random_mt_rand(mrb_state *mrb, mt_state *t, mrb_value max)
 { 
   mrb_value value;
 
@@ -120,7 +116,8 @@ static mrb_value mrb_random_mt_rand(mrb_state *mrb, mt_state *t, mrb_value max)
   return value;
 }
 
-static mrb_value get_opt(mrb_state* mrb)
+static mrb_value 
+get_opt(mrb_state* mrb)
 {
   mrb_value arg;
 
@@ -139,20 +136,29 @@ static mrb_value get_opt(mrb_state* mrb)
   return arg;
 }
 
-static mrb_value mrb_random_g_rand(mrb_state *mrb, mrb_value self)
+static void 
+mrb_random_g_rand_seed(mrb_state *mrb) 
 {
-  mrb_value max;
   mrb_value seed;
-
-  max = get_opt(mrb);
+  
   seed = mrb_gv_get(mrb, mrb_intern2(mrb, GLOBAL_RAND_SEED_KEY, GLOBAL_RAND_SEED_KEY_CSTR_LEN));
   if (mrb_nil_p(seed)) {
     mrb_random_mt_g_srand(mrb, mrb_nil_value());
   }
+}
+
+static mrb_value 
+mrb_random_g_rand(mrb_state *mrb, mrb_value self)
+{
+  mrb_value max;
+
+  max = get_opt(mrb);
+  mrb_random_g_rand_seed(mrb);
   return mrb_random_mt_g_rand(mrb, max);
 }
 
-static mrb_value mrb_random_g_srand(mrb_state *mrb, mrb_value self)
+static mrb_value 
+mrb_random_g_srand(mrb_state *mrb, mrb_value self)
 {
   mrb_value seed;
   mrb_value old_seed;
@@ -164,41 +170,62 @@ static mrb_value mrb_random_g_srand(mrb_state *mrb, mrb_value self)
   return old_seed;
 }
 
-static mrb_value mrb_random_init(mrb_state *mrb, mrb_value self)
+static mrb_value 
+mrb_random_init(mrb_state *mrb, mrb_value self)
 {
   mrb_value seed;
+  mt_state *t;
+  
+  DATA_TYPE(self) = &mt_state_type;
+  DATA_PTR(self) = NULL;
+  
+  /* avoid memory leaks */
+  t = (mt_state*)DATA_PTR(self);
+  if (t) {
+    mrb_free(mrb, t);
+  }  
 
-
-  mt_state *t = (mt_state *)mrb_malloc(mrb, sizeof(mt_state));
+  t = (mt_state *)mrb_malloc(mrb, sizeof(mt_state));
   t->mti = N + 1;
 
   seed = get_opt(mrb);
   seed = mrb_random_mt_srand(mrb, t, seed);
   mrb_iv_set(mrb, self, mrb_intern2(mrb, INSTANCE_RAND_SEED_KEY, INSTANCE_RAND_SEED_KEY_CSTR_LEN), seed);
-  mrb_iv_set(mrb, self, mrb_intern2(mrb, MT_STATE_KEY, MT_STATE_KEY_CSTR_LEN),
-    mrb_obj_value(Data_Wrap_Struct(mrb, mrb->object_class, &mt_state_type, (void*) t)));
+  
+  DATA_PTR(self) = t;
+  
   return self;
 }
 
-static mrb_value mrb_random_rand(mrb_state *mrb, mrb_value self)
+static void 
+mrb_random_rand_seed(mrb_state *mrb, mrb_value self)
 {
-  mrb_value max;
   mrb_value seed;
-  mt_state *t = mrb_mt_get_context(mrb, self);
-
-  max = get_opt(mrb);
+  mt_state *t = DATA_PTR(self);
+  
   seed = mrb_iv_get(mrb, self, mrb_intern2(mrb, INSTANCE_RAND_SEED_KEY, INSTANCE_RAND_SEED_KEY_CSTR_LEN));
   if (mrb_nil_p(seed)) {
     mrb_random_mt_srand(mrb, t, mrb_nil_value());
   }
+}
+
+static mrb_value 
+mrb_random_rand(mrb_state *mrb, mrb_value self)
+{
+  mrb_value max;
+  mt_state *t = DATA_PTR(self);
+
+  max = get_opt(mrb);
+  mrb_random_rand_seed(mrb, self);
   return mrb_random_mt_rand(mrb, t, max);
 }
 
-static mrb_value mrb_random_srand(mrb_state *mrb, mrb_value self)
+static mrb_value 
+mrb_random_srand(mrb_state *mrb, mrb_value self)
 {
   mrb_value seed;
   mrb_value old_seed;
-  mt_state *t = mrb_mt_get_context(mrb, self);
+  mt_state *t = DATA_PTR(self);
 
   seed = get_opt(mrb);
   seed = mrb_random_mt_srand(mrb, t, seed);
@@ -219,17 +246,29 @@ static mrb_value
 mrb_ary_shuffle_bang(mrb_state *mrb, mrb_value ary)
 {
   mrb_int i;
-  mrb_value seed;
-
-  seed = mrb_gv_get(mrb, mrb_intern2(mrb, GLOBAL_RAND_SEED_KEY, GLOBAL_RAND_SEED_KEY_CSTR_LEN));
-  if (mrb_nil_p(seed)) {
-    mrb_random_mt_g_srand(mrb, mrb_nil_value());
-  }
+  mrb_value random = mrb_nil_value();
   
   if (RARRAY_LEN(ary) > 1) {
+    mrb_get_args(mrb, "|o", &random);
+
+    if( mrb_nil_p(random) ) {
+      mrb_random_g_rand_seed(mrb);
+    } else {
+      mrb_data_check_type(mrb, random, &mt_state_type);
+      mrb_random_rand_seed(mrb, random);
+    }
+  
     mrb_ary_modify(mrb, mrb_ary_ptr(ary));
+    
     for (i = RARRAY_LEN(ary) - 1; i > 0; i--)  {
-      mrb_int j = mrb_fixnum(mrb_random_mt_g_rand(mrb, mrb_fixnum_value(RARRAY_LEN(ary))));
+      mrb_int j;
+      
+      if( mrb_nil_p(random) ) {
+        j = mrb_fixnum(mrb_random_mt_g_rand(mrb, mrb_fixnum_value(RARRAY_LEN(ary))));
+      } else {
+        j = mrb_fixnum(mrb_random_mt_rand(mrb, DATA_PTR(random), mrb_fixnum_value(RARRAY_LEN(ary))));
+      }
+      
       mrb_value t = RARRAY_PTR(ary)[i];
       RARRAY_PTR(ary)[i] = RARRAY_PTR(ary)[j];
       RARRAY_PTR(ary)[j] = t;
@@ -264,6 +303,7 @@ void mrb_mruby_random_gem_init(mrb_state *mrb)
   mrb_define_method(mrb, mrb->kernel_module, "srand", mrb_random_g_srand, MRB_ARGS_OPT(1));
 
   random = mrb_define_class(mrb, "Random", mrb->object_class);
+  MRB_SET_INSTANCE_TT(random, MRB_TT_DATA);
   mrb_define_class_method(mrb, random, "rand", mrb_random_g_rand, MRB_ARGS_OPT(1));
   mrb_define_class_method(mrb, random, "srand", mrb_random_g_srand, MRB_ARGS_OPT(1));
 
@@ -271,8 +311,8 @@ void mrb_mruby_random_gem_init(mrb_state *mrb)
   mrb_define_method(mrb, random, "rand", mrb_random_rand, MRB_ARGS_OPT(1));
   mrb_define_method(mrb, random, "srand", mrb_random_srand, MRB_ARGS_OPT(1));
   
-  mrb_define_method(mrb, array, "shuffle", mrb_ary_shuffle, MRB_ARGS_NONE());
-  mrb_define_method(mrb, array, "shuffle!", mrb_ary_shuffle_bang, MRB_ARGS_NONE());
+  mrb_define_method(mrb, array, "shuffle", mrb_ary_shuffle, MRB_ARGS_OPT(1));
+  mrb_define_method(mrb, array, "shuffle!", mrb_ary_shuffle_bang, MRB_ARGS_OPT(1));
 }
 
 void mrb_mruby_random_gem_final(mrb_state *mrb)

mrbgems/mruby-random/test/random.rb

@@ -43,4 +43,34 @@ assert('Array#shuffle!') do
   ary.shuffle!
   
   ary != [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] and 10.times { |x| ary.include? x }
+end
+
+assert("Array#shuffle(random)") do 
+  assert_raise(TypeError) do
+    # this will cause an exception due to the wrong argument
+    [1, 2].shuffle "Not a Random instance"
+  end
+  
+  # verify that the same seed causes the same results
+  ary1 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+  shuffle1 = ary1.shuffle Random.new 345
+  ary2 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+  shuffle2 = ary2.shuffle Random.new 345
+  
+  ary1 != shuffle1 and 10.times { |x| shuffle1.include? x } and shuffle1 == shuffle2
+end
+
+assert('Array#shuffle!(random)') do
+  assert_raise(TypeError) do
+    # this will cause an exception due to the wrong argument
+    [1, 2].shuffle! "Not a Random instance"
+  end
+    
+  # verify that the same seed causes the same results
+  ary1 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+  ary1.shuffle! Random.new 345
+  ary2 = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+  ary2.shuffle! Random.new 345
+  
+  ary1 != [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] and 10.times { |x| ary1.include? x } and ary1 == ary2
 end
\ No newline at end of file