Improve performance of suggestionDiagnostics for convertToAsyncFunction (#28089)

Author: Andy

src/services/codefixes/convertToAsyncFunction.ts

@@ -58,9 +58,9 @@ namespace ts.codefix {
         const allVarNames: SymbolAndIdentifier[] = [];
         const isInJavascript = isInJSFile(functionToConvert);
         const setOfExpressionsToReturn = getAllPromiseExpressionsToReturn(functionToConvert, checker);
-        const functionToConvertRenamed: FunctionLikeDeclaration = renameCollidingVarNames(functionToConvert, checker, synthNamesMap, context, setOfExpressionsToReturn, originalTypeMap, allVarNames);
+        const functionToConvertRenamed = renameCollidingVarNames(functionToConvert, checker, synthNamesMap, context, setOfExpressionsToReturn, originalTypeMap, allVarNames);
         const constIdentifiers = getConstIdentifiers(synthNamesMap);
-        const returnStatements = getReturnStatementsWithPromiseHandlers(functionToConvertRenamed);
+        const returnStatements = functionToConvertRenamed.body && isBlock(functionToConvertRenamed.body) ? getReturnStatementsWithPromiseHandlers(functionToConvertRenamed.body) : emptyArray;
         const transformer: Transformer = { checker, synthNamesMap, allVarNames, setOfExpressionsToReturn, constIdentifiers, originalTypeMap, isInJSFile: isInJavascript };
 
         if (!returnStatements.length) {
@@ -87,6 +87,14 @@ namespace ts.codefix {
         }
     }
 
+    function getReturnStatementsWithPromiseHandlers(body: Block): ReadonlyArray<ReturnStatement> {
+        const res: ReturnStatement[] = [];
+        forEachReturnStatement(body, ret => {
+            if (isReturnStatementWithFixablePromiseHandler(ret)) res.push(ret);
+        });
+        return res;
+    }
+
     // Returns the identifiers that are never reassigned in the refactor
     function getConstIdentifiers(synthNamesMap: ReadonlyMap<SynthIdentifier>): Identifier[] {
         const constIdentifiers: Identifier[] = [];
@@ -442,7 +450,7 @@ namespace ts.codefix {
                             seenReturnStatement = true;
                         }
 
-                        if (getReturnStatementsWithPromiseHandlers(statement).length) {
+                        if (isReturnStatementWithFixablePromiseHandler(statement)) {
                             refactoredStmts = refactoredStmts.concat(getInnerTransformationBody(transformer, [statement], prevArgName));
                         }
                         else {
@@ -458,7 +466,7 @@ namespace ts.codefix {
                             seenReturnStatement);
                 }
                 else {
-                    const innerRetStmts = getReturnStatementsWithPromiseHandlers(createReturn(funcBody));
+                    const innerRetStmts = isFixablePromiseHandler(funcBody) ? [createReturn(funcBody)] : emptyArray;
                     const innerCbBody = getInnerTransformationBody(transformer, innerRetStmts, prevArgName);
 
                     if (innerCbBody.length > 0) {

src/services/suggestionDiagnostics.ts

@@ -113,59 +113,41 @@ namespace ts {
         }
     }
 
-    function addConvertToAsyncFunctionDiagnostics(node: FunctionLikeDeclaration, checker: TypeChecker, diags: DiagnosticWithLocation[]): void {
-
-        if (isAsyncFunction(node) || !node.body) {
-            return;
+    function addConvertToAsyncFunctionDiagnostics(node: FunctionLikeDeclaration, checker: TypeChecker, diags: Push<DiagnosticWithLocation>): void {
+        if (!isAsyncFunction(node) &&
+            node.body &&
+            isBlock(node.body) &&
+            hasReturnStatementWithPromiseHandler(node.body) &&
+            returnsPromise(node, checker)) {
+            diags.push(createDiagnosticForNode(
+                !node.name && isVariableDeclaration(node.parent) && isIdentifier(node.parent.name) ? node.parent.name : node,
+                Diagnostics.This_may_be_converted_to_an_async_function));
         }
+    }
 
+    function returnsPromise(node: FunctionLikeDeclaration, checker: TypeChecker): boolean {
         const functionType = checker.getTypeAtLocation(node);
-
         const callSignatures = checker.getSignaturesOfType(functionType, SignatureKind.Call);
         const returnType = callSignatures.length ? checker.getReturnTypeOfSignature(callSignatures[0]) : undefined;
-
-        if (!returnType || !checker.getPromisedTypeOfPromise(returnType)) {
-            return;
-        }
-
-        // collect all the return statements
-        // check that a property access expression exists in there and that it is a handler
-        const returnStatements = getReturnStatementsWithPromiseHandlers(node);
-        if (returnStatements.length > 0) {
-            diags.push(createDiagnosticForNode(!node.name && isVariableDeclaration(node.parent) && isIdentifier(node.parent.name) ? node.parent.name : node, Diagnostics.This_may_be_converted_to_an_async_function));
-        }
+        return !!returnType && !!checker.getPromisedTypeOfPromise(returnType);
     }
 
     function getErrorNodeFromCommonJsIndicator(commonJsModuleIndicator: Node): Node {
         return isBinaryExpression(commonJsModuleIndicator) ? commonJsModuleIndicator.left : commonJsModuleIndicator;
     }
 
-    /** @internal */
-    export function getReturnStatementsWithPromiseHandlers(node: Node): ReturnStatement[] {
-        const returnStatements: ReturnStatement[] = [];
-        if (isFunctionLike(node)) {
-            forEachChild(node, visit);
-        }
-        else {
-            visit(node);
-        }
-
-        function visit(child: Node) {
-            if (isFunctionLike(child)) {
-                return;
-            }
-
-            if (isReturnStatement(child) && child.expression && isFixablePromiseHandler(child.expression)) {
-                returnStatements.push(child);
-            }
+    function hasReturnStatementWithPromiseHandler(body: Block): boolean {
+        return !!forEachReturnStatement(body, isReturnStatementWithFixablePromiseHandler);
+    }
 
-            forEachChild(child, visit);
-        }
-        return returnStatements;
+    /* @internal */
+    export function isReturnStatementWithFixablePromiseHandler(node: Node): node is ReturnStatement {
+        return isReturnStatement(node) && !!node.expression && isFixablePromiseHandler(node.expression);
     }
 
     // Should be kept up to date with transformExpression in convertToAsyncFunction.ts
-    function isFixablePromiseHandler(node: Node): boolean {
+    /* @internal */
+    export function isFixablePromiseHandler(node: Node): boolean {
         // ensure outermost call exists and is a promise handler
         if (!isPromiseHandler(node) || !node.arguments.every(isFixablePromiseArgument)) {
             return false;