fix nesting iterations over the same Accumulate instance

Author: Baptouuuu

proofs/accumulate.php

@@ -0,0 +1,71 @@
+<?php
+
+use Innmind\Immutable\Accumulate;
+use Innmind\BlackBox\Set;
+
+return static function() {
+    yield proof(
+        "Accumulate partial access doesn't skip values",
+        given(
+            Set\Sequence::of(Set\Type::any()),
+            Set\Sequence::of(Set\Type::any()),
+        ),
+        static function($assert, $prefix, $suffix) {
+            $accumulate = new Accumulate((static function() use ($prefix, $suffix) {
+                yield from $prefix;
+                yield from $suffix;
+            })());
+
+            foreach ($accumulate as $i => $_) {
+                if ($i === (\count($prefix) - 1)) {
+                    break;
+                }
+            }
+
+            $all = [];
+
+            foreach ($accumulate as $value) {
+                $all[] = $value;
+            }
+
+            $assert->same(
+                [...$prefix, ...$suffix],
+                $all,
+            );
+        },
+    );
+
+    yield test(
+        'Accumuluate nested iteration',
+        static function($assert) {
+            $accumulate = new Accumulate((static function() {
+                yield 1;
+                yield 2;
+                yield 3;
+            })());
+
+            $pairs = [];
+
+            foreach ($accumulate as $i) {
+                foreach ($accumulate as $j) {
+                    $pairs[] = [$i, $j];
+                }
+            }
+
+            $assert->same(
+                [
+                    [1, 1],
+                    [1, 2],
+                    [1, 3],
+                    [2, 1],
+                    [2, 2],
+                    [2, 3],
+                    [3, 1],
+                    [3, 2],
+                    [3, 3],
+                ],
+                $pairs,
+            );
+        },
+    );
+};

src/Accumulate.php

@@ -18,6 +18,8 @@ final class Accumulate implements \Iterator
     private \Generator $generator;
     /** @var list<S> */
     private array $values = [];
+    /** @var list<int<0, max>> */
+    private array $cursors = [];
     private bool $started = false;
 
     /**
@@ -70,6 +72,11 @@ public function next(): void
 
     public function rewind(): void
     {
+        if (\is_int($cursor = \key($this->values))) {
+            /** @psalm-suppress InaccessibleProperty */
+            $this->cursors[] = $cursor;
+        }
+
         /** @psalm-suppress InaccessibleProperty */
         $this->started = true;
         /** @psalm-suppress InaccessibleProperty */
@@ -84,10 +91,20 @@ public function valid(): bool
         $valid = !$this->reachedCacheEnd() || $this->generator->valid();
 
         if (!$valid) {
-            // once the "true" end has been reached we automatically rewind this
-            // iterator so it is always in a clean state
-            /** @psalm-suppress UnusedMethodCall */
-            $this->rewind();
+            // once the "true" end has been reached we automatically rewind to
+            // the previous cursor position to allow correct iteration when
+            // nesting loops on the same iterator
+            /** @psalm-suppress InaccessibleProperty */
+            $previousCursor = \array_pop($this->cursors);
+            /** @psalm-suppress InaccessibleProperty */
+            \reset($this->values);
+
+            if (\is_int($previousCursor)) {
+                while (\is_int(\key($this->values)) && \key($this->values) !== $previousCursor) {
+                    /** @psalm-suppress InaccessibleProperty */
+                    \next($this->values);
+                }
+            }
         }
 
         return $valid;