microsoft · brandonbloom · Sep 24, 2017
diff --git a/lib/lib.es2015.iterable.d.ts b/lib/lib.es2015.iterable.d.ts
@@ -35,7 +35,7 @@ interface IteratorResult<T> {
 
 interface Iterator<T> {
     next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
+    return?(value?: any): IteratorResult<any>;
     throw?(e?: any): IteratorResult<T>;
 }
 

diff --git a/lib/lib.es2016.full.d.ts b/lib/lib.es2016.full.d.ts
@@ -702,7 +702,7 @@ interface IteratorResult<T> {
 
 interface Iterator<T> {
     next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
+    return?(value?: any): IteratorResult<any>;
     throw?(e?: any): IteratorResult<T>;
 }
 

diff --git a/lib/lib.es2017.full.d.ts b/lib/lib.es2017.full.d.ts
@@ -706,7 +706,7 @@ interface IteratorResult<T> {
 
 interface Iterator<T> {
     next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
+    return?(value?: any): IteratorResult<any>;
     throw?(e?: any): IteratorResult<T>;
 }
 

diff --git a/lib/lib.es6.d.ts b/lib/lib.es6.d.ts
@@ -4763,7 +4763,7 @@ interface IteratorResult<T> {
 
 interface Iterator<T> {
     next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
+    return?(value?: any): IteratorResult<any>;
     throw?(e?: any): IteratorResult<T>;
 }
 

diff --git a/lib/lib.esnext.full.d.ts b/lib/lib.esnext.full.d.ts
@@ -703,7 +703,7 @@ interface IteratorResult<T> {
 
 interface Iterator<T> {
     next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
+    return?(value?: any): IteratorResult<any>;
     throw?(e?: any): IteratorResult<T>;
 }
 

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
@@ -20867,8 +20867,8 @@ namespace ts {
                     return undefined;
                 }
 
-                const returnType = getUnionType(map(signatures, getReturnTypeOfSignature), /*subtypeReduction*/ true);
-                const iteratedType = getIteratedTypeOfIterator(returnType, errorNode, /*isAsyncIterator*/ !!asyncMethodType);
+                const iteratorType = getUnionType(map(signatures, getReturnTypeOfSignature), /*subtypeReduction*/ true);
+                const iteratedType = getIteratedTypeOfIterator(iteratorType, errorNode, /*isAsyncIterator*/ !!asyncMethodType);
                 if (checkAssignability && errorNode && iteratedType) {
                     // If `checkAssignability` was specified, we were called from
                     // `checkIteratedTypeOrElementType`. As such, we need to validate that
@@ -20885,8 +20885,76 @@ namespace ts {
         }
 
         /**
-         * This function has very similar logic as getIteratedTypeOfIterable, except that it operates on
-         * Iterators instead of Iterables. Here is the structure:
+         * This function is to getIteratedTypeOfIterable as
+         * getReturnTypeOfIterator is to getReturnTypeOfIterable.
+         */
+        function getReturnTypeOfIterable(type: Type, errorNode: Node | undefined, allowAsyncIterables: boolean, allowSyncIterables: boolean): Type | undefined {
+            if (isTypeAny(type)) {
+                return undefined;
+            }
+
+            return mapType(type, getReturnType);
+
+            function getReturnType(type: Type) {
+                const typeAsIterable = <IterableOrIteratorType>type;
+                if (allowAsyncIterables) {
+                    if (typeAsIterable.returnTypeOfAsyncIterable) {
+                        return typeAsIterable.returnTypeOfAsyncIterable;
+                    }
+
+                    // As an optimization, if the type is an instantiation of the global `AsyncIterable<T>`
+                    // or the global `AsyncIterableIterator<T>` then just assume the type is any.
+                    if (isReferenceToType(type, getGlobalAsyncIterableType(/*reportErrors*/ false)) ||
+                        isReferenceToType(type, getGlobalAsyncIterableIteratorType(/*reportErrors*/ false))) {
+                        return typeAsIterable.returnTypeOfAsyncIterable = anyType;
+                    }
+                }
+
+                if (allowSyncIterables) {
+                    if (typeAsIterable.returnTypeOfIterable) {
+                        return typeAsIterable.returnTypeOfIterable;
+                    }
+
+                    // As an optimization, if the type is an instantiation of the global `Iterable<T>` or
+                    // `IterableIterator<T>` then just grab its type argument.
+                    if (isReferenceToType(type, getGlobalIterableType(/*reportErrors*/ false)) ||
+                        isReferenceToType(type, getGlobalIterableIteratorType(/*reportErrors*/ false))) {
+                        return typeAsIterable.returnTypeOfIterable = anyType;
+                    }
+                }
+
+                const asyncMethodType = allowAsyncIterables && getTypeOfPropertyOfType(type, getPropertyNameForKnownSymbolName("asyncIterator"));
+                const methodType = asyncMethodType || (allowSyncIterables && getTypeOfPropertyOfType(type, getPropertyNameForKnownSymbolName("iterator")));
+                if (isTypeAny(methodType)) {
+                    return undefined;
+                }
+
+                const signatures = methodType && getSignaturesOfType(methodType, SignatureKind.Call);
+                if (!some(signatures)) {
+                    if (errorNode) {
+                        error(errorNode,
+                            allowAsyncIterables
+                                ? Diagnostics.Type_must_have_a_Symbol_asyncIterator_method_that_returns_an_async_iterator
+                                : Diagnostics.Type_must_have_a_Symbol_iterator_method_that_returns_an_iterator);
+                        // only report on the first error
+                        errorNode = undefined;
+                    }
+                    return undefined;
+                }
+
+                const iteratorType = getUnionType(map(signatures, getReturnTypeOfSignature), /*subtypeReduction*/ true);
+                const returnType = getReturnTypeOfIterator(iteratorType, errorNode, /*isAsyncIterator*/ !!asyncMethodType);
+
+                return asyncMethodType
+                    ? typeAsIterable.returnTypeOfAsyncIterable = returnType
+                    : typeAsIterable.returnTypeOfIterable = returnType;
+            }
+        }
+
+        /**
+         * This function is analogous to getIteratedTypeOfIterator, except that
+         * it gets the result type, which is the argument to the return method.
+         * Here is the structure:
          *
          *  { // iterator
          *      next: { // nextMethod
@@ -20970,6 +21038,88 @@ namespace ts {
                 : typeAsIterator.iteratedTypeOfIterator = nextValue;
         }
 
+        /**
+         * This function is analogous to getIteratedTypeOfIterator, except that
+         * it gets the type from the return method instead of the next method.
+         * Here is the structure:
+         *
+         *  { // iterator
+         *      return: { // nextMethod
+         *          (value: T): { // nextResult
+         *              value: T // nextValue
+         *          }
+         *      }
+         *  }
+         *
+         * For an async iterator, we expect the following structure:
+         *
+         *  { // iterator
+         *      return: { // nextMethod
+         *          (value: T): PromiseLike<{ // nextResult
+         *              value: T // nextValue
+         *          }>
+         *      }
+         *  }
+         */
+        function getReturnTypeOfIterator(type: Type, errorNode: Node | undefined, isAsyncIterator: boolean): Type | undefined {
+            if (isTypeAny(type)) {
+                return undefined;
+            }
+
+            const typeAsIterator = <IterableOrIteratorType>type;
+            if (isAsyncIterator ? typeAsIterator.returnTypeOfAsyncIterator : typeAsIterator.returnTypeOfIterator) {
+                return isAsyncIterator ? typeAsIterator.returnTypeOfAsyncIterator : typeAsIterator.returnTypeOfIterator;
+            }
+
+            // As an optimization, if the type is an instantiation of the global `Iterator<T>` (for
+            // a non-async iterator) or the global `AsyncIterator<T>` (for an async-iterator) then
+            // the return type is any.
+            const getIteratorType = isAsyncIterator ? getGlobalAsyncIteratorType : getGlobalIteratorType;
+            if (isReferenceToType(type, getIteratorType(/*reportErrors*/ false))) {
+                return isAsyncIterator
+                    ? typeAsIterator.returnTypeOfAsyncIterator = anyType
+                    : typeAsIterator.returnTypeOfIterator = anyType;
+            }
+
+            // Both async and non-async iterators may have a `return` method.
+            const returnMethod = getTypeOfPropertyOfType(type, "return" as __String);
+            if (isTypeAny(returnMethod)) {
+                return undefined;
+            }
+
+            const returnMethodSignatures = returnMethod ? getSignaturesOfType(returnMethod, SignatureKind.Call) : emptyArray;
+            if (returnMethodSignatures.length === 0) {
+                return undefined;
+            }
+
+            let returnResult = getUnionType(map(returnMethodSignatures, getReturnTypeOfSignature), /*subtypeReduction*/ true);
+            if (isTypeAny(returnResult)) {
+                return undefined;
+            }
+
+            // For an async iterator, we must get the awaited type of the return type.
+            if (isAsyncIterator) {
+                returnResult = getAwaitedTypeOfPromise(returnResult, errorNode, Diagnostics.The_type_returned_by_the_return_method_of_an_async_iterator_must_be_a_promise_for_a_type_with_a_value_property);
+                if (isTypeAny(returnResult)) {
+                    return undefined;
+                }
+            }
+
+            const returnValue = returnResult && getTypeOfPropertyOfType(returnResult, "value" as __String);
+            if (!returnValue) {
+                if (errorNode) {
+                    error(errorNode, isAsyncIterator
+                        ? Diagnostics.The_type_returned_by_the_return_method_of_an_async_iterator_must_be_a_promise_for_a_type_with_a_value_property
+                        : Diagnostics.The_type_returned_by_the_return_method_of_an_iterator_must_have_a_value_property);
+                }
+                return undefined;
+            }
+
+            return isAsyncIterator
+                ? typeAsIterator.returnTypeOfAsyncIterator = returnValue
+                : typeAsIterator.returnTypeOfIterator = returnValue;
+        }
+
         /**
          * A generator may have a return type of `Iterator<T>`, `Iterable<T>`, or
          * `IterableIterator<T>`. An async generator may have a return type of `AsyncIterator<T>`,
@@ -20985,6 +21135,18 @@ namespace ts {
                 || getIteratedTypeOfIterator(returnType, /*errorNode*/ undefined, isAsyncGenerator);
         }
 
+        /**
+         * Like getIteratedTypeOfGenerator, but consults the iterator's return method.
+         */
+        function getReturnTypeOfGenerator(returnType: Type, isAsyncGenerator: boolean): Type {
+            if (isTypeAny(returnType)) {
+                return undefined;
+            }
+
+            return getReturnTypeOfIterable(returnType, /*errorNode*/ undefined, /*allowAsyncIterables*/ isAsyncGenerator, /*allowSyncIterables*/ !isAsyncGenerator)
+                || getReturnTypeOfIterator(returnType, /*errorNode*/ undefined, isAsyncGenerator);
+        }
+
         function checkBreakOrContinueStatement(node: BreakOrContinueStatement) {
             // Grammar checking
             checkGrammarStatementInAmbientContext(node) || checkGrammarBreakOrContinueStatement(node);
@@ -21021,14 +21183,10 @@ namespace ts {
                     const exprType = node.expression ? checkExpressionCached(node.expression) : undefinedType;
                     const functionFlags = getFunctionFlags(func);
                     if (functionFlags & FunctionFlags.Generator) { // 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 expressions of generators when return type tracking is added
-                        // for generators.
-                        return;
+                        const resultType = getReturnTypeOfGenerator(returnType, (functionFlags & FunctionFlags.Async) !== 0);
+                        checkTypeAssignableTo(exprType, resultType, node);
                     }
-
-                    if (func.kind === SyntaxKind.SetAccessor) {
+                    else if (func.kind === SyntaxKind.SetAccessor) {
                         if (node.expression) {
                             error(node, Diagnostics.Setters_cannot_return_a_value);
                         }

diff --git a/src/compiler/diagnosticMessages.json b/src/compiler/diagnosticMessages.json
@@ -2216,6 +2216,14 @@
         "category": "Error",
         "code": 2714
     },
+    "The type returned by the 'return()' method of an iterator must have a 'value' property.": {
+        "category": "Error",
+        "code": 2715
+    },
+    "The type returned by the 'return()' method of an async iterator must be a promise for a type with a 'value' property.": {
+        "category": "Error",
+        "code": 2716
+    },
 
     "Import declaration '{0}' is using private name '{1}'.": {
         "category": "Error",

diff --git a/src/compiler/types.ts b/src/compiler/types.ts
@@ -3419,6 +3419,10 @@ namespace ts {
         iteratedTypeOfIterator?: Type;
         iteratedTypeOfAsyncIterable?: Type;
         iteratedTypeOfAsyncIterator?: Type;
+        returnTypeOfIterable?: Type;
+        returnTypeOfIterator?: Type;
+        returnTypeOfAsyncIterable?: Type;
+        returnTypeOfAsyncIterator?: Type;
     }
 
     /* @internal */

diff --git a/src/lib/es2015.iterable.d.ts b/src/lib/es2015.iterable.d.ts
@@ -15,7 +15,7 @@ interface IteratorResult<T> {
 
 interface Iterator<T> {
     next(value?: any): IteratorResult<T>;
-    return?(value?: any): IteratorResult<T>;
+    return?(value?: any): IteratorResult<any>;
     throw?(e?: any): IteratorResult<T>;
 }
 

diff --git a/src/lib/esnext.asynciterable.d.ts b/src/lib/esnext.asynciterable.d.ts
@@ -11,7 +11,7 @@ interface SymbolConstructor {
 
 interface AsyncIterator<T> {
     next(value?: any): Promise<IteratorResult<T>>;
-    return?(value?: any): Promise<IteratorResult<T>>;
+    return?(value?: any): Promise<IteratorResult<any>>;
     throw?(e?: any): Promise<IteratorResult<T>>;
 }
 
@@ -21,4 +21,4 @@ interface AsyncIterable<T> {
 
 interface AsyncIterableIterator<T> extends AsyncIterator<T> {
     [Symbol.asyncIterator](): AsyncIterableIterator<T>;
-}
+}
diff --git a/tests/baselines/reference/for-of30.errors.txt b/tests/baselines/reference/for-of30.errors.txt
@@ -3,7 +3,7 @@ tests/cases/conformance/es6/for-ofStatements/for-of30.ts(16,15): error TS2322: T
     Type '() => StringIterator' is not assignable to type '() => Iterator<string>'.
       Type 'StringIterator' is not assignable to type 'Iterator<string>'.
         Types of property 'return' are incompatible.
-          Type 'number' is not assignable to type '(value?: any) => IteratorResult<string>'.
+          Type 'number' is not assignable to type '(value?: any) => IteratorResult<any>'.
 
 
 ==== tests/cases/conformance/es6/for-ofStatements/for-of30.ts (1 errors) ====
@@ -29,4 +29,4 @@ tests/cases/conformance/es6/for-ofStatements/for-of30.ts(16,15): error TS2322: T
 !!! error TS2322:     Type '() => StringIterator' is not assignable to type '() => Iterator<string>'.
 !!! error TS2322:       Type 'StringIterator' is not assignable to type 'Iterator<string>'.
 !!! error TS2322:         Types of property 'return' are incompatible.
-!!! error TS2322:           Type 'number' is not assignable to type '(value?: any) => IteratorResult<string>'.
+!!! error TS2322:           Type 'number' is not assignable to type '(value?: any) => IteratorResult<any>'.
diff --git a/tests/baselines/reference/iteratorReturn.errors.txt b/tests/baselines/reference/iteratorReturn.errors.txt
@@ -0,0 +1,27 @@
+tests/cases/compiler/iteratorReturn.ts(15,11): error TS2322: Type '"str"' is not assignable to type 'number'.
+tests/cases/compiler/iteratorReturn.ts(16,5): error TS2322: Type '1' is not assignable to type 'string'.
+
+
+==== tests/cases/compiler/iteratorReturn.ts (2 errors) ====
+    interface Coroutine extends Iterator<number> {
+        return?(effect?: string): IteratorResult<string>;
+    }
+
+    interface Process extends Coroutine {
+        [Symbol.iterator](): Coroutine;
+    }
+
+    let good = function*(): Process {
+        yield 1;
+        return "str";
+    };
+
+    let bad = function*(): Process {
+        yield "str";
+              ~~~~~
+!!! error TS2322: Type '"str"' is not assignable to type 'number'.
+        return 1;
+        ~~~~~~~~~
+!!! error TS2322: Type '1' is not assignable to type 'string'.
+    };
+
diff --git a/tests/baselines/reference/iteratorReturn.js b/tests/baselines/reference/iteratorReturn.js
@@ -0,0 +1,29 @@
+//// [iteratorReturn.ts]
+interface Coroutine extends Iterator<number> {
+    return?(effect?: string): IteratorResult<string>;
+}
+
+interface Process extends Coroutine {
+    [Symbol.iterator](): Coroutine;
+}
+
+let good = function*(): Process {
+    yield 1;
+    return "str";
+};
+
+let bad = function*(): Process {
+    yield "str";
+    return 1;
+};
+
+
+//// [iteratorReturn.js]
+let good = function* () {
+    yield 1;
+    return "str";
+};
+let bad = function* () {
+    yield "str";
+    return 1;
+};