Determine queryBuilderType if getQuery is called immediately

extension.neon

@@ -110,9 +110,7 @@ services:
 	-
 		class: PHPStan\Type\Doctrine\QueryBuilder\QueryBuilderMethodDynamicReturnTypeExtension
 		arguments:
-			parser: @defaultAnalysisParser
 			queryBuilderClass: %doctrine.queryBuilderClass%
-			descendIntoOtherMethods: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
 		tags:
 			- phpstan.broker.dynamicMethodReturnTypeExtension
 
@@ -147,6 +145,11 @@ services:
 		class: PHPStan\Rules\Doctrine\ORM\PropertiesExtension
 		tags:
 			- phpstan.properties.readWriteExtension
+	-
+		class: PHPStan\Type\Doctrine\QueryBuilder\OtherMethodQueryBuilderParser
+		arguments:
+			descendIntoOtherMethods: %doctrine.searchOtherMethodsForQueryBuilderBeginning%
+			parser: @defaultAnalysisParser
 
 	doctrineQueryBuilderArgumentsProcessor:
 		class: PHPStan\Type\Doctrine\ArgumentsProcessor

phpstan.neon

@@ -26,4 +26,4 @@ parameters:
 			path: src/Type/Doctrine/QueryBuilder/Expr/ExpressionBuilderDynamicReturnTypeExtension.php
 		-
 			message: '~^Variable property access on PhpParser\\Node\\Stmt\\Declare_\|PhpParser\\Node\\Stmt\\Namespace_\.$~'
-			path: src/Type/Doctrine/QueryBuilder/QueryBuilderMethodDynamicReturnTypeExtension.php
+			path: src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php

src/Type/Doctrine/QueryBuilder/OtherMethodQueryBuilderParser.php

@@ -0,0 +1,188 @@
+<?php declare(strict_types = 1);
+
+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;
+use PhpParser\Node\Stmt\Declare_;
+use PhpParser\Node\Stmt\Namespace_;
+use PhpParser\Node\Stmt\Return_;
+use PHPStan\Analyser\NodeScopeResolver;
+use PHPStan\Analyser\Scope;
+use PHPStan\Analyser\ScopeContext;
+use PHPStan\Analyser\ScopeFactory;
+use PHPStan\DependencyInjection\Container;
+use PHPStan\Parser\Parser;
+use PHPStan\Reflection\ReflectionProvider;
+use PHPStan\Type\Generic\TemplateTypeMap;
+use PHPStan\Type\TypeWithClassName;
+use function is_array;
+
+class OtherMethodQueryBuilderParser
+{
+
+	/** @var bool */
+	private $descendIntoOtherMethods;
+
+	/** @var ReflectionProvider */
+	private $broker;
+
+	/** @var Parser */
+	private $parser;
+
+	/** @var Container */
+	private $container;
+
+	public function __construct(bool $descendIntoOtherMethods, ReflectionProvider $broker, Parser $parser, Container $container)
+	{
+		$this->descendIntoOtherMethods = $descendIntoOtherMethods;
+		$this->broker = $broker;
+		$this->parser = $parser;
+		$this->container = $container;
+	}
+
+	/**
+	 * @return QueryBuilderType[]
+	 */
+	public function getQueryBuilderTypes(Scope $scope, MethodCall $methodCall): 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 [];
+		}
+
+		if (!$methodCalledOnType instanceof TypeWithClassName) {
+			return [];
+		}
+
+		if (!$this->broker->hasClass($methodCalledOnType->getClassName())) {
+			return [];
+		}
+
+		$classReflection = $this->broker->getClass($methodCalledOnType->getClassName());
+		$methodName = $methodCall->name->toString();
+		if (!$classReflection->hasNativeMethod($methodName)) {
+			return [];
+		}
+
+		$methodReflection = $classReflection->getNativeMethod($methodName);
+		$fileName = $methodReflection->getDeclaringClass()->getFileName();
+		if ($fileName === null) {
+			return [];
+		}
+
+		$nodes = $this->parser->parseFile($fileName);
+		$classNode = $this->findClassNode($methodReflection->getDeclaringClass()->getName(), $nodes);
+		if ($classNode === null) {
+			return [];
+		}
+
+		$methodNode = $this->findMethodNode($methodReflection->getName(), $classNode->stmts);
+		if ($methodNode === null || $methodNode->stmts === null) {
+			return [];
+		}
+
+		/** @var NodeScopeResolver $nodeScopeResolver */
+		$nodeScopeResolver = $this->container->getByType(NodeScopeResolver::class);
+
+		/** @var ScopeFactory $scopeFactory */
+		$scopeFactory = $this->container->getByType(ScopeFactory::class);
+
+		$methodScope = $scopeFactory->create(
+			ScopeContext::create($fileName),
+			$scope->isDeclareStrictTypes(),
+			[],
+			$methodReflection,
+			$scope->getNamespace()
+		)->enterClass($methodReflection->getDeclaringClass())->enterClassMethod($methodNode, TemplateTypeMap::createEmpty(), [], null, null, null, false, false, false);
+
+		$queryBuilderTypes = [];
+
+		$nodeScopeResolver->processNodes($methodNode->stmts, $methodScope, static function (Node $node, Scope $scope) use (&$queryBuilderTypes): void {
+			if (!$node instanceof Return_ || $node->expr === null) {
+				return;
+			}
+
+			$exprType = $scope->getType($node->expr);
+			if (!$exprType instanceof QueryBuilderType) {
+				return;
+			}
+
+			$queryBuilderTypes[] = $exprType;
+		});
+
+		return $queryBuilderTypes;
+	}
+
+	/**
+	 * @param Node[] $nodes
+	 */
+	private function findClassNode(string $className, array $nodes): ?Class_
+	{
+		foreach ($nodes as $node) {
+			if (
+				$node instanceof Class_
+				&& $node->namespacedName !== null
+				&& $node->namespacedName->toString() === $className
+			) {
+				return $node;
+			}
+
+			if (
+				!$node instanceof Namespace_
+				&& !$node instanceof Declare_
+			) {
+				continue;
+			}
+			$subNodeNames = $node->getSubNodeNames();
+			foreach ($subNodeNames as $subNodeName) {
+				$subNode = $node->{$subNodeName};
+				if (!is_array($subNode)) {
+					$subNode = [$subNode];
+				}
+
+				$result = $this->findClassNode($className, $subNode);
+				if ($result === null) {
+					continue;
+				}
+
+				return $result;
+			}
+		}
+
+		return null;
+	}
+
+	/**
+	 * @param Stmt[] $classStatements
+	 */
+	private function findMethodNode(string $methodName, array $classStatements): ?ClassMethod
+	{
+		foreach ($classStatements as $statement) {
+			if (
+				$statement instanceof ClassMethod
+				&& $statement->name->toString() === $methodName
+			) {
+				return $statement;
+			}
+		}
+
+		return null;
+	}
+
+}

src/Type/Doctrine/QueryBuilder/QueryBuilderGetQueryDynamicReturnTypeExtension.php

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