From 19482b75ead51f11da186eb0982a8bf14fee0f81 Mon Sep 17 00:00:00 2001
From: Ondrej Mirtes <ondrej@mirtes.cz>
Date: Tue, 17 Mar 2020 09:56:46 +0100
Subject: [PATCH] Arrow function return type inference

---
 src/Analyser/MutatingScope.php                    |  8 +++++++-
 tests/PHPStan/Analyser/NodeScopeResolverTest.php  | 10 ++++++++++
 .../Analyser/data/arrow-function-return-type.php  | 15 +++++++++++++++
 3 files changed, 32 insertions(+), 1 deletion(-)
 create mode 100644 tests/PHPStan/Analyser/data/arrow-function-return-type.php

diff --git a/src/Analyser/MutatingScope.php b/src/Analyser/MutatingScope.php
index b10db397f2..18e5e41921 100644
--- 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
index b3624d62b7..5efa42bbb4 100644
--- 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
new file mode 100644
index 0000000000..6cd49d1a74
--- /dev/null
+++ 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());
+};