Create a SourceFile-level indirection on children maps, store SyntaxList children directly on nodes. (#59154)

Author: DanielRosenwasser

src/compiler/factory/nodeChildren.ts

@@ -1,25 +1,61 @@
 import {
+    Debug,
     emptyArray,
     isNodeKind,
     Node,
+    SourceFileLike,
+    SyntaxKind,
+    SyntaxList,
 } from "../_namespaces/ts.js";
 
-const nodeChildren = new WeakMap<Node, readonly Node[] | undefined>();
+const sourceFileToNodeChildren = new WeakMap<SourceFileLike, WeakMap<Node, readonly Node[] | undefined>>();
 
 /** @internal */
-export function getNodeChildren(node: Node): readonly Node[] | undefined {
-    if (!isNodeKind(node.kind)) return emptyArray;
+export function getNodeChildren(node: Node, sourceFile: SourceFileLike): readonly Node[] | undefined {
+    const kind = node.kind;
+    if (!isNodeKind(kind)) {
+        return emptyArray;
+    }
+    if (kind === SyntaxKind.SyntaxList) {
+        return (node as SyntaxList)._children;
+    }
 
-    return nodeChildren.get(node);
+    return sourceFileToNodeChildren.get(sourceFile)?.get(node);
 }
 
 /** @internal */
-export function setNodeChildren(node: Node, children: readonly Node[]): readonly Node[] {
-    nodeChildren.set(node, children);
+export function setNodeChildren(node: Node, sourceFile: SourceFileLike, children: readonly Node[]): readonly Node[] {
+    if (node.kind === SyntaxKind.SyntaxList) {
+        // SyntaxList children are always eagerly created in the process of
+        // creating their parent's `children` list. We shouldn't need to set them here.
+        Debug.fail("Should not need to re-set the children of a SyntaxList.");
+    }
+
+    let map = sourceFileToNodeChildren.get(sourceFile);
+    if (map === undefined) {
+        map = new WeakMap();
+        sourceFileToNodeChildren.set(sourceFile, map);
+    }
+    map.set(node, children);
     return children;
 }
 
 /** @internal */
-export function unsetNodeChildren(node: Node) {
-    nodeChildren.delete(node);
+export function unsetNodeChildren(node: Node, origSourceFile: SourceFileLike) {
+    if (node.kind === SyntaxKind.SyntaxList) {
+        // Syntax lists are synthesized and we store their children directly on them.
+        // They are a special case where we expect incremental parsing to toss them away entirely
+        // if a change intersects with their containing parents.
+        Debug.fail("Did not expect to unset the children of a SyntaxList.");
+    }
+    sourceFileToNodeChildren.get(origSourceFile)?.delete(node);
+}
+
+/** @internal */
+export function transferSourceFileChildren(sourceFile: SourceFileLike, targetSourceFile: SourceFileLike) {
+    const map = sourceFileToNodeChildren.get(sourceFile);
+    if (map !== undefined) {
+        sourceFileToNodeChildren.delete(sourceFile);
+        sourceFileToNodeChildren.set(targetSourceFile, map);
+    }
 }

src/compiler/factory/nodeFactory.ts

@@ -388,7 +388,6 @@ import {
     setEmitFlags,
     setIdentifierAutoGenerate,
     setIdentifierTypeArguments,
-    setNodeChildren,
     setParent,
     setTextRange,
     ShorthandPropertyAssignment,
@@ -6212,7 +6211,7 @@ export function createNodeFactory(flags: NodeFactoryFlags, baseFactory: BaseNode
     // @api
     function createSyntaxList(children: readonly Node[]) {
         const node = createBaseNode<SyntaxList>(SyntaxKind.SyntaxList);
-        setNodeChildren(node, children);
+        node._children = children;
         return node;
     }
 

src/compiler/parser.ts

@@ -369,6 +369,7 @@ import {
     tokenIsIdentifierOrKeywordOrGreaterThan,
     tokenToString,
     tracing,
+    transferSourceFileChildren,
     TransformFlags,
     TryStatement,
     TupleTypeNode,
@@ -9969,6 +9970,7 @@ namespace IncrementalParser {
             aggressiveChecks,
         );
         result.impliedNodeFormat = sourceFile.impliedNodeFormat;
+        transferSourceFileChildren(sourceFile, result);
         return result;
     }
 
@@ -10021,9 +10023,9 @@ namespace IncrementalParser {
         }
     }
 
-    function moveElementEntirelyPastChangeRange(element: Node, isArray: false, delta: number, oldText: string, newText: string, aggressiveChecks: boolean): void;
-    function moveElementEntirelyPastChangeRange(element: NodeArray<Node>, isArray: true, delta: number, oldText: string, newText: string, aggressiveChecks: boolean): void;
-    function moveElementEntirelyPastChangeRange(element: Node | NodeArray<Node>, isArray: boolean, delta: number, oldText: string, newText: string, aggressiveChecks: boolean) {
+    function moveElementEntirelyPastChangeRange(element: Node, origSourceFile: SourceFile, isArray: false, delta: number, oldText: string, newText: string, aggressiveChecks: boolean): void;
+    function moveElementEntirelyPastChangeRange(element: NodeArray<Node>, origSourceFile: SourceFile, isArray: true, delta: number, oldText: string, newText: string, aggressiveChecks: boolean): void;
+    function moveElementEntirelyPastChangeRange(element: Node | NodeArray<Node>, origSourceFile: SourceFile, isArray: boolean, delta: number, oldText: string, newText: string, aggressiveChecks: boolean) {
         if (isArray) {
             visitArray(element as NodeArray<Node>);
         }
@@ -10040,7 +10042,7 @@ namespace IncrementalParser {
 
             // Ditch any existing LS children we may have created.  This way we can avoid
             // moving them forward.
-            unsetNodeChildren(node);
+            unsetNodeChildren(node, origSourceFile);
 
             setTextRangePosEnd(node, node.pos + delta, node.end + delta);
 
@@ -10187,7 +10189,7 @@ namespace IncrementalParser {
             if (child.pos > changeRangeOldEnd) {
                 // Node is entirely past the change range.  We need to move both its pos and
                 // end, forward or backward appropriately.
-                moveElementEntirelyPastChangeRange(child, /*isArray*/ false, delta, oldText, newText, aggressiveChecks);
+                moveElementEntirelyPastChangeRange(child, sourceFile, /*isArray*/ false, delta, oldText, newText, aggressiveChecks);
                 return;
             }
 
@@ -10197,7 +10199,7 @@ namespace IncrementalParser {
             const fullEnd = child.end;
             if (fullEnd >= changeStart) {
                 markAsIntersectingIncrementalChange(child);
-                unsetNodeChildren(child);
+                unsetNodeChildren(child, sourceFile);
 
                 // Adjust the pos or end (or both) of the intersecting element accordingly.
                 adjustIntersectingElement(child, changeStart, changeRangeOldEnd, changeRangeNewEnd, delta);
@@ -10220,7 +10222,7 @@ namespace IncrementalParser {
             if (array.pos > changeRangeOldEnd) {
                 // Array is entirely after the change range.  We need to move it, and move any of
                 // its children.
-                moveElementEntirelyPastChangeRange(array, /*isArray*/ true, delta, oldText, newText, aggressiveChecks);
+                moveElementEntirelyPastChangeRange(array, sourceFile, /*isArray*/ true, delta, oldText, newText, aggressiveChecks);
                 return;
             }
 

src/compiler/types.ts

@@ -9816,6 +9816,12 @@ export interface DiagnosticCollection {
 // SyntaxKind.SyntaxList
 export interface SyntaxList extends Node {
     kind: SyntaxKind.SyntaxList;
+
+    // Unlike other nodes which may or may not have their child nodes calculated,
+    // the entire purpose of a SyntaxList is to hold child nodes.
+    // Instead of using the WeakMap machinery in `nodeChildren.ts`,
+    // we just store the children directly on the SyntaxList.
+    /** @internal */ _children: readonly Node[];
 }
 
 // dprint-ignore

src/compiler/utilities.ts

@@ -1187,7 +1187,7 @@ export function getTokenPosOfNode(node: Node, sourceFile?: SourceFileLike, inclu
 
     if (isJSDocNode(node) || node.kind === SyntaxKind.JsxText) {
         // JsxText cannot actually contain comments, even though the scanner will think it sees comments
-        return skipTrivia((sourceFile || getSourceFileOfNode(node)).text, node.pos, /*stopAfterLineBreak*/ false, /*stopAtComments*/ true);
+        return skipTrivia((sourceFile ?? getSourceFileOfNode(node)).text, node.pos, /*stopAfterLineBreak*/ false, /*stopAtComments*/ true);
     }
 
     if (includeJsDoc && hasJSDocNodes(node)) {
@@ -1199,14 +1199,15 @@ export function getTokenPosOfNode(node: Node, sourceFile?: SourceFileLike, inclu
     // trivia for the list, we may have skipped the JSDocComment as well. So we should process its
     // first child to determine the actual position of its first token.
     if (node.kind === SyntaxKind.SyntaxList) {
-        const first = firstOrUndefined(getNodeChildren(node));
+        sourceFile ??= getSourceFileOfNode(node);
+        const first = firstOrUndefined(getNodeChildren(node, sourceFile));
         if (first) {
             return getTokenPosOfNode(first, sourceFile, includeJsDoc);
         }
     }
 
     return skipTrivia(
-        (sourceFile || getSourceFileOfNode(node)).text,
+        (sourceFile ?? getSourceFileOfNode(node)).text,
         node.pos,
         /*stopAfterLineBreak*/ false,
         /*stopAtComments*/ false,

src/services/services.ts

@@ -457,9 +457,9 @@ class NodeObject<TKind extends SyntaxKind> implements Node {
         return this.getChildren(sourceFile)[index];
     }
 
-    public getChildren(sourceFile?: SourceFileLike): readonly Node[] {
+    public getChildren(sourceFile: SourceFileLike = getSourceFileOfNode(this)): readonly Node[] {
         this.assertHasRealPosition("Node without a real position cannot be scanned and thus has no token nodes - use forEachChild and collect the result if that's fine");
-        return getNodeChildren(this) ?? setNodeChildren(this, createChildren(this, sourceFile));
+        return getNodeChildren(this, sourceFile) ?? setNodeChildren(this, sourceFile, createChildren(this, sourceFile));
     }
 
     public getFirstToken(sourceFile?: SourceFileLike): Node | undefined {
@@ -558,7 +558,7 @@ function createSyntaxList(nodes: NodeArray<Node>, parent: Node): Node {
         pos = node.end;
     }
     addSyntheticNodes(children, pos, nodes.end, parent);
-    setNodeChildren(list, children);
+    list._children = children;
     return list;
 }