Devirtualize ForEachChild and IterChildren

_scripts/generate-go-ast.ts

@@ -478,6 +478,54 @@ function generateForEachChild(w: CodeWriter, node: NodeType) {
     w.write("");
 }
 
+// ── Generate ForEachChild dispatch ─────────────────────────────────────────
+
+// Determines whether a node has a meaningful ForEachChild implementation (i.e.
+// one that is not the no-op inherited from NodeDefault). This is true for any
+// generated node with child members, or for hand-written nodes (e.g. SourceFile)
+// which define ForEachChild manually in ast.go.
+function hasForEachChild(node: NodeType): boolean {
+    if (node.handWritten) return true;
+    return schemaMembers(node).some(m => m.isChild());
+}
+
+// Generates a (*Node).ForEachChild method that dispatches on node.Kind to the
+// concrete node type's ForEachChild method.
+//
+// This deliberately avoids calling node.data.ForEachChild(v) through the
+// nodeData interface. An interface (or other indirect) call is opaque to escape
+// analysis, which must then assume the visitor `v` escapes; that forces caller
+// closures — and any locals they capture — onto the heap. Dispatching through a
+// Kind switch to a statically-resolved concrete method lets escape analysis
+// prove the visitor does not escape, keeping caller closures on the stack. The
+// integer switch over Kind also compiles to a jump table, making dispatch
+// cheaper than the interface call it replaces.
+//
+// Kinds whose node has no children fall through to `default` and return false,
+// matching NodeDefault.ForEachChild.
+function generateForEachChildDispatch(w: CodeWriter) {
+    w.write("func (n *Node) ForEachChild(v Visitor) bool {");
+    w.push();
+    w.write("switch n.Kind {");
+    for (const node of api.nodes()) {
+        if (!hasForEachChild(node)) continue;
+        const kinds = node.allKinds().map(k => k.formatGoConstant());
+        if (kinds.length === 0) continue;
+        w.write(`case ${kinds.join(", ")}:`);
+        w.push();
+        w.write(`return n.data.(*${node.name}).ForEachChild(v)`);
+        w.pop();
+    }
+    w.write("default:");
+    w.push();
+    w.write("return false");
+    w.pop();
+    w.write("}");
+    w.pop();
+    w.write("}");
+    w.write("");
+}
+
 // ── Generate VisitEachChild() ──────────────────────────────────────────────
 
 function generateVisitEachChild(w: CodeWriter, node: NodeType) {
@@ -824,6 +872,13 @@ function generate(): string {
         }
     }
 
+    // ForEachChild dispatch
+    w.write("// ──────────────────────────────────────────────────────────────────────");
+    w.write("// ForEachChild dispatch");
+    w.write("// ──────────────────────────────────────────────────────────────────────");
+    w.write("");
+    generateForEachChildDispatch(w);
+
     // As*() casts
     w.write("// ──────────────────────────────────────────────────────────────────────");
     w.write("// As*() cast methods");

internal/ast/ast.go

@@ -188,11 +188,21 @@ type Node struct {
 // type switches. Either approach is fine. Interface methods are likely more performant, but have higher
 // code size costs because we have hundreds of implementations of the NodeData interface.
 
-func (n *Node) AsNode() *Node                             { return n }
-func (n *Node) Pos() int                                  { return n.Loc.Pos() }
-func (n *Node) End() int                                  { return n.Loc.End() }
-func (n *Node) ForEachChild(v Visitor) bool               { return n.data.ForEachChild(v) }
-func (n *Node) IterChildren() iter.Seq[*Node]             { return n.data.IterChildren() }
+func (n *Node) AsNode() *Node { return n }
+func (n *Node) Pos() int      { return n.Loc.Pos() }
+func (n *Node) End() int      { return n.Loc.End() }
+func (n *Node) IterChildren() iter.Seq[*Node] {
+	// Implemented directly (rather than through the nodeData interface) so that the
+	// returned iterator and the visitor closure it passes to ForEachChild do not
+	// escape: an interface call is opaque to escape analysis. `true` stops a TS
+	// visitor early, whereas `false` stops a Go iterator yield, so the result is
+	// inverted.
+	return func(yield func(*Node) bool) {
+		n.ForEachChild(func(child *Node) bool {
+			return !yield(child)
+		})
+	}
+}
 func (n *Node) Clone(f NodeFactoryCoercible) *Node        { return n.data.Clone(f) }
 func (n *Node) VisitEachChild(v *NodeVisitor) *Node       { return n.data.VisitEachChild(v) }
 func (n *Node) Name() *DeclarationName                    { return n.data.Name() }
@@ -1170,7 +1180,6 @@ func (n *Node) AsFlowReduceLabelData() *FlowReduceLabelData {
 type nodeData interface {
 	AsNode() *Node
 	ForEachChild(v Visitor) bool
-	IterChildren() iter.Seq[*Node]
 	VisitEachChild(v *NodeVisitor) *Node
 	Clone(v NodeFactoryCoercible) *Node
 	Name() *DeclarationName
@@ -1197,21 +1206,8 @@ type NodeDefault struct {
 	Node
 }
 
-func invert(yield func(v *Node) bool) Visitor {
-	return func(n *Node) bool {
-		return !yield(n)
-	}
-}
-
 func (node *NodeDefault) AsNode() *Node               { return &node.Node }
 func (node *NodeDefault) ForEachChild(v Visitor) bool { return false }
-func (node *NodeDefault) forEachChildIter(yield func(v *Node) bool) {
-	node.data.ForEachChild(invert(yield)) // `true` is return early for a ts visitor, `false` is return early for a go iterator yield function
-}
-
-func (node *NodeDefault) IterChildren() iter.Seq[*Node] {
-	return node.forEachChildIter
-}
 
 func (node *NodeDefault) VisitEachChild(v *NodeVisitor) *Node                   { return node.AsNode() }
 func (node *NodeDefault) Clone(v NodeFactoryCoercible) *Node                    { return nil }