From 1750db32fff8299942c2fbb98b59247f0d79183f Mon Sep 17 00:00:00 2001
From: Simon Podlipsky <simon@podlipsky.net>
Date: Fri, 19 Jul 2024 12:14:14 +0200
Subject: [PATCH] feat: support unions in arrays

---
 .../DocBlockDriver/DocBlockTypeResolver.php   | 39 +++++++++++++++----
 src/Type/Lexer.php                            |  4 ++
 src/Type/Parser.php                           |  6 +++
 .../Collection/MapTypedAsGenericClass.php     | 10 +++++
 tests/Metadata/Driver/DocBlockDriverTest.php  | 21 ++++++++++
 5 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/src/Metadata/Driver/DocBlockDriver/DocBlockTypeResolver.php b/src/Metadata/Driver/DocBlockDriver/DocBlockTypeResolver.php
index b1ec7ad3c..46e829f51 100644
--- a/src/Metadata/Driver/DocBlockDriver/DocBlockTypeResolver.php
+++ b/src/Metadata/Driver/DocBlockDriver/DocBlockTypeResolver.php
@@ -21,6 +21,8 @@
 use PHPStan\PhpDocParser\Parser\TokenIterator;
 use PHPStan\PhpDocParser\Parser\TypeParser;
 
+use function sprintf;
+
 /**
  * @internal
  */
@@ -116,19 +118,42 @@ private function getDocBlocTypeHint($reflector): ?string
 
         // Generic array syntax: array<Product> | array<\Foo\Bar\Product> | array<int,Product>
         if ($type instanceof GenericTypeNode) {
-            if ($this->isSimpleType($type->type, 'array')) {
-                $resolvedTypes = array_map(fn (TypeNode $node) => $this->resolveTypeFromTypeNode($node, $reflector), $type->genericTypes);
+            $isSimpleTypeArray = $this->isSimpleType($type->type, 'array');
+            $isSimpleTypeList = $this->isSimpleType($type->type, 'list');
+            if (!$isSimpleTypeArray && !$isSimpleTypeList) {
+                throw new \InvalidArgumentException(sprintf("Can't use non-array generic type %s for collection in %s:%s", (string) $type->type, $reflector->getDeclaringClass()->getName(), $reflector->getName()));
+            }
+
+            if ($isSimpleTypeList) {
+                $keyType = 'int';
+                $valuesIndex = 0;
+            } else {
+                if (1 === count($type->genericTypes)) {
+                    $keyType = null;
+                    $valuesIndex = 0;
+                } else {
+                    $keyType = $this->resolveTypeFromTypeNode($type->genericTypes[0], $reflector);
+                    $valuesIndex = 1;
+                }
+            }
 
-                return 'array<' . implode(',', $resolvedTypes) . '>';
+            if ($type->genericTypes[$valuesIndex] instanceof UnionTypeNode) {
+                $valueTypes = array_map(
+                    fn (TypeNode $node) => $this->resolveTypeFromTypeNode($node, $reflector),
+                    $type->genericTypes[$valuesIndex]->types,
+                );
+            } else {
+                $valueType = $this->resolveTypeFromTypeNode($type->genericTypes[$valuesIndex], $reflector);
+                $valueTypes = [$valueType];
             }
 
-            if ($this->isSimpleType($type->type, 'list')) {
-                $resolvedTypes = array_map(fn (TypeNode $node) => $this->resolveTypeFromTypeNode($node, $reflector), $type->genericTypes);
+            $valueType = implode('|', $valueTypes);
 
-                return 'array<int, ' . implode(',', $resolvedTypes) . '>';
+            if (null === $keyType) {
+                return sprintf('array<%s>', $valueType);
             }
 
-            throw new \InvalidArgumentException(sprintf("Can't use non-array generic type %s for collection in %s:%s", (string) $type->type, $reflector->getDeclaringClass()->getName(), $reflector->getName()));
+            return sprintf('array<%s, %s>', $keyType, $valueType);
         }
 
         // Primitives and class names: Collection | \Foo\Bar\Product | string
diff --git a/src/Type/Lexer.php b/src/Type/Lexer.php
index 80dd048db..ee237ed61 100644
--- a/src/Type/Lexer.php
+++ b/src/Type/Lexer.php
@@ -23,6 +23,7 @@ final class Lexer extends AbstractLexer
     public const T_TYPE_END = 8;
     public const T_IDENTIFIER = 9;
     public const T_NULL = 10;
+    public const T_UNION_SEPARATOR = 11;
 
     public function parse(string $type)
     {
@@ -103,6 +104,9 @@ protected function getType(&$value)
             case '[' === $value:
                 return self::T_ARRAY_START;
 
+            case '|' === $value:
+                return self::T_UNION_SEPARATOR;
+
             // Default
             default:
                 // Do nothing
diff --git a/src/Type/Parser.php b/src/Type/Parser.php
index 560345c13..7b3803343 100644
--- a/src/Type/Parser.php
+++ b/src/Type/Parser.php
@@ -95,6 +95,12 @@ private function visitCompoundType(): array
                     break;
                 }
 
+                if ($this->lexer->isNextToken(Lexer::T_UNION_SEPARATOR)) {
+                    $this->match(Lexer::T_UNION_SEPARATOR);
+
+                    continue;
+                }
+
                 $this->match(Lexer::T_COMMA);
             }
         }
diff --git a/tests/Fixtures/DocBlockType/Collection/MapTypedAsGenericClass.php b/tests/Fixtures/DocBlockType/Collection/MapTypedAsGenericClass.php
index 5f3b3f7e0..fa0330d82 100644
--- a/tests/Fixtures/DocBlockType/Collection/MapTypedAsGenericClass.php
+++ b/tests/Fixtures/DocBlockType/Collection/MapTypedAsGenericClass.php
@@ -10,4 +10,14 @@ class MapTypedAsGenericClass
      * @var array<int, Product>
      */
     public array $productIds;
+
+    /**
+     * @var array<Product|Vehicle>
+     */
+    public array $productOrVehicleIds;
+
+    /**
+     * @var array<int, Product|Vehicle>
+     */
+    public array $productOrVehicleIdsWithKey;
 }
diff --git a/tests/Metadata/Driver/DocBlockDriverTest.php b/tests/Metadata/Driver/DocBlockDriverTest.php
index b7b5f8e79..edd60c937 100644
--- a/tests/Metadata/Driver/DocBlockDriverTest.php
+++ b/tests/Metadata/Driver/DocBlockDriverTest.php
@@ -131,6 +131,27 @@ public function testInferDocBlockMapFromGenericLikeClass()
             ['name' => 'array', 'params' => [['name' => 'int', 'params' => []], ['name' => Product::class, 'params' => []]]],
             $m->propertyMetadata['productIds']->type,
         );
+        self::assertEquals(
+            [
+                'name' => 'array',
+                'params' => [
+                    ['name' => Product::class, 'params' => []],
+                    ['name' => Vehicle::class, 'params' => []],
+                ]
+            ],
+            $m->propertyMetadata['productOrVehicleIds']->type,
+        );
+        self::assertEquals(
+            [
+                'name' => 'array',
+                'params' => [
+                    ['name' => 'int', 'params' => []],
+                    ['name' => Product::class, 'params' => []],
+                    ['name' => Vehicle::class, 'params' => []],
+                ]
+            ],
+            $m->propertyMetadata['productOrVehicleIdsWithKey']->type,
+        );
     }
 
     public function testInferDocBlockCollectionOfClassesIgnoringNullTypeHint()