Arrow function return type inference

ondrejmirtes · ondrejmirtes · commit 19482b75ead5 · 2020-03-17T10:02:35.000+01:00
diff --git a/src/Analyser/MutatingScope.php b/src/Analyser/MutatingScope.php
@@ -1183,9 +1183,15 @@ private function resolveType(Expr $node): Type
 				);
 			}
 
+			if ($node->returnType === null && $node instanceof Expr\ArrowFunction) {
+				$returnType = $this->getType($node->expr);
+			} else {
+				$returnType = $this->getFunctionType($node->returnType, $node->returnType === null, false);
+			}
+
 			return new ClosureType(
 				$parameters,
-				$this->getFunctionType($node->returnType, $node->returnType === null, false),
+				$returnType,
 				$isVariadic
 			);
 		} elseif ($node instanceof New_) {
diff --git a/tests/PHPStan/Analyser/NodeScopeResolverTest.php b/tests/PHPStan/Analyser/NodeScopeResolverTest.php
@@ -9806,6 +9806,15 @@ public function dataExtDs(): array
 		return $this->gatherAssertTypes(__DIR__ . '/data/ext-ds.php');
 	}
 
+	public function dataArrowFunctionReturnTypeInference(): array
+	{
+		if (PHP_VERSION_ID < 70400) {
+			return [];
+		}
+
+		return $this->gatherAssertTypes(__DIR__ . '/data/arrow-function-return-type.php');
+	}
+
 	/**
 	 * @dataProvider dataBug2574
 	 * @dataProvider dataBug2577
@@ -9839,6 +9848,7 @@ public function dataExtDs(): array
 	 * @dataProvider dataTypeChangeAfterArrayAccessAssignment
 	 * @dataProvider dataIteratorToArray
 	 * @dataProvider dataExtDs
+	 * @dataProvider dataArrowFunctionReturnTypeInference
 	 * @param ConstantStringType $expectedType
 	 * @param Type $actualType
 	 */
diff --git a/tests/PHPStan/Analyser/data/arrow-function-return-type.php b/tests/PHPStan/Analyser/data/arrow-function-return-type.php
@@ -0,0 +1,15 @@
+<?php
+
+namespace ArrowFunctionReturnTypeInference;
+
+use function PHPStan\Analyser\assertType;
+
+function (int $i): void {
+	$fn = fn () => $i;
+	assertType('int', $fn());
+};
+
+function (int $i): void {
+	$fn = fn (): string => $i;
+	assertType('string', $fn());
+};
