[GR-37235] Fix Enumerator::Lazy#with_index to start with new index for multiple enumerations

PullRequest: truffleruby/3230

Author: bjfish

CHANGELOG.md

@@ -39,6 +39,7 @@ Bug fixes:
 * Fix `Array#sample` for `[]` when called without `n` and a `Random` is given (#2612, @bjfish).
 * Fix `Module#const_get` to raise a `NameError` when nested modules do not exist (#2610, @bjfish).
 * Ensure native `VALUE`s returned from C are unwrapped before the objects can be collected (@aardvark179).
+* Fix `Enumerator::Lazy#with_index` to start with new index for multiple enumerations (@bjfish).
 
 Compatibility:
 

spec/ruby/core/enumerator/lazy/with_index_spec.rb

@@ -26,5 +26,13 @@
       (0..Float::INFINITY).lazy.with_index { |i, idx| result << [i * 2, idx] }.first(3)
       result.should == [[0,0],[2,1],[4,2]]
     end
+
+    it "resets after a new call to each" do
+      enum = (0..2).lazy.with_index.map { |i, idx| [i, idx] }
+      result = []
+      enum.each { |i, idx| result << [i, idx] }
+      enum.each { |i, idx| result << [i, idx] }
+      result.should == [[0,0], [1,1], [2,2], [0,0], [1,1], [2,2]]
+    end
   end
 end

src/main/ruby/truffleruby/core/enumerator.rb

@@ -230,10 +230,13 @@ def self.produce(initial = nil)
   end
 
   class Yielder
+    attr_accessor :memo
+
     def initialize(&block)
       raise LocalJumpError, 'Expected a block to be given' unless block_given?
 
       @proc = block
+      @memo = nil
 
       self
     end
@@ -507,16 +510,17 @@ def with_index(offset=0, &block)
       offset = if Primitive.nil?(offset)
                  0
                else
-                 Truffle::Type.coerce_to offset, Integer, :to_int
+                 offset
                end
 
       Lazy.new(self, enumerator_size) do |yielder, *args|
+        memo = yielder.memo || offset
         if block
-          yielder.yield yield(*args, offset)
+          yielder.yield yield(*args, memo)
         else
-          yielder.yield(*args, offset)
+          yielder.yield(*args, memo)
         end
-        offset += 1
+        yielder.memo = Primitive.rb_num2long(memo) + 1
       end
     end