Add contextual type for generator return type (#39772)

* WIP

* Add contextual type for generator return type

Author: andrewbranch

src/compiler/checker.ts

@@ -22977,18 +22977,25 @@ namespace ts {
         function getContextualTypeForReturnExpression(node: Expression): Type | undefined {
             const func = getContainingFunction(node);
             if (func) {
-                const functionFlags = getFunctionFlags(func);
-                if (functionFlags & FunctionFlags.Generator) { // AsyncGenerator function or Generator function
-                    return undefined;
-                }
-
-                const contextualReturnType = getContextualReturnType(func);
+                let contextualReturnType = getContextualReturnType(func);
                 if (contextualReturnType) {
-                    if (functionFlags & FunctionFlags.Async) { // Async function
-                        const contextualAwaitedType = mapType(contextualReturnType, getAwaitedTypeOfPromise);
+                    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) {
+                            return undefined;
+                        }
+                        contextualReturnType = iterationTypes.returnType;
+                        // falls through to unwrap Promise for AsyncGenerators
+                    }
+
+                    if (functionFlags & FunctionFlags.Async) { // Async function or AsyncGenerator function
+                        const contextualAwaitedType = mapType(contextualReturnType, getAwaitedType);
                         return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
                     }
-                    return contextualReturnType; // Regular function
+
+                    return contextualReturnType; // Regular function or Generator function
                 }
             }
             return undefined;
@@ -33439,6 +33446,8 @@ namespace ts {
                         reportTypeNotIterableError(errorNode, type, !!(use & IterationUse.AllowsAsyncIterablesFlag));
                         errorNode = undefined;
                     }
+                    setCachedIterationTypes(type, cacheKey, noIterationTypes);
+                    return undefined;
                 }
                 else {
                     allIterationTypes = append(allIterationTypes, iterationTypes);

tests/baselines/reference/destructuringArrayBindingPatternAndAssignment4.errors.txt

@@ -0,0 +1,12 @@
+tests/cases/conformance/es6/destructuring/destructuringArrayBindingPatternAndAssignment4.ts(5,7): error TS2548: Type 'number[] | null' is not an array type or does not have a '[Symbol.iterator]()' method that returns an iterator.
+
+
+==== tests/cases/conformance/es6/destructuring/destructuringArrayBindingPatternAndAssignment4.ts (1 errors) ====
+    // #35497
+    
+    
+    declare const data: number[] | null;
+    const [value] = data; // Error
+          ~~~~~~~
+!!! error TS2548: Type 'number[] | null' is not an array type or does not have a '[Symbol.iterator]()' method that returns an iterator.
+    

tests/baselines/reference/destructuringArrayBindingPatternAndAssignment4.js

@@ -0,0 +1,28 @@
+//// [destructuringArrayBindingPatternAndAssignment4.ts]
+// #35497
+
+
+declare const data: number[] | null;
+const [value] = data; // Error
+
+
+//// [destructuringArrayBindingPatternAndAssignment4.js]
+"use strict";
+// #35497
+var __read = (this && this.__read) || function (o, n) {
+    var m = typeof Symbol === "function" && o[Symbol.iterator];
+    if (!m) return o;
+    var i = m.call(o), r, ar = [], e;
+    try {
+        while ((n === void 0 || n-- > 0) && !(r = i.next()).done) ar.push(r.value);
+    }
+    catch (error) { e = { error: error }; }
+    finally {
+        try {
+            if (r && !r.done && (m = i["return"])) m.call(i);
+        }
+        finally { if (e) throw e.error; }
+    }
+    return ar;
+};
+var _a = __read(data, 1), value = _a[0]; // Error

tests/baselines/reference/destructuringArrayBindingPatternAndAssignment4.symbols

@@ -0,0 +1,11 @@
+=== tests/cases/conformance/es6/destructuring/destructuringArrayBindingPatternAndAssignment4.ts ===
+// #35497
+
+
+declare const data: number[] | null;
+>data : Symbol(data, Decl(destructuringArrayBindingPatternAndAssignment4.ts, 3, 13))
+
+const [value] = data; // Error
+>value : Symbol(value, Decl(destructuringArrayBindingPatternAndAssignment4.ts, 4, 7))
+>data : Symbol(data, Decl(destructuringArrayBindingPatternAndAssignment4.ts, 3, 13))
+

tests/baselines/reference/destructuringArrayBindingPatternAndAssignment4.types

@@ -0,0 +1,12 @@
+=== tests/cases/conformance/es6/destructuring/destructuringArrayBindingPatternAndAssignment4.ts ===
+// #35497
+
+
+declare const data: number[] | null;
+>data : number[] | null
+>null : null
+
+const [value] = data; // Error
+>value : any
+>data : number[] | null
+

tests/baselines/reference/generatorReturnContextualType.errors.txt

@@ -0,0 +1,29 @@
+tests/cases/conformance/generators/generatorReturnContextualType.ts(17,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) ====
+    // #35995
+    
+    function* f1(): Generator<any, { x: 'x' }, any> {
+      return { x: 'x' };
+    }
+    
+    async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
+      return { x: 'x' };
+    }
+    
+    async function* f3(): AsyncGenerator<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"'.
+    }
+    

