Consider optional keys in ConstantArrayType removeFirst and removeLast

phpstan · May 26, 2022 · b2ef801 · b2ef801
1 parent f22c92d
commit b2ef801
Show file tree

Hide file tree

Showing 5 changed files with 146 additions and 21 deletions.
diff --git a/src/Analyser/NodeScopeResolver.php b/src/Analyser/NodeScopeResolver.php
@@ -126,6 +126,7 @@
 use PHPStan\Type\Generic\TemplateTypeHelper;
 use PHPStan\Type\Generic\TemplateTypeMap;
 use PHPStan\Type\IntegerType;
+use PHPStan\Type\IntersectionType;
 use PHPStan\Type\MixedType;
 use PHPStan\Type\NeverType;
 use PHPStan\Type\NullType;
@@ -1813,29 +1814,27 @@ function (MutatingScope $scope) use ($expr, $nodeCallback, $context): Expression
 				&& count($expr->getArgs()) >= 1
 			) {
 				$arrayArg = $expr->getArgs()[0]->value;
-				$constantArrays = TypeUtils::getConstantArrays($scope->getType($arrayArg));
+				$arrayArgType = $scope->getType($arrayArg);
 				$scope = $scope->invalidateExpression($arrayArg);
-				if (count($constantArrays) > 0) {
-					$resultArrayTypes = [];
 
-					foreach ($constantArrays as $constantArray) {
-						if ($functionReflection->getName() === 'array_pop') {
-							$resultArrayTypes[] = $constantArray->removeLast();
-						} else {
-							$resultArrayTypes[] = $constantArray->removeFirst();
-						}
+				$functionName = $functionReflection->getName();
+				$arrayArgType = TypeTraverser::map($arrayArgType, static function (Type $type, callable $traverse) use ($functionName): Type {
+					if ($type instanceof UnionType || $type instanceof IntersectionType) {
+						return $traverse($type);
 					}
-
-					$scope = $scope->specifyExpressionType(
-						$arrayArg,
-						TypeCombinator::union(...$resultArrayTypes),
-					);
-				} else {
-					$arrays = TypeUtils::getAnyArrays($scope->getType($arrayArg));
-					if (count($arrays) > 0) {
-						$scope = $scope->specifyExpressionType($arrayArg, TypeCombinator::union(...$arrays));
+					if ($type instanceof ConstantArrayType) {
+						return $functionName === 'array_pop' ? $type->removeLast() : $type->removeFirst();
 					}
-				}
+					if ($type->isIterableAtLeastOnce()->yes()) {
+						return $type->toArray();
+					}
+					return $type;
+				});
+
+				$scope = $scope->specifyExpressionType(
+					$arrayArg,
+					$arrayArgType,
+				);
 			}
 
 			if (

diff --git a/src/Type/Constant/ConstantArrayType.php b/src/Type/Constant/ConstantArrayType.php
@@ -558,7 +558,17 @@ public function removeLast(): self
 		$keyTypes = $this->keyTypes;
 		$valueTypes = $this->valueTypes;
 		$optionalKeys = $this->optionalKeys;
-		unset($optionalKeys[$i]);
+
+		if ($this->isOptionalKey($i)) {
+			unset($optionalKeys[$i]);
+			// Removing the last optional element makes the previous non-optional element optional
+			for ($j = $i - 1; $j >= 0; $j--) {
+				if (!$this->isOptionalKey($j)) {
+					$optionalKeys[] = $j;
+					break;
+				}
+			}
+		}
 
 		$removedKeyType = array_pop($keyTypes);
 		array_pop($valueTypes);
@@ -577,17 +587,25 @@ public function removeLast(): self
 	public function removeFirst(): Type
 	{
 		$builder = ConstantArrayTypeBuilder::createEmpty();
+		$makeNextNonOptionalOptional = false;
 		foreach ($this->keyTypes as $i => $keyType) {
+			$isOptional = $this->isOptionalKey($i);
 			if ($i === 0) {
+				$makeNextNonOptionalOptional = $this->isOptionalKey($i);
 				continue;
 			}
 
+			if (!$isOptional && $makeNextNonOptionalOptional) {
+				$isOptional = true;
+				$makeNextNonOptionalOptional = false;
+			}
+
 			$valueType = $this->valueTypes[$i];
 			if ($keyType instanceof ConstantIntegerType) {
 				$keyType = null;
 			}
 
-			$builder->setOffsetValueType($keyType, $valueType);
+			$builder->setOffsetValueType($keyType, $valueType, $isOptional);
 		}
 
 		return $builder->getArray();

diff --git a/tests/PHPStan/Analyser/NodeScopeResolverTest.php b/tests/PHPStan/Analyser/NodeScopeResolverTest.php
@@ -179,6 +179,7 @@ public function dataFileAsserts(): iterable
 
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-3985.php');
 
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-shift.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-slice.php');
 
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-3990.php');
@@ -800,6 +801,7 @@ public function dataFileAsserts(): iterable
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-6439.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-6748.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-search-type-specifying.php');
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-pop.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-push.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-replace.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-reverse.php');

diff --git a/tests/PHPStan/Analyser/data/array-pop.php b/tests/PHPStan/Analyser/data/array-pop.php
@@ -0,0 +1,53 @@
+<?php declare(strict_types = 1);
+
+namespace ArrayPop;
+
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+
+	public function nonEmpty(array $arr): void
+	{
+		/** @var non-empty-array<string> $arr */
+		assertType('string', array_pop($arr));
+		assertType('array<string>', $arr);
+	}
+
+	public function normalArrays(array $arr): void
+	{
+		/** @var array<int, string> $arr */
+		assertType('string|null', array_pop($arr));
+		assertType('array<int, string>', $arr);
+	}
+
+	public function compoundTypes(array $arr): void
+	{
+		/** @var string[]|int[] $arr */
+		assertType('int|string|null', array_pop($arr));
+		assertType('array<int|string>', $arr);
+	}
+
+	public function constantArrays(array $arr): void
+	{
+		/** @var array{a: 0, b: 1, c: 2} $arr */
+		assertType('2', array_pop($arr));
+		assertType('array{a: 0, b: 1}', $arr);
+	}
+
+	public function constantArraysWithOptionalKeys(array $arr): void
+	{
+		/** @var array{a?: 0, b: 1, c: 2} $arr */
+		assertType('0|2', array_pop($arr)); // should be 2
+		assertType('array{a?: 0, b: 1}', $arr);
+
+		/** @var array{a: 0, b?: 1, c: 2} $arr */
+		assertType('1|2', array_pop($arr)); // should be 2
+		assertType('array{a: 0, b?: 1}', $arr);
+
+		/** @var array{a: 0, b: 1, c?: 2} $arr */
+		assertType('1|2', array_pop($arr));
+		assertType('array{a: 0, b?: 1}', $arr);
+	}
+
+}
diff --git a/tests/PHPStan/Analyser/data/array-shift.php b/tests/PHPStan/Analyser/data/array-shift.php
@@ -0,0 +1,53 @@
+<?php declare(strict_types = 1);
+
+namespace ArrayShift;
+
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+
+	public function nonEmpty(array $arr): void
+	{
+		/** @var non-empty-array<string> $arr */
+		assertType('string', array_shift($arr));
+		assertType('array<string>', $arr);
+	}
+
+	public function normalArrays(array $arr): void
+	{
+		/** @var array<int, string> $arr */
+		assertType('string|null', array_shift($arr));
+		assertType('array<int, string>', $arr);
+	}
+
+	public function compoundTypes(array $arr): void
+	{
+		/** @var string[]|int[] $arr */
+		assertType('int|string|null', array_shift($arr));
+		assertType('array<int|string>', $arr);
+	}
+
+	public function constantArrays(array $arr): void
+	{
+		/** @var array{a: 0, b: 1, c: 2} $arr */
+		assertType('0', array_shift($arr));
+		assertType('array{b: 1, c: 2}', $arr);
+	}
+
+	public function constantArraysWithOptionalKeys(array $arr): void
+	{
+		/** @var array{a?: 0, b: 1, c: 2} $arr */
+		assertType('1', array_shift($arr)); // should be 0|1
+		assertType('array{b?: 1, c: 2}', $arr);
+
+		/** @var array{a: 0, b?: 1, c: 2} $arr */
+		assertType('0', array_shift($arr));
+		assertType('array{b?: 1, c: 2}', $arr);
+
+		/** @var array{a: 0, b: 1, c?: 2} $arr */
+		assertType('0', array_shift($arr));
+		assertType('array{b: 1, c?: 2}', $arr);
+	}
+
+}