From ac903c5634c1d0d4d8058815db717e7c36a95ecd Mon Sep 17 00:00:00 2001
From: Bruce Weirdan <weirdan@gmail.com>
Date: Sun, 28 Nov 2021 05:49:39 +0200
Subject: [PATCH] Include enum cases in const wildcards

Refs vimeo/psalm#6471
---
 src/Psalm/Internal/Type/TypeExpander.php | 11 ++++-------
 tests/EnumTest.php                       | 17 ++++++++++++++++-
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/src/Psalm/Internal/Type/TypeExpander.php b/src/Psalm/Internal/Type/TypeExpander.php
index 678ebd32f6c..c6d73f1b318 100644
--- a/src/Psalm/Internal/Type/TypeExpander.php
+++ b/src/Psalm/Internal/Type/TypeExpander.php
@@ -215,9 +215,10 @@ public static function expandAtomic(
                 $class_storage = $codebase->classlike_storage_provider->get($return_type->fq_classlike_name);
 
                 if (strpos($return_type->const_name, '*') !== false) {
-                    $class_storage = $codebase->classlike_storage_provider->get($return_type->fq_classlike_name);
-
-                    $matching_constants = \array_keys($class_storage->constants);
+                    $matching_constants = \array_merge(
+                        \array_keys($class_storage->constants),
+                        \array_keys($class_storage->enum_cases)
+                    );
 
                     $const_name_part = \substr($return_type->const_name, 0, -1);
 
@@ -231,10 +232,6 @@ function ($constant_name) use ($const_name_part): bool {
                         );
                     }
                 } else {
-                    if ($class_storage->is_enum) {
-                        return new Type\Atomic\TEnumCase($return_type->fq_classlike_name, $return_type->const_name);
-                    }
-
                     $matching_constants = [$return_type->const_name];
                 }
 
diff --git a/tests/EnumTest.php b/tests/EnumTest.php
index 20ab198f545..7f5b10dd8b2 100644
--- a/tests/EnumTest.php
+++ b/tests/EnumTest.php
@@ -154,7 +154,7 @@ public function get(): string
                 [],
                 '8.1'
             ],
-            'SKIPPED-wildcardEnum' => [
+            'wildcardEnumAsParam' => [
                 '<?php
                     enum A {
                         case C_1;
@@ -174,6 +174,21 @@ public static function foo(self $i) : void {}
                 [],
                 '8.1',
             ],
+            'wildcardEnumAsReturn' => [
+                '<?php
+                    enum E {
+                        const A = 1;
+                        case B;
+                    }
+
+                    /** @return E::* */
+                    function f(): mixed {
+                        return E::B;
+                    }',
+                'assertions' => [],
+                [],
+                '8.1',
+            ],
             'wildcardConstantsOnEnum' => [
                 '<?php
                     enum A {