Merge pull request #4210 from dearblue/enum-chain

Add enumerator chain feature (CRuby-2.6 compatible)

3 files changed, 142 insertions, 0 deletions

diff --git a/mrbgems/mruby-enum-chain/mrbgem.rake b/mrbgems/mruby-enum-chain/mrbgem.rake
new file mode 100644
index 000000000..7294f2644
--- /dev/null
+++ b/mrbgems/mruby-enum-chain/mrbgem.rake
@@ -0,0 +1,6 @@
+MRuby::Gem::Specification.new('mruby-enum-chain') do |spec|
+  spec.license = 'MIT'
+  spec.author  = 'mruby developers'
+  spec.summary = 'Enumerator::Chain class'
+  spec.add_dependency('mruby-enumerator', :core => 'mruby-enumerator')
+end
diff --git a/mrbgems/mruby-enum-chain/mrblib/chain.rb b/mrbgems/mruby-enum-chain/mrblib/chain.rb
new file mode 100644
index 000000000..98515ea14
--- /dev/null
+++ b/mrbgems/mruby-enum-chain/mrblib/chain.rb
@@ -0,0 +1,60 @@
+##
+# chain.rb Enumerator::Chain class
+# See Copyright Notice in mruby.h
+
+module Enumerable
+  def chain(*args)
+    Enumerator::Chain.new(self, *args)
+  end
+
+  def +(other)
+    Enumerator::Chain.new(self, other)
+  end
+end
+
+class Enumerator
+  class Chain
+    include Enumerable
+
+    def initialize(*args)
+      @enums = args
+    end
+
+    def initialize_copy(orig)
+      @enums = orig.__copy_enums
+    end
+
+    def each(&block)
+      return to_enum unless block_given?
+
+      @enums.each { |e| e.each(&block) }
+
+      self
+    end
+
+    def size
+      @enums.reduce(0) do |a, e|
+        return nil unless e.respond_to?(:size)
+        a + e.size
+      end
+    end
+
+    def rewind
+      @enums.reverse_each do |e|
+        e.rewind if e.respond_to?(:rewind)
+      end
+
+      self
+    end
+
+    def inspect
+      "#<#{self.class}: #{@enums.inspect}>"
+    end
+
+    def __copy_enums
+      @enums.each_with_object([]) do |e, a|
+        a << e.clone
+      end
+    end
+  end
+end
diff --git a/mrbgems/mruby-enum-chain/test/enum_chain.rb b/mrbgems/mruby-enum-chain/test/enum_chain.rb
new file mode 100644
index 000000000..4dd59bd37
--- /dev/null
+++ b/mrbgems/mruby-enum-chain/test/enum_chain.rb
@@ -0,0 +1,76 @@
+##
+# Enumerator::Chain test
+
+assert("Enumerable#chain") do
+  a = []
+  b = {}
+  c = Object.new # not has #each method
+
+  assert_kind_of Enumerator::Chain, a.chain
+  assert_kind_of Enumerator::Chain, a.chain(b)
+  assert_kind_of Enumerator::Chain, a.chain(b, c)
+  assert_raise(NoMethodError) { c.chain }
+end
+
+assert("Enumerable#+") do
+  a = [].each
+  b = {}.reverse_each
+  c = Object.new # not has #each method
+
+  assert_kind_of Enumerator::Chain, a + b
+  assert_kind_of Enumerator::Chain, a + c
+  assert_kind_of Enumerator::Chain, b + a
+  assert_kind_of Enumerator::Chain, b + c
+  assert_raise(NoMethodError) { c + a }
+end
+
+assert("Enumerator.new") do
+  a = []
+  b = {}
+  c = Object.new # not has #each method
+
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(a, a)
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(a, b)
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(a, c)
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(b, a)
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(b, b)
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(b, c)
+  assert_kind_of Enumerator::Chain, Enumerator::Chain.new(c, a)
+end
+
+assert("Enumerator::Chain#each") do
+  a = [1, 2, 3]
+
+  aa = a.chain(a)
+  assert_kind_of Enumerator, aa.each
+  assert_equal [1, 2, 3, 1, 2, 3], aa.each.to_a
+
+  aa = a.chain(a.reverse_each)
+  assert_kind_of Enumerator, aa.each
+  assert_equal [1, 2, 3, 3, 2, 1], aa.each.to_a
+
+  aa = a.chain(a.reverse_each, a)
+  assert_kind_of Enumerator, aa.each
+  assert_equal [1, 2, 3, 3, 2, 1, 1, 2, 3], aa.each.to_a
+
+  aa = a.chain(Object.new)
+  assert_kind_of Enumerator, aa.each
+  assert_raise(NoMethodError) {  aa.each.to_a }
+end
+
+assert("Enumerator::Chain#size") do
+  a = [1, 2, 3]
+
+  aa = a.chain(a)
+  assert_equal 6, aa.size
+
+  aa = a.chain(a.reverse_each)
+  assert_nil aa.size
+
+  aa = a.chain(a.reverse_each, a)
+  assert_nil aa.size
+
+  aa = a.chain(Object.new)
+  assert_nil aa.size
+end