Use the enclosing body's send-decided template argument resolutions inside nested closure bodies

src/Analyser/ExprHandler/Helper/ClosureTypeResolver.php

@@ -229,6 +229,9 @@ static function (Node $node, Scope $scope) use ($arrowScope, &$arrowFunctionImpu
 		$closureImpurePoints = [];
 		$invalidateExpressions = [];
 
+		// a body walked while the enclosing body observes keeps its unresolved
+		// template arguments for the enclosing body to resolve
+		$resolveTemplateArguments = $scope->getCurrentTemplateArgumentFrame()?->isObserving() !== true;
 		try {
 			$walkStorage = new ExpressionResultStorage();
 			$closureScope->pushExpressionResultStorage($walkStorage);
@@ -268,7 +271,7 @@ static function (Node $node, Scope $scope) use ($arrowScope, &$arrowFunctionImpu
 				}
 
 				$closureYieldStatements[] = [$node, $scope];
-			}, StatementContext::createTopLevel())->toPublic();
+			}, StatementContext::createTopLevel($resolveTemplateArguments))->toPublic();
 		} finally {
 			$closureScope->popExpressionResultStorage();
 			self::$resolveClosureTypeDepth--;

src/Analyser/ExprHandler/NewHandler.php

@@ -783,9 +783,13 @@ private function unresolvedArgumentList(ClassReflection $classReflection, New_ $
 			if (!$type instanceof TemplateType) {
 				return $type;
 			}
-			if ($synthetic || ($allowUnresolved && $frame->isObserving())) {
+			if ($synthetic) {
 				return new UnresolvedTemplateArgumentType($site, $type, null);
 			}
+			if ($allowUnresolved && $frame->isObserving()) {
+				// a send in a body enclosing this one decided the site
+				return $frame->resolveSentFromEnclosingBody($site, $type->getName()) ?? new UnresolvedTemplateArgumentType($site, $type, null);
+			}
 
 			return $frame->resolveOrUnconstrained($site, $type);
 		}));
@@ -797,7 +801,8 @@ private function rekeyParentTemplateArgument(UnresolvedTemplateArgumentType $typ
 			return $type->getInitialType() ?? $template->getDefault() ?? $template->getBound();
 		}
 		if ($allowUnresolved && $frame->isObserving()) {
-			return $type->withSite($site, $template);
+			// a send in a body enclosing this one decided the site
+			return $frame->resolveSentFromEnclosingBody($site, $template->getName()) ?? $type->withSite($site, $template);
 		}
 
 		return $frame->resolve($site, $template->getName()) ?? $type->getInitialType() ?? $frame->resolveOrUnconstrained($site, $template);

src/Analyser/Generics/TemplateArgumentFrame.php

@@ -68,6 +68,7 @@ public static function returnTypeOfCall(ParametersAcceptor $acceptor, MutatingSc
 	 * @param array<int, array{Expr\Closure, int, bool}> $byRefSites spl_object_id() of the closure node => the node, its statement index, whether every invocation was seen
 	 * @param array<int, array{Expr\Variable, Type, Type}> $staticVariableTypes spl_object_id() of a `static` variable node => the node, its phpdoc and native type - see StaticVariableInference
 	 * @param array<int, array{Node\Stmt\Static_, array<string, ConditionalExpressionHolder[]>}> $staticVariableConditionalExpressions spl_object_id() of the last statement of a run of `static` statements => the statement, the conditional expressions between the run's variables - see StaticVariableInference::getRuns()
+	 * @param array<string, true> $sentResolutionKeys the keys of the resolutions a send decided - see resolveSentFromEnclosingBody()
 	 */
 	public function __construct(
 		private readonly ?self $parent,
@@ -80,6 +81,7 @@ public function __construct(
 		private readonly array $byRefSites = [],
 		private readonly array $staticVariableTypes = [],
 		private readonly array $staticVariableConditionalExpressions = [],
+		private readonly array $sentResolutionKeys = [],
 	)
 	{
 	}
@@ -131,6 +133,7 @@ public function withStaticVariableTypes(array $staticVariableTypes, array $state
 			$this->byRefSites,
 			$staticVariableTypes,
 			$conditionalExpressions,
+			$this->sentResolutionKeys,
 		);
 	}
 
@@ -190,6 +193,7 @@ public function withObservedClosures(array $closureResolutions, array $closureSi
 			byRefSites: $byRefSites,
 			staticVariableTypes: $this->staticVariableTypes,
 			staticVariableConditionalExpressions: $this->staticVariableConditionalExpressions,
+			sentResolutionKeys: $this->sentResolutionKeys,
 		);
 	}
 
