* enumerator.c: Finalize and document Lazy.new. [Bug #7248]

  Add Lazy#to_enum and simplify Lazy#size.

* test/ruby/test_lazy_enumerator.rb: tests for above

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@39057 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

1 files changed, 112 insertions, 84 deletions

diff --git a/enumerator.c b/enumerator.c
index c82608d796..682e9a71cb 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -457,6 +457,9 @@ rb_enumeratorize(VALUE obj, VALUE meth, int argc, VALUE *argv)
     return rb_enumeratorize_with_size(obj, meth, argc, argv, 0);
 }
 
+static VALUE
+lazy_to_enum_i(VALUE self, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS));
+
 VALUE
 rb_enumeratorize_with_size(VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS))
 {
@@ -1022,7 +1025,7 @@ enumerator_size(VALUE obj)
     struct enumerator *e = enumerator_ptr(obj);
 
     if (e->size_fn) {
-	return (*e->size_fn)(e->obj, e->args);
+	return (*e->size_fn)(e->obj, e->args, obj);
     }
     if (rb_obj_is_proc(e->size)) {
         if (e->args)
@@ -1271,20 +1274,22 @@ generator_each(int argc, VALUE *argv, VALUE obj)
 
 /* Lazy Enumerator methods */
 static VALUE
-lazy_receiver_size(VALUE self)
+enum_size(VALUE self)
 {
-    VALUE r = rb_check_funcall(rb_ivar_get(self, id_receiver), id_size, 0, 0);
+    VALUE r = rb_check_funcall(self, id_size, 0, 0);
     return (r == Qundef) ? Qnil : r;
 }
 
 static VALUE
 lazy_size(VALUE self)
 {
-    struct enumerator *e = enumerator_ptr(self);
-    if (e->size_fn) {
-	return (*e->size_fn)(self);
-    }
-    return Qnil;
+    return enum_size(rb_ivar_get(self, id_receiver));
+}
+
+static VALUE
+lazy_receiver_size(VALUE generator, VALUE args, VALUE lazy)
+{
+    return lazy_size(lazy);
 }
 
 static VALUE
@@ -1314,54 +1319,57 @@ lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv)
 }
 
 static VALUE
-lazy_init_yielder(VALUE val, VALUE m, int argc, VALUE *argv)
-{
-    VALUE result;
-    result = rb_funcall2(m, id_yield, argc, argv);
-    if (result == Qundef) rb_iter_break();
-    return Qnil;
-}
-
-static VALUE
 lazy_init_block_i(VALUE val, VALUE m, int argc, VALUE *argv)
 {
     rb_block_call(m, id_each, argc-1, argv+1, lazy_init_iterator, val);
     return Qnil;
 }
 
