Added shuffle and shuffle! to the Array class in the Random gem.

mruby · Nov 1, 2013 · d8cc37c · d8cc37c
1 parent 9282566
commit d8cc37c
Show file tree

Hide file tree

Showing 2 changed files with 66 additions and 0 deletions.
diff --git a/mrbgems/mruby-random/src/random.c b/mrbgems/mruby-random/src/random.c
@@ -7,6 +7,7 @@
 #include "mruby.h"
 #include "mruby/variable.h"
 #include "mruby/data.h"
+#include "mruby/array.h"
 #include "mt19937ar.h"
 
 #include <time.h>
@@ -207,9 +208,57 @@ static mrb_value mrb_random_srand(mrb_state *mrb, mrb_value self)
   return old_seed;
 }
 
+/*
+ *  call-seq:
+ *     ary.shuffle!   ->   ary
+ *
+ *  Shuffles elements in self in place.
+ */
+
+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());
+  }
+
+  if (RARRAY_LEN(ary) > 1) {
+    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_value t = RARRAY_PTR(ary)[i];
+      RARRAY_PTR(ary)[i] = RARRAY_PTR(ary)[j];
+      RARRAY_PTR(ary)[j] = t;
+    }    
+  }
+
+  return ary;
+}
+
+/*
+ *  call-seq:
+ *     ary.shuffle   ->   new_ary
+ *
+ *  Returns a new array with elements of self shuffled.
+ */
+
+static mrb_value
+mrb_ary_shuffle(mrb_state *mrb, mrb_value ary)
+{
+  mrb_value new_ary = mrb_ary_new_from_values(mrb, RARRAY_LEN(ary), RARRAY_PTR(ary));
+  mrb_ary_shuffle_bang(mrb, new_ary);
+
+  return new_ary;
+}
+
 void mrb_mruby_random_gem_init(mrb_state *mrb)
 {
   struct RClass *random;
+  struct RClass *array = mrb->array_class;
 
   mrb_define_method(mrb, mrb->kernel_module, "rand", mrb_random_g_rand, MRB_ARGS_OPT(1));
   mrb_define_method(mrb, mrb->kernel_module, "srand", mrb_random_g_srand, MRB_ARGS_OPT(1));
@@ -221,6 +270,9 @@ void mrb_mruby_random_gem_init(mrb_state *mrb)
   mrb_define_method(mrb, random, "initialize", mrb_random_init, MRB_ARGS_OPT(1));
   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());
 }
 
 void mrb_mruby_random_gem_final(mrb_state *mrb)

diff --git a/mrbgems/mruby-random/test/random.rb b/mrbgems/mruby-random/test/random.rb
@@ -30,3 +30,17 @@
 assert("float") do
   rand.class == Float
 end
+
+assert("Array#shuffle") do 
+  ary = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+  shuffled = ary.shuffle
+
+  ary == [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] and shuffled != ary and 10.times { |x| ary.include? x }
+end
+
+assert('Array#shuffle!') do
+  ary = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+  ary.shuffle!
+
+  ary != [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] and 10.times { |x| ary.include? x }
+end