refactor: Last operation improvements.

loophp · Sep 25, 2020 · 88350e2 · 88350e2
1 parent 13b5d6d
commit 88350e2
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 5 deletions.
diff --git a/spec/loophp/collection/CollectionSpec.php b/spec/loophp/collection/CollectionSpec.php
@@ -961,6 +961,19 @@ public function it_can_get_the_last_item(): void
         $this::fromIterable([])
             ->last()
             ->shouldIterateAs([]);
+
+        $input = [
+            ['a'],
+            ['b', 'a'],
+            ['c', 'b', 'a'],
+            ['d', 'c', 'b', 'a'],
+        ];
+
+        $this::fromIterable($input)
+            ->last()
+            ->shouldIterateAs([
+                3 => ['d', 'c', 'b', 'a'],
+            ]);
     }
 
     public function it_can_group()

diff --git a/src/Operation/Last.php b/src/Operation/Last.php
@@ -4,6 +4,7 @@
 
 namespace loophp\collection\Operation;
 
+use CachingIterator;
 use Closure;
 use Generator;
 use Iterator;
@@ -31,14 +32,17 @@ static function (Iterator $iterator): Generator {
                     return yield from [];
                 }
 
-                $key = $iterator->key();
-                $current = $iterator->current();
+                $cachingIterator = new CachingIterator($iterator, CachingIterator::FULL_CACHE);
 
-                for (; $iterator->valid(); $iterator->next()) {
-                    $key = $iterator->key();
-                    $current = $iterator->current();
+                while ($iterator->valid()) {
+                    $cachingIterator->next();
                 }
 
+                /** @psalm-var TKey $key */
+                $key = $cachingIterator->key();
+                /** @psalm-var T $current */
+                $current = $cachingIterator->current();
+
                 return yield $key => $current;
             };
     }