Support simple list assignment in foreach

Ref #4741

src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php

@@ -414,12 +414,22 @@ class ForeachAnalyzer
                     }
                     $iterator_atomic_type = $iterator_atomic_type->getGenericArrayType();
                 } elseif ($iterator_atomic_type instanceof Type\Atomic\TList) {
+                    $list_var_id = ExpressionIdentifier::getArrayVarId(
+                        $stmt->expr,
+                        $statements_analyzer->getFQCLN(),
+                        $statements_analyzer
+                    );
+
                     if (!$iterator_atomic_type instanceof Type\Atomic\TNonEmptyList) {
                         $always_non_empty_array = false;
                     }
 
                     $iterator_atomic_type = new Type\Atomic\TArray([
-                        Type::getInt(),
+                        $list_var_id
+                            ? new Type\Union([
+                                new Type\Atomic\TDependentListKey($list_var_id)
+                            ])
+                            : Type::getInt(),
                         $iterator_atomic_type->type_param
                     ]);
                 } elseif (!$iterator_atomic_type instanceof Type\Atomic\TNonEmptyArray) {

src/Psalm/Internal/Analyzer/Statements/Expression/Assignment/ArrayAssignmentAnalyzer.php

@@ -437,6 +437,16 @@ class ArrayAssignmentAnalyzer
                     $current_dim_type = Type::getArrayKey();
                 }
 
+                if ($current_dim_type->isSingle()) {
+                    $current_dim_type_type = \array_values($current_dim_type->getAtomicTypes())[0];
+
+                    if ($current_dim_type_type instanceof Type\Atomic\TDependentListKey
+                        && $current_dim_type_type->getVarId() === $parent_var_id
+                    ) {
+                        $offset_already_existed = true;
+                    }
+                }
+
                 $array_atomic_key_type = ArrayFetchAnalyzer::replaceOffsetTypeWithInts(
                     $current_dim_type
                 );
@@ -580,6 +590,7 @@ class ArrayAssignmentAnalyzer
                 $atomic_root_types['array']->count++;
             }
         }
+
         return $new_child_type;
     }
 
@@ -604,6 +615,7 @@ class ArrayAssignmentAnalyzer
         $reversed_child_stmts = [];
         $var_id_additions = [];
         $full_var_id = true;
+
         $child_stmt = null;
 
         // First go from the root element up, and go as far as we can to figure out what

src/Psalm/Internal/Analyzer/Statements/Expression/Call/ArgumentAnalyzer.php

@@ -1159,7 +1159,6 @@ class ArgumentAnalyzer
                                 if ($type_param->isEmpty() && isset($param_atomic_type->type_params[$i])) {
                                     $input_type_changed = true;
 
-                                    /** @psalm-suppress PropertyTypeCoercion */
                                     $input_atomic_type->type_params[$i] = clone $param_atomic_type->type_params[$i];
                                 }
                             }

src/Psalm/Internal/Analyzer/Statements/Expression/Call/ArrayFunctionArgumentsAnalyzer.php

@@ -532,7 +532,6 @@ class ArrayFunctionArgumentsAnalyzer
                         }
 
                         $array_type->addType($array_atomic_type);
-                        $context->removeDescendents($var_id, $array_type);
                     } elseif ($array_atomic_type instanceof TNonEmptyList) {
                         if (!$context->inside_loop && $array_atomic_type->count !== null) {
                             if ($array_atomic_type->count === 0) {
@@ -550,10 +549,10 @@ class ArrayFunctionArgumentsAnalyzer
                         }
 
                         $array_type->addType($array_atomic_type);
-                        $context->removeDescendents($var_id, $array_type);
                     }
                 }
 
+                $context->removeDescendents($var_id, $array_type);
                 $context->vars_in_scope[$var_id] = $array_type;
             }
         }

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

@@ -18,6 +18,7 @@ use Psalm\Type\Atomic\TTemplateParamClass;
 use Psalm\Type\Atomic\TDependentGetClass;
 use Psalm\Type\Atomic\TDependentGetDebugType;
 use Psalm\Type\Atomic\TDependentGetType;
