Improve accuracy of remove unnecessary await fix (microsoft#32384)

andrewbranch · web-flow · commit 7cdfbceb43e4 · 2019-07-15T15:17:32.000-07:00
diff --git a/src/compiler/factory.ts b/src/compiler/factory.ts
@@ -4701,7 +4701,7 @@ namespace ts {
         }
     }
 
-    function getLeftmostExpression(node: Expression, stopAtCallExpressions: boolean) {
+    export function getLeftmostExpression(node: Expression, stopAtCallExpressions: boolean) {
         while (true) {
             switch (node.kind) {
                 case SyntaxKind.PostfixUnaryExpression:
diff --git a/src/services/codefixes/removeUnnecessaryAwait.ts b/src/services/codefixes/removeUnnecessaryAwait.ts
@@ -26,8 +26,18 @@ namespace ts.codefix {
             return;
         }
 
-        const parenthesizedExpression = tryCast(awaitExpression.parent, isParenthesizedExpression);
-        const removeParens = parenthesizedExpression && (isIdentifier(awaitExpression.expression) || isCallExpression(awaitExpression.expression));
-        changeTracker.replaceNode(sourceFile, removeParens ? parenthesizedExpression || awaitExpression : awaitExpression, awaitExpression.expression);
+        let expressionToReplace: Node = awaitExpression;
+        const hasSurroundingParens = isParenthesizedExpression(awaitExpression.parent);
+        if (hasSurroundingParens) {
+            const leftMostExpression = getLeftmostExpression(awaitExpression.expression, /*stopAtCallExpressions*/ false);
+            if (isIdentifier(leftMostExpression)) {
+                const precedingToken = findPrecedingToken(awaitExpression.parent.pos, sourceFile);
+                if (precedingToken && precedingToken.kind !== SyntaxKind.NewKeyword) {
+                    expressionToReplace = awaitExpression.parent;
+                }
+            }
+        }
+
+        changeTracker.replaceNode(sourceFile, expressionToReplace, awaitExpression.expression);
     }
 }
diff --git a/tests/cases/fourslash/codeFixRemoveUnnecessaryAwait.ts b/tests/cases/fourslash/codeFixRemoveUnnecessaryAwait.ts
@@ -1,26 +1,32 @@
 /// <reference path="fourslash.ts" />
 ////declare class C { foo(): void }
+////declare function getC(): { Class: C };
 ////declare function foo(): string;
 ////async function f() {
 ////  await "";
 ////  await 0;
 ////  (await foo()).toLowerCase();
 ////  (await 0).toFixed();
 ////  (await new C).foo();
+////  (await function() { }());
+////  new (await getC()).Class();
 ////}
 
 verify.codeFix({
   description: ts.Diagnostics.Remove_unnecessary_await.message,
   index: 0,
   newFileContent:
 `declare class C { foo(): void }
+declare function getC(): { Class: C };
 declare function foo(): string;
 async function f() {
   "";
   await 0;
   (await foo()).toLowerCase();
   (await 0).toFixed();
   (await new C).foo();
+  (await function() { }());
+  new (await getC()).Class();
 }`
 });
 
@@ -29,12 +35,15 @@ verify.codeFixAll({
   fixId: "removeUnnecessaryAwait",
   newFileContent:
 `declare class C { foo(): void }
+declare function getC(): { Class: C };
 declare function foo(): string;
 async function f() {
   "";
   0;
   foo().toLowerCase();
   (0).toFixed();
   (new C).foo();
+  (function() { } ());
+  new (getC()).Class();
 }`
 });

Original file line number	Diff line number	Diff line change
`@@ -4701,7 +4701,7 @@ namespace ts {`
`4701`	`4701`	`}`
`4702`	`4702`	`}`
`4703`	`4703`
`4704`		`- function getLeftmostExpression(node: Expression, stopAtCallExpressions: boolean) {`
	`4704`	`+ export function getLeftmostExpression(node: Expression, stopAtCallExpressions: boolean) {`
`4705`	`4705`	`while (true) {`
`4706`	`4706`	`switch (node.kind) {`
`4707`	`4707`	`case SyntaxKind.PostfixUnaryExpression:`