1/************************************************
2
3  enumerator.c - provides Enumerator class
4
5  $Author: nagachika $
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 44150 2013-12-12 16:02:08Z nagachika $
12
13************************************************/
14
15#include "ruby/ruby.h"
16#include "node.h"
17#include "internal.h"
18
19/*
20 * Document-class: Enumerator
21 *
22 * A class which allows both internal and external iteration.
23 *
24 * An Enumerator can be created by the following methods.
25 * - Kernel#to_enum
26 * - Kernel#enum_for
27 * - Enumerator.new
28 *
29 * Most methods have two forms: a block form where the contents
30 * are evaluated for each item in the enumeration, and a non-block form
31 * which returns a new Enumerator wrapping the iteration.
32 *
33 *   enumerator = %w(one two three).each
34 *   puts enumerator.class # => Enumerator
35 *
36 *   enumerator.each_with_object("foo") do |item, obj|
37 *     puts "#{obj}: #{item}"
38 *   end
39 *
40 *   # foo: one
41 *   # foo: two
42 *   # foo: three
43 *
44 *   enum_with_obj = enumerator.each_with_object("foo")
45 *   puts enum_with_obj.class # => Enumerator
46 *
47 *   enum_with_obj.each do |item, obj|
48 *     puts "#{obj}: #{item}"
49 *   end
50 *
51 *   # foo: one
52 *   # foo: two
53 *   # foo: three
54 *
55 * This allows you to chain Enumerators together.  For example, you
56 * can map a list's elements to strings containing the index
57 * and the element as a string via:
58 *
59 *   puts %w[foo bar baz].map.with_index { |w, i| "#{i}:#{w}" }
60 *   # => ["0:foo", "1:bar", "2:baz"]
61 *
62 * An Enumerator can also be used as an external iterator.
63 * For example, Enumerator#next returns the next value of the iterator
64 * or raises StopIteration if the Enumerator is at the end.
65 *
66 *   e = [1,2,3].each   # returns an enumerator object.
67 *   puts e.next   # => 1
68 *   puts e.next   # => 2
69 *   puts e.next   # => 3
70 *   puts e.next   # raises StopIteration
71 *
72 * You can use this to implement an internal iterator as follows:
73 *
74 *   def ext_each(e)
75 *     while true
76 *       begin
77 *         vs = e.next_values
78 *       rescue StopIteration
79 *         return $!.result
80 *       end
81 *       y = yield(*vs)
82 *       e.feed y
83 *     end
84 *   end
85 *
86 *   o = Object.new
87 *
88 *   def o.each
89 *     puts yield
90 *     puts yield(1)
91 *     puts yield(1, 2)
92 *     3
93 *   end
94 *
95 *   # use o.each as an internal iterator directly.
96 *   puts o.each {|*x| puts x; [:b, *x] }
97 *   # => [], [:b], [1], [:b, 1], [1, 2], [:b, 1, 2], 3
98 *
99 *   # convert o.each to an external iterator for
100 *   # implementing an internal iterator.
101 *   puts ext_each(o.to_enum) {|*x| puts x; [:b, *x] }
102 *   # => [], [:b], [1], [:b, 1], [1, 2], [:b, 1, 2], 3
103 *
104 */
105VALUE rb_cEnumerator;
106VALUE rb_cLazy;
107static ID id_rewind, id_each, id_new, id_initialize, id_yield, id_call, id_size, id_to_enum;
108static ID id_eqq, id_next, id_result, id_lazy, id_receiver, id_arguments, id_memo, id_method, id_force;
109static VALUE sym_each, sym_cycle;
110
111VALUE rb_eStopIteration;
112
113struct enumerator {
114    VALUE obj;
115    ID    meth;
116    VALUE args;
117    VALUE fib;
118    VALUE dst;
119    VALUE lookahead;
120    VALUE feedvalue;
121    VALUE stop_exc;
122    VALUE size;
123    VALUE (*size_fn)(ANYARGS);
124};
125
126static VALUE rb_cGenerator, rb_cYielder;
127
128struct generator {
129    VALUE proc;
130};
131
132struct yielder {
133    VALUE proc;
134};
135
136static VALUE generator_allocate(VALUE klass);
137static VALUE generator_init(VALUE obj, VALUE proc);
138
139/*
140 * Enumerator
141 */
142static void
143enumerator_mark(void *p)
144{
145    struct enumerator *ptr = p;
146    rb_gc_mark(ptr->obj);
147    rb_gc_mark(ptr->args);
148    rb_gc_mark(ptr->fib);
149    rb_gc_mark(ptr->dst);
150    rb_gc_mark(ptr->lookahead);
151    rb_gc_mark(ptr->feedvalue);
152    rb_gc_mark(ptr->stop_exc);
153    rb_gc_mark(ptr->size);
154}
155
156#define enumerator_free RUBY_TYPED_DEFAULT_FREE
157
158static size_t
159enumerator_memsize(const void *p)
160{
161    return p ? sizeof(struct enumerator) : 0;
162}
163
164static const rb_data_type_t enumerator_data_type = {
165    "enumerator",
166    {
167	enumerator_mark,
168	enumerator_free,
169	enumerator_memsize,
170    },
171};
172
173static struct enumerator *
174enumerator_ptr(VALUE obj)
175{
176    struct enumerator *ptr;
177
178    TypedData_Get_Struct(obj, struct enumerator, &enumerator_data_type, ptr);
179    if (!ptr || ptr->obj == Qundef) {
180	rb_raise(rb_eArgError, "uninitialized enumerator");
181    }
182    return ptr;
183}
184
185/*
186 * call-seq:
187 *   obj.to_enum(method = :each, *args)                 -> enum
188 *   obj.enum_for(method = :each, *args)                -> enum
189 *   obj.to_enum(method = :each, *args) {|*args| block} -> enum
190 *   obj.enum_for(method = :each, *args){|*args| block} -> enum
191 *
192 * Creates a new Enumerator which will enumerate by calling +method+ on
193 * +obj+, passing +args+ if any.
194 *
195 * If a block is given, it will be used to calculate the size of
196 * the enumerator without the need to iterate it (see Enumerator#size).
197 *
198 * === Examples
199 *
200 *   str = "xyz"
201 *
202 *   enum = str.enum_for(:each_byte)
203 *   enum.each { |b| puts b }
204 *   # => 120
205 *   # => 121
206 *   # => 122
207 *
208 *   # protect an array from being modified by some_method
209 *   a = [1, 2, 3]
210 *   some_method(a.to_enum)
211 *
212 * It is typical to call to_enum when defining methods for
213 * a generic Enumerable, in case no block is passed.
214 *
215 * Here is such an example, with parameter passing and a sizing block:
216 *
217 *   module Enumerable
218 *     # a generic method to repeat the values of any enumerable
219 *     def repeat(n)
220 *       raise ArgumentError, "#{n} is negative!" if n < 0
221 *       unless block_given?
222 *         return to_enum(__method__, n) do # __method__ is :repeat here
223 *           sz = size     # Call size and multiply by n...
224 *           sz * n if sz  # but return nil if size itself is nil
225 *         end
226 *       end
227 *       each do |*val|
228 *         n.times { yield *val }
229 *       end
230 *     end
231 *   end
232 *
233 *   %i[hello world].repeat(2) { |w| puts w }
234 *     # => Prints 'hello', 'hello', 'world', 'world'
235 *   enum = (1..14).repeat(3)
236 *     # => returns an Enumerator when called without a block
237 *   enum.first(4) # => [1, 1, 1, 2]
238 *   enum.size # => 42
239 */
240static VALUE
241obj_to_enum(int argc, VALUE *argv, VALUE obj)
242{
243    VALUE enumerator, meth = sym_each;
244
245    if (argc > 0) {
246	--argc;
247	meth = *argv++;
248    }
249    enumerator = rb_enumeratorize_with_size(obj, meth, argc, argv, 0);
250    if (rb_block_given_p()) {
251	enumerator_ptr(enumerator)->size = rb_block_proc();
252    }
253    return enumerator;
254}
255
256static VALUE
257enumerator_allocate(VALUE klass)
258{
259    struct enumerator *ptr;
260    VALUE enum_obj;
261
262    enum_obj = TypedData_Make_Struct(klass, struct enumerator, &enumerator_data_type, ptr);
263    ptr->obj = Qundef;
264
265    return enum_obj;
266}
267
268static VALUE
269enumerator_init(VALUE enum_obj, VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS), VALUE size)
270{
271    struct enumerator *ptr;
272
273    TypedData_Get_Struct(enum_obj, struct enumerator, &enumerator_data_type, ptr);
274
275    if (!ptr) {
276	rb_raise(rb_eArgError, "unallocated enumerator");
277    }
278
279    ptr->obj  = obj;
280    ptr->meth = rb_to_id(meth);
281    if (argc) ptr->args = rb_ary_new4(argc, argv);
282    ptr->fib = 0;
283    ptr->dst = Qnil;
284    ptr->lookahead = Qundef;
285    ptr->feedvalue = Qundef;
286    ptr->stop_exc = Qfalse;
287    ptr->size = size;
288    ptr->size_fn = size_fn;
289
290    return enum_obj;
291}
292
293/*
294 * call-seq:
295 *   Enumerator.new(size = nil) { |yielder| ... }
296 *   Enumerator.new(obj, method = :each, *args)
297 *
298 * Creates a new Enumerator object, which can be used as an
299 * Enumerable.
300 *
301 * In the first form, iteration is defined by the given block, in
302 * which a "yielder" object, given as block parameter, can be used to
303 * yield a value by calling the +yield+ method (aliased as +<<+):
304 *
305 *   fib = Enumerator.new do |y|
306 *     a = b = 1
307 *     loop do
308 *       y << a
309 *       a, b = b, a + b
310 *     end
311 *   end
312 *
313 *   p fib.take(10) # => [1, 1, 2, 3, 5, 8, 13, 21, 34, 55]
314 *
315 * The optional parameter can be used to specify how to calculate the size
316 * in a lazy fashion (see Enumerator#size). It can either be a value or
317 * a callable object.
318 *
319 * In the second, deprecated, form, a generated Enumerator iterates over the
320 * given object using the given method with the given arguments passed.
321 *
322 * Use of this form is discouraged.  Use Kernel#enum_for or Kernel#to_enum
323 * instead.
324 *
325 *   e = Enumerator.new(ObjectSpace, :each_object)
326 *       #-> ObjectSpace.enum_for(:each_object)
327 *
328 *   e.select { |obj| obj.is_a?(Class) }  #=> array of all classes
329 *
330 */
331static VALUE
332enumerator_initialize(int argc, VALUE *argv, VALUE obj)
333{
334    VALUE recv, meth = sym_each;
335    VALUE size = Qnil;
336
337    if (rb_block_given_p()) {
338	rb_check_arity(argc, 0, 1);
339	recv = generator_init(generator_allocate(rb_cGenerator), rb_block_proc());
340	if (argc) {
341            if (NIL_P(argv[0]) || rb_obj_is_proc(argv[0]) ||
342                (RB_TYPE_P(argv[0], T_FLOAT) && RFLOAT_VALUE(argv[0]) == INFINITY)) {
343                size = argv[0];
344            } else {
345                size = rb_to_int(argv[0]);
346            }
347            argc = 0;
348        }
349    }
350    else {
351	rb_check_arity(argc, 1, UNLIMITED_ARGUMENTS);
352	rb_warn("Enumerator.new without a block is deprecated; use Object#to_enum");
353	recv = *argv++;
354	if (--argc) {
355	    meth = *argv++;
356	    --argc;
357	}
358    }
359
360    return enumerator_init(obj, recv, meth, argc, argv, 0, size);
361}
362
363/* :nodoc: */
364static VALUE
365enumerator_init_copy(VALUE obj, VALUE orig)
366{
367    struct enumerator *ptr0, *ptr1;
368
369    if (!OBJ_INIT_COPY(obj, orig)) return obj;
370    ptr0 = enumerator_ptr(orig);
371    if (ptr0->fib) {
372	/* Fibers cannot be copied */
373	rb_raise(rb_eTypeError, "can't copy execution context");
374    }
375
376    TypedData_Get_Struct(obj, struct enumerator, &enumerator_data_type, ptr1);
377
378    if (!ptr1) {
379	rb_raise(rb_eArgError, "unallocated enumerator");
380    }
381
382    ptr1->obj  = ptr0->obj;
383    ptr1->meth = ptr0->meth;
384    ptr1->args = ptr0->args;
385    ptr1->fib  = 0;
386    ptr1->lookahead  = Qundef;
387    ptr1->feedvalue  = Qundef;
388    ptr1->size  = ptr0->size;
389    ptr1->size_fn  = ptr0->size_fn;
390
391    return obj;
392}
393
394/*
395 * For backwards compatibility; use rb_enumeratorize_with_size
396 */
397VALUE
398rb_enumeratorize(VALUE obj, VALUE meth, int argc, VALUE *argv)
399{
400    return rb_enumeratorize_with_size(obj, meth, argc, argv, 0);
401}
402
403static VALUE
404lazy_to_enum_i(VALUE self, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS));
405
406VALUE
407rb_enumeratorize_with_size(VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS))
408{
409    /* Similar effect as calling obj.to_enum, i.e. dispatching to either
410       Kernel#to_enum vs Lazy#to_enum */
411    if (RTEST(rb_obj_is_kind_of(obj, rb_cLazy)))
412	return lazy_to_enum_i(obj, meth, argc, argv, size_fn);
413    else
414	return enumerator_init(enumerator_allocate(rb_cEnumerator),
415	    obj, meth, argc, argv, size_fn, Qnil);
416}
417
418static VALUE
419enumerator_block_call(VALUE obj, rb_block_call_func *func, VALUE arg)
420{
421    int argc = 0;
422    VALUE *argv = 0;
423    const struct enumerator *e = enumerator_ptr(obj);
424    ID meth = e->meth;
425
426    if (e->args) {
427	argc = RARRAY_LENINT(e->args);
428	argv = RARRAY_PTR(e->args);
429    }
430    return rb_block_call(e->obj, meth, argc, argv, func, arg);
431}
432
433/*
434 * call-seq:
435 *   enum.each {...}
436 *
437 * Iterates over the block according to how this Enumerable was constructed.
438 * If no block is given, returns self.
439 *
440 */
441static VALUE
442enumerator_each(int argc, VALUE *argv, VALUE obj)
443{
444    if (argc > 0) {
445	struct enumerator *e = enumerator_ptr(obj = rb_obj_dup(obj));
446	VALUE args = e->args;
447	if (args) {
448	    args = rb_ary_dup(args);
449	    rb_ary_cat(args, argv, argc);
450	}
451	else {
452	    args = rb_ary_new4(argc, argv);
453	}
454	e->args = args;
455    }
456    if (!rb_block_given_p()) return obj;
457    return enumerator_block_call(obj, 0, obj);
458}
459
460static VALUE
461enumerator_with_index_i(VALUE val, VALUE m, int argc, VALUE *argv)
462{
463    NODE *memo = (NODE *)m;
464    VALUE idx = memo->u1.value;
465    memo->u1.value = rb_int_succ(idx);
466
467    if (argc <= 1)
468	return rb_yield_values(2, val, idx);
469
470    return rb_yield_values(2, rb_ary_new4(argc, argv), idx);
471}
472
473static VALUE
474enumerator_size(VALUE obj);
475
476/*
477 * call-seq:
478 *   e.with_index(offset = 0) {|(*args), idx| ... }
479 *   e.with_index(offset = 0)
480 *
481 * Iterates the given block for each element with an index, which
482 * starts from +offset+.  If no block is given, returns a new Enumerator
483 * that includes the index, starting from +offset+
484 *
485 * +offset+:: the starting index to use
486 *
487 */
488static VALUE
489enumerator_with_index(int argc, VALUE *argv, VALUE obj)
490{
491    VALUE memo;
492
493    rb_scan_args(argc, argv, "01", &memo);
494    RETURN_SIZED_ENUMERATOR(obj, argc, argv, enumerator_size);
495    if (NIL_P(memo))
496	memo = INT2FIX(0);
497    else
498	memo = rb_to_int(memo);
499    return enumerator_block_call(obj, enumerator_with_index_i, (VALUE)NEW_MEMO(memo, 0, 0));
500}
501
502/*
503 * call-seq:
504 *   e.each_with_index {|(*args), idx| ... }
505 *   e.each_with_index
506 *
507 * Same as Enumerator#with_index(0), i.e. there is no starting offset.
508 *
509 * If no block is given, a new Enumerator is returned that includes the index.
510 *
511 */
512static VALUE
513enumerator_each_with_index(VALUE obj)
514{
515    return enumerator_with_index(0, NULL, obj);
516}
517
518static VALUE
519enumerator_with_object_i(VALUE val, VALUE memo, int argc, VALUE *argv)
520{
521    if (argc <= 1)
522	return rb_yield_values(2, val, memo);
523
524    return rb_yield_values(2, rb_ary_new4(argc, argv), memo);
525}
526
527/*
528 * call-seq:
529 *   e.with_object(obj) {|(*args), obj| ... }
530 *   e.with_object(obj)
531 *
532 * Iterates the given block for each element with an arbitrary object, +obj+,
533 * and returns +obj+
534 *
535 * If no block is given, returns a new Enumerator.
536 *
537 * === Example
538 *
539 *   to_three = Enumerator.new do |y|
540 *     3.times do |x|
541 *       y << x
542 *     end
543 *   end
544 *
545 *   to_three_with_string = to_three.with_object("foo")
546 *   to_three_with_string.each do |x,string|
547 *     puts "#{string}: #{x}"
548 *   end
549 *
550 *   # => foo:0
551 *   # => foo:1
552 *   # => foo:2
553 */
554static VALUE
555enumerator_with_object(VALUE obj, VALUE memo)
556{
557    RETURN_SIZED_ENUMERATOR(obj, 1, &memo, enumerator_size);
558    enumerator_block_call(obj, enumerator_with_object_i, memo);
559
560    return memo;
561}
562
563static VALUE
564next_ii(VALUE i, VALUE obj, int argc, VALUE *argv)
565{
566    struct enumerator *e = enumerator_ptr(obj);
567    VALUE feedvalue = Qnil;
568    VALUE args = rb_ary_new4(argc, argv);
569    rb_fiber_yield(1, &args);
570    if (e->feedvalue != Qundef) {
571        feedvalue = e->feedvalue;
572        e->feedvalue = Qundef;
573    }
574    return feedvalue;
575}
576
577static VALUE
578next_i(VALUE curr, VALUE obj)
579{
580    struct enumerator *e = enumerator_ptr(obj);
581    VALUE nil = Qnil;
582    VALUE result;
583
584    result = rb_block_call(obj, id_each, 0, 0, next_ii, obj);
585    e->stop_exc = rb_exc_new2(rb_eStopIteration, "iteration reached an end");
586    rb_ivar_set(e->stop_exc, id_result, result);
587    return rb_fiber_yield(1, &nil);
588}
589
590static void
591next_init(VALUE obj, struct enumerator *e)
592{
593    VALUE curr = rb_fiber_current();
594    e->dst = curr;
595    e->fib = rb_fiber_new(next_i, obj);
596    e->lookahead = Qundef;
597}
598
599static VALUE
600get_next_values(VALUE obj, struct enumerator *e)
601{
602    VALUE curr, vs;
603
604    if (e->stop_exc)
605	rb_exc_raise(e->stop_exc);
606
607    curr = rb_fiber_current();
608
609    if (!e->fib || !rb_fiber_alive_p(e->fib)) {
610	next_init(obj, e);
611    }
612
613    vs = rb_fiber_resume(e->fib, 1, &curr);
614    if (e->stop_exc) {
615	e->fib = 0;
616	e->dst = Qnil;
617	e->lookahead = Qundef;
618	e->feedvalue = Qundef;
619	rb_exc_raise(e->stop_exc);
620    }
621    return vs;
622}
623
624/*
625 * call-seq:
626 *   e.next_values   -> array
627 *
628 * Returns the next object as an array in the enumerator, and move the
629 * internal position forward.  When the position reached at the end,
630 * StopIteration is raised.
631 *
632 * This method can be used to distinguish <code>yield</code> and <code>yield
633 * nil</code>.
634 *
635 * === Example
636 *
637 *   o = Object.new
638 *   def o.each
639 *     yield
640 *     yield 1
641 *     yield 1, 2
642 *     yield nil
643 *     yield [1, 2]
644 *   end
645 *   e = o.to_enum
646 *   p e.next_values
647 *   p e.next_values
648 *   p e.next_values
649 *   p e.next_values
650 *   p e.next_values
651 *   e = o.to_enum
652 *   p e.next
653 *   p e.next
654 *   p e.next
655 *   p e.next
656 *   p e.next
657 *
658 *   ## yield args       next_values      next
659 *   #  yield            []               nil
660 *   #  yield 1          [1]              1
661 *   #  yield 1, 2       [1, 2]           [1, 2]
662 *   #  yield nil        [nil]            nil
663 *   #  yield [1, 2]     [[1, 2]]         [1, 2]
664 *
665 * Note that +next_values+ does not affect other non-external enumeration
666 * methods unless underlying iteration method itself has side-effect, e.g.
667 * IO#each_line.
668 *
669 */
670
671static VALUE
672enumerator_next_values(VALUE obj)
673{
674    struct enumerator *e = enumerator_ptr(obj);
675    VALUE vs;
676
677    if (e->lookahead != Qundef) {
678        vs = e->lookahead;
679        e->lookahead = Qundef;
680        return vs;
681    }
682
683    return get_next_values(obj, e);
684}
685
686static VALUE
687ary2sv(VALUE args, int dup)
688{
689    if (!RB_TYPE_P(args, T_ARRAY))
690        return args;
691
692    switch (RARRAY_LEN(args)) {
693      case 0:
694        return Qnil;
695
696      case 1:
697        return RARRAY_PTR(args)[0];
698
699      default:
700        if (dup)
701            return rb_ary_dup(args);
702        return args;
703    }
704}
705
706/*
707 * call-seq:
708 *   e.next   -> object
709 *
710 * Returns the next object in the enumerator, and move the internal position
711 * forward.  When the position reached at the end, StopIteration is raised.
712 *
713 * === Example
714 *
715 *   a = [1,2,3]
716 *   e = a.to_enum
717 *   p e.next   #=> 1
718 *   p e.next   #=> 2
719 *   p e.next   #=> 3
720 *   p e.next   #raises StopIteration
721 *
722 * Note that enumeration sequence by +next+ does not affect other non-external
723 * enumeration methods, unless the underlying iteration methods itself has
724 * side-effect, e.g. IO#each_line.
725 *
726 */
727
728static VALUE
729enumerator_next(VALUE obj)
730{
731    VALUE vs = enumerator_next_values(obj);
732    return ary2sv(vs, 0);
733}
734
735static VALUE
736enumerator_peek_values(VALUE obj)
737{
738    struct enumerator *e = enumerator_ptr(obj);
739
740    if (e->lookahead == Qundef) {
741        e->lookahead = get_next_values(obj, e);
742    }
743    return e->lookahead;
744}
745
746/*
747 * call-seq:
748 *   e.peek_values   -> array
749 *
750 * Returns the next object as an array, similar to Enumerator#next_values, but
751 * doesn't move the internal position forward.  If the position is already at
752 * the end, StopIteration is raised.
753 *
754 * === Example
755 *
756 *   o = Object.new
757 *   def o.each
758 *     yield
759 *     yield 1
760 *     yield 1, 2
761 *   end
762 *   e = o.to_enum
763 *   p e.peek_values    #=> []
764 *   e.next
765 *   p e.peek_values    #=> [1]
766 *   p e.peek_values    #=> [1]
767 *   e.next
768 *   p e.peek_values    #=> [1, 2]
769 *   e.next
770 *   p e.peek_values    # raises StopIteration
771 *
772 */
773
774static VALUE
775enumerator_peek_values_m(VALUE obj)
776{
777    return rb_ary_dup(enumerator_peek_values(obj));
778}
779
780/*
781 * call-seq:
782 *   e.peek   -> object
783 *
784 * Returns the next object in the enumerator, but doesn't move the internal
785 * position forward.  If the position is already at the end, StopIteration
786 * is raised.
787 *
788 * === Example
789 *
790 *   a = [1,2,3]
791 *   e = a.to_enum
792 *   p e.next   #=> 1
793 *   p e.peek   #=> 2
794 *   p e.peek   #=> 2
795 *   p e.peek   #=> 2
796 *   p e.next   #=> 2
797 *   p e.next   #=> 3
798 *   p e.next   #raises StopIteration
799 *
800 */
801
802static VALUE
803enumerator_peek(VALUE obj)
804{
805    VALUE vs = enumerator_peek_values(obj);
806    return ary2sv(vs, 1);
807}
808
809/*
810 * call-seq:
811 *   e.feed obj   -> nil
812 *
813 * Sets the value to be returned by the next yield inside +e+.
814 *
815 * If the value is not set, the yield returns nil.
816 *
817 * This value is cleared after being yielded.
818 *
819 *   o = Object.new
820 *   def o.each
821 *     x = yield         # (2) blocks
822 *     p x               # (5) => "foo"
823 *     x = yield         # (6) blocks
824 *     p x               # (8) => nil
825 *     x = yield         # (9) blocks
826 *     p x               # not reached w/o another e.next
827 *   end
828 *
829 *   e = o.to_enum
830 *   e.next              # (1)
831 *   e.feed "foo"        # (3)
832 *   e.next              # (4)
833 *   e.next              # (7)
834 *                       # (10)
835 */
836
837static VALUE
838enumerator_feed(VALUE obj, VALUE v)
839{
840    struct enumerator *e = enumerator_ptr(obj);
841
842    if (e->feedvalue != Qundef) {
843	rb_raise(rb_eTypeError, "feed value already set");
844    }
845    e->feedvalue = v;
846
847    return Qnil;
848}
849
850/*
851 * call-seq:
852 *   e.rewind   -> e
853 *
854 * Rewinds the enumeration sequence to the beginning.
855 *
856 * If the enclosed object responds to a "rewind" method, it is called.
857 */
858
859static VALUE
860enumerator_rewind(VALUE obj)
861{
862    struct enumerator *e = enumerator_ptr(obj);
863
864    rb_check_funcall(e->obj, id_rewind, 0, 0);
865
866    e->fib = 0;
867    e->dst = Qnil;
868    e->lookahead = Qundef;
869    e->feedvalue = Qundef;
870    e->stop_exc = Qfalse;
871    return obj;
872}
873
874static VALUE
875inspect_enumerator(VALUE obj, VALUE dummy, int recur)
876{
877    struct enumerator *e;
878    const char *cname;
879    VALUE eobj, eargs, str, method;
880    int tainted, untrusted;
881
882    TypedData_Get_Struct(obj, struct enumerator, &enumerator_data_type, e);
883
884    cname = rb_obj_classname(obj);
885
886    if (!e || e->obj == Qundef) {
887	return rb_sprintf("#<%s: uninitialized>", cname);
888    }
889
890    if (recur) {
891	str = rb_sprintf("#<%s: ...>", cname);
892	OBJ_TAINT(str);
893	return str;
894    }
895
896    eobj = rb_attr_get(obj, id_receiver);
897    if (NIL_P(eobj)) {
898	eobj = e->obj;
899    }
900
901    tainted   = OBJ_TAINTED(eobj);
902    untrusted = OBJ_UNTRUSTED(eobj);
903
904    /* (1..100).each_cons(2) => "#<Enumerator: 1..100:each_cons(2)>" */
905    str = rb_sprintf("#<%s: ", cname);
906    rb_str_concat(str, rb_inspect(eobj));
907    method = rb_attr_get(obj, id_method);
908    if (NIL_P(method)) {
909	rb_str_buf_cat2(str, ":");
910	rb_str_buf_cat2(str, rb_id2name(e->meth));
911    }
912    else if (method != Qfalse) {
913	Check_Type(method, T_SYMBOL);
914	rb_str_buf_cat2(str, ":");
915	rb_str_buf_cat2(str, rb_id2name(SYM2ID(method)));
916    }
917
918    eargs = rb_attr_get(obj, id_arguments);
919    if (NIL_P(eargs)) {
920	eargs = e->args;
921    }
922    if (eargs != Qfalse) {
923	long   argc = RARRAY_LEN(eargs);
924	VALUE *argv = RARRAY_PTR(eargs);
925
926	if (argc > 0) {
927	    rb_str_buf_cat2(str, "(");
928
929	    while (argc--) {
930		VALUE arg = *argv++;
931
932		rb_str_concat(str, rb_inspect(arg));
933		rb_str_buf_cat2(str, argc > 0 ? ", " : ")");
934
935		if (OBJ_TAINTED(arg)) tainted = TRUE;
936		if (OBJ_UNTRUSTED(arg)) untrusted = TRUE;
937	    }
938	}
939    }
940
941    rb_str_buf_cat2(str, ">");
942
943    if (tainted) OBJ_TAINT(str);
944    if (untrusted) OBJ_UNTRUST(str);
945    return str;
946}
947
948/*
949 * call-seq:
950 *   e.inspect  -> string
951 *
952 * Creates a printable version of <i>e</i>.
953 */
954
955static VALUE
956enumerator_inspect(VALUE obj)
957{
958    return rb_exec_recursive(inspect_enumerator, obj, 0);
959}
960
961/*
962 * call-seq:
963 *   e.size          -> int, Float::INFINITY or nil
964 *
965 * Returns the size of the enumerator, or +nil+ if it can't be calculated lazily.
966 *
967 *   (1..100).to_a.permutation(4).size # => 94109400
968 *   loop.size # => Float::INFINITY
969 *   (1..100).drop_while.size # => nil
970 */
971
972static VALUE
973enumerator_size(VALUE obj)
974{
975    struct enumerator *e = enumerator_ptr(obj);
976
977    if (e->size_fn) {
978	return (*e->size_fn)(e->obj, e->args, obj);
979    }
980    if (rb_obj_is_proc(e->size)) {
981        if (e->args)
982	    return rb_proc_call(e->size, e->args);
983        else
984            return rb_proc_call_with_block(e->size, 0, 0, Qnil);
985    }
986    return e->size;
987}
988
989/*
990 * Yielder
991 */
992static void
993yielder_mark(void *p)
994{
995    struct yielder *ptr = p;
996    rb_gc_mark(ptr->proc);
997}
998
999#define yielder_free RUBY_TYPED_DEFAULT_FREE
1000
1001static size_t
1002yielder_memsize(const void *p)
1003{
1004    return p ? sizeof(struct yielder) : 0;
1005}
1006
1007static const rb_data_type_t yielder_data_type = {
1008    "yielder",
1009    {
1010	yielder_mark,
1011	yielder_free,
1012	yielder_memsize,
1013    },
1014};
1015
1016static struct yielder *
1017yielder_ptr(VALUE obj)
1018{
1019    struct yielder *ptr;
1020
1021    TypedData_Get_Struct(obj, struct yielder, &yielder_data_type, ptr);
1022    if (!ptr || ptr->proc == Qundef) {
1023	rb_raise(rb_eArgError, "uninitialized yielder");
1024    }
1025    return ptr;
1026}
1027
1028/* :nodoc: */
1029static VALUE
1030yielder_allocate(VALUE klass)
1031{
1032    struct yielder *ptr;
1033    VALUE obj;
1034
1035    obj = TypedData_Make_Struct(klass, struct yielder, &yielder_data_type, ptr);
1036    ptr->proc = Qundef;
1037
1038    return obj;
1039}
1040
1041static VALUE
1042yielder_init(VALUE obj, VALUE proc)
1043{
1044    struct yielder *ptr;
1045
1046    TypedData_Get_Struct(obj, struct yielder, &yielder_data_type, ptr);
1047
1048    if (!ptr) {
1049	rb_raise(rb_eArgError, "unallocated yielder");
1050    }
1051
1052    ptr->proc = proc;
1053
1054    return obj;
1055}
1056
1057/* :nodoc: */
1058static VALUE
1059yielder_initialize(VALUE obj)
1060{
1061    rb_need_block();
1062
1063    return yielder_init(obj, rb_block_proc());
1064}
1065
1066/* :nodoc: */
1067static VALUE
1068yielder_yield(VALUE obj, VALUE args)
1069{
1070    struct yielder *ptr = yielder_ptr(obj);
1071
1072    return rb_proc_call(ptr->proc, args);
1073}
1074
1075/* :nodoc: */
1076static VALUE yielder_yield_push(VALUE obj, VALUE args)
1077{
1078    yielder_yield(obj, args);
1079    return obj;
1080}
1081
1082static VALUE
1083yielder_yield_i(VALUE obj, VALUE memo, int argc, VALUE *argv)
1084{
1085    return rb_yield_values2(argc, argv);
1086}
1087
1088static VALUE
1089yielder_new(void)
1090{
1091    return yielder_init(yielder_allocate(rb_cYielder), rb_proc_new(yielder_yield_i, 0));
1092}
1093
1094/*
1095 * Generator
1096 */
1097static void
1098generator_mark(void *p)
1099{
1100    struct generator *ptr = p;
1101    rb_gc_mark(ptr->proc);
1102}
1103
1104#define generator_free RUBY_TYPED_DEFAULT_FREE
1105
1106static size_t
1107generator_memsize(const void *p)
1108{
1109    return p ? sizeof(struct generator) : 0;
1110}
1111
1112static const rb_data_type_t generator_data_type = {
1113    "generator",
1114    {
1115	generator_mark,
1116	generator_free,
1117	generator_memsize,
1118    },
1119};
1120
1121static struct generator *
1122generator_ptr(VALUE obj)
1123{
1124    struct generator *ptr;
1125
1126    TypedData_Get_Struct(obj, struct generator, &generator_data_type, ptr);
1127    if (!ptr || ptr->proc == Qundef) {
1128	rb_raise(rb_eArgError, "uninitialized generator");
1129    }
1130    return ptr;
1131}
1132
1133/* :nodoc: */
1134static VALUE
1135generator_allocate(VALUE klass)
1136{
1137    struct generator *ptr;
1138    VALUE obj;
1139
1140    obj = TypedData_Make_Struct(klass, struct generator, &generator_data_type, ptr);
1141    ptr->proc = Qundef;
1142
1143    return obj;
1144}
1145
1146static VALUE
1147generator_init(VALUE obj, VALUE proc)
1148{
1149    struct generator *ptr;
1150
1151    TypedData_Get_Struct(obj, struct generator, &generator_data_type, ptr);
1152
1153    if (!ptr) {
1154	rb_raise(rb_eArgError, "unallocated generator");
1155    }
1156
1157    ptr->proc = proc;
1158
1159    return obj;
1160}
1161
1162/* :nodoc: */
1163static VALUE
1164generator_initialize(int argc, VALUE *argv, VALUE obj)
1165{
1166    VALUE proc;
1167
1168    if (argc == 0) {
1169	rb_need_block();
1170
1171	proc = rb_block_proc();
1172    }
1173    else {
1174	rb_scan_args(argc, argv, "1", &proc);
1175
1176	if (!rb_obj_is_proc(proc))
1177	    rb_raise(rb_eTypeError,
1178		     "wrong argument type %s (expected Proc)",
1179		     rb_obj_classname(proc));
1180
1181	if (rb_block_given_p()) {
1182	    rb_warn("given block not used");
1183	}
1184    }
1185
1186    return generator_init(obj, proc);
1187}
1188
1189/* :nodoc: */
1190static VALUE
1191generator_init_copy(VALUE obj, VALUE orig)
1192{
1193    struct generator *ptr0, *ptr1;
1194
1195    if (!OBJ_INIT_COPY(obj, orig)) return obj;
1196
1197    ptr0 = generator_ptr(orig);
1198
1199    TypedData_Get_Struct(obj, struct generator, &generator_data_type, ptr1);
1200
1201    if (!ptr1) {
1202	rb_raise(rb_eArgError, "unallocated generator");
1203    }
1204
1205    ptr1->proc = ptr0->proc;
1206
1207    return obj;
1208}
1209
1210/* :nodoc: */
1211static VALUE
1212generator_each(int argc, VALUE *argv, VALUE obj)
1213{
1214    struct generator *ptr = generator_ptr(obj);
1215    VALUE args = rb_ary_new2(argc + 1);
1216
1217    rb_ary_push(args, yielder_new());
1218    if (argc > 0) {
1219	rb_ary_cat(args, argv, argc);
1220    }
1221
1222    return rb_proc_call(ptr->proc, args);
1223}
1224
1225/* Lazy Enumerator methods */
1226static VALUE
1227enum_size(VALUE self)
1228{
1229    VALUE r = rb_check_funcall(self, id_size, 0, 0);
1230    return (r == Qundef) ? Qnil : r;
1231}
1232
1233static VALUE
1234lazy_size(VALUE self)
1235{
1236    return enum_size(rb_ivar_get(self, id_receiver));
1237}
1238
1239static VALUE
1240lazy_receiver_size(VALUE generator, VALUE args, VALUE lazy)
1241{
1242    return lazy_size(lazy);
1243}
1244
1245static VALUE
1246lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv)
1247{
1248    VALUE result;
1249    if (argc == 1) {
1250	VALUE args[2];
1251	args[0] = m;
1252	args[1] = val;
1253	result = rb_yield_values2(2, args);
1254    }
1255    else {
1256	VALUE args;
1257	int len = rb_long2int((long)argc + 1);
1258
1259	args = rb_ary_tmp_new(len);
1260	rb_ary_push(args, m);
1261	if (argc > 0) {
1262	    rb_ary_cat(args, argv, argc);
1263	}
1264	result = rb_yield_values2(len, RARRAY_PTR(args));
1265	RB_GC_GUARD(args);
1266    }
1267    if (result == Qundef) rb_iter_break();
1268    return Qnil;
1269}
1270
1271static VALUE
1272lazy_init_block_i(VALUE val, VALUE m, int argc, VALUE *argv)
1273{
1274    rb_block_call(m, id_each, argc-1, argv+1, lazy_init_iterator, val);
1275    return Qnil;
1276}
1277
1278/*
1279 * call-seq:
1280 *   Lazy.new(obj, size=nil) { |yielder, *values| ... }
1281 *
1282 * Creates a new Lazy enumerator. When the enumerator is actually enumerated
1283 * (e.g. by calling #force), +obj+ will be enumerated and each value passed
1284 * to the given block. The block can yield values back using +yielder+.
1285 * For example, to create a method +filter_map+ in both lazy and
1286 * non-lazy fashions:
1287 *
1288 *   module Enumerable
1289 *     def filter_map(&block)
1290 *       map(&block).compact
1291 *     end
1292 *   end
1293 *
1294 *   class Enumerator::Lazy
1295 *     def filter_map
1296 *       Lazy.new(self) do |yielder, *values|
1297 *         result = yield *values
1298 *         yielder << result if result
1299 *       end
1300 *     end
1301 *   end
1302 *
1303 *   (1..Float::INFINITY).lazy.filter_map{|i| i*i if i.even?}.first(5)
1304 *       # => [4, 16, 36, 64, 100]
1305 */
1306static VALUE
1307lazy_initialize(int argc, VALUE *argv, VALUE self)
1308{
1309    VALUE obj, size = Qnil;
1310    VALUE generator;
1311
1312    rb_check_arity(argc, 1, 2);
1313    if (!rb_block_given_p()) {
1314	rb_raise(rb_eArgError, "tried to call lazy new without a block");
1315    }
1316    obj = argv[0];
1317    if (argc > 1) {
1318	size = argv[1];
1319    }
1320    generator = generator_allocate(rb_cGenerator);
1321    rb_block_call(generator, id_initialize, 0, 0, lazy_init_block_i, obj);
1322    enumerator_init(self, generator, sym_each, 0, 0, 0, size);
1323    rb_ivar_set(self, id_receiver, obj);
1324
1325    return self;
1326}
1327
1328static VALUE
1329lazy_set_method(VALUE lazy, VALUE args, VALUE (*size_fn)(ANYARGS))
1330{
1331    ID id = rb_frame_this_func();
1332    struct enumerator *e = enumerator_ptr(lazy);
1333    rb_ivar_set(lazy, id_method, ID2SYM(id));
1334    if (NIL_P(args)) {
1335	/* Qfalse indicates that the arguments are empty */
1336	rb_ivar_set(lazy, id_arguments, Qfalse);
1337    }
1338    else {
1339	rb_ivar_set(lazy, id_arguments, args);
1340    }
1341    e->size_fn = size_fn;
1342    return lazy;
1343}
1344
1345/*
1346 * call-seq:
1347 *   e.lazy -> lazy_enumerator
1348 *
1349 * Returns a lazy enumerator, whose methods map/collect,
1350 * flat_map/collect_concat, select/find_all, reject, grep, zip, take,
1351 * take_while, drop, and drop_while enumerate values only on an
1352 * as-needed basis.  However, if a block is given to zip, values
1353 * are enumerated immediately.
1354 *
1355 * === Example
1356 *
1357 * The following program finds pythagorean triples:
1358 *
1359 *   def pythagorean_triples
1360 *     (1..Float::INFINITY).lazy.flat_map {|z|
1361 *       (1..z).flat_map {|x|
1362 *         (x..z).select {|y|
1363 *           x**2 + y**2 == z**2
1364 *         }.map {|y|
1365 *           [x, y, z]
1366 *         }
1367 *       }
1368 *     }
1369 *   end
1370 *   # show first ten pythagorean triples
1371 *   p pythagorean_triples.take(10).force # take is lazy, so force is needed
1372 *   p pythagorean_triples.first(10)      # first is eager
1373 *   # show pythagorean triples less than 100
1374 *   p pythagorean_triples.take_while { |*, z| z < 100 }.force
1375 */
1376static VALUE
1377enumerable_lazy(VALUE obj)
1378{
1379    VALUE result = lazy_to_enum_i(obj, sym_each, 0, 0, enum_size);
1380    /* Qfalse indicates that the Enumerator::Lazy has no method name */
1381    rb_ivar_set(result, id_method, Qfalse);
1382    return result;
1383}
1384
1385static VALUE
1386lazy_to_enum_i(VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS))
1387{
1388    return enumerator_init(enumerator_allocate(rb_cLazy),
1389	obj, meth, argc, argv, size_fn, Qnil);
1390}
1391
1392/*
1393 * call-seq:
1394 *   lzy.to_enum(method = :each, *args)                 -> lazy_enum
1395 *   lzy.enum_for(method = :each, *args)                -> lazy_enum
1396 *   lzy.to_enum(method = :each, *args) {|*args| block} -> lazy_enum
1397 *   lzy.enum_for(method = :each, *args){|*args| block} -> lazy_enum
1398 *
1399 * Similar to Kernel#to_enum, except it returns a lazy enumerator.
1400 * This makes it easy to define Enumerable methods that will
1401 * naturally remain lazy if called from a lazy enumerator.
1402 *
1403 * For example, continuing from the example in Kernel#to_enum:
1404 *
1405 *   # See Kernel#to_enum for the definition of repeat
1406 *   r = 1..Float::INFINITY
1407 *   r.repeat(2).first(5) # => [1, 1, 2, 2, 3]
1408 *   r.repeat(2).class # => Enumerator
1409 *   r.repeat(2).map{|n| n ** 2}.first(5) # => endless loop!
1410 *   # works naturally on lazy enumerator:
1411 *   r.lazy.repeat(2).class # => Enumerator::Lazy
1412 *   r.lazy.repeat(2).map{|n| n ** 2}.first(5) # => [1, 1, 4, 4, 9]
1413 */
1414
1415static VALUE
1416lazy_to_enum(int argc, VALUE *argv, VALUE self)
1417{
1418    VALUE lazy, meth = sym_each;
1419
1420    if (argc > 0) {
1421	--argc;
1422	meth = *argv++;
1423    }
1424    lazy = lazy_to_enum_i(self, meth, argc, argv, 0);
1425    if (rb_block_given_p()) {
1426	enumerator_ptr(lazy)->size = rb_block_proc();
1427    }
1428    return lazy;
1429}
1430
1431static VALUE
1432lazy_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
1433{
1434    VALUE result = rb_yield_values2(argc - 1, &argv[1]);
1435
1436    rb_funcall(argv[0], id_yield, 1, result);
1437    return Qnil;
1438}
1439
1440static VALUE
1441lazy_map(VALUE obj)
1442{
1443    if (!rb_block_given_p()) {
1444	rb_raise(rb_eArgError, "tried to call lazy map without a block");
1445    }
1446
1447    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1448					 lazy_map_func, 0),
1449			   Qnil, lazy_receiver_size);
1450}
1451
1452static VALUE
1453lazy_flat_map_i(VALUE i, VALUE yielder, int argc, VALUE *argv)
1454{
1455    return rb_funcall2(yielder, id_yield, argc, argv);
1456}
1457
1458static VALUE
1459lazy_flat_map_each(VALUE obj, VALUE yielder)
1460{
1461    rb_block_call(obj, id_each, 0, 0, lazy_flat_map_i, yielder);
1462    return Qnil;
1463}
1464
1465static VALUE
1466lazy_flat_map_to_ary(VALUE obj, VALUE yielder)
1467{
1468    VALUE ary = rb_check_array_type(obj);
1469    if (NIL_P(ary)) {
1470	rb_funcall(yielder, id_yield, 1, obj);
1471    }
1472    else {
1473	long i;
1474	for (i = 0; i < RARRAY_LEN(ary); i++) {
1475	    rb_funcall(yielder, id_yield, 1, RARRAY_PTR(ary)[i]);
1476	}
1477    }
1478    return Qnil;
1479}
1480
1481static VALUE
1482lazy_flat_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
1483{
1484    VALUE result = rb_yield_values2(argc - 1, &argv[1]);
1485    if (RB_TYPE_P(result, T_ARRAY)) {
1486	long i;
1487	for (i = 0; i < RARRAY_LEN(result); i++) {
1488	    rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(result)[i]);
1489	}
1490    }
1491    else {
1492	if (rb_respond_to(result, id_force) && rb_respond_to(result, id_each)) {
1493	    lazy_flat_map_each(result, argv[0]);
1494	}
1495	else {
1496	    lazy_flat_map_to_ary(result, argv[0]);
1497	}
1498    }
1499    return Qnil;
1500}
1501
1502/*
1503 *  call-seq:
1504 *     lazy.flat_map       { |obj| block } -> a_lazy_enumerator
1505 *
1506 *  Returns a new lazy enumerator with the concatenated results of running
1507 *  <i>block</i> once for every element in <i>lazy</i>.
1508 *
1509 *    ["foo", "bar"].lazy.flat_map {|i| i.each_char.lazy}.force
1510 *    #=> ["f", "o", "o", "b", "a", "r"]
1511 *
1512 *  A value <i>x</i> returned by <i>block</i> is decomposed if either of
1513 *  the following conditions is true:
1514 *
1515 *    a) <i>x</i> responds to both each and force, which means that
1516 *       <i>x</i> is a lazy enumerator.
1517 *    b) <i>x</i> is an array or responds to to_ary.
1518 *
1519 *  Otherwise, <i>x</i> is contained as-is in the return value.
1520 *
1521 *    [{a:1}, {b:2}].lazy.flat_map {|i| i}.force
1522 *    #=> [{:a=>1}, {:b=>2}]
1523 */
1524static VALUE
1525lazy_flat_map(VALUE obj)
1526{
1527    if (!rb_block_given_p()) {
1528	rb_raise(rb_eArgError, "tried to call lazy flat_map without a block");
1529    }
1530
1531    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1532					 lazy_flat_map_func, 0),
1533			   Qnil, 0);
1534}
1535
1536static VALUE
1537lazy_select_func(VALUE val, VALUE m, int argc, VALUE *argv)
1538{
1539    VALUE element = rb_enum_values_pack(argc - 1, argv + 1);
1540
1541    if (RTEST(rb_yield(element))) {
1542	return rb_funcall(argv[0], id_yield, 1, element);
1543    }
1544    return Qnil;
1545}
1546
1547static VALUE
1548lazy_select(VALUE obj)
1549{
1550    if (!rb_block_given_p()) {
1551	rb_raise(rb_eArgError, "tried to call lazy select without a block");
1552    }
1553
1554    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1555					 lazy_select_func, 0),
1556			   Qnil, 0);
1557}
1558
1559static VALUE
1560lazy_reject_func(VALUE val, VALUE m, int argc, VALUE *argv)
1561{
1562    VALUE element = rb_enum_values_pack(argc - 1, argv + 1);
1563
1564    if (!RTEST(rb_yield(element))) {
1565	return rb_funcall(argv[0], id_yield, 1, element);
1566    }
1567    return Qnil;
1568}
1569
1570static VALUE
1571lazy_reject(VALUE obj)
1572{
1573    if (!rb_block_given_p()) {
1574	rb_raise(rb_eArgError, "tried to call lazy reject without a block");
1575    }
1576
1577    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1578					 lazy_reject_func, 0),
1579			   Qnil, 0);
1580}
1581
1582static VALUE
1583lazy_grep_func(VALUE val, VALUE m, int argc, VALUE *argv)
1584{
1585    VALUE i = rb_enum_values_pack(argc - 1, argv + 1);
1586    VALUE result = rb_funcall(m, id_eqq, 1, i);
1587
1588    if (RTEST(result)) {
1589	rb_funcall(argv[0], id_yield, 1, i);
1590    }
1591    return Qnil;
1592}
1593
1594static VALUE
1595lazy_grep_iter(VALUE val, VALUE m, int argc, VALUE *argv)
1596{
1597    VALUE i = rb_enum_values_pack(argc - 1, argv + 1);
1598    VALUE result = rb_funcall(m, id_eqq, 1, i);
1599
1600    if (RTEST(result)) {
1601	rb_funcall(argv[0], id_yield, 1, rb_yield(i));
1602    }
1603    return Qnil;
1604}
1605
1606static VALUE
1607lazy_grep(VALUE obj, VALUE pattern)
1608{
1609    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1610					 rb_block_given_p() ?
1611					 lazy_grep_iter : lazy_grep_func,
1612					 pattern),
1613			   rb_ary_new3(1, pattern), 0);
1614}
1615
1616static VALUE
1617call_next(VALUE obj)
1618{
1619    return rb_funcall(obj, id_next, 0);
1620}
1621
1622static VALUE
1623next_stopped(VALUE obj)
1624{
1625    return Qnil;
1626}
1627
1628static VALUE
1629lazy_zip_arrays_func(VALUE val, VALUE arrays, int argc, VALUE *argv)
1630{
1631    VALUE yielder, ary, memo;
1632    long i, count;
1633
1634    yielder = argv[0];
1635    memo = rb_attr_get(yielder, id_memo);
1636    count = NIL_P(memo) ? 0 : NUM2LONG(memo);
1637
1638    ary = rb_ary_new2(RARRAY_LEN(arrays) + 1);
1639    rb_ary_push(ary, argv[1]);
1640    for (i = 0; i < RARRAY_LEN(arrays); i++) {
1641	rb_ary_push(ary, rb_ary_entry(RARRAY_PTR(arrays)[i], count));
1642    }
1643    rb_funcall(yielder, id_yield, 1, ary);
1644    rb_ivar_set(yielder, id_memo, LONG2NUM(++count));
1645    return Qnil;
1646}
1647
1648static VALUE
1649lazy_zip_func(VALUE val, VALUE zip_args, int argc, VALUE *argv)
1650{
1651    VALUE yielder, ary, arg, v;
1652    long i;
1653
1654    yielder = argv[0];
1655    arg = rb_attr_get(yielder, id_memo);
1656    if (NIL_P(arg)) {
1657	arg = rb_ary_new2(RARRAY_LEN(zip_args));
1658	for (i = 0; i < RARRAY_LEN(zip_args); i++) {
1659	    rb_ary_push(arg, rb_funcall(RARRAY_PTR(zip_args)[i], id_to_enum, 0));
1660	}
1661	rb_ivar_set(yielder, id_memo, arg);
1662    }
1663
1664    ary = rb_ary_new2(RARRAY_LEN(arg) + 1);
1665    v = Qnil;
1666    if (--argc > 0) {
1667	++argv;
1668	v = argc > 1 ? rb_ary_new4(argc, argv) : *argv;
1669    }
1670    rb_ary_push(ary, v);
1671    for (i = 0; i < RARRAY_LEN(arg); i++) {
1672	v = rb_rescue2(call_next, RARRAY_PTR(arg)[i], next_stopped, 0,
1673		       rb_eStopIteration, (VALUE)0);
1674	rb_ary_push(ary, v);
1675    }
1676    rb_funcall(yielder, id_yield, 1, ary);
1677    return Qnil;
1678}
1679
1680static VALUE
1681lazy_zip(int argc, VALUE *argv, VALUE obj)
1682{
1683    VALUE ary, v;
1684    long i;
1685    rb_block_call_func *func = lazy_zip_arrays_func;
1686
1687    if (rb_block_given_p()) {
1688	return rb_call_super(argc, argv);
1689    }
1690
1691    ary = rb_ary_new2(argc);
1692    for (i = 0; i < argc; i++) {
1693	v = rb_check_array_type(argv[i]);
1694	if (NIL_P(v)) {
1695	    for (; i < argc; i++) {
1696		if (!rb_respond_to(argv[i], id_each)) {
1697		    rb_raise(rb_eTypeError, "wrong argument type %s (must respond to :each)",
1698			rb_obj_classname(argv[i]));
1699		}
1700	    }
1701	    ary = rb_ary_new4(argc, argv);
1702	    func = lazy_zip_func;
1703	    break;
1704	}
1705	rb_ary_push(ary, v);
1706    }
1707
1708    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1709					 func, ary),
1710			   ary, lazy_receiver_size);
1711}
1712
1713static VALUE
1714lazy_take_func(VALUE val, VALUE args, int argc, VALUE *argv)
1715{
1716    long remain;
1717    VALUE memo = rb_attr_get(argv[0], id_memo);
1718    if (NIL_P(memo)) {
1719	memo = args;
1720    }
1721
1722    rb_funcall2(argv[0], id_yield, argc - 1, argv + 1);
1723    if ((remain = NUM2LONG(memo)-1) == 0) {
1724	return Qundef;
1725    }
1726    else {
1727	rb_ivar_set(argv[0], id_memo, LONG2NUM(remain));
1728	return Qnil;
1729    }
1730}
1731
1732static VALUE
1733lazy_take_size(VALUE generator, VALUE args, VALUE lazy)
1734{
1735    VALUE receiver = lazy_size(lazy);
1736    long len = NUM2LONG(RARRAY_PTR(rb_ivar_get(lazy, id_arguments))[0]);
1737    if (NIL_P(receiver) || (FIXNUM_P(receiver) && FIX2LONG(receiver) < len))
1738	return receiver;
1739    return LONG2NUM(len);
1740}
1741
1742static VALUE
1743lazy_take(VALUE obj, VALUE n)
1744{
1745    long len = NUM2LONG(n);
1746    VALUE lazy;
1747
1748    if (len < 0) {
1749	rb_raise(rb_eArgError, "attempt to take negative size");
1750    }
1751    if (len == 0) {
1752	VALUE len = INT2NUM(0);
1753	lazy = lazy_to_enum_i(obj, sym_cycle, 1, &len, 0);
1754    }
1755    else {
1756	lazy = rb_block_call(rb_cLazy, id_new, 1, &obj,
1757					 lazy_take_func, n);
1758    }
1759    return lazy_set_method(lazy, rb_ary_new3(1, n), lazy_take_size);
1760}
1761
1762static VALUE
1763lazy_take_while_func(VALUE val, VALUE args, int argc, VALUE *argv)
1764{
1765    VALUE result = rb_yield_values2(argc - 1, &argv[1]);
1766    if (!RTEST(result)) return Qundef;
1767    rb_funcall2(argv[0], id_yield, argc - 1, argv + 1);
1768    return Qnil;
1769}
1770
1771static VALUE
1772lazy_take_while(VALUE obj)
1773{
1774    if (!rb_block_given_p()) {
1775	rb_raise(rb_eArgError, "tried to call lazy take_while without a block");
1776    }
1777    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1778					 lazy_take_while_func, 0),
1779			   Qnil, 0);
1780}
1781
1782static VALUE
1783lazy_drop_size(VALUE generator, VALUE args, VALUE lazy)
1784{
1785    long len = NUM2LONG(RARRAY_PTR(rb_ivar_get(lazy, id_arguments))[0]);
1786    VALUE receiver = lazy_size(lazy);
1787    if (NIL_P(receiver))
1788	return receiver;
1789    if (FIXNUM_P(receiver)) {
1790	len = FIX2LONG(receiver) - len;
1791	return LONG2FIX(len < 0 ? 0 : len);
1792    }
1793    return rb_funcall(receiver, '-', 1, LONG2NUM(len));
1794}
1795
1796static VALUE
1797lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv)
1798{
1799    long remain;
1800    VALUE memo = rb_attr_get(argv[0], id_memo);
1801    if (NIL_P(memo)) {
1802	memo = args;
1803    }
1804    if ((remain = NUM2LONG(memo)) == 0) {
1805	rb_funcall2(argv[0], id_yield, argc - 1, argv + 1);
1806    }
1807    else {
1808	rb_ivar_set(argv[0], id_memo, LONG2NUM(--remain));
1809    }
1810    return Qnil;
1811}
1812
1813static VALUE
1814lazy_drop(VALUE obj, VALUE n)
1815{
1816    long len = NUM2LONG(n);
1817
1818    if (len < 0) {
1819	rb_raise(rb_eArgError, "attempt to drop negative size");
1820    }
1821    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1822					 lazy_drop_func, n),
1823			   rb_ary_new3(1, n), lazy_drop_size);
1824}
1825
1826static VALUE
1827lazy_drop_while_func(VALUE val, VALUE args, int argc, VALUE *argv)
1828{
1829    VALUE memo = rb_attr_get(argv[0], id_memo);
1830    if (NIL_P(memo) && !RTEST(rb_yield_values2(argc - 1, &argv[1]))) {
1831	rb_ivar_set(argv[0], id_memo, memo = Qtrue);
1832    }
1833    if (memo == Qtrue) {
1834	rb_funcall2(argv[0], id_yield, argc - 1, argv + 1);
1835    }
1836    return Qnil;
1837}
1838
1839static VALUE
1840lazy_drop_while(VALUE obj)
1841{
1842    if (!rb_block_given_p()) {
1843	rb_raise(rb_eArgError, "tried to call lazy drop_while without a block");
1844    }
1845    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
1846					 lazy_drop_while_func, 0),
1847			   Qnil, 0);
1848}
1849
1850static VALUE
1851lazy_super(int argc, VALUE *argv, VALUE lazy)
1852{
1853    return enumerable_lazy(rb_call_super(argc, argv));
1854}
1855
1856static VALUE
1857lazy_lazy(VALUE obj)
1858{
1859    return obj;
1860}
1861
1862/*
1863 * Document-class: StopIteration
1864 *
1865 * Raised to stop the iteration, in particular by Enumerator#next. It is
1866 * rescued by Kernel#loop.
1867 *
1868 *   loop do
1869 *     puts "Hello"
1870 *     raise StopIteration
1871 *     puts "World"
1872 *   end
1873 *   puts "Done!"
1874 *
1875 * <em>produces:</em>
1876 *
1877 *   Hello
1878 *   Done!
1879 */
1880
1881/*
1882 * call-seq:
1883 *   result       -> value
1884 *
1885 * Returns the return value of the iterator.
1886 *
1887 *   o = Object.new
1888 *   def o.each
1889 *     yield 1
1890 *     yield 2
1891 *     yield 3
1892 *     100
1893 *   end
1894 *
1895 *   e = o.to_enum
1896 *
1897 *   puts e.next                   #=> 1
1898 *   puts e.next                   #=> 2
1899 *   puts e.next                   #=> 3
1900 *
1901 *   begin
1902 *     e.next
1903 *   rescue StopIteration => ex
1904 *     puts ex.result              #=> 100
1905 *   end
1906 *
1907 */
1908
1909static VALUE
1910stop_result(VALUE self)
1911{
1912    return rb_attr_get(self, id_result);
1913}
1914
1915void
1916InitVM_Enumerator(void)
1917{
1918    rb_define_method(rb_mKernel, "to_enum", obj_to_enum, -1);
1919    rb_define_method(rb_mKernel, "enum_for", obj_to_enum, -1);
1920
1921    rb_cEnumerator = rb_define_class("Enumerator", rb_cObject);
1922    rb_include_module(rb_cEnumerator, rb_mEnumerable);
1923
1924    rb_define_alloc_func(rb_cEnumerator, enumerator_allocate);
1925    rb_define_method(rb_cEnumerator, "initialize", enumerator_initialize, -1);
1926    rb_define_method(rb_cEnumerator, "initialize_copy", enumerator_init_copy, 1);
1927    rb_define_method(rb_cEnumerator, "each", enumerator_each, -1);
1928    rb_define_method(rb_cEnumerator, "each_with_index", enumerator_each_with_index, 0);
1929    rb_define_method(rb_cEnumerator, "each_with_object", enumerator_with_object, 1);
1930    rb_define_method(rb_cEnumerator, "with_index", enumerator_with_index, -1);
1931    rb_define_method(rb_cEnumerator, "with_object", enumerator_with_object, 1);
1932    rb_define_method(rb_cEnumerator, "next_values", enumerator_next_values, 0);
1933    rb_define_method(rb_cEnumerator, "peek_values", enumerator_peek_values_m, 0);
1934    rb_define_method(rb_cEnumerator, "next", enumerator_next, 0);
1935    rb_define_method(rb_cEnumerator, "peek", enumerator_peek, 0);
1936    rb_define_method(rb_cEnumerator, "feed", enumerator_feed, 1);
1937    rb_define_method(rb_cEnumerator, "rewind", enumerator_rewind, 0);
1938    rb_define_method(rb_cEnumerator, "inspect", enumerator_inspect, 0);
1939    rb_define_method(rb_cEnumerator, "size", enumerator_size, 0);
1940
1941    /* Lazy */
1942    rb_cLazy = rb_define_class_under(rb_cEnumerator, "Lazy", rb_cEnumerator);
1943    rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0);
1944    rb_define_method(rb_cLazy, "initialize", lazy_initialize, -1);
1945    rb_define_method(rb_cLazy, "to_enum", lazy_to_enum, -1);
1946    rb_define_method(rb_cLazy, "enum_for", lazy_to_enum, -1);
1947    rb_define_method(rb_cLazy, "map", lazy_map, 0);
1948    rb_define_method(rb_cLazy, "collect", lazy_map, 0);
1949    rb_define_method(rb_cLazy, "flat_map", lazy_flat_map, 0);
1950    rb_define_method(rb_cLazy, "collect_concat", lazy_flat_map, 0);
1951    rb_define_method(rb_cLazy, "select", lazy_select, 0);
1952    rb_define_method(rb_cLazy, "find_all", lazy_select, 0);
1953    rb_define_method(rb_cLazy, "reject", lazy_reject, 0);
1954    rb_define_method(rb_cLazy, "grep", lazy_grep, 1);
1955    rb_define_method(rb_cLazy, "zip", lazy_zip, -1);
1956    rb_define_method(rb_cLazy, "take", lazy_take, 1);
1957    rb_define_method(rb_cLazy, "take_while", lazy_take_while, 0);
1958    rb_define_method(rb_cLazy, "drop", lazy_drop, 1);
1959    rb_define_method(rb_cLazy, "drop_while", lazy_drop_while, 0);
1960    rb_define_method(rb_cLazy, "lazy", lazy_lazy, 0);
1961    rb_define_method(rb_cLazy, "chunk", lazy_super, -1);
1962    rb_define_method(rb_cLazy, "slice_before", lazy_super, -1);
1963
1964    rb_define_alias(rb_cLazy, "force", "to_a");
1965
1966    rb_eStopIteration = rb_define_class("StopIteration", rb_eIndexError);
1967    rb_define_method(rb_eStopIteration, "result", stop_result, 0);
1968
1969    /* Generator */
1970    rb_cGenerator = rb_define_class_under(rb_cEnumerator, "Generator", rb_cObject);
1971    rb_include_module(rb_cGenerator, rb_mEnumerable);
1972    rb_define_alloc_func(rb_cGenerator, generator_allocate);
1973    rb_define_method(rb_cGenerator, "initialize", generator_initialize, -1);
1974    rb_define_method(rb_cGenerator, "initialize_copy", generator_init_copy, 1);
1975    rb_define_method(rb_cGenerator, "each", generator_each, -1);
1976
1977    /* Yielder */
1978    rb_cYielder = rb_define_class_under(rb_cEnumerator, "Yielder", rb_cObject);
1979    rb_define_alloc_func(rb_cYielder, yielder_allocate);
1980    rb_define_method(rb_cYielder, "initialize", yielder_initialize, 0);
1981    rb_define_method(rb_cYielder, "yield", yielder_yield, -2);
1982    rb_define_method(rb_cYielder, "<<", yielder_yield_push, -2);
1983
1984    rb_provide("enumerator.so");	/* for backward compatibility */
1985}
1986
1987void
1988Init_Enumerator(void)
1989{
1990    id_rewind = rb_intern("rewind");
1991    id_each = rb_intern("each");
1992    id_call = rb_intern("call");
1993    id_size = rb_intern("size");
1994    id_yield = rb_intern("yield");
1995    id_new = rb_intern("new");
1996    id_initialize = rb_intern("initialize");
1997    id_next = rb_intern("next");
1998    id_result = rb_intern("result");
1999    id_lazy = rb_intern("lazy");
2000    id_eqq = rb_intern("===");
2001    id_receiver = rb_intern("receiver");
2002    id_arguments = rb_intern("arguments");
2003    id_memo = rb_intern("memo");
2004    id_method = rb_intern("method");
2005    id_force = rb_intern("force");
2006    id_to_enum = rb_intern("to_enum");
2007    sym_each = ID2SYM(id_each);
2008    sym_cycle = ID2SYM(rb_intern("cycle"));
2009
2010    InitVM(Enumerator);
2011}
2012