From 7ff0ea472d28a05788dd39fbdb421c01d0fceea5 Mon Sep 17 00:00:00 2001 From: Chris Tate Date: Wed, 21 Oct 2026 14:59:26 -0500 Subject: [PATCH] Walk AST subtrees by their node ranges --- .../src/frontend/lowering/lower-stmts.test.ts | 2 +- .../compiler/src/frontend/program.test.ts | 2 +- .../src/frontend/ts7/ast-node.test.ts | 34 +++++++++ .../compiler/src/frontend/ts7/ast-node.ts | 53 +++++++++++--- .../compiler/src/frontend/ts7/ast-wire.ts | 71 +++++++++++++++++++ packages/compiler/src/frontend/ts7/ast.ts | 59 ++++++++++++++- 6 files changed, 207 insertions(+), 14 deletions(-) diff --git a/packages/compiler/src/frontend/lowering/lower-stmts.test.ts b/packages/compiler/src/frontend/lowering/lower-stmts.test.ts index 322b21b2a..8e0c809b7 100644 --- a/packages/compiler/src/frontend/lowering/lower-stmts.test.ts +++ b/packages/compiler/src/frontend/lowering/lower-stmts.test.ts @@ -54,7 +54,7 @@ function mutate(value) { captured = value; compound += value; let shadow = 0; sh ctx.bind(shadow, []); ctx.bind(untouched, []); ctx.bind(innerShadow, [at(ctx.assignments, 2)]); - const walk = vi.spyOn(ctx.file, "appendChildIndices"); + const walk = vi.spyOn(ctx.file.file.wire, "preorder"); for (let i = 0; i < 20; i++) { expect(jsBindingHasOpenWrites(ctx.lowerer, captured)).toBe(true); expect(jsBindingHasOpenWrites(ctx.lowerer, compound)).toBe(true); diff --git a/packages/compiler/src/frontend/program.test.ts b/packages/compiler/src/frontend/program.test.ts index c1bd183f2..bdd6260ed 100644 --- a/packages/compiler/src/frontend/program.test.ts +++ b/packages/compiler/src/frontend/program.test.ts @@ -351,7 +351,7 @@ test("repeated cycle edges reuse a file's binding index", () => { [peer, peerImports.map((stmt) => ({ dep: main, stmt }))], ]); const admit = makeCycleAdmission(load.program, (file) => edges.get(file) ?? []); - const walk = vi.spyOn(peer, "appendChildIndices"); + const walk = vi.spyOn(peer.file.wire, "preorder"); try { expect(admit(peer, edges.get(peer)![0]!)).toBeNull(); expect(walk).toHaveBeenCalled(); diff --git a/packages/compiler/src/frontend/ts7/ast-node.test.ts b/packages/compiler/src/frontend/ts7/ast-node.test.ts index 60bfaa005..6b2725b93 100644 --- a/packages/compiler/src/frontend/ts7/ast-node.test.ts +++ b/packages/compiler/src/frontend/ts7/ast-node.test.ts @@ -511,6 +511,40 @@ test("preorder traversal preserves depth, skipped subtrees and early termination expect(prefix).toEqual(expected.slice(0, 7)); }); +test("preorder walks from every node match a recursive visit of every file", () => { + const kinds = new Set([AstKind.Identifier, AstKind.JSDocTypeLiteral]); + for (const [name, { file, oracle }] of decoded) { + expect(file.wire.preorder(), name).not.toBeNull(); + for (let index = 1; index < file.wire.nodeCount; index++) { + if (file.wire.kind(index) === KIND_NODE_LIST) continue; + const expected: [number, number][] = []; + const selected: [number, number][] = []; + const visit = (node: OracleNode, depth: number): void => { + expected.push([node.index, depth]); + if (kinds.has(node.kind)) selected.push([node.index, depth]); + node.forEachChild((child) => { + visit(child as OracleNode, depth + 1); + }); + }; + visit(oracle.getOrCreateNodeAtIndex(index), 0); + const actual: [number, number][] = []; + walkPreorder(file.node(index), (node, depth) => { + actual.push([node.index, depth]); + }); + expect(actual, `${name}:${index}`).toEqual(expected); + const actualSelected: [number, number][] = []; + walkPreorder( + file.node(index), + (node, depth) => { + actualSelected.push([node.index, depth]); + }, + kinds, + ); + expect(actualSelected, `${name}:${index}`).toEqual(selected); + } + } +}); + test("child traversal still rejects a sibling belonging to another parent", () => { const { file, bytes } = decoded.get("main.ts")!; let parent = 0; diff --git a/packages/compiler/src/frontend/ts7/ast-node.ts b/packages/compiler/src/frontend/ts7/ast-node.ts index f337bd623..be2745616 100644 --- a/packages/compiler/src/frontend/ts7/ast-node.ts +++ b/packages/compiler/src/frontend/ts7/ast-node.ts @@ -9,6 +9,7 @@ import { import { astLineOfPosition, astLineStarts, astSkipTrivia } from "./ast-text.js"; import { AstWireFile, + PREORDER_LIST, parseAstNodeHandle, type AstFileReference, type AstNodeHandle, @@ -359,6 +360,32 @@ export class AstNode { forEachChild( visitNode: (node: AstNode) => T, visitList?: (nodes: AstNodeList) => T, + ): T | undefined { + const wire = this.file.wire; + const tree = wire.preorder(); + if (tree === null) return this.forEachLinkedChild(visitNode, visitList); + // In the checked preorder tree, the first child follows its parent and + // each later child starts where the previous child's subtree ends. + const kinds = tree.kinds; + const ends = tree.ends; + const stop = ends[this.index]!; + for (let index = this.index + 1; index < stop; index = ends[index]!) { + const kind = kinds[index]!; + if (kind === PREORDER_LIST) { + const result = this.visitChildList(index, visitNode, visitList); + if (result) return result; + } else if (kind !== AstKind.JSDoc) { + const result = visitNode(this.file.node(index)); + if (result) return result; + } + } + return undefined; + } + + /** forEachChild through the sibling and parent links. */ + private forEachLinkedChild( + visitNode: (node: AstNode) => T, + visitList?: (nodes: AstNodeList) => T, ): T | undefined { const wire = this.file.wire; // Walk the immutable sibling links directly. Repeated semantic scans @@ -368,16 +395,8 @@ export class AstNode { throw new AstDecodeError("sibling belongs to another parent"); const kind = wire.kind(index); if (kind === KIND_NODE_LIST) { - const list = this.file.list(index); - if (visitList !== undefined) { - const result = visitList(list); - if (result) return result; - } else { - for (const node of list.nodes) { - const result = visitNode(node); - if (result) return result; - } - } + const result = this.visitChildList(index, visitNode, visitList); + if (result) return result; } else if (kind !== AstKind.JSDoc) { const result = visitNode(this.file.node(index)); if (result) return result; @@ -386,6 +405,20 @@ export class AstNode { return undefined; } + private visitChildList( + index: number, + visitNode: (node: AstNode) => T, + visitList?: (nodes: AstNodeList) => T, + ): T | undefined { + const list = this.file.list(index); + if (visitList !== undefined) return visitList(list); + for (const node of list.nodes) { + const result = visitNode(node); + if (result) return result; + } + return undefined; + } + get jsDoc(): AstNode[] | undefined { const docs: AstNode[] = []; for (const index of this.file.wire.children(this.index)) { diff --git a/packages/compiler/src/frontend/ts7/ast-wire.ts b/packages/compiler/src/frontend/ts7/ast-wire.ts index a2466a743..ff70148df 100644 --- a/packages/compiler/src/frontend/ts7/ast-wire.ts +++ b/packages/compiler/src/frontend/ts7/ast-wire.ts @@ -32,6 +32,16 @@ import { astChildOrder, } from "./ast-schema.generated.js"; +/** A node table read as a preorder tree (AstWireFile.preorder): each + * node's kind, with PREORDER_LIST for lists, and the end of its subtree. */ +export interface AstPreorder { + readonly kinds: Uint16Array; + readonly ends: Uint32Array; +} + +/** The kind AstPreorder records for a list (KIND_NODE_LIST on the wire). */ +export const PREORDER_LIST = 0xffff; + export interface AstFileReference { pos: number; end: number; @@ -125,6 +135,7 @@ export class AstWireFile { private readonly words: Uint32Array; /** Decoded strings by string word index / 2. */ private readonly stringCache: (string | undefined)[]; + private preorderTree: AstPreorder | null | undefined; constructor(private readonly bytes: Uint8Array) { astBounds(bytes, 0, HEADER_SIZE); @@ -317,6 +328,66 @@ export class AstWireFile { return children; } + /** The node table as a preorder tree, built on first use: every node's + * subtree is the index range up to `ends[index]`, and `kinds[index]` is its + * kind. The pass checks the links the linked walk checks one child at a + * time: each node's parent is open, the first child follows its parent, + * each sibling link names the end of the previous sibling's subtree, a + * last child has no sibling link, and each list holds `data` nodes and no + * lists. null when any check fails; walks then follow the links and fail + * where they read the inconsistent node. */ + preorder(): AstPreorder | null { + if (this.preorderTree === undefined) this.preorderTree = this.buildPreorder(); + return this.preorderTree; + } + + private buildPreorder(): AstPreorder | null { + const count = this.nodeCount; + const words = this.words; + const kinds = new Uint16Array(count); + const ends = new Uint32Array(count); + // The nodes whose subtrees contain the current index, innermost last, + // with the number of children each has so far. + const open: number[] = [1]; + const children: number[] = [0]; + const rootKind = words[NODE_WORDS + WORD_KIND]!; + if (rootKind >= PREORDER_LIST) return null; + kinds[1] = rootKind; + for (let index = 2; index <= count; index++) { + const word = index * NODE_WORDS; + const parent = index < count ? words[word + WORD_PARENT]! : 0; + let previous = 0; + while (open.length > 0 && open[open.length - 1] !== parent) { + const closed = open.pop()!; + const closedChildren = children.pop()!; + ends[closed] = index; + const closedWord = closed * NODE_WORDS; + if (kinds[closed] === PREORDER_LIST && closedChildren !== words[closedWord + WORD_DATA]!) + return null; + if (open.length > 0 && open[open.length - 1] === parent) previous = closed; + else if (closed !== 1 && words[closedWord + WORD_NEXT]! !== 0) return null; + } + if (index === count) break; + if (open.length === 0 || parent >= index) return null; + if ( + previous !== 0 ? words[previous * NODE_WORDS + WORD_NEXT]! !== index : index !== parent + 1 + ) + return null; + const kind = words[word + WORD_KIND]!; + if (kind === KIND_NODE_LIST) { + if (kinds[parent] === PREORDER_LIST) return null; + kinds[index] = PREORDER_LIST; + } else { + if (kind >= PREORDER_LIST) return null; + kinds[index] = kind; + } + children[children.length - 1] = children[children.length - 1]! + 1; + open.push(index); + children.push(0); + } + return { kinds, ends }; + } + /** The semantic walk flattens lists and omits attached JSDoc trees. * Validate the same list links as list(), without materializing nodes * that a selective syntax scan will never inspect. */ diff --git a/packages/compiler/src/frontend/ts7/ast.ts b/packages/compiler/src/frontend/ts7/ast.ts index a92cfafd8..369213d89 100644 --- a/packages/compiler/src/frontend/ts7/ast.ts +++ b/packages/compiler/src/frontend/ts7/ast.ts @@ -47,6 +47,8 @@ import { } from "./ast-guards.generated.js"; import type { Ts7DiagnosticData as Diagnostic } from "./session-schema.generated.js"; import { ModifierFlags, SyntaxKind } from "./enums.js"; +import { AstKind } from "./ast-schema.generated.js"; +import { PREORDER_LIST } from "./ast-wire.js"; /* ---- everything 7 kept under the same name ---- */ @@ -101,7 +103,7 @@ export function forEachChild( return node.forEachChild(cbNode, cbNodes); } -/** Preorder walk over a whole subtree with an EXPLICIT stack. The idiomatic +/** Preorder walk over a whole subtree without recursion. The idiomatic * recursive forEachChild visit costs two JS frames per AST level, so a * pathologically deep tree — the ~6500-term left-nested binary chain of the * binderBinaryExpressionStress corpus pair — overflows the stack as an ICE @@ -110,11 +112,64 @@ export function forEachChild( * detection) walk this way instead, so absurd nesting always reaches the * named fence. The callback may answer "skip" to leave a node's children * unvisited (scope and depth cut-offs — `depth` counts levels below the - * root) or "stop" to end the walk early. */ + * root) or "stop" to end the walk early. + * + * The wire stores nodes in preorder, so the walk visits the subtree's index + * range in order: it enters lists without a level of their own, and passes + * over attached JSDoc trees and skipped subtrees by their end index. A table + * that is not a checked preorder tree (AstWireFile.preorder) is walked + * through its links instead. */ export function walkPreorder( root: Node, cb: (node: Node, depth: number) => void | "skip" | "stop", kinds?: ReadonlySet, +): void { + const file = root.file; + const tree = file.wire.preorder(); + if (tree === null) { + walkLinkedPreorder(root, cb, kinds); + return; + } + const nodeKinds = tree.kinds; + const ends = tree.ends; + const start = root.index; + const stop = ends[start]!; + // The subtree ends of the visited nodes that contain the current index. + const open: number[] = []; + let index = start; + while (index < stop) { + const kind = nodeKinds[index]!; + if (kind === PREORDER_LIST) { + index++; + continue; + } + if ( + kind === AstKind.JSDoc && + index !== start && + nodeKinds[file.wire.parent(index)] !== PREORDER_LIST + ) { + index = ends[index]!; + continue; + } + while (open.length > 0 && open[open.length - 1]! <= index) open.pop(); + const node = kinds === undefined || kinds.has(kind) ? file.node(index) : undefined; + const verdict = node === undefined ? undefined : cb(node, open.length); + if (verdict === "stop") return; + if (verdict === "skip") { + index = ends[index]!; + continue; + } + open.push(ends[index]!); + index++; + } +} + +/** walkPreorder through the sibling and parent links, with an explicit + * stack. */ +function walkLinkedPreorder( + root: Node, + cb: (node: Node, depth: number) => void | "skip" | "stop", + kinds?: ReadonlySet, ): void { const file = root.file; const stack: number[] = [root.index];