From dbc5f0043839dd4f8d2c854df1ea4d4e437249ff Mon Sep 17 00:00:00 2001
From: Tomas Votruba <tomas.vot@gmail.com>
Date: Sat, 9 Jul 2022 16:34:15 +0200
Subject: [PATCH] Add assigned new to variable in ReturnTypeFromReturnNewRector

---
 .../StrictReturnNewAnalyzer.php               | 137 ++++++++++++++++++
 .../ReturnTypeFromReturnNewRector.php         |  97 +++++++++----
 2 files changed, 202 insertions(+), 32 deletions(-)
 create mode 100644 rules/TypeDeclaration/NodeAnalyzer/ReturnTypeAnalyzer/StrictReturnNewAnalyzer.php

diff --git a/rules/TypeDeclaration/NodeAnalyzer/ReturnTypeAnalyzer/StrictReturnNewAnalyzer.php b/rules/TypeDeclaration/NodeAnalyzer/ReturnTypeAnalyzer/StrictReturnNewAnalyzer.php
new file mode 100644
index 00000000000..c1af50ee512
--- /dev/null
+++ b/rules/TypeDeclaration/NodeAnalyzer/ReturnTypeAnalyzer/StrictReturnNewAnalyzer.php
@@ -0,0 +1,137 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Rector\TypeDeclaration\NodeAnalyzer\ReturnTypeAnalyzer;
+
+use PhpParser\Node\Expr;
+use PhpParser\Node\Expr\Assign;
+use PhpParser\Node\Expr\Closure;
+use PhpParser\Node\Expr\New_;
+use PhpParser\Node\Expr\Variable;
+use PhpParser\Node\Expr\Yield_;
+use PhpParser\Node\Stmt\ClassMethod;
+use PhpParser\Node\Stmt\Expression;
+use PhpParser\Node\Stmt\Function_;
+use PhpParser\Node\Stmt\Return_;
+use Rector\Core\PhpParser\Node\BetterNodeFinder;
+use Rector\NodeNameResolver\NodeNameResolver;
+
+final class StrictReturnNewAnalyzer
+{
+    public function __construct(
+        private readonly BetterNodeFinder $betterNodeFinder,
+        private readonly NodeNameResolver $nodeNameResolver,
+    ) {
+    }
+
+    public function matchAlwaysReturnVariableNew(ClassMethod|Closure|Function_ $functionLike): ?string
+    {
+        if ($functionLike->stmts === null) {
+            return null;
+        }
+
+        if ($this->betterNodeFinder->hasInstancesOfInFunctionLikeScoped($functionLike, [Yield_::class])) {
+            return null;
+        }
+
+        /** @var Return_[] $returns */
+        $returns = $this->betterNodeFinder->findInstancesOfInFunctionLikeScoped($functionLike, Return_::class);
+        if ($returns === []) {
+            return null;
+        }
+
+        // is one statement depth 3?
+        if (! $this->areExclusiveExprReturns($returns)) {
+            return null;
+        }
+
+        // has root return?
+        if (! $this->hasClassMethodRootReturn($functionLike)) {
+            return null;
+        }
+
+        if (count($returns) !== 1) {
+            return null;
+        }
+
+        // exact one return of variable
+        $onlyReturn = $returns[0];
+        if (! $onlyReturn->expr instanceof Variable) {
+            return null;
+        }
+
+        $createdVariablesToTypes = $this->resolveCreatedVariablesToTypes($functionLike);
+
+        $returnedVariableName = $this->nodeNameResolver->getName($onlyReturn->expr);
+
+        return $createdVariablesToTypes[$returnedVariableName] ?? null;
+    }
+
+    /**
+     * @param Return_[] $returns
+     */
+    private function areExclusiveExprReturns(array $returns): bool
+    {
+        foreach ($returns as $return) {
+            if (! $return->expr instanceof Expr) {
+                return false;
+            }
+        }
+
+        return true;
+    }
+
+    private function hasClassMethodRootReturn(ClassMethod|Function_|Closure $functionLike): bool
+    {
+        foreach ((array) $functionLike->stmts as $stmt) {
+            if ($stmt instanceof Return_) {
+                return true;
+            }
+        }
+
+        return false;
+    }
+
+    /**
+     * @return array<string, string>
+     */
+    private function resolveCreatedVariablesToTypes(ClassMethod|Function_|Closure $functionLike): array
+    {
+        $createdVariablesToTypes = [];
+
+        // what new is assigned to it?
+        foreach ((array) $functionLike->stmts as $stmt) {
+            if (! $stmt instanceof Expression) {
+                continue;
+            }
+
+            if (! $stmt->expr instanceof Assign) {
+                continue;
+            }
+
+            $assign = $stmt->expr;
+            if (! $assign->expr instanceof New_) {
+                continue;
+            }
+
+            if (! $assign->var instanceof Variable) {
+                continue;
+            }
+
+            $variableName = $this->nodeNameResolver->getName($assign->var);
+            if (! is_string($variableName)) {
+                continue;
+            }
+
+            $className = $this->nodeNameResolver->getName($assign->expr->class);
+            if (! is_string($className)) {
+                continue;
+            }
+
+            $createdVariablesToTypes[$variableName] = $className;
+        }
+
+        return $createdVariablesToTypes;
+    }
+}
diff --git a/rules/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector.php b/rules/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector.php
index bf07b9b405a..89a5c82088c 100644
--- a/rules/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector.php
+++ b/rules/TypeDeclaration/Rector/ClassMethod/ReturnTypeFromReturnNewRector.php
@@ -9,6 +9,7 @@
 use PhpParser\Node\Expr\Closure;
 use PhpParser\Node\Expr\New_;
 use PhpParser\Node\Name;
