Retain comments inside return statements by including the return keyword in the parse tree

Author: weswigham

src/compiler/binder.ts

@@ -1001,6 +1001,7 @@ namespace ts {
         function bindReturnOrThrow(node: ReturnStatement | ThrowStatement): void {
             bind(node.expression);
             if (node.kind === SyntaxKind.ReturnStatement) {
+                bind(node.returnKeyword);
                 hasExplicitReturn = true;
                 if (currentReturnTarget) {
                     addAntecedent(currentReturnTarget, currentFlow);

src/compiler/emitter.ts

@@ -1571,7 +1571,7 @@ namespace ts {
         }
 
         function emitReturnStatement(node: ReturnStatement) {
-            writeToken(SyntaxKind.ReturnKeyword, node.pos, /*contextNode*/ node);
+            emit(node.returnKeyword);
             emitExpressionWithPrefix(" ", node.expression);
             write(";");
         }

src/compiler/factory.ts

@@ -1460,15 +1460,17 @@ namespace ts {
             : node;
     }
 
-    export function createReturn(expression?: Expression): ReturnStatement {
+    export function createReturn(expression?: Expression, returnKeyword: Token<SyntaxKind.ReturnKeyword> = createToken(SyntaxKind.ReturnKeyword)): ReturnStatement {
         const node = <ReturnStatement>createSynthesizedNode(SyntaxKind.ReturnStatement);
+        node.returnKeyword = returnKeyword;
         node.expression = expression;
         return node;
     }
 
-    export function updateReturn(node: ReturnStatement, expression: Expression | undefined) {
-        return node.expression !== expression
-            ? updateNode(createReturn(expression), node)
+    export function updateReturn(node: ReturnStatement, returnKeyword: Token<SyntaxKind.ReturnKeyword>, expression: Expression | undefined) {
+        return node.expression !== expression ||
+            node.returnKeyword !== returnKeyword
+            ? updateNode(createReturn(expression, returnKeyword), node)
             : node;
     }
 

src/compiler/parser.ts

@@ -259,7 +259,8 @@ namespace ts {
             case SyntaxKind.BreakStatement:
                 return visitNode(cbNode, (<BreakOrContinueStatement>node).label);
             case SyntaxKind.ReturnStatement:
-                return visitNode(cbNode, (<ReturnStatement>node).expression);
+                return visitNode(cbNode, (<ReturnStatement>node).returnKeyword) ||
+                    visitNode(cbNode, (<ReturnStatement>node).expression);
             case SyntaxKind.WithStatement:
                 return visitNode(cbNode, (<WithStatement>node).expression) ||
                     visitNode(cbNode, (<WithStatement>node).statement);
@@ -4687,7 +4688,7 @@ namespace ts {
         function parseReturnStatement(): ReturnStatement {
             const node = <ReturnStatement>createNode(SyntaxKind.ReturnStatement);
 
-            parseExpected(SyntaxKind.ReturnKeyword);
+            node.returnKeyword = parseExpectedToken(SyntaxKind.ReturnKeyword, /*reportAtCurrentPosition*/ true, Diagnostics._0_expected, "return");
             if (!canParseSemicolon()) {
                 node.expression = allowInAnd(parseExpression);
             }

src/compiler/types.ts

@@ -1750,6 +1750,7 @@ namespace ts {
 
     export interface ReturnStatement extends Statement {
         kind: SyntaxKind.ReturnStatement;
+        returnKeyword: Token<SyntaxKind.ReturnKeyword>;
         expression?: Expression;
     }
 

src/compiler/visitor.ts

@@ -634,6 +634,7 @@ namespace ts {
 
             case SyntaxKind.ReturnStatement:
                 return updateReturn(<ReturnStatement>node,
+                    visitNode((<ReturnStatement>node).returnKeyword, visitor, isToken),
                     visitNode((<ReturnStatement>node).expression, visitor, isExpression));
 
             case SyntaxKind.WithStatement:

src/harness/unittests/incrementalParser.ts

@@ -647,7 +647,7 @@ module m3 { }\
             const oldText = ScriptSnapshot.fromString(source);
             const newTextAndChange = withInsert(oldText, 0, "");
 
-            compareTrees(oldText, newTextAndChange.text, newTextAndChange.textChangeRange, 7);
+            compareTrees(oldText, newTextAndChange.text, newTextAndChange.textChangeRange, 8);
         });
 
         it("Class to interface", () => {

tests/baselines/reference/jsdocCastCommentEmit.js

@@ -0,0 +1,17 @@
+//// [jsdocCastCommentEmit.ts]
+// allowJs: true
+// checkJs: true
+// outDir: out/
+// filename: input.js
+function f() {
+    return /* @type {number} */ 42;
+}
+
+//// [jsdocCastCommentEmit.js]
+// allowJs: true
+// checkJs: true
+// outDir: out/
+// filename: input.js
+function f() {
+    return /* @type {number} */ 42;
+}

tests/baselines/reference/jsdocCastCommentEmit.symbols

@@ -0,0 +1,10 @@
+=== tests/cases/compiler/jsdocCastCommentEmit.ts ===
+// allowJs: true
+// checkJs: true
+// outDir: out/
+// filename: input.js
+function f() {
+>f : Symbol(f, Decl(jsdocCastCommentEmit.ts, 0, 0))
+
+    return /* @type {number} */ 42;
+}

tests/baselines/reference/jsdocCastCommentEmit.types

@@ -0,0 +1,11 @@
+=== tests/cases/compiler/jsdocCastCommentEmit.ts ===
+// allowJs: true
+// checkJs: true
+// outDir: out/
+// filename: input.js
+function f() {
+>f : () => number
+
+    return /* @type {number} */ 42;
+>42 : 42
+}

tests/cases/compiler/jsdocCastCommentEmit.ts

@@ -0,0 +1,7 @@
+// allowJs: true
+// checkJs: true
+// outDir: out/
+// filename: input.js
+function f() {
+    return /* @type {number} */ 42;
+}