More precise count($array) after array_* functions

Author: staabm

src/Analyser/NodeScopeResolver.php

@@ -38,6 +38,7 @@
 use PhpParser\Node\Expr\Variable;
 use PhpParser\Node\Identifier;
 use PhpParser\Node\Name;
+use PhpParser\Node\Scalar\Int_;
 use PhpParser\Node\Stmt\Break_;
 use PhpParser\Node\Stmt\Class_;
 use PhpParser\Node\Stmt\Continue_;
@@ -174,6 +175,7 @@
 use PHPStan\Type\Generic\TemplateTypeMap;
 use PHPStan\Type\Generic\TemplateTypeVariance;
 use PHPStan\Type\Generic\TemplateTypeVarianceMap;
+use PHPStan\Type\IntegerRangeType;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\IntersectionType;
 use PHPStan\Type\MixedType;
@@ -2628,29 +2630,72 @@ static function (): void {
 				&& in_array($functionReflection->getName(), ['array_pop', 'array_shift'], true)
 				&& count($expr->getArgs()) >= 1
 			) {
-				$arrayArg = $expr->getArgs()[0]->value;
+				$args = $expr->getArgs();
+				$arrayArg = $args[0]->value;
 
 				$arrayArgType = $scope->getType($arrayArg);
 				$arrayArgNativeType = $scope->getNativeType($arrayArg);
 
+				$countArrayExpr = new FuncCall(new Name('count'), [$args[0]]);
+				$hasCountExpr = $scope->hasExpressionType($countArrayExpr)->yes();
+				if ($hasCountExpr) {
+					$countType = $scope->getType(new BinaryOp\Minus($countArrayExpr, new Int_(1)));
+					$countNativeType = $scope->getType(new BinaryOp\Minus($countArrayExpr, new Int_(1)));
+				}
+
 				$isArrayPop = $functionReflection->getName() === 'array_pop';
 				$scope = $scope->invalidateExpression($arrayArg)->assignExpression(
 					$arrayArg,
 					$isArrayPop ? $arrayArgType->popArray() : $arrayArgType->shiftArray(),
 					$isArrayPop ? $arrayArgNativeType->popArray() : $arrayArgNativeType->shiftArray(),
 				);
+
+				if (
+					$hasCountExpr
+					&& IntegerRangeType::fromInterval(0, null)->isSuperTypeOf($countType)->yes()
+					&& IntegerRangeType::fromInterval(0, null)->isSuperTypeOf($countNativeType)->yes()
+				) {
+					$scope = $scope->assignExpression($countArrayExpr, $countType, $countNativeType);
+				}
 			}
 
 			if (
 				$functionReflection !== null
 				&& in_array($functionReflection->getName(), ['array_push', 'array_unshift'], true)
 				&& count($expr->getArgs()) >= 2
 			) {
+				$args = $expr->getArgs();
 				$arrayType = $this->getArrayFunctionAppendingType($functionReflection, $scope, $expr);
 				$arrayNativeType = $this->getArrayFunctionAppendingType($functionReflection, $scope->doNotTreatPhpDocTypesAsCertain(), $expr);
 
-				$arrayArg = $expr->getArgs()[0]->value;
+				$arrayArg = $args[0]->value;
+				$addedElementsCount = count($args) - 1;
+				for ($i = 1; $i < count($args); $i++) {
+					if ($args[$i]->unpack) {
+						$addedElementsCount = null;
+						break;
+					}
+				}
+
+				$countArrayExpr = new FuncCall(new Name('count'), [$args[0]]);
+				$hasCountExpr = $scope->hasExpressionType($countArrayExpr)->yes();
+				if ($hasCountExpr && $addedElementsCount !== null) {
+					$countType = $scope->getType(new BinaryOp\Plus($countArrayExpr, new Int_($addedElementsCount)));
+					$countNativeType = $scope->getType(new BinaryOp\Plus($countArrayExpr, new Int_($addedElementsCount)));
+				} else {
+					$countType = IntegerRangeType::fromInterval($addedElementsCount, null);
+					$countNativeType = IntegerRangeType::fromInterval($addedElementsCount, null);
+				}
+
 				$scope = $scope->invalidateExpression($arrayArg)->assignExpression($arrayArg, $arrayType, $arrayNativeType);
+
+				if (
+					IntegerRangeType::fromInterval(0, null)->isSuperTypeOf($countType)->yes()
+					&& IntegerRangeType::fromInterval(0, null)->isSuperTypeOf($countNativeType)->yes()
+				) {
+					$scope = $scope->assignExpression($countArrayExpr, $countType, $countNativeType);
+				}
+
 			}
 
 			if (

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

@@ -18,10 +18,10 @@ function (array $input) {
 	\assert(count($input) > 0);
 	assertType('int<1, max>', count($input));
 	array_unshift($input, 'test');
-	assertType('int<1, max>', count($input));
+	assertType('int<2, max>', count($input));
 
 	\assert(count($input) > 0);
-	assertType('int<1, max>', count($input));
+	assertType('int<2, max>', count($input));
 	array_push($input, 'nope');
-	assertType('int<1, max>', count($input));
+	assertType('int<3, max>', count($input));
 };

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

@@ -0,0 +1,64 @@
+<?php
+
+namespace Bug7804;
+
+use function PHPStan\Testing\assertType;
+
+/** @param array<int, string> $headers */
+function headers(array $headers): void
+{
+	assertType('int<0, max>', count($headers));
+	if (count($headers) >= 4) {
+		assertType('int<4, max>', count($headers));
+		array_pop($headers);
+		assertType('int<3, max>', count($headers));
+		array_pop($headers);
+		assertType('int<2, max>', count($headers));
+		array_pop($headers);
+		assertType('int<1, max>', count($headers));
+		array_pop($headers);
+		assertType('int<0, max>', count($headers));
+		array_pop($headers);
+		assertType('int<0, max>', count($headers));
+	}
+	assertType('int<0, max>', count($headers));
+}
+
+function doPop(array $arr) {
+	assertType('int<0, max>', count($arr));
+	array_pop($arr);
+	assertType('int<0, max>', count($arr));
+
+	if (count($arr) === 2) {
+		assertType('2', count($arr));
+		array_pop($arr);
+		assertType('1', count($arr));
+	}
+	assertType('int<0, 1>|int<3, max>', count($arr));
+}
+
+function doShift(array $arr) {
+	assertType('int<0, max>', count($arr));
+	array_shift($arr);
+	assertType('int<0, max>', count($arr));
+}
+
+function doPush(array $arr, int $i) {
+	assertType('int<0, max>', count($arr));
+	array_push($arr, $i);
+	assertType('int<1, max>', count($arr));
+	array_push($arr, 3, $i, false, null);
+	assertType('int<5, max>', count($arr));
+}
+
+function doPushVariadic(array $arr, array $arr2) {
+	assertType('int<0, max>', count($arr));
+	array_push($arr, ...$arr2);
+	assertType('int<0, max>', count($arr));
+}
+
+function doUnshift(array $arr, bool $b) {
+	assertType('int<0, max>', count($arr));
+	array_unshift($arr, $b);
+	assertType('int<1, max>', count($arr));
+}