Result infering: support indexBy and lists

composer.json

@@ -7,7 +7,7 @@
 	],
 	"require": {
 		"php": "^7.2 || ^8.0",
-		"phpstan/phpstan": "^1.8.11"
+		"phpstan/phpstan": "^1.9.0"
 	},
 	"conflict": {
 		"doctrine/collections": "<1.0",

src/Type/Doctrine/CreateQueryDynamicReturnTypeExtension.php

@@ -24,7 +24,7 @@
 use PHPStan\Type\UnionType;
 
 /**
- * Infers TResult in Query<TResult> on EntityManagerInterface::createQuery()
+ * Infers TResult and TKey in Query<TKey,TResult> on EntityManagerInterface::createQuery()
  */
 final class CreateQueryDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
 {
@@ -63,7 +63,7 @@ public function getTypeFromMethodCall(
 		if (!isset($args[$queryStringArgIndex])) {
 			return new GenericObjectType(
 				Query::class,
-				[new MixedType()]
+				[new MixedType(), new MixedType()]
 			);
 		}
 
@@ -78,7 +78,7 @@ public function getTypeFromMethodCall(
 
 				$em = $this->objectMetadataResolver->getObjectManager();
 				if (!$em instanceof EntityManagerInterface) {
-					return new QueryType($queryString, null);
+					return new QueryType($queryString, null, null);
 				}
 
 				$typeBuilder = new QueryResultTypeBuilder();
@@ -87,14 +87,14 @@ public function getTypeFromMethodCall(
 					$query = $em->createQuery($queryString);
 					QueryResultTypeWalker::walk($query, $typeBuilder, $this->descriptorRegistry);
 				} catch (ORMException | DBALException | NewDBALException | CommonException $e) {
-					return new QueryType($queryString, null);
+					return new QueryType($queryString, null, null);
 				}
 
-				return new QueryType($queryString, $typeBuilder->getResultType());
+				return new QueryType($queryString, $typeBuilder->getIndexType(), $typeBuilder->getResultType());
 			}
 			return new GenericObjectType(
 				Query::class,
-				[new MixedType()]
+				[new MixedType(), new MixedType()]
 			);
 		});
 	}

src/Type/Doctrine/Query/QueryResultDynamicReturnTypeExtension.php

@@ -8,8 +8,10 @@
 use PHPStan\Reflection\MethodReflection;
 use PHPStan\Reflection\ParametersAcceptorSelector;
 use PHPStan\ShouldNotHappenException;
+use PHPStan\Type\Accessory\AccessoryArrayListType;
 use PHPStan\Type\ArrayType;
 use PHPStan\Type\Constant\ConstantIntegerType;
+use PHPStan\Type\Doctrine\Type\ListIndexMarkerType;
 use PHPStan\Type\DynamicMethodReturnTypeExtension;
 use PHPStan\Type\Generic\GenericObjectType;
 use PHPStan\Type\IntegerType;
@@ -18,6 +20,7 @@
 use PHPStan\Type\NullType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
+use PHPStan\Type\TypeUtils;
 use PHPStan\Type\VoidType;
 
 final class QueryResultDynamicReturnTypeExtension implements DynamicMethodReturnTypeExtension
@@ -70,15 +73,17 @@ public function getTypeFromMethodCall(
 
 		$queryType = $scope->getType($methodCall->var);
 		$queryResultType = $this->getQueryResultType($queryType);
+		$queryIndexType = $this->getQueryIndexType($queryType);
 
 		return $this->getMethodReturnTypeForHydrationMode(
 			$methodReflection,
 			$hydrationMode,
-			$queryResultType
+			$queryResultType,
+			$queryIndexType
 		);
 	}
 
-	private function getQueryResultType(Type $queryType): Type
+	private function getQueryIndexType(Type $queryType): Type
 	{
 		if (!$queryType instanceof GenericObjectType) {
 			return new MixedType();
@@ -89,10 +94,22 @@ private function getQueryResultType(Type $queryType): Type
 		return $types[0] ?? new MixedType();
 	}
 
+	private function getQueryResultType(Type $queryType): Type
+	{
+		if (!$queryType instanceof GenericObjectType) {
+			return new MixedType();
+		}
+
+		$types = $queryType->getTypes();
+
+		return $types[1] ?? new MixedType();
+	}
+
 	private function getMethodReturnTypeForHydrationMode(
 		MethodReflection $methodReflection,
 		Type $hydrationMode,
-		Type $queryResultType
+		Type $queryResultType,
+		Type $queryIndexType
 	): Type
 	{
 		$isVoidType = (new VoidType())->isSuperTypeOf($queryResultType);
@@ -126,10 +143,17 @@ private function getMethodReturnTypeForHydrationMode(
 					$queryResultType
 				);
 			default:
-				return new ArrayType(
-					new MixedType(),
-					$queryResultType
-				);
+				return $queryIndexType instanceof ListIndexMarkerType
+					? AccessoryArrayListType::intersectWith(
+						TypeCombinator::intersect(
+							new ArrayType(new IntegerType(), $queryResultType),
+							...TypeUtils::getAccessoryTypes($queryResultType)
+						)
+					)
+					: new ArrayType(
+						new MixedType(),
+						$queryResultType
+					);
 		}
 	}
 

