Handle all hydration mode in QueryResultDynamicReturnTypeExtension

src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php

@@ -10,14 +10,22 @@
 use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\Accessory\AccessoryArrayListType;
 use PHPStan\Type\ArrayType;
+use PHPStan\Type\BenevolentUnionType;
 use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Doctrine\ObjectMetadataResolver;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\IntegerType;
 use PHPStan\Type\IterableType;
+use PHPStan\Type\MixedType;
 use PHPStan\Type\NullType;
+use PHPStan\Type\ObjectWithoutClassType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeTraverser;
+use PHPStan\Type\TypeUtils;
+use PHPStan\Type\TypeWithClassName;
 use PHPStan\Type\VoidType;
+use function count;
 
 final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -32,14 +40,32 @@ final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturn
 		'getSingleResult' => 0,
 	];
 
+	private const METHOD_HYDRATION_MODE = [
+		'getArrayResult' => AbstractQuery::HYDRATE_ARRAY,
+		'getScalarResult' => AbstractQuery::HYDRATE_SCALAR,
+		'getSingleColumnResult' => AbstractQuery::HYDRATE_SCALAR_COLUMN,
+		'getSingleScalarResult' => AbstractQuery::HYDRATE_SINGLE_SCALAR,
+	];
+
+	/** @var ObjectMetadataResolver */
+	private $objectMetadataResolver;
+
+	public function __construct(
+		ObjectMetadataResolver $objectMetadataResolver
+	)
+	{
+		$this->objectMetadataResolver = $objectMetadataResolver;
+	}
+
 	public function getClass(): string
 	{
 		return AbstractQuery::class;
 	}
 
 	public function isMethodSupported(MethodReflection $methodReflection): bool
 	{
-		return isset(self::METHOD_HYDRATION_MODE_ARG[$methodReflection->getName()]);
+		return isset(self::METHOD_HYDRATION_MODE_ARG[$methodReflection->getName()])
+			|| isset(self::METHOD_HYDRATION_MODE[$methodReflection->getName()]);
 	}
 
 	public function getTypeFromMethodCall(
@@ -50,21 +76,23 @@ public function getTypeFromMethodCall(
 	{
 		$methodName = $methodReflection->getName();
 
-		if (!isset(self::METHOD_HYDRATION_MODE_ARG[$methodName])) {
-			throw new ShouldNotHappenException();
-		}
-
-		$argIndex = self::METHOD_HYDRATION_MODE_ARG[$methodName];
-		$args = $methodCall->getArgs();
+		if (isset(self::METHOD_HYDRATION_MODE[$methodName])) {
+			$hydrationMode = new ConstantIntegerType(self::METHOD_HYDRATION_MODE[$methodName]);
+		} elseif (isset(self::METHOD_HYDRATION_MODE_ARG[$methodName])) {
+			$argIndex = self::METHOD_HYDRATION_MODE_ARG[$methodName];
+			$args = $methodCall->getArgs();
 
-		if (isset($args[$argIndex])) {
-			$hydrationMode = $scope->getType($args[$argIndex]->value);
+			if (isset($args[$argIndex])) {
+				$hydrationMode = $scope->getType($args[$argIndex]->value);
+			} else {
+				$parametersAcceptor = ParametersAcceptorSelector::selectSingle(
+					$methodReflection->getVariants()
+				);
+				$parameter = $parametersAcceptor->getParameters()[$argIndex];
+				$hydrationMode = $parameter->getDefaultValue() ?? new NullType();
+			}
 		} else {
-			$parametersAcceptor = ParametersAcceptorSelector::selectSingle(
-				$methodReflection->getVariants()
-			);
-			$parameter = $parametersAcceptor->getParameters()[$argIndex];
-			$hydrationMode = $parameter->getDefaultValue() ?? new NullType();
+			throw new ShouldNotHappenException();
 		}
 
 		$queryType = $scope->getType($methodCall->var);
@@ -98,23 +126,54 @@ private function getMethodReturnTypeForHydrationMode(
 			return $this->originalReturnType($methodReflection);
 		}
 
-		if (!$this->isObjectHydrationMode($hydrationMode)) {
-			// We support only HYDRATE_OBJECT. For other hydration modes, we
-			// return the declared return type of the method.
+		if (!$hydrationMode instanceof ConstantIntegerType) {
 			return $this->originalReturnType($methodReflection);
 		}
 
+		$singleResult = false;
+		switch ($hydrationMode->getValue()) {
+			case AbstractQuery::HYDRATE_OBJECT:
+				break;
+			case AbstractQuery::HYDRATE_ARRAY:
+				$queryResultType = $this->getArrayHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SCALAR:
+				$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SINGLE_SCALAR:
+				$singleResult = true;
+				$queryResultType = $this->getSingleScalarHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SIMPLEOBJECT:
+				$queryResultType = $this->getSimpleObjectHydratedReturnType($queryResultType);
+				break;
+			case AbstractQuery::HYDRATE_SCALAR_COLUMN:
+				$queryResultType = $this->getScalarColumnHydratedReturnType($queryResultType);
+				break;
+			default:
+				return $this->originalReturnType($methodReflection);
+		}
+
 		switch ($methodReflection->getName()) {
 			case 'getSingleResult':
 				return $queryResultType;
 			case 'getOneOrNullResult':
-				return TypeCombinator::addNull($queryResultType);
+				$nullableQueryResultType = TypeCombinator::addNull($queryResultType);
+				if ($queryResultType instanceof BenevolentUnionType) {
+					$nullableQueryResultType = TypeUtils::toBenevolentUnion($nullableQueryResultType);
+				}
+
+				return $nullableQueryResultType;
 			case 'toIterable':
 				return new IterableType(
 					$queryKeyType->isNull()->yes() ? new IntegerType() : $queryKeyType,
 					$queryResultType
 				);
 			default:
+				if ($singleResult) {
+					return $queryResultType;
+				}
+
 				if ($queryKeyType->isNull()->yes()) {
 					return AccessoryArrayListType::intersectWith(new ArrayType(
 						new IntegerType(),
@@ -128,13 +187,104 @@ private function getMethodReturnTypeForHydrationMode(
 		}
 	}
 
-	private function isObjectHydrationMode(Type $type): bool
+	private function getArrayHydratedReturnType(Type $queryResultType): Type
+	{
+		$objectManager = $this->objectMetadataResolver->getObjectManager();
+
+		return TypeTraverser::map(
+			$queryResultType,
+			static function (Type $type, callable $traverse) use ($objectManager): Type {
+				$isObject = (new ObjectWithoutClassType())->isSuperTypeOf($type);
+				if ($isObject->no()) {
+					return $traverse($type);
+				}
+				if (
+					$isObject->maybe()
+					|| !$type instanceof TypeWithClassName
+					|| $objectManager === null
+				) {
+					return new MixedType();
+				}
+
+				if (!$objectManager->getMetadataFactory()->hasMetadataFor($type->getClassName())) {
+					return $traverse($type);
+				}
+
+				// We could return `new ArrayTyp(new MixedType(), new MixedType())`
+				// but the lack of precision in the array keys/values would give false positive
+				// @see https://github.com/phpstan/phpstan-doctrine/pull/412#issuecomment-1497092934
+				return new MixedType();
+			}
+		);
+	}
+
+	private function getScalarHydratedReturnType(Type $queryResultType): Type
+	{
+		if (!$queryResultType->isArray()->yes()) {
+			return new ArrayType(new MixedType(), new MixedType());
+		}
+
+		foreach ($queryResultType->getArrays() as $arrayType) {
+			$itemType = $arrayType->getItemType();
+
+			if (
+				!(new ObjectWithoutClassType())->isSuperTypeOf($itemType)->no()
+				|| !$itemType->isArray()->no()
+			) {
+				return new ArrayType(new MixedType(), new MixedType());
+			}
+		}
+
+		return $queryResultType;
+	}
+
+	private function getSimpleObjectHydratedReturnType(Type $queryResultType): Type
+	{
+		if ((new ObjectWithoutClassType())->isSuperTypeOf($queryResultType)->yes()) {
+			return $queryResultType;
+		}
+
+		return new MixedType();
+	}
+
+	private function getSingleScalarHydratedReturnType(Type $queryResultType): Type
 	{
-		if (!$type instanceof ConstantIntegerType) {
-			return false;
+		$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+		if (!$queryResultType->isConstantArray()->yes()) {
+			return new MixedType();
+		}
+
+		$types = [];
+		foreach ($queryResultType->getConstantArrays() as $constantArrayType) {
+			$values = $constantArrayType->getValueTypes();
+			if (count($values) !== 1) {
+				return new MixedType();
+			}
+
+			$types[] = $constantArrayType->getFirstIterableValueType();
+		}
+
+		return TypeCombinator::union(...$types);
+	}
+
+	private function getScalarColumnHydratedReturnType(Type $queryResultType): Type
+	{
+		$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+		if (!$queryResultType->isConstantArray()->yes()) {
+			return new MixedType();
+		}
+
+		$types = [];
+		foreach ($queryResultType->getConstantArrays() as $constantArrayType) {
+			$values = $constantArrayType->getValueTypes();
+			if (count($values) !== 1) {
+				return new MixedType();
+			}
+
+			$types[] = $constantArrayType->getFirstIterableValueType();
 		}
 
-		return $type->getValue() === AbstractQuery::HYDRATE_OBJECT;
+		return TypeCombinator::union(...$types);
 	}
 
 	private function originalReturnType(MethodReflection $methodReflection): Type

src/Type/Doctrine/Query/QueryResultTypeWalker.php

@@ -107,6 +107,9 @@ class QueryResultTypeWalker extends SqlWalker
 	/** @var bool */
 	private $hasGroupByClause;
 
+	/** @var bool */
+	private $hasWhereClause;
+
 	/**
 	 * @param Query<mixed> $query
 	 */
@@ -135,6 +138,7 @@ public function __construct($query, $parserResult, array $queryComponents)
 		$this->nullableQueryComponents = [];
 		$this->hasAggregateFunction = false;
 		$this->hasGroupByClause = false;
+		$this->hasWhereClause = false;
 
 		// The object is instantiated by Doctrine\ORM\Query\Parser, so receiving
 		// dependencies through the constructor is not an option. Instead, we
@@ -177,6 +181,7 @@ public function walkSelectStatement(AST\SelectStatement $AST)
 		$this->typeBuilder->setSelectQuery();
 		$this->hasAggregateFunction = $this->hasAggregateFunction($AST);
 		$this->hasGroupByClause = $AST->groupByClause !== null;
+		$this->hasWhereClause = $AST->whereClause !== null;
 
 		$this->walkFromClause($AST->fromClause);
 
@@ -795,7 +800,7 @@ public function walkSelectExpression($selectExpression)
 
 			$type = $this->resolveDoctrineType($typeName, $enumType, $nullable);
 
-			$this->typeBuilder->addScalar($resultAlias, $type);
+			$this->addScalar($resultAlias, $type);
 
 			return '';
 		}
@@ -841,21 +846,32 @@ public function walkSelectExpression($selectExpression)
 				// the driver and PHP version.
 				// Here we assume that the value may or may not be casted to
 				// string by the driver.
-				$type = TypeTraverser::map($type, static function (Type $type, callable $traverse): Type {
+				$casted = false;
+				$originalType = $type;
+
+				$type = TypeTraverser::map($type, static function (Type $type, callable $traverse) use (&$casted): Type {
 					if ($type instanceof UnionType || $type instanceof IntersectionType) {
 						return $traverse($type);
 					}
 					if ($type instanceof IntegerType || $type instanceof FloatType) {
+						$casted = true;
 						return TypeCombinator::union($type->toString(), $type);
 					}
 					if ($type instanceof BooleanType) {
+						$casted = true;
 						return TypeCombinator::union($type->toInteger()->toString(), $type);
 					}
 					return $traverse($type);
 				});
+
+				// Since we made supposition about possibly casted values,
+				// we can only provide a benevolent union.
+				if ($casted && $type instanceof UnionType && !$originalType->equals($type)) {
+					$type = TypeUtils::toBenevolentUnion($type);
+				}
 			}
 
-			$this->typeBuilder->addScalar($resultAlias, $type);
+			$this->addScalar($resultAlias, $type);
 
 			return '';
 		}
@@ -1276,6 +1292,21 @@ public function walkResultVariable($resultVariable)
 		return $this->marshalType(new MixedType());
 	}
 
+	/**
+	 * @param array-key $alias
+	 */
+	private function addScalar($alias, Type $type): void
+	{
+		// Since we don't check the condition inside the WHERE
+		// conditions, we cannot be sure all the union types are correct.
+		// For exemple, a condition `WHERE foo.bar IS NOT NULL` could be added.
+		if ($this->hasWhereClause && $type instanceof UnionType) {
+			$type = TypeUtils::toBenevolentUnion($type);
+		}
+
+		$this->typeBuilder->addScalar($alias, $type);
+	}
+
 	private function unmarshalType(string $marshalledType): Type
 	{
 		$type = unserialize($marshalledType);