bachelor_bootstrap_compiler.py

from __future__ import annotations

import argparse
import re
from dataclasses import dataclass
from typing import List, Optional


class CompileError(Exception):
    pass


@dataclass
class Token:
    kind: str
    value: str
    pos: int


@dataclass
class Program:
    statements: List[object]


@dataclass
class Block:
    statements: List[object]


@dataclass
class LetStmt:
    name: str
    expr: object


@dataclass
class AssignStmt:
    name: str
    expr: object


@dataclass
class PrintStmt:
    expr: object


@dataclass
class IfStmt:
    condition: object
    then_block: Block
    else_block: Optional[Block]


@dataclass
class WhileStmt:
    condition: object
    body: Block


@dataclass
class FnStmt:
    name: str
    params: List[str]
    body: Block


@dataclass
class ReturnStmt:
    expr: object


@dataclass
class ExprStmt:
    expr: object


@dataclass
class BinaryExpr:
    left: object
    op: str
    right: object


@dataclass
class UnaryExpr:
    op: str
    expr: object


@dataclass
class NumberExpr:
    value: str


@dataclass
class StringExpr:
    value: str


@dataclass
class VarExpr:
    name: str


@dataclass
class CallExpr:
    callee: object
    args: List[object]


KEYWORDS = {
    "let": "LET",
    "fn": "FN",
    "return": "RETURN",
    "if": "IF",
    "else": "ELSE",
    "while": "WHILE",
    "print": "PRINT",
}

TOKEN_SPEC = [
    ("NUMBER", r"\d+(?:\.\d+)?"),
    ("STRING", r'"([^"\\]|\\.)*"'),
    ("ID", r"[A-Za-z_][A-Za-z0-9_]*"),
    ("EQ", r"=="),
    ("NE", r"!="),
    ("LE", r"<="),
    ("GE", r">="),
    ("ASSIGN", r"="),
    ("LT", r"<"),
    ("GT", r">"),
    ("PLUS", r"\+"),
    ("MINUS", r"-"),
    ("STAR", r"\*"),
    ("SLASH", r"/"),
    ("LPAREN", r"\("),
    ("RPAREN", r"\)"),
    ("LBRACE", r"\{"),
    ("RBRACE", r"\}"),
    ("COMMA", r","),
    ("SEMI", r";"),
    ("SKIP", r"[ \t\r\n]+"),
    ("MISMATCH", r"."),
]
TOKEN_RE = re.compile("|".join(f"(?P<{name}>{pattern})" for name, pattern in TOKEN_SPEC))


class Lexer:
    def __init__(self, source: str) -> None:
        self.source = source

    def tokenize(self) -> List[Token]:
        tokens: List[Token] = []
        for match in TOKEN_RE.finditer(self.source):
            kind = match.lastgroup
            value = match.group()
            pos = match.start()

            if kind == "SKIP":
                continue
            if kind == "MISMATCH":
                raise CompileError(f"Unexpected character '{value}' at position {pos}")
            if kind == "ID" and value in KEYWORDS:
                kind = KEYWORDS[value]

            tokens.append(Token(kind, value, pos))

        tokens.append(Token("EOF", "", len(self.source)))
        return tokens


