Improve type checking and inference for Generators and Async Generators (microsoft#30790)

* Improve typing for Generators and Async Generators

* Add TReturn and TNext to Iterator, IterableIterator, etc.

* Update ts internal Iterator to be assignable from global Iterator

* Make 'done' optional in IteratorYieldResult

* Revert Iterable and IterableIterator to simpler versions plus other fixes

* Add additional inference tests

* Added additional tests

* PR cleanup and minor async iteration type fix

* Updated diagnostics message and added non-strict tests

* Fix expected arity of Iterator/AsyncIterator

Author: rbuckton

src/compiler/builderState.ts

@@ -427,8 +427,8 @@ namespace ts.BuilderState {
                 const references = state.referencedMap.get(path);
                 if (references) {
                     const iterator = references.keys();
-                    for (let { value, done } = iterator.next(); !done; { value, done } = iterator.next()) {
-                        queue.push(value as Path);
+                    for (let iterResult = iterator.next(); !iterResult.done; iterResult = iterator.next()) {
+                        queue.push(iterResult.value as Path);
                     }
                 }
             }

src/compiler/checker.ts

@@ -40,6 +40,7 @@ namespace ts {
         ["es2017.string", "lib.es2017.string.d.ts"],
         ["es2017.intl", "lib.es2017.intl.d.ts"],
         ["es2017.typedarrays", "lib.es2017.typedarrays.d.ts"],
+        ["es2018.asyncgenerator", "lib.es2018.asyncgenerator.d.ts"],
         ["es2018.asynciterable", "lib.es2018.asynciterable.d.ts"],
         ["es2018.intl", "lib.es2018.intl.d.ts"],
         ["es2018.promise", "lib.es2018.promise.d.ts"],
@@ -1899,7 +1900,9 @@ namespace ts {
                 case "object":
                     return {};
                 default:
-                    return option.type.keys().next().value;
+                    const iterResult = option.type.keys().next();
+                    if (!iterResult.done) return iterResult.value;
+                    return Debug.fail("Expected 'option.type' to have entries.");
             }
         }
 

src/compiler/commandLineParser.ts

@@ -24,7 +24,6 @@ namespace ts {
         " __sortedArrayBrand": any;
     }
 
-
     /** ES6 Map interface, only read methods included. */
     export interface ReadonlyMap<T> {
         get(key: string): T | undefined;
@@ -45,7 +44,7 @@ namespace ts {
 
     /** ES6 Iterator type. */
     export interface Iterator<T> {
-        next(): { value: T, done: false } | { value: never, done: true };
+        next(): { value: T, done?: false } | { value: never, done: true };
     }
 
     /** Array that is only intended to be pushed to, never read. */
@@ -297,12 +296,13 @@ namespace ts {
             forEach(action: (value: T, key: string) => void): void {
                 const iterator = this.entries();
                 while (true) {
-                    const { value: entry, done } = iterator.next();
-                    if (done) {
+                    const iterResult = iterator.next();
+                    if (iterResult.done) {
                         break;
                     }
 
-                    action(entry[1], entry[0]);
+                    const [key, value] = iterResult.value;
+                    action(value, key);
                 }
             }
         };
@@ -346,11 +346,11 @@ namespace ts {
 
     export function firstDefinedIterator<T, U>(iter: Iterator<T>, callback: (element: T) => U | undefined): U | undefined {
         while (true) {
-            const { value, done } = iter.next();
-            if (done) {
+            const iterResult = iter.next();
+            if (iterResult.done) {
                 return undefined;
             }
-            const result = callback(value);
+            const result = callback(iterResult.value);
             if (result !== undefined) {
                 return result;
             }
@@ -375,7 +375,7 @@ namespace ts {
                     return { value: undefined as never, done: true };
                 }
                 i++;
-                return { value: [arrayA[i - 1], arrayB[i - 1]], done: false };
+                return { value: [arrayA[i - 1], arrayB[i - 1]] as [T, U], done: false };
             }
         };
     }
@@ -567,7 +567,7 @@ namespace ts {
         return {
             next() {
                 const iterRes = iter.next();
-                return iterRes.done ? iterRes : { value: mapFn(iterRes.value), done: false };
+                return iterRes.done ? iterRes as { done: true, value: never } : { value: mapFn(iterRes.value), done: false };
             }
         };
     }
@@ -678,7 +678,7 @@ namespace ts {
                     }
                     const iterRes = iter.next();
                     if (iterRes.done) {
-                        return iterRes;
+                        return iterRes as { done: true, value: never };
                     }
                     currentIter = getIterator(iterRes.value);
                 }
@@ -753,7 +753,7 @@ namespace ts {
                 while (true) {
                     const res = iter.next();
                     if (res.done) {
-                        return res;
+                        return res as { done: true, value: never };
                     }
                     const value = mapFn(res.value);
                     if (value !== undefined) {
@@ -1078,6 +1078,7 @@ namespace ts {
      * @param value The value to append to the array. If `value` is `undefined`, nothing is
      * appended.
      */
+    export function append<TArray extends any[] | undefined, TValue extends NonNullable<TArray>[number] | undefined>(to: TArray, value: TValue): [undefined, undefined] extends [TArray, TValue] ? TArray : NonNullable<TArray>[number][];
     export function append<T>(to: T[], value: T | undefined): T[];
     export function append<T>(to: T[] | undefined, value: T): T[];
     export function append<T>(to: T[] | undefined, value: T | undefined): T[] | undefined;
@@ -1402,8 +1403,8 @@ namespace ts {
     export function arrayFrom<T>(iterator: Iterator<T> | IterableIterator<T>): T[];
     export function arrayFrom<T, U>(iterator: Iterator<T> | IterableIterator<T>, map?: (t: T) => U): (T | U)[] {
         const result: (T | U)[] = [];
-        for (let { value, done } = iterator.next(); !done; { value, done } = iterator.next()) {
-            result.push(map ? map(value) : value);
+        for (let iterResult = iterator.next(); !iterResult.done; iterResult = iterator.next()) {
+            result.push(map ? map(iterResult.value) : iterResult.value);
         }
         return result;
     }

src/compiler/core.ts

@@ -1764,7 +1764,7 @@
         "category": "Error",
         "code": 2489
     },
-    "The type returned by the 'next()' method of an iterator must have a 'value' property.": {
+    "The type returned by the '{0}()' method of an iterator must have a 'value' property.": {
         "category": "Error",
         "code": 2490
     },
@@ -1992,7 +1992,7 @@
         "category": "Error",
         "code": 2546
     },
-    "The type returned by the 'next()' method of an async iterator must be a promise for a type with a 'value' property.": {
+    "The type returned by the '{0}()' method of an async iterator must be a promise for a type with a 'value' property.": {
         "category": "Error",
         "code": 2547
     },
@@ -2653,6 +2653,30 @@
         "category": "Error",
         "code": 2762
     },
+    "Cannot iterate value because the 'next' method of its iterator expects type '{1}', but for-of will always send '{0}'.": {
+        "category": "Error",
+        "code": 2763
+    },
+    "Cannot iterate value because the 'next' method of its iterator expects type '{1}', but array spread will always send '{0}'.": {
+        "category": "Error",
+        "code": 2764
+    },
+    "Cannot iterate value because the 'next' method of its iterator expects type '{1}', but array destructuring will always send '{0}'.": {
+        "category": "Error",
+        "code": 2765
+    },
+    "Cannot delegate iteration to value because the 'next' method of its iterator expects type '{1}', but the containing generator will always send '{0}'.": {
+        "category": "Error",
+        "code": 2766
+    },
+    "The '{0}' property of an iterator must be a method.": {
+        "category": "Error",
+        "code": 2767
+    },
+    "The '{0}' property of an async iterator must be a method.": {
+        "category": "Error",
+        "code": 2768
+    },
 
     "Import declaration '{0}' is using private name '{1}'.": {
         "category": "Error",
@@ -4214,7 +4238,7 @@
         "category": "Error",
         "code": 7024
     },
-    "Generator implicitly has type '{0}' because it does not yield any values. Consider supplying a return type.": {
+    "Generator implicitly has yield type '{0}' because it does not yield any values. Consider supplying a return type annotation.": {
         "category": "Error",
         "code": 7025
     },
@@ -4336,6 +4360,11 @@
         "category": "Error",
         "code": 7054
     },
+    "'{0}', which lacks return-type annotation, implicitly has an '{1}' yield type.": {
+        "category": "Error",
+        "code": 7055
+    },
+
     "You cannot rename this element.": {
         "category": "Error",
         "code": 8000

src/compiler/diagnosticMessages.json

@@ -148,7 +148,8 @@ namespace ts {
             const sourceIndexToNewSourceIndexMap: number[] = [];
             let nameIndexToNewNameIndexMap: number[] | undefined;
             const mappingIterator = decodeMappings(map.mappings);
-            for (let { value: raw, done } = mappingIterator.next(); !done; { value: raw, done } = mappingIterator.next()) {
+            for (let iterResult = mappingIterator.next(); !iterResult.done; iterResult = mappingIterator.next()) {
+                const raw = iterResult.value;
                 if (end && (
                     raw.generatedLine > end.line ||
                     (raw.generatedLine === end.line && raw.generatedCharacter > end.character))) {

src/compiler/sourcemap.ts

@@ -4273,13 +4273,23 @@ namespace ts {
         regularType: ResolvedType;  // Regular version of fresh type
     }
 
+    /* @internal */
+    export interface IterationTypes {
+        readonly yieldType: Type;
+        readonly returnType: Type;
+        readonly nextType: Type;
+    }
+
     // Just a place to cache element types of iterables and iterators
     /* @internal */
     export interface IterableOrIteratorType extends ObjectType, UnionType {
-        iteratedTypeOfIterable?: Type;
-        iteratedTypeOfIterator?: Type;
-        iteratedTypeOfAsyncIterable?: Type;
-        iteratedTypeOfAsyncIterator?: Type;
+        iterationTypesOfGeneratorReturnType?: IterationTypes;
+        iterationTypesOfAsyncGeneratorReturnType?: IterationTypes;
+        iterationTypesOfIterable?: IterationTypes;
+        iterationTypesOfIterator?: IterationTypes;
+        iterationTypesOfAsyncIterable?: IterationTypes;
+        iterationTypesOfAsyncIterator?: IterationTypes;
+        iterationTypesOfIteratorResult?: IterationTypes;
     }
 
     /* @internal */

src/compiler/types.ts

@@ -150,8 +150,8 @@ namespace ts {
     export function forEachEntry<T, U>(map: ReadonlyMap<T>, callback: (value: T, key: string) => U | undefined): U | undefined;
     export function forEachEntry<T, U>(map: ReadonlyUnderscoreEscapedMap<T> | ReadonlyMap<T>, callback: (value: T, key: (string & __String)) => U | undefined): U | undefined {
         const iterator = map.entries();
-        for (let { value: pair, done } = iterator.next(); !done; { value: pair, done } = iterator.next()) {
-            const [key, value] = pair;
+        for (let iterResult = iterator.next(); !iterResult.done; iterResult = iterator.next()) {
+            const [key, value] = iterResult.value;
             const result = callback(value, key as (string & __String));
             if (result) {
                 return result;
@@ -165,8 +165,8 @@ namespace ts {
     export function forEachKey<T>(map: ReadonlyMap<{}>, callback: (key: string) => T | undefined): T | undefined;
     export function forEachKey<T>(map: ReadonlyUnderscoreEscapedMap<{}> | ReadonlyMap<{}>, callback: (key: string & __String) => T | undefined): T | undefined {
         const iterator = map.keys();
-        for (let { value: key, done } = iterator.next(); !done; { value: key, done } = iterator.next()) {
-            const result = callback(key as string & __String);
+        for (let iterResult = iterator.next(); !iterResult.done; iterResult = iterator.next()) {
+            const result = callback(iterResult.value as string & __String);
             if (result) {
                 return result;
             }

src/compiler/utilities.ts

@@ -301,7 +301,8 @@ namespace Harness.SourceMapRecorder {
 
             SourceMapSpanWriter.initializeSourceMapSpanWriter(sourceMapRecorder, sourceMapData.sourceMap, currentFile);
             const mapper = ts.decodeMappings(sourceMapData.sourceMap.mappings);
-            for (let { value: decodedSourceMapping, done } = mapper.next(); !done; { value: decodedSourceMapping, done } = mapper.next()) {
+            for (let iterResult = mapper.next(); !iterResult.done; iterResult = mapper.next()) {
+                const decodedSourceMapping = iterResult.value;
                 const currentSourceFile = ts.isSourceMapping(decodedSourceMapping)
                     ? program.getSourceFile(sourceMapData.inputSourceFileNames[decodedSourceMapping.sourceIndex])
                     : undefined;
@@ -335,7 +336,8 @@ namespace Harness.SourceMapRecorder {
 
             SourceMapSpanWriter.initializeSourceMapSpanWriter(sourceMapRecorder, sourceMap, currentFile);
             const mapper = ts.decodeMappings(sourceMap.mappings);
-            for (let { value: decodedSourceMapping, done } = mapper.next(); !done; { value: decodedSourceMapping, done } = mapper.next()) {
+            for (let iterResult = mapper.next(); !iterResult.done; iterResult = mapper.next()) {
+                const decodedSourceMapping = iterResult.value;
                 const currentSourceFile = ts.isSourceMapping(decodedSourceMapping)
                     ? getFile(sourceFileAbsolutePaths[decodedSourceMapping.sourceIndex])
                     : undefined;

src/harness/sourceMapRecorder.ts

@@ -683,7 +683,7 @@ namespace vfs {
 
             if (isDirectory(node)) throw createIOError("EISDIR");
             if (!isFile(node)) throw createIOError("EBADF");
-            node.buffer = Buffer.isBuffer(data) ? data.slice() : ts.sys.bufferFrom!("" + data, encoding || "utf8");
+            node.buffer = Buffer.isBuffer(data) ? data.slice() : ts.sys.bufferFrom!("" + data, encoding || "utf8") as Buffer;
             node.size = node.buffer.byteLength;
             node.mtimeMs = time;
             node.ctimeMs = time;
@@ -1204,7 +1204,7 @@ namespace vfs {
                 }
             },
             readFileSync(path: string): Buffer {
-                return ts.sys.bufferFrom!(host.readFile(path)!, "utf8"); // TODO: GH#18217
+                return ts.sys.bufferFrom!(host.readFile(path)!, "utf8") as Buffer; // TODO: GH#18217
             }
         };
     }

src/harness/vfs.ts

@@ -1,4 +1,12 @@
-interface Generator extends Iterator<any> { }
+/// <reference lib="es2015.iterable" />
+
+interface Generator<T = unknown, TReturn = any, TNext = unknown> extends Iterator<T, TReturn, TNext> {
+    // NOTE: 'next' is defined using a tuple to ensure we report the correct assignability errors in all places.
+    next(...args: [] | [TNext]): IteratorResult<T, TReturn>;
+    return(value: TReturn): IteratorResult<T, TReturn>;
+    throw(e: any): IteratorResult<T, TReturn>;
+    [Symbol.iterator](): Generator<T, TReturn, TNext>;
+}
 
 interface GeneratorFunction {
     /**

src/lib/es2015.generator.d.ts

@@ -8,15 +8,23 @@ interface SymbolConstructor {
     readonly iterator: symbol;
 }
 
-interface IteratorResult<T> {
-    done: boolean;
-    value: T;
+interface IteratorYieldResult<TYield> {
+    done?: false;
+    value: TYield;
 }
 
-interface Iterator<T> {
-    next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
-    throw?(e?: any): IteratorResult<T>;
+interface IteratorReturnResult<TReturn> {
+    done: true;
+    value: TReturn;
+}
+
+type IteratorResult<T, TReturn = any> = IteratorYieldResult<T> | IteratorReturnResult<TReturn>;
+
+interface Iterator<T, TReturn = any, TNext = undefined> {
+    // NOTE: 'next' is defined using a tuple to ensure we report the correct assignability errors in all places.
+    next(...args: [] | [TNext]): IteratorResult<T, TReturn>;
+    return?(value?: TReturn): IteratorResult<T, TReturn>;
+    throw?(e?: any): IteratorResult<T, TReturn>;
 }
 
 interface Iterable<T> {