From c96764cec0f3aa9088535c949eb9192d0a3321f9 Mon Sep 17 00:00:00 2001
From: Vincent Langlet <vincentlanglet@hotmail.fr>
Date: Mon, 26 Dec 2022 21:06:02 +0100
Subject: [PATCH] Improve QueryResultDynamicReturnTypeExtension

---
 .../QueryResultDynamicReturnTypeExtension.php | 141 ++++++-
 .../Doctrine/data/QueryResult/queryResult.php | 366 +++++++++++++++++-
 2 files changed, 476 insertions(+), 31 deletions(-)

diff --git a/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php b/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php
index 1bd41a55..c9ae9593 100644
--- a/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php
+++ b/src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php
@@ -10,6 +10,7 @@
 use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\Accessory\AccessoryArrayListType;
 use PHPStan\Type\ArrayType;
+use PHPStan\Type\Constant\ConstantArrayType;
 use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\GenericTypeVariableResolver;
@@ -17,10 +18,13 @@
 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\TypeWithClassName;
 use PHPStan\Type\VoidType;
+use function count;
 
 final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -35,6 +39,13 @@ 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,
+	];
+
 	public function getClass(): string
 	{
 		return AbstractQuery::class;
@@ -42,7 +53,8 @@ public function getClass(): string
 
 	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(
@@ -53,21 +65,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);
@@ -131,12 +145,32 @@ 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);
 		}
 
+		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:
+				$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;
@@ -161,13 +195,78 @@ private function getMethodReturnTypeForHydrationMode(
 		}
 	}
 
