2 files changed, 60 insertions, 26 deletions

diff --git a/mrbgems/mruby-enum-chain/mrblib/chain.rb b/mrbgems/mruby-enum-chain/mrblib/chain.rb
index 98515ea14..43d0926c8 100644
--- a/mrbgems/mruby-enum-chain/mrblib/chain.rb
+++ b/mrbgems/mruby-enum-chain/mrblib/chain.rb
@@ -6,28 +6,30 @@ module Enumerable
   def chain(*args)
     Enumerator::Chain.new(self, *args)
   end
+end
 
+class Enumerator
   def +(other)
-    Enumerator::Chain.new(self, other)
+    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
+      @enums = args.freeze
+      @pos = -1
     end
 
     def each(&block)
-      return to_enum unless block_given?
+      return to_enum unless block
 
-      @enums.each { |e| e.each(&block) }
+      i = 0
+      while i < @enums.size
+        @pos = i
+        @enums[i].each(&block)
+        i += 1
+      end
 
       self
     end
@@ -40,21 +42,21 @@ class Enumerator
     end
 
     def rewind
-      @enums.reverse_each do |e|
+      while 0 <= @pos && @pos < @enums.size
+        e = @enums[@pos]
         e.rewind if e.respond_to?(:rewind)
+        @pos -= 1
       end
 
       self
     end
 
-    def inspect
-      "#<#{self.class}: #{@enums.inspect}>"
+    def +(other)
+      self.class.new(self, other)
     end
 
-    def __copy_enums
-      @enums.each_with_object([]) do |e, a|
-        a << e.clone
-      end
+    def inspect
+      "#<#{self.class}: #{@enums.inspect}>"
     end
   end
 end
diff --git a/mrbgems/mruby-enum-chain/test/enum_chain.rb b/mrbgems/mruby-enum-chain/test/enum_chain.rb
index 4dd59bd37..45bbc9a77 100644
--- a/mrbgems/mruby-enum-chain/test/enum_chain.rb
+++ b/mrbgems/mruby-enum-chain/test/enum_chain.rb
@@ -12,9 +12,9 @@ assert("Enumerable#chain") do
   assert_raise(NoMethodError) { c.chain }
 end
 
-assert("Enumerable#+") do
+assert("Enumerator#+") do
   a = [].each
-  b = {}.reverse_each
+  b = {}.each
   c = Object.new # not has #each method
 
   assert_kind_of Enumerator::Chain, a + b
@@ -24,7 +24,7 @@ assert("Enumerable#+") do
   assert_raise(NoMethodError) { c + a }
 end
 
-assert("Enumerator.new") do
+assert("Enumerator::Chain.new") do
   a = []
   b = {}
   c = Object.new # not has #each method
@@ -46,13 +46,13 @@ assert("Enumerator::Chain#each") do
   assert_kind_of Enumerator, aa.each
   assert_equal [1, 2, 3, 1, 2, 3], aa.each.to_a
 
-  aa = a.chain(a.reverse_each)
+  aa = a.chain(6..9)
   assert_kind_of Enumerator, aa.each
-  assert_equal [1, 2, 3, 3, 2, 1], aa.each.to_a
+  assert_equal [1, 2, 3, 6, 7, 8, 9], aa.each.to_a
 
-  aa = a.chain(a.reverse_each, a)
+  aa = a.chain((-3..-2).each_with_index, a)
   assert_kind_of Enumerator, aa.each
-  assert_equal [1, 2, 3, 3, 2, 1, 1, 2, 3], aa.each.to_a
+  assert_equal [1, 2, 3, [-3, 0], [-2, 1], 1, 2, 3], aa.each.to_a
 
   aa = a.chain(Object.new)
   assert_kind_of Enumerator, aa.each
@@ -65,12 +65,44 @@ assert("Enumerator::Chain#size") do
   aa = a.chain(a)
   assert_equal 6, aa.size
 
-  aa = a.chain(a.reverse_each)
+  aa = a.chain(3..4)
   assert_nil aa.size
 
-  aa = a.chain(a.reverse_each, a)
+  aa = a.chain(3..4, a)
   assert_nil aa.size
 
   aa = a.chain(Object.new)
   assert_nil aa.size
 end
+
+assert("Enumerator::Chain#rewind") do
+  rewound = nil
+  e1 = [1, 2]
+  e2 = (4..6)
+  (class << e1; self end).define_method(:rewind) { rewound << self }
+  (class << e2; self end).define_method(:rewind) { rewound << self }
+  c = e1.chain(e2)
+
+  rewound = []
+  c.rewind
+  assert_equal [], rewound
+
+  rewound = []
+  c.each{break c}.rewind
+  assert_equal [e1], rewound
+
+  rewound = []
+  c.each{}.rewind
+  assert_equal [e2, e1], rewound
+end
+
+assert("Enumerator::Chain#+") do
+  a = [].chain
+  b = {}.chain
+  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
+end