not-python-rust/tools/genast.py

#!/usr/bin/env python3
import abc
import sys
import time
from pathlib import Path
from typing import Self


class SB:
    "String builder"

    def __init__(self, init: str = ""):
        self.value = init
        self.padding = 0

    def indent(self) -> None:
        self.padding += 1

    def dedent(self) -> None:
        self.padding -= 1
        assert self.padding >= 0

    def line(self, line: str = ""):
        self.write(f"{line}\n")

    def write(self, value: str):
        if not self.value or self.value[-1] == "\n":
            self.value += "    " * self.padding
        self.value += value

    def __str__(self) -> str:
        return self.value


class Generator(metaclass=abc.ABCMeta):
    @abc.abstractmethod
    def generate_members(self) -> str:
        pass

    @abc.abstractmethod
    def generate_visitor(self) -> str:
        pass


class GenerateStruct(Generator):
    def __init__(self, base: str, name: str, rules: list[str]) -> None:
        self.base = base
        self.name = name
        self.rules = rules

    def generate_members(self) -> str:
        sb = SB()
        sb.line("#[derive(Debug)]")
        sb.line(f"pub struct {self.name}{self.base} " + "{")
        sb.indent()
        for member in self.rules:
            sb.write(f"pub {member}")
            sb.line(",")
        sb.dedent()
        sb.line("}")
        sb.line()

        sb.line(f"impl {self.base} for {self.name}{self.base} " + "{")
        sb.indent()
        sb.line(
            f"fn accept(&self, visitor: &mut dyn {self.base}Visitor) -> Result<(), Error> "
            + "{"
        )
        sb.indent()
        sb.line(f"visitor.visit_{self.name.lower()}_{self.base.lower()}(self)")

        sb.dedent()
        sb.line("}")
        sb.line()
        sb.line("fn as_any(self: Box<Self>) -> Box<dyn Any> { self }")
        sb.line("fn as_any_ref(&self) -> &dyn Any { self }")

        sb.dedent()
        sb.write("}")

        return str(sb)

    def generate_visitor(self) -> str:
        return f"fn visit_{self.name.lower()}_{self.base.lower()}(&mut self, {self.base.lower()}: &{self.name}{self.base}) -> Result<(), Error>;"


class GenerateEnum(Generator):
    def __init__(
        self, base: str, name: str, variants: list[tuple[str, str | list[str]]]
    ):
        self.base = base
        self.name = name
        self.variants = [
            (variant, ([members] if isinstance(members, str) else members))
            for (variant, members) in variants
        ]

    def generate_members(self) -> str:
        sb = SB()
        sb.line("#[derive(Debug)]")
        sb.line(f"pub enum {self.name}{self.base} " + "{")
        sb.indent()
        for variant, members in self.variants:
            sb.write(f"{variant}")
            if members:
                struct_like = any(":" in member for member in members)

                if struct_like:
                    sb.line(" {")
                    sb.indent()
                    for member in members:
                        sb.line(f"{member},")
                    sb.dedent()
                    sb.line("},")
                else:
                    sb.write("(")
                    sb.write(", ".join(members))
                    sb.line("),")
            else:
                sb.line(",")
        sb.dedent()
        sb.line("}")
        sb.line()

        sb.line(f"impl {self.base} for {self.name}{self.base} " + "{")
        sb.indent()
        sb.line(
            f"fn accept(&self, visitor: &mut dyn {self.base}Visitor) -> Result<(), Error> "
            + "{"
        )
        sb.indent()
        sb.line(f"visitor.visit_{self.name.lower()}_{self.base.lower()}(self)")

        sb.dedent()
        sb.line("}")
        sb.line()
        sb.line("fn as_any(self: Box<Self>) -> Box<dyn Any> { self }")
        sb.line("fn as_any_ref(&self) -> &dyn Any { self }")

        sb.dedent()
        sb.write("}")

        return str(sb)

    def generate_visitor(self) -> str:
        return f"fn visit_{self.name.lower()}_{self.base.lower()}(&mut self, {self.base.lower()}: &{self.name}{self.base}) -> Result<(), Error>;"


