Fix contextual typing on yield and return expressions in generator function

src/compiler/checker.ts

@@ -26777,12 +26777,15 @@ namespace ts {
                 if (contextualReturnType) {
                     const functionFlags = getFunctionFlags(func);
                     if (functionFlags & FunctionFlags.Generator) { // Generator or AsyncGenerator function
-                        const use = functionFlags & FunctionFlags.Async ? IterationUse.AsyncGeneratorReturnType : IterationUse.GeneratorReturnType;
-                        const iterationTypes = getIterationTypesOfIterable(contextualReturnType, use, /*errorNode*/ undefined);
-                        if (!iterationTypes) {
+                        const isAsyncGenerator = (functionFlags & FunctionFlags.Async) !== 0;
+                        if (contextualReturnType.flags & TypeFlags.Union) {
+                            contextualReturnType = filterType(contextualReturnType, type => !!getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Return, type, isAsyncGenerator));
+                        }
+                        const iterationReturnType = getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Return, contextualReturnType, (functionFlags & FunctionFlags.Async) !== 0);
+                        if (!iterationReturnType) {
                             return undefined;
                         }
-                        contextualReturnType = iterationTypes.returnType;
+                        contextualReturnType = iterationReturnType;
                         // falls through to unwrap Promise for AsyncGenerators
                     }
 
@@ -26811,11 +26814,15 @@ namespace ts {
             const func = getContainingFunction(node);
             if (func) {
                 const functionFlags = getFunctionFlags(func);
-                const contextualReturnType = getContextualReturnType(func, contextFlags);
+                let contextualReturnType = getContextualReturnType(func, contextFlags);
                 if (contextualReturnType) {
+                    const isAsyncGenerator = (functionFlags & FunctionFlags.Async) !== 0;
+                    if (!node.asteriskToken && contextualReturnType.flags & TypeFlags.Union) {
+                        contextualReturnType = filterType(contextualReturnType, type => !!getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Return, type, isAsyncGenerator));
+                    }
                     return node.asteriskToken
                         ? contextualReturnType
-                        : getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, contextualReturnType, (functionFlags & FunctionFlags.Async) !== 0);
+                        : getIterationTypeOfGeneratorFunctionReturnType(IterationTypeKind.Yield, contextualReturnType, isAsyncGenerator);
                 }
             }
 

tests/baselines/reference/contextualTypeOnYield1.js

@@ -0,0 +1,12 @@
+//// [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`.
+}
+
+//// [contextualTypeOnYield1.js]
+"use strict";
+const f = function* () {
+    yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+};

tests/baselines/reference/contextualTypeOnYield1.symbols

@@ -0,0 +1,17 @@
+=== tests/cases/compiler/contextualTypeOnYield1.ts ===
+type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
+>FuncOrGeneratorFunc : Symbol(FuncOrGeneratorFunc, Decl(contextualTypeOnYield1.ts, 0, 0))
+>Generator : Symbol(Generator, Decl(lib.es2015.generator.d.ts, --, --))
+>arg : Symbol(arg, Decl(contextualTypeOnYield1.ts, 0, 54))
+
+const f: FuncOrGeneratorFunc = function*() {
+>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(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(contextualTypeOnYield1.ts, 3, 9))
+}

tests/baselines/reference/contextualTypeOnYield1.types

