Feature 4890 lazy enumerator

enumerator.c

@@ -101,7 +101,8 @@
  *
  */
 VALUE rb_cEnumerator;
-static ID id_rewind, id_each;
+VALUE rb_cLazy;
+static ID id_rewind, id_each, id_new, id_initialize, id_yield, id_call;
 static VALUE sym_each;
 
 VALUE rb_eStopIteration;
@@ -1183,6 +1184,125 @@ generator_each(VALUE obj)
  *   end
  *
  */
+
+/* Lazy Enumerator methods */
+static VALUE
+lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    if (rb_block_given_p()) {
+        return rb_yield(rb_ary_new3(2, m, val));
+    } else {
+        return rb_funcall(m, id_yield, 1, val);
+    }
+}
+
+static VALUE
+lazy_init_block(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    return rb_block_call(m, id_each, 0, 0, lazy_init_iterator, val);
+}
+
+static VALUE
+lazy_initialize(int argc, VALUE *argv, VALUE obj)
+{
+    VALUE generator = generator_allocate(rb_cGenerator);
+
+    rb_block_call(generator, id_initialize, 0, 0, lazy_init_block, argv[0]);
+    enumerator_init(obj, generator, sym_each, 0, 0);
+
+    return obj;
+}
+
+static VALUE
+enumerable_lazy(VALUE obj)
+{
+    return rb_funcall(rb_cLazy, id_new, 1, obj);
+}
+
+static VALUE
+lazy_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    VALUE result = rb_funcall(rb_block_proc(), id_call, 1, rb_ary_entry(val, 1));
+
+    return rb_funcall(rb_ary_entry(val, 0), id_yield, 1, result);
+}
+
+static VALUE
+lazy_map(VALUE obj)
+{
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "tried to call lazy map without a block");
+    }
+
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_map_func, 0);
+}
+
+
+static VALUE
+lazy_select_func(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    VALUE element = rb_ary_entry(val, 1);
+    VALUE result = rb_funcall(rb_block_proc(), id_call, 1, element);
+
+    if (RTEST(result)) {
+        return rb_funcall(rb_ary_entry(val, 0), id_yield, 1, element);
+    } else {
+        return result;
+    }
+}
+
+static VALUE
+lazy_select(VALUE obj)
+{
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "tried to call lazy select without a block");
+    }
+
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_select_func, 0);
+}
+
+static VALUE
+lazy_reject_func(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    VALUE element = rb_ary_entry(val, 1);
+    VALUE result = rb_funcall(rb_block_proc(), id_call, 1, element);
+
+    if (!RTEST(result)) {
+        return rb_funcall(rb_ary_entry(val, 0), id_yield, 1, element);
+    } else {
+        return result;
+    }
+}
+
+static VALUE
+lazy_reject(VALUE obj)
+{
+    if (!rb_block_given_p()) {
+        rb_raise(rb_eArgError, "tried to call lazy reject without a block");
+    }
+
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_reject_func, 0);
+}
+
+static VALUE
+lazy_grep_func(VALUE val, VALUE m, int argc, VALUE *argv)
+{
+    VALUE element = rb_ary_entry(val, 1);
+    VALUE result = rb_funcall(m, rb_intern("=~"), 1, element);
+
+    if (result != Qnil) {
+        return rb_funcall(rb_ary_entry(val, 0), id_yield, 1, element);
+    } else {
+        return result;
+    }
+}
+
+static VALUE
+lazy_grep(VALUE obj, VALUE pattern)
+{
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_grep_func, pattern);
+}
+
 static VALUE
 stop_result(VALUE self)
 {
@@ -1214,6 +1334,18 @@ Init_Enumerator(void)
     rb_define_method(rb_cEnumerator, "rewind", enumerator_rewind, 0);
     rb_define_method(rb_cEnumerator, "inspect", enumerator_inspect, 0);
 
+    /* Enumerable::Lazy */
+    rb_cLazy = rb_define_class_under(rb_mEnumerable, "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, "map", lazy_map, 0);
+    rb_define_method(rb_cLazy, "select", lazy_select, 0);
+    rb_define_method(rb_cLazy, "reject", lazy_reject, 0);
+    rb_define_method(rb_cLazy, "grep", lazy_grep, 1);
+
+    rb_define_alias(rb_cLazy, "collect", "map");
+    rb_define_alias(rb_cLazy, "find_all", "select");
+
     rb_eStopIteration = rb_define_class("StopIteration", rb_eIndexError);
     rb_define_method(rb_eStopIteration, "result", stop_result, 0);
 
@@ -1234,6 +1366,10 @@ Init_Enumerator(void)
 
     id_rewind = rb_intern("rewind");
     id_each = rb_intern("each");
+    id_call = rb_intern("call");
+    id_yield = rb_intern("yield");
+    id_new = rb_intern("new");
+    id_initialize = rb_intern("initialize");
     sym_each = ID2SYM(id_each);
 
     rb_provide("enumerator.so");	/* for backward compatibility */

test/ruby/test_lazy_enumerator.rb

@@ -0,0 +1,34 @@
+require 'test/unit'
+
+class TestLazyEnumerator < Test::Unit::TestCase
+
+  def test_initialize
+    assert_equal([1, 2, 3], [1, 2, 3].lazy.to_a)
+    assert_equal([1, 2, 3], Enumerable::Lazy.new([1, 2, 3]).to_a)
+  end
+
+  def test_select
+    a = [1, 2, 3, 4, 5, 6]
+    assert_equal([4, 5, 6], a.lazy.select { |x| x > 3 }.to_a)
+    a = ['word', nil, 1]
+    assert_equal(['word', 1], a.lazy.select { |x| x }.to_a)
+  end
+
+  def test_map
+    a = [1, 2, 3]
+    assert_equal([2, 4, 6], a.lazy.map { |x| x * 2 }.to_a)
+  end
+
+  def test_reject
+    a = [1, 2, 3, 4, 5, 6]
+    assert_equal([1, 2, 3], a.lazy.reject { |x| x > 3 }.to_a)
+    a = ['word', nil, 1]
+    assert_equal([nil], a.lazy.reject { |x| x }.to_a)
+  end
+
+  def test_grep
+    a = ['a', 'b', 'c', 'd', 'f']
+    assert_equal(['c'], a.lazy.grep(/c/).to_a)
+  end
+
+end