Infer QueryBuilderType for any method returning QB

extension.neon

@@ -169,6 +169,11 @@ services:
 			descendIntoOtherMethods: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
 			parser: @defaultAnalysisParser
 
+	-
+		class: PHPStan\Type\Doctrine\QueryBuilder\ReturnQueryBuilderExpressionTypeResolverExtension
+		tags:
+			- phpstan.broker.expressionTypeResolverExtension
+
 	-
 		class: PHPStan\Stubs\Doctrine\StubFilesExtensionLoader
 		tags:

rules.neon

@@ -35,7 +35,6 @@ services:
 		class: PHPStan\Rules\Doctrine\ORM\QueryBuilderDqlRule
 		arguments:
 			reportDynamicQueryBuilders: %doctrine.reportDynamicQueryBuilders%
-			searchOtherMethodsForQueryBuilderBeginning: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
 		tags:
 			- phpstan.rules.rule
 	-

src/Rules/Doctrine/ORM/QueryBuilderDqlRule.php

@@ -12,7 +12,6 @@
 use PHPStan\Rules\RuleErrorBuilder;
 use PHPStan\Type\Doctrine\DoctrineTypeUtils;
 use PHPStan\Type\Doctrine\ObjectMetadataResolver;
-use PHPStan\Type\Doctrine\QueryBuilder\OtherMethodQueryBuilderParser;
 use PHPStan\Type\ObjectType;
 use PHPStan\Type\TypeUtils;
 use Throwable;
@@ -32,23 +31,13 @@ class QueryBuilderDqlRule implements Rule
 	/** @var bool */
 	private $reportDynamicQueryBuilders;
 
-	/** @var OtherMethodQueryBuilderParser */
-	private $otherMethodQueryBuilderParser;
-
-	/** @var bool */
-	private $searchOtherMethodsForQueryBuilderBeginning;
-
 	public function __construct(
 		ObjectMetadataResolver $objectMetadataResolver,
-		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser,
-		bool $reportDynamicQueryBuilders,
-		bool $searchOtherMethodsForQueryBuilderBeginning
+		bool $reportDynamicQueryBuilders
 	)
 	{
 		$this->objectMetadataResolver = $objectMetadataResolver;
-		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
 		$this->reportDynamicQueryBuilders = $reportDynamicQueryBuilders;
-		$this->searchOtherMethodsForQueryBuilderBeginning = $searchOtherMethodsForQueryBuilderBeginning;
 	}
 
 	public function getNodeType(): string