class Parser:
    def __init__(self, tokens: List[Token]) -> None:
        self.tokens = tokens
        self.i = 0

    def current(self) -> Token:
        return self.tokens[self.i]

    def advance(self) -> Token:
        tok = self.current()
        self.i += 1
        return tok

    def match(self, *kinds: str) -> Optional[Token]:
        if self.current().kind in kinds:
            return self.advance()
        return None

    def expect(self, kind: str, message: str) -> Token:
        tok = self.current()
        if tok.kind != kind:
            raise CompileError(f"{message} at position {tok.pos}")
        return self.advance()

    def parse_program(self) -> Program:
        statements: List[object] = []
        while self.current().kind != "EOF":
            statements.append(self.parse_statement())
        return Program(statements)

    def parse_block(self) -> Block:
        self.expect("LBRACE", "Expected '{'")
        statements: List[object] = []
        while self.current().kind != "RBRACE":
            if self.current().kind == "EOF":
                raise CompileError("Unclosed block")
            statements.append(self.parse_statement())
        self.expect("RBRACE", "Expected '}'")
        return Block(statements)

    def parse_statement(self) -> object:
        if self.match("LET"):
            name = self.expect("ID", "Expected variable name after let").value
            self.expect("ASSIGN", "Expected '=' after variable name")
            expr = self.parse_expression()
            self.expect("SEMI", "Expected ';' after let statement")
            return LetStmt(name, expr)

        if self.match("FN"):
            name = self.expect("ID", "Expected function name").value
            self.expect("LPAREN", "Expected '(' after function name")
            params: List[str] = []
            if self.current().kind != "RPAREN":
                while True:
                    params.append(self.expect("ID", "Expected parameter name").value)
                    if not self.match("COMMA"):
                        break
            self.expect("RPAREN", "Expected ')' after parameters")
            body = self.parse_block()
            return FnStmt(name, params, body)

        if self.match("RETURN"):
            expr = self.parse_expression()
            self.expect("SEMI", "Expected ';' after return")
            return ReturnStmt(expr)

        if self.match("IF"):
            condition = self.parse_expression()
            then_block = self.parse_block()
            else_block = self.parse_block() if self.match("ELSE") else None
            return IfStmt(condition, then_block, else_block)

        if self.match("WHILE"):
            condition = self.parse_expression()
            body = self.parse_block()
            return WhileStmt(condition, body)

        if self.match("PRINT"):
            expr = self.parse_expression()
            self.expect("SEMI", "Expected ';' after print")
            return PrintStmt(expr)

        if self.current().kind == "ID" and self.tokens[self.i + 1].kind == "ASSIGN":
            name = self.advance().value
            self.advance()
            expr = self.parse_expression()
            self.expect("SEMI", "Expected ';' after assignment")
            return AssignStmt(name, expr)

        expr = self.parse_expression()
        self.expect("SEMI", "Expected ';' after expression")
        return ExprStmt(expr)

    def parse_expression(self) -> object:
        return self.parse_equality()

    def parse_equality(self) -> object:
        expr = self.parse_comparison()
        while self.current().kind in ("EQ", "NE"):
            op = self.advance().value
            right = self.parse_comparison()
            expr = BinaryExpr(expr, op, right)
        return expr

    def parse_comparison(self) -> object:
        expr = self.parse_term()
        while self.current().kind in ("LT", "LE", "GT", "GE"):
            op = self.advance().value
            right = self.parse_term()
            expr = BinaryExpr(expr, op, right)
        return expr

    def parse_term(self) -> object:
        expr = self.parse_factor()
        while self.current().kind in ("PLUS", "MINUS"):
            op = self.advance().value
            right = self.parse_factor()
            expr = BinaryExpr(expr, op, right)
        return expr

    def parse_factor(self) -> object:
        expr = self.parse_unary()
        while self.current().kind in ("STAR", "SLASH"):
            op = self.advance().value
            right = self.parse_unary()
            expr = BinaryExpr(expr, op, right)
        return expr

    def parse_unary(self) -> object:
        if self.current().kind in ("MINUS",):
            op = self.advance().value
            return UnaryExpr(op, self.parse_unary())
        return self.parse_call()

    def parse_call(self) -> object:
        expr = self.parse_primary()
        while self.match("LPAREN"):
            args: List[object] = []
            if self.current().kind != "RPAREN":
                while True:
                    args.append(self.parse_expression())
                    if not self.match("COMMA"):
                        break
            self.expect("RPAREN", "Expected ')' after function call args")
            expr = CallExpr(expr, args)
        return expr

    def parse_primary(self) -> object:
        tok = self.current()
        if self.match("NUMBER"):
            return NumberExpr(tok.value)
        if self.match("STRING"):
            return StringExpr(tok.value)
        if self.match("ID"):
            return VarExpr(tok.value)
        if self.match("LPAREN"):
            expr = self.parse_expression()
            self.expect("RPAREN", "Expected ')' after expression")
            return expr

        raise CompileError(f"Expected expression at position {tok.pos}")


