From c758678022c2f112aa6ddf3d91ac5219b9bac795 Mon Sep 17 00:00:00 2001
From: Matthew Brown <github@muglug.com>
Date: Wed, 6 Feb 2019 15:52:43 -0500
Subject: [PATCH] Fix remaining issues

---
 .../FunctionLike/ReturnTypeCollector.php      |  8 ++++++
 .../Statements/ExpressionAnalyzer.php         | 25 +++++++++++++++++++
 src/Psalm/Internal/Analyzer/TypeAnalyzer.php  |  4 +++
 .../LanguageServer/Client/TextDocument.php    |  3 ++-
 .../LanguageServer/ProtocolStreamReader.php   |  3 +++
 src/Psalm/Internal/Stubs/Amp.php              |  4 +--
 6 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/Psalm/Internal/Analyzer/FunctionLike/ReturnTypeCollector.php b/src/Psalm/Internal/Analyzer/FunctionLike/ReturnTypeCollector.php
index b3da030bf..9e03649fd 100644
--- a/src/Psalm/Internal/Analyzer/FunctionLike/ReturnTypeCollector.php
+++ b/src/Psalm/Internal/Analyzer/FunctionLike/ReturnTypeCollector.php
@@ -182,6 +182,7 @@ class ReturnTypeCollector
                     )
                 );
             } elseif ($stmt instanceof PhpParser\Node\Stmt\While_) {
+                $yield_types = array_merge($yield_types, self::getYieldTypeFromExpression($stmt->cond));
                 $return_types = array_merge(
                     $return_types,
                     self::getReturnTypes(
@@ -303,6 +304,13 @@ class ReturnTypeCollector
             }
 
             return [new Atomic\TMixed()];
+        } elseif ($stmt instanceof PhpParser\Node\Expr\BinaryOp) {
+            return array_merge(
+                self::getYieldTypeFromExpression($stmt->left),
+                self::getYieldTypeFromExpression($stmt->right)
+            );
+        } elseif ($stmt instanceof PhpParser\Node\Expr\Assign) {
+            return self::getYieldTypeFromExpression($stmt->expr);
         }
 
         return [];
diff --git a/src/Psalm/Internal/Analyzer/Statements/ExpressionAnalyzer.php b/src/Psalm/Internal/Analyzer/Statements/ExpressionAnalyzer.php
index 961d6187e..5edcc3167 100644
--- a/src/Psalm/Internal/Analyzer/Statements/ExpressionAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/Statements/ExpressionAnalyzer.php
@@ -8,6 +8,7 @@ use Psalm\Internal\Analyzer\ClosureAnalyzer;
 use Psalm\Internal\Analyzer\CommentAnalyzer;
 use Psalm\Internal\Analyzer\FunctionLikeAnalyzer;
 use Psalm\Internal\Analyzer\ProjectAnalyzer;
+use Psalm\Internal\Analyzer\TraitAnalyzer;
 use Psalm\Internal\Analyzer\Statements\Expression\ArrayAnalyzer;
 use Psalm\Internal\Analyzer\Statements\Expression\AssertionFinder;
 use Psalm\Internal\Analyzer\Statements\Expression\AssignmentAnalyzer;
@@ -1220,6 +1221,30 @@ class ExpressionAnalyzer
             $stmt->inferredType = Type::getNull();
         }
 
+        $source = $statements_analyzer->getSource();
+
+        if ($source instanceof FunctionLikeAnalyzer
+            && !($source->getSource() instanceof TraitAnalyzer)
+        ) {
+            $source->addPossibleParamTypes($context, $codebase);
+
+            $storage = $source->getFunctionLikeStorage($statements_analyzer);
+
+            if ($storage->return_type) {
+                foreach ($storage->return_type->getTypes() as $atomic_return_type) {
+                    if ($atomic_return_type instanceof Type\Atomic\TGenericObject
+                        && $atomic_return_type->value === 'Generator'
+                    ) {
+                        if (!$atomic_return_type->type_params[2]->isMixed()
+                            && !$atomic_return_type->type_params[2]->isVoid()
+                        ) {
+                            $stmt->inferredType = clone $atomic_return_type->type_params[2];
+                        }
+                    }
+                }
+            }
+        }
+
         return null;
     }
 
diff --git a/src/Psalm/Internal/Analyzer/TypeAnalyzer.php b/src/Psalm/Internal/Analyzer/TypeAnalyzer.php
index 415252194..464ba5f45 100644
--- a/src/Psalm/Internal/Analyzer/TypeAnalyzer.php
+++ b/src/Psalm/Internal/Analyzer/TypeAnalyzer.php
@@ -1269,6 +1269,10 @@ class TypeAnalyzer
 
                 $container_param = $container_type_part->type_params[$i];
 
+                if ($input_type_part->value === 'Generator' && $i === 2) {
+                    continue;
+                }
+
                 if (!$input_param->isEmpty() &&
                     !self::isContainedBy(
                         $codebase,
diff --git a/src/Psalm/Internal/LanguageServer/Client/TextDocument.php b/src/Psalm/Internal/LanguageServer/Client/TextDocument.php
index db19d393c..690b0ecdd 100644
--- a/src/Psalm/Internal/LanguageServer/Client/TextDocument.php
+++ b/src/Psalm/Internal/LanguageServer/Client/TextDocument.php
@@ -56,7 +56,7 @@ class TextDocument
     {
         return call(
             /**
-             * @return \Generator<int, Promise, mixed, TextDocumentItem>
+             * @return \Generator<int, mixed, mixed, TextDocumentItem>
              */
             function () use ($textDocument) {
                 $result = yield $this->handler->request(
@@ -64,6 +64,7 @@ class TextDocument
                     ['textDocument' => $textDocument]
                 );
 
+                /** @var TextDocumentItem */
                 return $this->mapper->map($result, new TextDocumentItem);
             }
         );
diff --git a/src/Psalm/Internal/LanguageServer/ProtocolStreamReader.php b/src/Psalm/Internal/LanguageServer/ProtocolStreamReader.php
index 45418b5a3..17371666c 100644
--- a/src/Psalm/Internal/LanguageServer/ProtocolStreamReader.php
+++ b/src/Psalm/Internal/LanguageServer/ProtocolStreamReader.php
@@ -42,6 +42,9 @@ class ProtocolStreamReader implements ProtocolReader
     {
         $input = new ResourceInputStream($input);
         asyncCall(
+            /**
+             * @return \Generator<int, string, string, void>
+             */
             function () use ($input) : \Generator {
                 while ($this->is_accepting_new_requests && ($chunk = yield $input->read()) !== null) {
                     /** @var string $chunk */
diff --git a/src/Psalm/Internal/Stubs/Amp.php b/src/Psalm/Internal/Stubs/Amp.php
index 3f322dbd2..c527489b9 100644
--- a/src/Psalm/Internal/Stubs/Amp.php
+++ b/src/Psalm/Internal/Stubs/Amp.php
@@ -5,9 +5,9 @@ namespace Amp;
 /**
  * @template TReturn
  * @param callable():\Generator<mixed, mixed, mixed, TReturn> $gen
- * @return Promise<TReturn>
+ * @return callable():Promise<TReturn>
  */
-function coroutine(callable $gen) : Promise {}
+function coroutine(callable $gen) : callable {}
 
 /**
  * @template TReturn