Fix #1830 - infer key type after array_key_exists check

src/Psalm/Internal/Type/AssertionReconciler.php

@@ -53,6 +53,7 @@ use Psalm\Type\Atomic\TTrue;
 use function strpos;
 use function substr;
 use Psalm\Issue\InvalidDocblock;
+use Doctrine\Instantiator\Exception\UnexpectedValueException;
 
 class AssertionReconciler extends \Psalm\Type\Reconciler
 {
@@ -191,6 +192,14 @@ class AssertionReconciler extends \Psalm\Type\Reconciler
             return $existing_var_type;
         }
 
+        if (substr($assertion, 0, 9) === 'in-array-') {
+            return self::reconcileInArray(
+                $codebase,
+                $existing_var_type,
+                substr($assertion, 9)
+            );
+        }
+
         if ($assertion === 'falsy' || $assertion === 'empty') {
             return self::reconcileFalsyOrEmpty(
                 $assertion,
@@ -1371,6 +1380,44 @@ class AssertionReconciler extends \Psalm\Type\Reconciler
         return Type::getMixed();
     }
 
+    /**
+     * @param   string[]  $suppressed_issues
+     * @param   0|1|2    $failed_reconciliation
+     */
+    private static function reconcileInArray(
+        Codebase $codebase,
+        Union $existing_var_type,
+        string $assertion
+    ) : Union {
+        if (strpos($assertion, '::')) {
+            list($fq_classlike_name, $const_name) = explode('::', $assertion);
+
+            $class_constant_type = $codebase->classlikes->getConstantForClass(
+                $fq_classlike_name,
+                $const_name,
+                \ReflectionProperty::IS_PRIVATE
+            );
+
+            if ($class_constant_type) {
+                foreach ($class_constant_type->getTypes() as $const_type_atomic) {
+                    if ($const_type_atomic instanceof Type\Atomic\ObjectLike
+                        || $const_type_atomic instanceof Type\Atomic\TArray
+                    ) {
+                        if ($const_type_atomic instanceof Type\Atomic\ObjectLike) {
+                            $const_type_atomic = $const_type_atomic->getGenericArrayType();
+                        }
+
+                        return clone $const_type_atomic->type_params[0];
+                    }
+                }
+            }
+        }
+
+        $existing_var_type->removeType('null');
+
+        return $existing_var_type;
+    }
+
     /**
      * @param   string[]  $suppressed_issues
      * @param   0|1|2    $failed_reconciliation

src/Psalm/Internal/Type/NegatedAssertionReconciler.php

@@ -91,6 +91,8 @@ class NegatedAssertionReconciler extends Reconciler
                 return Type::getNull();
             } elseif ($assertion === 'array-key-exists') {
                 return Type::getEmpty();
+            } elseif (substr($assertion, 0, 9) === 'in-array-') {
+                return $existing_var_type;
             }
         }
 

src/Psalm/Type/Reconciler.php

@@ -84,76 +84,89 @@ class Reconciler
         $suppressed_issues = $statements_analyzer->getSuppressedIssues();
 
         foreach ($new_types as $nk => $type) {
-            if ((strpos($nk, '[') || strpos($nk, '->'))
+            if (strpos($nk, '[') || strpos($nk, '->')) {
-                && ($type[0][0] === '=isset'
+                if ($type[0][0] === '=isset'
                     || $type[0][0] === '!=empty'
                     || $type[0][0] === 'isset'
-                    || $type[0][0] === '!empty')
+                    || $type[0][0] === '!empty'
-            ) {
+                ) {
-                $isset_or_empty = $type[0][0] === 'isset' || $type[0][0] === '=isset'
+                    $isset_or_empty = $type[0][0] === 'isset' || $type[0][0] === '=isset'
-                    ? '=isset'
+                        ? '=isset'
-                    : '!=empty';
+                        : '!=empty';
 
-                $key_parts = Reconciler::breakUpPathIntoParts($nk);
+                    $key_parts = Reconciler::breakUpPathIntoParts($nk);
-
-                if (!$key_parts) {
-                    throw new \UnexpectedValueException('There should be some key parts');
-                }
-
-                $base_key = array_shift($key_parts);
-
-                if (!isset($existing_types[$base_key]) || $existing_types[$base_key]->isNullable()) {
-                    if (!isset($new_types[$base_key])) {
-                        $new_types[$base_key] = [['=isset']];
-                    } else {
-                        $new_types[$base_key][] = ['=isset'];
-                    }
-                }
-
-                while ($key_parts) {
-                    $divider = array_shift($key_parts);
-
-                    if ($divider === '[') {
-                        $array_key = array_shift($key_parts);
-                        array_shift($key_parts);
-
-                        $new_base_key = $base_key . '[' . $array_key . ']';
-
-                        if (strpos($array_key, '\'') !== false) {
-                            $new_types[$base_key][] = ['=string-array-access'];
-                        } else {
-                            $new_types[$base_key][] = ['=int-or-string-array-access'];
-                        }
-
-                        $base_key = $new_base_key;
-
-                        continue;
-                    }
-
-                    if ($divider === '->') {
-                        $property_name = array_shift($key_parts);
-                        $new_base_key = $base_key . '->' . $property_name;
-
-                        $base_key = $new_base_key;
-                    } else {
-                        throw new \InvalidArgumentException('Unexpected divider ' . $divider);
-                    }
 
                     if (!$key_parts) {
-                        break;
+                        throw new \UnexpectedValueException('There should be some key parts');
                     }
 
-                    if (!isset($new_types[$base_key])) {
+                    $base_key = array_shift($key_parts);
-                        $new_types[$base_key] = [['!~bool'], ['!~int'], ['=isset']];
+
-                    } else {
+                    if (!isset($existing_types[$base_key]) || $existing_types[$base_key]->isNullable()) {
-                        $new_types[$base_key][] = ['!~bool'];
+                        if (!isset($new_types[$base_key])) {
-                        $new_types[$base_key][] = ['!~int'];
+                            $new_types[$base_key] = [['=isset']];
-                        $new_types[$base_key][] = ['=isset'];
+                        } else {
+                            $new_types[$base_key][] = ['=isset'];
+                        }
                     }
+
+                    while ($key_parts) {
+                        $divider = array_shift($key_parts);
+
+                        if ($divider === '[') {
+                            $array_key = array_shift($key_parts);
+                            array_shift($key_parts);
+
+                            $new_base_key = $base_key . '[' . $array_key . ']';
+
+                            if (strpos($array_key, '\'') !== false) {
+                                $new_types[$base_key][] = ['=string-array-access'];
+                            } else {
+                                $new_types[$base_key][] = ['=int-or-string-array-access'];
+                            }
+
+                            $base_key = $new_base_key;
+
+                            continue;
+                        }
+
+                        if ($divider === '->') {
+                            $property_name = array_shift($key_parts);
+                            $new_base_key = $base_key . '->' . $property_name;
+
+                            $base_key = $new_base_key;
+                        } else {
+                            throw new \InvalidArgumentException('Unexpected divider ' . $divider);
+                        }
+
+                        if (!$key_parts) {
+                            break;
+                        }
+
+                        if (!isset($new_types[$base_key])) {
+                            $new_types[$base_key] = [['!~bool'], ['!~int'], ['=isset']];
+                        } else {
+                            $new_types[$base_key][] = ['!~bool'];
+                            $new_types[$base_key][] = ['!~int'];
+                            $new_types[$base_key][] = ['=isset'];
+                        }
+                    }
+
+                    // replace with a less specific check
+                    $new_types[$nk][0][0] = $isset_or_empty;
                 }
 
-                // replace with a less specific check
+                if ($type[0][0] === 'array-key-exists') {
-                $new_types[$nk][0][0] = $isset_or_empty;
+                    $key_parts = Reconciler::breakUpPathIntoParts($nk);
+
+                    if (count($key_parts) === 4 && $key_parts[1] === '[') {
+                        if (isset($new_types[$key_parts[2]])) {
+                            $new_types[$key_parts[2]][] = ['=in-array-' . $key_parts[0]];
+                        } else {
+                            $new_types[$key_parts[2]] = [['=in-array-' . $key_parts[0]]];
+                        }
+                    }
+                }
             }
         }
 

tests/AssertTest.php

@@ -379,6 +379,34 @@ class AssertTest extends TestCase
                         }
                     }'
             ],
+            'assertArrayKeyExistsRefinesType' => [
+                '<?php
+                    class Foo {
+                        /** @var array<int,string> */
+                        public const DAYS = [
+                            1 => "mon",
+                            2 => "tue",
+                            3 => "wed",
+                            4 => "thu",
+                            5 => "fri",
+                            6 => "sat",
+                            7 => "sun",
+                        ];
+
+                        /** @param key-of<self::DAYS> $dayNum*/
+                        private static function doGetDayName(int $dayNum): string {
+                            return self::DAYS[$dayNum];
+                        }
+
+                        /** @throws LogicException */
+                        public static function getDayName(int $dayNum): string {
+                            if (! array_key_exists($dayNum, self::DAYS)) {
+                                throw new \LogicException();
+                            }
+                            return self::doGetDayName($dayNum);
+                        }
+                    }'
+            ],
         ];
     }
 }