From 63adc5fb4035d5535953a490827c2fa7441c7ee5 Mon Sep 17 00:00:00 2001
From: Ron Buckton <ron.buckton@microsoft.com>
Date: Fri, 21 Sep 2018 10:17:18 -0700
Subject: [PATCH] Add contextual typing for await operand

---
 src/compiler/checker.ts                       | 13 +++++++-
 ...xtuallyTypeAsyncFunctionReturnType.symbols | 25 ++++++++++++++++
 ...textuallyTypeAsyncFunctionReturnType.types | 30 +++++++++++++++++++
 ...contextuallyTypeAsyncFunctionReturnType.ts | 10 +++++++
 4 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/src/compiler/checker.ts b/src/compiler/checker.ts
index fde33463942..8e448a10017 100644
--- a/src/compiler/checker.ts
+++ b/src/compiler/checker.ts
@@ -16409,6 +16409,15 @@ namespace ts {
             return undefined;
         }
 
+        function getContextualTypeForAwaitOperand(node: AwaitExpression): Type | undefined {
+            const contextualType = getContextualType(node);
+            if (contextualType) {
+                const contextualAwaitedType = getAwaitedType(contextualType);
+                return contextualAwaitedType && getUnionType([contextualAwaitedType, createPromiseLikeType(contextualAwaitedType)]);
+            }
+            return undefined;
+        }
+
         function getContextualTypeForYieldOperand(node: YieldExpression): Type | undefined {
             const func = getContainingFunction(node);
             if (func) {
@@ -16770,7 +16779,9 @@ namespace ts {
                     return getContextualTypeForReturnExpression(node);
                 case SyntaxKind.YieldExpression:
                     return getContextualTypeForYieldOperand(<YieldExpression>parent);
-                    case SyntaxKind.CallExpression:
+                case SyntaxKind.AwaitExpression:
+                    return getContextualTypeForAwaitOperand(<AwaitExpression>parent);
+                case SyntaxKind.CallExpression:
                 case SyntaxKind.NewExpression:
                     return getContextualTypeForArgument(<CallExpression | NewExpression>parent, node);
                 case SyntaxKind.TypeAssertionExpression:
diff --git a/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.symbols b/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.symbols
index 28f80dfa1a7..e85fa4d0d9a 100644
--- a/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.symbols
+++ b/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.symbols
@@ -27,3 +27,28 @@ async function fn2(): Promise<Obj> {
 
     });
 }
+
+async function fn3(): Promise<Obj> {
+>fn3 : Symbol(fn3, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 10, 1))
+>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, --, --))
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 0))
+
+    return await { key: "value" };
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 13, 18))
+}
+
+async function fn4(): Promise<Obj> {
+>fn4 : Symbol(fn4, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 14, 1))
+>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, --, --))
+>Obj : Symbol(Obj, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 0, 0))
+
+    return await new Promise(resolve => {
+>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(resolve, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 17, 29))
+
+        resolve({ key: "value" });
+>resolve : Symbol(resolve, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 17, 29))
+>key : Symbol(key, Decl(contextuallyTypeAsyncFunctionReturnType.ts, 18, 17))
+
+    });
+}
diff --git a/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.types b/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.types
index bcb58991c4d..d47a623ffc6 100644
--- a/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.types
+++ b/tests/baselines/reference/contextuallyTypeAsyncFunctionReturnType.types
@@ -29,3 +29,33 @@ async function fn2(): Promise<Obj> {
 
     });
 }
+
+async function fn3(): Promise<Obj> {
+>fn3 : () => Promise<Obj>
+
+    return await { key: "value" };
+>await { key: "value" } : { key: "value"; }
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+}
+
+async function fn4(): Promise<Obj> {
+>fn4 : () => Promise<Obj>
+
+    return await new Promise(resolve => {
+>await new Promise(resolve => {        resolve({ key: "value" });    }) : Obj
+>new Promise(resolve => {        resolve({ key: "value" });    }) : Promise<Obj>
+>Promise : PromiseConstructor
+>resolve => {        resolve({ key: "value" });    } : (resolve: (value?: Obj | PromiseLike<Obj>) => void) => void
+>resolve : (value?: Obj | PromiseLike<Obj>) => void
+
+        resolve({ key: "value" });
+>resolve({ key: "value" }) : void
+>resolve : (value?: Obj | PromiseLike<Obj>) => void
+>{ key: "value" } : { key: "value"; }
+>key : "value"
+>"value" : "value"
+
+    });
+}
diff --git a/tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts b/tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts
index c9912e0fae1..b2239b9bda6 100644
--- a/tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts
+++ b/tests/cases/conformance/types/contextualTypes/asyncFunctions/contextuallyTypeAsyncFunctionReturnType.ts
@@ -12,4 +12,14 @@ async function fn2(): Promise<Obj> {
     return new Promise(resolve => {
         resolve({ key: "value" });
     });
+}
+
+async function fn3(): Promise<Obj> {
+    return await { key: "value" };
+}
+
+async function fn4(): Promise<Obj> {
+    return await new Promise(resolve => {
+        resolve({ key: "value" });
+    });
 }
\ No newline at end of file