From 1d598ea68907ac94a12c0fe39246e55f1fae2270 Mon Sep 17 00:00:00 2001
From: Tomas Votruba <tomas.vot@gmail.com>
Date: Wed, 14 Aug 2024 20:12:30 +0200
Subject: [PATCH] add arrow function support

---
 .../Fixture/include_arrow_function.php.inc    | 30 ++++++++++++
 .../Fixture/multiple_types.php.inc            | 46 +++++++++++++++++++
 ...turnArrayDocblockBasedOnArrayMapRector.php |  9 ++--
 3 files changed, 81 insertions(+), 4 deletions(-)
 create mode 100644 rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/include_arrow_function.php.inc
 create mode 100644 rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/multiple_types.php.inc

diff --git a/rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/include_arrow_function.php.inc b/rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/include_arrow_function.php.inc
new file mode 100644
index 00000000000..daaa287326e
--- /dev/null
+++ b/rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/include_arrow_function.php.inc
@@ -0,0 +1,30 @@
+<?php
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\AddReturnArrayDocblockBasedOnArrayMapRector\Fixture;
+
+final class IncludeArrowFunction
+{
+    public function run(array $items)
+    {
+        return array_map(fn ($item): int => 1000, $items);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\AddReturnArrayDocblockBasedOnArrayMapRector\Fixture;
+
+final class IncludeArrowFunction
+{
+    /**
+     * @return int[]
+     */
+    public function run(array $items)
+    {
+        return array_map(fn ($item): int => 1000, $items);
+    }
+}
+
+?>
diff --git a/rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/multiple_types.php.inc b/rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/multiple_types.php.inc
new file mode 100644
index 00000000000..044c3d40c4a
--- /dev/null
+++ b/rules-tests/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector/Fixture/multiple_types.php.inc
@@ -0,0 +1,46 @@
+<?php
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\AddReturnArrayDocblockBasedOnArrayMapRector\Fixture;
+
+final class MultipleTypes
+{
+    public function run(array $items)
+    {
+        if (mt_rand(0, 1)) {
+            return array_map(function ($item): int {
+                return $item;
+            }, $items);
+        }
+
+        return array_map(function ($item): string {
+            return '1000';
+        }, $items);
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Tests\TypeDeclaration\Rector\ClassMethod\AddReturnArrayDocblockBasedOnArrayMapRector\Fixture;
+
+final class MultipleTypes
+{
+    /**
+     * @return int[]|string[]
+     */
+    public function run(array $items)
+    {
+        if (mt_rand(0, 1)) {
+            return array_map(function ($item): int {
+                return $item;
+            }, $items);
+        }
+
+        return array_map(function ($item): string {
+            return '1000';
+        }, $items);
+    }
+}
+
+?>
diff --git a/rules/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector.php b/rules/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector.php
index 2f68bfe9d5a..80b4497f152 100644
--- a/rules/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector.php
+++ b/rules/TypeDeclaration/Rector/ClassMethod/AddReturnArrayDocblockBasedOnArrayMapRector.php
@@ -5,8 +5,10 @@
 namespace Rector\TypeDeclaration\Rector\ClassMethod;
 
 use PhpParser\Node;
+use PhpParser\Node\Expr\ArrowFunction;
 use PhpParser\Node\Expr\Closure;
 use PhpParser\Node\Expr\FuncCall;
+use PhpParser\Node\FunctionLike;
 use PhpParser\Node\Identifier;
 use PhpParser\Node\Stmt\ClassMethod;
 use PhpParser\Node\Stmt\Function_;
@@ -31,7 +33,6 @@ public function __construct(
         private readonly BetterNodeFinder $betterNodeFinder,
         private readonly ReturnAnalyzer $returnAnalyzer,
         private readonly StaticTypeMapper $staticTypeMapper,
-        //        private readonly ReturnPhpDocDecorator $returnPhpDocDecorator,
         private readonly TypeFactory $typeFactory,
         private readonly PhpDocTypeChanger $phpDocTypeChanger,
         private readonly PhpDocInfoFactory $phpDocInfoFactory,
@@ -108,7 +109,7 @@ public function refactor(Node $node): null|Function_|ClassMethod
             }
 
             $arrayMapClosure = $this->matchArrayMapClosure($returnScoped->expr);
-            if (! $arrayMapClosure instanceof Closure) {
+            if (! $arrayMapClosure instanceof FunctionLike) {
                 return null;
             }
 
@@ -141,7 +142,7 @@ private function hasNonArrayReturnType(ClassMethod|Function_ $functionLike): boo
         return $functionLike->returnType->toLowerString() !== 'array';
     }
 
-    private function matchArrayMapClosure(FuncCall $funcCall): ?Closure
+    private function matchArrayMapClosure(FuncCall $funcCall): Closure|ArrowFunction|null
     {
         if (! $this->isName($funcCall, 'array_map')) {
             return null;
@@ -153,7 +154,7 @@ private function matchArrayMapClosure(FuncCall $funcCall): ?Closure
 
         // lets infer strict array_map() type
         $firstArg = $funcCall->getArgs()[0];
-        if (! $firstArg->value instanceof Closure) {
+        if (! $firstArg->value instanceof Closure && ! $firstArg->value instanceof ArrowFunction) {
             return null;
         }