Improve QueryResultDynamicReturnTypeExtension

Author: VincentLanglet

Diff for: src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php

@@ -9,15 +9,18 @@
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\ShouldNotHappenException;
 use PHPStan\Type\ArrayType;
+use PHPStan\Type\Constant\ConstantArrayType;
 use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\Generic\GenericObjectType;
 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\VoidType;
 
 final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
@@ -109,12 +112,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;
@@ -133,13 +156,78 @@ private function getMethodReturnTypeForHydrationMode(
 		}
 	}
 
-	private function isObjectHydrationMode(Type $type): bool
+	private function getArrayHydratedReturnType(Type $queryResultType): Type
+	{
+		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 (!$type instanceof ConstantIntegerType) {
-			return false;
+		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