summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
-rw-r--r--enumerator.c31
-rw-r--r--test/ruby/test_lazy_enumerator.rb8
2 files changed, 39 insertions, 0 deletions
diff --git a/enumerator.c b/enumerator.c
index c37c8af298..3fb953c7f3 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -2060,6 +2060,37 @@ lazy_filter_map(VALUE obj)
}
static struct MEMO *
+lazy_filter_map_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index)
+{
+ VALUE value = lazyenum_yield_values(proc_entry, result);
+ if (!RTEST(value)) return 0;
+ LAZY_MEMO_SET_VALUE(result, value);
+ LAZY_MEMO_RESET_PACKED(result);
+ return result;
+}
+
+static const lazyenum_funcs lazy_filter_map_funcs = {
+ lazy_filter_map_proc, 0,
+};
+
+/*
+ * call-seq:
+ * lazy.filter_map { |obj| block } -> lazy_enumerator
+ *
+ * Like Enumerable#filter_map, but chains operation to be lazy-evaluated.
+ */
+
+static VALUE
+lazy_filter_map(VALUE obj)
+{
+ if (!rb_block_given_p()) {
+ rb_raise(rb_eArgError, "tried to call lazy filter_map without a block");
+ }
+
+ return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_filter_map_funcs);
+}
+
+static struct MEMO *
lazy_reject_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index)
{
VALUE chain = lazyenum_yield(proc_entry, result);
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index e76a53e650..679e0291b1 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -116,6 +116,14 @@ class TestLazyEnumerator < Test::Unit::TestCase
assert_equal(expected, a.lazy.map {|*args| args}.map {|*args| args}.to_a, bug)
end
+ def test_filter_map
+ a = Step.new(1..3)
+ assert_equal(2, a.filter_map {|x| x.odd? && x * 2}.first)
+ assert_equal(3, a.current)
+ assert_equal(2, a.lazy.filter_map {|x| x.odd? && x * 2}.first)
+ assert_equal(1, a.current)
+ end
+
def test_flat_map
a = Step.new(1..3)
assert_equal(2, a.flat_map {|x| [x * 2]}.first)