enumerator.c

   1 /************************************************
   2
   3   enumerator.c - provides Enumerator class
   4
   5   $Author$
   6
   7   Copyright (C) 2001-2003 Akinori MUSHA
   8
   9   $Idaemons: /home/cvs/rb/enumerator/enumerator.c,v 1.1.1.1 2001/07/15 10:12:48 knu Exp $
  10   $RoughId: enumerator.c,v 1.6 2003/07/27 11:03:24 nobu Exp $
  11   $Id$
  12
  13 ************************************************/
  14
  15 #include "ruby/ruby.h"
  16 #include "debug.h"
  17
  18 /*
  19  * Document-class: Enumerator
  20  *
  21  * A class which provides a method `each' to be used as an Enumerable
  22  * object.
  23  */
  24 VALUE rb_cEnumerator;
  25 static VALUE sym_each;
  26
  27 VALUE rb_eStopIteration;
  28
  29 struct enumerator {
  30     VALUE obj;
  31     ID    meth;
  32     VALUE args;
  33     VALUE fib;
  34     VALUE dst;
  35     VALUE no_next;
  36 };
  37
  38 static void
  39 enumerator_mark(void *p)
  40 {
  41     struct enumerator *ptr = p;
  42     rb_gc_mark(ptr->obj);
  43     rb_gc_mark(ptr->args);
  44     rb_gc_mark(ptr->fib);
  45     rb_gc_mark(ptr->dst);
  46 }
  47
  48 static struct enumerator *
  49 enumerator_ptr(VALUE obj)
  50 {
  51     struct enumerator *ptr;
  52
  53     Data_Get_Struct(obj, struct enumerator, ptr);
  54     if (RDATA(obj)->dmark != enumerator_mark) {
  55         rb_raise(rb_eTypeError,
  56                  "wrong argument type %s (expected %s)",
  57                  rb_obj_classname(obj), rb_class2name(rb_cEnumerator));
  58     }
  59     if (!ptr || ptr->obj == Qundef) {
  60         rb_raise(rb_eArgError, "uninitialized enumerator");
  61     }
  62     return ptr;
  63 }
  64
  65 /*
  66  *  call-seq:
  67  *    obj.to_enum(method = :each, *args)
  68  *    obj.enum_for(method = :each, *args)
  69  *
  70  *  Returns Enumerator.new(self, method, *args).
  71  *
  72  *  e.g.:
  73  *
  74  *     str = "xyz"
  75  *
  76  *     enum = str.enum_for(:each_byte)
  77  *     a = enum.map {|b| '%02x' % b } #=> ["78", "79", "7a"]
  78  *
  79  *     # protects an array from being modified
  80  *     a = [1, 2, 3]
  81  *     some_method(a.to_enum)
  82  *
  83  */
  84 static VALUE
  85 obj_to_enum(int argc, VALUE *argv, VALUE obj)
  86 {
  87     VALUE meth = sym_each;
  88
  89     if (argc > 0) {
  90         --argc;
  91         meth = *argv++;
  92     }
  93     return rb_enumeratorize(obj, meth, argc, argv);
  94 }
  95
  96 static VALUE
  97 each_slice_i(VALUE val, VALUE *memo)
  98 {
  99     VALUE ary = memo[0];
 100     VALUE v = Qnil;
 101     long size = (long)memo[1];
 102
 103     rb_ary_push(ary, val);
 104
 105     if (RARRAY_LEN(ary) == size) {
 106         v = rb_yield(ary);
 107         memo[0] = rb_ary_new2(size);
 108     }
 109
 110     return v;
 111 }
 112
 113 /*
 114  *  call-seq:
 115  *    e.each_slice(n) {...}
 116  *    e.each_slice(n)
 117  *
 118  *  Iterates the given block for each slice of <n> elements.  If no
 119  *  block is given, returns an enumerator.
 120  *
 121  *  e.g.:
 122  *      (1..10).each_slice(3) {|a| p a}
 123  *      # outputs below
 124  *      [1, 2, 3]
 125  *      [4, 5, 6]
 126  *      [7, 8, 9]
 127  *      [10]
 128  *
 129  */
 130 static VALUE
 131 enum_each_slice(VALUE obj, VALUE n)
 132 {
 133     long size = NUM2LONG(n);
 134     VALUE args[2], ary;
 135
 136     if (size <= 0) rb_raise(rb_eArgError, "invalid slice size");
 137     RETURN_ENUMERATOR(obj, 1, &n);
 138     args[0] = rb_ary_new2(size);
 139     args[1] = (VALUE)size;
 140
 141     rb_block_call(obj, SYM2ID(sym_each), 0, 0, each_slice_i, (VALUE)args);
 142
 143     ary = args[0];
 144     if (RARRAY_LEN(ary) > 0) rb_yield(ary);
 145
 146     return Qnil;
 147 }
 148
 149 static VALUE
 150 each_cons_i(VALUE val, VALUE *memo)
 151 {
 152     VALUE ary = memo[0];
 153     VALUE v = Qnil;
 154     long size = (long)memo[1];
 155
 156     if (RARRAY_LEN(ary) == size) {
 157         rb_ary_shift(ary);
 158     }
 159     rb_ary_push(ary, val);
 160     if (RARRAY_LEN(ary) == size) {
 161         v = rb_yield(rb_ary_dup(ary));
 162     }
 163     return v;
 164 }
 165
 166 /*
 167  *  call-seq:
 168  *    each_cons(n) {...}
 169  *    each_cons(n)
 170  *
 171  *  Iterates the given block for each array of consecutive <n>
 172  *  elements.  If no block is given, returns an enumerator.
 173  *
 174  *  e.g.:
 175  *      (1..10).each_cons(3) {|a| p a}
 176  *      # outputs below
 177  *      [1, 2, 3]
 178  *      [2, 3, 4]
 179  *      [3, 4, 5]
 180  *      [4, 5, 6]
 181  *      [5, 6, 7]
 182  *      [6, 7, 8]
 183  *      [7, 8, 9]
 184  *      [8, 9, 10]
 185  *
 186  */
 187 static VALUE
 188 enum_each_cons(VALUE obj, VALUE n)
 189 {
 190     long size = NUM2LONG(n);
 191     VALUE args[2];
 192
 193     if (size <= 0) rb_raise(rb_eArgError, "invalid size");
 194     RETURN_ENUMERATOR(obj, 1, &n);
 195     args[0] = rb_ary_new2(size);
 196     args[1] = (VALUE)size;
 197
 198     rb_block_call(obj, SYM2ID(sym_each), 0, 0, each_cons_i, (VALUE)args);
 199
 200     return Qnil;
 201 }
 202
 203 static VALUE
 204 enumerator_allocate(VALUE klass)
 205 {
 206     struct enumerator *ptr;
 207     VALUE enum_obj;
 208
 209     enum_obj = Data_Make_Struct(klass, struct enumerator, enumerator_mark, -1, ptr);
 210     ptr->obj = Qundef;
 211
 212     return enum_obj;
 213 }
 214
 215 static VALUE
 216 enumerator_each_i(VALUE v, VALUE enum_obj, int argc, VALUE *argv)
 217 {
 218     return rb_yield_values2(argc, argv);
 219 }
 220
 221 static VALUE
 222 enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, VALUE *argv)
 223 {
 224     struct enumerator *ptr;
 225
 226     Data_Get_Struct(enum_obj, struct enumerator, ptr);
 227
 228     if (!ptr) {
 229         rb_raise(rb_eArgError, "unallocated enumerator");
 230     }
 231
 232     ptr->obj  = obj;
 233     ptr->meth = rb_to_id(meth);
 234     if (argc) ptr->args = rb_ary_new4(argc, argv);
 235     ptr->fib = 0;
 236     ptr->dst = Qnil;
 237     ptr->no_next = Qfalse;
 238
 239     return enum_obj;
 240 }
 241
 242 /*
 243  *  call-seq:
 244  *    Enumerator.new(obj, method = :each, *args)
 245  *
 246  *  Creates a new Enumerator object, which is to be used as an
 247  *  Enumerable object using the given object's given method with the
 248  *  given arguments.
 249  *
 250  *  Use of this method is discouraged.  Use Kernel#enum_for() instead.
 251  */
 252 static VALUE
 253 enumerator_initialize(int argc, VALUE *argv, VALUE obj)
 254 {
 255     VALUE recv, meth = sym_each;
 256
 257     if (argc == 0)
 258         rb_raise(rb_eArgError, "wrong number of argument (0 for 1)");
 259     recv = *argv++;
 260     if (--argc) {
 261         meth = *argv++;
 262         --argc;
 263     }
 264     return enumerator_init(obj, recv, meth, argc, argv);
 265 }
 266
 267 /* :nodoc: */
 268 static VALUE
 269 enumerator_init_copy(VALUE obj, VALUE orig)
 270 {
 271     struct enumerator *ptr0, *ptr1;
 272
 273     ptr0 = enumerator_ptr(orig);
 274     if (ptr0->fib) {
 275         /* Fibers cannot be copied */
 276         rb_raise(rb_eTypeError, "can't copy execution context");
 277     }
 278
 279     Data_Get_Struct(obj, struct enumerator, ptr1);
 280
 281     if (!ptr1) {
 282         rb_raise(rb_eArgError, "unallocated enumerator");
 283     }
 284
 285     ptr1->obj  = ptr0->obj;
 286     ptr1->meth = ptr0->meth;
 287     ptr1->args = ptr0->args;
 288     ptr1->fib  = 0;
 289
 290     return obj;
 291 }
 292
 293 VALUE
 294 rb_enumeratorize(VALUE obj, VALUE meth, int argc, VALUE *argv)
 295 {
 296     return enumerator_init(enumerator_allocate(rb_cEnumerator), obj, meth, argc, argv);
 297 }
 298
 299 /*
 300  *  call-seq:
 301  *    enum.each {...}
 302  *
 303  *  Iterates the given block using the object and the method specified
 304  *  in the first place.  If no block is given, returns self.
 305  *
 306  */
 307 static VALUE
 308 enumerator_each(VALUE obj)
 309 {
 310     struct enumerator *e;
 311     int argc = 0;
 312     VALUE *argv = 0;
 313
 314     if (!rb_block_given_p()) return obj;
 315     e = enumerator_ptr(obj);
 316     if (e->args) {
 317         argc = RARRAY_LEN(e->args);
 318         argv = RARRAY_PTR(e->args);
 319     }
 320     return rb_block_call(e->obj, e->meth, argc, argv,
 321                          enumerator_each_i, (VALUE)e);
 322 }
 323
 324 static VALUE
 325 enumerator_with_index_i(VALUE val, VALUE *memo)
 326 {
 327     val = rb_yield_values(2, val, INT2FIX(*memo));
 328     ++*memo;
 329     return val;
 330 }
 331
 332 /*
 333  *  call-seq:
 334  *    e.with_index {|(*args), idx| ... }
 335  *    e.with_index
 336  *
 337  *  Iterates the given block for each element with an index, which
 338  *  start from 0.  If no block is given, returns an enumerator.
 339  *
 340  */
 341 static VALUE
 342 enumerator_with_index(VALUE obj)
 343 {
 344     struct enumerator *e;
 345     VALUE memo = 0;
 346     int argc = 0;
 347     VALUE *argv = 0;
 348
 349     RETURN_ENUMERATOR(obj, 0, 0);
 350     e = enumerator_ptr(obj);
 351     if (e->args) {
 352         argc = RARRAY_LEN(e->args);
 353         argv = RARRAY_PTR(e->args);
 354     }
 355     return rb_block_call(e->obj, e->meth, argc, argv,
 356                          enumerator_with_index_i, (VALUE)&memo);
 357 }
 358
 359 static VALUE
 360 enumerator_with_object_i(VALUE val, VALUE memo)
 361 {
 362     return rb_yield_values(2, val, memo);
 363 }
 364
 365 /*
 366  *  call-seq:
 367  *    e.with_object(obj) {|(*args), memo_obj| ... }
 368  *    e.with_object(obj)
 369  *
 370  *  Iterates the given block for each element with an arbitrary
 371  *  object given, and returns the memo object.
 372  *
 373  *  If no block is given, returns an enumerator.
 374  *
 375  */
 376 static VALUE
 377 enumerator_with_object(VALUE obj, VALUE memo)
 378 {
 379     struct enumerator *e;
 380     int argc = 0;
 381     VALUE *argv = 0;
 382
 383     RETURN_ENUMERATOR(obj, 0, 0);
 384     e = enumerator_ptr(obj);
 385     if (e->args) {
 386         argc = RARRAY_LEN(e->args);
 387         argv = RARRAY_PTR(e->args);
 388     }
 389     rb_block_call(e->obj, e->meth, argc, argv,
 390                   enumerator_with_object_i, memo);
 391
 392     return memo;
 393 }
 394
 395 static VALUE
 396 next_ii(VALUE i, VALUE obj, int argc, VALUE *argv)
 397 {
 398     rb_fiber_yield(argc, argv);
 399     return Qnil;
 400 }
 401
 402 static VALUE
 403 next_i(VALUE curr, VALUE obj)
 404 {
 405     struct enumerator *e = enumerator_ptr(obj);
 406     VALUE nil = Qnil;
 407
 408     rb_block_call(obj, rb_intern("each"), 0, 0, next_ii, obj);
 409     e->no_next = Qtrue;
 410     return rb_fiber_yield(1, &nil);
 411 }
 412
 413 static void
 414 next_init(VALUE obj, struct enumerator *e)
 415 {
 416     VALUE curr = rb_fiber_current();
 417     e->dst = curr;
 418     e->fib = rb_fiber_new(next_i, obj);
 419 }
 420
 421 /*
 422  * call-seq:
 423  *   e.next   => object
 424  *
 425  * Returns the next object in the enumerator, and move the internal
 426  * position forward.  When the position reached at the end, internal
 427  * position is rewinded then StopIteration is raised.
 428  *
 429  * Note that enumeration sequence by next method does not affect other
 430  * non-external enumeration methods, unless underlying iteration
 431  * methods itself has side-effect, e.g. IO#each_line.
 432  *
 433  */
 434
 435 static VALUE
 436 enumerator_next(VALUE obj)
 437 {
 438     struct enumerator *e = enumerator_ptr(obj);
 439     VALUE curr, v;
 440     curr = rb_fiber_current();
 441
 442     if (!e->fib || !rb_fiber_alive_p(e->fib)) {
 443         next_init(obj, e);
 444     }
 445
 446     v = rb_fiber_resume(e->fib, 1, &curr);
 447     if (e->no_next) {
 448         e->fib = 0;
 449         e->dst = Qnil;
 450         e->no_next = Qfalse;
 451         rb_raise(rb_eStopIteration, "iteration reached at end");
 452     }
 453     return v;
 454 }
 455
 456 /*
 457  * call-seq:
 458  *   e.rewind   => e
 459  *
 460  * Rewinds the enumeration sequence by the next method.
 461  */
 462
 463 static VALUE
 464 enumerator_rewind(VALUE obj)
 465 {
 466     struct enumerator *e = enumerator_ptr(obj);
 467
 468     e->fib = 0;
 469     e->dst = Qnil;
 470     e->no_next = Qfalse;
 471     return obj;
 472 }
 473
 474 void
 475 Init_Enumerator(void)
 476 {
 477     rb_define_method(rb_mKernel, "to_enum", obj_to_enum, -1);
 478     rb_define_method(rb_mKernel, "enum_for", obj_to_enum, -1);
 479
 480     rb_define_method(rb_mEnumerable, "each_slice", enum_each_slice, 1);
 481     rb_define_method(rb_mEnumerable, "each_cons", enum_each_cons, 1);
 482
 483     rb_cEnumerator = rb_define_class("Enumerator", rb_cObject);
 484     rb_include_module(rb_cEnumerator, rb_mEnumerable);
 485
 486     rb_define_alloc_func(rb_cEnumerator, enumerator_allocate);
 487     rb_define_method(rb_cEnumerator, "initialize", enumerator_initialize, -1);
 488     rb_define_method(rb_cEnumerator, "initialize_copy", enumerator_init_copy, 1);
 489     rb_define_method(rb_cEnumerator, "each", enumerator_each, 0);
 490     rb_define_method(rb_cEnumerator, "each_with_index", enumerator_with_index, 0);
 491     rb_define_method(rb_cEnumerator, "with_index", enumerator_with_index, 0);
 492 #if 0
 493     rb_define_method(rb_cEnumerator, "with_object", enumerator_with_object, 1);
 494 #else
 495     (void)enumerator_with_object;
 496 #endif
 497     rb_define_method(rb_cEnumerator, "next", enumerator_next, 0);
 498     rb_define_method(rb_cEnumerator, "rewind", enumerator_rewind, 0);
 499
 500     rb_eStopIteration   = rb_define_class("StopIteration", rb_eIndexError);
 501
 502     sym_each            = ID2SYM(rb_intern("each"));
 503
 504     rb_provide("enumerator.so");        /* for backward compatibility */
 505 }