From d86fcf4a85bcf8a3b38a633fdf3d8303cb6c74b7 Mon Sep 17 00:00:00 2001
From: Anner Visser <mail@annervisser.nl>
Date: Sat, 21 Jan 2023 12:01:25 +0100
Subject: [PATCH] Narrow ->value of enum case(s) to only the possible values

Using $stmt_var_type to determine if we're dealing with a subset of the
enum cases. If so, we only consider those cases for the possible values.

Fixes #8568
Fixes #8812
---
 .../Fetch/AtomicPropertyFetchAnalyzer.php     | 29 ++++++++-
 tests/EnumTest.php                            | 60 ++++++++++++++++++-
 2 files changed, 85 insertions(+), 4 deletions(-)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/AtomicPropertyFetchAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/AtomicPropertyFetchAnalyzer.php
index 6509ddde1f4..9b3352c696e 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/AtomicPropertyFetchAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/AtomicPropertyFetchAnalyzer.php
@@ -62,7 +62,9 @@
 use Psalm\Type\Atomic\TTemplateParam;
 use Psalm\Type\Union;
 
+use function array_filter;
 use function array_keys;
+use function array_map;
 use function array_search;
 use function count;
 use function in_array;
@@ -70,6 +72,8 @@
 use function is_string;
 use function strtolower;
 
+use const ARRAY_FILTER_USE_KEY;
+
 /**
  * @internal
  */
@@ -198,7 +202,7 @@ public static function analyze(
                     ]),
                 );
             } elseif ($prop_name === 'value' && $class_storage->enum_type !== null && $class_storage->enum_cases) {
-                self::handleEnumValue($statements_analyzer, $stmt, $class_storage);
+                self::handleEnumValue($statements_analyzer, $stmt, $stmt_var_type, $class_storage);
             } elseif ($prop_name === 'name') {
                 self::handleEnumName($statements_analyzer, $stmt, $lhs_type_part);
             } else {
@@ -941,11 +945,31 @@ private static function handleEnumName(
     private static function handleEnumValue(
         StatementsAnalyzer $statements_analyzer,
         PropertyFetch $stmt,
+        Union $stmt_var_type,
         ClassLikeStorage $class_storage
     ): void {
+        $relevant_enum_cases = array_filter(
+            $stmt_var_type->getAtomicTypes(),
+            static fn(Atomic $type): bool => $type instanceof TEnumCase,
+        );
+        $relevant_enum_case_names = array_map(
+            static fn(TEnumCase $enumCase): string => $enumCase->case_name,
+            $relevant_enum_cases,
+        );
+
+        $enum_cases = $class_storage->enum_cases;
+        if (!empty($relevant_enum_case_names)) {
+            // If we have a known subset of enum cases, include only those
+            $enum_cases = array_filter(
+                $enum_cases,
+                static fn(string $key) => in_array($key, $relevant_enum_case_names, true),
+                ARRAY_FILTER_USE_KEY,
+            );
+        }
+
         $case_values = [];
 
-        foreach ($class_storage->enum_cases as $enum_case) {
+        foreach ($enum_cases as $enum_case) {
             if (is_string($enum_case->value)) {
                 $case_values[] = new TLiteralString($enum_case->value);
             } elseif (is_int($enum_case->value)) {
@@ -956,7 +980,6 @@ private static function handleEnumValue(
             }
         }
 
-        // todo: this is suboptimal when we reference enum directly, e.g. Status::Open->value
         /** @psalm-suppress ArgumentTypeCoercion */
         $statements_analyzer->node_data->setType(
             $stmt,
diff --git a/tests/EnumTest.php b/tests/EnumTest.php
index f6aa64ff565..abf03aebe8f 100644
--- a/tests/EnumTest.php
+++ b/tests/EnumTest.php
@@ -94,12 +94,70 @@ enum Mask: int {
                     $z = Mask::Two->value;
                 ',
                 'assertions' => [
-                    // xxx: we should be able to do better when we reference a case explicitly, like above
+                    '$z===' => '2',
+                ],
+                'ignored_issues' => [],
+                'php_version' => '8.1',
+            ],
+            'EnumCaseValue #8568' => [
+                'code' => '<?php
+                    enum Mask: int {
+                        case One = 1 << 0;
+                        case Two = 1 << 1;
+                    }
+                    /** @return Mask */
+                    function a() {
+                        return Mask::One;
+                    }
+
+                    $z = a()->value;
+                ',
+                'assertions' => [
                     '$z===' => '1|2',
                 ],
                 'ignored_issues' => [],
                 'php_version' => '8.1',
             ],
+            'EnumUnionAsCaseValue #8568' => [
+                'code' => '<?php
+                    enum Mask: int {
+                        case One = 1 << 0;
+                        case Two = 1 << 1;
+                        case Four = 1 << 2;
+                    }
+                    /** @return Mask::One|Mask::Two */
+                    function a() {
+                        return Mask::One;
+                    }
+
+                    $z = a()->value;
+                ',
+                'assertions' => [
+                    '$z===' => '1|2',
+                ],
+                'ignored_issues' => [],
+                'php_version' => '8.1',
+            ],
+            'matchCaseOnEnumValue #8812' => [
+                'code' => '<?php
+                    enum SomeType: string
+                    {
+                        case FOO = "FOO";
+                        case BAR = "BAR";
+                    }
+
+                    function getSomething(string $moduleString): int
+                    {
+                        return match ($moduleString) {
+                            SomeType::FOO->value => 1,
+                            SomeType::BAR->value => 2,
+                        };
+                    }
+                ',
+                'assertions' => [],
+                'ignored_issues' => [],
+                'php_version' => '8.1',
+            ],
             'namePropertyFromOutside' => [
                 'code' => '<?php
                     enum Status