@@ -0,0 +1,19 @@
+=== tests/cases/compiler/contextualTypeOnYield1.ts ===
+type FuncOrGeneratorFunc = () => (number | Generator<(arg: number) => void, any, void>)
+>FuncOrGeneratorFunc : () => number | Generator<(arg: number) => void, any, void>
+>arg : number
+
+const f: FuncOrGeneratorFunc = function*() {
+>f : FuncOrGeneratorFunc
+>function*() {  yield (num) => console.log(num); // `num` should be inferred to have type `number`.} : () => Generator<(num: number) => void, void, unknown>
+
+  yield (num) => console.log(num); // `num` should be inferred to have type `number`.
+>yield (num) => console.log(num) : any
+>(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/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, unknown>
+
+    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/baselines/reference/generatorReturnContextualType.errors.txt

@@ -1,29 +1,53 @@
-tests/cases/conformance/generators/generatorReturnContextualType.ts(17,3): error TS2322: Type '{ x: string; }' is not assignable to type '{ x: "x"; }'.
+tests/cases/conformance/generators/generatorReturnContextualType.ts(29,3): error TS2322: Type '{ x: string; }' is not assignable to type '{ x: "x"; }'.
+  Types of property 'x' are incompatible.
+    Type 'string' is not assignable to type '"x"'.
+tests/cases/conformance/generators/generatorReturnContextualType.ts(34,3): error TS2322: Type '{ x: string; }' is not assignable to type '{ x: "x"; }'.
   Types of property 'x' are incompatible.
     Type 'string' is not assignable to type '"x"'.
 
 
-==== tests/cases/conformance/generators/generatorReturnContextualType.ts (1 errors) ====
+==== tests/cases/conformance/generators/generatorReturnContextualType.ts (2 errors) ====
     // #35995
 
     function* f1(): Generator<any, { x: 'x' }, any> {
       return { x: 'x' };
     }
 
+    function* g1(): Iterator<any, { x: 'x' }, any> {
+      return { x: 'x' };
+    }
+
     async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
       return { x: 'x' };
     }
 
+    async function* g2(): AsyncIterator<any, { x: 'x' }, any> {
+      return { x: 'x' };
+    }
+
     async function* f3(): AsyncGenerator<any, { x: 'x' }, any> {
       return Promise.resolve({ x: 'x' });
     }
 
+    async function* g3(): AsyncIterator<any, { x: 'x' }, any> {
+      return Promise.resolve({ x: 'x' });
+    }
+
     async function* f4(): AsyncGenerator<any, { x: 'x' }, any> {
       const ret = { x: 'x' };
       return Promise.resolve(ret); // Error
       ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 !!! error TS2322: Type '{ x: string; }' is not assignable to type '{ x: "x"; }'.
 !!! error TS2322:   Types of property 'x' are incompatible.
+!!! error TS2322:     Type 'string' is not assignable to type '"x"'.
+    }
+
+    async function* g4(): AsyncIterator<any, { x: 'x' }, any> {
+      const ret = { x: 'x' };
+      return Promise.resolve(ret); // Error
+      ~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+!!! error TS2322: Type '{ x: string; }' is not assignable to type '{ x: "x"; }'.
+!!! error TS2322:   Types of property 'x' are incompatible.
 !!! error TS2322:     Type 'string' is not assignable to type '"x"'.
     }
 

tests/baselines/reference/generatorReturnContextualType.js

@@ -5,18 +5,35 @@ function* f1(): Generator<any, { x: 'x' }, any> {
   return { x: 'x' };
 }
 
+function* g1(): Iterator<any, { x: 'x' }, any> {
+  return { x: 'x' };
+}
+
 async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
   return { x: 'x' };
 }
 
+async function* g2(): AsyncIterator<any, { x: 'x' }, any> {
+  return { x: 'x' };
+}
+
 async function* f3(): AsyncGenerator<any, { x: 'x' }, any> {
   return Promise.resolve({ x: 'x' });
 }
 
+async function* g3(): AsyncIterator<any, { x: 'x' }, any> {
+  return Promise.resolve({ x: 'x' });
+}
+
 async function* f4(): AsyncGenerator<any, { x: 'x' }, any> {
   const ret = { x: 'x' };
   return Promise.resolve(ret); // Error
 }
+
+async function* g4(): AsyncIterator<any, { x: 'x' }, any> {
+  const ret = { x: 'x' };
+  return Promise.resolve(ret); // Error
+}
 
 
 //// [generatorReturnContextualType.js]
@@ -25,13 +42,26 @@ async function* f4(): AsyncGenerator<any, { x: 'x' }, any> {
 function* f1() {
     return { x: 'x' };
 }
+function* g1() {
+    return { x: 'x' };
+}
 async function* f2() {
     return { x: 'x' };
 }
+async function* g2() {
+    return { x: 'x' };
+}
 async function* f3() {
     return Promise.resolve({ x: 'x' });
 }
+async function* g3() {
+    return Promise.resolve({ x: 'x' });
+}
 async function* f4() {
     const ret = { x: 'x' };
     return Promise.resolve(ret); // Error
 }
+async function* g4() {
+    const ret = { x: 'x' };
+    return Promise.resolve(ret); // Error
+}

tests/baselines/reference/generatorReturnContextualType.symbols

@@ -10,40 +10,86 @@ function* f1(): Generator<any, { x: 'x' }, any> {
 >x : Symbol(x, Decl(generatorReturnContextualType.ts, 3, 10))
 }
 
+function* g1(): Iterator<any, { x: 'x' }, any> {
+>g1 : Symbol(g1, Decl(generatorReturnContextualType.ts, 4, 1))
+>Iterator : Symbol(Iterator, Decl(lib.es2015.iterable.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 6, 31))
+
+  return { x: 'x' };
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 7, 10))
+}
+
 async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