class PythonEmitter:
    def __init__(self) -> None:
        self.lines: List[str] = []
        self.indent_level = 0

    def emit(self, line: str = "") -> None:
        self.lines.append("    " * self.indent_level + line)

    def with_indent(self):
        class _Indent:
            def __init__(self, emitter: PythonEmitter):
                self.emitter = emitter

            def __enter__(self):
                self.emitter.indent_level += 1

            def __exit__(self, exc_type, exc, tb):
                self.emitter.indent_level -= 1

        return _Indent(self)

    def emit_program(self, program: Program) -> str:
        self.emit("# Generated by Bachelor Bootstrap Compiler")
        self.emit("from __future__ import annotations")
        self.emit("")
        for stmt in program.statements:
            self.emit_stmt(stmt)
        return "\n".join(self.lines) + "\n"

    def emit_stmt(self, stmt: object) -> None:
        if isinstance(stmt, LetStmt):
            self.emit(f"{stmt.name} = {self.emit_expr(stmt.expr)}")
            return
        if isinstance(stmt, AssignStmt):
            self.emit(f"{stmt.name} = {self.emit_expr(stmt.expr)}")
            return
        if isinstance(stmt, PrintStmt):
            self.emit(f"print({self.emit_expr(stmt.expr)})")
            return
        if isinstance(stmt, ExprStmt):
            self.emit(self.emit_expr(stmt.expr))
            return
        if isinstance(stmt, ReturnStmt):
            self.emit(f"return {self.emit_expr(stmt.expr)}")
            return
        if isinstance(stmt, IfStmt):
            self.emit(f"if {self.emit_expr(stmt.condition)}:")
            if stmt.then_block.statements:
                with self.with_indent():
                    for s in stmt.then_block.statements:
                        self.emit_stmt(s)
            else:
                with self.with_indent():
                    self.emit("pass")
            if stmt.else_block is not None:
                self.emit("else:")
                if stmt.else_block.statements:
                    with self.with_indent():
                        for s in stmt.else_block.statements:
                            self.emit_stmt(s)
                else:
                    with self.with_indent():
                        self.emit("pass")
            return
        if isinstance(stmt, WhileStmt):
            self.emit(f"while {self.emit_expr(stmt.condition)}:")
            if stmt.body.statements:
                with self.with_indent():
                    for s in stmt.body.statements:
                        self.emit_stmt(s)
            else:
                with self.with_indent():
                    self.emit("pass")
            return
        if isinstance(stmt, FnStmt):
            params = ", ".join(stmt.params)
            self.emit(f"def {stmt.name}({params}):")
            if stmt.body.statements:
                with self.with_indent():
                    for s in stmt.body.statements:
                        self.emit_stmt(s)
            else:
                with self.with_indent():
                    self.emit("pass")
            return

        raise CompileError(f"Unknown statement node: {type(stmt).__name__}")

    def emit_expr(self, expr: object) -> str:
        if isinstance(expr, NumberExpr):
            return expr.value
        if isinstance(expr, StringExpr):
            return expr.value
        if isinstance(expr, VarExpr):
            return expr.name
        if isinstance(expr, UnaryExpr):
            return f"({expr.op}{self.emit_expr(expr.expr)})"
        if isinstance(expr, BinaryExpr):
            return f"({self.emit_expr(expr.left)} {expr.op} {self.emit_expr(expr.right)})"
        if isinstance(expr, CallExpr):
            args = ", ".join(self.emit_expr(arg) for arg in expr.args)
            return f"{self.emit_expr(expr.callee)}({args})"

        raise CompileError(f"Unknown expression node: {type(expr).__name__}")


def compile_source(source: str) -> str:
    tokens = Lexer(source).tokenize()
    ast = Parser(tokens).parse_program()
    return PythonEmitter().emit_program(ast)


DEMO_SOURCE = """
fn fib(n) {


    if n <= 1 {
        return n;
    } else {
        return fib(n - 1) + fib(n - 2);
    }
}

let i = 0;
while i < 7 {
    print fib(i);
    i = i + 1;
}
""".strip()


def main() -> None:
    parser = argparse.ArgumentParser(description="Bachelor Bootstrap Compiler (Stage-0 in Python)")
    parser.add_argument("source", nargs="?", help="Path to source .btl file")
    parser.add_argument("--out", default="compiled_output.py", help="Output Python file")
    parser.add_argument("--run", action="store_true", help="Run the compiled Python code")
    parser.add_argument("--demo", action="store_true", help="Compile built-in demo program")
    args = parser.parse_args()

    if args.demo:
        source_code = DEMO_SOURCE
    elif args.source:
        with open(args.source, "r", encoding="utf-8") as f:
            source_code = f.read()
    else:
        parser.error("Provide a source file or use --demo")
        return

    python_code = compile_source(source_code)

    with open(args.out, "w", encoding="utf-8") as f:
        f.write(python_code)

    print(f"Compiled successfully -> {args.out}")

    if args.run:
        print("Running compiled program:")
        exec_globals = {"__name__": "__main__"}
        exec(python_code, exec_globals, exec_globals)


if __name__ == "__main__":
    main()