diff options
Diffstat (limited to 'src/ast_compiler.py')
| -rw-r--r-- | src/ast_compiler.py | 350 |
1 files changed, 350 insertions, 0 deletions
diff --git a/src/ast_compiler.py b/src/ast_compiler.py new file mode 100644 index 0000000..2208e43 --- /dev/null +++ b/src/ast_compiler.py @@ -0,0 +1,350 @@ +#!/usr/bin/env python3 +# Generate .atk16 assembly from a subset of Python + +from _ast import AnnAssign, Expr, Module +import sys +import ast + +from dataclasses import dataclass +from typing import Literal, Set, cast, Any, TypeAlias +from collections import OrderedDict + +if len(sys.argv) != 3: + print("usage: ast_compiler.py <infile.py> <outfile.atk16>") + sys.exit(1) + +infile_path = sys.argv[1] +outfile_path = sys.argv[2] + +Addr = int +Label = str +RegChar = Literal["A", "B", "C", "D", "E", "F", "G", "H"] + +@dataclass +class Reg: + reg: RegChar + + def __str__(self): + return f"R{self.reg}" + +ALL_REGS: list[RegChar] = ["A", "B", "C", "D", "E", "F", "G", "H"] +STACK_POINTER_REG = "G" +CSR_SCRATCH_REG = "H" +SPECIAL_REGS: list[RegChar] = [STACK_POINTER_REG, CSR_SCRATCH_REG] +GENERIC_REGS: OrderedDict[RegChar, None] = OrderedDict() +for char in ALL_REGS: + if char not in SPECIAL_REGS: + GENERIC_REGS[cast(RegChar, char)] = None + +def format_asm_row(asm: str) -> str: + if not (asm.startswith("@") or asm.startswith(";")) and not asm.startswith(" ") and len(asm) > 0: + return " " + asm + else: + return asm + +class Compiler(ast.NodeVisitor): + def __init__(self): + self.const_asm: list[str] = [] + self.program_asm: list[str] = [ + "@label main" + ] + self.const_bindings: dict[str, Label] = {} + self.call_depth: int = 0 + self.unique_name_counter = 0 + + self.reserved_regs: OrderedDict[RegChar, None] = OrderedDict() + + def get_unique_name(self, prefix: str): + ret = f"{prefix}_{self.unique_name_counter}" + self.unique_name_counter += 1 + return ret + + def assign_const(self, name: str, value: int): + self.const_asm.append(f"@label {name}") + self.const_asm.append(f" {value}") + self.const_bindings[name] = name + + def emit(self, asm: str): + asm = asm.strip() + asm = format_asm_row(asm) + + self.program_asm.append(asm) + + def alloc_reg(self) -> Reg: + for reg in GENERIC_REGS: + if not reg in self.reserved_regs: + self.reserved_regs[reg] = None + return Reg(reg) + + raise Exception("Ran out of registers, TODO use stack") + + def free_reg(self, reg: Reg): + self.reserved_regs.pop(reg.reg) + + def compile(self, bootstrap_asm: str, source: str) -> str: + tree = ast.parse(source) + print(ast.dump(tree, indent=4)) + self.visit(tree) + return "\n".join([ + bootstrap_asm, + "", + "\n".join(self.const_asm), + "", + "\n".join(self.program_asm) + ]) + + def emit_builtin_call(self, name: str, args: list[ast.expr]): + self.emit(f"; builtin call {name} {args}") + match name: + case "asm": + match args: + case [ast.Constant(value)]: + if type(value) != str: + raise Exception(f"asm: invalid arg type: {type(value)}") + + self.emit(value) + case _: raise Exception(f"asm: invalid args: {args}") + case "set_graphics_mode": + match args: + case [ast.Constant(value)]: + if type(value) != int: + raise Exception(f"set_graphics_mode: invalid arg type: {type(value)}") + + reg1 = self.alloc_reg() + reg2 = self.alloc_reg() + self.emit(f"ldi vt_gr_mode_addr {reg1}") + self.emit(f"ldr {reg1} {reg1}") + self.emit(f"ldi {value} {reg2}") + self.emit(f"str {reg2} {reg1}") + self.free_reg(reg1) + self.free_reg(reg2) + + case _: raise Exception(f"set_graphics_mode: invalid args: {args}") + case "store": + if len(args) != 2: + raise Exception("Invalid number of arguments to store: " + str(len(args))) + + self.eval_expr_and_spu(args[0]) + self.eval_expr_and_spu(args[1]) + arg1 = self.alloc_reg() + arg2 = self.alloc_reg() + self.emit(f"spo {arg2}") + self.emit(f"spo {arg1}") + self.emit(f"str {arg2} {arg1}") + self.free_reg(arg1) + self.free_reg(arg2) + + def eval_int_constant_and_spu(self, value: int): + reg = self.alloc_reg() + if value >= 0 and value < 8: + self.emit(f"ldi {value} {reg}") + else: + name = self.get_unique_name("int") + self.assign_const(name, value) + self.emit(f"ldi {name} {reg}") + self.emit(f"ldr {reg} {reg}") + + self.emit(f"spu {reg}") + self.free_reg(reg) + + def eval_expr_and_spu(self, expr: ast.expr): + match expr: + case ast.Constant(int(value)): + self.eval_int_constant_and_spu(value) + case ast.Constant(str(value)): + if len(value) > 1: + raise Exception("Invalid string, only single char values allowed: " + value) + + c = value[0] + int_value = ord(c) + self.eval_int_constant_and_spu(int_value) + + + def visit_Module(self, node: Module): + for stmt in node.body: + self.visit(stmt) + + def visit_Expr(self, expr: Expr): + if self.call_depth == 0 and type(expr.value) != ast.Call: + self.emit("; NOP top-level expression") + return + + match expr.value: + case ast.Call(func, args, keywords): + match func: + case ast.Attribute(ast.Name(id="atk16"), attr): + print("FOUND CAPTURED ATK16 CALL: " + attr) + self.emit_builtin_call(attr, args) + case ast.Name(name): + addr = self.const_bindings[name] + self.emit(f"csi {addr}") + case _: + raise NotImplementedError("Unhandled Call: " + str(func)) + case _: + raise NotImplementedError("Unhandled Expr.value: " + str(expr.value)) + + def visit_AnnAssign(self, node: AnnAssign): + match node: + case ast.AnnAssign( + target=ast.Name(name), + annotation=ast.Attribute( + value=ast.Name(id="atk16"), + attr="ConstInt"), + value=ast.Constant(value) # TODO: constant folding + ): + if type(value) == int: + self.assign_const(name, value) + return + + raise NotImplementedError("Unhandled AnnAssign:\n" + ast.dump(node, indent=4)) + +class Optimizer: + def __init__(self): + pass + + def optimize(self, asm_str: str): + asm = asm_str.split("\n") + asm = [row.strip() for row in asm] + asm = [self.strip_comment(row) for row in asm] + asm = [row for row in asm if not len(row) == 0] + asm = [row.split() for row in asm] + asm = self.compact_spu_spo_pattern(asm) + asm = self.compact_load_mov_pattern(asm) + asm = self.compact_load_mov_pattern(asm) + asm = self.compact_spu_load_spo_pattern(asm) + result = "\n".join([format_asm_row(" ".join(row)) for row in asm]) + + return result + + def strip_comment(self, row: str): + ret: str = "" + for c in row: + if c == ";": break + ret += c + + return ret + + def compact_spu_spo_pattern(self, asm: list[list[str]]) -> list[list[str]]: + i = 0 + result: list[list[str]] = [] + while i < len(asm): + current = asm[i] + next = asm[i + 1] if i + 1 < len(asm) else None + i += 1 + if current[0].startswith("@") or next is None: + result.append(current) + continue + + if current[0] == "spu" and next[0] == "spo": + arg_current = current[1] + arg_next = next[1] + + if arg_current == arg_next: + pass # remove both spu and spo + else: + mov = ["mov", arg_current, arg_next] + result.append(mov) + + i += 1 + continue + + result.append(current) + + return result + + def compact_load_mov_pattern(self, asm: list[list[str]]) -> list[list[str]]: + i = 0 + result: list[list[str]] = [] + while i < len(asm): + current = asm[i] + next = asm[i + 1] if i + 1 < len(asm) else None + i += 1 + if current[0].startswith("@") or next is None: + result.append(current) + continue + + if (current[0] == "ldi" or current[0] == "ldr") and next[0] == "mov": + load_op, load_from, load_target_reg = current[0], current[1], current[2] + mov_from_reg, mov_to_reg = next[1], next[2] + + if load_target_reg == mov_from_reg: + ret = [load_op, load_from, mov_to_reg] + result.append(ret) + i += 1 + continue + + result.append(current) + + return result + + def compact_spu_load_spo_pattern(self, asm: list[list[str]]) -> list[list[str]]: + # spu RA + # ldi int_7 RA + # ldr RA RB + # spo RA + # OR + # spu RA + # ldi 3 RB + # spo RA + + i = 0 + result: list[list[str]] = [] + while i < len(asm): + instr0 = asm[i] + instr1 = asm[i + 1] if i + 1 < len(asm) else None + instr2 = asm[i + 2] if i + 2 < len(asm) else None + instr3 = asm[i + 3] if i + 3 < len(asm) else None + i += 1 + if instr0[0].startswith("@") or instr1 is None or instr2 is None or instr3 is None: + result.append(instr0) + continue + + if instr0[0] == "spu" and instr1[0] == "ldi" and instr2[0] == "ldr" and instr3[0] == "spo": + spu_op, spu_reg = instr0 + ldi_op, ldi_imm, ldi_target_reg = instr1 + ldr_op, ldr_from_reg, ldr_to_reg = instr2 + spo_op, spo_reg = instr3 + + if spu_reg == spo_reg and ldi_target_reg == ldr_from_reg and ldr_from_reg != ldr_to_reg: + ret0 = f"ldi {ldi_imm} {ldr_to_reg}".split() + ret1 = f"ldr {ldr_to_reg} {ldr_to_reg}".split() + result.append(ret0) + result.append(ret1) + i += 3 + continue + + elif instr0[0] == "spu" and instr1[0] == "ldi" and instr2[0] == "spo": + spu_op, spu_reg = instr0 + ldi_op, ldi_imm, ldi_target_reg = instr1 + spo_op, spo_reg = instr2 + + if spu_reg == spo_reg and ldi_target_reg != spu_reg: + ret0 = f"ldi {ldi_imm} {ldi_target_reg}".split() + result.append(ret0) + i += 2 + continue + + result.append(instr0) + + return result + +with open(infile_path, "r") as f: + source_py = f.read() + +with open("asm/ast_compiler_bootstrap.atk16", "r") as f: + bootstrap_asm = f.read() + +compiler = Compiler() +asm_out = compiler.compile( + bootstrap_asm, + source_py, +) + +optimizer = Optimizer() +asm_out_optimized = optimizer.optimize(asm_out) + +with open(outfile_path, "w") as f: + f.write(asm_out) + +with open(f"{outfile_path}_optimized", "w") as f: + f.write(asm_out_optimized) |