From 69c22013f6d127aca89b2c18070a4e27119d94bf Mon Sep 17 00:00:00 2001
From: Abdul Malik Ikhsan <samsonasik@gmail.com>
Date: Fri, 20 Dec 2024 21:29:40 +0700
Subject: [PATCH] Rollback SimpleCallableNodeTraverser usage on
 ByRefReturnNodeVisitor (#6624)

---
 .../NodeVisitor/ByRefReturnNodeVisitor.php    | 51 ++++++++-----------
 1 file changed, 20 insertions(+), 31 deletions(-)

diff --git a/src/NodeTypeResolver/PHPStan/Scope/NodeVisitor/ByRefReturnNodeVisitor.php b/src/NodeTypeResolver/PHPStan/Scope/NodeVisitor/ByRefReturnNodeVisitor.php
index 4c0c3f459d..b3d9af4ace 100644
--- a/src/NodeTypeResolver/PHPStan/Scope/NodeVisitor/ByRefReturnNodeVisitor.php
+++ b/src/NodeTypeResolver/PHPStan/Scope/NodeVisitor/ByRefReturnNodeVisitor.php
@@ -6,16 +6,21 @@
 
 use PhpParser\Node;
 use PhpParser\Node\FunctionLike;
-use PhpParser\Node\Stmt;
+use PhpParser\Node\Stmt\Class_;
 use PhpParser\Node\Stmt\Return_;
-use PhpParser\Node\Stmt\Switch_;
+use PhpParser\NodeVisitor;
 use PhpParser\NodeVisitorAbstract;
-use Rector\Contract\PhpParser\Node\StmtsAwareInterface;
 use Rector\NodeTypeResolver\Node\AttributeKey;
 use Rector\NodeTypeResolver\PHPStan\Scope\Contract\NodeVisitor\ScopeResolverNodeVisitorInterface;
+use Rector\PhpDocParser\NodeTraverser\SimpleCallableNodeTraverser;
 
 final class ByRefReturnNodeVisitor extends NodeVisitorAbstract implements ScopeResolverNodeVisitorInterface
 {
+    public function __construct(
+        private readonly SimpleCallableNodeTraverser $simpleCallableNodeTraverser
+    ) {
+    }
+
     public function enterNode(Node $node): ?Node
     {
         if (! $node instanceof FunctionLike) {
@@ -31,37 +36,21 @@ public function enterNode(Node $node): ?Node
             return null;
         }
 
-        $this->setByRefAttribute($stmts);
-
-        return null;
-    }
-
-    /**
-     * @param Stmt[] $stmts
-     */
-    private function setByRefAttribute(array $stmts): void
-    {
-        foreach ($stmts as $stmt) {
-            if ($stmt instanceof FunctionLike) {
-                continue;
-            }
-
-            if ($stmt instanceof StmtsAwareInterface && $stmt->stmts !== null) {
-                $this->setByRefAttribute($stmt->stmts);
-                continue;
-            }
+        $this->simpleCallableNodeTraverser->traverseNodesWithCallable(
+            $stmts,
+            static function (Node $node): int|null|Node {
+                if ($node instanceof Class_ || $node instanceof FunctionLike) {
+                    return NodeVisitor::DONT_TRAVERSE_CURRENT_AND_CHILDREN;
+                }
 
-            if ($stmt instanceof Switch_) {
-                foreach ($stmt->cases as $case) {
-                    $this->setByRefAttribute($case->stmts);
+                if (! $node instanceof Return_) {
+                    return null;
                 }
 
-                continue;
-            }
+                $node->setAttribute(AttributeKey::IS_BYREF_RETURN, true);
+                return $node;
+            });
 
-            if ($stmt instanceof Return_) {
-                $stmt->setAttribute(AttributeKey::IS_BYREF_RETURN, true);
-            }
-        }
+        return null;
     }
 }