Use argument types as parameter types for inline closures when assigned

The functionality was introduced in #1628. It works. But as soon as you use an inline assign
expression it breaks. Let's support this case too.

Sometimes, you want to call something inline and also use the callback later.

Author: ruudk

src/Parser/ArrowFunctionArgVisitor.php

@@ -13,13 +13,28 @@ final class ArrowFunctionArgVisitor extends NodeVisitorAbstract
 
 	public function enterNode(Node $node): ?Node
 	{
-		if ($node instanceof Node\Expr\FuncCall && $node->name instanceof Node\Expr\ArrowFunction && !$node->isFirstClassCallable()) {
-			$args = $node->getArgs();
+		if (!$node instanceof Node\Expr\FuncCall) {
+			return null;
+		}
+
+		if ($node->isFirstClassCallable()) {
+			return null;
+		}
 
-			if (count($args) > 0) {
-				$node->name->setAttribute(self::ATTRIBUTE_NAME, $args);
-			}
+		if ($node->name instanceof Node\Expr\Assign && $node->name->expr instanceof Node\Expr\ArrowFunction) {
+			$arrow = $node->name->expr;
+		} elseif ($node->name instanceof Node\Expr\ArrowFunction) {
+			$arrow = $node->name;
+		} else {
+			return null;
 		}
+
+		$args = $node->getArgs();
+
+		if (count($args) > 0) {
+			$arrow->setAttribute(self::ATTRIBUTE_NAME, $args);
+		}
+
 		return null;
 	}
 

src/Parser/ClosureArgVisitor.php

@@ -13,13 +13,28 @@ final class ClosureArgVisitor extends NodeVisitorAbstract
 
 	public function enterNode(Node $node): ?Node
 	{
-		if ($node instanceof Node\Expr\FuncCall && $node->name instanceof Node\Expr\Closure && !$node->isFirstClassCallable()) {
-			$args = $node->getArgs();
+		if (!$node instanceof Node\Expr\FuncCall) {
+			return null;
+		}
+
+		if ($node->isFirstClassCallable()) {
+			return null;
+		}
 
-			if (count($args) > 0) {
-				$node->name->setAttribute(self::ATTRIBUTE_NAME, $args);
-			}
+		if ($node->name instanceof Node\Expr\Assign && $node->name->expr instanceof Node\Expr\Closure) {
+			$closure = $node->name->expr;
+		} elseif ($node->name instanceof Node\Expr\Closure) {
+			$closure = $node->name;
+		} else {
+			return null;
 		}
+
+		$args = $node->getArgs();
+
+		if (count($args) > 0) {
+			$closure->setAttribute(self::ATTRIBUTE_NAME, $args);
+		}
+
 		return null;
 	}
 

tests/PHPStan/Analyser/nsrt/arrow-function-argument-type.php

@@ -21,6 +21,8 @@ public function doFoo(int $integer, array $array, ?string $nullableString)
 		(fn($a, $b, $c) => assertType('array{int, array{a: int}, string|null}', [$a, $b, $c]))($integer, $array, $nullableString);
 
 		(fn($a, $b, $c = null) => assertType('array{int, array{a: int}, mixed}', [$a, $b, $c]))($integer, $array);
+
+		($callback = fn($context) => assertType('int', $context))($integer);
 	}
 
 }

tests/PHPStan/Analyser/nsrt/closure-argument-type.php

@@ -35,6 +35,10 @@ public function doFoo(int $integer, array $array, ?string $nullableString)
 			assertType('array{a: int}', $context2);
 			assertType('mixed', $context3);
 		})($integer, $array);
+
+		($callback = function($context) {
+			assertType('int', $context);
+		})($integer);
 	}
 
 }