+use PhpParser\Node\Name\FullyQualified;
 use PhpParser\Node\Stmt\ClassMethod;
 use PhpParser\Node\Stmt\Function_;
 use PhpParser\Node\Stmt\Return_;
@@ -24,6 +25,7 @@
 use Rector\NodeTypeResolver\PHPStan\Type\TypeFactory;
 use Rector\PHPStanStaticTypeMapper\Enum\TypeKind;
 use Rector\StaticTypeMapper\ValueObject\Type\SelfStaticType;
+use Rector\TypeDeclaration\NodeAnalyzer\ReturnTypeAnalyzer\StrictReturnNewAnalyzer;
 use Rector\VersionBonding\Contract\MinPhpVersionInterface;
 use Symplify\RuleDocGenerator\ValueObject\CodeSample\CodeSample;
 use Symplify\RuleDocGenerator\ValueObject\RuleDefinition;
@@ -36,7 +38,8 @@ final class ReturnTypeFromReturnNewRector extends AbstractRector implements MinP
     public function __construct(
         private readonly TypeFactory $typeFactory,
         private readonly ReflectionProvider $reflectionProvider,
-        private readonly ReflectionResolver $reflectionResolver
+        private readonly ReflectionResolver $reflectionResolver,
+        private readonly StrictReturnNewAnalyzer $strictReturnNewAnalyzer
     ) {
     }
 
@@ -84,41 +87,16 @@ public function refactor(Node $node): ?Node
             return null;
         }
 
-        if ($node instanceof ArrowFunction) {
-            $returns = [new Return_($node->expr)];
-        } else {
-            /** @var Return_[] $returns */
-            $returns = $this->betterNodeFinder->findInstancesOfInFunctionLikeScoped($node, Return_::class);
-        }
-
-        if ($returns === []) {
-            return null;
-        }
+        if (! $node instanceof ArrowFunction) {
+            $returnedNewClassName = $this->strictReturnNewAnalyzer->matchAlwaysReturnVariableNew($node);
+            if (is_string($returnedNewClassName)) {
+                $node->returnType = new FullyQualified($returnedNewClassName);
 
-        $newTypes = [];
-        foreach ($returns as $return) {
-            if (! $return->expr instanceof New_) {
-                return null;
-            }
-
-            $new = $return->expr;
-            if (! $new->class instanceof Name) {
-                return null;
+                return $node;
             }
-
-            $newTypes[] = $this->createObjectTypeFromNew($new);
-        }
-
-        $returnType = $this->typeFactory->createMixedPassedOrUnionType($newTypes);
-
-        $returnTypeNode = $this->staticTypeMapper->mapPHPStanTypeToPhpParserNode($returnType, TypeKind::RETURN);
-        if (! $returnTypeNode instanceof \PhpParser\Node) {
-            return null;
         }
 
-        $node->returnType = $returnTypeNode;
-
-        return $node;
+        return $this->refactorDirectReturnNew($node);
     }
 
     public function provideMinPhpVersion(): int
@@ -149,4 +127,59 @@ private function createObjectTypeFromNew(New_ $new): ObjectType|StaticType
         $classReflection = $this->reflectionProvider->getClass($className);
         return new ObjectType($className, null, $classReflection);
     }
+
+    private function refactorDirectReturnNew(
+        ClassMethod|Function_|ArrowFunction $node
+    ): null|ArrowFunction|Function_|ClassMethod
+    {
+        if ($node instanceof ArrowFunction) {
+            $returns = [new Return_($node->expr)];
+        } else {
+            /** @var Return_[] $returns */
+            $returns = $this->betterNodeFinder->findInstancesOfInFunctionLikeScoped($node, Return_::class);
+        }
+
+        if ($returns === []) {
+            return null;
+        }
+
+        $newTypes = $this->resolveReturnNewType($returns);
+        if ($newTypes === null) {
+            return null;
+        }
+
+        $returnType = $this->typeFactory->createMixedPassedOrUnionType($newTypes);
+
+        $returnTypeNode = $this->staticTypeMapper->mapPHPStanTypeToPhpParserNode($returnType, TypeKind::RETURN);
+        if (! $returnTypeNode instanceof \PhpParser\Node) {
+            return null;
+        }
+
+        $node->returnType = $returnTypeNode;
+
+        return $node;
+    }
+
+    /**
+     * @param Return_[] $returns
+     * @return \PHPStan\Type\Type[]|null
+     */
+    private function resolveReturnNewType(array $returns): ?array
+    {
+        $newTypes = [];
+        foreach ($returns as $return) {
+            if (! $return->expr instanceof New_) {
+                return null;
+            }
+
+            $new = $return->expr;
+            if (! $new->class instanceof Name) {
+                return null;
+            }
+
+            $newTypes[] = $this->createObjectTypeFromNew($new);
+        }
+
+        return $newTypes;
+    }
 }