Skip to content

Commit 67a1c76

Browse files
jakebaileyjakebailey
authored
Replace nodeData interface with generated AST accessors (microsoft#64773)
1 parent 6aace8b commit 67a1c76

7 files changed

Lines changed: 3653 additions & 306 deletions

File tree

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

Lines changed: 333 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@
1111
* - VisitEachChild() implementations
1212
* - Clone() implementations
1313
* - Polymorphic accessor switch statements (Expression(), Type(), etc.)
14+
* - Compact kind-indexed dispatch for shared node accessors
1415
* - Is*() type guard functions
1516
*/
1617

@@ -389,6 +390,27 @@ function isNodeFlagsMember(m: MemberInfo): boolean {
389390
return type.kind === "primitive" && type.name === "NodeFlags";
390391
}
391392

393+
function canonicalNodesByKind(): Map<NodeType, string[]> {
394+
const nodeByKind = new Map<string, NodeType>();
395+
for (const node of api.nodes()) {
396+
for (const kind of node.allKinds()) {
397+
nodeByKind.set(kind.formatGoConstant(), node);
398+
}
399+
}
400+
401+
const kindsByNode = new Map<NodeType, string[]>();
402+
for (const [kind, node] of nodeByKind) {
403+
const kinds = kindsByNode.get(node);
404+
if (kinds) {
405+
kinds.push(kind);
406+
}
407+
else {
408+
kindsByNode.set(node, [kind]);
409+
}
410+
}
411+
return kindsByNode;
412+
}
413+
392414
function emitNewFactory(
393415
w: CodeWriter,
394416
funcName: string,
@@ -425,7 +447,7 @@ function emitNewFactory(
425447
const kindArg = kindMember ? kindMember.goParamName() : `Kind${kindName}`;
426448

427449
if (nodeFlagsMembers.length > 0) {
428-
w.write(`node := f.newNode(${kindArg}, data)`);
450+
w.write(`node := f.newNode(${kindArg}, data.AsNode(), data)`);
429451
for (const m of nodeFlagsMembers) {
430452
const param = m.goParamName();
431453
if (m.bitmask) {
@@ -438,7 +460,7 @@ function emitNewFactory(
438460
w.write("return node");
439461
}
440462
else {
441-
w.write(`return f.newNode(${kindArg}, data)`);
463+
w.write(`return f.newNode(${kindArg}, data.AsNode(), data)`);
442464
}
443465

444466
w.pop();
@@ -451,7 +473,39 @@ function generateNewFactory(w: CodeWriter, node: NodeType) {
451473
const members = schemaMembers(node);
452474
const kindMember = members.find(m => m.isKindParam());
453475
const nodeFlagsMembers = members.filter(m => isNodeFlagsMember(m));
454-
emitNewFactory(w, `New${node.name}`, node.syntaxKindName, structName, node, members, kindMember, nodeFlagsMembers);
476+
if (node.name === "Token") {
477+
const tokenKinds = new Set(node.allKinds().map(kind => kind.formatGoConstant()));
478+
w.write("func (f *NodeFactory) NewToken(kind TokenSyntaxKind) *Node {");
479+
w.push();
480+
w.write("switch kind {");
481+
for (const [canonicalNode, kinds] of canonicalNodesByKind()) {
482+
if (canonicalNode === node) continue;
483+
const overlappingKinds = kinds.filter(kind => tokenKinds.has(kind));
484+
if (overlappingKinds.length === 0) continue;
485+
w.write(`case ${overlappingKinds.join(", ")}:`);
486+
w.push();
487+
if (canonicalNode.arena) {
488+
w.write(`data := f.${api.uncapitalize(canonicalNode.name)}Arena.New()`);
489+
}
490+
else {
491+
w.write(`data := &${canonicalNode.name}{}`);
492+
}
493+
w.write("return f.newNode(kind, data.AsNode(), data)");
494+
w.pop();
495+
}
496+
w.write("default:");
497+
w.push();
498+
w.write("data := f.tokenArena.New()");
499+
w.write("return f.newNode(kind, data.AsNode(), data)");
500+
w.pop();
501+
w.write("}");
502+
w.pop();
503+
w.write("}");
504+
w.write("");
505+
}
506+
else {
507+
emitNewFactory(w, `New${node.name}`, node.syntaxKindName, structName, node, members, kindMember, nodeFlagsMembers);
508+
}
455509
for (const alias of node.kindAliases) {
456510
emitNewFactory(w, `New${alias}`, alias, structName, node, members, kindMember, nodeFlagsMembers);
457511
}
@@ -604,14 +658,12 @@ function hasForEachChild(node: NodeType): boolean {
604658
// Generates a (*Node).ForEachChild method that dispatches on node.Kind to the
605659
// concrete node type's ForEachChild method.
606660
//
607-
// This deliberately avoids calling node.data.ForEachChild(v) through the
608-
// nodeData interface. An interface (or other indirect) call is opaque to escape
609-
// analysis, which must then assume the visitor `v` escapes; that forces caller
610-
// closures — and any locals they capture — onto the heap. Dispatching through a
611-
// Kind switch to a statically-resolved concrete method lets escape analysis
612-
// prove the visitor does not escape, keeping caller closures on the stack. The
613-
// integer switch over Kind also compiles to a jump table, making dispatch
614-
// cheaper than the interface call it replaces.
661+
// An indirect call is opaque to escape analysis, which must then assume the
662+
// visitor `v` escapes; that forces caller closures — and any locals they
663+
// capture — onto the heap. Dispatching through a Kind switch to a
664+
// statically-resolved concrete method lets escape analysis prove the visitor
665+
// does not escape, keeping caller closures on the stack. The integer switch
666+
// over Kind also compiles to a jump table.
615667
//
616668
// Kinds whose node has no children fall through to `default` and return false,
617669
// matching NodeDefault.ForEachChild.
@@ -638,6 +690,237 @@ function generateForEachChildDispatch(w: CodeWriter) {
638690
w.write("");
639691
}
640692

693+
function generateVisitEachChildDispatch(w: CodeWriter) {
694+
w.write("func (n *Node) VisitEachChild(v *NodeVisitor) *Node {");
695+
w.push();
696+
w.write("switch n.Kind {");
697+
for (const node of api.nodes()) {
698+
if (!hasForEachChild(node)) continue;
699+
const kinds = node.allKinds().map(kind => kind.formatGoConstant());
700+
if (kinds.length === 0) continue;
701+
w.write(`case ${kinds.join(", ")}:`);
702+
w.push();
703+
w.write(`return n.data.(*${node.name}).VisitEachChild(v)`);
704+
w.pop();
705+
}
706+
w.write("default:");
707+
w.push();
708+
w.write("return n");
709+
w.pop();
710+
w.write("}");
711+
w.pop();
712+
w.write("}");
713+
w.write("");
714+
}
715+
716+
function generateCloneDispatch(w: CodeWriter) {
717+
w.write("func (n *Node) Clone(f NodeFactoryCoercible) *Node {");
718+
w.push();
719+
w.write("switch n.Kind {");
720+
w.write("case kindFlowSwitchClauseData, kindFlowReduceLabelData:");
721+
w.push();
722+
w.write("return nil");
723+
w.pop();
724+
for (const [node, kinds] of canonicalNodesByKind()) {
725+
w.write(`case ${kinds.join(", ")}:`);
726+
w.push();
727+
w.write(`return n.data.(*${node.name}).Clone(f)`);
728+
w.pop();
729+
}
730+
w.write("default:");
731+
w.push();
732+
w.write("return nil");
733+
w.pop();
734+
w.write("}");
735+
w.pop();
736+
w.write("}");
737+
w.write("");
738+
}
739+
740+
function generateSubtreeFactsDispatch(w: CodeWriter) {
741+
w.write("func (n *Node) SubtreeFacts() SubtreeFacts {");
742+
w.push();
743+
w.write("switch n.Kind {");
744+
w.write("case kindFlowSwitchClauseData, kindFlowReduceLabelData:");
745+
w.push();
746+
w.write("return SubtreeFactsNone");
747+
w.pop();
748+
for (const [node, kinds] of canonicalNodesByKind()) {
749+
w.write(`case ${kinds.join(", ")}:`);
750+
w.push();
751+
if (transitiveBaseKeys(node).has("CompositeBase")) {
752+
w.write(`return n.data.(*${node.name}).subtreeFactsWorker(n)`);
753+
}
754+
else {
755+
w.write("return n.computeSubtreeFacts()");
756+
}
757+
w.pop();
758+
}
759+
w.write("default:");
760+
w.push();
761+
w.write("return SubtreeFactsNone");
762+
w.pop();
763+
w.write("}");
764+
w.pop();
765+
w.write("}");
766+
w.write("");
767+
}
768+
769+
function generateComputeSubtreeFactsDispatch(w: CodeWriter) {
770+
w.write("func (n *Node) computeSubtreeFacts() SubtreeFacts {");
771+
w.push();
772+
w.write("switch n.Kind {");
773+
for (const [node, kinds] of canonicalNodesByKind()) {
774+
w.write(`case ${kinds.join(", ")}:`);
775+
w.push();
776+
w.write(`return n.data.(*${node.name}).computeSubtreeFacts()`);
777+
w.pop();
778+
}
779+
w.write("default:");
780+
w.push();
781+
w.write("return SubtreeFactsNone");
782+
w.pop();
783+
w.write("}");
784+
w.pop();
785+
w.write("}");
786+
w.write("");
787+
}
788+
789+
function generatePropagateSubtreeFactsDispatch(w: CodeWriter) {
790+
w.write("func (n *Node) propagateSubtreeFacts() SubtreeFacts {");
791+
w.push();
792+
w.write("switch n.Kind {");
793+
for (const [node, kinds] of canonicalNodesByKind()) {
794+
w.write(`case ${kinds.join(", ")}:`);
795+
w.push();
796+
w.write(`return n.data.(*${node.name}).propagateSubtreeFacts()`);
797+
w.pop();
798+
}
799+
w.write("default:");
800+
w.push();
801+
w.write("return SubtreeFactsNone");
802+
w.pop();
803+
w.write("}");
804+
w.pop();
805+
w.write("}");
806+
w.write("");
807+
}
808+
809+
const NODE_ACCESSORS: { method: string; ret: string; base: string; }[] = [
810+
{ method: "FlowNodeData", ret: "*FlowNodeBase", base: "FlowNodeBase" },
811+
{ method: "DeclarationData", ret: "*DeclarationBase", base: "DeclarationBase" },
812+
{ method: "ExportableData", ret: "*ExportableBase", base: "ExportableBase" },
813+
{ method: "LocalsContainerData", ret: "*LocalsContainerBase", base: "LocalsContainerBase" },
814+
{ method: "FunctionLikeData", ret: "*FunctionLikeBase", base: "FunctionLikeBase" },
815+
{ method: "ClassLikeData", ret: "*ClassLikeBase", base: "ClassLikeBase" },
816+
{ method: "BodyData", ret: "*BodyBase", base: "BodyBase" },
817+
{ method: "LiteralLikeData", ret: "*LiteralLikeNodeBase", base: "LiteralLikeNodeBase" },
818+
{ method: "TemplateLiteralLikeData", ret: "*TemplateLiteralLikeNodeBase", base: "TemplateLiteralLikeNodeBase" },
819+
];
820+
821+
function transitiveBaseKeys(node: NodeType): Set<string> {
822+
const seen = new Set<string>();
823+
const visit = (n: NodeType) => {
824+
for (const base of n.extends) {
825+
if (!seen.has(base.key)) {
826+
seen.add(base.key);
827+
visit(base);
828+
}
829+
}
830+
};
831+
visit(node);
832+
return seen;
833+
}
834+
835+
function generateNodeAccessorDispatch(w: CodeWriter, method: string, ret: string, hasAccessor: (node: NodeType) => boolean) {
836+
const cases: { node: NodeType; kinds: string[]; }[] = [];
837+
for (const node of api.nodes()) {
838+
if (!hasAccessor(node)) continue;
839+
const kinds = node.allKinds().map(kind => kind.formatGoConstant());
840+
if (kinds.length !== 0) cases.push({ node, kinds });
841+
}
842+
843+
// Large sparse Kind switches become comparison trees. Dense indices let Go
844+
// emit jump tables; zero is reserved for kinds without the accessor.
845+
const useDispatchTable = cases.length >= 8;
846+
const table = `node${method}Dispatch`;
847+
if (useDispatchTable) {
848+
if (cases.length > 255) {
849+
throw new Error(`${method} dispatch exceeds the uint8 index range`);
850+
}
851+
w.write(`var ${table} = [kindFlowReduceLabelData + 1]uint8{`);
852+
w.push();
853+
for (const [index, { kinds }] of cases.entries()) {
854+
for (const kind of kinds) {
855+
w.write(`${kind}: ${index + 1},`);
856+
}
857+
}
858+
w.pop();
859+
w.write("}");
860+
w.write("");
861+
}
862+
863+
w.write(`func (n *Node) ${method}() ${ret} {`);
864+
w.push();
865+
if (useDispatchTable) {
866+
w.write(`switch ${table}[n.Kind] {`);
867+
}
868+
else {
869+
w.write("switch n.Kind {");
870+
}
871+
for (const [index, { node, kinds }] of cases.entries()) {
872+
w.write(`case ${useDispatchTable ? index + 1 : kinds.join(", ")}:`);
873+
w.push();
874+
w.write(`return n.data.(*${node.name}).${method}()`);
875+
w.pop();
876+
}
877+
w.write("default:");
878+
w.push();
879+
w.write("return nil");
880+
w.pop();
881+
w.write("}");
882+
w.pop();
883+
w.write("}");
884+
w.write("");
885+
}
886+
887+
function generateNodeAccessors(w: CodeWriter) {
888+
for (const { method, ret, base } of NODE_ACCESSORS) {
889+
generateNodeAccessorDispatch(w, method, ret, node => transitiveBaseKeys(node).has(base));
890+
}
891+
}
892+
893+
function hasMember(node: NodeType, name: string): boolean {
894+
return schemaMembers(node).some(member => member.name === name);
895+
}
896+
897+
function generateNameDispatch(w: CodeWriter) {
898+
generateNodeAccessorDispatch(w, "Name", "*DeclarationName", node => hasMember(node, "name"));
899+
}
900+
901+
function generateModifiersDispatch(w: CodeWriter) {
902+
generateNodeAccessorDispatch(w, "Modifiers", "*ModifierList", node => hasMember(node, "modifiers"));
903+
}
904+
905+
function generateSetModifiersDispatch(w: CodeWriter) {
906+
w.write("func (n *MutableNode) SetModifiers(modifiers *ModifierList) {");
907+
w.push();
908+
w.write("switch n.Kind {");
909+
for (const node of api.nodes()) {
910+
if (!hasMember(node, "modifiers")) continue;
911+
const kinds = node.allKinds().map(kind => kind.formatGoConstant());
912+
if (kinds.length === 0) continue;
913+
w.write(`case ${kinds.join(", ")}:`);
914+
w.push();
915+
w.write(`n.data.(*${node.name}).setModifiers(modifiers)`);
916+
w.pop();
917+
}
918+
w.write("}");
919+
w.pop();
920+
w.write("}");
921+
w.write("");
922+
}
923+
641924
// ── Generate VisitEachChild() ──────────────────────────────────────────────
642925

643926
function generateVisitEachChild(w: CodeWriter, node: NodeType) {
@@ -986,6 +1269,45 @@ function generate(): string {
9861269
w.write("");
9871270
generateForEachChildDispatch(w);
9881271

1272+
// VisitEachChild dispatch
1273+
w.write("// ──────────────────────────────────────────────────────────────────────");
1274+
w.write("// VisitEachChild dispatch");
1275+
w.write("// ──────────────────────────────────────────────────────────────────────");
1276+
w.write("");
1277+
generateVisitEachChildDispatch(w);
1278+
1279+
// Clone dispatch
1280+
w.write("// ──────────────────────────────────────────────────────────────────────");
1281+
w.write("// Clone dispatch");
1282+
w.write("// ──────────────────────────────────────────────────────────────────────");
1283+
w.write("");
1284+
generateCloneDispatch(w);
1285+
1286+
// Subtree facts dispatch
1287+
w.write("// ──────────────────────────────────────────────────────────────────────");
1288+
w.write("// Subtree facts dispatch");
1289+
w.write("// ──────────────────────────────────────────────────────────────────────");
1290+
w.write("");
1291+
generateSubtreeFactsDispatch(w);
1292+
generateComputeSubtreeFactsDispatch(w);
1293+
generatePropagateSubtreeFactsDispatch(w);
1294+
1295+
// Node accessor dispatch
1296+
w.write("// ──────────────────────────────────────────────────────────────────────");
1297+
w.write("// Node accessor dispatch");
1298+
w.write("// ──────────────────────────────────────────────────────────────────────");
1299+
w.write("");
1300+
generateNodeAccessors(w);
1301+
1302+
// Common node accessor dispatch
1303+
w.write("// ──────────────────────────────────────────────────────────────────────");
1304+
w.write("// Common node accessor dispatch");
1305+
w.write("// ──────────────────────────────────────────────────────────────────────");
1306+
w.write("");
1307+
generateNameDispatch(w);
1308+
generateModifiersDispatch(w);
1309+
generateSetModifiersDispatch(w);
1310+
9891311
// As*() casts
9901312
w.write("// ──────────────────────────────────────────────────────────────────────");
9911313
w.write("// As*() cast methods");

0 commit comments

Comments
 (0)