Files
endorphin-aster-vm/scripts/generate_ast.py
codex-bot f5268d2e09
Verify databases and deploy Aster VM / verify-and-deploy (push) Successful in 7m55s
feat: generate unified AST and editor contracts from ASDL
2026-10-19 14:31:28 +08:00

406 lines
17 KiB
Python

#!/usr/bin/env python3
"""Generate checked-in AST types, walkers and editor contracts from Endor.asdl.
Uses Python's standard library and rustfmt only. It never builds or runs Endor.
"""
from __future__ import annotations
import argparse
from dataclasses import dataclass
import hashlib
import json
from pathlib import Path
import re
import subprocess
BUILTINS = {"string": "String", "span": "Span", "schema_id": "String"}
HEADER = "// Generated by scripts/generate_ast.py from Endor.asdl. Do not edit.\n"
RESERVED = set("as break const continue crate else enum extern false fn for if impl in let loop match mod move mut pub ref return self Self static struct super trait true type unsafe use where while async await dyn abstract become box do final macro override priv typeof unsized virtual yield try gen".split())
@dataclass(frozen=True)
class Field:
type: str
name: str
quantifier: str = ""
@dataclass(frozen=True)
class Variant:
name: str
fields: tuple[Field, ...]
@dataclass(frozen=True)
class Definition:
name: str
fields: tuple[Field, ...] | None
variants: tuple[Variant, ...] = ()
def all_fields(self):
return self.fields if self.fields is not None else tuple(f for v in self.variants for f in v.fields)
def simple(self):
return self.fields is None and all(not v.fields for v in self.variants)
def parse_schema(text: str) -> dict[str, Definition]:
"""Parse the project's ASDL subset; reject unknown syntax and field types."""
text = re.sub(r"--[^\n]*", "", text)
tokens = []
cursor = 0
pattern = re.compile(r"\s*([A-Za-z_][A-Za-z_0-9]*|[{}()=|,*?])")
while cursor < len(text):
if not text[cursor:].strip():
break
match = pattern.match(text, cursor)
if not match:
raise ValueError(f"unsupported ASDL syntax near {text[cursor:cursor + 30]!r}")
tokens.append(match[1])
cursor = match.end()
index = 0
def peek():
if index == len(tokens):
raise ValueError("unexpected end of ASDL")
return tokens[index]
def take(expected=None):
nonlocal index
if index == len(tokens):
raise ValueError("unexpected end of ASDL")
value = tokens[index]
if expected is not None and value != expected:
raise ValueError(f"expected {expected!r}, found {value!r}")
index += 1
return value
def identifier():
value = take()
if not re.fullmatch(r"[A-Za-z_][A-Za-z_0-9]*", value) or value in RESERVED:
raise ValueError(f"invalid or reserved identifier: {value}")
return value
def fields():
take("(")
result = []
while peek() != ")":
type_name = identifier()
quantifier = take() if peek() in ("?", "*") else ""
result.append(Field(type_name, identifier(), quantifier))
if peek() != ",":
break
take(",")
take(")")
if len({f.name for f in result}) != len(result):
raise ValueError("duplicate field name")
return tuple(result)
take("module")
identifier()
take("{")
definitions = {}
while peek() != "}":
name = identifier()
if name in definitions or name in BUILTINS:
raise ValueError(f"duplicate or builtin type: {name}")
take("=")
if peek() == "(":
definition = Definition(name, fields())
else:
variants = []
while True:
variant = identifier()
members = fields() if peek() == "(" else ()
variants.append(Variant(variant, members))
if peek() != "|":
break
take("|")
if len({v.name for v in variants}) != len(variants):
raise ValueError(f"duplicate constructor in {name}")
definition = Definition(name, None, tuple(variants))
definitions[name] = definition
take("}")
if index != len(tokens):
raise ValueError("trailing ASDL input")
for definition in definitions.values():
for field in definition.all_fields():
if field.type not in definitions and field.type not in BUILTINS:
raise ValueError(f"undefined field type: {field.type}")
return definitions
def snake(name):
return re.sub(r"(?<!^)(?=[A-Z])", "_", name).lower()
def recursive_edges(definitions):
"""Box DFS back edges in the by-value graph; Vec already breaks recursion."""
active, done, boxed = set(), set(), set()
def visit(name):
active.add(name)
for field in definitions[name].all_fields():
if field.quantifier == "*" or field.type not in definitions:
continue
if field.type in active:
boxed.add((name, field.type))
elif field.type not in done:
visit(field.type)
active.remove(name)
done.add(name)
for name in definitions:
if name not in done:
visit(name)
return boxed
def spanned_types(definitions):
result = {d.name for d in definitions.values() if d.fields is not None and Field("span", "span") in d.fields}
while True:
expanded = result | {
d.name for d in definitions.values()
if d.variants and all(len(v.fields) == 1 and not v.fields[0].quantifier and v.fields[0].type in result for v in d.variants)
}
if expanded == result:
return result
result = expanded
def rust_types(definitions, schema_id):
boxed = recursive_edges(definitions)
spanned = spanned_types(definitions)
def field_type(owner, field):
value = BUILTINS.get(field.type, field.type)
if field.quantifier == "*":
return f"Vec<{value}>"
if (owner, field.type) in boxed:
value = f"Box<{value}>"
return f"Option<{value}>" if field.quantifier == "?" else value
lines = [HEADER, "use crate::Span;", "use super::Spanned;", f'pub const AST_SCHEMA_ID: &str = "{schema_id}";']
for definition in definitions.values():
derives = "Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize"
if definition.simple():
derives += ", Copy"
lines.append(f"#[derive({derives})]")
if not definition.simple():
tag = 'tag = "kind", content = "data", ' if definition.fields is None else ""
lines.append(f"#[serde({tag}deny_unknown_fields)]")
if definition.fields is not None:
lines.append(f"pub struct {definition.name} {{")
lines.extend(f"pub {f.name}: {field_type(definition.name, f)}," for f in definition.fields)
else:
lines.append(f"pub enum {definition.name} {{")
for variant in definition.variants:
if not variant.fields:
lines.append(f"{variant.name},")
elif len(variant.fields) == 1:
lines.append(f"{variant.name}({field_type(definition.name, variant.fields[0])}),")
else:
fields = ", ".join(f"{f.name}: {field_type(definition.name, f)}" for f in variant.fields)
lines.append(f"{variant.name} {{ {fields} }},")
lines.append("}")
if definition.name in spanned:
lines.append(f"impl Spanned for {definition.name} {{ fn span(&self) -> Span {{")
if definition.fields is not None:
lines.append("self.span")
else:
lines.append("match self {")
lines.extend(f"Self::{v.name}(node) => node.span()," for v in definition.variants)
lines.append("}")
lines.append("} }")
return "\n".join(lines) + "\n"
def rust_visitor(definitions):
spanned = spanned_types(definitions)
names = ", ".join(definitions)
lines = [HEADER, f"use super::{{{names}, Spanned}};", "use crate::Span;", "pub trait Visitor {", "type Error;",
"fn enter_node(&mut self, _span: Span) -> Result<(), Self::Error> { Ok(()) }",
"fn leave_node(&mut self) {}",
"fn visit_span(&mut self, _span: &Span) -> Result<(), Self::Error> { Ok(()) }"]
for name in definitions:
lines.append(f"fn visit_{snake(name)}(&mut self, node: &{name}) -> Result<(), Self::Error> {{")
if name in spanned:
lines += ["self.enter_node(node.span())?;", f"let result = walk_{snake(name)}(self, node);", "self.leave_node();", "result"]
else:
lines.append(f"walk_{snake(name)}(self, node)")
lines.append("}")
lines.append("}")
def children(field, access):
if field.type not in definitions and field.type != "span":
return []
method = f"visitor.visit_{snake(field.type)}"
if field.quantifier == "*":
return [f"for value in {access} {{ {method}(value)?; }}"]
if field.quantifier == "?":
return [f"if let Some(value) = {access} {{ {method}(value)?; }}"]
return [f"{method}({access})?;"]
for definition in definitions.values():
body = []
if definition.fields is not None:
for field in definition.fields:
body += children(field, f"&node.{field.name}")
elif not definition.simple():
body.append("match node {")
for variant in definition.variants:
visited = [f for f in variant.fields if f.type in definitions or f.type == "span"]
if not variant.fields:
pattern = f"{definition.name}::{variant.name}"
elif len(variant.fields) == 1:
pattern = f"{definition.name}::{variant.name}({'value' if visited else '_'})"
else:
members = ", ".join(f.name for f in visited)
pattern = f"{definition.name}::{variant.name} {{ " + (members + ", " if members else "") + ".. }"
body.append(pattern + " => {")
for field in visited:
body += children(field, "value" if len(variant.fields) == 1 else field.name)
body.append("}")
body.append("}")
visitor = "visitor" if any("visitor." in line for line in body) else "_visitor"
node = "node" if body else "_node"
lines.append(f"pub fn walk_{snake(definition.name)}<V: Visitor + ?Sized>({visitor}: &mut V, {node}: &{definition.name}) -> Result<(), V::Error> {{")
lines.extend(body)
lines += ["Ok(())", "}"]
return "\n".join(lines) + "\n"
def editor_types(definitions, schema_id):
def field_type(field):
result = {"string": "string", "span": "Span", "schema_id": "typeof AST_SCHEMA_ID"}.get(field.type, field.type)
if field.quantifier == "*":
result = f"Array<{result}>"
elif field.quantifier == "?":
result += " | null"
return result
def object_type(fields):
return "{ " + " ".join(f"{f.name}{'?' if f.quantifier == '?' else ''}: {field_type(f)};" for f in fields) + " }"
lines = [HEADER, f'export const AST_SCHEMA_ID = "{schema_id}" as const;', "export interface Span { start: number; end: number; }"]
for definition in definitions.values():
if definition.fields is not None:
lines.append(f"export type {definition.name} = {object_type(definition.fields)};")
else:
variants = []
for variant in definition.variants:
if definition.simple():
value = json.dumps(variant.name)
else:
data = ""
if variant.fields:
result = field_type(variant.fields[0]) if len(variant.fields) == 1 else object_type(variant.fields)
data = f" data: {result};"
value = '{ kind: ' + json.dumps(variant.name) + ';' + data + ' }'
variants.append(value)
lines.append(f"export type {definition.name} = " + " | ".join(variants) + ";")
lines += ["export interface Visitor {", "visitSpan?(node: Span): boolean | void;"]
lines.extend(f"visit{name}?(node: {name}): boolean | void;" for name in definitions)
lines.append("}")
def children(field, access):
if field.type not in definitions and field.type != "span":
return []
function = "walkSpan" if field.type == "span" else f"walk{field.type}"
if field.quantifier == "*":
return [f"for (const value of {access}) {function}(value, visitor);"]
if field.quantifier == "?":
return [f"if ({access} != null) {function}({access}, visitor);"]
return [f"{function}({access}, visitor);"]
lines.append("export function walkSpan(node: Span, visitor: Visitor): void { visitor.visitSpan?.(node); }")
for definition in definitions.values():
lines += [f"export function walk{definition.name}(node: {definition.name}, visitor: Visitor): void {{", f"if (visitor.visit{definition.name}?.(node) === false) return;"]
if definition.fields is not None:
for field in definition.fields:
lines += children(field, "node." + field.name)
elif not definition.simple():
lines.append("switch (node.kind) {")
for variant in definition.variants:
lines.append(f"case {json.dumps(variant.name)}:")
for field in variant.fields:
access = "node.data" + ("." + field.name if len(variant.fields) > 1 else "")
lines += children(field, access)
lines.append("break;")
lines.append("}")
lines.append("}")
return "\n".join(lines) + "\n"
def json_schema(definitions, schema_id):
def field_schema(field):
value = {"string": {"type": "string"}, "schema_id": {"const": schema_id}, "span": {"$ref": "#/$defs/Span"}}.get(field.type, {"$ref": f"#/$defs/{field.type}"})
if field.quantifier == "*":
value = {"type": "array", "items": value}
elif field.quantifier == "?":
value = {"anyOf": [value, {"type": "null"}]}
return value
def object_schema(properties, required=None):
return {"type": "object", "properties": properties, "required": list(properties) if required is None else required, "additionalProperties": False}
definitions_json = {"Span": object_schema({name: {"type": "integer", "minimum": 0, "maximum": 9007199254740991} for name in ["start", "end"]})}
for definition in definitions.values():
if definition.fields is not None:
value = object_schema({f.name: field_schema(f) for f in definition.fields}, [f.name for f in definition.fields if f.quantifier != "?"])
elif definition.simple():
value = {"enum": [v.name for v in definition.variants]}
else:
variants = []
for variant in definition.variants:
properties = {"kind": {"const": variant.name}}
if variant.fields:
properties["data"] = field_schema(variant.fields[0]) if len(variant.fields) == 1 else object_schema({f.name: field_schema(f) for f in variant.fields}, [f.name for f in variant.fields if f.quantifier != "?"])
variants.append(object_schema(properties))
value = {"oneOf": variants}
definitions_json[definition.name] = value
return json.dumps({"$schema": "https://json-schema.org/draft/2020-12/schema", "$id": schema_id, "$ref": "#/$defs/AstDocument", "$defs": definitions_json}, ensure_ascii=False, indent=2) + "\n"
def generate(text):
definitions = parse_schema(text)
for required in ("AstDocument", "AstRoot"):
if required not in definitions:
raise ValueError(f"missing required schema type: {required}")
schema_id = "urn:endor:ast:sha256:" + hashlib.sha256(text.encode("utf-8")).hexdigest()
return {
"crates/frontend/src/ast/generated.rs": rust_types(definitions, schema_id),
"crates/frontend/src/ast/visitor.rs": rust_visitor(definitions),
"crates/frontend/editor/ast.ts": editor_types(definitions, schema_id),
"crates/frontend/editor/ast.schema.json": json_schema(definitions, schema_id),
}
def main():
parser = argparse.ArgumentParser(description=__doc__)
mode = parser.add_mutually_exclusive_group(required=True)
mode.add_argument("--check", action="store_true")
mode.add_argument("--write", action="store_true")
args = parser.parse_args()
root = Path(__file__).resolve().parents[1]
outputs = generate((root / "crates/frontend/Endor.asdl").read_text(encoding="utf-8"))
stale = []
for relative, content in outputs.items():
if relative.endswith(".rs"):
content = subprocess.run(["rustfmt", "--edition", "2024", "--emit", "stdout"], input=content, text=True, capture_output=True, check=True).stdout
path = root / relative
if args.write:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(content, encoding="utf-8")
elif not path.is_file() or path.read_text(encoding="utf-8") != content:
stale.append(relative)
if stale:
parser.exit(1, "stale generated AST files:\n" + "\n".join(stale) + "\nRun python3 scripts/generate_ast.py --write\n")
print("AST artifacts generated." if args.write else "AST artifacts match Endor.asdl.")
if __name__ == "__main__":
main()