Support for Messenger HandleTrait return types

extension.neon

@@ -140,6 +140,13 @@ services:
 	-
 		factory: @symfony.parameterMapFactory::create()
 
+	# message map
+	symfony.messageMapFactory:
+		class: PHPStan\Symfony\MessageMapFactory
+		factory: PHPStan\Symfony\MessageMapFactory
+	-
+		factory: @symfony.messageMapFactory::create()
+
 	# ControllerTrait::get()/has() return type
 	-
 		factory: PHPStan\Type\Symfony\ServiceDynamicReturnTypeExtension(Symfony\Component\DependencyInjection\ContainerInterface)
@@ -203,6 +210,11 @@ services:
 		factory: PHPStan\Type\Symfony\EnvelopeReturnTypeExtension
 		tags: [phpstan.broker.dynamicMethodReturnTypeExtension]
 
+	# Messenger HandleTrait::handle() return type
+	-
+		class: PHPStan\Type\Symfony\MessengerHandleTraitReturnTypeExtension
+		tags: [phpstan.broker.expressionTypeResolverExtension]
+
 	# InputInterface::getArgument() return type
 	-
 		factory: PHPStan\Type\Symfony\InputInterfaceGetArgumentDynamicReturnTypeExtension

src/Symfony/MessageMap.php

@@ -0,0 +1,24 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Symfony;
+
+use PHPStan\Type\Type;
+
+final class MessageMap
+{
+
+	/** @var array<string, Type> */
+	private $messageMap;
+
+	/** @param array<string, Type> $messageMap */
+	public function __construct(array $messageMap)
+	{
+		$this->messageMap = $messageMap;
+	}
+
+	public function getTypeForClass(string $class): ?Type
+	{
+		return $this->messageMap[$class] ?? null;
+	}
+
+}

src/Symfony/MessageMapFactory.php

@@ -0,0 +1,154 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Symfony;
+
+use PHPStan\Reflection\ClassReflection;
+use PHPStan\Reflection\ReflectionProvider;
+use Symfony\Component\Messenger\Handler\MessageSubscriberInterface;
+use function class_exists;
+use function count;
+use function is_array;
+use function is_int;
+use function is_string;
+
+final class MessageMapFactory
+{
+
+	private const MESSENGER_HANDLER_TAG = 'messenger.message_handler';
+	private const DEFAULT_HANDLER_METHOD = '__invoke';
+
+	/** @var ReflectionProvider */
+	private $reflectionProvider;
+
+	/** @var ServiceMap */
+	private $serviceMap;
+
+	public function __construct(ServiceMap $symfonyServiceMap, ReflectionProvider $reflectionProvider)
+	{
+		$this->serviceMap = $symfonyServiceMap;
+		$this->reflectionProvider = $reflectionProvider;
+	}
+
+	public function create(): MessageMap
+	{
+		$returnTypesMap = [];
+
+		foreach ($this->serviceMap->getServices() as $service) {
+			$serviceClass = $service->getClass();
+
+			if ($serviceClass === null) {
+				continue;
+			}
+
+			foreach ($service->getTags() as $tag) {
+				if ($tag->getName() !== self::MESSENGER_HANDLER_TAG) {
+					continue;
+				}
+
+				if (!$this->reflectionProvider->hasClass($serviceClass)) {
+					continue;
+				}
+
+				$reflectionClass = $this->reflectionProvider->getClass($serviceClass);
+
+				/** @var array{handles?: class-string, method?: string} $tagAttributes */
+				$tagAttributes = $tag->getAttributes();
+
+				if (isset($tagAttributes['handles'])) {
+					$handles = [$tagAttributes['handles'] => ['method' => $tagAttributes['method'] ?? self::DEFAULT_HANDLER_METHOD]];
+				} else {
+					$handles = $this->guessHandledMessages($reflectionClass);
+				}
+
+				foreach ($handles as $messageClassName => $options) {
+					$methodName = $options['method'] ?? self::DEFAULT_HANDLER_METHOD;
+
+					if (!$reflectionClass->hasNativeMethod($methodName)) {
+						continue;
+					}
+
+					$methodReflection = $reflectionClass->getNativeMethod($methodName);
+
+					foreach ($methodReflection->getVariants() as $variant) {
+						$returnTypesMap[$messageClassName][] = $variant->getReturnType();
+					}
+				}
+			}
+		}
+
+		$messageMap = [];
+		foreach ($returnTypesMap as $messageClassName => $returnTypes) {
+			if (count($returnTypes) !== 1) {
+				continue;
+			}
+
+			$messageMap[$messageClassName] = $returnTypes[0];
+		}
+
+		return new MessageMap($messageMap);
+	}
+
+	/** @return iterable<string, array<string, string>> */
+	private function guessHandledMessages(ClassReflection $reflectionClass): iterable
+	{
+		if ($reflectionClass->implementsInterface(MessageSubscriberInterface::class)) {
+			$className = $reflectionClass->getName();
+
+			foreach ($className::getHandledMessages() as $index => $value) {
+				$containOptions = self::containOptions($index, $value);
+				if ($containOptions === true) {
+					yield $index => $value;
+				} elseif ($containOptions === false) {
+					yield $value => ['method' => self::DEFAULT_HANDLER_METHOD];
+				}
+			}
+
+			return;
+		}
+
+		if (!$reflectionClass->hasNativeMethod(self::DEFAULT_HANDLER_METHOD)) {
+			return;
+		}
+
+		$methodReflection = $reflectionClass->getNativeMethod(self::DEFAULT_HANDLER_METHOD);
+
+		$variants = $methodReflection->getVariants();
+		if (count($variants) !== 1) {
+			return;
+		}
+
+		$parameters = $variants[0]->getParameters();
+
+		if (count($parameters) !== 1) {
+			return;
+		}
+
+		$classNames = $parameters[0]->getType()->getObjectClassNames();
+
+		if (count($classNames) !== 1) {
+			return;
+		}
+
+		yield $classNames[0] => ['method' => self::DEFAULT_HANDLER_METHOD];
+	}
+
+	/**
+	 * @param mixed $index
+	 * @param mixed $value
+	 * @phpstan-assert-if-true =class-string $index
+	 * @phpstan-assert-if-true =array<string, mixed> $value
+	 * @phpstan-assert-if-false =int $index
+	 * @phpstan-assert-if-false =class-string $value
+	 */
+	private static function containOptions($index, $value): ?bool
+	{
+		if (is_string($index) && class_exists($index) && is_array($value)) {
+			return true;
+		} elseif (is_int($index) && is_string($value) && class_exists($value)) {
+			return false;
+		}
+
+		return null;
+	}
+
+}

