Retain comments inside return statements (#17557)

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

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

This reverts commit 5d2142edb1.

* Readd test

* Function for handling printing comments on a token

src/compiler/comments.ts

@@ -8,7 +8,7 @@ namespace ts {
         setWriter(writer: EmitTextWriter): void;
         emitNodeWithComments(hint: EmitHint, node: Node, emitCallback: (hint: EmitHint, node: Node) => void): void;
         emitBodyWithDetachedComments(node: Node, detachedRange: TextRange, emitCallback: (node: Node) => void): void;
-        emitTrailingCommentsOfPosition(pos: number): void;
+        emitTrailingCommentsOfPosition(pos: number, prefixSpace?: boolean): void;
         emitLeadingCommentsOfPosition(pos: number): void;
     }
 
@@ -306,7 +306,7 @@ namespace ts {
             }
         }
 
-        function emitTrailingCommentsOfPosition(pos: number) {
+        function emitTrailingCommentsOfPosition(pos: number, prefixSpace?: boolean) {
             if (disabled) {
                 return;
             }
@@ -315,7 +315,7 @@ namespace ts {
                 performance.mark("beforeEmitTrailingCommentsOfPosition");
             }
 
-            forEachTrailingCommentToEmit(pos, emitTrailingCommentOfPosition);
+            forEachTrailingCommentToEmit(pos, prefixSpace ? emitTrailingComment : emitTrailingCommentOfPosition);
 
             if (extendedDiagnostics) {
                 performance.measure("commentTime", "beforeEmitTrailingCommentsOfPosition");

src/compiler/emitter.ts

@@ -1572,8 +1572,20 @@ namespace ts {
             write(";");
         }
 
+        function emitTokenWithComment(token: SyntaxKind, pos: number, contextNode?: Node) {
+            const node = contextNode && getParseTreeNode(contextNode);
+            if (node && node.kind === contextNode.kind) {
+                pos = skipTrivia(currentSourceFile.text, pos);
+            }
+            pos = writeToken(token, pos, /*contextNode*/ contextNode);
+            if (node && node.kind === contextNode.kind) {
+                emitTrailingCommentsOfPosition(pos, /*prefixSpace*/ true);
+            }
+            return pos;
+        }
+
         function emitReturnStatement(node: ReturnStatement) {
-            writeToken(SyntaxKind.ReturnKeyword, node.pos, /*contextNode*/ node);
+            emitTokenWithComment(SyntaxKind.ReturnKeyword, node.pos, /*contextNode*/ node);
             emitExpressionWithPrefix(" ", node.expression);
             write(";");
         }

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;
+}