+use Psalm\Type\Atomic\TDependentListKey;
 use Psalm\Type\Atomic\THtmlEscapedString;
 use Psalm\Type\Atomic\TInt;
 use Psalm\Type\Atomic\TLiteralClassString;
@@ -245,11 +246,11 @@ class ScalarTypeComparator
             return true;
         }
 
-        if (get_class($container_type_part) === TInt::class && $input_type_part instanceof TLiteralInt) {
-            return true;
-        }
-
-        if (get_class($container_type_part) === TInt::class && $input_type_part instanceof TPositiveInt) {
+        if ((get_class($container_type_part) === TInt::class
+                || get_class($container_type_part) === TDependentListKey::class)
+            && ($input_type_part instanceof TLiteralInt
+                || $input_type_part instanceof TPositiveInt)
+        ) {
             return true;
         }
 
@@ -272,7 +273,8 @@ class ScalarTypeComparator
             return true;
         }
 
-        if (get_class($container_type_part) === TInt::class
+        if ((get_class($container_type_part) === TInt::class
+                || get_class($container_type_part) === TDependentListKey::class)
             && $input_type_part instanceof TInt
         ) {
             return true;

src/Psalm/Type/Atomic/TDependentListKey.php

@@ -0,0 +1,48 @@
+<?php
+namespace Psalm\Type\Atomic;
+
+/**
+ * Represents a list key created from foreach ($list as $key => $value)
+ */
+class TDependentListKey extends TInt implements DependentType
+{
+    /**
+     * Used to hold information as to what list variable this refers to
+     *
+     * @var string
+     */
+    public $var_id;
+
+    /**
+     * @param string $var_id the variable id
+     */
+    public function __construct(string $var_id)
+    {
+        $this->var_id = $var_id;
+    }
+
+    public function getId(bool $nested = false): string
+    {
+        return 'list-key<' . $this->var_id . '>';
+    }
+
+    public function getVarId() : string
+    {
+        return $this->var_id;
+    }
+
+    public function getAssertionString(): string
+    {
+        return 'int';
+    }
+
+    public function getReplacement() : \Psalm\Type\Atomic
+    {
+        return new TInt();
+    }
+
+    public function canBeFullyExpressedInPhp(int $php_major_version, int $php_minor_version): bool
+    {
+        return false;
+    }
+}

tests/ArrayAssignmentTest.php

@@ -1509,6 +1509,20 @@ class ArrayAssignmentTest extends TestCase
                     ];
                 ',
             ],
+            'assignToListWithForeachKey' => [
+                '<?php
+                    /**
+                     * @param list<string> $list
+                     * @return list<string>
+                     */
+                    function getList(array $list): array {
+                        foreach ($list as $key => $value) {
+                            $list[$key] = $value . "!";
+                        }
+
+                        return $list;
+                    }'
+            ],
         ];
     }
 
@@ -1825,6 +1839,40 @@ class ArrayAssignmentTest extends TestCase
                     }',
                 'error_message' => 'RedundantCast',
             ],
+            'assignToListWithUpdatedForeachKey' => [
+                '<?php
+                    /**
+                     * @param list<string> $list
+                     * @return list<string>
+                     */
+                    function getList(array $list): array {
+                        foreach ($list as $key => $value) {
+                            $list[$key + 1] = $value . "!";
+                        }
+
+                        return $list;
+                    }',
+                'error_message' => 'LessSpecificReturnStatement',
+            ],
+            'assignToListWithAlteredForeachKeyVar' => [
+                '<?php
+                    /**
+                     * @param list<string> $list
+                     * @return list<string>
+                     */
+                    function getList(array $list): array {
+                        foreach ($list as $key => $value) {
+                            if (rand(0, 1)) {
+                                array_pop($list);
+                            }
+
+                            $list[$key] = $value . "!";
+                        }
+
+                        return $list;
+                    }',
+                'error_message' => 'LessSpecificReturnStatement',
+            ],
         ];
     }
 }