From f55167e5658f92061b767e4cee839833a022d327 Mon Sep 17 00:00:00 2001
From: Kanchalai Tanglertsampan <yuisu@microsoft.com>
Date: Wed, 22 Feb 2017 10:23:47 -0800
Subject: [PATCH] Collect type from return statment in generator function

---
 src/compiler/checker.ts | 16 +++++++++-------
 1 file changed, 9 insertions(+), 7 deletions(-)

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index 12b55aa818b..56306fb772f 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -14963,10 +14963,12 @@ namespace ts {
                 }
             }
             else {
-                let types: Type[];
+                let yieldTypes: Type[];
+                let returnTypes: Type[];
                 if (functionFlags & FunctionFlags.Generator) { // Generator or AsyncGenerator function
-                    types = checkAndAggregateYieldOperandTypes(func, contextualMapper);
-                    if (types.length === 0) {
+                    yieldTypes = checkAndAggregateYieldOperandTypes(func, contextualMapper);
+                    returnTypes = checkAndAggregateReturnExpressionTypes(func, contextualMapper);
+                    if (yieldTypes.length === 0 && (!returnTypes || returnTypes.length === 0)) {
                         const iterableIteratorAny = functionFlags & FunctionFlags.Async
                             ? createAsyncIterableIteratorType(anyType) // AsyncGenerator function
                             : createIterableIteratorType(anyType); // Generator function
@@ -14978,14 +14980,14 @@ namespace ts {
                     }
                 }
                 else {
-                    types = checkAndAggregateReturnExpressionTypes(func, contextualMapper);
-                    if (!types) {
+                    returnTypes = checkAndAggregateReturnExpressionTypes(func, contextualMapper);
+                    if (!returnTypes) {
                         // For an async function, the return type will not be never, but rather a Promise for never.
                         return functionFlags & FunctionFlags.Async
                             ? createPromiseReturnType(func, neverType) // Async function
                             : neverType; // Normal function
                     }
-                    if (types.length === 0) {
+                    if (returnTypes.length === 0) {
                         // For an async function, the return type will not be void, but rather a Promise for void.
                         return functionFlags & FunctionFlags.Async
                             ? createPromiseReturnType(func, voidType) // Async function
@@ -14993,7 +14995,7 @@ namespace ts {
                     }
                 }
                 // Return a union of the return expression types.
-                type = getUnionType(types, /*subtypeReduction*/ true);
+                type = getUnionType(yieldTypes ? yieldTypes.concat(returnTypes) : returnTypes, /*subtypeReduction*/ true);
 
                 if (functionFlags & FunctionFlags.Generator) { // AsyncGenerator function or Generator function
                     type = functionFlags & FunctionFlags.Async