-static VALUE
-lazy_init_block(VALUE val, VALUE m, int argc, VALUE *argv)
-{
-    rb_block_call(m, id_each, argc-1, argv+1, lazy_init_yielder, val);
-    return Qnil;
-}
-
+/*
+ * call-seq:
+ *   Lazy.new(obj, size=nil) { |yielder, *values| ... }
+ *
+ * Creates a new Lazy enumerator. When the enumerator is actually enumerated
+ * (e.g. by calling #force), +obj+ will be enumerated and each value passed
+ * to the given block. The block can yield values back using +yielder+.
+ * For example, to create a method +filter_map+ in both lazy and
+ * non-lazy fashions:
+ *
+ *   module Enumerable
+ *     def filter_map(&block)
+ *       map(&block).compact
+ *     end
+ *   end
+ *
+ *   class Enumerator::Lazy
+ *     def filter_map
+ *       Lazy.new(self) do |yielder, *values|
+ *         result = yield *values
+ *         yielder << result if result
+ *       end
+ *     end
+ *   end
+ *
+ *   (1..Float::INFINITY).lazy.filter_map{|i| i*i if i.even?}.first(5)
+ *       # => [4, 16, 36, 64, 100]
+ */
 static VALUE
 lazy_initialize(int argc, VALUE *argv, VALUE self)
 {
-    VALUE obj, meth;
+    VALUE obj, size = Qnil;
     VALUE generator;
-    int offset;
 
-    if (argc < 1) {
-	rb_raise(rb_eArgError, "wrong number of arguments (%d for 1..)", argc);
+    rb_check_arity(argc, 1, 2);
+    if (!rb_block_given_p()) {
+	rb_raise(rb_eArgError, "tried to call lazy new without a block");
     }
-    else {
-	obj = argv[0];
-	if (argc == 1) {
-	    meth = sym_each;
-	    offset = 1;
-	}
-	else {
-	    meth = argv[1];
-	    offset = 2;
-	}
+    obj = argv[0];
+    if (argc > 1) {
+	size = argv[1];
     }
     generator = generator_allocate(rb_cGenerator);
-    rb_block_call(generator, id_initialize, 0, 0,
-		  (rb_block_given_p() ? lazy_init_block_i : lazy_init_block),
-		  obj);
-    enumerator_init(self, generator, meth, argc - offset, argv + offset, lazy_receiver_size, Qnil);
+    rb_block_call(generator, id_initialize, 0, 0, lazy_init_block_i, obj);
+    enumerator_init(self, generator, sym_each, 0, 0, 0, size);
     rb_ivar_set(self, id_receiver, obj);
 
     return self;
@@ -1418,15 +1426,59 @@ lazy_set_method(VALUE lazy, VALUE args, VALUE (*size_fn)(ANYARGS))
 static VALUE
 enumerable_lazy(VALUE obj)
 {
-    VALUE result;
-
-    result = rb_class_new_instance(1, &obj, rb_cLazy);
+    VALUE result = lazy_to_enum_i(obj, sym_each, 0, 0, enum_size);
     /* Qfalse indicates that the Enumerator::Lazy has no method name */
     rb_ivar_set(result, id_method, Qfalse);
     return result;
 }
 
 static VALUE
+lazy_to_enum_i(VALUE obj, VALUE meth, int argc, VALUE *argv, VALUE (*size_fn)(ANYARGS))
+{
+    return enumerator_init(enumerator_allocate(rb_cLazy),
+	obj, meth, argc, argv, size_fn, Qnil);
+}
+
+/*
+ * call-seq:
+ *   lzy.to_enum(method = :each, *args)                 -> lazy_enum
+ *   lzy.enum_for(method = :each, *args)                -> lazy_enum
+ *   lzy.to_enum(method = :each, *args) {|*args| block} -> lazy_enum
+ *   lzy.enum_for(method = :each, *args){|*args| block} -> lazy_enum
+ *
+ * Similar to Kernel#to_enum, except it returns a lazy enumerator.
+ * This makes it easy to define Enumerable methods that will
+ * naturally remain lazy if called from a lazy enumerator.
+ *
+ * For example, continuing from the example in Kernel#to_enum:
+ *
+ *   # See Kernel#to_enum for the definition of repeat
+ *   r = 1..Float::INFINITY
+ *   r.repeat(2).first(5) # => [1, 1, 2, 2, 3]
+ *   r.repeat(2).class # => Enumerator
+ *   r.repeat(2).map{|n| n ** 2}.first(5) # => endless loop!
+ *   # works naturally on lazy enumerator:
+ *   r.lazy.repeat(2).class # => Enumerator::Lazy
+ *   r.lazy.repeat(2).map{|n| n ** 2}.first(5) # => [1, 1, 4, 4, 9]
+ */
+
+static VALUE
+lazy_to_enum(int argc, VALUE *argv, VALUE self)
+{
+    VALUE lazy, meth = sym_each;
+
+    if (argc > 0) {
+	--argc;
+	meth = *argv++;
+    }
+    lazy = lazy_to_enum_i(self, meth, argc, argv, 0);
+    if (rb_block_given_p()) {
+	enumerator_ptr(lazy)->size = rb_block_proc();
+    }
+    return lazy;
+}
+
+static VALUE
 lazy_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
     VALUE result = rb_yield_values2(argc - 1, &argv[1]);
@@ -1723,10 +1775,10 @@ lazy_take_func(VALUE val, VALUE args, int argc, VALUE *argv)
 }
 
 static VALUE
-lazy_take_size(VALUE lazy)
+lazy_take_size(VALUE generator, VALUE args, VALUE lazy)
 {
+    VALUE receiver = lazy_size(lazy);
     long len = NUM2LONG(RARRAY_PTR(rb_ivar_get(lazy, id_arguments))[0]);
-    VALUE receiver = lazy_receiver_size(lazy);
     if (NIL_P(receiver) || (FIXNUM_P(receiver) && FIX2LONG(receiver) < len))
 	return receiver;
     return LONG2NUM(len);
@@ -1736,21 +1788,20 @@ static VALUE
 lazy_take(VALUE obj, VALUE n)
 {
     long len = NUM2LONG(n);
-    int argc = 1;
-    VALUE argv[3];
+    VALUE lazy;
 
     if (len < 0) {
 	rb_raise(rb_eArgError, "attempt to take negative size");
     }
-    argv[0] = obj;
     if (len == 0) {
-	argv[1] = sym_cycle;
-	argv[2] = INT2NUM(0);
-	argc = 3;
+	VALUE len = INT2NUM(0);
+	lazy = lazy_to_enum_i(obj, sym_cycle, 1, &len, 0);
+    }
+    else {
+	lazy = rb_block_call(rb_cLazy, id_new, 1, &obj,
+					 lazy_take_func, n);
     }
-    return lazy_set_method(rb_block_call(rb_cLazy, id_new, argc, argv,
-					 lazy_take_func, n),
-			   rb_ary_new3(1, n), lazy_take_size);
+    return lazy_set_method(lazy, rb_ary_new3(1, n), lazy_take_size);
 }
 
 static VALUE
@@ -1774,10 +1825,10 @@ lazy_take_while(VALUE obj)
 }
 
 static VALUE
-lazy_drop_size(VALUE lazy)
+lazy_drop_size(VALUE generator, VALUE args, VALUE lazy)
 {
     long len = NUM2LONG(RARRAY_PTR(rb_ivar_get(lazy, id_arguments))[0]);
-    VALUE receiver = lazy_receiver_size(lazy);
+    VALUE receiver = lazy_size(lazy);
     if (NIL_P(receiver))
 	return receiver;
     if (FIXNUM_P(receiver)) {
@@ -1842,36 +1893,12 @@ lazy_drop_while(VALUE obj)
 }
 
 static VALUE
-lazy_cycle_size(VALUE lazy)
-{
-    return rb_enum_cycle_size(rb_ivar_get(lazy, id_receiver), rb_ivar_get(lazy, id_arguments));
-}
-
-static VALUE
-lazy_cycle_func(VALUE val, VALUE m, int argc, VALUE *argv)
-{
-    return rb_funcall2(argv[0], id_yield, argc - 1, argv + 1);
-}
-
-static VALUE
 lazy_cycle(int argc, VALUE *argv, VALUE obj)
 {
-    VALUE args;
-    int len = rb_long2int((long)argc + 2);
-
     if (rb_block_given_p()) {
 	return rb_call_super(argc, argv);
     }
-    args = rb_ary_tmp_new(len);
-    rb_ary_push(args, obj);
-    rb_ary_push(args, sym_cycle);
-    if (argc > 0) {
-	rb_ary_cat(args, argv, argc);
-    }
-    return lazy_set_method(rb_block_call(rb_cLazy, id_new, len,
-					 RARRAY_PTR(args), lazy_cycle_func,
-					 args /* prevent from GC */),
-			   rb_ary_new4(argc, argv), lazy_cycle_size);
+    return lazy_to_enum_i(obj, sym_cycle, argc, argv, rb_enum_cycle_size);
 }
 
 static VALUE
@@ -1963,12 +1990,13 @@ InitVM_Enumerator(void)
     rb_cLazy = rb_define_class_under(rb_cEnumerator, "Lazy", rb_cEnumerator);
     rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0);
     rb_define_method(rb_cLazy, "initialize", lazy_initialize, -1);
+    rb_define_method(rb_cLazy, "to_enum", lazy_to_enum, -1);
+    rb_define_method(rb_cLazy, "enum_for", lazy_to_enum, -1);
     rb_define_method(rb_cLazy, "map", lazy_map, 0);
     rb_define_method(rb_cLazy, "collect", lazy_map, 0);
     rb_define_method(rb_cLazy, "flat_map", lazy_flat_map, 0);
     rb_define_method(rb_cLazy, "collect_concat", lazy_flat_map, 0);
     rb_define_method(rb_cLazy, "select", lazy_select, 0);
-    rb_define_method(rb_cLazy, "size", lazy_size, 0);
     rb_define_method(rb_cLazy, "find_all", lazy_select, 0);
     rb_define_method(rb_cLazy, "reject", lazy_reject, 0);
     rb_define_method(rb_cLazy, "grep", lazy_grep, 1);