@@ -331,6 +335,26 @@ public function resolve(Expr $site, string $templateName): ?Type
 		return null;
 	}
 
+	/**
+	 * The resolution a body enclosing this one gave a site of a closure body it
+	 * encloses, when a send in it decided the argument - the closure's own walk
+	 * cannot see that send. An argument the enclosing body merely inferred is
+	 * left to the closure's walk, which infers it from its own state.
+	 */
+	public function resolveSentFromEnclosingBody(Expr $site, string $templateName): ?Type
+	{
+		$key = self::key($site, $templateName);
+		if (isset($this->sentResolutionKeys[$key], $this->resolutions[$key])) {
+			return $this->resolutions[$key];
+		}
+
+		if ($this->parent !== null) {
+			return $this->parent->resolveSentFromEnclosingBody($site, $templateName);
+		}
+
+		return null;
+	}
+
 	/**
 	 * Distinguishes cache entries computed while observing from those computed
 	 * with the resolutions installed (the closure type cache keys on scope state

src/Analyser/Generics/TemplateArgumentResolver.php

@@ -58,14 +58,14 @@ public function resolve(TemplateArgumentConstraints $constraints, ?TemplateArgum
 				TemplateArgumentStats::increment('closureObservationPasses');
 			}
 
-			return new TemplateArgumentFrame($parent, array_diff_key($resolutions, $closureKeys), $templateSiteStatementIndexes, $closureSignatureBody, $closureSignatureStmts, observingClosures: true);
+			return new TemplateArgumentFrame($parent, array_diff_key($resolutions, $closureKeys), $templateSiteStatementIndexes, $closureSignatureBody, $closureSignatureStmts, observingClosures: true, sentResolutionKeys: $this->collectSentResolutionKeys($observations));
 		}
 
 		if (TemplateArgumentStats::$enabled && count($settledClosureSites) > 0) {
 			TemplateArgumentStats::increment('closureSitesSettled', count($settledClosureSites));
 		}
 
-		return new TemplateArgumentFrame($parent, $resolutions, $this->collectSiteStatementIndexes($siteIndexes, $settledClosureSites), $closureSignatureBody, $closureSignatureStmts, $settledClosureSites, byRefSites: $this->collectByRefSites($observations, $siteIndexes));
+		return new TemplateArgumentFrame($parent, $resolutions, $this->collectSiteStatementIndexes($siteIndexes, $settledClosureSites), $closureSignatureBody, $closureSignatureStmts, $settledClosureSites, byRefSites: $this->collectByRefSites($observations, $siteIndexes), sentResolutionKeys: $this->collectSentResolutionKeys($observations));
 	}
 
 	/**
@@ -225,6 +225,23 @@ private function collectByRefSites(array $observations, array $siteIndexes): arr
 		return $byRefSites;
 	}
 
+	/**
+	 * @param array<string, array{marker: UnresolvedTemplateArgumentType, initial: Type|null, sends: list<array{Type, TemplateTypeVariance}>, lowerBounds: list<Type>, unconstrainingSend: bool}> $observations
+	 * @return array<string, true>
+	 */
+	private function collectSentResolutionKeys(array $observations): array
+	{
+		$keys = [];
+		foreach ($observations as $key => $observation) {
+			if ($observation['sends'] === [] || ClosureSignatureInference::isClosureSignatureMarker($observation['marker'])) {
+				continue;
+			}
+			$keys[$key] = true;
+		}
+
+		return $keys;
+	}
+
 	/**
 	 * @param array<string, array{marker: UnresolvedTemplateArgumentType, initial: Type|null, sends: list<array{Type, TemplateTypeVariance}>, lowerBounds: list<Type>, unconstrainingSend: bool}> $observations
 	 * @param array<string, Type> $resolutions

src/Analyser/Traverser/GenericTypeTemplateTraverser.php

@@ -64,6 +64,13 @@ public function traverse(Type $type, callable $traverse): Type
 			// markers - the real site re-keys and resolves them
 			$synthetic = $this->site->getAttribute(TemplateArgumentFrame::SYNTHETIC_SITE_ATTRIBUTE) === true;
 			if ($synthetic || ($this->allowUnresolved && $this->frame->isObserving())) {
+				if (!$synthetic) {
+					// a send in a body enclosing this one decided the site
+					$resolved = $this->frame->resolveSentFromEnclosingBody($this->site, $type->getName());
+					if ($resolved !== null) {
+						return $resolved;
+					}
+				}
 				if ($initialType instanceof UnresolvedTemplateArgumentType) {
 					return $initialType;
 				}

src/Reflection/ResolvedFunctionVariantWithOriginal.php

@@ -400,6 +400,11 @@ private function resolveResolvableTemplateTypes(Type $type, TemplateTypeVariance
 	private function unresolvedOrResolvedTemplateArgument(TemplateType $template, Type $inferred, Expr $site, TemplateArgumentFrame $frame, bool $allowUnresolved, bool $onlyCovariant): Type
 	{
 		if ($allowUnresolved && $frame->isObserving()) {
+			// a send in a body enclosing this one decided the site
+			$resolved = $frame->resolveSentFromEnclosingBody($site, $template->getName());
+			if ($resolved !== null) {
+				return $resolved;
+			}
 			if ($inferred instanceof UnresolvedTemplateArgumentType) {
 				return $inferred;
 			}

tests/PHPStan/Analyser/nsrt/bug-15415.php

@@ -0,0 +1,67 @@
+<?php declare(strict_types = 1);
+
+namespace Bug15415Types;
+
+use function PHPStan\Testing\assertType;
+
+/**
+ * @template TKey of array-key
+ * @template TValue
+ */
+final class Coll
+{
+	/**
+	 * @template TNewKey of array-key
+	 * @param callable(TValue): TNewKey $callback
+	 * @return Coll<TNewKey, TValue>
+	 */
+	public function keyBy(callable $callback): Coll
+	{
+		return new Coll();
+	}
+}
+
+/**
+ * @template T
+ * @param callable(): T $callback
+ * @return T
+ */
+function run(callable $callback)
+{
+	return $callback();
+}
+
+final class Test
+{
+	/**
+	 * @param Coll<int, int> $c
+	 * @return Coll<string, int>
+	 */
+	public function blockClosure(Coll $c): Coll
+	{
+		$r = run(function () use ($c): Coll {
+			$keyed = $c->keyBy(fn (int $i): string => $i > 0 ? 'positive' : 'negative');
+			assertType('Bug15415Types\Coll<string, int>', $keyed);
+
+			return $keyed;
+		});
+		assertType('Bug15415Types\Coll<string, int>', $r);
+
+		return $r;
+	}
+
+	/**
+	 * @param Coll<int, int> $c
+	 * @return Coll<string, int>
+	 */
+	public function closureInVariable(Coll $c): Coll
+	{
+		$f = function () use ($c): Coll {
+			return $c->keyBy(fn (int $i): string => $i > 0 ? 'positive' : 'negative');
+		};
+		$r = $f();
+		assertType('Bug15415Types\Coll<string, int>', $r);
+
+		return $r;
+	}
+}

tests/PHPStan/Rules/Methods/ReturnTypeRuleTest.php

@@ -1419,4 +1419,10 @@ public function testBug13325(): void
 		$this->analyse([__DIR__ . '/data/bug-13325.php'], []);
 	}
 
+	#[RequiresPhp('>= 8.0.0')]
+	public function testBug15415(): void
+	{
+		$this->analyse([__DIR__ . '/data/bug-15415.php'], []);
+	}
+
 }