feat(printer): implement JSDoc node emission in AST printer

packages/typescript/src/api/async/api.ts

@@ -67,6 +67,7 @@ import {
     toRootedPath,
 } from "../path.ts";
 import { prepareCompilerOptions } from "../prepareCompilerOptions.generated.ts";
+import type { PrintNodeComment } from "../proto.generated.ts";
 import type {
     BuildResponse,
     CleanBuildResponse,
@@ -3129,24 +3130,28 @@ export class Printer {
     }
 
     async printNode(node: Node, options: PrintNodeOptions = {}): Promise<string> {
-        const encoded = encodeNode(node);
+        const comments: PrintNodeComment[] = [];
+        const encoded = encodeNode(node, comments);
         const base64 = uint8ArrayToBase64(encoded);
         return this.client.apiRequest("printNode", {
             data: base64,
             preserveSourceNewlines: options.preserveSourceNewlines,
             neverAsciiEscape: options.neverAsciiEscape,
             terminateUnterminatedLiterals: options.terminateUnterminatedLiterals,
+            comments: comments.length ? comments : undefined,
         });
     }
 
     async printFile(sourceFile: SourceFile, options: PrintNodeOptions = {}): Promise<string> {
-        const encoded = encodeNode(sourceFile);
+        const comments: PrintNodeComment[] = [];
+        const encoded = encodeNode(sourceFile, comments);
         const base64 = uint8ArrayToBase64(encoded);
         return this.client.apiRequest("printNode", {
             data: base64,
             preserveSourceNewlines: options.preserveSourceNewlines,
             neverAsciiEscape: options.neverAsciiEscape,
             terminateUnterminatedLiterals: options.terminateUnterminatedLiterals,
+            comments: comments.length ? comments : undefined,
         });
     }
 }

packages/typescript/src/api/node/encoder.ts

@@ -6,7 +6,12 @@ import type {
     SourceFile,
     TemplateLiteralLikeNode,
 } from "../../ast/index.ts";
