array_filter w/ arrow functions return type extension

Author: b1rdex

src/Type/Php/ArrayFilterFunctionReturnTypeReturnTypeExtension.php

@@ -2,6 +2,7 @@
 
 namespace PHPStan\Type\Php;
 
+use PhpParser\Node\Expr\ArrowFunction;
 use PhpParser\Node\Expr\Closure;
 use PhpParser\Node\Expr\FuncCall;
 use PhpParser\Node\Expr\Variable;
@@ -53,18 +54,29 @@ public function getTypeFromFunctionCall(FunctionReflection $functionReflection,
 				);
 			}
 
-			if ($flagArg === null && $callbackArg instanceof Closure && count($callbackArg->stmts) === 1) {
-				$statement = $callbackArg->stmts[0];
-				if ($statement instanceof Return_ && $statement->expr !== null && count($callbackArg->params) > 0) {
-					if (!$callbackArg->params[0]->var instanceof Variable || !is_string($callbackArg->params[0]->var->name)) {
+			if ($flagArg === null) {
+				$var = null;
+				$expr = null;
+				if ($callbackArg instanceof Closure && count($callbackArg->stmts) === 1 && count($callbackArg->params) > 0) {
+					$statement = $callbackArg->stmts[0];
+					if ($statement instanceof Return_ && $statement->expr !== null) {
+						$var = $callbackArg->params[0]->var;
+						$expr = $statement->expr;
+					}
+				} elseif ($callbackArg instanceof ArrowFunction && count($callbackArg->params) > 0) {
+					$var = $callbackArg->params[0]->var;
+					$expr = $callbackArg->expr;
+				}
+				if ($var !== null && $expr !== null) {
+					if (!$var instanceof Variable || !is_string($var->name)) {
 						throw new \PHPStan\ShouldNotHappenException();
 					}
-					$itemVariableName = $callbackArg->params[0]->var->name;
+					$itemVariableName = $var->name;
 					if (!$scope instanceof MutatingScope) {
 						throw new \PHPStan\ShouldNotHappenException();
 					}
 					$scope = $scope->assignVariable($itemVariableName, $itemType);
-					$scope = $scope->filterByTruthyValue($statement->expr);
+					$scope = $scope->filterByTruthyValue($expr);
 					$itemType = $scope->getVariableType($itemVariableName);
 				}
 			}

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -10585,6 +10585,15 @@ public function dataBug3986(): array
 		return $this->gatherAssertTypes(__DIR__ . '/data/bug-3986.php');
 	}
 
+	public function dataBug4188(): array
+	{
+		if (PHP_VERSION_ID < 70400) {
+			return [];
+		}
+
+		return $this->gatherAssertTypes(__DIR__ . '/data/bug-4188.php');
+	}
+
 	/**
 	 * @param string $file
 	 * @return array<string, mixed[]>
@@ -10787,6 +10796,7 @@ private function gatherAssertTypes(string $file): array
 	 * @dataProvider dataBug1209
 	 * @dataProvider dataBug2980
 	 * @dataProvider dataBug3986
+	 * @dataProvider dataBug4188
 	 * @param string $assertType
 	 * @param string $file
 	 * @param mixed ...$args

tests/PHPStan/Analyser/data/bug-4188.php

@@ -0,0 +1,40 @@
+<?php // lint >= 7.4
+
+namespace Bug4188;
+
+interface A {}
+interface B {}
+
+use function PHPStan\Analyser\assertType;
+
+class Test
+{
+	/** @param array<A|B> $data */
+	public function set(array $data): void
+	{
+		$filtered = array_filter(
+			$data,
+			function ($param): bool {
+				return $param instanceof B;
+			},
+		);
+		assertType('array<Bug4188\B>', $filtered);
+
+		$this->onlyB($filtered);
+	}
+
+	/** @param array<A|B> $data */
+	public function setShort(array $data): void
+	{
+		$filtered = array_filter(
+			$data,
+			fn($param): bool => $param instanceof B,
+		);
+		assertType('array<Bug4188\B>', $filtered);
+
+		$this->onlyB($filtered);
+	}
+
+	/** @param B[] $data */
+	public function onlyB(array $data): void {}
+}