1 /*
2 ** random.c - Random module
3 **
4 ** See Copyright Notice in mruby.h
5 */
6
7 #include <mruby.h>
8 #include <mruby/variable.h>
9 #include <mruby/class.h>
10 #include <mruby/data.h>
11 #include <mruby/array.h>
12 #include <mruby/istruct.h>
13 #if INT32_MAX <= INTPTR_MAX
14 # define XORSHIFT96
15 # define NSEEDS 3
16 #else
17 # define NSEEDS 4
18 #endif
19 #define LASTSEED (NSEEDS-1)
20
21 #include <time.h>
22
23 typedef struct rand_state {
24 uint32_t seed[NSEEDS];
25 } rand_state;
26
27 static void
rand_init(rand_state * t)28 rand_init(rand_state *t)
29 {
30 t->seed[0] = 123456789;
31 t->seed[1] = 362436069;
32 t->seed[2] = 521288629;
33 #ifndef XORSHIFT96
34 t->seed[3] = 88675123;
35 #endif
36 }
37
38 static uint32_t
rand_seed(rand_state * t,uint32_t seed)39 rand_seed(rand_state *t, uint32_t seed)
40 {
41 uint32_t old_seed = t->seed[LASTSEED];
42 rand_init(t);
43 t->seed[LASTSEED] = seed;
44 return old_seed;
45 }
46
47 #ifdef XORSHIFT96
48 static uint32_t
rand_uint32(rand_state * state)49 rand_uint32(rand_state *state)
50 {
51 uint32_t *seed = state->seed;
52 uint32_t x = seed[0];
53 uint32_t y = seed[1];
54 uint32_t z = seed[2];
55 uint32_t t;
56
57 t = (x ^ (x << 3)) ^ (y ^ (y >> 19)) ^ (z ^ (z << 6));
58 x = y; y = z; z = t;
59 seed[0] = x;
60 seed[1] = y;
61 seed[2] = z;
62
63 return z;
64 }
65 #else /* XORSHIFT96 */
66 static uint32_t
rand_uint32(rand_state * state)67 rand_uint32(rand_state *state)
68 {
69 uint32_t *seed = state->seed;
70 uint32_t x = seed[0];
71 uint32_t y = seed[1];
72 uint32_t z = seed[2];
73 uint32_t w = seed[3];
74 uint32_t t;
75
76 t = x ^ (x << 11);
77 x = y; y = z; z = w;
78 w = (w ^ (w >> 19)) ^ (t ^ (t >> 8));
79 seed[0] = x;
80 seed[1] = y;
81 seed[2] = z;
82 seed[3] = w;
83
84 return w;
85 }
86 #endif /* XORSHIFT96 */
87
88 #ifndef MRB_WITHOUT_FLOAT
89 static double
rand_real(rand_state * t)90 rand_real(rand_state *t)
91 {
92 uint32_t x = rand_uint32(t);
93 return x*(1.0/4294967295.0);
94 }
95 #endif
96
97 static mrb_value
random_rand(mrb_state * mrb,rand_state * t,mrb_value max)98 random_rand(mrb_state *mrb, rand_state *t, mrb_value max)
99 {
100 mrb_value value;
101
102 if (mrb_fixnum(max) == 0) {
103 #ifndef MRB_WITHOUT_FLOAT
104 value = mrb_float_value(mrb, rand_real(t));
105 #else
106 mrb_raise(mrb, E_ARGUMENT_ERROR, "Float not supported");
107 #endif
108 }
109 else {
110 value = mrb_fixnum_value(rand_uint32(t) % mrb_fixnum(max));
111 }
112
113 return value;
114 }
115
116 static mrb_value
get_opt(mrb_state * mrb)117 get_opt(mrb_state* mrb)
118 {
119 mrb_value arg;
120
121 arg = mrb_nil_value();
122 mrb_get_args(mrb, "|o", &arg);
123
124 if (!mrb_nil_p(arg)) {
125 mrb_int i;
126
127 arg = mrb_to_int(mrb, arg);
128 i = mrb_fixnum(arg);
129 if (i < 0) {
130 arg = mrb_fixnum_value(0 - i);
131 }
132 }
133 return arg;
134 }
135
136 static void
random_check(mrb_state * mrb,mrb_value random)137 random_check(mrb_state *mrb, mrb_value random) {
138 struct RClass *c = mrb_class_get(mrb, "Random");
139 if (!mrb_obj_is_kind_of(mrb, random, c) || !mrb_istruct_p(random)) {
140 mrb_raise(mrb, E_TYPE_ERROR, "Random instance required");
141 }
142 }
143
144 static mrb_value
random_default(mrb_state * mrb)145 random_default(mrb_state *mrb) {
146 struct RClass *c = mrb_class_get(mrb, "Random");
147 mrb_value d = mrb_const_get(mrb, mrb_obj_value(c), mrb_intern_lit(mrb, "DEFAULT"));
148 if (!mrb_obj_is_kind_of(mrb, d, c)) {
149 mrb_raise(mrb, E_TYPE_ERROR, "Random::DEFAULT replaced");
150 }
151 return d;
152 }
153
154 #define random_ptr(v) (rand_state*)mrb_istruct_ptr(v)
155 #define random_default_state(mrb) random_ptr(random_default(mrb))
156
157 static mrb_value
random_m_init(mrb_state * mrb,mrb_value self)158 random_m_init(mrb_state *mrb, mrb_value self)
159 {
160 mrb_value seed;
161 rand_state *t;
162
163 seed = get_opt(mrb);
164 /* avoid memory leaks */
165 t = random_ptr(self);
166 if (mrb_nil_p(seed)) {
167 rand_init(t);
168 }
169 else {
170 rand_seed(t, (uint32_t)mrb_fixnum(seed));
171 }
172
173 return self;
174 }
175
176 static mrb_value
random_m_rand(mrb_state * mrb,mrb_value self)177 random_m_rand(mrb_state *mrb, mrb_value self)
178 {
179 mrb_value max;
180 rand_state *t = random_ptr(self);
181
182 max = get_opt(mrb);
183 return random_rand(mrb, t, max);
184 }
185
186 static mrb_value
random_m_srand(mrb_state * mrb,mrb_value self)187 random_m_srand(mrb_state *mrb, mrb_value self)
188 {
189 uint32_t seed;
190 uint32_t old_seed;
191 mrb_value sv;
192 rand_state *t = random_ptr(self);
193
194 sv = get_opt(mrb);
195 if (mrb_nil_p(sv)) {
196 seed = (uint32_t)time(NULL) + rand_uint32(t);
197 }
198 else {
199 seed = (uint32_t)mrb_fixnum(sv);
200 }
201 old_seed = rand_seed(t, seed);
202
203 return mrb_fixnum_value((mrb_int)old_seed);
204 }
205
206 /*
207 * call-seq:
208 * ary.shuffle! -> ary
209 *
210 * Shuffles elements in self in place.
211 */
212
213 static mrb_value
mrb_ary_shuffle_bang(mrb_state * mrb,mrb_value ary)214 mrb_ary_shuffle_bang(mrb_state *mrb, mrb_value ary)
215 {
216 mrb_int i;
217 mrb_value max;
218 mrb_value r = mrb_nil_value();
219 rand_state *random;
220
221 if (RARRAY_LEN(ary) > 1) {
222 mrb_get_args(mrb, "|o", &r);
223
224 if (mrb_nil_p(r)) {
225 random = random_default_state(mrb);
226 }
227 else {
228 random_check(mrb, r);
229 random = random_ptr(r);
230 }
231 mrb_ary_modify(mrb, mrb_ary_ptr(ary));
232 max = mrb_fixnum_value(RARRAY_LEN(ary));
233 for (i = RARRAY_LEN(ary) - 1; i > 0; i--) {
234 mrb_int j;
235 mrb_value *ptr = RARRAY_PTR(ary);
236 mrb_value tmp;
237
238 j = mrb_fixnum(random_rand(mrb, random, max));
239
240 tmp = ptr[i];
241 ptr[i] = ptr[j];
242 ptr[j] = tmp;
243 }
244 }
245
246 return ary;
247 }
248
249 /*
250 * call-seq:
251 * ary.shuffle -> new_ary
252 *
253 * Returns a new array with elements of self shuffled.
254 */
255
256 static mrb_value
mrb_ary_shuffle(mrb_state * mrb,mrb_value ary)257 mrb_ary_shuffle(mrb_state *mrb, mrb_value ary)
258 {
259 mrb_value new_ary = mrb_ary_new_from_values(mrb, RARRAY_LEN(ary), RARRAY_PTR(ary));
260 mrb_ary_shuffle_bang(mrb, new_ary);
261
262 return new_ary;
263 }
264
265 /*
266 * call-seq:
267 * ary.sample -> obj
268 * ary.sample(n) -> new_ary
269 *
270 * Choose a random element or +n+ random elements from the array.
271 *
272 * The elements are chosen by using random and unique indices into the array
273 * in order to ensure that an element doesn't repeat itself unless the array
274 * already contained duplicate elements.
275 *
276 * If the array is empty the first form returns +nil+ and the second form
277 * returns an empty array.
278 */
279
280 static mrb_value
mrb_ary_sample(mrb_state * mrb,mrb_value ary)281 mrb_ary_sample(mrb_state *mrb, mrb_value ary)
282 {
283 mrb_int n = 0;
284 mrb_bool given;
285 mrb_value r = mrb_nil_value();
286 rand_state *random;
287 mrb_int len;
288
289 mrb_get_args(mrb, "|i?o", &n, &given, &r);
290 if (mrb_nil_p(r)) {
291 random = random_default_state(mrb);
292 }
293 else {
294 random_check(mrb, r);
295 random = random_ptr(r);
296 }
297 len = RARRAY_LEN(ary);
298 if (!given) { /* pick one element */
299 switch (len) {
300 case 0:
301 return mrb_nil_value();
302 case 1:
303 return RARRAY_PTR(ary)[0];
304 default:
305 return RARRAY_PTR(ary)[rand_uint32(random) % len];
306 }
307 }
308 else {
309 mrb_value result;
310 mrb_int i, j;
311
312 if (n < 0) mrb_raise(mrb, E_ARGUMENT_ERROR, "negative sample number");
313 if (n > len) n = len;
314 result = mrb_ary_new_capa(mrb, n);
315 for (i=0; i<n; i++) {
316 mrb_int r;
317
318 for (;;) {
319 retry:
320 r = (mrb_int)(rand_uint32(random) % len);
321
322 for (j=0; j<i; j++) {
323 if (mrb_fixnum(RARRAY_PTR(result)[j]) == r) {
324 goto retry; /* retry if duplicate */
325 }
326 }
327 break;
328 }
329 mrb_ary_push(mrb, result, mrb_fixnum_value(r));
330 }
331 for (i=0; i<n; i++) {
332 mrb_ary_set(mrb, result, i, RARRAY_PTR(ary)[mrb_fixnum(RARRAY_PTR(result)[i])]);
333 }
334 return result;
335 }
336 }
337
338 static mrb_value
random_f_rand(mrb_state * mrb,mrb_value self)339 random_f_rand(mrb_state *mrb, mrb_value self)
340 {
341 rand_state *t = random_default_state(mrb);
342 return random_rand(mrb, t, get_opt(mrb));
343 }
344
345 static mrb_value
random_f_srand(mrb_state * mrb,mrb_value self)346 random_f_srand(mrb_state *mrb, mrb_value self)
347 {
348 mrb_value random = random_default(mrb);
349 return random_m_srand(mrb, random);
350 }
351
352
mrb_mruby_random_gem_init(mrb_state * mrb)353 void mrb_mruby_random_gem_init(mrb_state *mrb)
354 {
355 struct RClass *random;
356 struct RClass *array = mrb->array_class;
357
358 mrb_assert(sizeof(rand_state) <= ISTRUCT_DATA_SIZE);
359
360 mrb_define_method(mrb, mrb->kernel_module, "rand", random_f_rand, MRB_ARGS_OPT(1));
361 mrb_define_method(mrb, mrb->kernel_module, "srand", random_f_srand, MRB_ARGS_OPT(1));
362
363 random = mrb_define_class(mrb, "Random", mrb->object_class);
364 MRB_SET_INSTANCE_TT(random, MRB_TT_ISTRUCT);
365 mrb_define_class_method(mrb, random, "rand", random_f_rand, MRB_ARGS_OPT(1));
366 mrb_define_class_method(mrb, random, "srand", random_f_srand, MRB_ARGS_OPT(1));
367
368 mrb_define_method(mrb, random, "initialize", random_m_init, MRB_ARGS_OPT(1));
369 mrb_define_method(mrb, random, "rand", random_m_rand, MRB_ARGS_OPT(1));
370 mrb_define_method(mrb, random, "srand", random_m_srand, MRB_ARGS_OPT(1));
371
372 mrb_define_method(mrb, array, "shuffle", mrb_ary_shuffle, MRB_ARGS_OPT(1));
373 mrb_define_method(mrb, array, "shuffle!", mrb_ary_shuffle_bang, MRB_ARGS_OPT(1));
374 mrb_define_method(mrb, array, "sample", mrb_ary_sample, MRB_ARGS_OPT(2));
375
376 mrb_const_set(mrb, mrb_obj_value(random), mrb_intern_lit(mrb, "DEFAULT"),
377 mrb_obj_new(mrb, random, 0, NULL));
378 }
379
mrb_mruby_random_gem_final(mrb_state * mrb)380 void mrb_mruby_random_gem_final(mrb_state *mrb)
381 {
382 }
383