Apply FunctionParameterClosureTypeExtension override to the in-function-call stack so the closure's inferred return type matches its analysed body

src/Analyser/MutatingScope.php

@@ -57,6 +57,7 @@
 use PHPStan\Reflection\InitializerExprContext;
 use PHPStan\Reflection\InitializerExprTypeResolver;
 use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\Native\NativeParameterReflection;
 use PHPStan\Reflection\ParameterReflection;
 use PHPStan\Reflection\Php\PhpFunctionFromParserNodeReflection;
 use PHPStan\Reflection\Php\PhpMethodFromParserNodeReflection;
@@ -1398,6 +1399,56 @@ public function pushInFunctionCall($reflection, ?ParameterReflection $parameter,
 		return $functionScope;
 	}
 
+	/**
+	 * Overrides the type of the parameter on top of the in-function-call stack.
+	 *
+	 * Used when a FunctionParameterClosureTypeExtension changes the type of a
+	 * closure/callable parameter, so that the closure's return type inferred via
+	 * ClosureTypeResolver (which reads the parameter type off this stack) matches
+	 * the parameter type the closure body is actually analysed with.
+	 */
+	public function replaceInFunctionCallStackParameterType(Type $type): self
+	{
+		if (count($this->inFunctionCallsStack) === 0) {
+			return $this;
+		}
+
+		$stack = $this->inFunctionCallsStack;
+		$lastIndex = count($stack) - 1;
+		[$reflection, $parameter] = $stack[$lastIndex];
+		if ($parameter === null) {
+			return $this;
+		}
+
+		$stack[$lastIndex] = [$reflection, new NativeParameterReflection(
+			$parameter->getName(),
+			$parameter->isOptional(),
+			$type,
+			$parameter->passedByReference(),
+			$parameter->isVariadic(),
+			$parameter->getDefaultValue(),
+		)];
+
+		return $this->scopeFactory->create(
+			$this->context,
+			$this->isDeclareStrictTypes(),
+			$this->getFunction(),
+			$this->getNamespace(),
+			$this->expressionTypes,
+			$this->nativeExpressionTypes,
+			$this->conditionalExpressions,
+			$this->inClosureBindScopeClasses,
+			$this->anonymousFunctionReflection,
+			$this->isInFirstLevelStatement(),
+			$this->currentlyAssignedExpressions,
+			$this->currentlyAllowedUndefinedExpressions,
+			$stack,
+			$this->afterExtractCall,
+			$this->parentScope,
+			$this->nativeTypesPromoted,
+		);
+	}
+
 	public function popInFunctionCall(): self
 	{
 		$stack = $this->inFunctionCallsStack;

src/Analyser/NodeScopeResolver.php

@@ -3645,6 +3645,7 @@ public function processArgs(
 
 					if ($overwritingParameterType !== null) {
 						$parameterType = $overwritingParameterType;
+						$scopeToPass = $scopeToPass->replaceInFunctionCallStackParameterType($overwritingParameterType);
 					}
 				}
 
@@ -3711,6 +3712,7 @@ public function processArgs(
 
 					if ($overwritingParameterType !== null) {
 						$parameterType = $overwritingParameterType;
+						$scopeToPass = $scopeToPass->replaceInFunctionCallStackParameterType($overwritingParameterType);
 					}
 				}
 

tests/PHPStan/Rules/Functions/ArrowFunctionReturnTypeRuleTest.php

@@ -77,4 +77,9 @@ public function testBugFunctionMethodConstants(): void
 		$this->analyse([__DIR__ . '/data/bug-anonymous-function-method-constant.php'], []);
 	}
 
+	public function testBug14914(): void
+	{
+		$this->analyse([__DIR__ . '/data/bug-14914.php'], []);
+	}
+
 }

tests/PHPStan/Rules/Functions/ClosureReturnTypeParameterClosureExtensionRuleTest.php

