Generalize an inferred template type by what a conditional branch knows about it

src/Reflection/ResolvedFunctionVariantWithOriginal.php

@@ -13,6 +13,7 @@
 use PHPStan\Type\Generic\TemplateType;
 use PHPStan\Type\Generic\TemplateTypeHelper;
 use PHPStan\Type\Generic\TemplateTypeMap;
+use PHPStan\Type\Generic\TemplateTypeReference;
 use PHPStan\Type\Generic\TemplateTypeVariance;
 use PHPStan\Type\Generic\TemplateTypeVarianceMap;
 use PHPStan\Type\Generic\UnresolvedTemplateArgumentType;
@@ -290,51 +291,7 @@ private function resolveResolvableTemplateTypes(Type $type, TemplateTypeVariance
 			}
 		}
 
-		$objectCb = function (Type $type, callable $traverse) use ($references, $onlyCovariant, $site, $frame, $allowUnresolved): Type {
-			if (
-				$type instanceof TemplateType
-				&& !$type instanceof NarrowedSubjectType
-				&& !$type->isArgument()
-				&& $type->getScope()->getFunctionName() !== null
-			) {
-				$newType = $this->resolvedTemplateTypeMap->getType($type->getName());
-				if ($newType === null || $newType instanceof ErrorType) {
-					return $traverse($type);
-				}
-
-				if ($site !== null && $frame !== null) {
-					$newType = $this->unresolvedOrResolvedTemplateArgument($type, $newType, $site, $frame, $allowUnresolved, $onlyCovariant[$type->getName()] ?? false);
-				} else {
-					$newType = TemplateTypeHelper::generalizeInferredTemplateType($type, $newType);
-				}
-				$variance = TemplateTypeVariance::createInvariant();
-				foreach ($references as $reference) {
-					// this uses identity to distinguish between different occurrences of the same template type
-					// see https://github.com/phpstan/phpstan-src/pull/2485#discussion_r1328555397 for details
-					if ($reference->getType() === $type) {
-						$variance = $reference->getPositionVariance();
-						break;
-					}
-				}
-
-				$callSiteVariance = $this->callSiteVarianceMap->getVariance($type->getName());
-				if ($callSiteVariance === null || $callSiteVariance->invariant()) {
-					return $newType;
-				}
-
-				if (!$callSiteVariance->covariant() && $variance->covariant()) {
-					return $traverse($type->getBound());
-				}
-
-				if (!$callSiteVariance->contravariant() && $variance->contravariant()) {
-					return new NonAcceptingNeverType();
-				}
-
-				return $newType;
-			}
-
-			return $traverse($type);
-		};
+		$objectCb = fn (Type $type, callable $traverse): Type => $this->resolveTemplateTypeInGenericType($type, $traverse, $references, $onlyCovariant, $site, $frame, $allowUnresolved, null);
 
 		return TypeTraverser::map($type, function (Type $type, callable $traverse) use ($references, $objectCb): Type {
 			if ($type instanceof GenericObjectType || $type instanceof GenericStaticType) {
@@ -383,6 +340,87 @@ private function resolveResolvableTemplateTypes(Type $type, TemplateTypeVariance
 		});
 	}
 
+	/**
+	 * @param callable(Type): Type $traverse
+	 * @param list<TemplateTypeReference> $references
+	 * @param array<string, bool> $onlyCovariant
+	 * @param string|null $keepInferredName a template type whose inferred type is not generalized here
+	 */
+	private function resolveTemplateTypeInGenericType(
+		Type $type,
+		callable $traverse,
+		array $references,
+		array $onlyCovariant,
+		?Expr $site,
+		?TemplateArgumentFrame $frame,
+		bool $allowUnresolved,
+		?string $keepInferredName,
+	): Type
+	{
+		if (
+			$type instanceof NarrowedSubjectType
+			&& $keepInferredName === null
+			&& !$type->isArgument()
+			&& $type->getScope()->getFunctionName() !== null
+		) {
+			// generalize what the branch knows about the subject, not the subject: in the
+			// else branch of `(T is \UnitEnum ? array-key : T)` with T of array-key|\UnitEnum,
+			// T is a scalar and an inferred 'foo' stays 'foo'. Only the subject is taken as
+			// inferred - a template type in the target is generalized as the condition sees it.
+			$subjectName = $type->getName();
+			$resolved = TypeTraverser::map($type, fn (Type $type, callable $traverse): Type => $this->resolveTemplateTypeInGenericType($type, $traverse, $references, $onlyCovariant, $site, $frame, $allowUnresolved, $subjectName));
+
+			return TemplateTypeHelper::generalizeInferredTemplateType($type, $resolved);
+		}
+
+		if (
+			$type instanceof TemplateType
+			&& !$type instanceof NarrowedSubjectType
+			&& !$type->isArgument()
+			&& $type->getScope()->getFunctionName() !== null
+		) {
+			$newType = $this->resolvedTemplateTypeMap->getType($type->getName());
+			if ($newType === null || $newType instanceof ErrorType) {
+				return $traverse($type);
+			}
+
+			// the subject of a narrowed subject is generalized as a whole by the caller
+			if ($type->getName() !== $keepInferredName) {
+				if ($site !== null && $frame !== null) {
+					$newType = $this->unresolvedOrResolvedTemplateArgument($type, $newType, $site, $frame, $allowUnresolved, $onlyCovariant[$type->getName()] ?? false);
+				} else {
+					$newType = TemplateTypeHelper::generalizeInferredTemplateType($type, $newType);
+				}
+			}
+			$variance = TemplateTypeVariance::createInvariant();
+			foreach ($references as $reference) {
+				// this uses identity to distinguish between different occurrences of the same template type
+				// see https://github.com/phpstan/phpstan-src/pull/2485#discussion_r1328555397 for details
+				if ($reference->getType() === $type) {
+					$variance = $reference->getPositionVariance();
+					break;
+				}
+			}
+
+			$callSiteVariance = $this->callSiteVarianceMap->getVariance($type->getName());
+			if ($callSiteVariance === null || $callSiteVariance->invariant()) {
+				return $newType;
+			}
+
+			if (!$callSiteVariance->covariant() && $variance->covariant()) {
+				return $traverse($type->getBound());
+			}
+
+			if (!$callSiteVariance->contravariant() && $variance->contravariant()) {
+				return new NonAcceptingNeverType();
+			}
+
+			return $newType;
+		}
+
+		return $traverse($type);
+	}
+
 	/**
 	 * An inferred template argument inside a generic object of the return type:
 	 * during the body's observation pass a marker keyed by the call (an inferred

tests/PHPStan/Analyser/NarrowedSubjectInferredLiteralTest.php

@@ -0,0 +1,36 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Analyser;
+
+use PHPStan\Testing\TypeInferenceTestCase;
+use PHPUnit\Framework\Attributes\DataProvider;
+use PHPUnit\Framework\Attributes\RequiresPhp;
+
+/**
+ * Without bleedingEdge: the inferred template arguments are generalized right away
+ * (featureToggles.unresolvedTemplateArguments is off), which is where a conditional
+ * branch's subject has to keep what the branch knows about it.
+ */
+#[RequiresPhp('>= 8.1.0')]
+class NarrowedSubjectInferredLiteralTest extends TypeInferenceTestCase
+{
+
+	public static function dataFileAsserts(): iterable
+	{
+		yield from self::gatherAssertTypes(__DIR__ . '/data/narrowed-subject-inferred-literal.php');
+	}
+
+	/**
+	 * @param mixed ...$args
+	 */
+	#[DataProvider('dataFileAsserts')]
+	public function testFileAsserts(
+		string $assertType,
+		string $file,
+		...$args,
+	): void
+	{
+		$this->assertFileAsserts($assertType, $file, ...$args);
+	}
+
+}

tests/PHPStan/Analyser/data/narrowed-subject-inferred-literal.php

@@ -0,0 +1,136 @@
+<?php // lint >= 8.1
+
+declare(strict_types = 1);
+
+namespace NarrowedSubjectInferredLiteral;
+
+use function PHPStan\Testing\assertType;
+
+class User
+{
+
+}
+
+enum Digit: int
+{
+
+	case One = 1;
+
+}
+
+/**
+ * @template TKey of array-key
+ * @template TValue
+ */
+class Collection
+{
+
+	/**
+	 * @template TNewKey of array-key|\UnitEnum
+	 *
+	 * @param  (callable(TValue, TKey): TNewKey)|array<mixed>|string  $keyBy
+	 * @return static<($keyBy is (array|string) ? array-key : (TNewKey is \UnitEnum ? array-key : TNewKey)), TValue>
+	 */
+	public function keyBy($keyBy)
+	{
+		throw new \Exception();
+	}
+
+	/**
+	 * @template TNewKey of array-key|\UnitEnum
+	 *
+	 * @param  callable(TValue, TKey): TNewKey  $keyBy
+	 * @return static<(TNewKey is \UnitEnum ? array-key : TNewKey), TValue>
+	 */
+	public function keyByCallback($keyBy)
+	{
+		throw new \Exception();
+	}
+
+}
+
+/**
+ * @template T of array-key
+ * @template U of array-key
+ * @param T $a
+ * @param U $b
+ * @return Collection<(T is U ? T : int), int>
+ */
+function templateTarget($a, $b)
+{
+	throw new \Exception();
+}
+
+/**
+ * @template T of array-key
+ * @template U of array-key
+ * @param T $a
+ * @param U $b
+ * @return Collection<(T is not U ? int : T), int>
+ */
+function negatedTemplateTarget($a, $b)
+{
+	throw new \Exception();
+}
+
+/** @param Collection<int, User> $collection */
+function test(Collection $collection): void
+{
+	assertType("NarrowedSubjectInferredLiteral\\Collection<'foo', NarrowedSubjectInferredLiteral\\User>", $collection->keyBy(fn ($user) => 'foo'));
+	assertType('NarrowedSubjectInferredLiteral\\Collection<0, NarrowedSubjectInferredLiteral\\User>', $collection->keyBy(static fn ($user): int => 0));
+	assertType('NarrowedSubjectInferredLiteral\\Collection<(int|string), NarrowedSubjectInferredLiteral\\User>', $collection->keyBy('name'));
+	assertType('NarrowedSubjectInferredLiteral\\Collection<(int|string), NarrowedSubjectInferredLiteral\\User>', $collection->keyBy(static fn ($user) => Digit::One));
+
+	assertType("NarrowedSubjectInferredLiteral\\Collection<'foo', NarrowedSubjectInferredLiteral\\User>", $collection->keyByCallback(fn ($user) => 'foo'));
+	assertType('NarrowedSubjectInferredLiteral\\Collection<(int|string), NarrowedSubjectInferredLiteral\\User>', $collection->keyByCallback(static fn ($user) => Digit::One));
+
+	// the condition compares the generalized T and U, the branch keeps the inferred T
+	assertType("NarrowedSubjectInferredLiteral\\Collection<'foo', int>", templateTarget('foo', 'bar'));
+	assertType("NarrowedSubjectInferredLiteral\\Collection<'foo', int>", negatedTemplateTarget('foo', 'bar'));
+}
+
+/**
+ * @template TKey of array-key
+ * @template TValue
+ */
+interface Enumerable
+{
+
+	/**
+	 * @template TNewKey of array-key|\UnitEnum
+	 *
+	 * @param  (callable(TValue, TKey): TNewKey)|array<mixed>|string  $keyBy
+	 * @return static<($keyBy is (array|string) ? array-key : (TNewKey is \UnitEnum ? array-key : TNewKey)), TValue>
+	 */
+	public function keyBy($keyBy);
+
+}
+
+/**
+ * @template TKey of array-key
+ * @template TValue
+ * @implements Enumerable<TKey, TValue>
+ */
+class InheritedCollection implements Enumerable
+{
+
+	/**
+	 * {@inheritDoc}
+	 */
+	public function keyBy($keyBy)
+	{
+		throw new \Exception();
+	}
+
+}
+
+/** @param InheritedCollection<int, User> $collection */
+function testInherited(InheritedCollection $collection): void
+{
+	assertType("NarrowedSubjectInferredLiteral\\InheritedCollection<'foo', NarrowedSubjectInferredLiteral\\User>", $collection->keyBy(function ($user, $int) {
+		return 'foo';
+	}));
+	assertType('NarrowedSubjectInferredLiteral\\InheritedCollection<0, NarrowedSubjectInferredLiteral\\User>', $collection->keyBy(static fn ($user): int => 0));
+	assertType('NarrowedSubjectInferredLiteral\\InheritedCollection<(int|string), NarrowedSubjectInferredLiteral\\User>', $collection->keyBy('name'));
+	assertType('NarrowedSubjectInferredLiteral\\InheritedCollection<(int|string), NarrowedSubjectInferredLiteral\\User>', $collection->keyBy(static fn ($user) => Digit::One));
+}