src/Type/Doctrine/Query/QueryResultTypeBuilder.php

@@ -5,6 +5,7 @@
 use PHPStan\Type\Constant\ConstantArrayTypeBuilder;
 use PHPStan\Type\Constant\ConstantIntegerType;
 use PHPStan\Type\Constant\ConstantStringType;
+use PHPStan\Type\Doctrine\Type\ListIndexMarkerType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeCombinator;
 use PHPStan\Type\VoidType;
@@ -62,6 +63,14 @@ final class QueryResultTypeBuilder
 	 */
 	private $newObjects = [];
 
+	/** @var Type */
+	private $indexedBy;
+
+	public function __construct()
+	{
+		$this->indexedBy = new ListIndexMarkerType();
+	}
+
 	public function setSelectQuery(): void
 	{
 		$this->selectQuery = true;
@@ -230,4 +239,18 @@ private function resolveOffsetType($alias): Type
 		return new ConstantStringType($alias);
 	}
 
+	public function setIndexedBy(Type $type): void
+	{
+		$this->indexedBy = $type;
+	}
+
+	public function getIndexType(): Type
+	{
+		if (!$this->selectQuery) {
+			return new VoidType();
+		}
+
+		return $this->indexedBy;
+	}
+
 }

src/Type/Doctrine/Query/QueryResultTypeWalker.php

@@ -78,7 +78,7 @@ class QueryResultTypeWalker extends SqlWalker
 	 */
 	private $newObjectCounter = 0;
 
-	/** @var Query<mixed> */
+	/** @var Query<mixed, mixed> */
 	private $query;
 
 	/** @var EntityManagerInterface */
