95 lines
2.5 KiB
Python
95 lines
2.5 KiB
Python
import ast
|
|
|
|
from .wasm import Function, Import, Module, Statement
|
|
|
|
class Visitor(ast.NodeVisitor):
|
|
def __init__(self, module: Module) -> None:
|
|
self._stack = []
|
|
self.module = module
|
|
|
|
# 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):
|
|
if node.decorator_list:
|
|
# TODO: Support normal decorators
|
|
assert 1 == len(node.decorator_list)
|
|
call = node.decorator_list[0]
|
|
assert 'external' == call.func.id
|
|
assert 1 == len(call.args)
|
|
assert isinstance(call.args[0].value, str)
|
|
|
|
import_ = Import(
|
|
call.args[0].value,
|
|
node.name,
|
|
node.name,
|
|
)
|
|
|
|
import_.params = [
|
|
arg.annotation.id
|
|
for arg in node.args.args
|
|
]
|
|
|
|
self.module.imports.append(import_)
|
|
return
|
|
|
|
func = Function(
|
|
node.name,
|
|
)
|
|
|
|
self._stack.append(func)
|
|
self.generic_visit(node)
|
|
self._stack.pop()
|
|
|
|
self.module.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:
|
|
raise NotImplementedError
|
|
|
|
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)
|