Merge pull request #23423 from Kingwl/add-braces

add support for add or remove braces to arrow function

src/compiler/diagnosticMessages.json

@@ -4402,5 +4402,17 @@
     "Convert named imports to namespace import": {
         "category": "Message",
         "code": 95057
+    },
+    "Add or remove braces in an arrow function": {
+        "category": "Message",
+        "code": 95058
+    },
+    "Add braces to arrow function": {
+        "category": "Message",
+        "code": 95059
+    },
+    "Remove braces from arrow function": {
+        "category": "Message",
+        "code": 95060
     }
 }

src/harness/tsconfig.json

@@ -125,6 +125,7 @@
         "../services/refactors/extractSymbol.ts",
         "../services/refactors/generateGetAccessorAndSetAccessor.ts",
         "../services/refactors/moveToNewFile.ts",
+        "../services/refactors/addOrRemoveBracesToArrowFunction.ts",
         "../services/sourcemaps.ts",
         "../services/services.ts",
         "../services/breakpoints.ts",

src/server/tsconfig.json

@@ -120,6 +120,7 @@
         "../services/refactors/extractSymbol.ts",
         "../services/refactors/generateGetAccessorAndSetAccessor.ts",
         "../services/refactors/moveToNewFile.ts",
+        "../services/refactors/addOrRemoveBracesToArrowFunction.ts",
         "../services/sourcemaps.ts",
         "../services/services.ts",
         "../services/breakpoints.ts",

src/server/tsconfig.library.json

@@ -126,6 +126,7 @@
         "../services/refactors/extractSymbol.ts",
         "../services/refactors/generateGetAccessorAndSetAccessor.ts",
         "../services/refactors/moveToNewFile.ts",
+        "../services/refactors/addOrRemoveBracesToArrowFunction.ts",
         "../services/sourcemaps.ts",
         "../services/services.ts",
         "../services/breakpoints.ts",

src/services/codefixes/convertFunctionToEs6Class.ts

@@ -202,22 +202,6 @@ namespace ts.codefix {
         }
     }
 
-    function copyComments(sourceNode: Node, targetNode: Node, sourceFile: SourceFile) {
-        forEachLeadingCommentRange(sourceFile.text, sourceNode.pos, (pos, end, kind, htnl) => {
-            if (kind === SyntaxKind.MultiLineCommentTrivia) {
-                // Remove leading /*
-                pos += 2;
-                // Remove trailing */
-                end -= 2;
-            }
-            else {
-                // Remove leading //
-                pos += 2;
-            }
-            addSyntheticLeadingComment(targetNode, kind, sourceFile.text.slice(pos, end), htnl);
-        });
-    }
-
     function getModifierKindFromSource(source: Node, kind: SyntaxKind): ReadonlyArray<Modifier> | undefined {
         return filter(source.modifiers, modifier => modifier.kind === kind);
     }

src/services/refactors/addOrRemoveBracesToArrowFunction.ts

