phasm/phasm/wasm.py

"""
Python classes for storing the representation of Web Assembly code,
and being able to conver it to Web Assembly Text Format
"""

from typing import Iterable, List, Optional, Tuple


class WatSerializable:
    """
    Mixin for clases that can be serialized as WebAssembly Text
    """
    def to_wat(self) -> str:
        """
        Renders this object as WebAssembly Text
        """
        raise NotImplementedError(self, 'to_wat')

class WasmType(WatSerializable):
    """
    Type base class
    """

class WasmTypeNone(WasmType):
    """
    Type when there is no type
    """
    def to_wat(self) -> str:
        raise Exception('None type is only a placeholder')

class WasmTypeInt32(WasmType):
    """
    i32 value

    Signed or not depends on the operations, not the type
    """
    def to_wat(self) -> str:
        return 'i32'

class WasmTypeInt64(WasmType):
    """
    i64 value

    Signed or not depends on the operations, not the type
    """
    def to_wat(self) -> str:
        return 'i64'

class WasmTypeFloat32(WasmType):
    """
    f32 value
    """
    def to_wat(self) -> str:
        return 'f32'

class WasmTypeFloat64(WasmType):
    """
    f64 value
    """
    def to_wat(self) -> str:
        return 'f64'

class WasmTypeVector(WasmType):
    """
    A vector is a 128-bit value
    """
    def to_wat(self) -> str:
        return 'v128'

class WasmTypeVectorInt32x4(WasmTypeVector):
    """
    4 Int32 values in a single vector
    """

Param = Tuple[str, WasmType]

class Import(WatSerializable):
    """
    Represents a Web Assembly import
    """
    def __init__(
            self,
            module: str,
            name: str,
            intname: str,
            params: Iterable[Param],
            result: WasmType,
        ) -> None:
        self.module = module
        self.name = name
        self.intname = intname
        self.params = [*params]
        self.result = result

    def to_wat(self) -> str:
        return '(import "{}" "{}" (func ${}{}{}))'.format(
            self.module,
            self.name,
            self.intname,
            ''.join(
                f' (param {typ.to_wat()})'
                for _, typ in self.params
            ),
            '' if isinstance(self.result, WasmTypeNone)
               else f' (result {self.result.to_wat()})'
        )

class StatementBase(WatSerializable):
    pass

class Statement(StatementBase   ):
    """
    Represents a Web Assembly statement
    """
    def __init__(self, name: str, *args: str, comment: Optional[str] = None):
        assert ' ' not in name, 'Please pass argument separately'
        assert name != 'call', 'Please use StatementCall'

        self.name = name
        self.args = args
        self.comment = comment

    def to_wat(self) -> str:
        args = ' '.join(self.args)
        comment = f' ;; {self.comment}' if self.comment else ''

        return f'{self.name} {args}{comment}'

class StatementCall(StatementBase   ):
    def __init__(self, func_name: str, comment: str | None = None):
        self.func_name = func_name
        self.comment = comment

    def to_wat(self) -> str:
        comment = f' ;; {self.comment}' if self.comment else ''

        return f'call ${self.func_name} {comment}'

class Function(WatSerializable):
    """
    Represents a Web Assembly function
    """
    def __init__(
            self,
            name: str,
            exported_name: Optional[str],
            params: Iterable[Param],
            locals_: Iterable[Param],
            result: WasmType,
            statements: Iterable[StatementBase],
        ) -> None:
        self.name = name
        self.exported_name = exported_name
        self.params = [*params]
        self.locals = [*locals_]
        self.result = result
        self.statements = [*statements]

    def to_wat(self) -> str:
        header = f'${self.name}' # Name for internal use

        if self.exported_name is not None:
            # Name for external use
            header += f' (export "{self.exported_name}")'

        for nam, typ in self.params:
            header += f' (param ${nam} {typ.to_wat()})'

        if not isinstance(self.result, WasmTypeNone):
            header += f' (result {self.result.to_wat()})'

        for nam, typ in self.locals:
            header += f' (local ${nam} {typ.to_wat()})'

        return '(func {}\n    {}\n    )'.format(
            header,
            '\n    '.join(x.to_wat() for x in self.statements),
        )

class ModuleMemory(WatSerializable):
    """
    Represents a WebAssembly module's memory
    """
    def __init__(self, data: bytes = b'') -> None:
        self.data = data

    def to_wat(self) -> str:
        data = ''.join(
            f'\\{x:02x}'
            for x in self.data
        )

        return (
            '(memory 1)\n  '
            f'(data (memory 0) (i32.const 0) "{data}")\n  '
            '(export "memory" (memory 0))\n'
        )

class Module(WatSerializable):
    """
    Represents a Web Assembly module
    """
    def __init__(self) -> None:
        self.imports: List[Import] = []
        self.table: dict[int, str] = {}
        self.functions: List[Function] = []
        self.memory = ModuleMemory()

    def to_wat(self) -> str:
        """
        Generates the text version
        """
        return '(module\n  {}\n  {}\n  {}\n  {}\n  {})\n'.format(
            '\n  '.join(x.to_wat() for x in self.imports),
            f'(table {len(self.table)} funcref)',
            '\n  '.join(f'(elem (i32.const {k}) ${v})' for k, v in self.table.items()),
            self.memory.to_wat(),
            '\n  '.join(x.to_wat() for x in self.functions),
        )