https://git.ruby-lang.org/ruby.git/commit/?id=85e43e1dfe

From 85e43e1dfecef69b935c48c235cc20f21bd4f0d4 Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@j...>
Date: Tue, 10 Dec 2019 13:38:12 -0800
Subject: 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]

diff --git a/enumerator.c b/enumerator.c
index 67b9302..e85c911 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -2645,6 +2645,62 @@ lazy_uniq(VALUE obj) https://github.com/ruby/ruby/blob/trunk/enumerator.c#L2645
     return lazy_add_method(obj, 0, 0, Qnil, Qnil, funcs);
 }
 
+static struct MEMO *
+lazy_with_index_proc(VALUE proc_entry, struct MEMO* result, VALUE memos, long memo_index)
+{
+    struct proc_entry *entry = proc_entry_ptr(proc_entry);
+    VALUE memo = rb_ary_entry(memos, memo_index);
+    VALUE argv[2];
+
+    if (NIL_P(memo)) {
+        memo = entry->memo;
+    }
+
+    argv[0] = result->memo_value;
+    argv[1] = memo;
+    if (entry->proc) {
+        rb_proc_call_with_block(entry->proc, 2, argv, Qnil);
+        LAZY_MEMO_RESET_PACKED(result);
+    } else {
+        LAZY_MEMO_SET_VALUE(result, rb_ary_new_from_values(2, argv));
+        LAZY_MEMO_SET_PACKED(result);
+    }
+    rb_ary_store(memos, memo_index, LONG2NUM(NUM2LONG(memo) + 1));
+    return result;
+}
+
+static const lazyenum_funcs lazy_with_index_funcs = {
+    lazy_with_index_proc, 0,
+};
+
+/*
+ *  call-seq:
+ *     lazy.with_index(offset = 0) {|(*args), idx| ... }
+ *     lazy.with_index(offset = 0)
+ *
+ *  If a block is given, iterates the given block for each element
+ *  with an index, which starts from +offset+, and returns a
+ *  lazy enumerator that yields the same values (without the index).
+ *
+ *  If a block is not given, returns a new lazy enumerator that
+ *  includes the index, starting from +offset+.
+ *
+ * +offset+:: the starting index to use
+ *
+ * see Enumerator#with_index.
+ */
+static VALUE
+lazy_with_index(int argc, VALUE *argv, VALUE obj)
+{
+    VALUE memo;
+
+    rb_scan_args(argc, argv, "01", &memo);
+    if (NIL_P(memo))
+        memo = LONG2NUM(0);
+
+    return lazy_add_method(obj, 0, 0, memo, rb_ary_new_from_values(1, &memo), &lazy_with_index_funcs);
+}
+
 #if 0 /* for RDoc */
 
 /*
@@ -3841,71 +3897,6 @@ arith_seq_size(VALUE self) https://github.com/ruby/ruby/blob/trunk/enumerator.c#L3897
     return len;
 }
 
-static VALUE
-lazy_with_index_func(RB_BLOCK_CALL_FUNC_ARGLIST(val, offset))
-{
-    VALUE yielder, memo, result;
-    VALUE e = rb_enum_values_pack(argc - 1, argv + 1);
-    long idx;
-
-    yielder = argv[0];
-    memo = rb_attr_get(yielder, id_memo);
-    if (NIL_P(memo))
-        memo = offset;
-    idx = NUM2LONG(memo);
-    result = rb_assoc_new(e, memo);
-    rb_funcall(yielder, idLTLT, 1, result);
-    rb_ivar_set(yielder, id_memo, LONG2NUM(++idx));
-    return Qnil;
-}
-
-static VALUE
-lazy_with_index_iter(RB_BLOCK_CALL_FUNC_ARGLIST(val, offset))
-{
-    VALUE yielder, memo, result;
-    VALUE e = rb_enum_values_pack(argc - 1, argv + 1);
-    long idx;
-
-    yielder = argv[0];
-    memo = rb_attr_get(yielder, id_memo);
-    if (NIL_P(memo))
-        memo = offset;
-    idx = NUM2LONG(memo);
-    result = rb_yield(rb_assoc_new(e, memo));
-    rb_funcall(yielder, idLTLT, 1, result);
-    rb_ivar_set(yielder, id_memo, LONG2NUM(++idx));
-    return Qnil;
-}
-
-/*
- *  call-seq:
- *     lazy.with_index(offset = 0) {|(*args), idx| ... }
- *     lazy.with_index(offset = 0)
- *
- *  Iterates the given block for each element with an index, which
- *  starts from +offset+.  If no block is given, returns a new
- *  lazy enumerator that includes the index, starting from +offset+
- *
- * +offset+:: the starting index to use
- *
- * see Enumerator#with_index.
- */
-static VALUE
-lazy_with_index(int argc, VALUE *argv, VALUE obj)
-{
-    VALUE memo;
-
-    rb_scan_args(argc, argv, "01", &memo);
-    if (NIL_P(memo))
-        memo = LONG2NUM(0);
-
-    return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj,
-                                         rb_block_given_p() ?
-                                         lazy_with_index_iter : lazy_with_index_func,
-                                         memo),
-                           rb_ary_new_from_values(argc, argv), 0);
-}
-
 void
 InitVM_Enumerator(void)
 {
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index caa9a9c..6e5c171 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -647,7 +647,7 @@ EOS https://github.com/ruby/ruby/blob/trunk/test/ruby/test_lazy_enumerator.rb#L647
   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 https://github.com/ruby/ruby/blob/trunk/test/ruby/test_lazy_enumerator.rb#L656
       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
-- 
cgit v0.10.2


--
ML: ruby-changes@q...
Info: http://www.atdot.net/~ko1/quickml/