fix return statement contextual typing

Author: gabritto

src/compiler/checker.ts

@@ -26638,6 +26638,12 @@ namespace ts {
                 if (contextualReturnType) {
                     const functionFlags = getFunctionFlags(func);
                     if (functionFlags & FunctionFlags.Generator) { // Generator or AsyncGenerator function
+                        if (contextualReturnType.flags & TypeFlags.Union) {
+                            contextualReturnType = filterType(contextualReturnType, type => {
+                                const generator = createGeneratorReturnType(anyType, anyType, anyType, (functionFlags & FunctionFlags.Async) !== 0);
+                                return checkTypeAssignableTo(generator, type, /*errorNode*/ undefined);
+                            });
+                        }
                         const iterationReturnType = getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Return, contextualReturnType, (functionFlags & FunctionFlags.Async) !== 0);
                         if (!iterationReturnType) {
                             return undefined;
@@ -26676,14 +26682,10 @@ namespace ts {
                     const isAsyncGenerator = (functionFlags & FunctionFlags.Async) !== 0;
                     if (!node.asteriskToken && contextualReturnType.flags & TypeFlags.Union) {
                         contextualReturnType = filterType(contextualReturnType, type => {
-                            // >> What should we use to filter?
-                            // >> What will be the perf impacts? Is there caching? Is it enough to make this ok?
-                            // return !!getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, type, isAsyncGenerator);
-                            // >> ^^^ doesn't work, because we get { yield: any, return: any } for bogus types
                             const generator = createGeneratorReturnType(anyType, anyType, anyType, isAsyncGenerator);
                             return checkTypeAssignableTo(generator, type, /*errorNode*/ undefined);
                         });
-                    } // >> TODO: should we do something when `node.asteriskToken`?
+                    }
                     return node.asteriskToken
                         ? contextualReturnType
                         : getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, contextualReturnType, isAsyncGenerator);

tests/baselines/reference/contextualTypeOnYield.js renamed to tests/baselines/reference/contextualTypeOnYield1.js

@@ -1,11 +1,11 @@
-//// [contextualTypeOnYield.ts]
+//// [contextualTypeOnYield1.ts]
 type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
 
 const f: FuncOrGeneratorFunc = function*() {
   yield (num) => console.log(num); // `num` should be inferred to have type `number`.
 }
 
-//// [contextualTypeOnYield.js]
+//// [contextualTypeOnYield1.js]
 "use strict";
 const f = function* () {
     yield (num) => console.log(num); // `num` should be inferred to have type `number`.

tests/baselines/reference/contextualTypeOnYield.symbols renamed to tests/baselines/reference/contextualTypeOnYield1.symbols

@@ -1,17 +1,17 @@
-=== tests/cases/compiler/contextualTypeOnYield.ts ===
+=== tests/cases/compiler/contextualTypeOnYield1.ts ===
 type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
->FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield.ts, 0, 0))
+>FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield1.ts, 0, 0))
 >Generator : Symbol(Generator, Decl(lib.es2015.generator.d.ts, --, --))
->arg : Symbol(arg, Decl(contextualTypeOnYield.ts, 0, 54))
+>arg : Symbol(arg, Decl(contextualTypeOnYield1.ts, 0, 54))
 
 const f: FuncOrGeneratorFunc = function*() {
->f : Symbol(f, Decl(contextualTypeOnYield.ts, 2, 5))
->FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield.ts, 0, 0))
+>f : Symbol(f, Decl(contextualTypeOnYield1.ts, 2, 5))
+>FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield1.ts, 0, 0))
 
   yield (num) => console.log(num); // `num` should be inferred to have type `number`.
->num : Symbol(num, Decl(contextualTypeOnYield.ts, 3, 9))
+>num : Symbol(num, Decl(contextualTypeOnYield1.ts, 3, 9))
 >console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
 >console : Symbol(console, Decl(lib.dom.d.ts, --, --))
 >log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
->num : Symbol(num, Decl(contextualTypeOnYield.ts, 3, 9))
+>num : Symbol(num, Decl(contextualTypeOnYield1.ts, 3, 9))
 }

tests/baselines/reference/contextualTypeOnYield.types renamed to tests/baselines/reference/contextualTypeOnYield1.types

@@ -1,4 +1,4 @@
-=== tests/cases/compiler/contextualTypeOnYield.ts ===
+=== tests/cases/compiler/contextualTypeOnYield1.ts ===
 type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
 >FuncOrGeneratorFunc : () => number | Generator<(arg: number) => void, any, void>
 >arg : number

tests/baselines/reference/contextualTypeOnYield2.js

@@ -0,0 +1,11 @@
+//// [contextualTypeOnYield2.ts]
+type OrGen = () => (number | Generator<string, (arg: number) => void, undefined>);
+const g: OrGen = function* () {
+    return (num) => console.log(num);
+}
+
+//// [contextualTypeOnYield2.js]
+"use strict";
+const g = function* () {
+    return (num) => console.log(num);
+};

tests/baselines/reference/contextualTypeOnYield2.symbols

@@ -0,0 +1,17 @@
+=== tests/cases/compiler/contextualTypeOnYield2.ts ===
+type OrGen = () => (number | Generator<string, (arg: number) => void, undefined>);
+>OrGen : Symbol(OrGen, Decl(contextualTypeOnYield2.ts, 0, 0))
+>Generator : Symbol(Generator, Decl(lib.es2015.generator.d.ts, --, --))
+>arg : Symbol(arg, Decl(contextualTypeOnYield2.ts, 0, 48))
+
+const g: OrGen = function* () {
+>g : Symbol(g, Decl(contextualTypeOnYield2.ts, 1, 5))
+>OrGen : Symbol(OrGen, Decl(contextualTypeOnYield2.ts, 0, 0))
+
+    return (num) => console.log(num);
+>num : Symbol(num, Decl(contextualTypeOnYield2.ts, 2, 12))
+>console.log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>console : Symbol(console, Decl(lib.dom.d.ts, --, --))
+>log : Symbol(Console.log, Decl(lib.dom.d.ts, --, --))
+>num : Symbol(num, Decl(contextualTypeOnYield2.ts, 2, 12))
+}

tests/baselines/reference/contextualTypeOnYield2.types

@@ -0,0 +1,18 @@
+=== tests/cases/compiler/contextualTypeOnYield2.ts ===
+type OrGen = () => (number | Generator<string, (arg: number) => void, undefined>);
+>OrGen : () => number | Generator<string, (arg: number) => void, undefined>
+>arg : number
+
+const g: OrGen = function* () {
+>g : OrGen
+>function* () {    return (num) => console.log(num);} : () => Generator<never, (num: number) => void, any>
+
+    return (num) => console.log(num);
+>(num) => console.log(num) : (num: number) => void
+>num : number
+>console.log(num) : void
+>console.log : (...data: any[]) => void
+>console : Console
+>log : (...data: any[]) => void
+>num : number
+}

tests/cases/compiler/contextualTypeOnYield.ts renamed to tests/cases/compiler/contextualTypeOnYield1.ts

@@ -0,0 +1,7 @@
+// @strict: true
+// @target: ES6
+
+type OrGen = () => (number | Generator<string, (arg: number) => void, undefined>);
+const g: OrGen = function* () {
+    return (num) => console.log(num);
+}