Improved collections performance in common use cases

This improves the performance of most collection methods sometimes
by making it 100% faster compared to the previous implementation.

The improvement is not for free, with this change we are giving up
the lazyness feature (not iterate unless requested). The giving
up of the feature only happens when the collection has been initialized
with an array, as oposed to being initialized with another iterator or
generator; and this is the reson I consider this change a safe one.

For the curious, the improvement comes from the (sad) fact that calling
functions in php is extremenly expensive, specially when using iterators
since each iteration will call at least 4 functions (valid, next, current, key).

This becomes even worse as `IteratorIterator` does not have any optimizations,
so for each wrapped iterator, the number of functions is multiplied by 2
for each iteration.

The change proposed here will unwrap nested iterators as much as possible
to avoid the function call explosion, and in some (safe) cases, will
iterator the collection immediately as an array before wrapping it again
in an iterator.

I was inspired by Haskell when implementing this, as the language is lazy
by default, but the compiler optimizes the cases where code is safe to be
called strictly. Thats is called strictness analysis.

src/Collection/CollectionTrait.php

@@ -47,7 +47,7 @@ trait CollectionTrait
      */
     public function each(callable $c)
     {
-        foreach ($this->unwrap() as $k => $v) {
+        foreach ($this->optimizeUnwrap() as $k => $v) {
             $c($v, $k);
         }
 
@@ -87,7 +87,7 @@ trait CollectionTrait
      */
     public function every(callable $c)
     {
-        foreach ($this->unwrap() as $key => $value) {
+        foreach ($this->optimizeUnwrap() as $key => $value) {
             if (!$c($value, $key)) {
                 return false;
             }
@@ -101,7 +101,7 @@ trait CollectionTrait
      */
     public function some(callable $c)
     {
-        foreach ($this->unwrap() as $key => $value) {
+        foreach ($this->optimizeUnwrap() as $key => $value) {
             if ($c($value, $key) === true) {
                 return true;
             }
@@ -115,7 +115,7 @@ trait CollectionTrait
      */
     public function contains($value)
     {
-        foreach ($this->unwrap() as $v) {
+        foreach ($this->optimizeUnwrap() as $v) {
             if ($value === $v) {
                 return true;
             }
@@ -145,7 +145,7 @@ trait CollectionTrait
         }
 
         $result = $zero;
-        foreach ($this->unwrap() as $k => $value) {
+        foreach ($this->optimizeUnwrap() as $k => $value) {
             if ($isFirst) {
                 $result = $value;
                 $isFirst = false;
@@ -205,7 +205,7 @@ trait CollectionTrait
     {
         $callback = $this->_propertyExtractor($callback);
         $group = [];
-        foreach ($this as $value) {
+        foreach ($this->optimizeUnwrap() as $value) {
             $group[$callback($value)][] = $value;
         }
 
@@ -219,7 +219,7 @@ trait CollectionTrait
     {
         $callback = $this->_propertyExtractor($callback);
         $group = [];
-        foreach ($this as $value) {
+        foreach ($this->optimizeUnwrap() as $value) {
             $group[$callback($value)] = $value;
         }
 
@@ -255,7 +255,7 @@ trait CollectionTrait
 
         $callback = $this->_propertyExtractor($matcher);
         $sum = 0;
-        foreach ($this as $k => $v) {
+        foreach ($this->optimizeUnwrap() as $k => $v) {
             $sum += $callback($v, $k);
         }
 
@@ -500,7 +500,7 @@ trait CollectionTrait
      */
     public function buffered()
     {
-        return new BufferedIterator($this);
+        return new BufferedIterator($this->unwrap());
     }
 
     /**
@@ -534,7 +534,7 @@ trait CollectionTrait
             $condition = $this->_createMatcherFilter($condition);
         }
 
-        return new StoppableIterator($this, $condition);
+        return new StoppableIterator($this->unwrap(), $condition);
     }
 
     /**
@@ -550,7 +550,7 @@ trait CollectionTrait
 
         return new Collection(
             new RecursiveIteratorIterator(
-                new UnfoldIterator($this, $transformer),
+                new UnfoldIterator($this->unwrap(), $transformer),
                 RecursiveIteratorIterator::LEAVES_ONLY
             )
         );
@@ -571,7 +571,7 @@ trait CollectionTrait
      */
     public function zip($items)
     {
-        return new ZipIterator(array_merge([$this], func_get_args()));
+        return new ZipIterator(array_merge([$this->unwrap()], func_get_args()));
     }
 
     /**
@@ -586,7 +586,7 @@ trait CollectionTrait
             $items = [$items];
         }
 
-        return new ZipIterator(array_merge([$this], $items), $callable);
+        return new ZipIterator(array_merge([$this->unwrap()], $items), $callable);
     }
 
     /**
@@ -640,7 +640,7 @@ trait CollectionTrait
      */
     public function isEmpty()
     {
-        foreach ($this->unwrap() as $el) {
+        foreach ($this as $el) {
             return false;
         }
 
@@ -657,6 +657,10 @@ trait CollectionTrait
             $iterator = $iterator->getInnerIterator();
         }
 
+        if ($iterator !== $this && $iterator instanceof CollectionInterface) {
+            $iterator = $iterator->unwrap();
+        }
+
         return $iterator;
     }
 
@@ -747,4 +751,21 @@ trait CollectionTrait
 
         return new Collection($result);
     }
+
+     /**
+      * Unwraps this iterator and returns the simplest
+      * traversable that can be used for getting the data out
+      *
+      * @return \Traversable|array
+      */
+    protected function optimizeUnwrap()
+    {
+        $iterator = $this->unwrap();
+
+        if ($iterator instanceof ArrayIterator) {
+            $iterator = $iterator->getArrayCopy();
+        }
+
+        return $iterator;
+    }
 }

src/Collection/Iterator/ExtractIterator.php

@@ -14,7 +14,9 @@
  */
 namespace Cake\Collection\Iterator;
 
+use ArrayIterator;
 use Cake\Collection\Collection;
+use Cake\Collection\CollectionInterface;
 
 /**
  * Creates an iterator from another iterator that extract the requested column
@@ -69,4 +71,37 @@ class ExtractIterator extends Collection
 
         return $extractor(parent::current());
     }
+
+    /**
+     * {@inheritDoc}
+     *
+     * We perform here some stricness analysys so that the
+     * iterator logic is bypassed entirely.
+     *
+     * @return \Iterator
+     */
+    public function unwrap()
+    {
+        $iterator = $this->getInnerIterator();
+
+        if ($iterator instanceof CollectionInterface) {
+            $iterator = $iterator->unwrap();
+        }
+
+        if (!$iterator instanceof ArrayIterator) {
+            return $this;
+        }
+
+        // ArrayIterator can be traversed strictly.
+        // Let's do that for performance gains
+
+        $callback = $this->_extractor;
+        $res = [];
+
+        foreach ($iterator->getArrayCopy() as $k => $v) {
+            $res[$k] = $callback($v);
+        }
+
+        return new ArrayIterator($res);
+    }
 }

src/Collection/Iterator/FilterIterator.php

@@ -14,7 +14,9 @@
  */
 namespace Cake\Collection\Iterator;
 
+use ArrayIterator;
 use Cake\Collection\Collection;
+use Cake\Collection\CollectionInterface;
 use CallbackFilterIterator;
 use Iterator;
 
@@ -27,6 +29,13 @@ class FilterIterator extends Collection
 {
 
     /**
+     * The callback used to filter the elements in this collection
+     *
+     * @var callable
+     */
+    protected $_callback;
+
+    /**
      * Creates a filtered iterator using the callback to determine which items are
      * accepted or rejected.
      *
@@ -37,9 +46,50 @@ class FilterIterator extends Collection
      * @param \Iterator $items The items to be filtered.
      * @param callable $callback Callback.
      */
-    public function __construct(Iterator $items, callable $callback)
+    public function __construct($items, callable $callback)
     {
+        if (!$items instanceof Iterator) {
+            $items = new Collection($items);
+        }
+
+        $this->_callback = $callback;
         $wrapper = new CallbackFilterIterator($items, $callback);
         parent::__construct($wrapper);
     }
+
+    /**
+     * {@inheritDoc}
+     *
+     * We perform here some stricness analysys so that the
+     * iterator logic is bypassed entirely.
+     *
+     * @return \Iterator
+     */
+    public function unwrap()
+    {
+        $filter = $this->getInnerIterator();
+        $iterator = $filter->getInnerIterator();
+
+        if ($iterator instanceof CollectionInterface) {
+            $iterator = $iterator->unwrap();
+        }
+
+        if (!$iterator instanceof ArrayIterator) {
+            return $filter;
+        }
+
+        // ArrayIterator can be traversed strictly.
+        // Let's do that for performance gains
+
+        $callback = $this->_callback;
+        $res = [];
+
+        foreach ($iterator as $k => $v) {
+            if ($callback($v, $k, $iterator)) {
+                $res[$k] = $v;
+            }
+        }
+
+        return new ArrayIterator($res);
+    }
 }

src/Collection/Iterator/ReplaceIterator.php

@@ -14,7 +14,9 @@
  */
 namespace Cake\Collection\Iterator;
 
+use ArrayIterator;
 use Cake\Collection\Collection;
+use Cake\Collection\CollectionInterface;
 
 /**
  * Creates an iterator from another iterator that will modify each of the values
@@ -24,7 +26,7 @@ class ReplaceIterator extends Collection
 {
 
     /**
-     * The callback function to be used to modify each of the values
+     * The callback function to be used to transform values
      *
      * @var callable
      */
@@ -67,4 +69,38 @@ class ReplaceIterator extends Collection
 
         return $callback(parent::current(), $this->key(), $this->_innerIterator);
     }
+
+
+    /**
+     * {@inheritDoc}
+     *
+     * We perform here some stricness analysys so that the
+     * iterator logic is bypassed entirely.
+     *
+     * @return \Iterator
+     */
+    public function unwrap()
+    {
+        $iterator = $this->_innerIterator;
+
+        if ($iterator instanceof CollectionInterface) {
+            $iterator = $iterator->unwrap();
+        }
+
+        if (!$iterator instanceof ArrayIterator) {
+            return $this;
+        }
+
+        // ArrayIterator can be traversed strictly.
+        // Let's do that for performance gains
+
+        $callback = $this->_callback;
+        $res = [];
+
+        foreach ($iterator as $k => $v) {
+            $res[$k] = $callback($v, $k, $iterator);
+        }
+
+        return new ArrayIterator($res);
+    }
 }

src/Collection/Iterator/SortIterator.php

@@ -59,11 +59,10 @@ class SortIterator extends Collection
      */
     public function __construct($items, $callback, $dir = SORT_DESC, $type = SORT_NUMERIC)
     {
-        if (is_array($items)) {
-            $items = new Collection($items);
+        if (!is_array($items)) {
+            $items = iterator_to_array((new Collection($items))->unwrap(), false);
         }
 
-        $items = iterator_to_array($items, false);
         $callback = $this->_propertyExtractor($callback);
         $results = [];
         foreach ($items as $key => $value) {
@@ -81,4 +80,14 @@ class SortIterator extends Collection
         }
         parent::__construct($results);
     }
+
+    /**
+     * {@inheritDoc}
+     *
+     * @return \Iterator
+     */
+    public function unwrap()
+    {
+        return $this->getInnerIterator();
+    }
 }

src/Collection/Iterator/StoppableIterator.php

@@ -14,7 +14,9 @@
  */
 namespace Cake\Collection\Iterator;
 
+use ArrayIterator;
 use Cake\Collection\Collection;
+use Cake\Collection\CollectionInterface;
 
 /**
  * Creates an iterator from another iterator that will verify a condition on each
@@ -78,4 +80,41 @@ class StoppableIterator extends Collection
 
         return !$condition($current, $key, $this->_innerIterator);
     }
+
+    /**
+     * {@inheritDoc}
+     *
+     * We perform here some stricness analysys so that the
+     * iterator logic is bypassed entirely.
+     *
+     * @return \Iterator
+     */
+    public function unwrap()
+    {
+        $iterator = $this->_innerIterator;
+
+        if ($iterator instanceof CollectionInterface) {
+            $iterator = $iterator->unwrap();
+        }
+
+        if (!$iterator instanceof ArrayIterator) {
+            return $this;
+        }
+
+        // ArrayIterator can be traversed strictly.
+        // Let's do that for performance gains
+
+        $callback = $this->_condition;
+        $res = [];
+        $stop = false;
+
+        foreach ($iterator as $k => $v) {
+            if ($callback($v, $k, $iterator)) {
+                break;
+            }
+            $res[$k] = $v;
+        }
+
+        return new ArrayIterator($res);
+    }
 }

tests/TestCase/Collection/CollectionTest.php

@@ -95,14 +95,23 @@ class CollectionTest extends TestCase
         $collection->each($callable);
     }
 
+    public function filterProvider()
+    {
+        $items = [1, 2, 0, 3, false, 4, null, 5, ''];
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
     /**
      * Test filter() with no callback.
      *
+     * @dataProvider filterProvider
      * @return void
      */
-    public function testFilterNoCallback()
+    public function testFilterNoCallback($items)
     {
-        $items = [1, 2, 0, 3, false, 4, null, 5, ''];
         $collection = new Collection($items);
         $result = $collection->filter()->toArray();
         $expected = [1, 2, 3, 4, 5];
@@ -298,13 +307,28 @@ class CollectionTest extends TestCase
     }
 
     /**
+     * Provider for some simple tests
+     *
+     * @return array
+     */
+    public function simpleProvider()
+    {
+        $items = ['a' => 1, 'b' => 2, 'c' => 3, 'd' => 4];
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
      * Tests map
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testMap()
+    public function testMap($items)
     {
-        $items = ['a' => 1, 'b' => 2, 'c' => 3];
         $collection = new Collection($items);
         $map = $collection->map(function ($v, $k, $it) use ($collection) {
             $this->assertSame($collection->getInnerIterator(), $it);
@@ -312,17 +336,17 @@ class CollectionTest extends TestCase
             return $v * $v;
         });
         $this->assertInstanceOf('Cake\Collection\Iterator\ReplaceIterator', $map);
-        $this->assertEquals(['a' => 1, 'b' => 4, 'c' => 9], iterator_to_array($map));
+        $this->assertEquals(['a' => 1, 'b' => 4, 'c' => 9, 'd' => 16], iterator_to_array($map));
     }
 
     /**
      * Tests reduce with initial value
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testReduceWithInitialValue()
+    public function testReduceWithInitialValue($items)
     {
-        $items = ['a' => 1, 'b' => 2, 'c' => 3];
         $collection = new Collection($items);
         $callable = $this->getMockBuilder(\StdClass::class)
             ->setMethods(['__invoke'])
@@ -340,17 +364,21 @@ class CollectionTest extends TestCase
             ->method('__invoke')
             ->with(13, 3, 'c')
             ->will($this->returnValue(16));
-        $this->assertEquals(16, $collection->reduce($callable, 10));
+        $callable->expects($this->at(3))
+            ->method('__invoke')
+            ->with(16, 4, 'd')
+            ->will($this->returnValue(20));
+        $this->assertEquals(20, $collection->reduce($callable, 10));
     }
 
     /**
      * Tests reduce without initial value
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testReduceWithoutInitialValue()
+    public function testReduceWithoutInitialValue($items)
     {
-        $items = ['a' => 1, 'b' => 2, 'c' => 3, 'd' => 4];
         $collection = new Collection($items);
         $callable = $this->getMockBuilder(\StdClass::class)
             ->setMethods(['__invoke'])
@@ -372,13 +400,28 @@ class CollectionTest extends TestCase
     }
 
     /**
+     * Provider for some extract tests
+     *
+     * @return array
+     */
+    public function extractProvider()
+    {
+        $items = [['a' => ['b' => ['c' => 1]]], 2];
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
      * Tests extract
      *
+     * @dataProvider extractProvider
      * @return void
      */
-    public function testExtract()
+    public function testExtract($items)
     {
-        $items = [['a' => ['b' => ['c' => 1]]], 2];
         $collection = new Collection($items);
         $map = $collection->extract('a.b.c');
         $this->assertInstanceOf('Cake\Collection\Iterator\ExtractIterator', $map);
@@ -386,17 +429,32 @@ class CollectionTest extends TestCase
     }
 
     /**
-     * Tests sort
+     * Provider for some sort tests
      *
-     * @return void
+     * @return array
      */
-    public function testSortString()
+    public function sortProvider()
     {
         $items = [
             ['a' => ['b' => ['c' => 4]]],
             ['a' => ['b' => ['c' => 10]]],
             ['a' => ['b' => ['c' => 6]]]
         ];
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
+     * Tests sort
+     *
+     * @dataProvider sortProvider
+     * @return void
+     */
+    public function testSortString($items)
+    {
         $collection = new Collection($items);
         $map = $collection->sortBy('a.b.c');
         $this->assertInstanceOf('Cake\Collection\Collection', $map);
@@ -411,18 +469,39 @@ class CollectionTest extends TestCase
     /**
      * Tests max
      *
+     * @dataProvider sortProvider
      * @return void
      */
-    public function testMax()
+    public function testMax($items)
     {
-        $items = [
-            ['a' => ['b' => ['c' => 4]]],
-            ['a' => ['b' => ['c' => 10]]],
-            ['a' => ['b' => ['c' => 6]]]
-        ];
         $collection = new Collection($items);
         $this->assertEquals(['a' => ['b' => ['c' => 10]]], $collection->max('a.b.c'));
+    }
+
+    /**
+     * Tests max
+     *
+     * @dataProvider sortProvider
+     * @return void
+     */
+    public function testMaxCallback($items)
+    {
+        $collection = new Collection($items);
+        $callback = function ($e) {
+            return $e['a']['b']['c'] * - 1;
+        };
+        $this->assertEquals(['a' => ['b' => ['c' => 4]]], $collection->max($callback));
+    }
 
+    /**
+     * Tests max
+     *
+     * @dataProvider sortProvider
+     * @return void
+     */
+    public function testMaxCallable($items)
+    {
+        $collection = new Collection($items);
         $callback = function ($e) {
             return $e['a']['b']['c'] * - 1;
         };
@@ -432,31 +511,42 @@ class CollectionTest extends TestCase
     /**
      * Tests min
      *
+     * @dataProvider sortProvider
      * @return void
      */
-    public function testMin()
+    public function testMin($items)
     {
-        $items = [
-            ['a' => ['b' => ['c' => 4]]],
-            ['a' => ['b' => ['c' => 10]]],
-            ['a' => ['b' => ['c' => 6]]]
-        ];
         $collection = new Collection($items);
         $this->assertEquals(['a' => ['b' => ['c' => 4]]], $collection->min('a.b.c'));
     }
 
     /**
-     * Tests groupBy
+     * Provider for some groupBy tests
      *
-     * @return void
+     * @return array
      */
-    public function testGroupBy()
+    public function groupByProvider()
     {
         $items = [
             ['id' => 1, 'name' => 'foo', 'parent_id' => 10],
             ['id' => 2, 'name' => 'bar', 'parent_id' => 11],
             ['id' => 3, 'name' => 'baz', 'parent_id' => 10],
         ];
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
+     * Tests groupBy
+     *
+     * @dataProvider groupByProvider
+     * @return void
+     */
+    public function testGroupBy($items)
+    {
         $collection = new Collection($items);
         $grouped = $collection->groupBy('parent_id');
         $expected = [
@@ -470,7 +560,26 @@ class CollectionTest extends TestCase
         ];
         $this->assertEquals($expected, iterator_to_array($grouped));
         $this->assertInstanceOf('Cake\Collection\Collection', $grouped);
+    }
 
+    /**
+     * Tests groupBy
+     *
+     * @dataProvider groupByProvider
+     * @return void
+     */
+    public function testGroupByCallback($items)
+    {
+        $collection = new Collection($items);
+        $expected = [
+            10 => [
+                ['id' => 1, 'name' => 'foo', 'parent_id' => 10],
+                ['id' => 3, 'name' => 'baz', 'parent_id' => 10],
+            ],
+            11 => [
+                ['id' => 2, 'name' => 'bar', 'parent_id' => 11],
+            ]
+        ];
         $grouped = $collection->groupBy(function ($element) {
             return $element['parent_id'];
         });
@@ -504,17 +613,32 @@ class CollectionTest extends TestCase
     }
 
     /**
-     * Tests indexBy
+     * Provider for some indexBy tests
      *
-     * @return void
+     * @return array
      */
-    public function testIndexBy()
+    public function indexByProvider()
     {
         $items = [
             ['id' => 1, 'name' => 'foo', 'parent_id' => 10],
             ['id' => 2, 'name' => 'bar', 'parent_id' => 11],
             ['id' => 3, 'name' => 'baz', 'parent_id' => 10],
         ];
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
+     * Tests indexBy
+     *
+     * @dataProvider indexByProvider
+     * @return void
+     */
+    public function testIndexBy($items)
+    {
         $collection = new Collection($items);
         $grouped = $collection->indexBy('id');
         $expected = [
@@ -524,10 +648,25 @@ class CollectionTest extends TestCase
         ];
         $this->assertEquals($expected, iterator_to_array($grouped));
         $this->assertInstanceOf('Cake\Collection\Collection', $grouped);
+    }
 
+    /**
+     * Tests indexBy
+     *
+     * @dataProvider indexByProvider
+     * @return void
+     */
+    public function testIndexByCallback($items)
+    {
+        $collection = new Collection($items);
         $grouped = $collection->indexBy(function ($element) {
             return $element['id'];
         });
+        $expected = [
+            1 => ['id' => 1, 'name' => 'foo', 'parent_id' => 10],
+            3 => ['id' => 3, 'name' => 'baz', 'parent_id' => 10],
+            2 => ['id' => 2, 'name' => 'bar', 'parent_id' => 11],
+        ];
         $this->assertEquals($expected, iterator_to_array($grouped));
     }
 
@@ -555,24 +694,35 @@ class CollectionTest extends TestCase
     /**
      * Tests countBy
      *
+     * @dataProvider groupByProvider
      * @return void
      */
-    public function testCountBy()
+    public function testCountBy($items)
     {
-        $items = [
-            ['id' => 1, 'name' => 'foo', 'parent_id' => 10],
-            ['id' => 2, 'name' => 'bar', 'parent_id' => 11],
-            ['id' => 3, 'name' => 'baz', 'parent_id' => 10],
-        ];
         $collection = new Collection($items);
         $grouped = $collection->countBy('parent_id');
         $expected = [
             10 => 2,
             11 => 1
         ];
-        $this->assertEquals($expected, iterator_to_array($grouped));
+        $result = iterator_to_array($grouped);
         $this->assertInstanceOf('Cake\Collection\Collection', $grouped);
+        $this->assertEquals($expected, $result);
+    }
 
+    /**
+     * Tests countBy
+     *
+     * @dataProvider groupByProvider
+     * @return void
+     */
+    public function testCountByCallback($items)
+    {
+        $expected = [
+            10 => 2,
+            11 => 1
+        ];
+        $collection = new Collection($items);
         $grouped = $collection->countBy(function ($element) {
             return $element['parent_id'];
         });
@@ -582,32 +732,32 @@ class CollectionTest extends TestCase
     /**
      * Tests shuffle
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testShuffle()
+    public function testShuffle($data)
     {
-        $data = [1, 2, 3, 4];
         $collection = (new Collection($data))->shuffle();
-        $this->assertCount(count($data), iterator_to_array($collection));
-
-        foreach ($collection as $value) {
-            $this->assertContains($value, $data);
-        }
+        $result = $collection->toArray();
+        $this->assertCount(4, $result);
+        $this->assertContains(1, $result);
+        $this->assertContains(2, $result);
+        $this->assertContains(3, $result);
+        $this->assertContains(4, $result);
     }
 
     /**
      * Tests sample
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testSample()
+    public function testSample($data)
     {
-        $data = [1, 2, 3, 4];
-        $collection = (new Collection($data))->sample(2);
-        $this->assertCount(2, iterator_to_array($collection));
-
-        foreach ($collection as $value) {
-            $this->assertContains($value, $data);
+        $result = (new Collection($data))->sample(2)->toArray();
+        $this->assertCount(2, $result);
+        foreach ($result as $number) {
+            $this->assertContains($number, [1, 2, 3, 4]);
         }
     }
 
@@ -618,7 +768,7 @@ class CollectionTest extends TestCase
      */
     public function testToArray()
     {
-        $data = [1, 2, 3, 4];
+        $data = ['a' => 1, 'b' => 2, 'c' => 3, 'd' => 4];
         $collection = new Collection($data);
         $this->assertEquals($data, $collection->toArray());
     }
@@ -626,13 +776,13 @@ class CollectionTest extends TestCase
     /**
      * Test toList method
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testToList()
+    public function testToList($data)
     {
-        $data = [100 => 1, 300 => 2, 500 => 3, 1 => 4];
         $collection = new Collection($data);
-        $this->assertEquals(array_values($data), $collection->toList());
+        $this->assertEquals([1, 2, 3, 4], $collection->toList());
     }
 
     /**
@@ -1299,18 +1449,42 @@ class CollectionTest extends TestCase
     }
 
     /**
-     * Tests the sumOf method
+     * Provider for sumOf tests
      *
-     * @return void
+     * @return array
      */
-    public function testSumOf()
+    public function sumOfProvider()
     {
         $items = [
             ['invoice' => ['total' => 100]],
             ['invoice' => ['total' => 200]]
         ];
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
+     * Tests the sumOf method
+     *
+     * @dataProvider sumOfProvider
+     * @return void
+     */
+    public function testSumOf($items)
+    {
         $this->assertEquals(300, (new Collection($items))->sumOf('invoice.total'));
+    }
 
+    /**
+     * Tests the sumOf method
+     *
+     * @dataProvider sumOfProvider
+     * @return void
+     */
+    public function testSumOfCallable($items)
+    {
         $sum = (new Collection($items))->sumOf(function ($v) {
             return $v['invoice']['total'] * 2;
         });
@@ -1320,15 +1494,15 @@ class CollectionTest extends TestCase
     /**
      * Tests the stopWhen method with a callable
      *
+     * @dataProvider simpleProvider
      * @return void
      */
-    public function testStopWhenCallable()
+    public function testStopWhenCallable($items)
     {
-        $items = [10, 20, 40, 10, 5];
         $collection = (new Collection($items))->stopWhen(function ($v) {
-            return $v > 20;
+            return $v > 3;
         });
-        $this->assertEquals([10, 20], $collection->toArray());
+        $this->assertEquals(['a' => 1, 'b' => 2, 'c' => 3], $collection->toArray());
     }
 
     /**
@@ -1746,13 +1920,29 @@ class CollectionTest extends TestCase
     }
 
     /**
+     * Provider for some chunk tests
+     *
+     * @return array
+     */
+    public function chunkProvider()
+    {
+        $items = range(1, 10);
+
+        return [
+            'array' => [$items],
+            'iterator' => [$this->yieldItems($items)]
+        ];
+    }
+
+    /**
      * Tests the chunk method with exact chunks
      *
+     * @dataProvider chunkProvider
      * @return void
      */
-    public function testChunk()
+    public function testChunk($items)
     {
-        $collection = new Collection(range(1, 10));
+        $collection = new Collection($items);
         $chunked = $collection->chunk(2)->toList();
         $expected = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]];
         $this->assertEquals($expected, $chunked);
@@ -2028,4 +2218,17 @@ class CollectionTest extends TestCase
 
         $collection->transpose();
     }
+
+    /**
+     * Yields all the elmentes as passed
+     *
+     * @param array $itmes the elements to be yielded
+     * @return void
+     */
+    function yieldItems(array $items)
+    {
+        foreach ($items as $k => $v) {
+            yield $k => $v;
+        }
+    }
 }