Add native return type declaration support to NarrowReturnUnionToCollectionRector (#440)

* add array-filter to ArrayMapOnCollectionToArrayRector

* Add native return type declaration support to NarrowReturnUnionToCollectionRector

Author: TomasVotruba

rules-tests/TypedCollections/Rector/ClassMethod/NarrowReturnUnionToCollectionRector/Fixture/array_and_collection.php.inc

@@ -11,7 +11,6 @@ final class ArrayAndCollection
      */
     public function someMethod()
     {
-
     }
 }
 
@@ -30,7 +29,6 @@ final class ArrayAndCollection
      */
     public function someMethod()
     {
-
     }
 }
 

rules-tests/TypedCollections/Rector/ClassMethod/NarrowReturnUnionToCollectionRector/Fixture/nullable_union_collection.php.inc

@@ -0,0 +1,37 @@
+<?php
+
+namespace Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowReturnUnionToCollectionRector\Fixture;
+
+use Doctrine\Common\Collections\Collection;
+use Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowParamUnionToCollectionRector\Source\AnyEntity;
+
+final class NullableUnionCollection
+{
+    /**
+     * @return Collection<int, AnyEntity>
+     */
+    public function someMethod(): Collection|null
+    {
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowReturnUnionToCollectionRector\Fixture;
+
+use Doctrine\Common\Collections\Collection;
+use Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowParamUnionToCollectionRector\Source\AnyEntity;
+
+final class NullableUnionCollection
+{
+    /**
+     * @return Collection<int, AnyEntity>
+     */
+    public function someMethod(): Collection
+    {
+    }
+}
+
+?>

rules-tests/TypedCollections/Rector/ClassMethod/NarrowReturnUnionToCollectionRector/Fixture/return_non_nullable_collection.php.inc

@@ -0,0 +1,37 @@
+<?php
+
+namespace Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowReturnUnionToCollectionRector\Fixture;
+
+use Doctrine\Common\Collections\Collection;
+use Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowParamUnionToCollectionRector\Source\AnyEntity;
+
+final class ReturnNonNullableCollection
+{
+    /**
+     * @return Collection<int, AnyEntity>
+     */
+    public function someMethod(): ?Collection
+    {
+    }
+}
+
+?>
+-----
+<?php
+
+namespace Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowReturnUnionToCollectionRector\Fixture;
+
+use Doctrine\Common\Collections\Collection;
+use Rector\Doctrine\Tests\TypedCollections\Rector\ClassMethod\NarrowParamUnionToCollectionRector\Source\AnyEntity;
+
+final class ReturnNonNullableCollection
+{
+    /**
+     * @return Collection<int, AnyEntity>
+     */
+    public function someMethod(): Collection
+    {
+    }
+}
+
+?>

rules/TypedCollections/Rector/ClassMethod/NarrowReturnUnionToCollectionRector.php

@@ -5,8 +5,11 @@
 namespace Rector\Doctrine\TypedCollections\Rector\ClassMethod;
 
 use PhpParser\Node;
+use PhpParser\Node\NullableType;
 use PhpParser\Node\Stmt\ClassMethod;
+use PhpParser\Node\UnionType;
 use PHPStan\PhpDocParser\Ast\PhpDoc\ReturnTagValueNode;
+use Rector\BetterPhpDocParser\PhpDocInfo\PhpDocInfo;
 use Rector\BetterPhpDocParser\PhpDocInfo\PhpDocInfoFactory;
 use Rector\Comments\NodeDocBlock\DocBlockUpdater;
 use Rector\Doctrine\Enum\DoctrineClass;
@@ -74,26 +77,97 @@ public function getNodeTypes(): array
      */
     public function refactor(Node $node): ?ClassMethod
     {
-        $classMethodPhpDocInfo = $this->phpDocInfoFactory->createFromNodeOrEmpty($node);
-        $returnTagValueNode = $classMethodPhpDocInfo->getReturnTagValue();
+        $hasChanged = false;
+        if ($this->refactorReturnDocblockTag($node)) {
+            $hasChanged = true;
+        }
 
-        if (! $returnTagValueNode instanceof ReturnTagValueNode) {
+        if ($this->refactorNativeReturn($node)) {
+            $hasChanged = true;
+        }
+
+        if (! $hasChanged) {
             return null;
         }
 
-        if ($node->returnType !== null && $this->isName($node->returnType, DoctrineClass::COLLECTION)) {
-            $hasNativeCollectionType = true;
-        } else {
-            $hasNativeCollectionType = false;
+        return $node;
+    }
+
+    private function refactorReturnDocblockTag(ClassMethod $classMethod): bool
+    {
+        $phpDocInfo = $this->phpDocInfoFactory->createFromNode($classMethod);
+        if (! $phpDocInfo instanceof PhpDocInfo) {
+            return false;
         }
 
+        $returnTagValueNode = $phpDocInfo->getReturnTagValue();
+        if (! $returnTagValueNode instanceof ReturnTagValueNode) {
+            return false;
+        }
+
+        $hasNativeCollectionType = $this->hasNativeCollectionType($classMethod);
+
         $hasChanged = $this->unionCollectionTagValueNodeNarrower->narrow($returnTagValueNode, $hasNativeCollectionType);
         if ($hasChanged === false) {
-            return null;
+            return false;
         }
 
-        $this->docBlockUpdater->updateRefactoredNodeWithPhpDocInfo($node);
+        $this->docBlockUpdater->updateRefactoredNodeWithPhpDocInfo($classMethod);
 
-        return $node;
+        return true;
+    }
+
+    private function hasNativeCollectionType(ClassMethod $classMethod): bool
+    {
+        if (! $classMethod->returnType instanceof Node) {
+            return false;
+        }
+
+        return $this->isName($classMethod->returnType, DoctrineClass::COLLECTION);
+    }
+
+    private function refactorNativeReturn(ClassMethod $classMethod): bool
+    {
+        if (! $classMethod->returnType instanceof Node) {
+            return false;
+        }
+
+        if ($classMethod->returnType instanceof NullableType && $this->isName(
+            $classMethod->returnType->type,
+            DoctrineClass::COLLECTION
+        )) {
+            // unwrap nullable type
+            $classMethod->returnType = $classMethod->returnType->type;
+            return true;
+        }
+
+        if (! $classMethod->returnType instanceof UnionType) {
+            return false;
+        }
+
+        $unionType = $classMethod->returnType;
+        if (! $this->hasNativeReturnCollectionType($unionType)) {
+            return false;
+        }
+
+        // remove null from union type
+        foreach ($unionType->types as $key => $unionedType) {
+            if ($this->isName($unionedType, 'null')) {
+                unset($unionType->types[$key]);
+            }
+        }
+
+        return true;
+    }
+
+    private function hasNativeReturnCollectionType(UnionType $unionType): bool
+    {
+        foreach ($unionType->types as $unionedType) {
+            if ($this->isName($unionedType, DoctrineClass::COLLECTION)) {
+                return true;
+            }
+        }
+
+        return false;
     }
 }