Fix return type inference from immediately invoked closure

Author: ondrejmirtes

src/Analyser/MutatingScope.php

@@ -50,6 +50,7 @@
 use PHPStan\Node\Printer\ExprPrinter;
 use PHPStan\Node\PropertyAssignNode;
 use PHPStan\Parser\ArrayMapArgVisitor;
+use PHPStan\Parser\ImmediatelyInvokedClosureVisitor;
 use PHPStan\Parser\NewAssignedToPropertyVisitor;
 use PHPStan\Parser\Parser;
 use PHPStan\Php\PhpVersion;
@@ -1374,11 +1375,16 @@ private function resolveType(string $exprString, Expr $node): Type
 
 			$callableParameters = null;
 			$arrayMapArgs = $node->getAttribute(ArrayMapArgVisitor::ATTRIBUTE_NAME);
+			$immediatelyInvokedArgs = $node->getAttribute(ImmediatelyInvokedClosureVisitor::ARGS_ATTRIBUTE_NAME);
 			if ($arrayMapArgs !== null) {
 				$callableParameters = [];
 				foreach ($arrayMapArgs as $funcCallArg) {
 					$callableParameters[] = new DummyParameter('item', $this->getType($funcCallArg->value)->getIterableValueType(), false, PassedByReference::createNo(), false, null);
 				}
+			} elseif ($immediatelyInvokedArgs !== null) {
+				foreach ($immediatelyInvokedArgs as $immediatelyInvokedArg) {
+					$callableParameters[] = new DummyParameter('item', $this->getType($immediatelyInvokedArg->value), false, PassedByReference::createNo(), false, null);
+				}
 			} else {
 				$inFunctionCallsStackCount = count($this->inFunctionCallsStack);
 				if ($inFunctionCallsStackCount > 0) {

src/Parser/ImmediatelyInvokedClosureVisitor.php

@@ -12,12 +12,18 @@ final class ImmediatelyInvokedClosureVisitor extends NodeVisitorAbstract
 {
 
 	public const ATTRIBUTE_NAME = 'isImmediatelyInvokedClosure';
+	public const ARGS_ATTRIBUTE_NAME = 'immediatelyInvokedClosureArgs';
 
 	#[Override]
 	public function enterNode(Node $node): ?Node
 	{
-		if ($node instanceof Node\Expr\FuncCall && $node->name instanceof Node\Expr\Closure && !$node->isFirstClassCallable()) {
+		if (
+			$node instanceof Node\Expr\FuncCall
+			&& ($node->name instanceof Node\Expr\Closure || $node->name instanceof Node\Expr\ArrowFunction)
+			&& !$node->isFirstClassCallable()
+		) {
 			$node->name->setAttribute(self::ATTRIBUTE_NAME, true);
+			$node->name->setAttribute(self::ARGS_ATTRIBUTE_NAME, $node->args);
 		}
 
 		return null;

tests/PHPStan/Analyser/nsrt/pipe-operator.php

@@ -81,4 +81,125 @@ public function doBaz2(): void
 		(fn (int $x) => assertType('int<1, max>', $x))(doFoo());
 	}
 
+	/**
+	 * @param array<positive-int> $ints
+	 */
+	public function doArrayMap(array $ints): void
+	{
+		assertType('array<int<1, max>>', $ints |>
+			(fn ($x) => array_map(static fn ($i) => $i, $x)),
+		);
+		assertType('array<int<1, max>>', $ints |>
+			(fn ($x) => array_map(static fn (int $i) => $i, $x)),
+		);
+		assertType('array<int<1, max>>', $ints |>
+			(fn (array $x) => array_map(static fn ($i) => $i, $x)),
+		);
+		assertType('array<int<1, max>>', $ints |>
+			(fn (array $x) => array_map(static fn (int $i) => $i, $x)),
+		);
+
+		assertType('array<\'foo\'>', $ints |>
+			(fn (array $x) => array_map(static fn (int $i) => 'foo', $x)),
+		);
+
+		assertType('array<int<1, max>>', $ints |> function ($x) {
+			assertType('array<int<1, max>>', $x);
+			return array_map(function ($i) {
+				assertType('int<1, max>', $i);
+
+				return $i;
+			}, $x);
+		});
+
+		assertType('array<int<1, max>>', $ints |> function (array $x) {
+		assertType('array<int<1, max>>', $x);
+		return array_map(function (int $i) {
+			assertType('int<1, max>', $i);
+
+			return $i;
+		}, $x);
+	});
+	}
+
+	/**
+	 * @param array<positive-int> $ints
+	 */
+	public function doArrayFilter(array $ints): void
+	{
+		assertType('array<int<1, max>>', $ints |> function ($x) {
+			assertType('array<int<1, max>>', $x);
+			return array_filter($x, function ($i) {
+				assertType('int<1, max>', $i);
+
+				return true;
+			});
+		});
+		assertType('array<int<1, max>>', $ints |> function (array $x) {
+			assertType('array<int<1, max>>', $x);
+			return array_filter($x, function (int $i) {
+				assertType('int<1, max>', $i);
+
+				return true;
+			});
+		});
+		assertType('array<int<1, max>>', $ints |> fn ($x) => array_filter($x, function ($i) {
+			assertType('int<1, max>', $i);
+
+			return true;
+		}));
+		assertType('array<int<1, max>>', $ints |> fn (array $x) => array_filter($x, function (int $i) {
+			assertType('int<1, max>', $i);
+
+			return true;
+		}));
+		assertType('array<0|1|2, 1|2|3>', (function (array $x) {
+			assertType('array{1, 2, 3}', $x);
+			return array_filter($x, function (int $i) {
+				assertType('1|2|3', $i);
+
+				return true;
+			});
+		})([1, 2, 3]));
+		assertType('array<0|1|2, 1|2|3>', (function ($x) {
+			assertType('array{1, 2, 3}', $x);
+			return array_filter($x, function ($i) {
+				assertType('1|2|3', $i);
+
+				return true;
+			});
+		})([1, 2, 3]));
+		assertType('array<0|1|2, 1|2|3>', (function (array $x) {
+			assertType('array{1, 2, 3}', $x);
+			return array_filter($x, function (int $i) {
+				assertType('1|2|3', $i);
+
+				return true;
+			});
+		})([1, 2, 3]));
+		assertType('array<0|1|2, 1|2|3>', (function ($x) {
+			assertType('array{1, 2, 3}', $x);
+			return array_filter($x, function ($i) {
+				assertType('1|2|3', $i);
+
+				return true;
+			});
+		})([1, 2, 3]));
+		assertType('array<0|1|2, 1|2|3>', (fn (array $x) => array_filter($x, function (int $i) {
+			assertType('1|2|3', $i);
+
+			return true;
+		}))([1, 2, 3]));
+		assertType('array<0|1|2, 1|2|3>', (fn ($x) => array_filter($x, function ($i) {
+			assertType('1|2|3', $i);
+
+			return true;
+		}))([1, 2, 3]));
+		assertType('1', (fn ($x) => $x)(1));
+		assertType('1', (function ($x) {
+			assertType('1', $x);
+			return $x;
+		})(1));
+	}
+
 }