Walk AST subtrees by their node ranges

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);

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();

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<AstNode["kind"]>([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;

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<T>(
     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<T>(
+    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<T>(
+    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)) {

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. */

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<T>(
   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<T>(
  * 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<SyntaxKind>,
+): 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<SyntaxKind>,
 ): void {
   const file = root.file;
   const stack: number[] = [root.index];