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