155 lines
5.0 KiB
Python
155 lines
5.0 KiB
Python
"""
|
|
Helper functions to generate WASM code by writing Python functions
|
|
"""
|
|
from typing import Any, Callable, Dict, List, Optional, Type
|
|
|
|
import functools
|
|
|
|
from . import wasm
|
|
|
|
# pylint: disable=C0103,C0115,C0116,R0902
|
|
|
|
class VarType_Base:
|
|
wasm_type: Type[wasm.WasmType]
|
|
|
|
def __init__(self, name: str) -> None:
|
|
self.name = name
|
|
self.name_ref = f'${name}'
|
|
|
|
class VarType_i32(VarType_Base):
|
|
wasm_type = wasm.WasmTypeInt32
|
|
|
|
class Generator_i32:
|
|
def __init__(self, generator: 'Generator') -> None:
|
|
self.generator = generator
|
|
|
|
# 2.4.1. Numeric Instructions
|
|
# ibinop
|
|
self.add = functools.partial(self.generator.add, 'i32.add')
|
|
|
|
# irelop
|
|
self.eq = functools.partial(self.generator.add, 'i32.eq')
|
|
self.ne = functools.partial(self.generator.add, 'i32.ne')
|
|
|
|
# 2.4.4. Memory Instructions
|
|
self.load = functools.partial(self.generator.add, 'i32.load')
|
|
self.store = functools.partial(self.generator.add, 'i32.store')
|
|
|
|
def const(self, value: int, comment: Optional[str] = None) -> None:
|
|
self.generator.add('i32.const', f'0x{value:08x}', comment=comment)
|
|
|
|
class Generator_Local:
|
|
def __init__(self, generator: 'Generator') -> None:
|
|
self.generator = generator
|
|
|
|
# 2.4.3. Variable Instructions
|
|
def get(self, variable: VarType_Base, comment: Optional[str] = None) -> None:
|
|
self.generator.add('local.get', variable.name_ref, comment=comment)
|
|
|
|
def set(self, variable: VarType_Base, comment: Optional[str] = None) -> None:
|
|
self.generator.locals.setdefault(variable.name, variable)
|
|
|
|
self.generator.add('local.set', variable.name_ref, comment=comment)
|
|
|
|
def tee(self, variable: VarType_Base, comment: Optional[str] = None) -> None:
|
|
self.generator.locals.setdefault(variable.name, variable)
|
|
|
|
self.generator.add('local.tee', variable.name_ref, comment=comment)
|
|
|
|
class GeneratorBlock:
|
|
def __init__(self, generator: 'Generator', name: str) -> None:
|
|
self.generator = generator
|
|
self.name = name
|
|
|
|
def __enter__(self) -> None:
|
|
self.generator.add(self.name)
|
|
|
|
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
|
if not exc_type:
|
|
self.generator.add('end')
|
|
|
|
class Generator:
|
|
def __init__(self) -> None:
|
|
self.statements: List[wasm.Statement] = []
|
|
self.locals: Dict[str, VarType_Base] = {}
|
|
|
|
self.i32 = Generator_i32(self)
|
|
|
|
# 2.4.3 Variable Instructions
|
|
self.local = Generator_Local(self)
|
|
|
|
# 2.4.5 Control Instructions
|
|
self.nop = functools.partial(self.add, 'nop')
|
|
self.unreachable = functools.partial(self.add, 'unreachable')
|
|
# block
|
|
# loop
|
|
self.if_ = functools.partial(GeneratorBlock, self, 'if')
|
|
# br
|
|
# br_if
|
|
# br_table
|
|
self.return_ = functools.partial(self.add, 'return')
|
|
# call
|
|
# call_indirect
|
|
|
|
def call(self, function: wasm.Function) -> None:
|
|
self.add('call', f'${function.name}')
|
|
|
|
def add(self, name: str, *args: str, comment: Optional[str] = None) -> None:
|
|
self.statements.append(wasm.Statement(name, *args, comment=comment))
|
|
|
|
def func_wrapper(exported: bool = True) -> Callable[[Any], wasm.Function]:
|
|
"""
|
|
This wrapper will execute the function and return
|
|
a wasm Function with the generated Statements
|
|
"""
|
|
def inner(func: Any) -> wasm.Function:
|
|
func_name_parts = func.__module__.split('.') + [func.__name__]
|
|
if 'phasm' == func_name_parts[0]:
|
|
func_name_parts.pop(0)
|
|
func_name = '.'.join(func_name_parts)
|
|
|
|
annot = dict(func.__annotations__)
|
|
|
|
# Check if we can pass the generator
|
|
assert Generator is annot.pop('g')
|
|
|
|
# Convert return type to WasmType
|
|
annot_return = annot.pop('return')
|
|
if annot_return is None:
|
|
return_type = wasm.WasmTypeNone()
|
|
else:
|
|
assert issubclass(annot_return, VarType_Base)
|
|
return_type = annot_return.wasm_type()
|
|
|
|
# Load the argument types, and generate instances
|
|
args: Dict[str, VarType_Base] = {}
|
|
params: List[wasm.Param] = []
|
|
for param_name, param_type in annot.items():
|
|
assert issubclass(param_type, VarType_Base)
|
|
|
|
params.append((param_name, param_type.wasm_type(), ))
|
|
|
|
args[param_name] = VarType_Base(param_name)
|
|
|
|
# Make a generator, and run the function on that generator,
|
|
# so the Statements get added
|
|
generator = Generator()
|
|
func(g=generator, **args)
|
|
|
|
# Check what locals were used, and define them
|
|
locals_: List[wasm.Param] = []
|
|
for local_name, local_type in generator.locals.items():
|
|
locals_.append((local_name, local_type.wasm_type(), ))
|
|
|
|
# Complete function definition
|
|
return wasm.Function(
|
|
func_name,
|
|
func_name if exported else None,
|
|
params,
|
|
locals_,
|
|
return_type,
|
|
generator.statements,
|
|
)
|
|
|
|
return inner
|