Fix #3234 - infer iterator key types properly

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

@@ -854,46 +854,33 @@ class ForeachAnalyzer
                         )
                     )
                 ) {
-                    $old_data_provider = $statements_analyzer->node_data;
+                    $iterator_value_type = self::getFakeMethodCallType(
-
-                    $statements_analyzer->node_data = clone $statements_analyzer->node_data;
-
-                    $fake_method_call = new PhpParser\Node\Expr\MethodCall(
-                        $foreach_expr,
-                        new PhpParser\Node\Identifier('current', $foreach_expr->getAttributes())
-                    );
-
-                    $suppressed_issues = $statements_analyzer->getSuppressedIssues();
-
-                    if (!in_array('PossiblyInvalidMethodCall', $suppressed_issues, true)) {
-                        $statements_analyzer->addSuppressedIssues(['PossiblyInvalidMethodCall']);
-                    }
-
-                    $was_inside_call = $context->inside_call;
-
-                    $context->inside_call = true;
-
-                    \Psalm\Internal\Analyzer\Statements\Expression\Call\MethodCallAnalyzer::analyze(
                         $statements_analyzer,
-                        $fake_method_call,
+                        $foreach_expr,
-                        $context
+                        $context,
+                        'current'
                     );
 
-                    $context->inside_call = $was_inside_call;
+                    $iterator_key_type = self::getFakeMethodCallType(
+                        $statements_analyzer,
+                        $foreach_expr,
+                        $context,
+                        'current'
+                    );
 
-                    if (!in_array('PossiblyInvalidMethodCall', $suppressed_issues, true)) {
+                    if ($iterator_value_type && !$iterator_value_type->isMixed()) {
-                        $statements_analyzer->removeSuppressedIssues(['PossiblyInvalidMethodCall']);
+                        if (!$value_type) {
+                            $value_type = $iterator_value_type;
+                        } else {
+                            $value_type = Type::combineUnionTypes($value_type, $iterator_value_type);
+                        }
                     }
 
-                    $iterator_class_type = $statements_analyzer->node_data->getType($fake_method_call) ?: null;
+                    if ($iterator_key_type && !$iterator_key_type->isMixed()) {
-
+                        if (!$key_type) {
-                    $statements_analyzer->node_data = $old_data_provider;
+                            $key_type = $iterator_key_type;
-
-                    if ($iterator_class_type && !$iterator_class_type->isMixed()) {
-                        if (!$value_type) {
-                            $value_type = $iterator_class_type;
                         } else {
-                            $value_type = Type::combineUnionTypes($value_type, $iterator_class_type);
+                            $key_type = Type::combineUnionTypes($key_type, $iterator_key_type);
                         }
                     }
                 }
@@ -1015,6 +1002,58 @@ class ForeachAnalyzer
         }
     }
 
+    private static function getFakeMethodCallType(
+        StatementsAnalyzer $statements_analyzer,
+        PhpParser\Node\Expr $foreach_expr,
+        Context $context,
+        string $method_name
+    ) : ?Type\Union {
+        $old_data_provider = $statements_analyzer->node_data;
+
+        $statements_analyzer->node_data = clone $statements_analyzer->node_data;
+
+        $fake_method_call = new PhpParser\Node\Expr\MethodCall(
+            $foreach_expr,
+            new PhpParser\Node\Identifier($method_name, $foreach_expr->getAttributes())
+        );
+
+        $suppressed_issues = $statements_analyzer->getSuppressedIssues();
+
+        if (!in_array('PossiblyInvalidMethodCall', $suppressed_issues, true)) {
+            $statements_analyzer->addSuppressedIssues(['PossiblyInvalidMethodCall']);
+        }
+
+        if (!in_array('PossiblyUndefinedMethod', $suppressed_issues, true)) {
+            $statements_analyzer->addSuppressedIssues(['PossiblyUndefinedMethod']);
+        }
+
+        $was_inside_call = $context->inside_call;
+
+        $context->inside_call = true;
+
+        \Psalm\Internal\Analyzer\Statements\Expression\Call\MethodCallAnalyzer::analyze(
+            $statements_analyzer,
+            $fake_method_call,
+            $context
+        );
+
+        $context->inside_call = $was_inside_call;
+
+        if (!in_array('PossiblyInvalidMethodCall', $suppressed_issues, true)) {
+            $statements_analyzer->removeSuppressedIssues(['PossiblyInvalidMethodCall']);
+        }
+
+        if (!in_array('PossiblyUndefinedMethod', $suppressed_issues, true)) {
+            $statements_analyzer->removeSuppressedIssues(['PossiblyUndefinedMethod']);
+        }
+
+        $iterator_class_type = $statements_analyzer->node_data->getType($fake_method_call) ?: null;
+
+        $statements_analyzer->node_data = $old_data_provider;
+
+        return $iterator_class_type;
+    }
+
     /**
      * @param  string $template_name
      * @param  array<string, array<int|string, Type\Union>>  $template_type_extends

tests/Loop/ForeachTest.php

@@ -999,6 +999,48 @@ class ForeachTest extends \Psalm\Tests\TestCase
                         }
                     }'
             ],
+            'iteratorForeach' => [
+                '<?php
+                    /**
+                     * @implements Iterator<int, string>
+                     */
+                    class FooIterator implements \Iterator {
+                        private ?int $key = null;
+
+                        public function current(): string
+                        {
+                            return "a";
+                        }
+
+                        public function next(): void
+                        {
+                            $this->key = $this->key === null ? 0 : $this->key + 1;
+                        }
+
+                        public function key(): int
+                        {
+                            if ($this->key === null) {
+                                throw new \Exception();
+                            }
+                            return $this->key;
+                        }
+
+                        public function valid(): bool
+                        {
+                            return $this->key !== null && $this->key <= 3;
+                        }
+
+                        public function rewind(): void
+                        {
+                            $this->key = null;
+                            $this->next();
+                        }
+                    }
+
+                    foreach (new FooIterator() as $key => $value) {
+                        echo $key . " " . $value;
+                    }'
+            ],
         ];
     }