feat: SimpleCachingIteratorAggregate to fail when iterating multipl…

…e times.
loophp · Dec 19, 2021 · c4e92f5 · c4e92f5
1 parent 85a5047
commit c4e92f5
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 8 deletions.
diff --git a/src/SimpleCachingIteratorAggregate.php b/src/SimpleCachingIteratorAggregate.php
@@ -22,6 +22,8 @@
  */
 final class SimpleCachingIteratorAggregate implements IteratorAggregate
 {
+    private bool $isDone = false;
+
     /**
      * @var CachingIterator<TKey, T>
      */
@@ -43,11 +45,18 @@ public function __construct(Iterator $iterator)
      */
     public function getIterator(): Traversable
     {
-        /** @var iterable<TKey, T> $traversable */
-        $traversable = $this->iterator->getInnerIterator()->valid()
-            ? $this->iterator
-            : $this->iterator->getCache();
+        if (true === $this->isDone) {
+            return yield from $this->iterator->getCache();
+        }
+
+        $this->iterator->next();
+
+        for (; $this->iterator->valid(); $this->iterator->next()) {
+            $this->iterator->current();
+        }
+
+        $this->isDone = true;
 
-        yield from $traversable;
+        return yield from $this->iterator->getCache();
     }
 }
diff --git a/tests/benchmarks/GeneratorCacheBench.php b/tests/benchmarks/GeneratorCacheBench.php
@@ -58,14 +58,24 @@ private function test(Traversable $input): void
     {
         $a = $b = [];
 
+        $i = 0;
+
         foreach ($input as $key => $value) {
-            $a[] = [$key, $value];
+            $a = [$key, $value];
+
+            if (2 === $i++) {
+                break;
+            }
         }
 
         foreach ($input as $key => $value) {
             $b[] = [$key, $value];
         }
 
+        foreach ($input as $key => $value) {
+            $c[] = [$key, $value];
+        }
+
         if ($a !== $b) {
             throw new Exception('$a !== $b');
         }

diff --git a/tests/unit/CachingIteratorAggregateTest.php b/tests/unit/CachingIteratorAggregateTest.php
@@ -35,15 +35,37 @@ public function testWithAGenerator(): void
 
         $a = $b = [];
 
+        $i = 0;
+
         foreach ($iterator as $key => $value) {
             $a[] = [$key, $value];
+            $d[] = [$key, $value];
+
+            if (2 === $i++) {
+                break;
+            }
         }
 
         foreach ($iterator as $key => $value) {
             $b[] = [$key, $value];
+            $d[] = [$key, $value];
+        }
+
+        foreach ($iterator as $key => $value) {
+            $c[] = [$key, $value];
+            $d[] = [$key, $value];
         }
 
         $expected = [
+            [true, true],
+            [false, false],
+            [['a'], ['a']],
+            [true, true],
+            [false, false],
+            [['a'], ['a']],
+            [0, 'a'],
+            [1, 'b'],
+            [2, 'c'],
             [true, true],
             [false, false],
             [['a'], ['a']],
@@ -52,7 +74,7 @@ public function testWithAGenerator(): void
             [2, 'c'],
         ];
 
-        self::assertEquals($expected, $a);
-        self::assertTrue($a === $b);
+        self::assertTrue($b === $c);
+        self::assertEquals($expected, $d);
     }
 }