@@ -0,0 +1,44 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Rules\Functions;
+
+use PHPStan\Rules\FunctionReturnTypeCheck;
+use PHPStan\Rules\Rule;
+use PHPStan\Rules\RuleLevelHelper;
+use PHPStan\Testing\RuleTestCase;
+
+/**
+ * @extends RuleTestCase<ClosureReturnTypeRule>
+ */
+class ClosureReturnTypeParameterClosureExtensionRuleTest extends RuleTestCase
+{
+
+	protected function getRule(): Rule
+	{
+		return new ClosureReturnTypeRule(new FunctionReturnTypeCheck(
+			new RuleLevelHelper(
+				self::createReflectionProvider(),
+				checkNullables: true,
+				checkThisOnly: false,
+				checkUnionTypes: true,
+				checkExplicitMixed: false,
+				checkImplicitMixed: false,
+				checkBenevolentUnionTypes: false,
+				discoveringSymbolsTip: true,
+			),
+		));
+	}
+
+	public function testRule(): void
+	{
+		$this->analyse([__DIR__ . '/data/closure-return-type-parameter-closure-extension.php'], []);
+	}
+
+	public static function getAdditionalConfigFiles(): array
+	{
+		return [
+			__DIR__ . '/data/closure-return-type-parameter-closure-extension.neon',
+		];
+	}
+
+}

tests/PHPStan/Rules/Functions/ClosureReturnTypeRuleTest.php

@@ -149,4 +149,9 @@ public function testBug13964(): void
 		$this->analyse([__DIR__ . '/data/bug-13964.php'], []);
 	}
 
+	public function testBug14914(): void
+	{
+		$this->analyse([__DIR__ . '/data/bug-14914.php'], []);
+	}
+
 }

tests/PHPStan/Rules/Functions/data/bug-14914.php

@@ -0,0 +1,30 @@
+<?php // lint >= 8.0
+
+declare(strict_types = 1);
+
+namespace Bug14914;
+
+function doFoo(): void
+{
+	preg_replace_callback(
+		'/a|(?<b>b)/',
+		function (array $match): string {
+			if ($match['b'] !== null) {
+				return 'aa';
+			}
+			return 'possible?';
+		},
+		'abcd',
+		flags: PREG_UNMATCHED_AS_NULL,
+	);
+}
+
+function doBar(): void
+{
+	preg_replace_callback(
+		'/a|(?<b>b)/',
+		fn (array $match) => $match['b'] !== null ? 'aa' : 'possible?',
+		'abcd',
+		flags: PREG_UNMATCHED_AS_NULL,
+	);
+}

tests/PHPStan/Rules/Functions/data/closure-return-type-parameter-closure-extension.neon

@@ -0,0 +1,15 @@
+services:
+	-
+		class: ClosureReturnTypeParameterClosureExtension\FunctionParameterClosureTypeExtension
+		tags:
+			- phpstan.functionParameterClosureTypeExtension
+
+	-
+		class: ClosureReturnTypeParameterClosureExtension\MethodParameterClosureTypeExtension
+		tags:
+			- phpstan.methodParameterClosureTypeExtension
+
+	-
+		class: ClosureReturnTypeParameterClosureExtension\StaticMethodParameterClosureTypeExtension
+		tags:
+			- phpstan.staticMethodParameterClosureTypeExtension

tests/PHPStan/Rules/Functions/data/closure-return-type-parameter-closure-extension.php

