fix void return statement

Author: Kingwl

src/services/refactors/addOrRemoveBracesToArrowFunction.ts

@@ -10,7 +10,8 @@ namespace ts.refactor.addOrRemoveBracesToArrowFunction {
 
     interface Info {
         func: ArrowFunction;
-        expression: Expression;
+        expression: Expression | undefined;
+        returnStatement?: ReturnStatement;
         addBraces: boolean;
     }
 
@@ -40,23 +41,23 @@ namespace ts.refactor.addOrRemoveBracesToArrowFunction {
         const info = getConvertibleArrowFunctionAtPosition(file, startPosition);
         if (!info) return undefined;
 
-        const { expression, func } = info;
+        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, true, true);
+            copyComments(expression!, returnStatement, file, SyntaxKind.MultiLineCommentTrivia, /* explicitHtnl */ true);
         }
-        else if (actionName === removeBracesActionName) {
-            const returnStatement = <ReturnStatement>expression.parent;
-            body = needsParentheses(expression) ? createParen(expression) : expression;
+        else if (actionName === removeBracesActionName && returnStatement) {
+            const actualExpression = expression || createVoidZero();
+            body = needsParentheses(actualExpression) ? createParen(actualExpression) : actualExpression;
             suppressLeadingAndTrailingTrivia(body);
-            copyComments(returnStatement, body, file, SyntaxKind.MultiLineCommentTrivia, false);
+            copyComments(returnStatement, body, file, SyntaxKind.MultiLineCommentTrivia, /* explicitHtnl */ false);
         }
         else {
-            Debug.fail('invalid action');
+            Debug.fail("invalid action");
         }
 
         const edits = textChanges.ChangeTracker.with(context, t => updateBody(t, file, func, body));
@@ -89,7 +90,8 @@ namespace ts.refactor.addOrRemoveBracesToArrowFunction {
                 return {
                     func,
                     addBraces: false,
-                    expression: firstStatement.expression
+                    expression: firstStatement.expression,
+                    returnStatement: firstStatement
                 };
             }
         }

src/services/utilities.ts

@@ -1643,7 +1643,7 @@ namespace ts {
         return lastPos;
     }
 
-    export function copyComments(sourceNode: Node, targetNode: Node, sourceFile: SourceFile, explicitKind?: CommentKind, explicitHtnl?: boolean, inline?: boolean) {
+    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 /*
@@ -1656,6 +1656,6 @@ namespace ts {
                 pos += 2;
             }
             addSyntheticLeadingComment(targetNode, explicitKind || kind, sourceFile.text.slice(pos, end), explicitHtnl !== undefined ? explicitHtnl : htnl);
-        })
+        });
     }
 }

tests/cases/fourslash/refactorAddOrRemoveBracesToArrowFunction23.ts

@@ -0,0 +1,13 @@
+/// <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 => void 0`,
+});