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) {