@@ -0,0 +1,96 @@
+/* @internal */
+namespace ts.refactor.addOrRemoveBracesToArrowFunction {
+    const refactorName = "Add or remove braces in an arrow function";
+    const refactorDescription = Diagnostics.Add_or_remove_braces_in_an_arrow_function.message;
+    const addBracesActionName = "Add braces to arrow function";
+    const removeBracesActionName = "Remove braces from arrow function";
+    const addBracesActionDescription = Diagnostics.Add_braces_to_arrow_function.message;
+    const removeBracesActionDescription = Diagnostics.Remove_braces_from_arrow_function.message;
+    registerRefactor(refactorName, { getEditsForAction, getAvailableActions });
+
+    interface Info {
+        func: ArrowFunction;
+        expression: Expression | undefined;
+        returnStatement?: ReturnStatement;
+        addBraces: boolean;
+    }
+
+    function getAvailableActions(context: RefactorContext): ApplicableRefactorInfo[] | undefined {
+        const { file, startPosition } = context;
+        const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
+        if (!info) return undefined;
+
+        return [{
+            name: refactorName,
+            description: refactorDescription,
+            actions: [
+                info.addBraces ?
+                    {
+                        name: addBracesActionName,
+                        description: addBracesActionDescription
+                    } : {
+                        name: removeBracesActionName,
+                        description: removeBracesActionDescription
+                    }
+            ]
+        }];
+    }
+
+    function getEditsForAction(context: RefactorContext, actionName: string): RefactorEditInfo | undefined {
+        const { file, startPosition } = context;
+        const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
+        if (!info) return undefined;
+
+        const { expression, returnStatement, func } = info;
+
+        let body: ConciseBody;
+        if (actionName === addBracesActionName) {
+            const returnStatement = createReturn(expression);
+            body = createBlock([returnStatement], /* multiLine */ true);
+            suppressLeadingAndTrailingTrivia(body);
+            copyComments(expression!, returnStatement, file, SyntaxKind.MultiLineCommentTrivia, /* hasTrailingNewLine */ true);
+        }
+        else if (actionName === removeBracesActionName && returnStatement) {
+            const actualExpression = expression || createVoidZero();
+            body = needsParentheses(actualExpression) ? createParen(actualExpression) : actualExpression;
+            suppressLeadingAndTrailingTrivia(body);
+            copyComments(returnStatement, body, file, SyntaxKind.MultiLineCommentTrivia, /* hasTrailingNewLine */ false);
+        }
+        else {
+            Debug.fail("invalid action");
+        }
+
+        const edits = textChanges.ChangeTracker.with(context, t => t.replaceNode(file, func.body, body));
+        return { renameFilename: undefined, renameLocation: undefined, edits };
+    }
+
+    function needsParentheses(expression: Expression) {
+        return isBinaryExpression(expression) && expression.operatorToken.kind === SyntaxKind.CommaToken || isObjectLiteralExpression(expression);
+    }
+
+    function getConvertibleArrowFunctionAtPosition(file: SourceFile, startPosition: number): Info | undefined {
+        const node = getTokenAtPosition(file, startPosition, /*includeJsDocComment*/ false);
+        const func = getContainingFunction(node);
+        if (!func || !isArrowFunction(func) || (!rangeContainsRange(func, node) || rangeContainsRange(func.body, node))) return undefined;
+
+        if (isExpression(func.body)) {
+            return {
+                func,
+                addBraces: true,
+                expression: func.body
+            };
+        }
+        else if (func.body.statements.length === 1) {
+            const firstStatement = first(func.body.statements);
+            if (isReturnStatement(firstStatement)) {
+                return {
+                    func,
+                    addBraces: false,
+                    expression: firstStatement.expression,
+                    returnStatement: firstStatement
+                };
+            }
+        }
+        return undefined;
+    }
+}

src/services/tsconfig.json

@@ -117,6 +117,7 @@
         "refactors/extractSymbol.ts",
         "refactors/generateGetAccessorAndSetAccessor.ts",
         "refactors/moveToNewFile.ts",
+        "refactors/addOrRemoveBracesToArrowFunction.ts",
         "sourcemaps.ts",
         "services.ts",
         "breakpoints.ts",

src/services/utilities.ts

@@ -1720,6 +1720,22 @@ namespace ts {
         return lastPos;
     }
 
+    export function copyComments(sourceNode: Node, targetNode: Node, sourceFile: SourceFile, commentKind?: CommentKind, hasTrailingNewLine?: boolean) {
+        forEachLeadingCommentRange(sourceFile.text, sourceNode.pos, (pos, end, kind, htnl) => {
+            if (kind === SyntaxKind.MultiLineCommentTrivia) {
+                // Remove leading /*
+                pos += 2;
+                // Remove trailing */
+                end -= 2;
+            }
+            else {
+                // Remove leading //
+                pos += 2;
+            }
+            addSyntheticLeadingComment(targetNode, commentKind || kind, sourceFile.text.slice(pos, end), hasTrailingNewLine !== undefined ? hasTrailingNewLine : htnl);
+        });
+    }
+
     function indexInTextChange(change: string, name: string): number {
         if (startsWith(change, name)) return 0;
         // Add a " " to avoid references inside words