Allow class constant as offset key (#5943)

* allow class-constant as an array offset when the type match

* remove object with __toString as a possible offset as it's not valid

src/Psalm/Internal/Analyzer/Statements/Expression/Fetch/ArrayFetchAnalyzer.php

@@ -14,6 +14,7 @@ use Psalm\Internal\Type\Comparator\UnionTypeComparator;
 use Psalm\Internal\Type\TemplateInferredTypeReplacer;
 use Psalm\Internal\Type\TemplateResult;
 use Psalm\Internal\Type\TypeCombiner;
+use Psalm\Internal\Type\TypeExpander;
 use Psalm\Issue\EmptyArrayAccess;
 use Psalm\Issue\InvalidArrayAccess;
 use Psalm\Issue\InvalidArrayAssignment;
@@ -42,6 +43,7 @@ use Psalm\Node\VirtualIdentifier;
 use Psalm\Node\VirtualName;
 use Psalm\Plugin\EventHandler\Event\AddRemoveTaintsEvent;
 use Psalm\Type;
+use Psalm\Type\Atomic;
 use Psalm\Type\Atomic\TArray;
 use Psalm\Type\Atomic\TArrayKey;
 use Psalm\Type\Atomic\TClassStringMap;
@@ -450,7 +452,6 @@ class ArrayFetchAnalyzer
         $non_array_types = [];
 
         $has_valid_expected_offset = false;
-        $has_valid_absolute_offset = false;
         $expected_offset_types = [];
 
         $key_values = [];
@@ -519,28 +520,11 @@ class ArrayFetchAnalyzer
         }
 
         if ($array_type->isArray()) {
-            foreach ($offset_type->getAtomicTypes() as $atomic_offset_type) {
+            $has_valid_absolute_offset = self::checkArrayOffsetType(
-                if ($atomic_offset_type instanceof Type\Atomic\TFalse &&
+                $offset_type,
-                    $offset_type->ignore_falsable_issues === true
+                $offset_type->getAtomicTypes(),
-                ) {
+                $codebase
-                    //do nothing
+            );
-                } elseif ($atomic_offset_type instanceof Type\Atomic\TNull &&
-                    $offset_type->ignore_nullable_issues === true
-                ) {
-                    //do nothing
-                } elseif ($atomic_offset_type instanceof Type\Atomic\TString ||
-                    $atomic_offset_type instanceof Type\Atomic\TInt ||
-                    $atomic_offset_type instanceof Type\Atomic\TArrayKey ||
-                    $atomic_offset_type instanceof Type\Atomic\TMixed ||
-                    $atomic_offset_type instanceof Type\Atomic\TTemplateParam ||
-                    (
-                        $atomic_offset_type instanceof Type\Atomic\TObjectWithProperties
-                        && isset($atomic_offset_type->methods['__toString'])
-                    )
-                ) {
-                    $has_valid_absolute_offset = true;
-                }
-            }
 
             if ($has_valid_absolute_offset === false) {
                 //we didn't find a single type that could be valid
@@ -2069,4 +2053,74 @@ class ArrayFetchAnalyzer
             }
         }
     }
