Add caching iterator

Author: alcaeus

src/CachingIterator.php

@@ -0,0 +1,119 @@
+<?php
+
+namespace MongoDB;
+
+class CachingIterator implements \Iterator, \Countable
+{
+    /**
+     * @var \Traversable
+     */
+    private $iterator;
+
+    /**
+     * @var array
+     */
+    private $items = [];
+
+    /**
+     * @var bool
+     */
+    private $iteratorExhausted = false;
+
+    /**
+     * @param \Traversable $iterator
+     */
+    public function __construct(\Traversable $iterator)
+    {
+        $this->iterator = $this->wrapTraversable($iterator);
+        $this->storeCurrentItem();
+    }
+
+    /**
+     * @return int
+     */
+    public function count()
+    {
+        $this->exhaustIterator();
+        return count($this->items);
+    }
+
+    /**
+     * @return mixed
+     */
+    public function current()
+    {
+        return current($this->items);
+    }
+
+    /**
+     * @return mixed
+     */
+    public function key()
+    {
+        return key($this->items);
+    }
+
+    /**
+     * @return void
+     */
+    public function next()
+    {
+        if (! $this->iteratorExhausted) {
+            $this->iterator->next();
+            $this->storeCurrentItem();
+        }
+
+        next($this->items);
+    }
+
+    /**
+     * @return void
+     */
+    public function rewind()
+    {
+        $this->exhaustIterator();
+        reset($this->items);
+    }
+
+    /**
+     * @return bool
+     */
+    public function valid()
+    {
+        return $this->key() !== null;
+    }
+
+    /**
+     * Ensures the original iterator is fully consumed and all items cached
+     */
+    private function exhaustIterator()
+    {
+        while (!$this->iteratorExhausted) {
+            $this->next();
+        }
+    }
+
+    /**
+     * Stores the current item
+     */
+    private function storeCurrentItem()
+    {
+        if (null === $key = $this->iterator->key()) {
+            return;
+        }
+
+        $this->items[$key] = $this->iterator->current();
+    }
+
+    /**
+     * @param \Traversable $traversable
+     * @return \Generator
+     */
+    private function wrapTraversable(\Traversable $traversable)
+    {
+        foreach ($traversable as $key => $value) {
+            yield $key => $value;
+        }
+        $this->iteratorExhausted = true;
+    }
+}

tests/CachingIteratorTest.php

@@ -0,0 +1,63 @@
+<?php
+
+namespace MongoDB\Tests;
+
+use MongoDB\CachingIterator;
+
+class CachingIteratorTest extends \PHPUnit_Framework_TestCase
+{
+    /**
+     * Sanity check for all following tests
+     * @expectedException \Exception
+     * @expectedExceptionMessage Cannot traverse an already closed generator
+     */
+    public function testTraverseGeneratorConsumesIt()
+    {
+        $iterator = $this->getTraversable([1, 2, 3]);
+        $this->assertSame([1, 2, 3], iterator_to_array($iterator));
+        $this->assertSame([1, 2, 3], iterator_to_array($iterator));
+    }
+
+    public function testIterateOverItems()
+    {
+        $iterator = new CachingIterator($this->getTraversable([1, 2, 3]));
+
+        $expectedKey = 0;
+        $expectedItem = 1;
+        foreach ($iterator as $key => $item) {
+            $this->assertSame($expectedKey++, $key);
+            $this->assertSame($expectedItem++, $item);
+        }
+        $this->assertFalse($iterator->valid());
+    }
+
+    public function testIteratePartiallyThenRewind()
+    {
+        $iterator = new CachingIterator($this->getTraversable([1, 2, 3]));
+
+        $this->assertSame(1, $iterator->current());
+        $iterator->next();
+
+        $this->assertSame([1, 2, 3], iterator_to_array($iterator));
+    }
+
+    public function testCount()
+    {
+        $iterator = new CachingIterator($this->getTraversable([1, 2, 3]));
+        $this->assertCount(3, $iterator);
+    }
+
+    public function testCountAfterPartiallyIterating()
+    {
+        $iterator = new CachingIterator($this->getTraversable([1, 2, 3]));
+        $iterator->next();
+        $this->assertCount(3, $iterator);
+    }
+
+    private function getTraversable($items)
+    {
+        foreach ($items as $item) {
+            yield $item;
+        }
+    }
+}