refactor: Improve Intersperse operation.

Avoid using local variables and use iterators.
loophp · Dec 11, 2021 · a208d4c · a208d4c
1 parent c109e6e
commit a208d4c
Showing 1 changed file with 14 additions and 3 deletions.
diff --git a/src/Operation/Intersperse.php b/src/Operation/Intersperse.php
@@ -9,8 +9,11 @@
 
 namespace loophp\collection\Operation;
 
+use AppendIterator;
+use ArrayIterator;
 use Closure;
 use Generator;
+use InfiniteIterator;
 use InvalidArgumentException;
 use Iterator;
 
@@ -62,15 +65,23 @@ static function (Iterator $iterator) use ($element, $atEvery, $startAt): Generat
                                 );
                             }
 
-                            $current = 0;
+                            $intersperse = new AppendIterator();
+                            $intersperse->append(
+                                new ArrayIterator(array_fill(0, $startAt, 1))
+                            );
+                            $intersperse->append(
+                                new InfiniteIterator(new ArrayIterator(range(0, $atEvery - 1)))
+                            );
+                            $intersperse->rewind();
 
                             foreach ($iterator as $key => $value) {
-                                if ($startAt === $current++) {
-                                    $startAt += $atEvery;
+                                if (0 === $intersperse->current()) {
                                     yield $element;
                                 }
 
                                 yield $key => $value;
+
+                                $intersperse->next();
                             }
                         };
     }