Implement Enumerator::Lazy#with_index

Author: bjfish

CHANGELOG.md

@@ -30,6 +30,7 @@ Compatibility:
 * Allow `Fiber#raise` after `Fiber#transfer` like Ruby 3.0 (#2342).
 * Fix `ObjectSpace._id2ref` for Symbols and frozen String literals (#2358).
 * Implemented `Enumerator::Lazy#filter_map` (#2356).
+* Implemented `Enumerator::Lazy#with_index` (#2356).
 
 Performance:
 

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

@@ -0,0 +1,31 @@
+# -*- encoding: us-ascii -*-
+
+require_relative '../../../spec_helper'
+require_relative 'fixtures/classes'
+
+ruby_version_is "2.7" do
+  describe "Enumerator::Lazy#with_index" do
+    it "enumerates with an index" do
+      (0..Float::INFINITY).lazy.with_index.map { |i, idx| [i, idx] }.first(3).should == [[0, 0], [1, 1], [2, 2]]
+    end
+
+    it "enumerates with an index starting at a given offset" do
+      (0..Float::INFINITY).lazy.with_index(3).map { |i, idx| [i, idx] }.first(3).should == [[0, 3], [1, 4], [2, 5]]
+    end
+
+    it "enumerates with an index starting at 0 when offset is nil" do
+      (0..Float::INFINITY).lazy.with_index(nil).map { |i, idx| [i, idx] }.first(3).should == [[0, 0], [1, 1], [2, 2]]
+    end
+
+    it "raises TypeError when offset does not convert to Integer" do
+      -> { (0..Float::INFINITY).lazy.with_index(false).map { |i, idx| i }.first(3) }.should raise_error(TypeError)
+    end
+
+    it "enumerates with a given block" do
+      result = []
+      (0..Float::INFINITY).lazy.with_index { |i, idx| result << [i * 2, idx] }.first(3)
+      result.should == [[0,0],[2,1],[4,2]]
+    end
+  end
+end
+

spec/tags/truffle/methods_tags.txt

@@ -57,7 +57,6 @@ fails:Public methods on TracePoint should include parameters
 fails:Public methods on TracePoint should include raised_exception
 fails:Public methods on TracePoint should include return_value
 fails:Public methods on ENV.singleton_class should include freeze
-fails:Public methods on Enumerator::Lazy should include with_index
 fails:Public methods on IO should include set_encoding_by_bom
 fails:Public methods on BasicSocket should include read_nonblock
 fails:Public methods on BasicSocket should include write_nonblock

src/main/java/org/truffleruby/core/CoreLibrary.java

@@ -979,6 +979,7 @@ public boolean isTruffleBootMainMethod(SharedMethodInfo info) {
             "/core/truffle/debug.rb",
             "/core/truffle/diggable.rb",
             "/core/truffle/encoding_operations.rb",
+            "/core/truffle/enumerator_operations.rb",
             "/core/truffle/hash_operations.rb",
             "/core/truffle/interop_operations.rb",
             "/core/truffle/numeric_operations.rb",

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

@@ -275,6 +275,14 @@ def each(*args, &block)
   end
 
   class Lazy < Enumerator
+
+    aliases = Truffle::EnumeratorOperations::LAZY_OVERRIDE_METHODS.map do |m|
+      name = :"_enumerable_#{m}"
+      alias_method name, :"#{m}"
+      name
+    end
+    private(*aliases)
+
     class StopLazyError < Exception # rubocop:disable Lint/InheritException
     end
 
@@ -299,6 +307,7 @@ def initialize(receiver, size=nil)
     def to_enum(method_name=:each, *method_args, &block)
       size = block_given? ? block : nil
       ret = Lazy.allocate
+      method_name = Truffle::EnumeratorOperations.lazy_method(method_name)
 
       ret.__send__ :initialize_enumerator, self, size, method_name, *method_args
 
@@ -492,6 +501,23 @@ def flat_map
     end
     alias_method :collect_concat, :flat_map
 
+    def with_index(offset=0, &block)
+      offset = if Primitive.nil?(offset)
+                 0
+               else
+                 Truffle::Type.coerce_to offset, Integer, :to_int
+               end
+
+      Lazy.new(self, enumerator_size) do |yielder, *args|
+        if block
+          yielder.yield yield(*args, offset)
+        else
+          yielder.yield(*args, offset)
+        end
+        offset += 1
+      end
+    end
+
     def zip(*lists)
       return super(*lists) { |entry| yield entry } if block_given?
 

src/main/ruby/truffleruby/core/truffle/enumerator_operations.rb

@@ -0,0 +1,24 @@
+# frozen_string_literal: true
+
+# Copyright (c) 2021, 2021 Oracle and/or its affiliates. All rights reserved. This
+# code is released under a tri EPL/GPL/LGPL license. You can use it,
+# redistribute it and/or modify it under the terms of the:
+#
+# Eclipse Public License version 2.0, or
+# GNU General Public License version 2, or
+# GNU Lesser General Public License version 2.1.
+
+module Truffle
+  module EnumeratorOperations
+    LAZY_OVERRIDE_METHODS = %i[map collect flat_map collect_concat select find_all filter filter_map reject grep grep_v
+                               zip take take_while drop drop_while uniq with_index]
+
+    def self.lazy_method(meth)
+      if LAZY_OVERRIDE_METHODS.include?(meth)
+        :"_enumerable_#{meth}"
+      else
+        meth
+      end
+    end
+  end
+end