From 829a6701fc9a54bce003f6417055355e1a157310 Mon Sep 17 00:00:00 2001
From: Bruce Weirdan <weirdan@gmail.com>
Date: Mon, 12 Feb 2024 05:14:17 +0100
Subject: [PATCH] Forbid iterating over generators with non-nullable `send()`

Fixes vimeo/psalm#6985
---
 .../Statements/Block/ForeachAnalyzer.php      | 41 ++++++++++++++++++-
 tests/Loop/ForeachTest.php                    | 34 +++++++++++++++
 2 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
index 8f0ec8a713f..8c16dbaf25f 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
@@ -657,6 +657,7 @@ public static function checkIteratorType(
                         $key_type,
                         $value_type,
                         $has_valid_iterator,
+                        $invalid_iterator_types,
                     );
                 } else {
                     $raw_object_types[] = $iterator_atomic_type->value;
@@ -725,6 +726,7 @@ public static function checkIteratorType(
         return null;
     }
 
+    /** @param list<string> $invalid_iterator_types */
     public static function handleIterable(
         StatementsAnalyzer $statements_analyzer,
         TNamedObject $iterator_atomic_type,
@@ -733,7 +735,8 @@ public static function handleIterable(
         Context $context,
         ?Union &$key_type,
         ?Union &$value_type,
-        bool &$has_valid_iterator
+        bool &$has_valid_iterator,
+        array &$invalid_iterator_types = []
     ): void {
         if ($iterator_atomic_type->extra_types) {
             $iterator_atomic_types = array_merge(
@@ -753,7 +756,6 @@ public static function handleIterable(
             }
 
 
-            $has_valid_iterator = true;
 
             if ($iterator_atomic_type instanceof TIterable
                 || (strtolower($iterator_atomic_type->value) === 'traversable'
@@ -781,6 +783,8 @@ public static function handleIterable(
                         )
                     )
                 ) {
+                    $has_valid_iterator = true;
+
                     $old_data_provider = $statements_analyzer->node_data;
 
                     $statements_analyzer->node_data = clone $statements_analyzer->node_data;
@@ -867,6 +871,7 @@ public static function handleIterable(
                                             $key_type,
                                             $value_type,
                                             $has_valid_iterator,
+                                            $invalid_iterator_types,
                                         );
 
                                         continue;
@@ -899,6 +904,37 @@ public static function handleIterable(
                             $value_type = Type::combineUnionTypes($value_type, $value_type_part);
                         }
                     }
+                } elseif ($iterator_atomic_type instanceof TGenericObject
+                    && strtolower($iterator_atomic_type->value) === 'generator'
+                ) {
+                    $type_params = $iterator_atomic_type->type_params;
+                    if (isset($type_params[2]) && !$type_params[2]->isNullable() && !$type_params[2]->isMixed()) {
+                        $invalid_iterator_types[] = $iterator_atomic_type->getKey();
+                    } else {
+                        $has_valid_iterator = true;
+                    }
+
+                    $iterator_value_type = self::getFakeMethodCallType(
+                        $statements_analyzer,
+                        $foreach_expr,
+                        $context,
+                        'current',
+                    );
+
+                    $iterator_key_type = self::getFakeMethodCallType(
+                        $statements_analyzer,
+                        $foreach_expr,
+                        $context,
+                        'key',
+                    );
+
+                    if ($iterator_value_type && !$iterator_value_type->isMixed()) {
+                        $value_type = Type::combineUnionTypes($value_type, $iterator_value_type);
+                    }
+
+                    if ($iterator_key_type && !$iterator_key_type->isMixed()) {
+                        $key_type = Type::combineUnionTypes($key_type, $iterator_key_type);
+                    }
                 } elseif ($codebase->classImplements(
                     $iterator_atomic_type->value,
                     'Iterator',
@@ -911,6 +947,7 @@ public static function handleIterable(
                         )
                     )
                 ) {
+                    $has_valid_iterator = true;
                     $iterator_value_type = self::getFakeMethodCallType(
                         $statements_analyzer,
                         $foreach_expr,
diff --git a/tests/Loop/ForeachTest.php b/tests/Loop/ForeachTest.php
index d7b56e5fa48..ba1ca170b88 100644
--- a/tests/Loop/ForeachTest.php
+++ b/tests/Loop/ForeachTest.php
@@ -1169,6 +1169,28 @@ function f(array $a): array {
                     }
                     PHP,
             ],
+            'generatorWithUnspecifiedSend' => [
+                'code' => <<<'PHP'
+                    <?php
+                    /** @return Generator<int,int> */
+                    function gen() : Generator {
+                        return yield 1;
+                    }
+                    $gen = gen();
+                    foreach ($gen as $i) {}
+                PHP,
+            ],
+            'generatorWithMixedSend' => [
+                'code' => <<<'PHP'
+                    <?php
+                    /** @return Generator<int,int, mixed, mixed> */
+                    function gen() : Generator {
+                        return yield 1;
+                    }
+                    $gen = gen();
+                    foreach ($gen as $i) {}
+                PHP,
+            ],
         ];
     }
 
@@ -1395,6 +1417,18 @@ function f(array $a): array {
                     PHP,
                 'error_message' => 'LessSpecificReturnStatement',
             ],
+            'generatorWithNonNullableSend' => [
+                'code' => <<<'PHP'
+                    <?php
+                    /** @return Generator<int,int,string,string> */
+                    function gen() : Generator {
+                        return yield 1;
+                    }
+                    $gen = gen();
+                    foreach ($gen as $i) {}
+                PHP,
+                'error_message' => 'InvalidIterator',
+            ],
         ];
     }
 }