API: add getChildren and token getters to Node

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

@@ -1,6 +1,9 @@
 // Code generated by tools/scripts/tsc/generate-encoder.ts. DO NOT EDIT.
 
 import {
+    getChildren,
+    getFirstToken,
+    getLastToken,
     getTokenPosOfNode,
     ModifierFlags,
     type Node,
@@ -319,6 +322,26 @@ export class RemoteNode extends RemoteNodeBase implements Node {
         return sourceFile.text.substring(this.getStart(sourceFile), this.end);
     }
 
+    getChildCount(sourceFile?: SourceFile): number {
+        return this.getChildren(sourceFile).length;
+    }
+
+    getChildAt(index: number, sourceFile?: SourceFile): Node {
+        return this.getChildren(sourceFile)[index];
+    }
+
+    getChildren(sourceFile?: SourceFile): readonly Node[] {
+        return getChildren(this as unknown as Node, sourceFile ?? this.getSourceFile());
+    }
+
+    getFirstToken(sourceFile?: SourceFile): Node | undefined {
+        return getFirstToken(this as unknown as Node, sourceFile ?? this.getSourceFile());
+    }
+
+    getLastToken(sourceFile?: SourceFile): Node | undefined {
+        return getLastToken(this as unknown as Node, sourceFile ?? this.getSourceFile());
+    }
+
     protected getString(index: number): string {
         const offsetStringTableOffsets = this.sourceFile._offsetStringTableOffsets;
         const start = this.view.getUint32(offsetStringTableOffsets + index * 4, true);

packages/typescript/src/ast/ast.ts

@@ -100,6 +100,11 @@ export interface Node extends ReadonlyTextRange {
     getLeadingTriviaWidth(sourceFile?: SourceFile): number;
     getFullText(sourceFile?: SourceFile): string;
     getText(sourceFile?: SourceFile): string;
+    getChildCount(sourceFile?: SourceFile): number;
+    getChildAt(index: number, sourceFile?: SourceFile): Node;
+    getChildren(sourceFile?: SourceFile): readonly Node[];
+    getFirstToken(sourceFile?: SourceFile): Node | undefined;
+    getLastToken(sourceFile?: SourceFile): Node | undefined;
 }
 
 export interface FileReference extends TextRange {
@@ -159,6 +164,8 @@ export interface SourceFile extends Node {
     getPositionOfLineAndCharacter(line: number, character: number): number;
     /** @internal */
     tokenCache?: Map<string, Node>;
+    /** @internal */
+    childrenCache?: WeakMap<Node, readonly Node[]>;
 }
 
 // ── Token hierarchy ──

packages/typescript/src/ast/astnav.ts

@@ -6,7 +6,10 @@ import type {
     NodeArray,
     SourceFile,
 } from "./ast.ts";
-import { createToken } from "./factory.generated.ts";
+import {
+    createSyntaxList,
+    createToken,
+} from "./factory.generated.ts";
 import {
     isJSDocNodeKind,
     isKeywordKind,
@@ -616,6 +619,149 @@ function getOrCreateToken(sourceFile: SourceFile, kind: SyntaxKind, pos: number,
     return token;
 }
 
+const emptyArray: readonly Node[] = [];
+
+function assertHasRealPosition(node: Node): void {
+    if (node.pos < 0 || node.end < 0) {
+        throw new Error("Node without a real position cannot be scanned and thus has no token nodes - use forEachChild and collect the result if that's fine");
+    }
+}
+
+export function getChildren(node: Node, sourceFile: SourceFile = node.getSourceFile()): readonly Node[] {
+    // A SyntaxList already holds its (pre-materialized) children.
+    if (node.kind === SyntaxKind.SyntaxList) {
+        return (node as unknown as { children: readonly Node[]; }).children;
+    }
+
+    if (isTokenKind(node.kind)) {
+        // EndOfFile may carry leading JSDoc; every other token has no children. The EndOfFile
+        // result must go through the cache: remote nodes rebuild .jsDoc on every access.
+        if (node.kind !== SyntaxKind.EndOfFile) {
+            return emptyArray;
+        }
+    }
+    else {
+        assertHasRealPosition(node);
+    }
+
+    const cache = (sourceFile.childrenCache ??= new WeakMap<Node, readonly Node[]>());
+    const cached = cache.get(node);
+
+    if (cached !== undefined) {
+        return cached;
+    }
+
+    const children = node.kind === SyntaxKind.EndOfFile
+        ? node.jsDoc ?? emptyArray
+        : createChildren(node, sourceFile);
+    cache.set(node, children);
+    return children;
+}
+
+function createChildren(node: Node, sourceFile: SourceFile): readonly Node[] {
+    const children: Node[] = [];
+
+    // Inside a JSDoc comment there are no real tokens to synthesize.
+    if (shouldSkipChild(node)) {
+        node.forEachChild(child => void children.push(child));
+        return children;
+    }
+
+    let pos = node.pos;
+    const consumed = new Set<Node>();
+    const processNode = (child: Node): undefined => {
+        if (consumed.has(child)) {
+            return;
+        }
+        addSyntheticNodes(children, pos, child.pos, node, sourceFile);
+        children.push(child);
+        pos = child.end;
+    };
+    const processNodes = (nodes: NodeArray<Node>): undefined => {
+        addSyntheticNodes(children, pos, nodes.pos, node, sourceFile);
+        children.push(createSyntaxListNode(nodes, node, sourceFile));
+        pos = nodes.end;
+        for (const element of nodes) {
+            consumed.add(element);
+        }
+    };
+
+    // JSDoc attached to the node is leading content, processed first.
+    if (node.jsDoc) {
+        for (const jsDoc of node.jsDoc) {
+            processNode(jsDoc);
+        }
+    }
+    pos = node.pos;
+    node.forEachChild(processNode, processNodes);
+    addSyntheticNodes(children, pos, node.end, node, sourceFile);
+    return children;
+}
+
+function addSyntheticNodes(children: Node[], pos: number, end: number, parent: Node, sourceFile: SourceFile): void {
+    if (pos >= end) {
+        return;
+    }
+    const scanner = getScannerForSourceFile(sourceFile, pos);
+    while (pos < end) {
+        const token = scanner.getToken();
+        const tokenEnd = scanner.getTokenEnd();
+        if (tokenEnd <= end) {
+            // An identifier should never appear as trivia between AST children; skip defensively.
+            if (token !== SyntaxKind.Identifier) {
+                children.push(getOrCreateToken(sourceFile, token, pos, tokenEnd, parent, scanner.getTokenFlags()));
+            }
+        }
+        pos = tokenEnd;
+        if (token === SyntaxKind.EndOfFile) {
+            break;
+        }
+        scanner.scan();
+    }
+}
+
+function createSyntaxListNode(nodes: NodeArray<Node>, parent: Node, sourceFile: SourceFile): Node {
+    const listChildren: Node[] = [];
+    let pos = nodes.pos;
+    for (const child of nodes) {
+        addSyntheticNodes(listChildren, pos, child.pos, parent, sourceFile);
+        listChildren.push(child);
+        pos = child.end;
+    }
+    addSyntheticNodes(listChildren, pos, nodes.end, parent, sourceFile);
+    const list = createSyntaxList(listChildren) as Mutable<Node>;
+    list.pos = nodes.pos;
+    list.end = nodes.end;
+    list.parent = parent;
+    return list as Node;
+}
+
+export function getFirstToken(node: Node, sourceFile: SourceFile = node.getSourceFile()): Node | undefined {
+    if (isTokenKind(node.kind)) {
+        return undefined;
+    }
+    assertHasRealPosition(node);
+    const children = getChildren(node, sourceFile);
+    const child = children.find(kid => kid.kind < SyntaxKind.FirstJSDocNode || kid.kind > SyntaxKind.LastJSDocNode);
+    if (child === undefined) {
+        return undefined;
+    }
+    return child.kind < SyntaxKind.FirstNode ? child : getFirstToken(child, sourceFile);
+}
+
+export function getLastToken(node: Node, sourceFile: SourceFile = node.getSourceFile()): Node | undefined {
+    if (isTokenKind(node.kind)) {
+        return undefined;
+    }
+    assertHasRealPosition(node);
+    const children = getChildren(node, sourceFile);
+    const child = children.length ? children[children.length - 1] : undefined;
+    if (child === undefined) {
+        return undefined;
+    }
+    return child.kind < SyntaxKind.FirstNode ? child : getLastToken(child, sourceFile);
+}
+
 /** Binary search a node list for the node containing position. */
 function binarySearchNodeList(
     nodes: NodeArray<Node>,

packages/typescript/src/ast/factory.generated.ts

@@ -258,7 +258,12 @@ import type {
     WithStatement,
     YieldExpression,
 } from "./ast.ts";
-import { getTokenPosOfNode } from "./astnav.ts";
+import {
+    getChildren,
+    getFirstToken,
+    getLastToken,
+    getTokenPosOfNode,
+} from "./astnav.ts";
 import { cloneSourceFileData } from "./utils.ts";
 import {
     forEachChildOfJSDocParameterTag,
@@ -724,6 +729,26 @@ export class NodeObject {
         sourceFile ??= this.getSourceFile();
         return sourceFile.text.substring(this.getStart(sourceFile), this.end);
     }
+
+    getChildCount(sourceFile?: SourceFile): number {
+        return this.getChildren(sourceFile).length;
+    }
+
+    getChildAt(index: number, sourceFile?: SourceFile): Node {
+        return this.getChildren(sourceFile)[index];
+    }
+
+    getChildren(sourceFile?: SourceFile): readonly Node[] {
+        return getChildren(this as unknown as Node, sourceFile ?? this.getSourceFile());
+    }
+
+    getFirstToken(sourceFile?: SourceFile): Node | undefined {
+        return getFirstToken(this as unknown as Node, sourceFile ?? this.getSourceFile());
+    }
+
+    getLastToken(sourceFile?: SourceFile): Node | undefined {
+        return getLastToken(this as unknown as Node, sourceFile ?? this.getSourceFile());
+    }
 }
 
 function isNodeArray<T extends Node>(array: readonly T[]): array is NodeArray<T> {