@@ -69,14 +58,6 @@ public function processNode(Node $node, Scope $scope): array
 		$calledOnType = $scope->getType($node->var);
 		$queryBuilderTypes = DoctrineTypeUtils::getQueryBuilderTypes($calledOnType);
 		if (count($queryBuilderTypes) === 0) {
-
-			if ($this->searchOtherMethodsForQueryBuilderBeginning) {
-				$queryBuilderTypes = $this->otherMethodQueryBuilderParser->getQueryBuilderTypes($scope, $node);
-				if (count($queryBuilderTypes) !== 0) {
-					return [];
-				}
-			}
-
 			if (
 				$this->reportDynamicQueryBuilders
 				&& (new ObjectType('Doctrine\ORM\QueryBuilder'))->isSuperTypeOf($calledOnType)->yes()

src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php

@@ -3,8 +3,6 @@
 namespace PHPStan\Type\Doctrine\QueryBuilder;
 
 use PhpParser\Node;
-use PhpParser\Node\Expr\MethodCall;
-use PhpParser\Node\Identifier;
 use PhpParser\Node\Stmt;
 use PhpParser\Node\Stmt\Class_;
 use PhpParser\Node\Stmt\ClassMethod;
@@ -17,13 +15,12 @@
 use PHPStan\Analyser\ScopeFactory;
 use PHPStan\DependencyInjection\Container;
 use PHPStan\Parser\Parser;
-use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Reflection\MethodReflection;
 use PHPStan\Type\Generic\TemplateTypeMap;
 use PHPStan\Type\IntersectionType;
 use PHPStan\Type\Type;
 use PHPStan\Type\TypeTraverser;
 use PHPStan\Type\UnionType;
-use function count;
 use function is_array;
 
 class OtherMethodQueryBuilderParser
@@ -32,61 +29,28 @@ class OtherMethodQueryBuilderParser
 	/** @var bool */
 	private $descendIntoOtherMethods;
 
-	/** @var ReflectionProvider */
-	private $reflectionProvider;
-
 	/** @var Parser */
 	private $parser;
 
 	/** @var Container */
 	private $container;
 
-	public function __construct(bool $descendIntoOtherMethods, ReflectionProvider $reflectionProvider, Parser $parser, Container $container)
+	public function __construct(bool $descendIntoOtherMethods, Parser $parser, Container $container)
 	{
 		$this->descendIntoOtherMethods = $descendIntoOtherMethods;
-		$this->reflectionProvider = $reflectionProvider;
 		$this->parser = $parser;
 		$this->container = $container;
 	}
 
 	/**
-	 * @return QueryBuilderType[]
+	 * @return list<QueryBuilderType>
 	 */
-	public function getQueryBuilderTypes(Scope $scope, MethodCall $methodCall): array
+	public function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodReflection $methodReflection): array
 	{
-		if (!$this->descendIntoOtherMethods || !$methodCall->var instanceof MethodCall) {
-			return [];
-		}
-
-		return $this->findQueryBuilderTypesInCalledMethod($scope, $methodCall->var);
-	}
-	/**
-	 * @return QueryBuilderType[]
-	 */
-	private function findQueryBuilderTypesInCalledMethod(Scope $scope, MethodCall $methodCall): array
-	{
-		$methodCalledOnType = $scope->getType($methodCall->var);
-		if (!$methodCall->name instanceof Identifier) {
-			return [];
-		}
-
-		$methodCalledOnTypeClassNames = $methodCalledOnType->getObjectClassNames();
-
-		if (count($methodCalledOnTypeClassNames) !== 1) {
-			return [];
-		}
-
-		if (!$this->reflectionProvider->hasClass($methodCalledOnTypeClassNames[0])) {
-			return [];
-		}
-
-		$classReflection = $this->reflectionProvider->getClass($methodCalledOnTypeClassNames[0]);
-		$methodName = $methodCall->name->toString();
-		if (!$classReflection->hasNativeMethod($methodName)) {
+		if (!$this->descendIntoOtherMethods) {
 			return [];
 		}
 
-		$methodReflection = $classReflection->getNativeMethod($methodName);
 		$fileName = $methodReflection->getDeclaringClass()->getFileName();
 		if ($fileName === null) {
 			return [];

src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php

@@ -65,22 +65,17 @@ class QueryBuilderGetQueryDynamicReturnTypeExtension implements DynamicMethodRet
 	/** @var DescriptorRegistry */
 	private $descriptorRegistry;
 
-	/** @var OtherMethodQueryBuilderParser */
-	private $otherMethodQueryBuilderParser;
-
 	public function __construct(
 		ObjectMetadataResolver $objectMetadataResolver,
 		ArgumentsProcessor $argumentsProcessor,
 		?string $queryBuilderClass,
-		DescriptorRegistry $descriptorRegistry,
-		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser
+		DescriptorRegistry $descriptorRegistry
 	)
 	{
 		$this->objectMetadataResolver = $objectMetadataResolver;
 		$this->argumentsProcessor = $argumentsProcessor;
 		$this->queryBuilderClass = $queryBuilderClass;
 		$this->descriptorRegistry = $descriptorRegistry;
-		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
 	}
 
 	public function getClass(): string
@@ -107,10 +102,7 @@ public function getTypeFromMethodCall(
 		)->getReturnType();
 		$queryBuilderTypes = DoctrineTypeUtils::getQueryBuilderTypes($calledOnType);
 		if (count($queryBuilderTypes) === 0) {
-			$queryBuilderTypes = $this->otherMethodQueryBuilderParser->getQueryBuilderTypes($scope, $methodCall);
-			if (count($queryBuilderTypes) === 0) {
-				return $defaultReturnType;
-			}
+			return $defaultReturnType;
 		}
 
 		$objectManager = $this->objectMetadataResolver->getObjectManager();

src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php

@@ -26,16 +26,11 @@ class QueryBuilderMethodDynamicReturnTypeExtension implements DynamicMethodRetur
 	/** @var string|null */
 	private $queryBuilderClass;
 
-	/** @var OtherMethodQueryBuilderParser */
-	private $otherMethodQueryBuilderParser;
-
 	public function __construct(
-		?string $queryBuilderClass,
-		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser
+		?string $queryBuilderClass
 	)
 	{
 		$this->queryBuilderClass = $queryBuilderClass;
-		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
 	}
 
 	public function getClass(): string
@@ -74,10 +69,7 @@ public function getTypeFromMethodCall(
 
 		$queryBuilderTypes = DoctrineTypeUtils::getQueryBuilderTypes($calledOnType);
 		if (count($queryBuilderTypes) === 0) {
-			$queryBuilderTypes = $this->otherMethodQueryBuilderParser->getQueryBuilderTypes($scope, $methodCall);
-			if (count($queryBuilderTypes) === 0) {
-				return $calledOnType;
-			}
+			return $calledOnType;
 		}
 
 		if (count($queryBuilderTypes) > self::MAX_COMBINATIONS) {

src/Type/Doctrine/QueryBuilder/ReturnQueryBuilderExpressionTypeResolverExtension.php

@@ -0,0 +1,103 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Doctrine\QueryBuilder;
+
+use Doctrine\ORM\EntityManagerInterface;
+use Doctrine\ORM\EntityRepository;
+use Doctrine\ORM\QueryBuilder;
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\CallLike;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Expr\StaticCall;
+use PhpParser\Node\Identifier;
+use PhpParser\Node\Name;
+use PHPStan\Analyser\Scope;
+use PHPStan\Reflection\MethodReflection;
+use PHPStan\Reflection\ParametersAcceptorSelector;
+use PHPStan\Type\ExpressionTypeResolverExtension;
+use PHPStan\Type\ObjectType;
+use PHPStan\Type\Type;
+use PHPStan\Type\TypeCombinator;
+use function count;
+
+class ReturnQueryBuilderExpressionTypeResolverExtension implements ExpressionTypeResolverExtension
+{
+
+	/** @var OtherMethodQueryBuilderParser */
+	private $otherMethodQueryBuilderParser;
+
+	public function __construct(
+		OtherMethodQueryBuilderParser $otherMethodQueryBuilderParser
+	)
+	{
+		$this->otherMethodQueryBuilderParser = $otherMethodQueryBuilderParser;
+	}
+
+	public function getType(Expr $expr, Scope $scope): ?Type
+	{
+		if (!$expr instanceof MethodCall && !$expr instanceof StaticCall) {
+			return null;
+		}
+
+		if ($expr->isFirstClassCallable()) {
+			return null;
+		}
+
+		$methodReflection = $this->getMethodReflection($expr, $scope);
+
+		if ($methodReflection === null) {
+			return null;
+		}
+
+		$returnType = ParametersAcceptorSelector::selectFromArgs($scope, $expr->getArgs(), $methodReflection->getVariants())->getReturnType();
+
+		$returnsQueryBuilder = (new ObjectType(QueryBuilder::class))->isSuperTypeOf($returnType)->yes();
+
+		if (!$returnsQueryBuilder) {
+			return null;
+		}
+
+		$queryBuilderTypes = $this->otherMethodQueryBuilderParser->findQueryBuilderTypesInCalledMethod($scope, $methodReflection);
+		if (count($queryBuilderTypes) === 0) {
+			return null;
+		}
+
+		return TypeCombinator::union(...$queryBuilderTypes);
+	}
+
+	/**
+	 * @param StaticCall|MethodCall $call
+	 */
+	private function getMethodReflection(CallLike $call, Scope $scope): ?MethodReflection
+	{
+		if (!$call->name instanceof Identifier) {
+			return null;
+		}
+
+		if ($call instanceof MethodCall) {
+			$callerType = $scope->getType($call->var);
+		} else {
+			if (!$call->class instanceof Name) {
+				return null;
+			}
+			$callerType = $scope->resolveTypeByName($call->class);
+		}
+
+		$methodName = $call->name->name;
+
+		foreach ($callerType->getObjectClassReflections() as $callerClassReflection) {
+			if ($callerClassReflection->is(QueryBuilder::class)) {
+				return null; // covered by QueryBuilderMethodDynamicReturnTypeExtension
+			}
+			if ($callerClassReflection->is(EntityRepository::class) && $methodName === 'createQueryBuilder') {
+				return null; // covered by EntityRepositoryCreateQueryBuilderDynamicReturnTypeExtension
+			}
+			if ($callerClassReflection->is(EntityManagerInterface::class) && $methodName === 'createQueryBuilder') {
+				return null; // no need to dive there
+			}
+		}
+
+		return $scope->getMethodReflection($callerType, $methodName);
+	}
+
+}

tests/Rules/Doctrine/ORM/QueryBuilderDqlRuleSlowTest.php

@@ -5,7 +5,6 @@
 use PHPStan\Rules\Rule;
 use PHPStan\Testing\RuleTestCase;
 use PHPStan\Type\Doctrine\ObjectMetadataResolver;
-use PHPStan\Type\Doctrine\QueryBuilder\OtherMethodQueryBuilderParser;
 
 /**
  * @extends RuleTestCase<QueryBuilderDqlRule>
@@ -17,8 +16,6 @@ protected function getRule(): Rule
 	{
 		return new QueryBuilderDqlRule(
 			new ObjectMetadataResolver(__DIR__ . '/entity-manager.php', __DIR__ . '/../../../../tmp'),
-			self::getContainer()->getByType(OtherMethodQueryBuilderParser::class),
-			true,
 			true
 		);
 	}

tests/Rules/Doctrine/ORM/QueryBuilderDqlRuleTest.php

@@ -5,7 +5,6 @@
 use PHPStan\Rules\Rule;
 use PHPStan\Testing\RuleTestCase;
 use PHPStan\Type\Doctrine\ObjectMetadataResolver;
-use PHPStan\Type\Doctrine\QueryBuilder\OtherMethodQueryBuilderParser;
 
 /**
  * @extends RuleTestCase<QueryBuilderDqlRule>
@@ -17,8 +16,6 @@ protected function getRule(): Rule
 	{
 		return new QueryBuilderDqlRule(
 			new ObjectMetadataResolver(__DIR__ . '/entity-manager.php', __DIR__ . '/../../../../tmp'),
-			self::getContainer()->getByType(OtherMethodQueryBuilderParser::class),
-			true,
 			true
 		);
 	}

tests/Rules/Doctrine/ORM/data/query-builder-dql.php

@@ -84,14 +84,14 @@ public function selectArray(): void
 			->getQuery();
 	}
 
-	public function analyseQueryBuilderUnknownBeginning(): void
+	public function analyseQueryBuilderOtherMethodBeginning(): void
 	{
 		$this->createQb()->getQuery();
 	}
 
 	private function createQb(): \Doctrine\ORM\QueryBuilder
 	{
-		return $this->entityManager->createQueryBuilder();
+		return $this->entityManager->createQueryBuilder()->select('e')->from(MyEntity::class, 'e');
 	}
 
 	public function analyseQueryBuilderDynamicArgs(string $entity): void