src/Symfony/Service.php

@@ -20,19 +20,25 @@ final class Service implements ServiceDefinition
 	/** @var string|null */
 	private $alias;
 
+	/** @var ServiceTag[] */
+	private $tags;
+
+	/** @param ServiceTag[] $tags */
 	public function __construct(
 		string $id,
 		?string $class,
 		bool $public,
 		bool $synthetic,
-		?string $alias
+		?string $alias,
+		array $tags = []
 	)
 	{
 		$this->id = $id;
 		$this->class = $class;
 		$this->public = $public;
 		$this->synthetic = $synthetic;
 		$this->alias = $alias;
+		$this->tags = $tags;
 	}
 
 	public function getId(): string
@@ -60,4 +66,9 @@ public function getAlias(): ?string
 		return $this->alias;
 	}
 
+	public function getTags(): array
+	{
+		return $this->tags;
+	}
+
 }

src/Symfony/ServiceDefinition.php

@@ -18,4 +18,7 @@ public function isSynthetic(): bool;
 
 	public function getAlias(): ?string;
 
+	/** @return ServiceTag[] */
+	public function getTags(): array;
+
 }

src/Symfony/ServiceTag.php

@@ -0,0 +1,31 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Symfony;
+
+final class ServiceTag implements ServiceTagDefinition
+{
+
+	/** @var string */
+	private $name;
+
+	/** @var array<string, string> */
+	private $attributes;
+
+	/** @param array<string, string> $attributes */
+	public function __construct(string $name, array $attributes = [])
+	{
+		$this->name = $name;
+		$this->attributes = $attributes;
+	}
+
+	public function getName(): string
+	{
+		return $this->name;
+	}
+
+	public function getAttributes(): array
+	{
+		return $this->attributes;
+	}
+
+}

src/Symfony/ServiceTagDefinition.php

@@ -0,0 +1,13 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Symfony;
+
+interface ServiceTagDefinition
+{
+
+	public function getName(): string;
+
+	/** @return array<string, string> */
+	public function getAttributes(): array;
+
+}

src/Symfony/XmlServiceMapFactory.php

@@ -47,12 +47,22 @@ public function create(): ServiceMap
 				continue;
 			}
 