@@ -107,7 +107,7 @@ class QueryResultTypeWalker extends SqlWalker
 	private $hasGroupByClause;
 
 	/**
-	 * @param Query<mixed> $query
+	 * @param Query<mixed, mixed> $query
 	 */
 	public static function walk(Query $query, QueryResultTypeBuilder $typeBuilder, DescriptorRegistry $descriptorRegistry): void
 	{
@@ -122,7 +122,7 @@ public static function walk(Query $query, QueryResultTypeBuilder $typeBuilder, D
 	/**
 	 * {@inheritDoc}
 	 *
-	 * @param Query<mixed> $query
+	 * @param Query<mixed, mixed> $query
 	 * @param ParserResult $parserResult
 	 * @param array<QueryComponent> $queryComponents
 	 */
@@ -312,6 +312,10 @@ public function walkFromClause($fromClause)
 	 */
 	public function walkIdentificationVariableDeclaration($identificationVariableDecl)
 	{
+		if ($identificationVariableDecl->indexBy !== null) {
+			$identificationVariableDecl->indexBy->dispatch($this);
+		}
+
 		foreach ($identificationVariableDecl->joins as $join) {
 			assert($join instanceof AST\Node);
 
@@ -326,6 +330,8 @@ public function walkIdentificationVariableDeclaration($identificationVariableDec
 	 */
 	public function walkIndexBy($indexBy): void
 	{
+		$type = $this->unmarshalType($indexBy->singleValuedPathExpression->dispatch($this));
+		$this->typeBuilder->setIndexedBy($type);
 	}
 
 	/**

src/Type/Doctrine/Query/QueryType.php

@@ -14,10 +14,16 @@ class QueryType extends GenericObjectType
 	/** @var string */
 	private $dql;
 
-	public function __construct(string $dql, ?Type $resultType = null)
+	public function __construct(
+		string $dql,
+		?Type $indexType,
+		?Type $resultType
+	)
 	{
-		$resultType = $resultType ?? new MixedType();
-		parent::__construct('Doctrine\ORM\Query', [$resultType]);
+		parent::__construct('Doctrine\ORM\Query', [
+			$indexType ?? new MixedType(),
+			$resultType ?? new MixedType(),
+		]);
 		$this->dql = $dql;
 	}
 

src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php

@@ -153,7 +153,7 @@ private function getQueryType(string $dql): Type
 	{
 		$em = $this->objectMetadataResolver->getObjectManager();
 		if (!$em instanceof EntityManagerInterface) {
-			return new QueryType($dql, null);
+			return new QueryType($dql, null, null);
 		}
 
 		$typeBuilder = new QueryResultTypeBuilder();
@@ -162,10 +162,10 @@ private function getQueryType(string $dql): Type
 			$query = $em->createQuery($dql);
 			QueryResultTypeWalker::walk($query, $typeBuilder, $this->descriptorRegistry);
 		} catch (ORMException | DBALException | CommonException $e) {
-			return new QueryType($dql, null);
+			return new QueryType($dql, null, null);
 		}
 
-		return new QueryType($dql, $typeBuilder->getResultType());
+		return new QueryType($dql, $typeBuilder->getIndexType(), $typeBuilder->getResultType());
 	}
 
 }

src/Type/Doctrine/Type/ListIndexMarkerType.php

@@ -0,0 +1,10 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\Type;
+
+use PHPStan\Type\IntegerType;
+
+class ListIndexMarkerType extends IntegerType
+{
+
+}

stubs/ORM/AbstractQuery.stub

@@ -5,6 +5,7 @@ namespace Doctrine\ORM;
 use Doctrine\Common\Collections\ArrayCollection;
 
 /**
+ * @template TKey The type of column used in indexBy
  * @template TResult The type of results returned by this query in HYDRATE_OBJECT mode
  */
 abstract class AbstractQuery

stubs/ORM/Query.stub

@@ -3,9 +3,10 @@
 namespace Doctrine\ORM;
 
 /**
+ * @template TKey The type of column used in indexBy
  * @template TResult The type of results returned by this query in HYDRATE_OBJECT mode
  *
- * @extends AbstractQuery<TResult>
+ * @extends AbstractQuery<TKey, TResult>
  */
 final class Query extends AbstractQuery
 {

stubs/ORM/QueryBuilder.stub

@@ -15,7 +15,7 @@ class QueryBuilder
 	}
 
 	/**
-	 * @return Query<mixed>
+	 * @return Query<mixed, mixed>
 	 */
 	public function getQuery()
 	{

stubs/bleedingEdge/ORM/QueryBuilder.stub

@@ -8,7 +8,7 @@ class QueryBuilder
 {
 
 	/**
-	 * @return Query<mixed>
+	 * @return Query<mixed, mixed>
 	 */
 	public function getQuery()
 	{
@@ -143,7 +143,7 @@ class QueryBuilder
 	{
 
 	}
-	
+
 	/**
 	 * @param literal-string|object|array<mixed> $predicates
 	 * @return $this

tests/Type/Doctrine/data/QueryResult/config.neon

@@ -3,3 +3,5 @@ includes:
 parameters:
 	doctrine:
 		objectManagerLoader: entity-manager.php
+	featureToggles:
+		listType: true

tests/Type/Doctrine/data/QueryResult/createQuery.php

@@ -15,28 +15,28 @@ public function testQueryTypeParametersAreInfered(EntityManagerInterface $em): v
 			FROM		QueryResult\Entities\Many m
 		');
 
-		assertType('Doctrine\ORM\Query<QueryResult\Entities\Many>', $query);
+		assertType('Doctrine\ORM\Query<int, QueryResult\Entities\Many>', $query);
 
 		$query = $em->createQuery('
 			SELECT		m.intColumn, m.stringNullColumn
 			FROM		QueryResult\Entities\Many m
 		');
 
-		assertType('Doctrine\ORM\Query<array{intColumn: int, stringNullColumn: string|null}>', $query);
+		assertType('Doctrine\ORM\Query<int, array{intColumn: int, stringNullColumn: string|null}>', $query);
 	}
 
 	public function testQueryResultTypeIsMixedWhenDQLIsNotKnown(EntityManagerInterface $em, string $dql): void
 	{
 		$query = $em->createQuery($dql);
 
-		assertType('Doctrine\ORM\Query<mixed>', $query);
+		assertType('Doctrine\ORM\Query<mixed, mixed>', $query);
 	}
 
 	public function testQueryResultTypeIsMixedWhenDQLIsInvalid(EntityManagerInterface $em, string $dql): void
 	{
 		$query = $em->createQuery('invalid');
 
-		assertType('Doctrine\ORM\Query<mixed>', $query);
+		assertType('Doctrine\ORM\Query<mixed, mixed>', $query);
 	}
 
 }