fix(refactor): keep comments after refactor (#35937)

* add two tests: Refactor: Remove braces from arrow function

* refactor: simplify test and add another test

* fix: copyTrailingAsLeadingComments in addOrRemoveBracesToArrowFunction

* test: add additional test

* fix: clean up changes

* fix: add check for newEdit

* fix: add function for semi colon modifier

* feat: grab all comments during refactor

* refactor: update addOrRemoveBraces logic

* fix: remove duplicate function call

* Update src/services/refactors/addOrRemoveBracesToArrowFunction.ts

* remove blank line

remove blank line

Co-authored-by: Jesse Trinity <42591254+jessetrinity@users.noreply.github.com>

src/services/refactors/addOrRemoveBracesToArrowFunction.ts

@@ -44,6 +44,7 @@ namespace ts.refactor.addOrRemoveBracesToArrowFunction {
         const { expression, returnStatement, func } = info;
 
         let body: ConciseBody;
+
         if (actionName === addBracesActionName) {
             const returnStatement = createReturn(expression);
             body = createBlock([returnStatement], /* multiLine */ true);
@@ -54,13 +55,18 @@ namespace ts.refactor.addOrRemoveBracesToArrowFunction {
             const actualExpression = expression || createVoidZero();
             body = needsParentheses(actualExpression) ? createParen(actualExpression) : actualExpression;
             suppressLeadingAndTrailingTrivia(body);
+            copyTrailingAsLeadingComments(returnStatement, body, file, SyntaxKind.MultiLineCommentTrivia, /* hasTrailingNewLine */ false);
             copyLeadingComments(returnStatement, body, file, SyntaxKind.MultiLineCommentTrivia, /* hasTrailingNewLine */ false);
+            copyTrailingComments(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));
+        const edits = textChanges.ChangeTracker.with(context, t => {
+            t.replaceNode(file, func.body, body);
+        });
+
         return { renameFilename: undefined, renameLocation: undefined, edits };
     }
 

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction24.ts

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

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction25.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const a = (a: number) /*a*/=>/*b*/ { return a; /* trailing */};
+
+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 a = (a: number) => a /* trailing */;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction26.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const a = (a: number) /*a*/=>/*b*/ {/* leading */ return a; /* trailing */};
+
+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 a = (a: number) => /* leading */ a /* trailing */;`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction27.ts

@@ -0,0 +1,13 @@
+/// <reference path='fourslash.ts' />
+
+//// const b = (a: number) /*a*/=>/*b*/ { /* leading */
+////  return a; /* trailing */
+//// }
+
+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 b = (a: number) => /* leading */ a /* trailing */`,
+});

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction28.ts

@@ -0,0 +1,11 @@
+/// <reference path='fourslash.ts' />
+
+//// const a = (a: number) /*a*/=>/*b*/ { return a; /* c */ /* d */ };
+
+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 a = (a: number) => a /* c */ /* d */;`,
+});