From 27461c98f6f9def9d1a6a4e4643c6066e65a5f9d Mon Sep 17 00:00:00 2001
From: Bruce Weirdan <weirdan@gmail.com>
Date: Mon, 12 Feb 2024 07:47:58 +0100
Subject: [PATCH] Strip null used to signify completed iterations in foreach
 context

Even though `Generator::current()` can return `null` once generator is
exhausted, `foreach()` never iterates after iterator ends, so we can
safely remove `null` (unless, of course, generator can yield `null`).
---
 .../Statements/Block/ForeachAnalyzer.php      | 14 ++++++++
 tests/Loop/ForeachTest.php                    | 35 +++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
index 8c16dbaf2..19099c9c2 100644
--- a/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
@@ -929,10 +929,24 @@ final class ForeachAnalyzer
                     );
 
                     if ($iterator_value_type && !$iterator_value_type->isMixed()) {
+                        // remove null coming from current() to signify invalid iterations
+                        // we're in a foreach context, so we know we're not going iterate past the end
+                        if (isset($type_params[1]) && !$type_params[1]->isNullable()) {
+                            $iterator_value_type = $iterator_value_type->getBuilder();
+                            $iterator_value_type->removeType('null');
+                            $iterator_value_type = $iterator_value_type->freeze();
+                        }
                         $value_type = Type::combineUnionTypes($value_type, $iterator_value_type);
                     }
 
                     if ($iterator_key_type && !$iterator_key_type->isMixed()) {
+                        // remove null coming from key() to signify invalid iterations
+                        // we're in a foreach context, so we know we're not going iterate past the end
+                        if (isset($type_params[0]) && !$type_params[0]->isNullable()) {
+                            $iterator_key_type = $iterator_key_type->getBuilder();
+                            $iterator_key_type->removeType('null');
+                            $iterator_key_type = $iterator_key_type->freeze();
+                        }
                         $key_type = Type::combineUnionTypes($key_type, $iterator_key_type);
                     }
                 } elseif ($codebase->classImplements(
diff --git a/tests/Loop/ForeachTest.php b/tests/Loop/ForeachTest.php
index ba1ca170b..0777815cb 100644
--- a/tests/Loop/ForeachTest.php
+++ b/tests/Loop/ForeachTest.php
@@ -1191,6 +1191,41 @@ class ForeachTest extends TestCase
                     foreach ($gen as $i) {}
                 PHP,
             ],
+            'nullableGenerator' => [
+                'code' => <<<'PHP'
+                    <?php
+                    /** @return Generator<int,int|null> */
+                    function gen() : Generator {
+                        yield null;
+                        yield 1;
+                    }
+                    $gen = gen();
+                    $a = "";
+                    foreach ($gen as $i) {
+                        $a = $i;
+                    }
+                PHP,
+                'assertions' => [
+                    '$a===' => "''|int|null",
+                ],
+            ],
+            'nonNullableGenerator' => [
+                'code' => <<<'PHP'
+                    <?php
+                    /** @return Generator<int,int> */
+                    function gen() : Generator {
+                        yield 1;
+                    }
+                    $gen = gen();
+                    $a = "";
+                    foreach ($gen as $i) {
+                $a = $i;
+                    }
+                PHP,
+                'assertions' => [
+                    '$a===' => "''|int",
+                ],
+            ],
         ];
     }