phasm/py2wasm/wasm.py

"""
Python classes for storing the representation of Web Assembly code
"""

from typing import Any, Iterable, List, Optional, Tuple, Union

###
### This part is more intermediate code
###

class OurType:
    def to_wasm(self) -> str:
        raise NotImplementedError

    def alloc_size(self) -> int:
        raise NotImplementedError

class OurTypeNone(OurType):
    pass

class OurTypeBool(OurType):
    pass

class OurTypeInt32(OurType):
    def to_wasm(self) -> str:
        return 'i32'

    def alloc_size(self) -> int:
        return 4

class OurTypeInt64(OurType):
    def to_wasm(self) -> str:
        return 'i64'

    def alloc_size(self) -> int:
        return 8

class OurTypeFloat32(OurType):
    def to_wasm(self) -> str:
        return 'f32'

    def alloc_size(self) -> int:
        return 4

class OurTypeFloat64(OurType):
    def to_wasm(self) -> str:
        return 'f64'

    def alloc_size(self) -> int:
        return 8

class Constant:
    """
    TODO
    """
    def __init__(self, value: Union[None, bool, int, float]) -> None:
        self.value = value

class ClassMember:
    """
    Represents a class member
    """
    def __init__(self, name: str, type_: OurType, offset: int, default: Optional[Constant]) -> None:
        self.name = name
        self.type = type_
        self.offset = offset
        self.default = default

class OurTypeClass(OurType):
    """
    Represents a class
    """
    def __init__(self, name: str, members: List[ClassMember]) -> None:
        self.name = name
        self.members = members

    def to_wasm(self) -> str:
        return 'i32' # WASM uses 32 bit pointers

    def alloc_size(self) -> int:
        return sum(x.type.alloc_size() for x in self.members)

class TupleMember:
    """
    Represents a tuple member
    """
    def __init__(self, type_: OurType, offset: int) -> None:
        self.type = type_
        self.offset = offset

class OurTypeTuple(OurType):
    """
    Represents a tuple
    """
    def __init__(self, members: List[TupleMember]) -> None:
        self.members = members

    def to_wasm(self) -> str:
        return 'i32' # WASM uses 32 bit pointers

    def alloc_size(self) -> int:
        return sum(x.type.alloc_size() for x in self.members)


Param = Tuple[str, OurType]

###
## This part is more actual web assembly
###

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

    def generate(self) -> str:
        """
        Generates the text version
        """
        return '(import "{}" "{}" (func ${}{}))'.format(
            self.module,
            self.name,
            self.intname,
            ''.join(' (param {})'.format(x[1]) for x in self.params)
        )

class Statement:
    """
    Represents a Web Assembly statement
    """
    def __init__(self, name: str, *args: str, comment: Optional[str] = None):
        self.name = name
        self.args = args
        self.comment = comment

    def generate(self) -> str:
        """
        Generates the text version
        """
        args = ' '.join(self.args)
        comment = f' ;; {self.comment}' if self.comment else ''

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

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

    def generate(self) -> str:
        """
        Generates the text version
        """
        header = ('(export "{}")' if self.exported else '${}').format(self.name)

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

        if self.result:
            header += f' (result {self.result.to_wasm()})'

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

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

class Module:
    """
    Represents a Web Assembly module
    """
    def __init__(self) -> None:
        self.imports: List[Import] = []
        self.functions: List[Function] = []

    def generate(self) -> str:
        """
        Generates the text version
        """
        return '(module\n  (memory 1)\n  (data (memory 0) (i32.const 0) {})\n  {}\n  {})\n'.format(
            '"\\04\\00\\00\\00"',
            '\n  '.join(x.generate() for x in self.imports),
            '\n  '.join(x.generate() for x in self.functions),
        )