summaryrefslogtreecommitdiff
diff options
context:
space:
mode:
authorNobuyoshi Nakada <nobu@ruby-lang.org>2019-06-21 16:28:39 +0900
committerNobuyoshi Nakada <nobu@ruby-lang.org>2019-06-21 16:28:39 +0900
commit06a622b76e20b8996cd6d66c11ac7819a0996fae (patch)
tree4876d7dd948f43d3d0185000034a337bb6cb7e40
parent4149fd60786527121d42e3d6f12d1447505d22fa (diff)
Lazy filter_map
-rw-r--r--enumerator.c35
-rw-r--r--test/ruby/test_lazy_enumerator.rb9
2 files changed, 44 insertions, 0 deletions
diff --git a/enumerator.c b/enumerator.c
index 808ab67..983e14c 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -2026,6 +2026,40 @@ lazy_select(VALUE obj)
return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_select_funcs);
}
+static VALUE
+lazy_filter_map_proc(RB_BLOCK_CALL_FUNC_ARGLIST(val, m))
+{
+ VALUE result = rb_yield_values2(argc - 1, &argv[1]);
+ if (RTEST(result)) {
+ rb_funcall(argv[0], idLTLT, 1, result);
+ }
+ return Qnil;
+}
+
+/*
+ * call-seq:
+ * enum.filter_map { |obj| block } -> an_enumerator
+ * enum.filter_map -> an_enumerator
+ *
+ * Returns an enumerator which will return a new array containing the
+ * truthy results (everything except +false+ or +nil+) of running the
+ * +block+ for every element in +enum+.
+ *
+ * (1..).lazy.filter_map { |i| i * 2 if i.even? }.take(5) #=> [4, 8, 12, 16, 20]
+ *
+ */
+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_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
+ lazy_filter_map_proc, 0),
+ Qnil, 0);
+}
+
static struct MEMO *
lazy_reject_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index)
{
@@ -3555,6 +3589,7 @@ InitVM_Enumerator(void)
rb_define_method(rb_cLazy, "select", lazy_select, 0);
rb_define_method(rb_cLazy, "find_all", lazy_select, 0);
rb_define_method(rb_cLazy, "filter", lazy_select, 0);
+ rb_define_method(rb_cLazy, "filter_map", lazy_filter_map, 0);
rb_define_method(rb_cLazy, "reject", lazy_reject, 0);
rb_define_method(rb_cLazy, "grep", lazy_grep, 1);
rb_define_method(rb_cLazy, "grep_v", lazy_grep_v, 1);
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index d42c8d3..e76a53e 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -578,4 +578,13 @@ EOS
assert_equal([1, 2, 3, 4, 5, 10], u.first(6))
assert_equal([1, 2, 3, 4, 5, 10], u.first(6))
end
+
+ def test_filter_map
+ e = (1..Float::INFINITY).lazy.filter_map do |x|
+ raise "too big" if x > 10000
+ (x**2) % 10 if x.even?
+ end
+ assert_equal([4, 6, 6, 4, 0, 4], e.first(6))
+ assert_equal([4, 6, 6, 4, 0, 4], e.first(6))
+ end
end