Memoize generators in order to be able to rewind them

CachedIterator.php

<?php

use Iterator;
use IteratorAggregate;

final class CachedIterator implements IteratorAggregate
{
    private $cached;
    private $iterator;

    public function __construct(Iterator $iterator)
    {
        $this->iterator = $iterator;
    }

    public function getIterator()
    {
        if ($this->cached === null) {
            $this->iterator->rewind();
            $this->cached = [];
        } else {
            yield from $this->cached();
        }

        yield from $this->memoize(
            $this->iterate()
        );
    }

    private function cached()
    {
        foreach ($this->cached as list($key, $value)) {
            yield $key => $value;
        }
    }

    private function iterate()
    {
        while ($this->iterator->valid()) {
            yield $this->iterator->key() => $this->iterator->current();
            $this->iterator->next();
        }
    }

    private function memoize($iterator)
    {
        foreach ($iterator as $key => $value) {
            $this->cached[] = [ $key, $value ];
            yield $key => $value;
        }
    }
}

CachedIteratorTest.php

<?php

class CachedIteratorTest extends \PHPUnit_Framework_TestCase
{
    public function test_can_rewind()
    {
        $generator = function () {
            yield 1;
        };

        $iterator = new CachedIterator($generator());

        // First iteration uses generator
        $result = iterator_to_array($iterator);
        // Further iterations uses cached values
        $cachedResult = iterator_to_array($iterator);

        self::assertContains(1, $result);
        self::assertContains(1, $cachedResult);
    }

    public function test_inner_generator_is_called_only_once()
    {
        $enterIteration = 0;
        $leaveIteration = 0;
        $generator = function () use (&$enterIteration, &$leaveIteration) {
            ++$enterIteration;
            yield from [];
            ++$leaveIteration;
        };

        $iterator = new CachedIterator($generator());

        // Iterate many times
        iterator_to_array($iterator);
        iterator_to_array($iterator);
        iterator_to_array($iterator);
        iterator_to_array($iterator);

        self::assertEquals(1, $enterIteration);
        self::assertEquals(1, $leaveIteration);
    }

    public function test_allow_resume_of_partial_iteration()
    {
        $expected = [ 1, 2, 3 ];
        $generator = function () use ($expected) {
            yield from $expected;
        };

        $iterator = new CachedIterator($generator());

        // Iterate partially over the generator
        $value = null;
        foreach ($iterator as $value) {
            break;
        }

        // Full iteration returns partial cached values and resumes generator execution
        $result = iterator_to_array($iterator);
        // Next iteration uses cached values
        $cachedResult = iterator_to_array($iterator);

        self::assertEquals(1, $value);
        self::assertEquals($expected, $result);
        self::assertEquals($expected, $cachedResult);
    }

    public function test_allows_multiple_values_with_same_key()
    {
        $generator = function () {
            yield from [ 1, 2, 3 ];
            yield from [ 1, 2, 3 ];
            yield from [ 1, 2, 3 ];
        };

        $iterator = new CachedIterator($generator());

        $count = iterator_count($iterator);
        $cachedCount = iterator_count($iterator);

        self::assertEquals(9, $count);
        self::assertEquals(9, $cachedCount);
    }
}