tests/baselines/reference/generatorReturnContextualType.js

@@ -0,0 +1,37 @@
+//// [generatorReturnContextualType.ts]
+// #35995
+
+function* f1(): Generator<any, { x: 'x' }, any> {
+  return { x: 'x' };
+}
+
+async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
+  return { x: 'x' };
+}
+
+async function* f3(): AsyncGenerator<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
+}
+
+
+//// [generatorReturnContextualType.js]
+"use strict";
+// #35995
+function* f1() {
+    return { x: 'x' };
+}
+async function* f2() {
+    return { x: 'x' };
+}
+async function* f3() {
+    return Promise.resolve({ x: 'x' });
+}
+async function* f4() {
+    const ret = { x: 'x' };
+    return Promise.resolve(ret); // Error
+}

tests/baselines/reference/generatorReturnContextualType.symbols

@@ -0,0 +1,49 @@
+=== tests/cases/conformance/generators/generatorReturnContextualType.ts ===
+// #35995
+
+function* f1(): Generator<any, { x: 'x' }, any> {
+>f1 : Symbol(f1, Decl(generatorReturnContextualType.ts, 0, 0))
+>Generator : Symbol(Generator, Decl(lib.es2015.generator.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 2, 32))
+
+  return { x: 'x' };
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 3, 10))
+}
+
+async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
+>f2 : Symbol(f2, Decl(generatorReturnContextualType.ts, 4, 1))
+>AsyncGenerator : Symbol(AsyncGenerator, Decl(lib.es2018.asyncgenerator.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 6, 43))
+
+  return { x: 'x' };
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 7, 10))
+}
+
+async function* f3(): AsyncGenerator<any, { x: 'x' }, any> {
+>f3 : Symbol(f3, Decl(generatorReturnContextualType.ts, 8, 1))
+>AsyncGenerator : Symbol(AsyncGenerator, Decl(lib.es2018.asyncgenerator.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 10, 43))
+
+  return Promise.resolve({ x: 'x' });
+>Promise.resolve : Symbol(PromiseConstructor.resolve, 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, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 11, 26))
+}
+
+async function* f4(): AsyncGenerator<any, { x: 'x' }, any> {
+>f4 : Symbol(f4, Decl(generatorReturnContextualType.ts, 12, 1))
+>AsyncGenerator : Symbol(AsyncGenerator, Decl(lib.es2018.asyncgenerator.d.ts, --, --))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 14, 43))
+
+  const ret = { x: 'x' };
+>ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 15, 7))
+>x : Symbol(x, Decl(generatorReturnContextualType.ts, 15, 15))
+
+  return Promise.resolve(ret); // Error
+>Promise.resolve : Symbol(PromiseConstructor.resolve, 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, --, --))
+>ret : Symbol(ret, Decl(generatorReturnContextualType.ts, 15, 7))
+}
+

