Improve the return type of array_map() for constant arrays

Author: zonuexe

src/Type/Php/ArrayMapFunctionReturnTypeExtension.php

@@ -2,14 +2,21 @@
 
 namespace PHPStan\Type\Php;
 
+use PhpParser\Node\Arg;
 use PhpParser\Node\Expr\FuncCall;
+use PhpParser\Node\Name;
+use PhpParser\Node\Scalar\String_;
 use PHPStan\Analyser\Scope;
+use PHPStan\Node\Expr\TypeExpr;
+use PHPStan\Parser\ArrayMapArgVisitor;
 use PHPStan\Reflection\FunctionReflection;
 use PHPStan\Type\Accessory\AccessoryArrayListType;
 use PHPStan\Type\Accessory\NonEmptyArrayType;
 use PHPStan\Type\ArrayType;
+use PHPStan\Type\ClosureType;
 use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
 use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Constant\ConstantStringType;
 use PHPStan\Type\DynamicFunctionReturnTypeExtension;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\MixedType;
@@ -37,13 +44,27 @@ public function getTypeFromFunctionCall(FunctionReflection $functionReflection,
 		$singleArrayArgument = !isset($functionCall->getArgs()[2]);
 		$callableType = $scope->getType($functionCall->getArgs()[0]->value);
 		$callableIsNull = $callableType->isNull()->yes();
+		$callback = null;
 
 		if ($callableType->isCallable()->yes()) {
 			$valueTypes = [new NeverType()];
 			foreach ($callableType->getCallableParametersAcceptors($scope) as $parametersAcceptor) {
 				$valueTypes[] = $parametersAcceptor->getReturnType();
 			}
 			$valueType = TypeCombinator::union(...$valueTypes);
+			$callback = $functionCall->getArgs()[0]->value;
+			if ($callback instanceof String_) {
+				/** @var non-falsy-string $callName */
+				$callName = $callback->value;
+				$callback = new Name($callName);
+			} elseif ($callback instanceof FuncCall && $callback->isFirstClassCallable() &&
+				$callback->getAttribute('phpstan_cache_printer') !== null &&
+				preg_match('/\A(?<name>\\\\?[^()]+)\(...\)\z/', $callback->getAttribute('phpstan_cache_printer'), $m) === 1
+			) {
+				/** @var non-falsy-string $callName */
+				$callName = $m['name'];
+				$callback = new Name($callName);
+			}
 		} elseif ($callableIsNull) {
 			$arrayBuilder = ConstantArrayTypeBuilder::createEmpty();
 			foreach (array_slice($functionCall->getArgs(), 1) as $index => $arg) {
@@ -73,7 +94,11 @@ public function getTypeFromFunctionCall(FunctionReflection $functionReflection,
 						foreach ($constantArray->getKeyTypes() as $i => $keyType) {
 							$returnedArrayBuilder->setOffsetValueType(
 								$keyType,
-								$valueType,
+								$callback === null
+									? $valueType
+									: $scope->getType(new FuncCall($callback, [
+										new Arg(new TypeExpr($constantArray->getValueTypes()[$i])),
+									])),
 								$constantArray->isOptionalKey($i),
 							);
 						}

tests/PHPStan/Analyser/nsrt/array-map-callable81.php

@@ -0,0 +1,43 @@
+<?php
+
+namespace ArrayMapCallable81;
+
+use function PHPStan\Testing\assertType;
+use function strval as str;
+
+class Foo
+{
+	/**
+	 * @template T of int
+	 * @param T $n
+	 * @return (T is 3 ? 'Fizz' : (T is 5 ? 'Buzz' : T))
+	 */
+	public static function fizzbuzz(int $n): int|string
+	{
+		return match ($n) {
+			3 => 'Fizz',
+			5 => 'Buzz',
+			default => $n,
+		};
+	}
+
+	public function doFoo(): void
+	{
+		$a = range(0, 1);
+
+		assertType("array{'0', '1'}", array_map('strval', $a));
+		assertType("array{'0', '1'}", array_map(strval(...), $a));
+		assertType("array{'0', '1'}", array_map(str(...), $a));
+		assertType("array{'0'|'1', '0'|'1'}", array_map(fn ($v) => strval($v), $a));
+		assertType("array{'0'|'1', '0'|'1'}", array_map(fn ($v) => (string)$v, $a));
+	}
+
+	public function doFizzBuzz(): void
+	{
+		assertType("array{1, 2, 'Fizz', 4, 'Buzz', 6}", array_map([__CLASS__, 'fizzbuzz'], range(1, 6)));
+		assertType("array{1, 2, 'Fizz', 4, 'Buzz', 6}", array_map([$this, 'fizzbuzz'], range(1, 6)));
+		assertType("array{1|'Buzz'|'Fizz', 2|'Buzz'|'Fizz', 3|'Buzz'|'Fizz', 4|'Buzz'|'Fizz', 5|'Buzz'|'Fizz', 6|'Buzz'|'Fizz'}", array_map(self::fizzbuzz(...), range(1, 6)));
+		assertType("array{1|'Buzz'|'Fizz', 2|'Buzz'|'Fizz', 3|'Buzz'|'Fizz', 4|'Buzz'|'Fizz', 5|'Buzz'|'Fizz', 6|'Buzz'|'Fizz'}", array_map($this->fizzbuzz(...), range(1, 6)));
+	}
+
+}