0e8540c611
Had to implement both functions as arguments and type place holders (variables) for type constructors. Had to implement functions as a type as well. Still have to figure out how to pass functions around.
303 lines
13 KiB
Python
303 lines
13 KiB
Python
"""
|
|
Helper functions to generate WASM code by writing Python functions
|
|
"""
|
|
import functools
|
|
from typing import Any, Callable, Dict, Iterable, List, Optional, Type
|
|
|
|
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_u8(VarType_Base):
|
|
wasm_type = wasm.WasmTypeInt32
|
|
|
|
class VarType_i32(VarType_Base):
|
|
wasm_type = wasm.WasmTypeInt32
|
|
|
|
class VarType_i64(VarType_Base):
|
|
wasm_type = wasm.WasmTypeInt64
|
|
|
|
class VarType_f32(VarType_Base):
|
|
wasm_type = wasm.WasmTypeFloat32
|
|
|
|
class VarType_f64(VarType_Base):
|
|
wasm_type = wasm.WasmTypeFloat64
|
|
|
|
class Generator_i32i64:
|
|
def __init__(self, prefix: str, generator: 'Generator') -> None:
|
|
self.prefix = prefix
|
|
self.generator = generator
|
|
|
|
# 2.4.1. Numeric Instructions
|
|
# ibinop
|
|
self.add = functools.partial(self.generator.add_statement, f'{prefix}.add')
|
|
self.sub = functools.partial(self.generator.add_statement, f'{prefix}.sub')
|
|
self.mul = functools.partial(self.generator.add_statement, f'{prefix}.mul')
|
|
self.div_s = functools.partial(self.generator.add_statement, f'{prefix}.div_s')
|
|
self.div_u = functools.partial(self.generator.add_statement, f'{prefix}.div_u')
|
|
self.rem_s = functools.partial(self.generator.add_statement, f'{prefix}.rem_s')
|
|
self.rem_u = functools.partial(self.generator.add_statement, f'{prefix}.rem_u')
|
|
self.and_ = functools.partial(self.generator.add_statement, f'{prefix}.and')
|
|
self.or_ = functools.partial(self.generator.add_statement, f'{prefix}.or')
|
|
self.xor = functools.partial(self.generator.add_statement, f'{prefix}.xor')
|
|
self.shl = functools.partial(self.generator.add_statement, f'{prefix}.shl')
|
|
self.shr_s = functools.partial(self.generator.add_statement, f'{prefix}.shr_s')
|
|
self.shr_u = functools.partial(self.generator.add_statement, f'{prefix}.shr_u')
|
|
self.rotl = functools.partial(self.generator.add_statement, f'{prefix}.rotl')
|
|
self.rotr = functools.partial(self.generator.add_statement, f'{prefix}.rotr')
|
|
|
|
# itestop
|
|
self.eqz = functools.partial(self.generator.add_statement, f'{prefix}.eqz')
|
|
|
|
# irelop
|
|
self.eq = functools.partial(self.generator.add_statement, f'{prefix}.eq')
|
|
self.ne = functools.partial(self.generator.add_statement, f'{prefix}.ne')
|
|
self.lt_s = functools.partial(self.generator.add_statement, f'{prefix}.lt_s')
|
|
self.lt_u = functools.partial(self.generator.add_statement, f'{prefix}.lt_u')
|
|
self.gt_s = functools.partial(self.generator.add_statement, f'{prefix}.gt_s')
|
|
self.gt_u = functools.partial(self.generator.add_statement, f'{prefix}.gt_u')
|
|
self.le_s = functools.partial(self.generator.add_statement, f'{prefix}.le_s')
|
|
self.le_u = functools.partial(self.generator.add_statement, f'{prefix}.le_u')
|
|
self.ge_s = functools.partial(self.generator.add_statement, f'{prefix}.ge_s')
|
|
self.ge_u = functools.partial(self.generator.add_statement, f'{prefix}.ge_u')
|
|
|
|
# 2.4.4. Memory Instructions
|
|
self.load = functools.partial(self.generator.add_statement, f'{prefix}.load')
|
|
self.load8_u = functools.partial(self.generator.add_statement, f'{prefix}.load8_u')
|
|
self.store = functools.partial(self.generator.add_statement, f'{prefix}.store')
|
|
|
|
def const(self, value: int, comment: Optional[str] = None) -> None:
|
|
self.generator.add_statement(f'{self.prefix}.const', f'{value}', comment=comment)
|
|
|
|
class Generator_i32(Generator_i32i64):
|
|
def __init__(self, generator: 'Generator') -> None:
|
|
super().__init__('i32', generator)
|
|
|
|
# 2.4.1. Numeric Instructions
|
|
self.wrap_i64 = functools.partial(self.generator.add_statement, 'i32.wrap_i64')
|
|
|
|
class Generator_i64(Generator_i32i64):
|
|
def __init__(self, generator: 'Generator') -> None:
|
|
super().__init__('i64', generator)
|
|
|
|
# 2.4.1. Numeric Instructions
|
|
self.extend_i32_s = functools.partial(self.generator.add_statement, 'i64.extend_i32_s')
|
|
self.extend_i32_u = functools.partial(self.generator.add_statement, 'i64.extend_i32_u')
|
|
|
|
class Generator_f32f64:
|
|
def __init__(self, prefix: str, generator: 'Generator') -> None:
|
|
self.prefix = prefix
|
|
self.generator = generator
|
|
|
|
# 2.4.1. Numeric Instructions
|
|
# funop
|
|
self.abs = functools.partial(self.generator.add_statement, f'{prefix}.abs')
|
|
self.neg = functools.partial(self.generator.add_statement, f'{prefix}.neg')
|
|
self.sqrt = functools.partial(self.generator.add_statement, f'{prefix}.sqrt')
|
|
self.ceil = functools.partial(self.generator.add_statement, f'{prefix}.ceil')
|
|
self.floor = functools.partial(self.generator.add_statement, f'{prefix}.floor')
|
|
self.trunc = functools.partial(self.generator.add_statement, f'{prefix}.trunc')
|
|
self.nearest = functools.partial(self.generator.add_statement, f'{prefix}.nearest')
|
|
|
|
# fbinop
|
|
self.add = functools.partial(self.generator.add_statement, f'{prefix}.add')
|
|
self.sub = functools.partial(self.generator.add_statement, f'{prefix}.sub')
|
|
self.mul = functools.partial(self.generator.add_statement, f'{prefix}.mul')
|
|
self.div = functools.partial(self.generator.add_statement, f'{prefix}.div')
|
|
self.min = functools.partial(self.generator.add_statement, f'{prefix}.min')
|
|
self.max = functools.partial(self.generator.add_statement, f'{prefix}.max')
|
|
self.copysign = functools.partial(self.generator.add_statement, f'{prefix}.copysign')
|
|
|
|
# frelop
|
|
self.eq = functools.partial(self.generator.add_statement, f'{prefix}.eq')
|
|
self.ne = functools.partial(self.generator.add_statement, f'{prefix}.ne')
|
|
self.lt = functools.partial(self.generator.add_statement, f'{prefix}.lt')
|
|
self.gt = functools.partial(self.generator.add_statement, f'{prefix}.gt')
|
|
self.le = functools.partial(self.generator.add_statement, f'{prefix}.le')
|
|
self.ge = functools.partial(self.generator.add_statement, f'{prefix}.ge')
|
|
|
|
# Other instr - convert
|
|
self.convert_i32_s = functools.partial(self.generator.add_statement, f'{prefix}.convert_i32_s')
|
|
self.convert_i32_u = functools.partial(self.generator.add_statement, f'{prefix}.convert_i32_u')
|
|
self.convert_i64_s = functools.partial(self.generator.add_statement, f'{prefix}.convert_i64_s')
|
|
self.convert_i64_u = functools.partial(self.generator.add_statement, f'{prefix}.convert_i64_u')
|
|
|
|
# 2.4.4. Memory Instructions
|
|
self.load = functools.partial(self.generator.add_statement, f'{prefix}.load')
|
|
self.store = functools.partial(self.generator.add_statement, f'{prefix}.store')
|
|
|
|
def const(self, value: float, comment: Optional[str] = None) -> None:
|
|
# FIXME: Is this sufficient to guarantee the float comes across properly?
|
|
self.generator.add_statement(f'{self.prefix}.const', f'{value}', comment=comment)
|
|
|
|
class Generator_f32(Generator_f32f64):
|
|
def __init__(self, generator: 'Generator') -> None:
|
|
super().__init__('f32', generator)
|
|
|
|
# 2.4.1 Numeric Instructions
|
|
self.demote_f64 = functools.partial(self.generator.add_statement, 'f32.demote_f64')
|
|
|
|
class Generator_f64(Generator_f32f64):
|
|
def __init__(self, generator: 'Generator') -> None:
|
|
super().__init__('f64', generator)
|
|
|
|
# 2.4.1 Numeric Instructions
|
|
self.promote_f32 = functools.partial(self.generator.add_statement, 'f64.promote_f32')
|
|
|
|
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_statement('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_statement('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_statement('local.tee', variable.name_ref, comment=comment)
|
|
|
|
class GeneratorBlock:
|
|
def __init__(self, generator: 'Generator', name: str, params: Iterable[str] = (), result: str | None = None, comment: str | None = None) -> None:
|
|
self.generator = generator
|
|
self.name = name
|
|
self.params = params
|
|
self.result = result
|
|
self.comment = comment
|
|
|
|
def __enter__(self) -> None:
|
|
stmt = self.name
|
|
if self.params:
|
|
stmt = f'{stmt} ' + ' '.join(
|
|
f'(param {typ})'
|
|
for typ in self.params
|
|
)
|
|
if self.result:
|
|
stmt = f'{stmt} (result {self.result})'
|
|
|
|
self.generator.add_statement(stmt, comment=self.comment)
|
|
|
|
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
|
if not exc_type:
|
|
self.generator.add_statement('end')
|
|
|
|
class Generator:
|
|
def __init__(self) -> None:
|
|
self.statements: List[wasm.Statement] = []
|
|
self.locals: Dict[str, VarType_Base] = {}
|
|
|
|
self.i32 = Generator_i32(self)
|
|
self.i64 = Generator_i64(self)
|
|
self.f32 = Generator_f32(self)
|
|
self.f64 = Generator_f64(self)
|
|
|
|
# 2.4.3 Variable Instructions
|
|
self.local = Generator_Local(self)
|
|
|
|
# 2.4.5 Control Instructions
|
|
self.nop = functools.partial(self.add_statement, 'nop')
|
|
self.unreachable = functools.partial(self.add_statement, 'unreachable')
|
|
self.block = functools.partial(GeneratorBlock, self, 'block')
|
|
self.loop = functools.partial(GeneratorBlock, self, 'loop')
|
|
self.if_ = functools.partial(GeneratorBlock, self, 'if')
|
|
# br
|
|
# br_if - see below
|
|
# br_table
|
|
self.return_ = functools.partial(self.add_statement, 'return')
|
|
# call - see below
|
|
# call_indirect
|
|
|
|
def br_if(self, idx: int) -> None:
|
|
self.add_statement('br_if', f'{idx}')
|
|
|
|
def call(self, function: wasm.Function) -> None:
|
|
self.add_statement('call', f'${function.name}')
|
|
|
|
def add_statement(self, name: str, *args: str, comment: Optional[str] = None) -> None:
|
|
self.statements.append(wasm.Statement(name, *args, comment=comment))
|
|
|
|
def temp_var[T: VarType_Base](self, var: T) -> T:
|
|
idx = 0
|
|
while (varname := f'__{var.name}_tmp_var_{idx}__') in self.locals:
|
|
idx += 1
|
|
|
|
return var.__class__(varname)
|
|
|
|
def temp_var_i32(self, infix: str) -> VarType_i32:
|
|
return self.temp_var(VarType_i32(infix))
|
|
|
|
def temp_var_u8(self, infix: str) -> VarType_u8:
|
|
return self.temp_var(VarType_u8(infix))
|
|
|
|
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():
|
|
if local_name in args:
|
|
# Already defined as a local by wasm itself
|
|
continue
|
|
|
|
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
|