Update Flatten operation.

loophp · Aug 27, 2020 · f5ceeb8 · f5ceeb8
1 parent 6fce398
commit f5ceeb8
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 34 deletions.
diff --git a/spec/loophp/collection/CollectionSpec.php b/spec/loophp/collection/CollectionSpec.php
@@ -692,43 +692,49 @@ static function ($value) {
 
     public function it_can_flatten(): void
     {
-        $input = [];
+        $input = [
+            ['a', 'b', 'c'],
+            'd',
+            ['d', ['e', 'f']],
+        ];
 
-        for ($j = 0; 5 > $j; ++$j) {
-            $items = [];
+        $output = static function (): Generator {
+            yield 0 => 'a';
 
-            for ($i = 0; 2 > $i; ++$i) {
-                $items[] = $j * 2 + $i;
-            }
-            $input[] = $items;
-        }
+            yield 1 => 'b';
 
-        $input = array_pad([], 5, $input);
+            yield 2 => 'c';
 
-        $output = [];
+            yield 1 => 'd';
 
-        for ($i = 0; 5 > $i; ++$i) {
-            $output = array_merge($output, range(0, 9));
-        }
+            yield 0 => 'd';
+
+            yield 0 => 'e';
+
+            yield 1 => 'f';
+        };
 
         $this::fromIterable($input)
             ->flatten()
-            ->shouldIterateAs($output);
+            ->shouldIterateAs($output());
 
-        $output = [];
+        $output = static function (): Generator {
+            yield 0 => 'a';
+
+            yield 1 => 'b';
 
-        $j = 0;
+            yield 2 => 'c';
 
-        for ($i = 0; 25 > $i; ++$i) {
-            $output[] = [
-                $j++ % 10,
-                $j++ % 10,
-            ];
-        }
+            yield 1 => 'd';
+
+            yield 0 => 'd';
+
+            yield 1 => ['e', 'f'];
+        };
 
         $this::fromIterable($input)
             ->flatten(1)
-            ->shouldIterateAs($output);
+            ->shouldIterateAs($output());
     }
 
     public function it_can_flip(): void

diff --git a/src/Operation/Flatten.php b/src/Operation/Flatten.php
@@ -4,15 +4,13 @@
 
 namespace loophp\collection\Operation;
 
-use ArrayIterator;
 use Closure;
 use Generator;
 use Iterator;
 use loophp\collection\Contract\Operation;
+use loophp\collection\Iterator\IterableIterator;
 use loophp\collection\Transformation\Run;
 
-use function is_array;
-
 /**
  * @psalm-template TKey
  * @psalm-template TKey of array-key
@@ -34,18 +32,20 @@ public function __invoke(): Closure
              * @psalm-return Generator<int, T>
              */
             static function (Iterator $iterator, int $depth): Generator {
-                foreach ($iterator as $value) {
+                foreach ($iterator as $key => $value) {
                     if (false === is_iterable($value)) {
-                        yield $value;
+                        yield $key => $value;
                     } elseif (1 === $depth) {
                         /** @psalm-var T $subValue */
-                        foreach ($value as $subValue) {
-                            yield $subValue;
+                        foreach ($value as $subKey => $subValue) {
+                            yield $subKey => $subValue;
                         }
-                    } elseif (is_array($value)) {
-                        /** @psalm-var T $subValue */
-                        foreach ((new Run(new Flatten($depth - 1)))(new ArrayIterator($value)) as $subValue) {
-                            yield $subValue;
+                    } else {
+                        /** @psalm-var IterableIterator<TKey, T> $iterable */
+                        $iterable = (new Run(new Flatten($depth - 1)))(new IterableIterator($value));
+
+                        foreach ($iterable as $subKey => $subValue) {
+                            yield $subKey => $subValue;
                         }
                     }
                 }