phasm/phasm/codestyle.py

"""
This module generates source code based on the parsed AST

It's intented to be a "any color, as long as it's black" kind of renderer
"""
from typing import Generator

from . import ourlang
from . import typing

def phasm_render(inp: ourlang.Module) -> str:
    """
    Public method for rendering a Phasm module into Phasm code
    """
    return module(inp)

Statements = Generator[str, None, None]

def type_(inp: typing.TypeBase) -> str:
    """
    Render: Type (name)
    """
    if isinstance(inp, typing.TypeNone):
        return 'None'

    if isinstance(inp, typing.TypeBool):
        return 'bool'

    if isinstance(inp, typing.TypeUInt8):
        return 'u8'

    if isinstance(inp, typing.TypeUInt32):
        return 'u32'

    if isinstance(inp, typing.TypeUInt64):
        return 'u64'

    if isinstance(inp, typing.TypeInt32):
        return 'i32'

    if isinstance(inp, typing.TypeInt64):
        return 'i64'

    if isinstance(inp, typing.TypeFloat32):
        return 'f32'

    if isinstance(inp, typing.TypeFloat64):
        return 'f64'

    if isinstance(inp, typing.TypeBytes):
        return 'bytes'

    if isinstance(inp, typing.TypeTuple):
        mems = ', '.join(
            type_(x.type)
            for x in inp.members
        )

        return f'({mems}, )'

    if isinstance(inp, typing.TypeStaticArray):
        return f'{type_(inp.member_type)}[{len(inp.members)}]'

    if isinstance(inp, typing.TypeStruct):
        return inp.name

    raise NotImplementedError(type_, inp)

def struct_definition(inp: typing.TypeStruct) -> str:
    """
    Render: TypeStruct's definition
    """
    result = f'class {inp.name}:\n'
    for mem in inp.members:
        result += f'    {mem.name}: {type_(mem.type)}\n'

    return result

def constant_definition(inp: ourlang.ModuleConstantDef) -> str:
    """
    Render: Module Constant's  definition
    """
    return f'{inp.name}: {type_(inp.type)} = {expression(inp.constant)}\n'

def expression(inp: ourlang.Expression) -> str:
    """
    Render: A Phasm expression
    """
    if isinstance(inp, (
            ourlang.ConstantUInt8, ourlang.ConstantUInt32, ourlang.ConstantUInt64,
            ourlang.ConstantInt32, ourlang.ConstantInt64,
        )):
        return str(inp.value)

    if isinstance(inp, (ourlang.ConstantFloat32, ourlang.ConstantFloat64, )):
        # These might not round trip if the original constant
        # could not fit in the given float type
        return str(inp.value)

    if isinstance(inp, (ourlang.ConstantTuple, ourlang.ConstantStaticArray, )):
        return '(' + ', '.join(
            expression(x)
            for x in inp.value
        ) + ', )'

    if isinstance(inp, ourlang.VariableReference):
        return str(inp.name)

    if isinstance(inp, ourlang.UnaryOp):
        if (
            inp.operator in ourlang.WEBASSEMBLY_BUILDIN_FLOAT_OPS
            or inp.operator in ourlang.WEBASSEMBLY_BUILDIN_BYTES_OPS):
            return f'{inp.operator}({expression(inp.right)})'

        if inp.operator == 'cast':
            return f'{type_(inp.type)}({expression(inp.right)})'

        return f'{inp.operator}{expression(inp.right)}'

    if isinstance(inp, ourlang.BinaryOp):
        return f'{expression(inp.left)} {inp.operator} {expression(inp.right)}'

    if isinstance(inp, ourlang.FunctionCall):
        args = ', '.join(
            expression(arg)
            for arg in inp.arguments
        )

        if isinstance(inp.function, ourlang.StructConstructor):
            return f'{inp.function.struct.name}({args})'

        if isinstance(inp.function, ourlang.TupleConstructor):
            return f'({args}, )'

        return f'{inp.function.name}({args})'

    if isinstance(inp, ourlang.AccessBytesIndex):
        return f'{expression(inp.varref)}[{expression(inp.index)}]'

    if isinstance(inp, ourlang.AccessStructMember):
        return f'{expression(inp.varref)}.{inp.member.name}'

    if isinstance(inp, (ourlang.AccessTupleMember, ourlang.AccessStaticArrayMember, )):
        if isinstance(inp.member, ourlang.Expression):
            return f'{expression(inp.varref)}[{expression(inp.member)}]'

        return f'{expression(inp.varref)}[{inp.member.idx}]'

    if isinstance(inp, ourlang.Fold):
        fold_name = 'foldl' if ourlang.Fold.Direction.LEFT == inp.dir else 'foldr'
        return f'{fold_name}({inp.func.name}, {expression(inp.base)}, {expression(inp.iter)})'

    if isinstance(inp, ourlang.ModuleConstantReference):
        return inp.definition.name

    raise NotImplementedError(expression, inp)

def statement(inp: ourlang.Statement) -> Statements:
    """
    Render: A list of Phasm statements
    """
    if isinstance(inp, ourlang.StatementPass):
        yield 'pass'
        return

    if isinstance(inp, ourlang.StatementReturn):
        yield f'return {expression(inp.value)}'
        return

    if isinstance(inp, ourlang.StatementIf):
        yield f'if {expression(inp.test)}:'

        for stmt in inp.statements:
            for line in statement(stmt):
                yield f'    {line}' if line else ''

        yield ''
        return

    raise NotImplementedError(statement, inp)

def function(inp: ourlang.Function) -> str:
    """
    Render: Function body

    Imported functions only have "pass" as a body. Later on we might replace
    this by the function documentation, if any.
    """
    result = ''
    if inp.exported:
        result += '@exported\n'
    if inp.imported:
        result += '@imported\n'

    args = ', '.join(
        f'{x}: {type_(y)}'
        for x, y in inp.posonlyargs
    )

    result += f'def {inp.name}({args}) -> {type_(inp.returns)}:\n'

    if inp.imported:
        result += '    pass\n'
    else:
        for stmt in inp.statements:
            for line in statement(stmt):
                result += f'    {line}\n' if line else '\n'

    return result


def module(inp: ourlang.Module) -> str:
    """
    Render: Module
    """
    result = ''

    for struct in inp.structs.values():
        if result:
            result += '\n'
        result += struct_definition(struct)

    for cdef in inp.constant_defs.values():
        if result:
            result += '\n'
        result += constant_definition(cdef)

    for func in inp.functions.values():
        if func.lineno < 0:
            # Buildin (-2) or auto generated (-1)
            continue

        if result:
            result += '\n'
        result += function(func)

    return result