add support for add or remove braces to arrow function

src/compiler/diagnosticMessages.json

@@ -4284,5 +4284,17 @@
     "Convert '{0}' to mapped object type": {
         "category": "Message",
         "code": 95055
+    },
+    "Add or remove braces in an arrow function": {
+        "category": "Message",
+        "code": 95056
+    },
+    "Add braces to arrow function": {
+        "category": "Message",
+        "code": 95057
+    },
+    "Remove braces from arrow function": {
+        "category": "Message",
+        "code": 95058
     }
 }

src/compiler/scanner.ts

@@ -727,7 +727,7 @@ namespace ts {
     export function forEachLeadingCommentRange<U>(text: string, pos: number, cb: (pos: number, end: number, kind: CommentKind, hasTrailingNewLine: boolean) => U): U | undefined;
     export function forEachLeadingCommentRange<T, U>(text: string, pos: number, cb: (pos: number, end: number, kind: CommentKind, hasTrailingNewLine: boolean, state: T) => U, state: T): U | undefined;
     export function forEachLeadingCommentRange<T, U>(text: string, pos: number, cb: (pos: number, end: number, kind: CommentKind, hasTrailingNewLine: boolean, state: T) => U, state?: T): U | undefined {
-        return iterateCommentRanges(/*reduce*/ false, text, pos, /*trailing*/ false, cb, state);
+        return iterateCommentRanges(/*reduce*/ false, text, pos, /*trailing*/ false, cb, state, /* initial */ undefined);
     }
 
     export function forEachTrailingCommentRange<U>(text: string, pos: number, cb: (pos: number, end: number, kind: CommentKind, hasTrailingNewLine: boolean) => U): U | undefined;

src/harness/tsconfig.json

@@ -123,6 +123,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

@@ -119,6 +119,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

@@ -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/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,100 @@
+/* @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, /* explicitHtnl */ true);
+        }
+        else if (actionName === removeBracesActionName && returnStatement) {
+            const actualExpression = expression || createVoidZero();
+            body = needsParentheses(actualExpression) ? createParen(actualExpression) : actualExpression;
+            suppressLeadingAndTrailingTrivia(body);
+            copyComments(returnStatement, body, file, SyntaxKind.MultiLineCommentTrivia, /* explicitHtnl */ false);
+        }
+        else {
+            Debug.fail("invalid action");
+        }
+
+        const edits = textChanges.ChangeTracker.with(context, t => updateBody(t, file, func, body));
+        return { renameFilename: undefined, renameLocation: undefined, edits };
+    }
+
+    function needsParentheses(expression: Expression) {
+        return isBinaryExpression(expression) && expression.operatorToken.kind === SyntaxKind.CommaToken || isObjectLiteralExpression(expression);
+    }
+
+    function updateBody(changeTracker: textChanges.ChangeTracker, file: SourceFile, container: ArrowFunction, body: ConciseBody) {
+        changeTracker.replaceNode(file, container.body, body);
+    }
+
+    function getConvertibleArrowFunctionAtPosition(file: SourceFile, startPosition: number): Info | undefined {
+        const node = getTokenAtPosition(file, startPosition, /*includeJsDocComment*/ false);
+        const func = getContainingFunction(node);
+        if (!func || !isArrowFunction(func)) 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

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

src/services/utilities.ts

@@ -1649,4 +1649,20 @@ namespace ts {
         Debug.assert(lastPos >= 0);
         return lastPos;
     }
+
+    export function copyComments(sourceNode: Node, targetNode: Node, sourceFile: SourceFile, explicitKind?: CommentKind, explicitHtnl?: 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, explicitKind || kind, sourceFile.text.slice(pos, end), explicitHtnl !== undefined ? explicitHtnl : htnl);
+        });
+    }
 }

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction1.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => a + 1;
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Add braces to arrow function",
+    actionDescription: "Add braces to arrow function",
+    newContent: `const foo = a => {
+    return a + 1;
+};`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction10.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return (1, 2, 3); };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => (1, 2, 3);`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction11.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return 1, 2, 3; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => (1, 2, 3);`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction12.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return "foo"; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => "foo";`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction13.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return null; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => null;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction14.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return undefined; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => undefined;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction15.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return void 0; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => void 0;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction16.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return {}; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => ({});`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction17.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return `abc{a}`; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => \`abc{a}\`;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction18.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return `abc`; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => \`abc\`;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction19.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return a; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => a;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction2.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => ({ a: 1 });
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Add braces to arrow function",
+    actionDescription: "Add braces to arrow function",
+    newContent: `const foo = a => {
+    return ({ a: 1 });
+};`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction20.ts

@@ -0,0 +1,14 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { 
+////     // return comment
+////     return a; 
+//// };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => /* return comment*/ a;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction3.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => 1;
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Add braces to arrow function",
+    actionDescription: "Add braces to arrow function",
+    newContent: `const foo = a => {
+    return 1;
+};`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction4.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => { return a + 1; };
+
+goTo.select("a", "b");
+edit.applyRefactor({
+    refactorName: "Add or remove braces in an arrow function",
+    actionName: "Remove braces from arrow function",
+    actionDescription: "Remove braces from arrow function",
+    newContent: `const foo = a => a + 1;`,
+});