diff --git a/enumerator.c b/enumerator.c index 0d54058215..0e8665bcd2 100644 --- a/enumerator.c +++ b/enumerator.c @@ -2734,6 +2734,52 @@ lazy_with_index(int argc, VALUE *argv, VALUE obj) return lazy_add_method(obj, 0, 0, memo, rb_ary_new_from_values(1, &memo), &lazy_with_index_funcs); } +static struct MEMO * +lazy_lazy_each_proc(VALUE proc_entry, struct MEMO *result, VALUE memos, long memo_index) +{ + struct proc_entry *entry = proc_entry_ptr(proc_entry); + + rb_proc_call_with_block(entry->proc, 1, &result->memo_value, Qnil); + + return result; +} + +static const lazyenum_funcs lazy_lazy_each_funcs = { + lazy_lazy_each_proc, 0, +}; + +/* + * call-seq: + * lazy.lazy_each { |item| ... } -> lazy_enumerator + * + * Passes each element through to the block for side effects only, + * without modifying the element or affecting the enumeration. + * Returns a new lazy enumerator. + * + * This is useful for debugging or logging inside lazy chains, + * without breaking laziness or misusing +map+. + * + * (1..).lazy + * .lazy_each { |x| puts "got #{x}" } + * .select(&:even?) + * .first(3) + * # prints: got 1, got 2, ..., got 6 + * # returns: [2, 4, 6] + * + * Similar in intent to Java's Stream#peek. + */ + +static VALUE +lazy_lazy_each(VALUE obj) +{ + if (!rb_block_given_p()) + { + rb_raise(rb_eArgError, "tried to call lazy lazy_each without a block"); + } + + return lazy_add_method(obj, 0, 0, Qnil, Qnil, &lazy_lazy_each_funcs); +} + #if 0 /* for RDoc */ /* @@ -4561,6 +4607,7 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "uniq", lazy_uniq, 0); rb_define_method(rb_cLazy, "compact", lazy_compact, 0); rb_define_method(rb_cLazy, "with_index", lazy_with_index, -1); + rb_define_method(rb_cLazy, "lazy_each", lazy_lazy_each, 0); lazy_use_super_method = rb_hash_new_with_size(18); rb_hash_aset(lazy_use_super_method, sym("map"), sym("_enumerable_map")); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 4dddbab50c..3e1046d2aa 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -608,7 +608,7 @@ EOS end def test_require_block - %i[select reject drop_while take_while map flat_map].each do |method| + %i[select reject drop_while take_while map flat_map lazy_each].each do |method| assert_raise(ArgumentError){ [].lazy.send(method) } end end @@ -715,4 +715,23 @@ EOS def test_with_index_size assert_equal(3, Enumerator::Lazy.new([1, 2, 3], 3){|y, v| y << v}.with_index.size) end + + def test_lazy_lazy_each + out = [] + + e = (1..Float::INFINITY).lazy + .lazy_each { |x| out << x } + .select(&:even?) + .first(5) + + assert_equal([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], out) + assert_equal([2, 4, 6, 8, 10], e) + end + + def test_lazy_lazy_each_is_not_intrusive + s = Step.new(1..3) + + assert_equal(2, s.lazy.lazy_each { |x| x }.map { |x| x * 2 }.first) + assert_equal(1, s.current) + end end