From 4ed0fe934f3902ab8d7a3e26bcdd41d9ea3cf6eb Mon Sep 17 00:00:00 2001
From: Daniil Gentili <daniil@daniil.it>
Date: Wed, 6 Dec 2023 14:12:19 +0100
Subject: [PATCH] Fix shaped array class string key combination

---
 .../Assignment/ArrayAssignmentAnalyzer.php     |  4 +++-
 src/Psalm/Internal/Type/TypeCombination.php    |  3 +++
 src/Psalm/Internal/Type/TypeCombiner.php       | 12 +++++++++++-
 src/Psalm/Type/Reconciler.php                  |  4 +++-
 tests/ArrayAssignmentTest.php                  | 18 ++++++++++++++++++
 5 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/ArrayAssignmentAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/ArrayAssignmentAnalyzer.php
index 576fe45a0..01150233a 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/ArrayAssignmentAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/ArrayAssignmentAnalyzer.php
@@ -351,7 +351,9 @@ final class ArrayAssignmentAnalyzer
         if (!$has_matching_objectlike_property && !$has_matching_string) {
             $properties = [];
             $classStrings = [];
-            $current_type = $current_type->setPossiblyUndefined(count($key_values) > 1);
+            $current_type = $current_type->setPossiblyUndefined(
+                $current_type->possibly_undefined || count($key_values) > 1,
+            );
             foreach ($key_values as $key_value) {
                 $properties[$key_value->value] = $current_type;
                 if ($key_value instanceof TLiteralClassString) {
diff --git a/src/Psalm/Internal/Type/TypeCombination.php b/src/Psalm/Internal/Type/TypeCombination.php
index 0706315d1..0e4cc2253 100644
--- a/src/Psalm/Internal/Type/TypeCombination.php
+++ b/src/Psalm/Internal/Type/TypeCombination.php
@@ -55,6 +55,9 @@ final class TypeCombination
     /** @var array<string|int, Union> */
     public array $objectlike_entries = [];
 
+    /** @var array<string, true> */
+    public array $objectlike_class_string_keys = [];
+
     public bool $objectlike_sealed = true;
 
     public ?Union $objectlike_key_type = null;
diff --git a/src/Psalm/Internal/Type/TypeCombiner.php b/src/Psalm/Internal/Type/TypeCombiner.php
index 3f841da4f..b48387acb 100644
--- a/src/Psalm/Internal/Type/TypeCombiner.php
+++ b/src/Psalm/Internal/Type/TypeCombiner.php
@@ -667,6 +667,7 @@ final class TypeCombiner
 
             $has_defined_keys = false;
 
+            $class_strings = $type->class_strings ?? [];
             foreach ($type->properties as $candidate_property_name => $candidate_property_type) {
                 $value_type = $combination->objectlike_entries[$candidate_property_name] ?? null;
 
@@ -705,6 +706,15 @@ final class TypeCombiner
                     );
                 }
 
+                if (isset($combination->objectlike_class_string_keys[$candidate_property_name])) {
+                    $combination->objectlike_class_string_keys[$candidate_property_name] =
+                        $combination->objectlike_class_string_keys[$candidate_property_name]
+                        && ($class_strings[$candidate_property_name] ?? false);
+                } else {
+                    $combination->objectlike_class_string_keys[$candidate_property_name] =
+                        ($class_strings[$candidate_property_name] ?? false);
+                }
+
                 unset($missing_entries[$candidate_property_name]);
             }
 
@@ -1421,7 +1431,7 @@ final class TypeCombiner
                 } else {
                     $objectlike = new TKeyedArray(
                         $combination->objectlike_entries,
-                        null,
+                        $combination->objectlike_class_string_keys,
                         $sealed || $fallback_key_type === null || $fallback_value_type === null
                             ? null
                             : [$fallback_key_type, $fallback_value_type],
diff --git a/src/Psalm/Type/Reconciler.php b/src/Psalm/Type/Reconciler.php
index f5e288639..baae33964 100644
--- a/src/Psalm/Type/Reconciler.php
+++ b/src/Psalm/Type/Reconciler.php
@@ -1132,7 +1132,9 @@ class Reconciler
 
         $base_key = implode($key_parts);
 
-        $result_type = $result_type->setPossiblyUndefined(count($array_key_offsets) > 1);
+        $result_type = $result_type->setPossiblyUndefined(
+            $result_type->possibly_undefined || count($array_key_offsets) > 1,
+        );
 
         foreach ($array_key_offsets as $array_key_offset) {
             if (isset($existing_types[$base_key]) && $array_key_offset !== false) {
diff --git a/tests/ArrayAssignmentTest.php b/tests/ArrayAssignmentTest.php
index 364b6def3..790f901b3 100644
--- a/tests/ArrayAssignmentTest.php
+++ b/tests/ArrayAssignmentTest.php
@@ -55,6 +55,24 @@ class ArrayAssignmentTest extends TestCase
                     '$resultOpt===' => 'array{a?: true, b?: true}',
                 ],
             ],
+            'assignUnionOfLiteralsClassKeys' => [
+                'code' => '<?php
+                    class a {}
+                    class b {}
+
+                    $result = [];
+
+                    foreach ([a::class, b::class] as $k) {
+                        $result[$k] = true;
+                    }
+
+                    foreach ($result as $k => $v) {
+                        $vv = new $k;
+                    }',
+                'assertions' => [
+                    '$result===' => 'array{a::class: true, b::class: true}',
+                ],
+            ],
             'genericArrayCreationWithSingleIntValue' => [
                 'code' => '<?php
                     $out = [];