tests/baselines/reference/generatorReturnContextualType.types

@@ -0,0 +1,55 @@
+=== tests/cases/conformance/generators/generatorReturnContextualType.ts ===
+// #35995
+
+function* f1(): Generator<any, { x: 'x' }, any> {
+>f1 : () => Generator<any, {    x: 'x';}, any>
+>x : "x"
+
+  return { x: 'x' };
+>{ x: 'x' } : { x: "x"; }
+>x : "x"
+>'x' : "x"
+}
+
+async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
+>f2 : () => AsyncGenerator<any, {    x: 'x';}, any>
+>x : "x"
+
+  return { x: 'x' };
+>{ x: 'x' } : { x: "x"; }
+>x : "x"
+>'x' : "x"
+}
+
+async function* f3(): AsyncGenerator<any, { x: 'x' }, any> {
+>f3 : () => AsyncGenerator<any, {    x: 'x';}, any>
+>x : "x"
+
+  return Promise.resolve({ x: 'x' });
+>Promise.resolve({ x: 'x' }) : Promise<{ x: "x"; }>
+>Promise.resolve : { <T>(value: T | PromiseLike<T>): Promise<T>; (): Promise<void>; }
+>Promise : PromiseConstructor
+>resolve : { <T>(value: T | PromiseLike<T>): Promise<T>; (): Promise<void>; }
+>{ x: 'x' } : { x: "x"; }
+>x : "x"
+>'x' : "x"
+}
+
+async function* f4(): AsyncGenerator<any, { x: 'x' }, any> {
+>f4 : () => AsyncGenerator<any, {    x: 'x';}, any>
+>x : "x"
+
+  const ret = { x: 'x' };
+>ret : { x: string; }
+>{ x: 'x' } : { x: string; }
+>x : string
+>'x' : "x"
+
+  return Promise.resolve(ret); // Error
+>Promise.resolve(ret) : Promise<{ x: string; }>
+>Promise.resolve : { <T>(value: T | PromiseLike<T>): Promise<T>; (): Promise<void>; }
+>Promise : PromiseConstructor
+>resolve : { <T>(value: T | PromiseLike<T>): Promise<T>; (): Promise<void>; }
+>ret : { x: string; }
+}
+

tests/baselines/reference/generatorYieldContextualType.types

@@ -25,11 +25,11 @@ declare function f2<T, R, S>(gen: () => Generator<R, T, S> | AsyncGenerator<R, T
 f2<0, 0, 1>(async function* () {
 >f2<0, 0, 1>(async function* () {	const a = yield 0;	return 0;}) : void
 >f2 : <T, R, S>(gen: () => Generator<R, T, S> | AsyncGenerator<R, T, S>) => void
->async function* () {	const a = yield 0;	return 0;} : () => AsyncGenerator<0, 0, 1>
+>async function* () {	const a = yield 0;	return 0;} : () => AsyncGenerator<0, 0, 1 | undefined>
 
 	const a = yield 0;
->a : 1
->yield 0 : 1
+>a : 1 | undefined
+>yield 0 : 1 | undefined
 >0 : 0
 
 	return 0;

tests/cases/conformance/es6/destructuring/destructuringArrayBindingPatternAndAssignment4.ts

@@ -0,0 +1,9 @@
+// #35497
+
+// @target: es5
+// @downlevelIteration: true
+// @lib: es6
+// @strict: true
+
+declare const data: number[] | null;
+const [value] = data; // Error

tests/cases/conformance/generators/generatorReturnContextualType.ts

@@ -0,0 +1,21 @@
+// @target: esnext
+// @strict: true
+
+// #35995
+
+function* f1(): Generator<any, { x: 'x' }, any> {
+  return { x: 'x' };
+}
+
+async function* f2(): AsyncGenerator<any, { x: 'x' }, any> {
+  return { x: 'x' };
+}
+
+async function* f3(): AsyncGenerator<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
+}