Infer Return Types for Arrow Functions #3376 (#3385)

* Infer Return Types for Arrow Functions #3376

- Made a small patch to check for closure or arrow
  function when attempting to infer the functions
  params
- Added new isExprClosureLike to start to consolidate
  all checks on closure/arrow fns

Signed-off-by: RJ Garcia <ragboyjr@icloud.com>

* Use better check

* Remove unused method

Co-authored-by: Matthew Brown <github@muglug.com>

src/Psalm/Internal/Analyzer/Statements/ReturnAnalyzer.php

@@ -129,7 +129,9 @@ class ReturnAnalyzer
         if ($stmt->expr) {
             $context->inside_call = true;
 
-            if ($stmt->expr instanceof PhpParser\Node\Expr\Closure) {
+            if ($stmt->expr instanceof PhpParser\Node\Expr\Closure
+                || $stmt->expr instanceof PhpParser\Node\Expr\ArrowFunction
+            ) {
                 self::potentiallyInferTypesOnClosureFromParentReturnType(
                     $statements_analyzer,
                     $stmt->expr,
@@ -591,7 +593,7 @@ class ReturnAnalyzer
      */
     private static function potentiallyInferTypesOnClosureFromParentReturnType(
         StatementsAnalyzer $statements_analyzer,
-        PhpParser\Node\Expr\Closure $expr,
+        PhpParser\Node\FunctionLike $expr,
         Context $context
     ): void {
         // if not returning from inside of a function, return

tests/ReturnTypeTest.php

@@ -721,6 +721,22 @@ class ReturnTypeTest extends TestCase
                     '$res' => 'iterable<mixed, string>',
                 ],
             ],
+            'infersArrowClosureReturnTypes' => [
+                '<?php
+                    /**
+                     * @param Closure(int, int): bool $op
+                     * @return Closure(int): bool
+                     */
+                    function reflexive(Closure $op): Closure {
+                        return fn ($x) => $op($x, $x) === true;
+                    }
+
+                    $res = reflexive(fn(int $a, int $b): bool => $a === $b);
+                ',
+                'assertions' => [
+                    '$res' => 'Closure(int):bool',
+                ],
+            ],
             'infersClosureReturnTypesWithPartialTypehinting' => [
                 '<?php
                     /**