->f2 : Symbol(f2, Decl(generatorReturnContextualType.ts, 4, 1))
+>f2 : Symbol(f2, Decl(generatorReturnContextualType.ts, 8, 1))
 >AsyncGenerator : Symbol(AsyncGenerator, Decl(lib.es2018.asyncgenerator.d.ts, --, --))
->x : Symbol(x, Decl(generatorReturnContextualType.ts, 6, 43))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 10, 43))
 
   return { x: 'x' };
->x : Symbol(x, Decl(generatorReturnContextualType.ts, 7, 10))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 11, 10))
+}
+
+async function* g2(): AsyncIterator<any, { x: 'x' }, any> {
+>g2 : Symbol(g2, Decl(generatorReturnContextualType.ts, 12, 1))
+>AsyncIterator : Symbol(AsyncIterator, Decl(lib.es2018.asynciterable.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 14, 42))
+
+  return { x: 'x' };
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 15, 10))
 }
 
 async function* f3(): AsyncGenerator<any, { x: 'x' }, any> {
->f3 : Symbol(f3, Decl(generatorReturnContextualType.ts, 8, 1))
+>f3 : Symbol(f3, Decl(generatorReturnContextualType.ts, 16, 1))
 >AsyncGenerator : Symbol(AsyncGenerator, Decl(lib.es2018.asyncgenerator.d.ts, --, --))
->x : Symbol(x, Decl(generatorReturnContextualType.ts, 10, 43))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 18, 43))
+
+  return Promise.resolve({ x: 'x' });
+>Promise.resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
+>Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
+>resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 19, 26))
+}
+
+async function* g3(): AsyncIterator<any, { x: 'x' }, any> {
+>g3 : Symbol(g3, Decl(generatorReturnContextualType.ts, 20, 1))
+>AsyncIterator : Symbol(AsyncIterator, Decl(lib.es2018.asynciterable.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 22, 42))
 
   return Promise.resolve({ x: 'x' });
 >Promise.resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
 >Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
 >resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
->x : Symbol(x, Decl(generatorReturnContextualType.ts, 11, 26))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 23, 26))
 }
 
 async function* f4(): AsyncGenerator<any, { x: 'x' }, any> {
->f4 : Symbol(f4, Decl(generatorReturnContextualType.ts, 12, 1))
+>f4 : Symbol(f4, Decl(generatorReturnContextualType.ts, 24, 1))
 >AsyncGenerator : Symbol(AsyncGenerator, Decl(lib.es2018.asyncgenerator.d.ts, --, --))
->x : Symbol(x, Decl(generatorReturnContextualType.ts, 14, 43))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 26, 43))
+
+  const ret = { x: 'x' };
+>ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 27, 7))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 27, 15))
+
+  return Promise.resolve(ret); // Error
+>Promise.resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
+>Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
+>resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
+>ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 27, 7))
+}
+
+async function* g4(): AsyncIterator<any, { x: 'x' }, any> {
+>g4 : Symbol(g4, Decl(generatorReturnContextualType.ts, 29, 1))
+>AsyncIterator : Symbol(AsyncIterator, Decl(lib.es2018.asynciterable.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 31, 42))
 
   const ret = { x: 'x' };
->ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 15, 7))
->x : Symbol(x, Decl(generatorReturnContextualType.ts, 15, 15))
+>ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 32, 7))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 32, 15))
 
   return Promise.resolve(ret); // Error
 >Promise.resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
 >Promise : Symbol(Promise, Decl(lib.es5.d.ts, --, --), Decl(lib.es2015.iterable.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.symbol.wellknown.d.ts, --, --), Decl(lib.es2018.promise.d.ts, --, --))
 >resolve : Symbol(PromiseConstructor.resolve, Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --), Decl(lib.es2015.promise.d.ts, --, --))
->ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 15, 7))
+>ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 32, 7))
 }