Reject mismatched kinds in multi-kind node factories

Node casts now trust Kind to recover the concrete struct, so a factory
that accepts any kind could hand out a header whose cast reinterprets
the wrong allocation. Previously the payload type assertion caught
this. Restore the guarantee by rejecting kinds the struct cannot
represent at construction time.

Author: jakebailey

tools/scripts/tsc/generate-go-ast.ts

@@ -374,27 +374,31 @@ function generateBaseStructDefs(w: CodeWriter) {
 
 // ── Generate As*() cast methods ────────────────────────────────────────────
 
-function generateAsCast(w: CodeWriter, node: NodeType) {
-    const structName = node.name;
+function generateKindCheck(w: CodeWriter, node: NodeType, expression: string, message: string) {
     const kinds = canonicalNodesByKind().get(node) ?? [];
-    w.write(`func (n *Node) As${structName}() *${structName} {`);
-    w.push();
     if (kinds.length === 1) {
-        w.write(`if n.Kind != ${kinds[0]} {`);
+        w.write(`if ${expression} != ${kinds[0]} {`);
         w.push();
-        w.write(`panic("Invalid node cast to ${structName}")`);
+        w.write(`panic("${message}")`);
         w.pop();
         w.write("}");
     }
     else {
-        w.write("switch n.Kind {");
+        w.write(`switch ${expression} {`);
         w.write(`case ${kinds.join(", ")}:`);
         w.write("default:");
         w.push();
-        w.write(`panic("Invalid node cast to ${structName}")`);
+        w.write(`panic("${message}")`);
         w.pop();
         w.write("}");
     }
+}
+
+function generateAsCast(w: CodeWriter, node: NodeType) {
+    const structName = node.name;
+    w.write(`func (n *Node) As${structName}() *${structName} {`);
+    w.push();
+    generateKindCheck(w, node, "n.Kind", `Invalid node cast to ${structName}`);
     w.write(`return (*${structName})(unsafe.Pointer(n))`);
     w.pop();
     w.write("}");
@@ -444,6 +448,10 @@ function emitNewFactory(
     w.write(`func (f *NodeFactory) ${funcName}(${params}) *Node {`);
     w.push();
 
+    if (kindMember) {
+        generateKindCheck(w, node, kindMember.goParamName(), `Invalid node kind for ${structName}`);
+    }
+
     if (node.arena) {
         w.write(`data := f.${api.uncapitalize(structName)}Arena.New()`);
     }

tsc/internal/ast/ast_generated.go

@@ -1,7 +1,9 @@
 package ast_test
 
 import (
+	"reflect"
 	"runtime"
+	"slices"
 	"testing"
 	"unsafe"
 
@@ -10,6 +12,48 @@ import (
 	"gotest.tools/v3/assert"
 )
 
+func TestNodeFactoryKinds(t *testing.T) {
+	t.Parallel()
+	for _, test := range []struct {
+		name  string
+		kinds []ast.Kind
+	}{
+		{"ForInOrOfStatement", []ast.Kind{ast.KindForInStatement, ast.KindForOfStatement}},
+		{"CaseOrDefaultClause", []ast.Kind{ast.KindCaseClause, ast.KindDefaultClause}},
+		{"BindingPattern", []ast.Kind{ast.KindArrayBindingPattern, ast.KindObjectBindingPattern}},
+		{"KeywordExpression", []ast.Kind{ast.KindFalseKeyword, ast.KindImportKeyword, ast.KindNullKeyword, ast.KindSuperKeyword, ast.KindThisKeyword, ast.KindTrueKeyword}},
+		{"KeywordTypeNode", []ast.Kind{ast.KindVoidKeyword, ast.KindAnyKeyword, ast.KindBooleanKeyword, ast.KindIntrinsicKeyword, ast.KindNeverKeyword, ast.KindNumberKeyword, ast.KindObjectKeyword, ast.KindStringKeyword, ast.KindSymbolKeyword, ast.KindUndefinedKeyword, ast.KindUnknownKeyword, ast.KindBigIntKeyword}},
+		{"JSDocParameterOrPropertyTag", []ast.Kind{ast.KindJSDocParameterTag, ast.KindJSDocPropertyTag}},
+	} {
+		t.Run(test.name, func(t *testing.T) {
+			t.Parallel()
+			factory := reflect.ValueOf(ast.NewNodeFactory(ast.NodeFactoryHooks{}))
+			constructor := factory.MethodByName("New" + test.name)
+			args := make([]reflect.Value, constructor.Type().NumIn())
+			for i := range args {
+				args[i] = reflect.Zero(constructor.Type().In(i))
+			}
+			for kind := ast.Kind(-1); kind <= ast.KindCount+2; kind++ {
+				args[0] = reflect.ValueOf(kind)
+				var node *ast.Node
+				panicked := func() (panicked bool) {
+					defer func() { panicked = recover() != nil }()
+					node = constructor.Call(args)[0].Interface().(*ast.Node)
+					return false
+				}()
+				valid := slices.Contains(test.kinds, kind)
+				if panicked == valid {
+					t.Fatalf("New%s(%v): valid=%v, panicked=%v", test.name, kind, valid, panicked)
+				}
+				if valid {
+					assert.Equal(t, node.Kind, kind)
+					assert.Assert(t, reflect.ValueOf(node).MethodByName("As" + test.name).Call(nil)[0].Pointer() != 0)
+				}
+			}
+		})
+	}
+}
+
 func TestNodeAccessors(t *testing.T) {
 	t.Parallel()
 	factory := ast.NewNodeFactory(ast.NodeFactoryHooks{})