-	private function isObjectHydrationMode(Type $type): bool
+	private function getArrayHydratedReturnType(Type $queryResultType): Type
 	{
-		if (!$type instanceof ConstantIntegerType) {
-			return false;
+		return TypeTraverser::map(
+			$queryResultType,
+			static function (Type $type, callable $traverse): Type {
+				$isObject = (new ObjectWithoutClassType())->isSuperTypeOf($type);
+				if ($isObject->yes()) {
+					return new ArrayType(new MixedType(), new MixedType());
+				}
+				if ($isObject->maybe()) {
+					return new MixedType();
+				}
+
+				return $traverse($type);
+			}
+		);
+	}
+
+	private function getScalarHydratedReturnType(Type $queryResultType): Type
+	{
+		if (!$queryResultType instanceof ArrayType) {
+			return new ArrayType(new MixedType(), new MixedType());
+		}
+
+		$itemType = $queryResultType->getItemType();
+		$hasNoObject = (new ObjectWithoutClassType())->isSuperTypeOf($itemType)->no();
+		$hasNoArray = $itemType->isArray()->no();
+
+		if ($hasNoArray && $hasNoObject) {
+			return $queryResultType;
+		}
+
+		return new ArrayType(new MixedType(), new MixedType());
+	}
+
+	private function getSimpleObjectHydratedReturnType(Type $queryResultType): Type
+	{
+		if ((new ObjectWithoutClassType())->isSuperTypeOf($queryResultType)->yes()) {
+			return $queryResultType;
+		}
+
+		return new MixedType();
+	}
+
+	private function getSingleScalarHydratedReturnType(Type $queryResultType): Type
+	{
+		$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+		if (!$queryResultType instanceof ConstantArrayType) {
+			return new ArrayType(new MixedType(), new MixedType());
+		}
+
+		$values = $queryResultType->getValueTypes();
+		if (count($values) !== 1) {
+			return new ArrayType(new MixedType(), new MixedType());
+		}
+
+		return $queryResultType;
+	}
+
+	private function getScalarColumnHydratedReturnType(Type $queryResultType): Type
+	{
+		$queryResultType = $this->getScalarHydratedReturnType($queryResultType);
+		if (!$queryResultType instanceof ConstantArrayType) {
+			return new MixedType();
+		}
+
+		$values = $queryResultType->getValueTypes();
+		if (count($values) !== 1) {
+			return new MixedType();
 		}
 
-		return $type->getValue() === AbstractQuery::HYDRATE_OBJECT;
+		return $queryResultType->getFirstIterableValueType();
 	}
 
 	private function originalReturnType(MethodReflection $methodReflection): Type
diff --git a/tests/Type/Doctrine/data/QueryResult/queryResult.php b/tests/Type/Doctrine/data/QueryResult/queryResult.php
index 18a1faa9..3589b685 100644
--- a/tests/Type/Doctrine/data/QueryResult/queryResult.php
+++ b/tests/Type/Doctrine/data/QueryResult/queryResult.php
@@ -143,11 +143,11 @@ public function testReturnTypeOfQueryMethodsWithExplicitObjectHydrationMode(Enti
 	}
 
 	/**
-	 * Test that we properly infer the return type of Query methods with explicit hydration mode that is not HYDRATE_OBJECT
+	 * Test that we properly infer the return type of Query methods with explicit hydration mode of HYDRATE_ARRAY
 	 *
-	 * We are never able to infer the return type here
+	 * We can infer the return type by changing every object by an array
 	 */
-	public function testReturnTypeOfQueryMethodsWithExplicitNonObjectHydrationMode(EntityManagerInterface $em): void
+	public function testReturnTypeOfQueryMethodsWithExplicitArrayHydrationMode(EntityManagerInterface $em): void
 	{
 		$query = $em->createQuery('
 			SELECT		m
@@ -155,35 +155,381 @@ public function testReturnTypeOfQueryMethodsWithExplicitNonObjectHydrationMode(E
 		');
 
 		assertType(
-			'mixed',
+			'array<array>',
 			$query->getResult(AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'iterable',
+			'array<array>',
+			$query->getArrayResult()
+		);
+		assertType(
+			'iterable<array>',
 			$query->toIterable([], AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'mixed',
+			'array<array>',
 			$query->execute(null, AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'mixed',
+			'array<array>',
 			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'mixed',
+			'array<array>',
 			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'mixed',
+			'array',
 			$query->getSingleResult(AbstractQuery::HYDRATE_ARRAY)
 		);
 		assertType(
-			'mixed',
+			'array|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_ARRAY)
+		);
+
+		$query = $em->createQuery('
+			SELECT		m.intColumn, m.stringNullColumn
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->getResult(AbstractQuery::HYDRATE_ARRAY)
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->getArrayResult()
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->execute(null, AbstractQuery::HYDRATE_ARRAY)
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_ARRAY)
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_ARRAY)
+		);
+		assertType(
+			'array{intColumn: int, stringNullColumn: string|null}',
+			$query->getSingleResult(AbstractQuery::HYDRATE_ARRAY)
+		);
+		assertType(
+			'array{intColumn: int, stringNullColumn: string|null}|null',
 			$query->getOneOrNullResult(AbstractQuery::HYDRATE_ARRAY)
 		);
 	}
 
+	/**
+	 * Test that we properly infer the return type of Query methods with explicit hydration mode of HYDRATE_SCALAR
+	 */
+	public function testReturnTypeOfQueryMethodsWithExplicitScalarHydrationMode(EntityManagerInterface $em): void
+	{
+		$query = $em->createQuery('
+			SELECT		m
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<array>',
+			$query->getResult(AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->getScalarResult()
+		);
+		assertType(
+			'iterable<array>',
+			$query->toIterable([], AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->execute(null, AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SCALAR)
+		);
+
+		$query = $em->createQuery('
+			SELECT		m.intColumn, m.stringNullColumn
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->getResult(AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->getScalarResult()
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->execute(null, AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array<array{intColumn: int, stringNullColumn: string|null}>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array{intColumn: int, stringNullColumn: string|null}',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SCALAR)
+		);
+		assertType(
+			'array{intColumn: int, stringNullColumn: string|null}|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SCALAR)
+		);
+	}
+
+	/**
+	 * Test that we properly infer the return type of Query methods with explicit hydration mode of HYDRATE_SCALAR
+	 */
+	public function testReturnTypeOfQueryMethodsWithExplicitSingleScalarHydrationMode(EntityManagerInterface $em): void
+	{
+		$query = $em->createQuery('
+			SELECT		m
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<array>',
+			$query->getResult(AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->getSingleScalarResult()
+		);
+		assertType(
+			'iterable<array>',
+			$query->toIterable([], AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->execute(null, AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+
+		$query = $em->createQuery('
+			SELECT		m.intColumn
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<array{intColumn: int}>',
+			$query->getResult(AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array{intColumn: int}>',
+			$query->getSingleScalarResult()
+		);
+		assertType(
+			'array<array{intColumn: int}>',
+			$query->execute(null, AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array{intColumn: int}>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array<array{intColumn: int}>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array{intColumn: int}',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+		assertType(
+			'array{intColumn: int}|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SINGLE_SCALAR)
+		);
+	}
+
+	/**
+	 * Test that we properly infer the return type of Query methods with explicit hydration mode of HYDRATE_SIMPLEOBJECT
+	 *
+	 * We are never able to infer the return type here
+	 */
+	public function testReturnTypeOfQueryMethodsWithExplicitSimpleObjectHydrationMode(EntityManagerInterface $em): void
+	{
+		$query = $em->createQuery('
+			SELECT		m
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<QueryResult\Entities\Many>',
+			$query->getResult(AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'iterable<QueryResult\Entities\Many>',
+			$query->toIterable([], AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'array<QueryResult\Entities\Many>',
+			$query->execute(null, AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'array<QueryResult\Entities\Many>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'array<QueryResult\Entities\Many>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'QueryResult\Entities\Many',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'QueryResult\Entities\Many|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+
+		$query = $em->createQuery('
+			SELECT		m.intColumn, m.stringNullColumn
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array',
+			$query->getResult(AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'array',
+			$query->execute(null, AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'array',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'array',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'mixed',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+		assertType(
+			'mixed',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SIMPLEOBJECT)
+		);
+	}
+
+	/**
+	 * Test that we properly infer the return type of Query methods with explicit hydration mode of HYDRATE_SCALAR_COLUMN
+	 *
+	 * We are never able to infer the return type here
+	 */
+	public function testReturnTypeOfQueryMethodsWithExplicitScalarColumnHydrationMode(EntityManagerInterface $em): void
+	{
+		$query = $em->createQuery('
+			SELECT		m
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array',
+			$query->getResult(AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array',
+			$query->getSingleColumnResult()
+		);
+		assertType(
+			'iterable',
+			$query->toIterable([], AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array',
+			$query->execute(null, AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'mixed',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'mixed',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+
+		$query = $em->createQuery('
+			SELECT		m.intColumn
+			FROM		QueryResult\Entities\Many m
+		');
+
+		assertType(
+			'array<int>',
+			$query->getResult(AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array<int>',
+			$query->getSingleColumnResult()
+		);
+		assertType(
+			'array<int>',
+			$query->execute(null, AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array<int>',
+			$query->executeIgnoreQueryCache(null, AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'array<int>',
+			$query->executeUsingQueryCache(null, AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'int',
+			$query->getSingleResult(AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+		assertType(
+			'int|null',
+			$query->getOneOrNullResult(AbstractQuery::HYDRATE_SCALAR_COLUMN)
+		);
+	}
+
 	/**
 	 * Test that we properly infer the return type of Query methods with explicit hydration mode that is not a constant value
 	 *