Fix async generator return as well

Author: rbuckton

src/compiler/binder.ts

@@ -3712,6 +3712,10 @@ namespace ts {
                 break;
 
             case SyntaxKind.ReturnStatement:
+                // Return statements may require an `await` in ESNext.
+                transformFlags |= TransformFlags.ContainsHoistedDeclarationOrCompletion | TransformFlags.AssertESNext;
+                break;
+
             case SyntaxKind.ContinueStatement:
             case SyntaxKind.BreakStatement:
                 transformFlags |= TransformFlags.ContainsHoistedDeclarationOrCompletion;

src/compiler/transformers/esnext.ts

@@ -63,6 +63,8 @@ namespace ts {
                     return visitAwaitExpression(node as AwaitExpression);
                 case SyntaxKind.YieldExpression:
                     return visitYieldExpression(node as YieldExpression);
+                case SyntaxKind.ReturnStatement:
+                    return visitReturnStatement(node as ReturnStatement);
                 case SyntaxKind.LabeledStatement:
                     return visitLabeledStatement(node as LabeledStatement);
                 case SyntaxKind.ObjectLiteralExpression:
@@ -161,6 +163,16 @@ namespace ts {
             return visitEachChild(node, visitor, context);
         }
 
+        function visitReturnStatement(node: ReturnStatement) {
+            if (enclosingFunctionFlags & FunctionFlags.Async && enclosingFunctionFlags & FunctionFlags.Generator) {
+                return updateReturn(node, createDownlevelAwait(
+                    node.expression ? visitNode(node.expression, visitor, isExpression) : createVoidZero()
+                ));
+            }
+
+            return visitEachChild(node, visitor, context);
+        }
+
         function visitLabeledStatement(node: LabeledStatement) {
             if (enclosingFunctionFlags & FunctionFlags.Async) {
                 const statement = unwrapInnermostStatementOfLabel(node);

src/harness/evaluator.ts

@@ -0,0 +1,55 @@
+namespace evaluator {
+    declare var Symbol: SymbolConstructor;
+
+    const sourceFile = vpath.combine(vfs.srcFolder, "source.ts");
+
+    function compile(sourceText: string, options?: ts.CompilerOptions) {
+        const fs = vfs.createFromFileSystem(Harness.IO, /*ignoreCase*/ false);
+        fs.writeFileSync(sourceFile, sourceText);
+        const compilerOptions: ts.CompilerOptions = { target: ts.ScriptTarget.ES5, module: ts.ModuleKind.CommonJS, lib: ["lib.esnext.d.ts"], ...options };
+        const host = new fakes.CompilerHost(fs, compilerOptions);
+        return compiler.compileFiles(host, [sourceFile], compilerOptions);
+    }
+
+    function noRequire(id: string) {
+        throw new Error(`Module '${id}' could not be found.`);
+    }
+
+    // Define a custom "Symbol" constructor to attach missing built-in symbols without
+    // modifying the global "Symbol" constructor
+    // tslint:disable-next-line:variable-name
+    const FakeSymbol: SymbolConstructor = ((description?: string) => Symbol(description)) as any;
+    (<any>FakeSymbol).prototype = Symbol.prototype;
+    for (const key of Object.getOwnPropertyNames(Symbol)) {
+        Object.defineProperty(FakeSymbol, key, Object.getOwnPropertyDescriptor(Symbol, key)!);
+    }
+
+    // Add "asyncIterator" if missing
+    if (!ts.hasProperty(FakeSymbol, "asyncIterator")) Object.defineProperty(FakeSymbol, "asyncIterator", { value: Symbol.for("Symbol.asyncIterator"), configurable: true });
+
+    function evaluate(result: compiler.CompilationResult, globals?: Record<string, any>) {
+        globals = { Symbol: FakeSymbol, ...globals };
+
+        const output = result.getOutput(sourceFile, "js")!;
+        assert.isDefined(output);
+
+        const globalNames: string[] = [];
+        const globalArgs: any[] = [];
+        for (const name in globals) {
+            if (ts.hasProperty(globals, name)) {
+                globalNames.push(name);
+                globalArgs.push(globals[name]);
+            }
+        }
+
+        const evaluateText = `(function (module, exports, require, __dirname, __filename, ${globalNames.join(", ")}) { ${output.text} })`;
+        const evaluateThunk = eval(evaluateText) as (module: any, exports: any, require: (id: string) => any, dirname: string, filename: string, ...globalArgs: any[]) => void;
+        const module: { exports: any; } = { exports: {} };
+        evaluateThunk.call(globals, module, module.exports, noRequire, vpath.dirname(output.file), output.file, FakeSymbol, ...globalArgs);
+        return module.exports;
+    }
+
+    export function evaluateTypeScript(sourceText: string, options?: ts.CompilerOptions, globals?: Record<string, any>) {
+        return evaluate(compile(sourceText, options), globals);
+    }
+}

src/harness/tsconfig.json

@@ -148,6 +148,7 @@
         "vpath.ts",
         "vfs.ts",
         "compiler.ts",
+        "evaluator.ts",
         "fakes.ts",
 
         "sourceMapRecorder.ts",
@@ -169,5 +170,5 @@
         "parallel/shared.ts",
         "runner.ts"
     ],
-    "include": ["unittests/**.ts"]
+    "include": ["unittests/**/*.ts"]
 }

src/harness/unittests/evaluation/asyncGenerator.ts

@@ -0,0 +1,30 @@
+describe("asyncGeneratorEvaluation", () => {
+    it("return (es5)", async () => {
+        const result = evaluator.evaluateTypeScript(`
+        async function * g() {
+            return Promise.resolve(0);
+        }
+        export const output: any[] = [];
+        export async function main() {
+            output.push(await g().next());
+        }`);
+        await result.main();
+        assert.deepEqual(result.output, [
+            { value: 0, done: true }
+        ]);
+    });
+    it("return (es2015)", async () => {
+        const result = evaluator.evaluateTypeScript(`
+        async function * g() {
+            return Promise.resolve(0);
+        }
+        export const output: any[] = [];
+        export async function main() {
+            output.push(await g().next());
+        }`, { target: ts.ScriptTarget.ES2015 });
+        await result.main();
+        assert.deepEqual(result.output, [
+            { value: 0, done: true }
+        ]);
+    });
+});

src/harness/unittests/evaluation/forAwaitOf.ts

@@ -0,0 +1,105 @@
+describe("forAwaitOfEvaluation", () => {
+    it("sync (es5)", async () => {
+        const result = evaluator.evaluateTypeScript(`
+        let i = 0;
+        const iterator = {
+            [Symbol.iterator]() { return this; },
+            next() {
+                switch (i++) {
+                    case 0: return { value: 1, done: false };
+                    case 1: return { value: Promise.resolve(2), done: false };
+                    case 2: return { value: new Promise<number>(resolve => setTimeout(resolve, 100, 3)), done: false };
+                    default: return { value: undefined: done: true };
+                }
+            }
+        };
+        export const output: any[] = [];
+        export async function main() {
+            for await (const item of iterator) {
+                output.push(item);
+            }
+        }`);
+        await result.main();
+        assert.strictEqual(result.output[0], 1);
+        assert.strictEqual(result.output[1], 2);
+        assert.strictEqual(result.output[2], 3);
+    });
+
+    it("sync (es2015)", async () => {
+        const result = evaluator.evaluateTypeScript(`
+        let i = 0;
+        const iterator = {
+            [Symbol.iterator]() { return this; },
+            next() {
+                switch (i++) {
+                    case 0: return { value: 1, done: false };
+                    case 1: return { value: Promise.resolve(2), done: false };
+                    case 2: return { value: new Promise<number>(resolve => setTimeout(resolve, 100, 3)), done: false };
+                    default: return { value: undefined: done: true };
+                }
+            }
+        };
+        export const output: any[] = [];
+        export async function main() {
+            for await (const item of iterator) {
+                output.push(item);
+            }
+        }`, { target: ts.ScriptTarget.ES2015 });
+        await result.main();
+        assert.strictEqual(result.output[0], 1);
+        assert.strictEqual(result.output[1], 2);
+        assert.strictEqual(result.output[2], 3);
+    });
+
+    it("async (es5)", async () => {
+        const result = evaluator.evaluateTypeScript(`
+        let i = 0;
+        const iterator = {
+            [Symbol.asyncIterator]() { return this; },
+            async next() {
+                switch (i++) {
+                    case 0: return { value: 1, done: false };
+                    case 1: return { value: Promise.resolve(2), done: false };
+                    case 2: return { value: new Promise<number>(resolve => setTimeout(resolve, 100, 3)), done: false };
+                    default: return { value: undefined: done: true };
+                }
+            }
+        };
+        export const output: any[] = [];
+        export async function main() {
+            for await (const item of iterator) {
+                output.push(item);
+            }
+        }`);
+        await result.main();
+        assert.strictEqual(result.output[0], 1);
+        assert.instanceOf(result.output[1], Promise);
+        assert.instanceOf(result.output[2], Promise);
+    });
+
+    it("async (es2015)", async () => {
+        const result = evaluator.evaluateTypeScript(`
+        let i = 0;
+        const iterator = {
+            [Symbol.asyncIterator]() { return this; },
+            async next() {
+                switch (i++) {
+                    case 0: return { value: 1, done: false };
+                    case 1: return { value: Promise.resolve(2), done: false };
+                    case 2: return { value: new Promise<number>(resolve => setTimeout(resolve, 100, 3)), done: false };
+                    default: return { value: undefined: done: true };
+                }
+            }
+        };
+        export const output: any[] = [];
+        export async function main() {
+            for await (const item of iterator) {
+                output.push(item);
+            }
+        }`, { target: ts.ScriptTarget.ES2015 });
+        await result.main();
+        assert.strictEqual(result.output[0], 1);
+        assert.instanceOf(result.output[1], Promise);
+        assert.instanceOf(result.output[2], Promise);
+    });
+});