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