From 2a2682a05a47b392e430be97e7d592254483a8c8 Mon Sep 17 00:00:00 2001
From: Artem Tyurin <artem.tyurin@uber.com>
Date: Fri, 20 Apr 2018 18:38:55 +0200
Subject: [PATCH] Check return type inside generators

---
 src/compiler/checker.ts                             | 13 +++++++------
 .../reference/generatorTypeCheck12.errors.txt       |  9 +++++++++
 .../reference/generatorTypeCheck13.errors.txt       | 10 ++++++++++
 3 files changed, 26 insertions(+), 6 deletions(-)
 create mode 100644 tests/baselines/reference/generatorTypeCheck12.errors.txt
 create mode 100644 tests/baselines/reference/generatorTypeCheck13.errors.txt

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index c1cb71660150d..dffa3879d9275 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -23219,14 +23219,15 @@ namespace ts {
             const returnType = getReturnTypeOfSignature(signature);
             const functionFlags = getFunctionFlags(func);
             const isGenerator = functionFlags & FunctionFlags.Generator;
+            const isAsync = (functionFlags & FunctionFlags.Async) !== 0;
             if (strictNullChecks || node.expression || returnType.flags & TypeFlags.Never) {
                 const exprType = node.expression ? checkExpressionCached(node.expression) : undefinedType;
-                if (isGenerator) { // AsyncGenerator function or Generator function
-                    // A generator does not need its return expressions checked against its return type.
-                    // Instead, the yield expressions are checked against the element type.
-                    // TODO: Check return types of generators when return type tracking is added
-                    // for generators.
-                    return;
+                if (isGenerator) {
+                    const returnType = getEffectiveReturnTypeNode(func);
+                    if (returnType) {
+                        const signatureElementType = getIteratedTypeOfGenerator(getTypeFromTypeNode(returnType), isAsync) || anyType;
+                        checkTypeAssignableTo(exprType, signatureElementType, node);
+                    }
                 }
                 else if (func.kind === SyntaxKind.SetAccessor) {
                     if (node.expression) {
diff --git a/tests/baselines/reference/generatorTypeCheck12.errors.txt b/tests/baselines/reference/generatorTypeCheck12.errors.txt
new file mode 100644
index 0000000000000..f7e71442dec36
--- /dev/null
+++ b/tests/baselines/reference/generatorTypeCheck12.errors.txt
@@ -0,0 +1,9 @@
+tests/cases/conformance/es6/yieldExpressions/generatorTypeCheck12.ts(2,5): error TS2322: Type '""' is not assignable to type 'number'.
+
+
+==== tests/cases/conformance/es6/yieldExpressions/generatorTypeCheck12.ts (1 errors) ====
+    function* g(): IterableIterator<number> {
+        return "";
+        ~~~~~~~~~~
+!!! error TS2322: Type '""' is not assignable to type 'number'.
+    }
\ No newline at end of file
diff --git a/tests/baselines/reference/generatorTypeCheck13.errors.txt b/tests/baselines/reference/generatorTypeCheck13.errors.txt
new file mode 100644
index 0000000000000..3914beb02f7a1
--- /dev/null
+++ b/tests/baselines/reference/generatorTypeCheck13.errors.txt
@@ -0,0 +1,10 @@
+tests/cases/conformance/es6/yieldExpressions/generatorTypeCheck13.ts(3,5): error TS2322: Type '""' is not assignable to type 'number'.
+
+
+==== tests/cases/conformance/es6/yieldExpressions/generatorTypeCheck13.ts (1 errors) ====
+    function* g(): IterableIterator<number> {
+        yield 0;
+        return "";
+        ~~~~~~~~~~
+!!! error TS2322: Type '""' is not assignable to type 'number'.
+    }
\ No newline at end of file