Add TypeSpecifierContext::getConditionType() and move count/strlen narrowing to extensions

src/Analyser/TypeSpecifierContext.php

@@ -3,6 +3,7 @@
 namespace PHPStan\Analyser;
 
 use PHPStan\ShouldNotHappenException;
+use PHPStan\Type\Type;
 
 /**
  * @api
@@ -21,7 +22,7 @@ final class TypeSpecifierContext
 	/** @var self[] */
 	private static array $registry;
 
-	private function __construct(private ?int $value)
+	private function __construct(private ?int $value, private ?Type $narrowedReturnType = null)
 	{
 	}
 
@@ -90,4 +91,32 @@ public function null(): bool
 		return $this->value === null;
 	}
 
+	/**
+	 * If non-null, a constraint on the analyzed function/method's return value for the branch this
+	 * context represents (the truthy()=true branch). Type-specifying extensions may use it to narrow
+	 * arguments more precisely than the truthy/falsey/true/false buckets allow - e.g. for
+	 * `count($x) >= 2` the count() extension receives IntegerRangeType(2, max) here.
+	 *
+	 * It is null in every standard context; only comparison-rewriting code attaches one via
+	 * withNarrowedReturnType(), and negate() drops it again.
+	 *
+	 * @api
+	 */
+	public function getNarrowedReturnType(): ?Type
+	{
+		return $this->narrowedReturnType;
+	}
+
+	/**
+	 * Returns a context carrying the given narrowed return type. Intentionally bypasses the flyweight
+	 * registry. The narrowed return type is a one-directional hint valid only for this exact (bitmask, type)
+	 * pair - negate() discards it and falls back to bitmask-only behavior.
+	 *
+	 * @api
+	 */
+	public function withNarrowedReturnType(Type $narrowedReturnType): self
+	{
+		return new self($this->value, $narrowedReturnType);
+	}
+
 }

src/Analyser/TypeSpecifierFactory.php

@@ -8,6 +8,7 @@
 use PHPStan\Node\Printer\ExprPrinter;
 use PHPStan\Php\PhpVersion;
 use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Type\Php\CountFuncCallTypeSpecifier;
 use function array_merge;
 
 #[AutowiredService(name: 'typeSpecifierFactory')]
@@ -36,6 +37,7 @@ public function create(): TypeSpecifier
 			$methodTypeSpecifying,
 			$staticMethodTypeSpecifying,
 			$this->container->getParameter('rememberPossiblyImpureFunctionValues'),
