fix ArrayKeyFirstLastRector for return and other cases

packages/Php/src/Rector/FuncCall/ArrayKeyFirstLastRector.php

@@ -5,7 +5,9 @@ namespace Rector\Php\Rector\FuncCall;
 use PhpParser\Node;
 use PhpParser\Node\Expr\FuncCall;
 use PhpParser\Node\Name;
+use PhpParser\Node\Stmt\Expression;
 use Rector\NodeTypeResolver\Node\Attribute;
+use Rector\PhpParser\Node\BetterNodeFinder;
 use Rector\Rector\AbstractRector;
 use Rector\RectorDefinition\CodeSample;
 use Rector\RectorDefinition\RectorDefinition;
@@ -26,6 +28,16 @@ final class ArrayKeyFirstLastRector extends AbstractRector
         'end' => 'array_key_last',
     ];
 
+    /**
+     * @var BetterNodeFinder
+     */
+    private $betterNodeFinder;
+
+    public function __construct(BetterNodeFinder $betterNodeFinder)
+    {
+        $this->betterNodeFinder = $betterNodeFinder;
+    }
+
     public function getDefinition(): RectorDefinition
     {
         return new RectorDefinition(
@@ -69,32 +81,52 @@ CODE_SAMPLE
      */
     public function refactor(Node $node): ?Node
     {
-        $previousExpression = $node->getAttribute(Attribute::PREVIOUS_EXPRESSION);
-        if ($previousExpression === null) {
+        if (! $this->isNames($node, ['reset', 'end'])) {
             return null;
         }
 
-        $previousFuncCall = $previousExpression->expr;
-        if (! $previousFuncCall instanceof FuncCall) {
+        $nextExpression = $this->getNextExpression($node);
+        if ($nextExpression === null) {
             return null;
         }
 
-        if (! $this->isNames($previousFuncCall, array_keys($this->previousToNewFunctions))) {
+        $resetOrEndFuncCall = $node;
+
+        /** @var FuncCall|null $keyFuncCall */
+        $keyFuncCall = $this->betterNodeFinder->findFirst($nextExpression, function (Node $node) use (
+            $resetOrEndFuncCall
+        ) {
+            if (! $node instanceof FuncCall) {
+                return false;
+            }
+
+            if (! $this->isName($node, 'key')) {
+                return false;
+            }
+
+            return $this->areNodesEqual($resetOrEndFuncCall->args[0], $node->args[0]);
+        });
+
+        if ($keyFuncCall === null) {
             return null;
         }
 
-        if (! $this->isName($node, 'key')) {
-            return null;
-        }
+        $newName = $this->previousToNewFunctions[$this->getName($node)];
+        $keyFuncCall->name = new Name($newName);
 
-        if (! $this->areNodesEqual($previousFuncCall->args[0], $node->args[0])) {
-            return null;
-        }
-
-        $this->removeNode($previousFuncCall);
-
-        $node->name = new Name($this->previousToNewFunctions[$this->getName($previousFuncCall)]);
+        $this->removeNode($node);
 
         return $node;
     }
+
+    private function getNextExpression(Node $node): ?Node
+    {
+        /** @var Expression|null $currentExpression */
+        $currentExpression = $node->getAttribute(Attribute::CURRENT_EXPRESSION);
+        if ($currentExpression === null) {
+            return null;
+        }
+
+        return $currentExpression->getAttribute(Attribute::NEXT_NODE);
+    }
 }

packages/Php/tests/Rector/FuncCall/ArrayKeyFirstLastRector/ArrayKeyFirstLastRectorTest.php

@@ -10,9 +10,9 @@ final class ArrayKeyFirstLastRectorTest extends AbstractRectorTestCase
     public function test(): void
     {
         $this->doTestFiles([
-            __DIR__ . '/Fixture/fixture.php.inc',
-            __DIR__ . '/Fixture/fixture2.php.inc',
-            __DIR__ . '/Fixture/fixture3.php.inc',
+            __DIR__ . '/Fixture/array_key_first.php.inc',
+            __DIR__ . '/Fixture/array_key_last.php.inc',
+            __DIR__ . '/Fixture/both.php.inc',
         ]);
     }
 

packages/Php/tests/Rector/FuncCall/ArrayKeyFirstLastRector/Fixture/array_key_last.php.inc

@@ -2,6 +2,13 @@
 
 namespace Rector\Php\Tests\Rector\FuncCall\ArrayKeyFirstLastRector\Fixture;
 
+function someFunction()
+{
+    $items = [1, 2, 3];
+    end($items);
+    $lastKey = key($items);
+}
+
 function someFunction3()
 {
     if ($this) {
@@ -19,6 +26,12 @@ function someFunction3()
 
 namespace Rector\Php\Tests\Rector\FuncCall\ArrayKeyFirstLastRector\Fixture;
 
+function someFunction()
+{
+    $items = [1, 2, 3];
+    $lastKey = array_key_last($items);
+}
+
 function someFunction3()
 {
     if ($this) {

packages/Php/tests/Rector/FuncCall/ArrayKeyFirstLastRector/Fixture/both.php.inc

@@ -0,0 +1,39 @@
+<?php
+
+namespace Rector\Php\Tests\Rector\FuncCall\ArrayKeyFirstLastRector\Fixture;
+
+class Both
+{
+    public function getFirstKey($items)
+    {
+        reset($items);
+        return key($items);
+    }
+
+    public function getLastKey($items2)
+    {
+        end($items2);
+        return key($items2);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Php\Tests\Rector\FuncCall\ArrayKeyFirstLastRector\Fixture;
+
+class Both
+{
+    public function getFirstKey($items)
+    {
+        return array_key_first($items);
+    }
+
+    public function getLastKey($items2)
+    {
+        return array_key_last($items2);
+    }
+}
+
+?>

packages/Php/tests/Rector/FuncCall/ArrayKeyFirstLastRector/Fixture/fixture2.php.inc

@@ -1,24 +0,0 @@
-<?php
-
-namespace Rector\Php\Tests\Rector\FuncCall\ArrayKeyFirstLastRector\Fixture;
-
-function someFunction()
-{
-    $items = [1, 2, 3];
-    end($items);
-    $lastKey = key($items);
-}
-
-?>
------
-<?php
-
-namespace Rector\Php\Tests\Rector\FuncCall\ArrayKeyFirstLastRector\Fixture;
-
-function someFunction()
-{
-    $items = [1, 2, 3];
-    $lastKey = array_key_last($items);
-}
-
-?>