class GenerateGroup(Generator):
    def __init__(self, name: str, group: list[Generator] | None = None) -> None:
        self.name = name
        if not group:
            group = []
        self.group = group

    def generate(self) -> str:
        sb = SB()

        sb.line(self.generate_visitor())
        sb.line()
        sb.line(self.generate_members())

        return str(sb)

    def add_enum(self, name: str, variants: list[tuple[str, str | list[str]]]) -> Self:
        self.group += [GenerateEnum(self.name, name, variants)]
        return self

    def add_struct(self, name: str, rules: list[str]) -> Self:
        self.group += [GenerateStruct(self.name, name, rules)]
        return self

    def generate_members(self) -> str:
        sb = SB()

        # Trait and acceptor
        sb.line(f"pub trait {self.name}: Debug + Any " + "{")
        sb.indent()
        sb.line(
            f"fn accept(&self, visitor: &mut dyn {self.name}Visitor) -> Result<(), Error>;"
        )
        sb.line("fn as_any(self: Box<Self>) -> Box<dyn Any>;")
        sb.line("fn as_any_ref(&self) -> &dyn Any;")
        sb.dedent()
        sb.line("}")
        sb.line()

        # Type
        sb.line(f"pub type {self.name}P = Box<dyn {self.name} + 'static>;")
        sb.line()

        # All members and rules
        for g in self.group:
            sb.line(g.generate_members())
            sb.line()

        return str(sb).strip()

    def generate_visitor(self) -> str:
        sb = SB()

        # Visitor
        sb.line(f"pub trait {self.name}Visitor " + "{")
        sb.indent()
        for g in self.group:
            sb.line(g.generate_visitor())
        sb.dedent()
        sb.line("}")

        return str(sb).strip()


GENERATE = [
    # Expr
    GenerateGroup("Expr")
    .add_struct("Binary", ["lhs: ExprP", "op: Token", "rhs: ExprP"])
    .add_struct("Unary", ["op: Token", "expr: ExprP"])
    .add_struct("Call", ["expr: ExprP", "args: Vec<ExprP>", "rparen: Token"])
    .add_struct("Get", ["expr: ExprP", "name: Token"])
    .add_struct("Index", ["expr: ExprP", "index: ExprP", "rbracket: Token"])
    .add_struct("Primary", ["token: Token"])
    .add_struct(
        "Function",
        [
            "lparen: Token",
            "params: Vec<(Token, Option<ExprP>)>",
            "return_type: Option<ExprP>",
            "body: Vec<StmtP>",
            "rbrace: Token",
        ],
    )
    .add_struct("List", ["lbracket: Token", "exprs: Vec<ExprP>", "rbracket: Token"]),
    # Stmt
    GenerateGroup("Stmt")
    .add_struct("Import", ["import_kw: Token", "what: Vec<Token>", "module: Token"])
    .add_struct("Expr", ["expr: ExprP"])
    .add_struct("Assign", ["lhs: Token", "rhs: ExprP"])
    .add_struct("Set", ["expr: ExprP", "name: Token", "rhs: ExprP"])
    .add_struct("Block", ["lbrace: Token", "stmts: Vec<StmtP>", "rbrace: Token"])
    .add_struct("Return", ["return_kw: Token", "expr: Option<ExprP>"])
    .add_struct(
        "If",
        [
            "if_kw: Token",
            "condition: ExprP",
            "then_branch: BlockStmt",
            "else_branch: Vec<StmtP>",
        ],
    ),
]


def main():
    if len(sys.argv) == 1:
        path = Path(__file__).parent / "../src/ast.rs"
    else:
        path = Path(sys.argv[1])

    path = path.resolve()
    print(f"Using {path}")

    sb = SB()

    # File headers
    sb.line(
        "// This is an auto-generated file. Any changes made to this file may be overwritten."
    )
    sb.line("// This file was created at: " + time.strftime("%Y-%m-%d %H:%M:%S"))
    sb.line("#![allow(dead_code)]")

    # Imports
    sb.line("use std::fmt::Debug;")
    sb.line("use std::any::Any;")
    sb.line()
    sb.line("use crate::token::Token;")
    sb.line()
    sb.line("type Error = Box<dyn std::error::Error>;")
    sb.line()

    for g in GENERATE:
        sb.line(g.generate())

    with open(path, "w") as fp:
        fp.write(str(sb).strip())


if __name__ == "__main__":
    main()