Remember function return type extensions (#4644)

Author: staabm

src/Analyser/MutatingScope.php

@@ -2501,16 +2501,13 @@ static function (Node $node, Scope $scope) use ($arrowScope, &$arrowFunctionImpu
 
 	private function getDynamicFunctionReturnType(FuncCall $normalizedNode, FunctionReflection $functionReflection): ?Type
 	{
-		foreach ($this->dynamicReturnTypeExtensionRegistry->getDynamicFunctionReturnTypeExtensions() as $dynamicFunctionReturnTypeExtension) {
-			if (!$dynamicFunctionReturnTypeExtension->isFunctionSupported($functionReflection)) {
-				continue;
-			}
-
+		foreach ($this->dynamicReturnTypeExtensionRegistry->getDynamicFunctionReturnTypeExtensions($functionReflection) as $dynamicFunctionReturnTypeExtension) {
 			$resolvedType = $dynamicFunctionReturnTypeExtension->getTypeFromFunctionCall(
 				$functionReflection,
 				$normalizedNode,
 				$this,
 			);
+
 			if ($resolvedType !== null) {
 				return $resolvedType;
 			}

src/Type/DynamicReturnTypeExtensionRegistry.php

@@ -2,6 +2,7 @@
 
 namespace PHPStan\Type;
 
+use PHPStan\Reflection\FunctionReflection;
 use PHPStan\Reflection\ReflectionProvider;
 use function array_merge;
 use function strtolower;
@@ -15,6 +16,9 @@ final class DynamicReturnTypeExtensionRegistry
 	/** @var DynamicStaticMethodReturnTypeExtension[][]|null */
 	private ?array $dynamicStaticMethodReturnTypeExtensionsByClass = null;
 
+	/** @var array<string, list<DynamicFunctionReturnTypeExtension>> */
+	private array $dynamicReturnTypeExtensionsByFunction = [];
+
 	/**
 	 * @param DynamicMethodReturnTypeExtension[] $dynamicMethodReturnTypeExtensions
 	 * @param DynamicStaticMethodReturnTypeExtension[] $dynamicStaticMethodReturnTypeExtensions
@@ -88,9 +92,23 @@ private function getDynamicExtensionsForType(array $extensions, string $classNam
 	/**
 	 * @return DynamicFunctionReturnTypeExtension[]
 	 */
-	public function getDynamicFunctionReturnTypeExtensions(): array
+	public function getDynamicFunctionReturnTypeExtensions(FunctionReflection $functionReflection): array
 	{
-		return $this->dynamicFunctionReturnTypeExtensions;
+		$functionName = $functionReflection->getName();
+		if (isset($this->dynamicReturnTypeExtensionsByFunction[$functionName])) {
+			return $this->dynamicReturnTypeExtensionsByFunction[$functionName];
+		}
+
+		$supportedFunctions = [];
+		foreach ($this->dynamicFunctionReturnTypeExtensions as $dynamicFunctionReturnTypeExtension) {
+			if (!$dynamicFunctionReturnTypeExtension->isFunctionSupported($functionReflection)) {
+				continue;
+			}
+
+			$supportedFunctions[] = $dynamicFunctionReturnTypeExtension;
+		}
+
+		return $this->dynamicReturnTypeExtensionsByFunction[$functionName] = $supportedFunctions;
 	}
 
 }