@@ -0,0 +1,130 @@
+<?php declare(strict_types = 1);
+
+namespace ClosureReturnTypeParameterClosureExtension;
+
+use Closure;
+use PhpParser\Node\Expr\FuncCall;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Expr\StaticCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\Native\NativeParameterReflection;
+use PHPStan\Reflection\ParameterReflection;
+use PHPStan\Type\ClosureType;
+use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
+use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\NullType;
+use PHPStan\Type\StringType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+
+/**
+ * Overrides the callback parameter with an array shape whose `k` offset is
+ * nullable. The declared callback signature is `Closure(array{k: string}): string`,
+ * so without the override the `$x['k'] !== null` branch would look always-true
+ * and the closure's inferred return type would collapse to a single constant.
+ */
+function nullableOffsetClosureType(ParameterReflection $parameter): ClosureType
+{
+	$builder = ConstantArrayTypeBuilder::createEmpty();
+	$builder->setOffsetValueType(new ConstantStringType('k'), TypeCombinator::union(new StringType(), new NullType()));
+
+	return new ClosureType(
+		[
+			new NativeParameterReflection($parameter->getName(), $parameter->isOptional(), $builder->getArray(), $parameter->passedByReference(), $parameter->isVariadic(), $parameter->getDefaultValue()),
+		],
+		new StringType(),
+	);
+}
+
+class FunctionParameterClosureTypeExtension implements \PHPStan\Type\FunctionParameterClosureTypeExtension
+{
+
+	public function isFunctionSupported(FunctionReflection $functionReflection, ParameterReflection $parameter): bool
+	{
+		return $functionReflection->getName() === 'ClosureReturnTypeParameterClosureExtension\functionWithClosure' && $parameter->getName() === 'callback';
+	}
+
+	public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, ParameterReflection $parameter, Scope $scope): ?Type
+	{
+		return nullableOffsetClosureType($parameter);
+	}
+
+}
+
+class MethodParameterClosureTypeExtension implements \PHPStan\Type\MethodParameterClosureTypeExtension
+{
+
+	public function isMethodSupported(MethodReflection $methodReflection, ParameterReflection $parameter): bool
+	{
+		return $methodReflection->getDeclaringClass()->getName() === Foo::class && $methodReflection->getName() === 'methodWithClosure' && $parameter->getName() === 'callback';
+	}
+
+	public function getTypeFromMethodCall(MethodReflection $methodReflection, MethodCall $methodCall, ParameterReflection $parameter, Scope $scope): ?Type
+	{
+		return nullableOffsetClosureType($parameter);
+	}
+
+}
+
+class StaticMethodParameterClosureTypeExtension implements \PHPStan\Type\StaticMethodParameterClosureTypeExtension
+{
+
+	public function isStaticMethodSupported(MethodReflection $methodReflection, ParameterReflection $parameter): bool
+	{
+		return $methodReflection->getDeclaringClass()->getName() === Foo::class && $methodReflection->getName() === 'staticMethodWithClosure' && $parameter->getName() === 'callback';
+	}
+
+	public function getTypeFromStaticMethodCall(MethodReflection $methodReflection, StaticCall $methodCall, ParameterReflection $parameter, Scope $scope): ?Type
+	{
+		return nullableOffsetClosureType($parameter);
+	}
+
+}
+
+class Foo
+{
+
+	/** @param Closure(array{k: string}): string $callback */
+	public function methodWithClosure(Closure $callback): void
+	{
+	}
+
+	/** @param Closure(array{k: string}): string $callback */
+	public static function staticMethodWithClosure(Closure $callback): void
+	{
+	}
+
+}
+
+/** @param Closure(array{k: string}): string $callback */
+function functionWithClosure(Closure $callback): void
+{
+}
+
+function test(Foo $foo): void
+{
+	functionWithClosure(function (array $x): string {
+		if ($x['k'] !== null) {
+			return 'a';
+		}
+		return 'b';
+	});
+
+	$foo->methodWithClosure(function (array $x): string {
+		if ($x['k'] !== null) {
+			return 'a';
+		}
+		return 'b';
+	});
+
+	Foo::staticMethodWithClosure(function (array $x): string {
+		if ($x['k'] !== null) {
+			return 'a';
+		}
+		return 'b';
+	});
+
+	functionWithClosure(fn (array $x): string => $x['k'] !== null ? 'a' : 'b');
+}