+
+    /**
+     * @param Atomic[] $offset_types
+     */
+    private static function checkArrayOffsetType(
+        Type\Union $offset_type,
+        array $offset_types,
+        \Psalm\Codebase $codebase
+    ): bool {
+        $has_valid_absolute_offset = false;
+        foreach ($offset_types as $atomic_offset_type) {
+            if ($atomic_offset_type instanceof Type\Atomic\TClassConstant) {
+                $expanded = TypeExpander::expandAtomic(
+                    $codebase,
+                    $atomic_offset_type,
+                    $atomic_offset_type->fq_classlike_name,
+                    $atomic_offset_type->fq_classlike_name,
+                    null,
+                    true,
+                    true
+                );
+
+                if ($expanded instanceof Atomic) {
+                    $has_valid_absolute_offset = self::checkArrayOffsetType(
+                        $offset_type,
+                        [$expanded],
+                        $codebase
+                    );
+                } else {
+                    $has_valid_absolute_offset = self::checkArrayOffsetType(
+                        $offset_type,
+                        $expanded,
+                        $codebase
+                    );
+                }
+
+                if ($has_valid_absolute_offset) {
+                    break;
+                }
+            }
+
+            if ($atomic_offset_type instanceof Type\Atomic\TFalse &&
+                $offset_type->ignore_falsable_issues === true
+            ) {
+                //do nothing
+            } elseif ($atomic_offset_type instanceof Type\Atomic\TNull &&
+                $offset_type->ignore_nullable_issues === true
+            ) {
+                //do nothing
+            } elseif ($atomic_offset_type instanceof Type\Atomic\TString ||
+                $atomic_offset_type instanceof Type\Atomic\TInt ||
+                $atomic_offset_type instanceof Type\Atomic\TArrayKey ||
+                $atomic_offset_type instanceof Type\Atomic\TMixed
+            ) {
+                $has_valid_absolute_offset = true;
+                break;
+            } elseif ($atomic_offset_type instanceof Type\Atomic\TTemplateParam) {
+                $has_valid_absolute_offset = self::checkArrayOffsetType(
+                    $offset_type,
+                    $atomic_offset_type->as->getAtomicTypes(),
+                    $codebase
+                );
+
+                if ($has_valid_absolute_offset) {
+                    break;
+                }
+            }
+        }
+        return $has_valid_absolute_offset;
+    }
 }

src/Psalm/Internal/Type/Comparator/UnionTypeComparator.php

@@ -3,7 +3,9 @@
 namespace Psalm\Internal\Type\Comparator;
 
 use Psalm\Codebase;
+use Psalm\Internal\Type\TypeExpander;
 use Psalm\Type;
+use Psalm\Type\Atomic;
 use Psalm\Type\Atomic\TArrayKey;
 use Psalm\Type\Atomic\TFalse;
 use Psalm\Type\Atomic\TMixed;
@@ -67,6 +69,26 @@ class UnionTypeComparator
                 continue;
             }
 
+            if ($input_type_part instanceof Type\Atomic\TClassConstant) {
+                $expanded = TypeExpander::expandAtomic(
+                    $codebase,
+                    $input_type_part,
+                    $input_type_part->fq_classlike_name,
+                    $input_type_part->fq_classlike_name,
+                    null,
+                    true,
+                    true
+                );
+
+                if ($expanded instanceof Atomic) {
+                    $input_atomic_types[] = $expanded;
+                } else {
+                    $input_atomic_types = array_merge($expanded, $input_atomic_types);
+                }
+
+                continue;
+            }
+
             $type_match_found = false;
             $scalar_type_match_found = false;
             $all_to_string_cast = true;

tests/ArrayAssignmentTest.php

@@ -1579,6 +1579,37 @@ class ArrayAssignmentTest extends TestCase
                         return (array)$a;
                     }',
             ],
+            'ClassConstantAsKey'  => [
+                '<?php
+                    /**
+                     * @property Foo::C_* $aprop
+                     */
+                    class Foo {
+                        public const C_ONE = 1;
+                        public const C_TWO = 2;
+
+                        public function __get(string $prop) {
+                            if ($prop === "aprop")
+                                return self::C_ONE;
+                            throw new \RuntimeException("Unsupported property: $prop");
+                        }
+
+                        /** @return array<Foo::C_*, string> */
+                        public static function getNames(): array {
+                            return [
+                                self::C_ONE => "One",
+                                self::C_TWO => "Two",
+                            ];
+                        }
+
+                        public function getThisName(): string {
+                            $names = self::getNames();
+                            $aprop = $this->aprop;
+
+                            return $names[$aprop];
+                        }
+                    }',
+            ],
         ];
     }
 
@@ -1991,6 +2022,22 @@ class ArrayAssignmentTest extends TestCase
                     }',
                 'error_message' => 'InvalidArrayOffset'
             ],
+            'TemplateAsKey' => [
+                '<?php
+
+                class Foo {
+
+                    /**
+                     * @psalm-template T of array
+                     * @param T $offset
+                     * @param array<array, string> $weird_array
+                     */
+                    public function getThisName($offset, $weird_array): string {
+                        return $weird_array[$offset];
+                    }
+                }',
+                'error_message' => 'InvalidArrayOffset'
+            ],
         ];
     }
 }