Fix phpstan/phpstan#14464: Narrow list types when count() result is stored in variable

src/Analyser/ExprHandler/AssignHandler.php

@@ -51,6 +51,7 @@
 use PHPStan\Type\Accessory\HasOffsetValueType;
 use PHPStan\Type\Accessory\NonEmptyArrayType;
 use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
 use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\Constant\ConstantStringType;
 use PHPStan\Type\ConstantTypeHelper;
@@ -73,6 +74,8 @@
 use function in_array;
 use function is_int;
 use function is_string;
+use function min;
+use function strtolower;
 
 /**
  * @implements ExprHandler<Assign|AssignRef>
@@ -313,6 +316,36 @@
 					$conditionalExpressions = $this->processSureNotTypesForConditionalExpressionsAfterAssign($scope, $var->name, $conditionalExpressions, $identicalSpecifiedTypes, $falseyType);
 				}
 
+				if (
+					$assignedExpr instanceof FuncCall
+					&& $assignedExpr->name instanceof Name
+					&& in_array(strtolower((string) $assignedExpr->name), ['count', 'sizeof'], true)
+					&& count($assignedExpr->getArgs()) >= 1
+				) {
+					$countArgType = $scope->getType($assignedExpr->getArgs()[0]->value);
+					if ($countArgType->isList()->yes() || $countArgType->isConstantArray()->yes()) {
+						$arraySize = $countArgType->getArraySize();
+						$maxSize = ConstantArrayTypeBuilder::ARRAY_COUNT_LIMIT;
+						if ($arraySize instanceof ConstantIntegerType) {
+							$maxSize = $arraySize->getValue();
+						} elseif ($arraySize instanceof IntegerRangeType && $arraySize->getMax() !== null) {
+							$maxSize = min($maxSize, $arraySize->getMax());
+						}
+
+						for ($i = 1; $i <= $maxSize; $i++) {
+							$sizeType = new ConstantIntegerType($i);
+							if (!$type->isSuperTypeOf($sizeType)->yes()) {
+								continue;
+							}
+
+							$identicalConditionExpr = new Expr\BinaryOp\Identical($assignedExpr, new Node\Scalar\Int_($i));
+							$identicalSpecifiedTypes = $this->typeSpecifier->specifyTypesInCondition($scope, $identicalConditionExpr, TypeSpecifierContext::createTrue());
+							$conditionalExpressions = $this->processSureTypesForConditionalExpressionsAfterAssign($scope, $var->name, $conditionalExpressions, $identicalSpecifiedTypes, $sizeType);
+							$conditionalExpressions = $this->processSureNotTypesForConditionalExpressionsAfterAssign($scope, $var->name, $conditionalExpressions, $identicalSpecifiedTypes, $sizeType);
+						}
+					}
+				}
+
 				$nodeScopeResolver->callNodeCallback($nodeCallback, new VariableAssignNode($var, $assignedExpr), $scopeBeforeAssignEval, $storage);
 				$scope = $scope->assignVariable($var->name, $type, $scope->getNativeType($assignedExpr), TrinaryLogic::createYes());
 				foreach ($conditionalExpressions as $exprString => $holders) {

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

@@ -0,0 +1,109 @@
+<?php declare(strict_types = 1);
+
+namespace Bug14464;
+
+use function PHPStan\Testing\assertType;
+
+class HelloWorld
+{
+	protected function columnOrAlias(string $columnName): void
+	{
+		$colParts = preg_split('/\s+/', $columnName, -1, \PREG_SPLIT_NO_EMPTY);
+		if ($colParts === false) {
+			throw new \RuntimeException('preg error');
+		}
+		assertType('list<non-empty-string>', $colParts);
+		$numParts = count($colParts);
+
+		if ($numParts == 3) {
+			assertType('array{non-empty-string, non-empty-string, non-empty-string}', $colParts);
+		} elseif ($numParts == 2) {
+			assertType('array{non-empty-string, non-empty-string}', $colParts);
+		} elseif ($numParts == 1) {
+			assertType('array{non-empty-string}', $colParts);
+		}
+	}
+
+	/** @param list<string> $list */
+	public function indirectCountCheck(array $list): void
+	{
+		$n = count($list);
+		if ($n === 3) {
+			assertType('array{string, string, string}', $list);
+		}
+		if ($n === 2) {
+			assertType('array{string, string}', $list);
+		}
+		if ($n === 1) {
+			assertType('array{string}', $list);
+		}
+	}
+
+	/** @param list<string> $list */
+	public function directCountCheck(array $list): void
+	{
+		if (count($list) === 3) {
+			assertType('array{string, string, string}', $list);
+		}
+		if (count($list) === 2) {
+			assertType('array{string, string}', $list);
+		}
+		if (count($list) === 1) {
+			assertType('array{string}', $list);
+		}
+	}
+
+	/** @param list<string> $list */
+	public function sizeofIndirect(array $list): void
+	{
+		$n = sizeof($list);
+		if ($n === 2) {
+			assertType('array{string, string}', $list);
+		}
+	}
+
+	/** @param list<int> $list */
+	public function looseEqualityCheck(array $list): void
+	{
+		$n = count($list);
+		if ($n == 3) {
+			assertType('array{int, int, int}', $list);
+		}
+	}
+
+	/**
+	 * Non-list arrays should not get specific shapes since keys are unknown
+	 * @param array<string, int> $map
+	 */
+	public function nonListArray(array $map): void
+	{
+		$n = count($map);
+		if ($n === 2) {
+			assertType('non-empty-array<string, int>', $map);
+		}
+	}
+
+	/** @param array{string}|array{string, string}|array{string, string, string} $list */
+	public function constantArrayUnionIndirect(array $list): void
+	{
+		$n = count($list);
+		if ($n === 2) {
+			assertType('array{string, string}', $list);
+		}
+		if ($n === 3) {
+			assertType('array{string, string, string}', $list);
+		}
+	}
+
+	/** @param array{a: string, b: int}|array{x: float, y: float, z: float} $map */
+	public function constantNonListDifferentShapes(array $map): void
+	{
+		$n = count($map);
+		if ($n === 2) {
+			assertType('array{a: string, b: int}', $map);
+		}
+		if ($n === 3) {
+			assertType('array{x: float, y: float, z: float}', $map);
+		}
+	}
+}