Fix Enumerator::Lazy#with_index

* Make it correctly handle lambdas
* Make it iterate over the block if block is given

The original implementation was flawed, based on lazy_set_method
instead of lazy_add_method.

Note that there is no implicit map when passing a block, the return
value of the block passed to with_index is ignored, just as it
is for Enumerator#with_index. Also like Enumerator#with_index,
when called with a block, the return value is an enumerator without
the index.

Fixes [Bug #16414]

1 files changed, 15 insertions, 2 deletions

diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index caa9a9c713..6e5c1714a5 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -647,7 +647,7 @@ EOS
   def test_with_index
     feature7877 = '[ruby-dev:47025] [Feature #7877]'
     leibniz = ->(n) {
-      (0..Float::INFINITY).lazy.with_index {|i, j|
+      (0..Float::INFINITY).lazy.with_index.map {|i, j|
         raise IndexError, "limit exceeded (#{n})" unless j < n
         ((-1) ** j) / (2*i+1).to_f
       }.take(n).reduce(:+)
@@ -656,7 +656,20 @@ EOS
       assert_in_epsilon(Math::PI/4, leibniz[1000])
     }
 
-    ary = (0..Float::INFINITY).lazy.with_index(2) {|i, j| [i-1, j] }.take(2).to_a
+    a = []
+    ary = (0..Float::INFINITY).lazy.with_index(2) {|i, j| a << [i-1, j] }.take(2).to_a
+    assert_equal([[-1, 2], [0, 3]], a)
+    assert_equal([0, 1], ary)
+
+    a = []
+    ary = (0..Float::INFINITY).lazy.with_index(2, &->(i,j) { a << [i-1, j] }).take(2).to_a
+    assert_equal([[-1, 2], [0, 3]], a)
+    assert_equal([0, 1], ary)
+
+    ary = (0..Float::INFINITY).lazy.with_index(2).map {|i, j| [i-1, j] }.take(2).to_a
+    assert_equal([[-1, 2], [0, 3]], ary)
+
+    ary = (0..Float::INFINITY).lazy.with_index(2).map(&->(i, j) { [i-1, j] }).take(2).to_a
     assert_equal([[-1, 2], [0, 3]], ary)
 
     ary = (0..Float::INFINITY).lazy.with_index(2).take(2).to_a