Improve count() narrowing of constant arrays

src/Analyser/TypeSpecifier.php

@@ -269,22 +269,21 @@ public function specifyTypesInCondition(
 			) {
 				$argType = $scope->getType($expr->right->getArgs()[0]->value);
 
-				if ($argType instanceof UnionType) {
-					$sizeType = null;
-					if ($leftType instanceof ConstantIntegerType) {
-						if ($orEqual) {
-							$sizeType = IntegerRangeType::createAllGreaterThanOrEqualTo($leftType->getValue());
-						} else {
-							$sizeType = IntegerRangeType::createAllGreaterThan($leftType->getValue());
-						}
-					} elseif ($leftType instanceof IntegerRangeType) {
-						$sizeType = $leftType;
+				if ($leftType instanceof ConstantIntegerType) {
+					if ($orEqual) {
+						$sizeType = IntegerRangeType::createAllGreaterThanOrEqualTo($leftType->getValue());
+					} else {
+						$sizeType = IntegerRangeType::createAllGreaterThan($leftType->getValue());
 					}
+				} elseif ($leftType instanceof IntegerRangeType) {
+					$sizeType = $leftType->shift($offset);
+				} else {
+					$sizeType = $leftType;
+				}
 
-					$narrowed = $this->narrowUnionByArraySize($expr->right, $argType, $sizeType, $context, $scope, $expr);
-					if ($narrowed !== null) {
-						return $narrowed;
-					}
+				$specifiedTypes = $this->specifyTypesForCountFuncCall($expr->right, $argType, $sizeType, $context, $scope, $expr);
+				if ($specifiedTypes !== null) {
+					$result = $result->unionWith($specifiedTypes);
 				}
 
 				if (
@@ -969,115 +968,95 @@ public function specifyTypesInCondition(
 		return (new SpecifiedTypes([], []))->setRootExpr($expr);
 	}
 
-	private function narrowUnionByArraySize(FuncCall $countFuncCall, UnionType $argType, ?Type $sizeType, TypeSpecifierContext $context, Scope $scope, ?Expr $rootExpr): ?SpecifiedTypes
+	private function specifyTypesForCountFuncCall(
+		FuncCall $countFuncCall,
+		Type $type,
+		Type $sizeType,
+		TypeSpecifierContext $context,
+		Scope $scope,
+		Expr $rootExpr,
+	): ?SpecifiedTypes
 	{
-		if ($sizeType === null) {
-			return null;
-		}
-
 		if (count($countFuncCall->getArgs()) === 1) {
 			$isNormalCount = TrinaryLogic::createYes();
 		} else {
 			$mode = $scope->getType($countFuncCall->getArgs()[1]->value);
-			$isNormalCount = (new ConstantIntegerType(COUNT_NORMAL))->isSuperTypeOf($mode)->result->or($argType->getIterableValueType()->isArray()->negate());
+			$isNormalCount = (new ConstantIntegerType(COUNT_NORMAL))->isSuperTypeOf($mode)->result->or($type->getIterableValueType()->isArray()->negate());
 		}
 
+		$isList = $type->isList();
 		if (
-			$isNormalCount->yes()
-			&& $argType->isConstantArray()->yes()
+			!$isNormalCount->yes()
+			|| (!$type->isConstantArray()->yes() && !$isList->yes())
+			|| $type->isIterableAtLeastOnce()->no() // array{} cannot be used for further narrowing
 		) {
-			$result = [];
-			foreach ($argType->getTypes() as $innerType) {
-				$arraySize = $innerType->getArraySize();
-				$isSize = $sizeType->isSuperTypeOf($arraySize);
-				if ($context->truthy()) {
-					if ($isSize->no()) {
-						continue;
-					}
-
-					$constArray = $this->turnListIntoConstantArray($countFuncCall, $innerType, $sizeType, $scope);
-					if ($constArray !== null) {
-						$innerType = $constArray;
-					}
-				}
-				if ($context->falsey()) {
-					if (!$isSize->yes()) {
-						continue;
-					}
-				}
-
-				$result[] = $innerType;
-			}
-
-			return $this->create($countFuncCall->getArgs()[0]->value, TypeCombinator::union(...$result), $context, $scope)->setRootExpr($rootExpr);
+			return null;
 		}
 
-		return null;
-	}
-
-	private function turnListIntoConstantArray(FuncCall $countFuncCall, Type $type, Type $sizeType, Scope $scope): ?Type
-	{
-		$argType = $scope->getType($countFuncCall->getArgs()[0]->value);
-
-		if (count($countFuncCall->getArgs()) === 1) {
-			$isNormalCount = TrinaryLogic::createYes();
-		} else {
-			$mode = $scope->getType($countFuncCall->getArgs()[1]->value);
-			$isNormalCount = (new ConstantIntegerType(COUNT_NORMAL))->isSuperTypeOf($mode)->result->or($argType->getIterableValueType()->isArray()->negate());
-		}
+		$resultTypes = [];
+		foreach ($type->getArrays() as $arrayType) {
+			$isSizeSuperTypeOfArraySize = $sizeType->isSuperTypeOf($arrayType->getArraySize());
+			if ($isSizeSuperTypeOfArraySize->no()) {
+				continue;
+			}
 
-		if (
-			$isNormalCount->yes()
-			&& $type->isList()->yes()
-			&& $sizeType instanceof ConstantIntegerType
-			&& $sizeType->getValue() < ConstantArrayTypeBuilder::ARRAY_COUNT_LIMIT
-		) {
-			// turn optional offsets non-optional
-			$valueTypesBuilder = ConstantArrayTypeBuilder::createEmpty();
-			for ($i = 0; $i < $sizeType->getValue(); $i++) {
-				$offsetType = new ConstantIntegerType($i);
-				$valueTypesBuilder->setOffsetValueType($offsetType, $type->getOffsetValueType($offsetType));
+			if ($context->falsey() && $isSizeSuperTypeOfArraySize->maybe()) {
+				continue;
 			}
-			return $valueTypesBuilder->getArray();
-		}
 
-		if (
-			$isNormalCount->yes()
-			&& $type->isList()->yes()
-			&& $sizeType instanceof IntegerRangeType
-			&& $sizeType->getMin() !== null
-		) {
-			// turn optional offsets non-optional
-			$valueTypesBuilder = ConstantArrayTypeBuilder::createEmpty();
-			for ($i = 0; $i < $sizeType->getMin(); $i++) {
-				$offsetType = new ConstantIntegerType($i);
-				$valueTypesBuilder->setOffsetValueType($offsetType, $type->getOffsetValueType($offsetType));
-			}
-			if ($sizeType->getMax() !== null) {
-				for ($i = $sizeType->getMin(); $i < $sizeType->getMax(); $i++) {
+			if (
+				$isList->yes()
+				&& $sizeType instanceof ConstantIntegerType
+				&& $sizeType->getValue() < ConstantArrayTypeBuilder::ARRAY_COUNT_LIMIT
+			) {
+				// turn optional offsets non-optional
+				$valueTypesBuilder = ConstantArrayTypeBuilder::createEmpty();
+				for ($i = 0; $i < $sizeType->getValue(); $i++) {
 					$offsetType = new ConstantIntegerType($i);
-					$valueTypesBuilder->setOffsetValueType($offsetType, $type->getOffsetValueType($offsetType), true);
+					$valueTypesBuilder->setOffsetValueType($offsetType, $arrayType->getOffsetValueType($offsetType));
 				}
-			} elseif ($type->isConstantArray()->yes()) {
-				for ($i = $sizeType->getMin();; $i++) {
+				$resultTypes[] = $valueTypesBuilder->getArray();
+				continue;
+			}
+
+			if (
+				$isList->yes()
+				&& $sizeType instanceof IntegerRangeType
+				&& $sizeType->getMin() !== null
+			) {
+				// turn optional offsets non-optional
+				$valueTypesBuilder = ConstantArrayTypeBuilder::createEmpty();
+				for ($i = 0; $i < $sizeType->getMin(); $i++) {
 					$offsetType = new ConstantIntegerType($i);
-					$hasOffset = $type->hasOffsetValueType($offsetType);
-					if ($hasOffset->no()) {
-						break;
+					$valueTypesBuilder->setOffsetValueType($offsetType, $arrayType->getOffsetValueType($offsetType));
+				}
+				if ($sizeType->getMax() !== null) {
+					for ($i = $sizeType->getMin(); $i < $sizeType->getMax(); $i++) {
+						$offsetType = new ConstantIntegerType($i);
+						$valueTypesBuilder->setOffsetValueType($offsetType, $arrayType->getOffsetValueType($offsetType), true);
 					}
-					$valueTypesBuilder->setOffsetValueType($offsetType, $type->getOffsetValueType($offsetType), !$hasOffset->yes());
+				} elseif ($arrayType->isConstantArray()->yes()) {
+					for ($i = $sizeType->getMin();; $i++) {
+						$offsetType = new ConstantIntegerType($i);
+						$hasOffset = $arrayType->hasOffsetValueType($offsetType);
+						if ($hasOffset->no()) {
+							break;
+						}
+						$valueTypesBuilder->setOffsetValueType($offsetType, $arrayType->getOffsetValueType($offsetType), !$hasOffset->yes());
+					}
+				} else {
+					$resultTypes[] = TypeCombinator::intersect($arrayType, new NonEmptyArrayType());
+					continue;
 				}
-			} else {
-				return null;
-			}
 
-			$arrayType = $valueTypesBuilder->getArray();
-			if ($arrayType->isIterableAtLeastOnce()->yes()) {
-				return $arrayType;
+				$resultTypes[] = $valueTypesBuilder->getArray();
+				continue;
 			}
+
+			$resultTypes[] = $arrayType;
 		}
 
-		return null;
+		return $this->create($countFuncCall->getArgs()[0]->value, TypeCombinator::union(...$resultTypes), $context, $scope)->setRootExpr($rootExpr);
 	}
 
 	private function specifyTypesForConstantBinaryExpression(
@@ -2109,36 +2088,20 @@ public function resolveIdentical(Expr\BinaryOp\Identical $expr, Scope $scope, Ty
 				);
 			}
 
-			if ($argType instanceof UnionType) {
-				$narrowed = $this->narrowUnionByArraySize($unwrappedLeftExpr, $argType, $rightType, $context, $scope, $expr);
-				if ($narrowed !== null) {
-					return $narrowed;
-				}
+			$specifiedTypes = $this->specifyTypesForCountFuncCall($unwrappedLeftExpr, $argType, $rightType, $context, $scope, $expr);
+			if ($specifiedTypes !== null) {
+				return $specifiedTypes;
 			}
 
-			if ($context->truthy()) {
-				if ($argType->isArray()->yes()) {
-					if (
-						$argType->isConstantArray()->yes()
-						&& $rightType->isSuperTypeOf($argType->getArraySize())->no()
-					) {
-						return $this->create($unwrappedLeftExpr->getArgs()[0]->value, new NeverType(), $context, $scope)->setRootExpr($expr);
-					}
-
-					$funcTypes = $this->create($unwrappedLeftExpr, $rightType, $context, $scope)->setRootExpr($expr);
-					$constArray = $this->turnListIntoConstantArray($unwrappedLeftExpr, $argType, $rightType, $scope);
-					if ($constArray !== null) {
-						return $funcTypes->unionWith(
-							$this->create($unwrappedLeftExpr->getArgs()[0]->value, $constArray, $context, $scope)->setRootExpr($expr),
-						);
-					} elseif (IntegerRangeType::fromInterval(1, null)->isSuperTypeOf($rightType)->yes()) {
-						return $funcTypes->unionWith(
-							$this->create($unwrappedLeftExpr->getArgs()[0]->value, new NonEmptyArrayType(), $context, $scope)->setRootExpr($expr),
-						);
-					}
-
-					return $funcTypes;
+			if ($context->truthy() && $argType->isArray()->yes()) {
+				$funcTypes = $this->create($unwrappedLeftExpr, $rightType, $context, $scope)->setRootExpr($expr);
+				if (IntegerRangeType::fromInterval(1, null)->isSuperTypeOf($rightType)->yes()) {
+					return $funcTypes->unionWith(
+						$this->create($unwrappedLeftExpr->getArgs()[0]->value, new NonEmptyArrayType(), $context, $scope)->setRootExpr($expr),
+					);
 				}
+
+				return $funcTypes;
 			}
 		}
 

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

@@ -40,10 +40,10 @@ function(array $array, int $count): void {
 	if (isset($array['d'])) $a[] = $array['d'];
 	if (isset($array['e'])) $a[] = $array['e'];
 	if (count($a) > $count) {
-		assertType('int<1, 5>', count($a));
-		assertType('array{0: mixed~null, 1?: mixed~null, 2?: mixed~null, 3?: mixed~null, 4?: mixed~null}', $a);
+		assertType('int<2, 5>', count($a));
+		assertType('list{0: mixed~null, 1: mixed~null, 2?: mixed~null, 3?: mixed~null, 4?: mixed~null}', $a);
 	} else {
-		assertType('0', count($a));
-		assertType('array{}', $a);
+		assertType('int<0, 5>', count($a)); // Could be int<0, 1>
+		assertType('array{}|array{0: mixed~null, 1?: mixed~null, 2?: mixed~null, 3?: mixed~null, 4?: mixed~null}', $a); // Could be array{}|array{0: mixed~null}
 	}
 };

tests/PHPStan/Analyser/nsrt/bug11480.php

@@ -84,7 +84,7 @@ public function intUnionCount(): void
 		if (count($x) >= $count) {
 			assertType("array{'xy'}|array{0: 'ab', 1?: 'xy'}", $x);
 		} else {
-			assertType("array{}|array{'xy'}|array{0: 'ab', 1?: 'xy'}", $x);
+			assertType("array{}", $x);
 		}
 		assertType("array{}|array{'xy'}|array{0: 'ab', 1?: 'xy'}", $x);
 	}

tests/PHPStan/Analyser/nsrt/count-type.php

@@ -64,6 +64,29 @@ public function doFooBar(
 		}
 	}
 
+	/** @param array{0: string, 1?: string} $arr */
+	public function doBar(array $arr): void
+	{
+		if (count($arr) <= 1) {
+			assertType('1', count($arr));
+			return;
+		}
+
+		assertType('2', count($arr));
+		assertType('array{string, string}', $arr);
+	}
+
+	/** @param array{0: string, 1?: string} $arr */
+	public function doBaz(array $arr): void
+	{
+		if (count($arr) > 1) {
+			assertType('2', count($arr));
+			assertType('array{string, string}', $arr);
+		}
+
+		assertType('1|2', count($arr));
+	}
+
 }
 
 /**

tests/PHPStan/Rules/Functions/CallToFunctionParametersRuleTest.php

@@ -545,6 +545,11 @@ public function testBug3608(): void
 		$this->analyse([__DIR__ . '/data/bug-3608.php'], []);
 	}
 
+	public function testBug3631(): void
+	{
+		$this->analyse([__DIR__ . '/data/bug-3631.php'], []);
+	}
+
 	public function testBug3920(): void
 	{
 		$this->analyse([__DIR__ . '/data/bug-3920.php'], []);

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

@@ -0,0 +1,27 @@
+<?php declare(strict_types = 1);
+
+namespace Bug3631;
+
+/**
+ * @return array<string>
+ */
+function someFunc(bool $flag): array
+{
+	$ids = [
+		['fa', 'foo', 'baz']
+	];
+
+	if ($flag) {
+		$ids[] = ['foo', 'bar', 'baz'];
+
+	}
+
+	if (count($ids) > 1) {
+		return array_intersect(...$ids);
+	}
+
+	return $ids[0];
+}
+
+var_dump(someFunc(true));
+var_dump(someFunc(false));