From c399230767a470e703e361333bac0cb6be513b8c Mon Sep 17 00:00:00 2001
From: Wesley Wigham <wwigham@gmail.com>
Date: Tue, 8 Aug 2017 19:53:53 -0700
Subject: [PATCH] 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 5d2142edb1ffb9f6cb150b815aff6e627ae80449.

* Readd test

* Function for handling printing comments on a token
---
 src/compiler/comments.ts                        |  6 +++---
 src/compiler/emitter.ts                         | 14 +++++++++++++-
 .../baselines/reference/jsdocCastCommentEmit.js | 17 +++++++++++++++++
 .../reference/jsdocCastCommentEmit.symbols      | 10 ++++++++++
 .../reference/jsdocCastCommentEmit.types        | 11 +++++++++++
 tests/cases/compiler/jsdocCastCommentEmit.ts    |  7 +++++++
 6 files changed, 61 insertions(+), 4 deletions(-)
 create mode 100644 tests/baselines/reference/jsdocCastCommentEmit.js
 create mode 100644 tests/baselines/reference/jsdocCastCommentEmit.symbols
 create mode 100644 tests/baselines/reference/jsdocCastCommentEmit.types
 create mode 100644 tests/cases/compiler/jsdocCastCommentEmit.ts

diff --git a/src/compiler/comments.ts b/src/compiler/comments.ts
index 06285309f64..50d706d3c42 100644
--- a/src/compiler/comments.ts
+++ b/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");
diff --git a/src/compiler/emitter.ts b/src/compiler/emitter.ts
index 9eb1af2f187..923e6da4215 100644
--- a/src/compiler/emitter.ts
+++ b/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(";");
         }
diff --git a/tests/baselines/reference/jsdocCastCommentEmit.js b/tests/baselines/reference/jsdocCastCommentEmit.js
new file mode 100644
index 00000000000..d071f8f6b2b
--- /dev/null
+++ b/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;
+}
diff --git a/tests/baselines/reference/jsdocCastCommentEmit.symbols b/tests/baselines/reference/jsdocCastCommentEmit.symbols
new file mode 100644
index 00000000000..5490315bc19
--- /dev/null
+++ b/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;
+}
diff --git a/tests/baselines/reference/jsdocCastCommentEmit.types b/tests/baselines/reference/jsdocCastCommentEmit.types
new file mode 100644
index 00000000000..3d4c3e47e46
--- /dev/null
+++ b/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
+}
diff --git a/tests/cases/compiler/jsdocCastCommentEmit.ts b/tests/cases/compiler/jsdocCastCommentEmit.ts
new file mode 100644
index 00000000000..5e7230bd049
--- /dev/null
+++ b/tests/cases/compiler/jsdocCastCommentEmit.ts
@@ -0,0 +1,7 @@
+// allowJs: true
+// checkJs: true
+// outDir: out/
+// filename: input.js
+function f() {
+    return /* @type {number} */ 42;
+}
\ No newline at end of file