Fix enum size/counting for each_slice & friends.

As it turns out we can't always use Enumerable#count as a fallback as this doesn't work for either very large enumerators or lazy ones that go on to infinity (e.g. 0..Float::INFINITY). To work around this we _only_ use Enumerable#count as a fallback for Enumerable#each_slice and _only_ if the enumerator is not a Enumerator::Lazy instance. The resulting code is a bit messy but I can't think of any better place to put it.
rubinius · Feb 17, 2015 · 7ebdd69 · 7ebdd69
1 parent 33d8568
commit 7ebdd69
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 10 deletions.
diff --git a/kernel/common/enumerable.rb b/kernel/common/enumerable.rb
@@ -424,7 +424,17 @@ def each_cons(num)
   end
 
   def each_slice(slice_size)
-    return to_enum(:each_slice, slice_size) unless block_given?
+    unless block_given?
+      enum = to_enum(:each_slice, slice_size)
+
+      # For non lazy enumerators the enum size must be set to the size of the
+      # slice, _not_ the size of the containing enumerator.
+      unless enum.kind_of?(Enumerator::Lazy)
+        Rubinius.privately { enum.size = enum.count }
+      end
+
+      return enum
+    end
 
     n = Rubinius::Type.coerce_to_collection_index slice_size
     raise ArgumentError, "invalid slice size: #{n}" if n <= 0

diff --git a/kernel/common/enumerator.rb b/kernel/common/enumerator.rb
@@ -148,15 +148,7 @@ def rewind
     end
 
     def size
-      if @size.kind_of?(Proc)
-        @size.call
-      elsif @size
-        @size
-      # Lazy enums don't have a size and in case of Rbx can go into an infinite
-      # loop when using "count".
-      elsif !kind_of?(Enumerator::Lazy)
-        count
-      end
+      @size.kind_of?(Proc) ? @size.call : @size
     end
 
     def with_index(offset=0)