Add `Enumerator.produce` from Ruby2.7

2 files changed, 75 insertions, 1 deletions

diff --git a/mrbgems/mruby-enumerator/mrblib/enumerator.rb b/mrbgems/mruby-enumerator/mrblib/enumerator.rb
index 89472ef01..d02caf5a0 100644
--- a/mrbgems/mruby-enumerator/mrblib/enumerator.rb
+++ b/mrbgems/mruby-enumerator/mrblib/enumerator.rb
@@ -556,6 +556,42 @@ class Enumerator
       self
     end
   end
+
+  ##
+  # call-seq:
+  #    Enumerator.produce(initial = nil) { |val| } -> enumerator
+  #
+  # Creates an infinite enumerator from any block, just called over and
+  # over.  Result of the previous iteration is passed to the next one.
+  # If +initial+ is provided, it is passed to the first iteration, and
+  # becomes the first element of the enumerator; if it is not provided,
+  # first iteration receives +nil+, and its result becomes first
+  # element of the iterator.
+  #
+  # Raising StopIteration from the block stops an iteration.
+  #
+  # Examples of usage:
+  #
+  #   Enumerator.produce(1, &:succ)   # => enumerator of 1, 2, 3, 4, ....
+  #
+  #   Enumerator.produce { rand(10) } # => infinite random number sequence
+  #
+  #   ancestors = Enumerator.produce(node) { |prev| node = prev.parent or raise StopIteration }
+  #   enclosing_section = ancestors.find { |n| n.type == :section }
+  def Enumerator.produce(init=NONE, &block)
+    raise ArgumentError, "no block given" if block.nil?
+    Enumerator.new do |y|
+      if init == NONE
+        val = nil
+      else
+        val = init
+        y.yield(val)
+      end
+      loop do
+        y.yield(val = block.call(val))
+      end
+    end
+  end
 end
 
 module Kernel
diff --git a/mrbgems/mruby-enumerator/test/enumerator.rb b/mrbgems/mruby-enumerator/test/enumerator.rb
index dce0c2cf2..7cff49c3b 100644
--- a/mrbgems/mruby-enumerator/test/enumerator.rb
+++ b/mrbgems/mruby-enumerator/test/enumerator.rb
@@ -10,9 +10,9 @@ def assert_take(exp, enumerator)
   result = []
   n = exp.size
   enumerator.each do |v|
-    break if n == 0
     result << v
     n -= 1
+    break if n == 0
   end
   assert_equal exp, result
 end
@@ -560,3 +560,41 @@ assert 'Enumerable#zip' do
 
   assert_raise(TypeError) { [1].zip(1) }
 end
+
+assert 'Enumerator.produce' do
+  assert_raise(ArgumentError) { Enumerator.produce }
+
+  # Without initial object
+  passed_args = []
+  enum = Enumerator.produce {|obj| passed_args << obj; (obj || 0).succ }
+  assert_equal Enumerator, enum.class 
+  assert_take [1, 2, 3], enum
+  assert_equal [nil, 1, 2], passed_args
+
+  # With initial object
+  passed_args = []
+  enum = Enumerator.produce(1) {|obj| passed_args << obj; obj.succ }
+  assert_take [1, 2, 3], enum
+  assert_equal [1, 2], passed_args
+
+  # Raising StopIteration
+  words = %w[The quick brown fox jumps over the lazy dog]
+  enum = Enumerator.produce { words.shift or raise StopIteration }
+  assert_equal %w[The quick brown fox jumps over the lazy dog], enum.to_a
+
+  # Raising StopIteration
+  object = [[[["abc", "def"], "ghi", "jkl"], "mno", "pqr"], "stuv", "wxyz"]
+  enum = Enumerator.produce(object) {|obj|
+    obj.respond_to?(:first) or raise StopIteration
+    obj.first
+  }
+  assert_nothing_raised {
+    assert_equal [
+      [[[["abc", "def"], "ghi", "jkl"], "mno", "pqr"], "stuv", "wxyz"],
+      [[["abc", "def"], "ghi", "jkl"], "mno", "pqr"],
+      [["abc", "def"], "ghi", "jkl"],
+      ["abc", "def"],
+      "abc",
+    ], enum.to_a
+  }
+end