-import { SyntaxKind } from "../../ast/index.ts";
+import {
+    getSyntheticLeadingComments,
+    getSyntheticTrailingComments,
+    SyntaxKind,
+} from "../../ast/index.ts";
+import type { PrintNodeComment } from "../proto.generated.ts";
 import {
     getNodeCommonData,
     getNodeDataType,
@@ -199,7 +204,7 @@ export function encodeSourceFile(sourceFile: SourceFile): Uint8Array {
  * Encode an arbitrary AST node into the binary format.
  * When encoding a non-SourceFile node, the header hash and parse options fields will be zero.
  */
-export function encodeNode(node: Node): Uint8Array {
+export function encodeNode(node: Node, comments?: PrintNodeComment[]): Uint8Array {
     const strs = new StringTable();
     const extendedDataValues: number[] = [];
     const structuredWriter = new MsgpackWriter();
@@ -214,9 +219,20 @@ export function encodeNode(node: Node): Uint8Array {
     let parentIndex = 0;
     let prevIndex = 0;
 
+    function collectComments(node: Node, index: number): void {
+        if (!comments) return;
+        for (const comment of getSyntheticLeadingComments(node) ?? []) {
+            comments.push({ node: index, kind: comment.kind, text: comment.text, hasTrailingNewLine: comment.hasTrailingNewLine || undefined });
+        }
+        for (const comment of getSyntheticTrailingComments(node) ?? []) {
+            comments.push({ node: index, trailing: true, kind: comment.kind, text: comment.text, hasTrailingNewLine: comment.hasTrailingNewLine || undefined });
+        }
+    }
+
     function visitNode(node: Node): void {
         nodeCount++;
         const currentIndex = nodeCount;
+        collectComments(node, currentIndex);
 
         if (prevIndex !== 0) {
             // Set next pointer on previous sibling
@@ -300,6 +316,7 @@ export function encodeNode(node: Node): Uint8Array {
 
     // Encode root node
     nodeCount++;
+    collectComments(node, 1);
     parentIndex++;
     const rootData = getNodeData(node, strs, extendedDataValues, structuredWriter);
     nodeValues.push(

packages/typescript/src/api/proto.generated.ts

@@ -1056,6 +1056,8 @@ export interface PrintNodeParams {
     preserveSourceNewlines?: boolean | undefined;
     neverAsciiEscape?: boolean | undefined;
     terminateUnterminatedLiterals?: boolean | undefined;
+    /** Comments are synthetic comments attached to nodes of the encoded AST. */
+    comments?: readonly PrintNodeComment[] | undefined;
 }
 
 export interface EmitParams {
@@ -1696,6 +1698,17 @@ export interface DiagnosticSourceLineResponse {
     text: string;
 }
 
+/** PrintNodeComment is a synthetic comment attached to a node of the encoded AST. */
+export interface PrintNodeComment {
+    /** index of the node in the encoded AST data */
+    node: number;
+    trailing?: boolean | undefined;
+    /** SingleLineCommentTrivia or MultiLineCommentTrivia */
+    kind: number;
+    text: string;
+    hasTrailingNewLine?: boolean | undefined;
+}
+
 export interface EmitOutputFile {
     fileName: RootedFilePath;
     text: string;

packages/typescript/src/api/sync/api.ts

@@ -84,6 +84,7 @@ import {
     toRootedPath,
 } from "../path.ts";
 import { prepareCompilerOptions } from "../prepareCompilerOptions.generated.ts";
+import type { PrintNodeComment } from "../proto.generated.ts";
 import type {
     BuildResponse,
     CleanBuildResponse,
@@ -6484,23 +6485,27 @@ export class Printer {
             owner,
             "printNode",
             function (node: Node, options: PrintNodeOptions = {}): string {
-                const encoded = encodeNode(node);
+                const comments: PrintNodeComment[] = [];
+                const encoded = encodeNode(node, comments);
                 const base64 = uint8ArrayToBase64(encoded);
                 return owner.client.apiRequest("printNode", {
                     data: base64,
                     preserveSourceNewlines: options.preserveSourceNewlines,
                     neverAsciiEscape: options.neverAsciiEscape,
                     terminateUnterminatedLiterals: options.terminateUnterminatedLiterals,
+                    comments: comments.length ? comments : undefined,
                 });
             },
             function* (node: Node, options: PrintNodeOptions = {}): Generator<ProtocolRequest, string, ProtocolResponse["result"]> {
-                const encoded = encodeNode(node);
+                const comments: PrintNodeComment[] = [];
+                const encoded = encodeNode(node, comments);
                 const base64 = uint8ArrayToBase64(encoded);
                 return yield* apiRequest("printNode", {
                     data: base64,
                     preserveSourceNewlines: options.preserveSourceNewlines,
                     neverAsciiEscape: options.neverAsciiEscape,
                     terminateUnterminatedLiterals: options.terminateUnterminatedLiterals,
+                    comments: comments.length ? comments : undefined,
                 });
             },
         );
@@ -6515,23 +6520,27 @@ export class Printer {
             owner,
             "printFile",
             function (sourceFile: SourceFile, options: PrintNodeOptions = {}): string {
-                const encoded = encodeNode(sourceFile);
+                const comments: PrintNodeComment[] = [];
+                const encoded = encodeNode(sourceFile, comments);
                 const base64 = uint8ArrayToBase64(encoded);
                 return owner.client.apiRequest("printNode", {
                     data: base64,
                     preserveSourceNewlines: options.preserveSourceNewlines,
                     neverAsciiEscape: options.neverAsciiEscape,
                     terminateUnterminatedLiterals: options.terminateUnterminatedLiterals,
+                    comments: comments.length ? comments : undefined,
                 });
             },
             function* (sourceFile: SourceFile, options: PrintNodeOptions = {}): Generator<ProtocolRequest, string, ProtocolResponse["result"]> {
-                const encoded = encodeNode(sourceFile);
+                const comments: PrintNodeComment[] = [];
+                const encoded = encodeNode(sourceFile, comments);
                 const base64 = uint8ArrayToBase64(encoded);
                 return yield* apiRequest("printNode", {
                     data: base64,
                     preserveSourceNewlines: options.preserveSourceNewlines,
                     neverAsciiEscape: options.neverAsciiEscape,
                     terminateUnterminatedLiterals: options.terminateUnterminatedLiterals,
+                    comments: comments.length ? comments : undefined,
                 });
             },
         );

packages/typescript/src/ast/index.ts

@@ -20,5 +20,6 @@ export * from "./is.ts";
 export * from "./jsdoc.ts";
 export * from "./scanner.ts";
 export * from "./spanMap.ts";
+export * from "./syntheticComments.ts";
 export * from "./utils.ts";
 export * from "./visitor.ts";

packages/typescript/src/ast/syntheticComments.ts

@@ -0,0 +1,66 @@
+import { SyntaxKind } from "#enums/syntaxKind";
+import type { Node } from "./ast.ts";
+
+export interface SynthesizedComment {
+    readonly kind: SyntaxKind.SingleLineCommentTrivia | SyntaxKind.MultiLineCommentTrivia;
+    readonly text: string;
+    readonly hasTrailingNewLine?: boolean | undefined;
+}
+
+const leadingComments = new WeakMap<Node, SynthesizedComment[]>();
+const trailingComments = new WeakMap<Node, SynthesizedComment[]>();
+
+function add(store: WeakMap<Node, SynthesizedComment[]>, node: Node, comment: SynthesizedComment): void {
+    const existing = store.get(node);
+    if (existing) {
+        existing.push(comment);
+    }
+    else {
+        store.set(node, [comment]);
+    }
+}
+
+/**
+ * Attaches a synthetic comment before `node`. It is emitted by the printer when the node is printed.
+ * For multi-line comments `text` excludes the `/*` and `*\/` delimiters, so a JSDoc block is `"*\n * text\n "`.
+ */
+export function addSyntheticLeadingComment<T extends Node>(
+    node: T,
+    kind: SynthesizedComment["kind"],
+    text: string,
+    hasTrailingNewLine?: boolean,
+): T {
+    add(leadingComments, node, { kind, text, hasTrailingNewLine });
+    return node;
+}
+
+/** Attaches a synthetic comment after `node`. */
+export function addSyntheticTrailingComment<T extends Node>(
+    node: T,
+    kind: SynthesizedComment["kind"],
+    text: string,
+    hasTrailingNewLine?: boolean,
+): T {
+    add(trailingComments, node, { kind, text, hasTrailingNewLine });
+    return node;
+}
+
+export function getSyntheticLeadingComments(node: Node): readonly SynthesizedComment[] | undefined {
+    return leadingComments.get(node);
+}
+
+export function getSyntheticTrailingComments(node: Node): readonly SynthesizedComment[] | undefined {
+    return trailingComments.get(node);
+}
+
+export function setSyntheticLeadingComments<T extends Node>(node: T, comments: readonly SynthesizedComment[] | undefined): T {
+    if (comments) leadingComments.set(node, [...comments]);
+    else leadingComments.delete(node);
+    return node;
+}
+
+export function setSyntheticTrailingComments<T extends Node>(node: T, comments: readonly SynthesizedComment[] | undefined): T {
+    if (comments) trailingComments.set(node, [...comments]);
+    else trailingComments.delete(node);
+    return node;
+}

packages/typescript/test/async/api.test.ts

@@ -1,5 +1,7 @@
 import {
     type __String,
+    addSyntheticLeadingComment,
+    addSyntheticTrailingComment,
     cast,
     escapeLeadingUnderscores,
     type Expression,
@@ -45,6 +47,7 @@ import {
 import {
     cloneNode,
     createArrayTypeNode,
+    createBlock,
     createFunctionTypeNode,
     createIdentifier,
     createKeywordTypeNode,
@@ -7937,6 +7940,51 @@ export const obj = { m: 1, s: "hi", b: true };
         assert.strictEqual(text, "(x: string) => number");
     });
 
+    function createConstStatement(name: string, value: string) {
+        return createVariableStatement(
+            undefined,
+            createVariableDeclarationList(
+                [createVariableDeclaration(createIdentifier(name), undefined, undefined, createNumericLiteral(value, 0))],
+                NodeFlags.Const,
+            ),
+        );
+    }
+
+    test("printNode emits synthetic leading JSDoc comment", async () => {
+        await using api = spawnAPI(emitterFiles);
+
+        const statement = addSyntheticLeadingComment(createConstStatement("x", "1"), SyntaxKind.MultiLineCommentTrivia, "*\n * Hello.\n ", true);
+        const text = await api.printer.printNode(statement);
+        assert.strictEqual(text.trimEnd(), "/**\n * Hello.\n */\nconst x = 1;");
+    });
+
+    test("printNode emits synthetic leading and trailing comments of both kinds", async () => {
+        await using api = spawnAPI(emitterFiles);
+
+        const statement = createConstStatement("y", "2");
+        addSyntheticLeadingComment(statement, SyntaxKind.SingleLineCommentTrivia, " first", true);
+        addSyntheticLeadingComment(statement, SyntaxKind.MultiLineCommentTrivia, " second ", true);
+        addSyntheticTrailingComment(statement, SyntaxKind.SingleLineCommentTrivia, " after");
+        const text = await api.printer.printNode(statement);
+        assert.strictEqual(text.trimEnd(), "// first\n/* second */\nconst y = 2; // after");
+    });
+
+    test("printNode emits synthetic comments on nested nodes", async () => {
+        await using api = spawnAPI(emitterFiles);
+
+        const inner = createConstStatement("a", "1");
+        const outer = createBlock([inner, createConstStatement("b", "2")], true);
+        addSyntheticLeadingComment(inner, SyntaxKind.SingleLineCommentTrivia, " inner", true);
+        const text = await api.printer.printNode(outer);
+        assert.match(text, /\{\n\s+\/\/ inner\n\s+const a = 1;\n\s+const b = 2;\n\}/);
+    });
+
+    test("printNode without synthetic comments is unchanged", async () => {
+        await using api = spawnAPI(emitterFiles);
+
+        assert.strictEqual((await api.printer.printNode(createConstStatement("z", "3"))).trimEnd(), "const z = 3;");
+    });
+
     test("printNode with factory-created type reference", async () => {
         await using api = spawnAPI(emitterFiles);
 

packages/typescript/test/sync/api.test.ts

@@ -8,6 +8,8 @@
 //
 import {
     type __String,
+    addSyntheticLeadingComment,
+    addSyntheticTrailingComment,
     cast,
     escapeLeadingUnderscores,
     type Expression,
@@ -53,6 +55,7 @@ import {
 import {
     cloneNode,
     createArrayTypeNode,
+    createBlock,
     createFunctionTypeNode,
     createIdentifier,
     createKeywordTypeNode,
@@ -7738,6 +7741,51 @@ export const obj = { m: 1, s: "hi", b: true };
         assert.strictEqual(text, "(x: string) => number");
     });
 
+    function createConstStatement(name: string, value: string) {
+        return createVariableStatement(
+            undefined,
+            createVariableDeclarationList(
+                [createVariableDeclaration(createIdentifier(name), undefined, undefined, createNumericLiteral(value, 0))],
+                NodeFlags.Const,
+            ),
+        );
+    }
+
+    test("printNode emits synthetic leading JSDoc comment", () => {
+        using api = spawnAPI(emitterFiles);
+
+        const statement = addSyntheticLeadingComment(createConstStatement("x", "1"), SyntaxKind.MultiLineCommentTrivia, "*\n * Hello.\n ", true);
+        const text = api.printer.printNode(statement);
+        assert.strictEqual(text.trimEnd(), "/**\n * Hello.\n */\nconst x = 1;");
+    });
+
+    test("printNode emits synthetic leading and trailing comments of both kinds", () => {
+        using api = spawnAPI(emitterFiles);
+
+        const statement = createConstStatement("y", "2");
+        addSyntheticLeadingComment(statement, SyntaxKind.SingleLineCommentTrivia, " first", true);
+        addSyntheticLeadingComment(statement, SyntaxKind.MultiLineCommentTrivia, " second ", true);
+        addSyntheticTrailingComment(statement, SyntaxKind.SingleLineCommentTrivia, " after");
+        const text = api.printer.printNode(statement);
+        assert.strictEqual(text.trimEnd(), "// first\n/* second */\nconst y = 2; // after");
+    });
+
+    test("printNode emits synthetic comments on nested nodes", () => {
+        using api = spawnAPI(emitterFiles);
+
+        const inner = createConstStatement("a", "1");
+        const outer = createBlock([inner, createConstStatement("b", "2")], true);
+        addSyntheticLeadingComment(inner, SyntaxKind.SingleLineCommentTrivia, " inner", true);
+        const text = api.printer.printNode(outer);
+        assert.match(text, /\{\n\s+\/\/ inner\n\s+const a = 1;\n\s+const b = 2;\n\}/);
+    });
+
+    test("printNode without synthetic comments is unchanged", () => {
+        using api = spawnAPI(emitterFiles);
+
+        assert.strictEqual((api.printer.printNode(createConstStatement("z", "3"))).trimEnd(), "const z = 3;");
+    });
+
     test("printNode with factory-created type reference", () => {
         using api = spawnAPI(emitterFiles);