From 94c5c3ff47469c3b92342192a6a8353d6c7005ac Mon Sep 17 00:00:00 2001
From: Daniel Rosenwasser <DanielRosenwasser@users.noreply.github.com>
Date: Mon, 4 May 2020 19:32:27 -0700
Subject: [PATCH] Add missing arguments to typeToTypeNode. (#38336)

* Add missing arguments to typeToTypeNode.

* Use returnTypeNode as the enclosingDeclaration.

* Add a test.
---
 .../codefixes/fixReturnTypeInAsyncFunction.ts | 22 +++++++------------
 .../codeFixReturnTypeInAsyncFunction18.ts     | 20 +++++++++++++++++
 2 files changed, 28 insertions(+), 14 deletions(-)
 create mode 100644 tests/cases/fourslash/codeFixReturnTypeInAsyncFunction18.ts

diff --git a/src/services/codefixes/fixReturnTypeInAsyncFunction.ts b/src/services/codefixes/fixReturnTypeInAsyncFunction.ts
index 39e2c1900f..4bb9162b32 100644
--- a/src/services/codefixes/fixReturnTypeInAsyncFunction.ts
+++ b/src/services/codefixes/fixReturnTypeInAsyncFunction.ts
@@ -39,14 +39,20 @@ namespace ts.codefix {
     });
 
     function getInfo(sourceFile: SourceFile, checker: TypeChecker, pos: number): Info | undefined {
-        const returnTypeNode = getReturnTypeNode(sourceFile, pos);
+        if (isInJSFile(sourceFile)) {
+            return undefined;
+        }
+
+        const token = getTokenAtPosition(sourceFile, pos);
+        const func = findAncestor(token, isFunctionLikeDeclaration);
+        const returnTypeNode = func?.type;
         if (!returnTypeNode) {
             return undefined;
         }
 
         const returnType = checker.getTypeFromTypeNode(returnTypeNode);
         const promisedType = checker.getAwaitedType(returnType) || checker.getVoidType();
-        const promisedTypeNode = checker.typeToTypeNode(promisedType);
+        const promisedTypeNode = checker.typeToTypeNode(promisedType, /*enclosingDeclaration*/ returnTypeNode, /*flags*/ undefined);
         if (promisedTypeNode) {
             return { returnTypeNode, returnType, promisedTypeNode, promisedType };
         }
@@ -55,16 +61,4 @@ namespace ts.codefix {
     function doChange(changes: textChanges.ChangeTracker, sourceFile: SourceFile, returnTypeNode: TypeNode, promisedTypeNode: TypeNode): void {
         changes.replaceNode(sourceFile, returnTypeNode, createTypeReferenceNode("Promise", [promisedTypeNode]));
     }
-
-    function getReturnTypeNode(sourceFile: SourceFile, pos: number): TypeNode | undefined {
-        if (isInJSFile(sourceFile)) {
-            return undefined;
-        }
-
-        const token = getTokenAtPosition(sourceFile, pos);
-        const parent = findAncestor(token, isFunctionLikeDeclaration);
-        if (parent?.type) {
-            return parent.type;
-        }
-    }
 }
diff --git a/tests/cases/fourslash/codeFixReturnTypeInAsyncFunction18.ts b/tests/cases/fourslash/codeFixReturnTypeInAsyncFunction18.ts
new file mode 100644
index 0000000000..75b6727fd0
--- /dev/null
+++ b/tests/cases/fourslash/codeFixReturnTypeInAsyncFunction18.ts
@@ -0,0 +1,20 @@
+/// <reference path='fourslash.ts' />
+
+// @target: es2015
+////
+////interface A {}
+////export { A as PublicA };
+////async function foo(): A {
+////    return {}
+////}
+
+verify.codeFix({
+    index: 0,
+    description: [ts.Diagnostics.Replace_0_with_Promise_1.message, "A", "A"],
+    newFileContent: `
+interface A {}
+export { A as PublicA };
+async function foo(): Promise<A> {
+    return {}
+}`
+});