From 702cf3ec9051914c8ef5a353580e0e74f7ec7005 Mon Sep 17 00:00:00 2001
From: Shugo Maeda <shugo@ruby-lang.org>
Date: Fri, 21 Jun 2019 17:27:20 +0900
Subject: [PATCH] Enumerator::Lazy should support filter_map

Fixes [Bug #15949]
---
 enumerator.c                      | 31 +++++++++++++++++++++++++++++++
 test/ruby/test_lazy_enumerator.rb |  8 ++++++++
 2 files changed, 39 insertions(+)

diff --git a/enumerator.c b/enumerator.c
index c37c8af298..3fb953c7f3 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -2059,6 +2059,37 @@ lazy_filter_map(VALUE obj)
                            Qnil, 0);
 }
 
+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)
 {
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)