From e94ac03eb0d07af3cbff20194acf479bf8851c39 Mon Sep 17 00:00:00 2001
From: Jeremy Evans <code@jeremyevans.net>
Date: Thu, 8 Aug 2019 13:09:17 -0700
Subject: [PATCH] Make Enumerator::Lazy#with_index be lazy

Previously, Enumerator::Lazy#with_index was not defined, so it
picked up the default implementation from Enumerator, which was
not lazy.

Based on earlier patch from nobu.

Fixes [Bug #7877]
---
 enumerator.c                      | 66 +++++++++++++++++++++++++++++++
 test/ruby/test_lazy_enumerator.rb | 22 +++++++++++
 2 files changed, 88 insertions(+)

diff --git a/enumerator.c b/enumerator.c
index 96daad2b4b..d4022aefc4 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -3597,6 +3597,71 @@ arith_seq_size(VALUE self)
     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)
 {
@@ -3654,6 +3719,7 @@ InitVM_Enumerator(void)
     rb_define_method(rb_cLazy, "slice_when", lazy_super, -1);
     rb_define_method(rb_cLazy, "chunk_while", lazy_super, -1);
     rb_define_method(rb_cLazy, "uniq", lazy_uniq, 0);
+    rb_define_method(rb_cLazy, "with_index", lazy_with_index, -1);
 
 #if 0 /* for RDoc */
     rb_define_method(rb_cLazy, "to_a", lazy_to_a, 0);
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index e76a53e650..80977a0bbd 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -587,4 +587,26 @@ EOS
     assert_equal([4, 6, 6, 4, 0, 4], e.first(6))
     assert_equal([4, 6, 6, 4, 0, 4], e.first(6))
   end
+
+  def test_with_index
+    feature7877 = '[ruby-dev:47025] [Feature #7877]'
+    leibniz = ->(n) {
+      (0..Float::INFINITY).lazy.with_index {|i, j|
+        raise IndexError, "limit exceeded (#{n})" unless j < n
+        ((-1) ** j) / (2*i+1).to_f
+      }.take(n).reduce(:+)
+    }
+    assert_nothing_raised(IndexError, feature7877) {
+      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
+    assert_equal([[-1, 2], [0, 3]], ary)
+
+    ary = (0..Float::INFINITY).lazy.with_index(2).take(2).to_a
+    assert_equal([[0, 2], [1, 3]], ary)
+
+    ary = (0..Float::INFINITY).lazy.with_index.take(2).to_a
+    assert_equal([[0, 0], [1, 1]], ary)
+  end
 end