+			$this->container->getByType(CountFuncCallTypeSpecifier::class),
 		);
 
 		foreach (array_merge(

src/Type/Php/ArrayKeyFirstLastFunctionTypeSpecifyingExtension.php

@@ -0,0 +1,67 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Php;
+
+use PhpParser\Node\Expr\FuncCall;
+use PhpParser\Node\Name;
+use PHPStan\Analyser\Scope;
+use PHPStan\Analyser\SpecifiedTypes;
+use PHPStan\Analyser\TypeSpecifier;
+use PHPStan\Analyser\TypeSpecifierAwareExtension;
+use PHPStan\Analyser\TypeSpecifierContext;
+use PHPStan\DependencyInjection\AutowiredService;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\Accessory\NonEmptyArrayType;
+use PHPStan\Type\FunctionTypeSpecifyingExtension;
+use function in_array;
+use function strtolower;
+
+/**
+ * Narrows the array argument of array_key_first()/array_key_last()/array_find_key() when its result is
+ * compared against null: a non-null key means the array is non-empty. array_key_first()/array_key_last()
+ * narrow in both directions (a null key means the array is empty); array_find_key() only narrows the
+ * non-null direction. Driven by the narrowed return type carried by the comparison.
+ */
+#[AutowiredService]
+final class ArrayKeyFirstLastFunctionTypeSpecifyingExtension implements FunctionTypeSpecifyingExtension, TypeSpecifierAwareExtension
+{
+
+	private TypeSpecifier $typeSpecifier;
+
+	public function isFunctionSupported(FunctionReflection $functionReflection, FuncCall $node, TypeSpecifierContext $context): bool
+	{
+		return $context->getNarrowedReturnType() !== null
+			&& $node->name instanceof Name
+			&& !$node->isFirstClassCallable()
+			&& isset($node->getArgs()[0])
+			&& in_array(strtolower($functionReflection->getName()), ['array_key_first', 'array_key_last', 'array_find_key'], true);
+	}
+
+	public function specifyTypes(FunctionReflection $functionReflection, FuncCall $node, Scope $scope, TypeSpecifierContext $context): SpecifiedTypes
+	{
+		$narrowedReturnType = $context->getNarrowedReturnType();
+		if ($narrowedReturnType === null || !$narrowedReturnType->isNull()->yes()) {
+			return new SpecifiedTypes();
+		}
+
+		$argValue = $node->getArgs()[0]->value;
+		if (!$scope->getType($argValue)->isArray()->yes()) {
+			return new SpecifiedTypes();
+		}
+
+		$functionName = strtolower($functionReflection->getName());
+		$bothDirections = in_array($functionName, ['array_key_first', 'array_key_last'], true);
+
+		if ($bothDirections || $context->falsey()) {
+			return $this->typeSpecifier->create($argValue, new NonEmptyArrayType(), $context->negate(), $scope);
+		}
+
+		return new SpecifiedTypes();
+	}
+
+	public function setTypeSpecifier(TypeSpecifier $typeSpecifier): void
+	{
+		$this->typeSpecifier = $typeSpecifier;
+	}
+
+}

src/Type/Php/CountFuncCallTypeSpecifier.php

@@ -0,0 +1,183 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Php;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Analyser\SpecifiedTypes;
+use PHPStan\Analyser\TypeSpecifier;
+use PHPStan\Analyser\TypeSpecifierContext;
+use PHPStan\DependencyInjection\AutowiredService;
+use PHPStan\Node\Printer\ExprPrinter;
+use PHPStan\TrinaryLogic;
+use PHPStan\Type\Accessory\HasOffsetValueType;
+use PHPStan\Type\Accessory\NonEmptyArrayType;
+use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\IntegerRangeType;
+use PHPStan\Type\MixedType;
+use PHPStan\Type\NeverType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use function count;
+use const COUNT_NORMAL;
+
+/**
+ * Narrows the array argument of a count()/sizeof() call based on a known constraint on its result
+ * (the size type). Shared between CountFunctionTypeSpecifyingExtension - which obtains the size type
+ * from TypeSpecifierContext::getNarrowedReturnType() - and TypeSpecifier's count()-comparison handling.
+ *
+ * TypeSpecifier is passed in per call (rather than injected) to avoid a construction cycle, since
+ * TypeSpecifier itself depends on this service.
+ */
+#[AutowiredService]
+final class CountFuncCallTypeSpecifier
+{
+
+	public function __construct(private ExprPrinter $exprPrinter)
+	{
+	}
+
+	/**
+	 * Whether the count()/sizeof() call counts in COUNT_NORMAL mode, so its result equals the array's
+	 * top-level size and can be used to narrow the counted array.
+	 */
+	public function isNormalCountCall(FuncCall $countFuncCall, Type $typeToCount, Scope $scope): TrinaryLogic
+	{
+		if (count($countFuncCall->getArgs()) === 1) {
+			return TrinaryLogic::createYes();
+		}
+
+		$mode = $scope->getType($countFuncCall->getArgs()[1]->value);
+		return (new ConstantIntegerType(COUNT_NORMAL))->isSuperTypeOf($mode)->result->or($typeToCount->getIterableValueType()->isArray()->negate());
+	}
+
+	public function specifyTypesForCountFuncCall(
+		TypeSpecifier $typeSpecifier,
+		FuncCall $countFuncCall,
+		Type $type,
+		Type $sizeType,
+		TypeSpecifierContext $context,
+		Scope $scope,
+		Expr $rootExpr,
+	): ?SpecifiedTypes
+	{
+		$isConstantArray = $type->isConstantArray();
+		$isList = $type->isList();
+		$oneOrMore = IntegerRangeType::fromInterval(1, null);
+		if (
+			!$this->isNormalCountCall($countFuncCall, $type, $scope)->yes()
+			|| (!$isConstantArray->yes() && !$isList->yes())
+			|| !$oneOrMore->isSuperTypeOf($sizeType)->yes()
+			|| $sizeType->isSuperTypeOf($type->getArraySize())->yes()
+		) {
+			return null;
+		}
+
+		if ($context->falsey() && $isConstantArray->yes()) {
+			$remainingSize = TypeCombinator::remove($type->getArraySize(), $sizeType);
+			if (!$remainingSize instanceof NeverType) {
+				$negatedContext = $context->false()
+					? TypeSpecifierContext::createTrue()
+					: TypeSpecifierContext::createTruthy();
+				$result = $this->specifyTypesForCountFuncCall(
+					$typeSpecifier,
+					$countFuncCall,
+					$type,
+					$remainingSize,
+					$negatedContext,
+					$scope,
+					$rootExpr,
+				);
+				if ($result !== null) {
+					return $result;
+				}
+			}
+
+			// Fallback: directly filter constant arrays by their exact sizes.
+			// This avoids using TypeCombinator::remove() with falsey context,
+			// which can incorrectly remove arrays whose count doesn't match
+			// but whose shape is a subtype of the matched array.
+			$keptTypes = [];
+			foreach ($type->getConstantArrays() as $arrayType) {
+				if ($sizeType->isSuperTypeOf($arrayType->getArraySize())->yes()) {
+					continue;
+				}
+
+				$keptTypes[] = $arrayType;
+			}
+			if ($keptTypes !== []) {
+				return $typeSpecifier->create(
+					$countFuncCall->getArgs()[0]->value,
+					TypeCombinator::union(...$keptTypes),
+					$context->negate(),
+					$scope,
+				)->setRootExpr($rootExpr);
+			}
+		}
+
+		$resultTypes = [];
+		foreach ($type->getArrays() as $arrayType) {
+			$isSizeSuperTypeOfArraySize = $sizeType->isSuperTypeOf($arrayType->getArraySize());
+			if ($isSizeSuperTypeOfArraySize->no()) {
+				continue;
+			}
+
+			if ($context->falsey() && $isSizeSuperTypeOfArraySize->maybe()) {
+				continue;
+			}
+
+			$resultTypes[] = $isList->yes()
+				? $arrayType->truncateListToSize($sizeType)
+				: TypeCombinator::intersect($arrayType, new NonEmptyArrayType());
+		}
+
+		if ($context->truthy() && $isConstantArray->yes() && $isList->yes()) {
+			$hasOptionalKeysOrUnsealed = false;
+			foreach ($type->getConstantArrays() as $arrayType) {
+				if ($arrayType->getOptionalKeys() !== [] || $arrayType->isUnsealed()->yes()) {
+					// Unsealed CATs can't be narrowed via the
+					// `HasOffsetValueType`-only shortcut below — the
+					// intersection of an unsealed shape with a single-slot
+					// constraint produces `NeverType`. Fall through to
+					// the full builder-based narrowing, which carries the
+					// unsealed slot via the loop above.
+					$hasOptionalKeysOrUnsealed = true;
+					break;
+				}
+			}
+
+			if (!$hasOptionalKeysOrUnsealed) {
+				$argExpr = $countFuncCall->getArgs()[0]->value;
+				$argExprString = $this->exprPrinter->printExpr($argExpr);
+
+				$sizeMin = null;
+				$sizeMax = null;
+				if ($sizeType instanceof ConstantIntegerType) {
+					$sizeMin = $sizeType->getValue();
+					$sizeMax = $sizeType->getValue();
+				} elseif ($sizeType instanceof IntegerRangeType) {
+					$sizeMin = $sizeType->getMin();
+					$sizeMax = $sizeType->getMax();
+				}
+
+				$sureTypes = [];
+				$sureNotTypes = [];
+
+				if ($sizeMin !== null && $sizeMin >= 1) {
+					$sureTypes[$argExprString] = [$argExpr, new HasOffsetValueType(new ConstantIntegerType($sizeMin - 1), new MixedType())];
+				}
+				if ($sizeMax !== null) {
+					$sureNotTypes[$argExprString] = [$argExpr, new HasOffsetValueType(new ConstantIntegerType($sizeMax), new MixedType())];
+				}
+
+				if ($sureTypes !== [] || $sureNotTypes !== []) {
+					return (new SpecifiedTypes($sureTypes, $sureNotTypes))->setRootExpr($rootExpr);
+				}
+			}
+		}
+
+		return $typeSpecifier->create($countFuncCall->getArgs()[0]->value, TypeCombinator::union(...$resultTypes), $context, $scope)->setRootExpr($rootExpr);
+	}
+
+}

src/Type/Php/CountFunctionTypeSpecifyingExtension.php

@@ -21,6 +21,10 @@ final class CountFunctionTypeSpecifyingExtension implements FunctionTypeSpecifyi
 
 	private TypeSpecifier $typeSpecifier;
 
+	public function __construct(private CountFuncCallTypeSpecifier $countFuncCallTypeSpecifier)
+	{
+	}
+
 	public function isFunctionSupported(
 		FunctionReflection $functionReflection,
 		FuncCall $node,
@@ -39,7 +43,27 @@ public function specifyTypes(
 		TypeSpecifierContext $context,
 	): SpecifiedTypes
 	{
-		if (!$scope->getType($node->getArgs()[0]->value)->isArray()->yes()) {
+		$argType = $scope->getType($node->getArgs()[0]->value);
+
+		$narrowedReturnType = $context->getNarrowedReturnType();
+		if ($narrowedReturnType !== null) {
+			$specifiedTypes = $this->countFuncCallTypeSpecifier->specifyTypesForCountFuncCall(
+				$this->typeSpecifier,
+				$node,
+				$argType,
+				$narrowedReturnType,
+				$context,
+				$scope,
+				$node,
+			);
+			if ($specifiedTypes !== null) {
+				return $specifiedTypes;
+			}
+
+			return new SpecifiedTypes([], []);
+		}
+
+		if (!$argType->isArray()->yes()) {
 			return new SpecifiedTypes([], []);
 		}