+			$serviceTags = [];
+			foreach ($def->tag as $tag) {
+				$tagAttrs = ((array) $tag->attributes())['@attributes'] ?? [];
+				$tagName = $tagAttrs['name'];
+				unset($tagAttrs['name']);
+
+				$serviceTags[] = new ServiceTag($tagName, $tagAttrs);
+			}
+
 			$service = new Service(
 				$this->cleanServiceId((string) $attrs->id),
 				isset($attrs->class) ? (string) $attrs->class : null,
 				isset($attrs->public) && (string) $attrs->public === 'true',
 				isset($attrs->synthetic) && (string) $attrs->synthetic === 'true',
-				isset($attrs->alias) ? $this->cleanServiceId((string) $attrs->alias) : null
+				isset($attrs->alias) ? $this->cleanServiceId((string) $attrs->alias) : null,
+				$serviceTags
 			);
 
 			if ($service->getAlias() !== null) {

src/Type/Symfony/MessengerHandleTraitReturnTypeExtension.php

@@ -0,0 +1,91 @@
+<?php declare(strict_types = 1);
+
+namespace PHPStan\Type\Symfony;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\MethodCall;
+use PhpParser\Node\Identifier;
+use PHPStan\Analyser\Scope;
+use PHPStan\Symfony\MessageMap;
+use PHPStan\Symfony\MessageMapFactory;
+use PHPStan\Type\ExpressionTypeResolverExtension;
+use PHPStan\Type\Type;
+use function count;
+use function is_null;
+
+final class MessengerHandleTraitReturnTypeExtension implements ExpressionTypeResolverExtension
+{
+
+	private const TRAIT_NAME = 'Symfony\Component\Messenger\HandleTrait';
+	private const TRAIT_METHOD_NAME = 'handle';
+
+	/** @var MessageMapFactory */
+	private $messageMapFactory;
+
+	/** @var MessageMap|null */
+	private $messageMap;
+
+	public function __construct(MessageMapFactory $symfonyMessageMapFactory)
+	{
+		$this->messageMapFactory = $symfonyMessageMapFactory;
+	}
+
+	public function getType(Expr $expr, Scope $scope): ?Type
+	{
+		if ($this->isSupported($expr, $scope)) {
+			$args = $expr->getArgs();
+			if (count($args) !== 1) {
+				return null;
+			}
+
+			$arg = $args[0]->value;
+			$argClassNames = $scope->getType($arg)->getObjectClassNames();
+
+			if (count($argClassNames) === 1) {
+				$messageMap = $this->getMessageMap();
+				$returnType = $messageMap->getTypeForClass($argClassNames[0]);
+
+				if (!is_null($returnType)) {
+					return $returnType;
+				}
+			}
+		}
+
+		return null;
+	}
+
+	private function getMessageMap(): MessageMap
+	{
+		if ($this->messageMap === null) {
+			$this->messageMap = $this->messageMapFactory->create();
+		}
+
+		return $this->messageMap;
+	}
+
+	/**
+	 * @phpstan-assert-if-true =MethodCall $expr
+	 */
+	private function isSupported(Expr $expr, Scope $scope): bool
+	{
+		if (!($expr instanceof MethodCall) || !($expr->name instanceof Identifier) || $expr->name->name !== self::TRAIT_METHOD_NAME) {
+			return false;
+		}
+
+		if (!$scope->isInClass()) {
+			return false;
+		}
+
+		$reflectionClass = $scope->getClassReflection()->getNativeReflection();
+
+		if (!$reflectionClass->hasMethod(self::TRAIT_METHOD_NAME)) {
+			return false;
+		}
+
+		$methodReflection = $reflectionClass->getMethod(self::TRAIT_METHOD_NAME);
+		$declaringClassReflection = $methodReflection->getBetterReflection()->getDeclaringClass();
+
+		return $declaringClassReflection->getName() === self::TRAIT_NAME;
+	}
+
+}

tests/Type/Symfony/ExtensionTest.php

@@ -14,6 +14,7 @@ class ExtensionTest extends TypeInferenceTestCase
 	/** @return mixed[] */
 	public function dataFileAsserts(): iterable
 	{
+		yield from $this->gatherAssertTypes(__DIR__ . '/data/messenger_handle_trait.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/envelope_all.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/header_bag_get.php');
 		yield from $this->gatherAssertTypes(__DIR__ . '/data/response_header_bag_get_cookies.php');

tests/Type/Symfony/container.xml

@@ -354,5 +354,23 @@
         <service id="parameterised_foo" class="%app.class%"></service>
         <service id="parameterised_bar" class="%app.class%\Bar"></service>
         <service id="synthetic" class="Synthetic" public="true" synthetic="true" />
+        <service id="regular_query_handler" class="MessengerHandleTrait\RegularQueryHandler">
+            <tag name="messenger.message_handler"/>
+        </service>
+        <service id="multi_query_handler" class="MessengerHandleTrait\MultiQueryHandler">
+            <tag name="messenger.message_handler"/>
+        </service>
+        <service id="tagged_handler" class="MessengerHandleTrait\TaggedHandler">
+            <tag name="messenger.message_handler" handles="MessengerHandleTrait\TaggedQuery" method="handle"/>
+        </service>
+        <service id="multi_handles_for_in_the_same_handler" class="MessengerHandleTrait\MultiHandlesForInTheSameHandler">
+            <tag name="messenger.message_handler"/>
+        </service>
+        <service id="multi_handlers_for_the_same_message_handler_1" class="MessengerHandleTrait\MultiHandlersForTheSameMessageHandler1">
+            <tag name="messenger.message_handler"/>
+        </service>
+        <service id="multi_handlers_for_the_same_message_handler_2" class="MessengerHandleTrait\MultiHandlersForTheSameMessageHandler2">
+            <tag name="messenger.message_handler"/>
+        </service>
     </services>
 </container>

tests/Type/Symfony/data/messenger_handle_trait.php

@@ -0,0 +1,113 @@
+<?php declare(strict_types = 1);
+
+namespace MessengerHandleTrait;
+
+use Symfony\Component\Messenger\Handler\MessageSubscriberInterface;
+use Symfony\Component\Messenger\HandleTrait;
+use function PHPStan\Testing\assertType;
+
+class RegularQuery {}
+class RegularQueryResult {}
+class RegularQueryHandler
+{
+    public function __invoke(RegularQuery $query): RegularQueryResult
+    {
+        return new RegularQueryResult();
+    }
+}
+
+class BooleanQuery {}
+class StringQuery {}
+class IntQuery {}
+class FloatQuery {}
+class MultiQueryHandler implements MessageSubscriberInterface
+{
+    public static function getHandledMessages(): iterable
+    {
+        yield BooleanQuery::class;
+        yield IntQuery::class => ['method' => 'handleInt'];
+        yield FloatQuery::class => ['method' => 'handleFloat'];
+        yield StringQuery::class => ['method' => 'handleString'];
+    }
+
+    public function __invoke(BooleanQuery $query): bool
+    {
+        return true;
+    }
+
+    public function handleInt(IntQuery $query): int
+    {
+        return 0;
+    }
+
+    public function handleFloat(FloatQuery $query): float
+    {
+        return 0.0;
+    }
+
+    public function handleString(StringQuery $query): string
+    {
+        return 'string result';
+    }
+}
+
+class TaggedQuery {}
+class TaggedResult {}
+class TaggedHandler
+{
+    public function handle(TaggedQuery $query): TaggedResult
+    {
+        return new TaggedResult();
+    }
+}
+
+class MultiHandlesForInTheSameHandlerQuery {}
+class MultiHandlesForInTheSameHandler implements MessageSubscriberInterface
+{
+    public static function getHandledMessages(): iterable
+    {
+        yield MultiHandlesForInTheSameHandlerQuery::class;
+        yield MultiHandlesForInTheSameHandlerQuery::class => ['priority' => '0'];
+    }
+
+    public function __invoke(MultiHandlesForInTheSameHandlerQuery $query): bool
+    {
+        return true;
+    }
+}
+
+class MultiHandlersForTheSameMessageQuery {}
+class MultiHandlersForTheSameMessageHandler1
+{
+    public function __invoke(MultiHandlersForTheSameMessageQuery $query): bool
+    {
+        return true;
+    }
+}
+class MultiHandlersForTheSameMessageHandler2
+{
+    public function __invoke(MultiHandlersForTheSameMessageQuery $query): bool
+    {
+        return false;
+    }
+}
+
+class HandleTraitClass {
+    use HandleTrait;
+
+    public function __invoke()
+    {
+        assertType(RegularQueryResult::class, $this->handle(new RegularQuery()));
+
+        assertType('bool', $this->handle(new BooleanQuery()));
+        assertType('int', $this->handle(new IntQuery()));
+        assertType('float', $this->handle(new FloatQuery()));
+        assertType('string', $this->handle(new StringQuery()));
+
+        assertType(TaggedResult::class, $this->handle(new TaggedQuery()));
+
+        // HandleTrait will throw exception in fact due to multiple handle methods/handlers per single query
+        assertType('mixed', $this->handle(new MultiHandlesForInTheSameHandlerQuery()));
+        assertType('mixed', $this->handle(new MultiHandlersForTheSameMessageQuery()));
+    }
+}