phasm/compile.py
2021-04-03 12:26:06 +02:00

134 lines
3.3 KiB
Python

import _ast
import ast
import sys
class Import:
def __init__(self, module, name, intname):
self.module = module
self.name = name
self.intname = intname
self.params = None
def generate(self):
return '(import "{}" "{}" (func ${}{}))'.format(
self.module,
self.name,
self.intname,
''.join(' (param {})'.format(x) for x in self.params)
)
class Statement:
def __init__(self, name, *args):
self.name = name
self.args = args
def generate(self):
return '{} {}'.format(self.name, ' '.join(self.args))
class Function:
def __init__(self, name, exported=True):
self.name = name
self.exported = exported # TODO: Use __all__!
self.statements = []
def generate(self):
return '(func {}\n {})'.format(
('(export "{}")' if self.exported else '${}').format(self.name),
'\n '.join(x.generate() for x in self.statements),
)
class Visitor(ast.NodeVisitor):
def __init__(self):
self._stack = []
self.imports = []
self.functions = []
def visit_ImportFrom(self, node):
for alias in node.names:
self.imports.append(Import(
node.module,
alias.name,
alias.asname,
))
def visit_FunctionDef(self, node):
func = Function(
node.name,
)
self._stack.append(func)
self.generic_visit(node)
self._stack.pop()
self.functions.append(func)
def visit_Expr(self, node):
self.generic_visit(node)
def visit_Call(self, node):
self.generic_visit(node)
func = self._stack[-1]
func.statements.append(
Statement('call', '$' + node.func.id)
)
def visit_BinOp(self, node):
self.generic_visit(node)
func = self._stack[-1]
if 'Add' == node.op.__class__.__name__:
func.statements.append(
Statement('i32.add')
)
elif 'Mult' == node.op.__class__.__name__:
func.statements.append(
Statement('i32.mul')
)
else:
err(node.op)
def visit_Constant(self, node):
if not self._stack:
# Constant outside of any function
imp = self.imports[-1]
prefix = imp.name + '('
val = node.value.strip()
if val.startswith(prefix) and val.endswith(')'):
imp.params = val[len(prefix):-1].split(',')
else:
func = self._stack[-1]
if isinstance(node.value, int):
func.statements.append(
Statement('i32.const', str(node.value))
)
self.generic_visit(node)
def generate(self):
return '(module\n {}\n {})'.format(
'\n '.join(x.generate() for x in self.imports),
'\n '.join(x.generate() for x in self.functions),
)
def err(msg: str) -> None:
sys.stderr.write('{}\n'.format(msg))
def main(source: str) -> int:
with open(source, 'r') as fil:
code = fil.read()
res = ast.parse(code, source)
visitor = Visitor()
visitor.visit(res)
print(visitor.generate())
return 0
if __name__ == '__main__':
sys.exit(main(*sys.argv[1:]))