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`).
vimeo · Feb 12, 2024 · ff66d24 · ff66d24
1 parent 829a670
commit ff66d24
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 0 deletions.
diff --git a/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/Block/ForeachAnalyzer.php
@@ -929,10 +929,24 @@ public static function handleIterable(
                     );
 
                     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
@@ -1191,6 +1191,41 @@ function gen() : Generator {
                     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",
+                ],
+            ],
         ];
     }