Add ArrayChunkFunctionReturnTypeExtension

Author: herndlm

conf/config.neon

@@ -979,6 +979,11 @@ services:
 		tags:
 			- phpstan.broker.dynamicFunctionReturnTypeExtension
 
+	-
+		class: PHPStan\Type\Php\ArrayChunkFunctionReturnTypeExtension
+		tags:
+			- phpstan.broker.dynamicFunctionReturnTypeExtension
+
 	-
 		class: PHPStan\Type\Php\ArrayColumnFunctionReturnTypeExtension
 		tags:

src/Type/Constant/ConstantArrayType.php

@@ -730,6 +730,25 @@ public function reverse(bool $preserveKeys = false): self
 		return $preserveKeys ? $reversed : $reversed->reindex();
 	}
 
+	/** @param positive-int $length */
+	public function chunk(int $length, bool $preserveKeys = false): self
+	{
+		$builder = ConstantArrayTypeBuilder::createEmpty();
+
+		$keyTypesCount = count($this->keyTypes);
+		for ($i = 0; $i < $keyTypesCount; $i += $length) {
+			$chunk = $this->slice($i, $length, true);
+			$builder->setOffsetValueType(null, $preserveKeys ? $chunk : $chunk->getValuesArray());
+		}
+
+		$chunks = $builder->getArray();
+		if (!$chunks instanceof self) {
+			throw new ShouldNotHappenException();
+		}
+
+		return $chunks;
+	}
+
 	private function reindex(): self
 	{
 		$keyTypes = [];

src/Type/Php/ArrayChunkFunctionReturnTypeExtension.php

@@ -0,0 +1,65 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Php;
+
+use PhpParser\Node\Expr\FuncCall;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\FunctionReflection;
+use PHPStan\Type\Accessory\NonEmptyArrayType;
+use PHPStan\Type\ArrayType;
+use PHPStan\Type\Constant\ConstantArrayType;
+use PHPStan\Type\Constant\ConstantBooleanType;
+use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\DynamicFunctionReturnTypeExtension;
+use PHPStan\Type\IntegerType;
+use PHPStan\Type\IntersectionType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeTraverser;
+use PHPStan\Type\UnionType;
+use function count;
+
+final class ArrayChunkFunctionReturnTypeExtension implements DynamicFunctionReturnTypeExtension
+{
+
+	public function isFunctionSupported(FunctionReflection $functionReflection): bool
+	{
+		return $functionReflection->getName() === 'array_chunk';
+	}
+
+	public function getTypeFromFunctionCall(FunctionReflection $functionReflection, FuncCall $functionCall, Scope $scope): ?Type
+	{
+		if (count($functionCall->getArgs()) < 2) {
+			return null;
+		}
+
+		$arrayType = $scope->getType($functionCall->getArgs()[0]->value);
+		$lengthType = $scope->getType($functionCall->getArgs()[1]->value);
+		$preserveKeysType = isset($functionCall->getArgs()[2]) ? $scope->getType($functionCall->getArgs()[2]->value) : null;
+		$preserveKeys = $preserveKeysType instanceof ConstantBooleanType ? $preserveKeysType->getValue() : false;
+
+		if (!$arrayType->isArray()->yes() || !$lengthType instanceof ConstantIntegerType || $lengthType->getValue() < 1) {
+			return null;
+		}
+
+		return TypeTraverser::map($arrayType, static function (Type $type, callable $traverse) use ($lengthType, $preserveKeys): Type {
+			if ($type instanceof UnionType || $type instanceof IntersectionType) {
+				return $traverse($type);
+			}
+
+			if ($type instanceof ConstantArrayType) {
+				return $type->chunk($lengthType->getValue(), $preserveKeys);
+			}
+
+			$chunkType = $preserveKeys ? $type : new ArrayType(new IntegerType(), $type->getIterableValueType());
+			$chunkType = TypeCombinator::intersect($chunkType, new NonEmptyArrayType());
+
+			$resultType = new ArrayType(new IntegerType(), $chunkType);
+			if ($type->isIterableAtLeastOnce()->yes()) {
+				$resultType = TypeCombinator::intersect($type, new NonEmptyArrayType());
+			}
+			return $resultType;
+		});
+	}
+
+}

tests/PHPStan/Analyser/NodeScopeResolverTest.php

@@ -650,6 +650,7 @@ public function dataFileAsserts(): iterable
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-4357.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/bug-5817.php');
 
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-chunk.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/array-column.php');
 		if (PHP_VERSION_ID >= 80000) {
 			yield from $this->gatherAssertTypes(__DIR__ . '/data/array-column-php8.php');

tests/PHPStan/Analyser/data/array-chunk.php

@@ -0,0 +1,46 @@
+<?php declare(strict_types = 1);
+
+namespace ArrayChunk;
+
+use function PHPStan\Testing\assertType;
+
+class Foo
+{
+
+	public function generalArrays(array $arr): void
+	{
+		/** @var mixed[] $arr */
+		assertType('array<int, non-empty-array<int, mixed>>', array_chunk($arr, 2));
+		assertType('array<int, non-empty-array>', array_chunk($arr, 2, true));
+
+		/** @var array<string, int> $arr */
+		assertType('array<int, non-empty-array<int, int>>', array_chunk($arr, 2));
+		assertType('array<int, non-empty-array<string, int>>', array_chunk($arr, 2, true));
+
+		/** @var non-empty-array<int|string, bool> $arr */
+		assertType('non-empty-array<int, non-empty-array<int, bool>>', array_chunk($arr, 1));
+		assertType('non-empty-array<int, non-empty-array<int|string, bool>>', array_chunk($arr, 1, true));
+	}
+
+
+	public function constantArrays(array $arr): void
+	{
+		/** @var array{a: 0, 17: 1, b: 2} $arr */
+		assertType('array{array{0, 1}, array{2}}', array_chunk($arr, 2));
+		assertType('array{array{a: 0, 17: 1}, array{b: 2}}', array_chunk($arr, 2, true));
+		assertType('array{array{0}, array{1}, array{2}}', array_chunk($arr, 1));
+		assertType('array{array{a: 0}, array{17: 1}, array{b: 2}}', array_chunk($arr, 1, true));
+	}
+
+	public function constantArraysWithOptionalKeys(array $arr): void
+	{
+		/** @var array{a: 0, b?: 1, c: 2} $arr */
+		assertType('array{array{a: 0, b?: 1, c?: 2}, array{c?: 2}}', array_chunk($arr, 2, true));
+		assertType('array{array{a: 0, b?: 1, c: 2}}', array_chunk($arr, 3, true));
+		assertType('array{array{a: 0}, array{b?: 1, c?: 2}, array{c?: 2}}', array_chunk($arr, 1, true));
+
+		/** @var array{a?: 0, b?: 1, c?: 2} $arr */
+		assertType('array{array{a?: 0, b?: 1, c?: 2}, array{c?: 2}}', array_chunk($arr, 2, true));
+	}
+
+}