Fix #2951 - use reflection to identify trait when there’s more than one per file

src/Psalm/Internal/Codebase/ClassLikes.php

@@ -711,27 +711,19 @@ class ClassLikes
 
         $file_statements = $this->statements_provider->getStatementsForFile($storage->location->file_path);
 
-        foreach ($file_statements as $file_statement) {
-            if ($file_statement instanceof PhpParser\Node\Stmt\Trait_
-                && $file_statement->name
-                && strtolower($file_statement->name->name) === $fq_trait_name_lc
-            ) {
-                $this->trait_nodes[$fq_trait_name_lc] = $file_statement;
-                return $file_statement;
-            }
+        $trait_finder = new \Psalm\Internal\Visitor\TraitFinder($fq_trait_name);
 
-            if ($file_statement instanceof PhpParser\Node\Stmt\Namespace_) {
-                $namespace_stub = $file_statement->name ? $file_statement->name . '\\' : '';
+        $traverser = new \PhpParser\NodeTraverser();
+        $traverser->addVisitor(
+            $trait_finder
+        );
 
-                foreach ($file_statement->stmts as $namespace_stmt) {
-                    if ($namespace_stmt instanceof PhpParser\Node\Stmt\Trait_
-                        && strtolower($namespace_stub . $namespace_stmt->name) === $fq_trait_name_lc
-                    ) {
-                        $this->trait_nodes[$fq_trait_name_lc] = $namespace_stmt;
-                        return $namespace_stmt;
-                    }
-                }
-            }
+        $traverser->traverse($file_statements);
+
+        $trait_node = $trait_finder->getNode();
+
+        if ($trait_node) {
+            return $trait_node;
         }
 
         throw new \UnexpectedValueException('Could not locate trait statement');

src/Psalm/Internal/Visitor/SimpleNameResolver.php

@@ -116,6 +116,8 @@ class SimpleNameResolver extends NodeVisitorAbstract
             }
         } elseif ($node instanceof Expr\ConstFetch) {
             $node->name = $this->resolveName($node->name, Stmt\Use_::TYPE_CONSTANT);
+        } elseif ($node instanceof Stmt\Trait_) {
+            $this->resolveTrait($node);
         } elseif ($node instanceof Stmt\TraitUse) {
             foreach ($node->traits as &$trait) {
                 $trait = $this->resolveClassName($trait);
@@ -218,4 +220,16 @@ class SimpleNameResolver extends NodeVisitorAbstract
     {
         return $this->resolveName($name, Stmt\Use_::TYPE_NORMAL);
     }
+
+    /**
+     * @return void
+     */
+    protected function resolveTrait(Stmt\Trait_ $node)
+    {
+        $resolvedName = Name::concat($this->nameContext->getNamespace(), (string) $node->name);
+
+        if (null !== $resolvedName) {
+            $node->setAttribute('resolvedName', $resolvedName->toString());
+        }
+    }
 }

src/Psalm/Internal/Visitor/TraitFinder.php

@@ -0,0 +1,70 @@
+<?php
+namespace Psalm\Internal\Visitor;
+
+use PhpParser;
+use function count;
+
+/**
+ * Given a list of file diffs, this scans an AST to find the sections it can replace, and parses
+ * just those methods.
+ */
+class TraitFinder extends PhpParser\NodeVisitorAbstract implements PhpParser\NodeVisitor
+{
+    /** @var list<PhpParser\Node\Stmt\Trait_> */
+    private $matching_trait_nodes = [];
+
+    private $fq_trait_name;
+
+    public function __construct(string $fq_trait_name)
+    {
+        $this->fq_trait_name = $fq_trait_name;
+    }
+
+    /**
+     * @param  PhpParser\Node $node
+     * @param  bool $traverseChildren
+     *
+     * @return null|int|PhpParser\Node
+     */
+    public function enterNode(PhpParser\Node $node, &$traverseChildren = true)
+    {
+        if ($node instanceof PhpParser\Node\Stmt\Trait_
+            && $node->getAttribute('resolvedName') === $this->fq_trait_name
+        ) {
+            $this->matching_trait_nodes[] = $node;
+        }
+
+        if ($node instanceof PhpParser\Node\Stmt\ClassLike
+            || $node instanceof PhpParser\Node\FunctionLike
+        ) {
+            return PhpParser\NodeTraverser::DONT_TRAVERSE_CHILDREN;
+        }
+
+        return null;
+    }
+
+    public function getNode() : ?PhpParser\Node\Stmt\Trait_
+    {
+        if (!count($this->matching_trait_nodes)) {
+            return null;
+        }
+
+        if (count($this->matching_trait_nodes) === 1 || !\trait_exists($this->fq_trait_name)) {
+            return $this->matching_trait_nodes[0];
+        }
+
+        try {
+            $reflection_trait = new \ReflectionClass($this->fq_trait_name);
+        } catch (\Throwable $t) {
+            return null;
+        }
+
+        foreach ($this->matching_trait_nodes as $node) {
+            if ($node->getLine() === $reflection_trait->getStartLine()) {
+                return $node;
+            }
+        }
+
+        return null;
+    }
+}