add tests and fix

src/compiler/diagnosticMessages.json

@@ -4281,7 +4281,7 @@
         "category": "Message",
         "code": 95054
     },
-    "Convert arrow function": {
+    "Add or remove braces in an arrow function": {
         "category": "Message",
         "code": 95055
     },

src/harness/tsconfig.json

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

src/server/tsconfig.json

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

src/server/tsconfig.library.json

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

src/services/refactors/addOrRemoveBracesToArrowFunction.ts

@@ -1,7 +1,7 @@
 /* @internal */
-namespace ts.refactor.convertArrowFunction {
-    const refactorName = "Convert arrow function";
-    const refactorDescription = Diagnostics.Convert_arrow_function.message;
+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;
@@ -19,21 +19,19 @@ namespace ts.refactor.convertArrowFunction {
         const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
         if (!info) return undefined;
 
-        const actions: RefactorActionInfo[] = [
-            info.addBraces ?
-                {
-                    name: addBracesActionName,
-                    description: addBracesActionDescription
-                } : {
-                    name: removeBracesActionName,
-                    description: removeBracesActionDescription
-                }
-        ];
-
         return [{
             name: refactorName,
             description: refactorDescription,
-            actions
+            actions: [
+                info.addBraces ?
+                    {
+                        name: addBracesActionName,
+                        description: addBracesActionDescription
+                    } : {
+                        name: removeBracesActionName,
+                        description: removeBracesActionDescription
+                    }
+            ]
         }];
     }
 
@@ -42,9 +40,18 @@ namespace ts.refactor.convertArrowFunction {
         const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
         if (!info) return undefined;
 
-        const { addBraces, expression, container } = info;
+        const { expression, container } = info;
         const changeTracker = textChanges.ChangeTracker.fromContext(context);
-        updateBraces(changeTracker, file, container, expression, addBraces);
+
+        if (_actionName === addBracesActionName) {
+            addBraces(changeTracker, file, container, expression);
+        }
+        else if (_actionName === removeBracesActionName) {
+            removeBraces(changeTracker, file, container, expression);
+        }
+        else {
+            Debug.fail("invalid action");
+        }
 
         return {
             renameFilename: undefined,
@@ -53,9 +60,18 @@ namespace ts.refactor.convertArrowFunction {
         };
     }
 
-    function updateBraces(changeTracker: textChanges.ChangeTracker, file: SourceFile, container: ArrowFunction, expression: Expression, addBraces: boolean) {
-        const body = addBraces ? createBlock([createReturn(expression)]) : expression;
+    function addBraces(changeTracker: textChanges.ChangeTracker, file: SourceFile, container: ArrowFunction, expression: Expression) {
+        updateBraces(changeTracker, file, container, createBlock([createReturn(expression)]));
+    }
 
+    function removeBraces(changeTracker: textChanges.ChangeTracker, file: SourceFile, container: ArrowFunction, expression: Expression) {
+        if (!isLiteralExpression(expression) && !isIdentifier(expression) && !isParenthesizedExpression(expression) && expression.kind !== SyntaxKind.NullKeyword) {
+            expression = createParen(expression);
+        }
+        updateBraces(changeTracker, file, container, expression);
+    }
+
+    function updateBraces(changeTracker: textChanges.ChangeTracker, file: SourceFile, container: ArrowFunction, body: ConciseBody) {
         const arrowFunction = updateArrowFunction(
             container,
             container.modifiers,
@@ -78,12 +94,15 @@ namespace ts.refactor.convertArrowFunction {
                 expression: container.body
             };
         }
-        else if (container.body.statements.length === 1 && isReturnStatement(first(container.body.statements))) {
-            return {
-                container,
-                addBraces: false,
-                expression: (<ReturnStatement>first(container.body.statements)).expression
-            };
+        else if (container.body.statements.length === 1) {
+            const firstStatement = first(container.body.statements);
+            if (isReturnStatement(firstStatement)) {
+                return {
+                    container,
+                    addBraces: false,
+                    expression: firstStatement.expression
+                };
+            }
         }
         return undefined;
     }

src/services/tsconfig.json

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

tests/cases/fourslash/refactorAddBracesToArrowFunction7.ts

@@ -1,6 +0,0 @@
-/// <reference path='fourslash.ts' />
-
-//// const foo = /*a*/a/*b*/ => {  };
-
-goTo.select("a", "b");
-verify.not.refactorAvailable("Convert arrow function");

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction1.ts

@@ -4,7 +4,7 @@
 
 goTo.select("a", "b");
 edit.applyRefactor({
-    refactorName: "Convert arrow function",
+    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

@@ -4,7 +4,7 @@
 
 goTo.select("a", "b");
 edit.applyRefactor({
-    refactorName: "Convert arrow function",
+    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/refactorAddOrRemoveBracesToArrowFunction3.ts

@@ -4,7 +4,7 @@
 
 goTo.select("a", "b");
 edit.applyRefactor({
-    refactorName: "Convert arrow function",
+    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

@@ -4,8 +4,8 @@
 
 goTo.select("a", "b");
 edit.applyRefactor({
-    refactorName: "Convert arrow function",
+    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;`,
+    newContent: `const foo = a => (a + 1);`,
 });

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction5.ts

@@ -4,7 +4,7 @@
 
 goTo.select("a", "b");
 edit.applyRefactor({
-    refactorName: "Convert arrow function",
+    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 });`,

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction6.ts

@@ -4,7 +4,7 @@
 
 goTo.select("a", "b");
 edit.applyRefactor({
-    refactorName: "Convert arrow function",
+    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;`,

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction8.ts

@@ -6,4 +6,4 @@
 //// };
 
 goTo.select("a", "b");
-verify.not.refactorAvailable("Convert arrow function");
+verify.not.refactorAvailable("Add or remove braces in an arrow function");

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction9.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const foo = /*a*/a/*b*/ => (1, 2, 3);
+
+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, 2, 3); };`,
+});