<filename>backend/backend.py
"""C6T VM codegen backend.
"""

import re
import collections.abc
from string import whitespace
from collections import deque
from ctypes import Union
from typing import Callable, Dict, Deque, Optional, Tuple, overload, Iterator, Iterable, Any, List


class ValType:
    """Base class for a int,float,char,or double.
    """


class IntType(ValType):
    pass


class CharType(ValType):
    pass


class FltType(ValType):
    pass


class DblType(ValType):
    pass


class Flags(collections.abc.MutableMapping):
    """A collection of flags, which default to a given value.
    """

    def __init__(self, default: bool, **kwargs) -> None:
        super().__init__()
        if not isinstance(default, bool):
            raise TypeError("default value must be a bool")
        self.default = default
        self.flags = {}
        for name in kwargs:
            self[name] = True

    def __delitem__(self, __v: str) -> None:
        del self.flags[__v]

    def __getitem__(self, __k: str) -> bool:
        if __k in self.flags:
            return self.flags[__k]
        return self.default

    def __setitem__(self, __k: str, __v: bool) -> None:
        if not isinstance(__v, bool):
            raise TypeError("Flag values must be bools")
        self.flags[__k] = __v

    def __iter__(self) -> Iterator[str]:
        for name in self.flags:
            yield name

    def __len__(self) -> int:
        return len(self.flags)

    def __getattr__(self, name: str) -> bool:
        return self.flags[name]

    def __setattr__(self, __name: str, __value: Any) -> None:
        if __name in self.__dict__:
            super().__setattr__(__name, __value)
        self[__name] = __value


class Opcode:
    """A C6T opcode representation, with information on its properties.
    """

    def __init__(self, code: str, stk_args: int = None, inline_args: int = None, **flags) -> None:
        if stk_args is None:
            stk_args = 0
        if inline_args is None:
            inline_args = 0
        self.code = code
        self.stk_args = stk_args
        self.inline_args = inline_args
        self.flags = Flags(False, **flags)

    def build(self, stack: Deque, inline_args: Optional[list] = None):
        """Create the representation of the opcode.
        """
        return None


class Leaf:
    """Base class for a C6T leaf value (constant or memory cell).
    """


class FramePointer(Leaf):
    """Representation of the frame pointer.
    """

    def __str__(self) -> str:
        return "fp"


class MemoryCell(Leaf):
    """Base class for C6T memory cells.
    """


class Register(MemoryCell):
    """The value at the given register number.
    """

    def __init__(self, regnum: int) -> None:
        super().__init__()
        self.regnum = regnum

    def __str__(self) -> str:
        return f"reg{self.regnum}"


class Static(MemoryCell):
    """A static C6T memory cell.
    """

    def __init__(self, name: str):
        self.name = name

    def __str__(self) -> str:
        return self.name


class Con(Leaf):
    """A C6T integer constant.
    """

    def __init__(self, value: int):
        self.value = value

    def __str__(self) -> str:
        return str(self.value)


class Temp(MemoryCell):
    """A C6T temporary memory cell.
    """
    cur_index = 0

    def __init__(self):
        self.index = self.__class__.cur_index
        self.__class__.cur_index += 1

    def reset(self):
        """Reset the index of all temporaries.
        """
        self.__class__.cur_index = 0


class Node:
    """A C6T stack node.

    Raises:
        ValueError: Not enough inline arguments for the opcode
        ValueError: Not enough stack arguments for the opcode
    """
    @overload
    def __init__(self, opcode: Opcode, *children, args: Optional[list] = None) -> None:
        if args is None:
            args = []
        self.opcode = opcode
        self.children = list(children)
        self.args = args
        self.regnum = None  # type:Union[None, int]

    def __init__(self, opcode: Opcode, stack: deque, args: Optional[deque] = None) -> None:
        if args is None:
            args = []
        self.opcode = opcode
        self.children = []
        if len(args) != opcode.inline_args:
            raise ValueError(
                f"Not enough inline arguments (need {opcode.inline_args}, have {len(args)})")
        self.args = args
        if len(stack) < opcode.stk_args:
            raise ValueError(
                f"Not enough stack arguments (need {opcode.stk_args}, have {len(stack)})")
        for i in range(opcode.stk_args):
            self.children.insert(0, stack.pop())
        self.regnum = None  # type:Union[None, int]


class Statement:
    """A C6T VM statement node.
    """

    def __init__(self, statecode: str):
        self.statecode = statecode


class Move(Statement):
    """Statement to compute an expression and store its value.
    """

    def __init__(self, dest: Union[Leaf, Node], expr: Union[Node, Leaf], dest_type: ValType):
        super().__init__('move')
        self.dest = dest
        self.expr = expr
        self.dest_type = dest_type


class Eval(Statement):
    """Compute an expression and ignore the result.
    """

    def __init__(self, expr: Union[Node, Leaf]):
        super().__init__('eval')
        self.expr = expr


class Jump(Statement):
    """Jump unconditionall to an address.
    """

    def __init__(self, target: Union[Node, Leaf]):
        super().__init__('jump')
        self.target = target


class Brz(Statement):
    """Compute an expression and jump to the address if the expression is 0.
    """

    def __init__(self, expr: Union[Node, Leaf], target: Leaf):
        super().__init__('brz')
        self.expr = expr
        self.target = target


class Label(Statement):
    """Define a label here.
    """

    def __init__(self, name: str):
        super().__init__('label')
        self.name = name


class SeqState(Statement):
    """Perform multiple statements in order.
    """

    def __init__(self, *statements: Statement):
        super().__init__('seq')
        self.statements = list(statements)

    def __iter__(self):
        for statement in self.statements:
            if isinstance(statement, SeqState):
                for substate in statement:
                    yield substate
            else:
                yield statement

    def __len__(self):
        return len(list(self))

    def leaders(self) -> Iterable[int]:
        """Returns iterable of the indexes of all statements.
        """
        next_is_leader = False
        for index, statement in list(self):
            leader = None
            if next_is_leader:
                leader = index
            if index == 0:
                leader = index
            if isinstance(statement, (Brz, Jump)):
                next_is_leader = True
            if isinstance(statement, Label):
                leader = index
            else:
                next_is_leader = False
            if leader is not None:
                yield leader

    def basic_blocks(self) -> Iterable[List[Statement]]:
        """Iterator over all basic blocks.
        """
        statements = list(self)
        leaders = list(self.leaders())
        while len(leaders) > 2:
            leader = leaders.pop(0)
            tail = leaders[0]
            yield statements[leader:tail]
        if len(leaders) > 1:
            yield statements[leaders[0]:]


class StkOp(Opcode):
    """A opcode which performs stack operations.
    """

    def build(self, stack: deque, inline_args: Optional[list] = None) -> None:
        """Create an expression node from this opcode."""
        if inline_args is None:
            inline_args = []
        stack.append(Node(self, stack, inline_args))
        return None


class LexName:
    """A lexical token.
    """

    def __init__(self, name: str) -> None:
        self.name = name


class PushOp(Opcode):
    """Opcode to push a literal value.
    """

    def __init__(self) -> None:
        super().__init__('push', 0, 1)

    def build(self, stack: Deque, inline_args: Optional[list] = None):
        if len(inline_args) != self.inline_args:
            raise ValueError("Not enough inline args")
        literal = inline_args[0]
        if isinstance(literal, LexName):
            literal = Static(literal.name)
        stack.append(inline_args[0])
        return None


add_op = StkOp('add', 2)
sub_op = StkOp('sub', 2)


class PushframeOp(Opcode):
    """Opcode to push an address relative to the frame pointer.
    """

    def __init__(self) -> None:
        super().__init__('pushframe', 0, 1)

    def build(self, stack: Deque, inline_args: Optional[list] = None):
        if len(inline_args) != self.inline_args:
            raise ValueError("Not enough inline args")
        stack.append(FramePointer())
        stack.append(inline_args[0])
        return add_op.build(stack)


class JmpOp(Opcode):
    """Opcode for a jump instruction.
    """

    def __init__(self) -> None:
        super().__init__('jmp', 0, 1)

    def build(self, stack: Deque, inline_args: Optional[list] = None) -> Jump:
        if len(inline_args) != self.stk_args:
            raise ValueError("Not enough inline arguments.")
        return Jump(Static(inline_args[0]))


class MoveOp(Opcode):
    """Opcode for storing in a target location.
    """

    def __init__(self, code: str, dest_type: ValType) -> None:
        super().__init__(code, 2, 0)
        self.dest_type = dest_type

    def build(self, stack: Deque, inline_args: Optional[list] = None) -> Move:
        if len(stack) < self.stk_args:
            raise ValueError("Not enough stack arguments")
        dest = stack.pop()
        expr = stack.pop()
        return Move(dest, expr, self.dest_type)


class BrzOp(Opcode):
    """Opcode for branching conditionally.
    """

    def __init__(self) -> None:
        super().__init__('brz', 1, 1)

    def build(self, stack: Deque, inline_args: Optional[list] = None) -> Brz:
        if len(stack) < self.stk_args:
            raise ValueError("Not enough stack args")
        if len(inline_args) != self.inline_args:
            raise ValueError("Not enough inline args")
        return Brz(stack.popleft(), Static(inline_args[0]))


class StkJmp(Opcode):
    """Opcode for jumping to a stack evaluating location.
    """

    def __init__(self) -> None:
        super().__init__('stkjmp', 1, 0)

    def build(self, stack: Deque, inline_args: Optional[list] = None) -> Jump:
        if len(stack) < self.stk_args:
            raise ValueError("Not enough stack args")
        return Jump(stack.pop())


class ClearOp(Opcode):
    """Opcode for a stack clear instruction.
    """

    def __init__(self) -> None:
        super().__init__('clear', 0, 0)

    def build(self, stack: Deque, inline_args: Optional[list] = None):
        stack.clear()
        return None


def for_effect(node: Union[Leaf, Node]) -> list[Union[Leaf, Node]]:
    """Recursively computes those Nodes and Leafs which need to be computed
    for effect only.
    """
    if isinstance(node, Node):
        if node.opcode.flags.has_effect:
            return [node]
        children = []
        for child in node.children:
            children += for_effect(child)
        return children
    return []


RE_ALPHA = r"([a-zA-Z_.])"
RE_NUM = r"(\d)"
RE_ALPHANUM = r"(("+RE_ALPHA+r")|("+RE_NUM+r"))"
RE_LABEL = r"(?P<label>"+RE_ALPHA+RE_ALPHANUM+r"*)"
RE_NAME = r"(?P<name>"+RE_ALPHA+RE_ALPHANUM+r"*)"
RE_CON = r"(?P<con>"+RE_NUM+r")"


class Grabber:
    """Allows text to be grabbed.
    """

    def __init__(self, text: str):
        self.text = text
        self.__index = 0

    def __iter__(self) -> Iterator[str]:
        return self

    def bounds(self) -> None:
        """Makes sure our index >= 0.
        """
        self.__index = max(0, self.__index)

    def grab(self, count: int = 1) -> str:
        """Return the next count characters from me.
        """
        old_index = self.__index
        self.__index += count
        return self.text[old_index:old_index+count]

    def __next__(self) -> str:
        return self.grab(1)

    def match(self, pattern: str) -> Union[str, None]:
        """Try to match against a given regular expression.
        """
        match = re.match(pattern, self.text[self.__index:])
        if match is None:
            return None
        self.seek(len(match[0]))
        return match[0]

    def match_any(self, *args: Tuple[str, str]) -> Union[Tuple[str, str], None]:
        """Try to match, in order, the given arguments. Returns the one
        matched, or None of none matched. Ends on the first match.
        """
        for name, pattern in args:
            match = self.match(pattern)
            if match is not None:
                return name, match
        return None

    def seek(self, offset: int, relative: bool = True) -> None:
        """Seek to a given offset, either relative from thecurrent position
        or from the start.
        """
        if relative:
            self.__index += offset
        else:
            self.__index = offset
        self.bounds()

    def __len__(self):
        return len(self.text[self.__index:])

    def peek(self, count: int = 1):
        """Return the next count elements w/o modifying my index.
        """
        return self.text[self.__index:self.__index+count]


class LexChar:
    """A lexical literal character.
    """

    def __init__(self, char: str) -> None:
        self.char = char


class LexNL:
    """A lexical newline.
    """


def lex(source_text: str) -> Iterable:
    """Lex over a source assembly.
    """
    src = Grabber(source_text)
    while len(src) > 0:
        if src.peek() == ';':
            src.grab()
            while len(src) > 0 and src.peek() != '\n':
                src.grab()
        elif src.peek() == '\n':
            src.grab()
            yield LexNL()
        elif src.peek() in whitespace:
            src.grab()
        else:
            match = src.match_any(
                'label', RE_LABEL, 'name', RE_NAME, 'con', RE_CON)
            if match is not None:
                match_type, match_str = match
                if match_type == 'label':
                    yield Label(match_str.rstrip(':'))
                elif match_type == 'name':
                    yield LexName(match_str)
                elif match_type == 'con':
                    yield Con(int(match_str, base=8))
            else:
                yield LexChar(src.grab())


def build(source: str, opcodes: Dict[str, Opcode]) -> SeqState:
    stack = deque()
    statements = []
    tokens = list(lex(source))
    while len(tokens) > 0:
        token = tokens.pop(0)
        if isinstance(token, Label):
            statements.append(token)
        elif isinstance(token, LexName):
            if token.name in opcodes:
                opcode = opcodes[token.name]
                inline_args = []
                while len(tokens) > 0 and not isinstance(tokens[0], LexNL):
                    argtkn = tokens.pop(0)
                    if isinstance(argtkn, LexName):
                        inline_args.append(Static(argtkn.name))
                    elif isinstance(argtkn, Con):
                        inline_args.append(argtkn)
                    elif isinstance(argtkn, LexChar):
                        if argtkn.char in ('+', '-'):
                            right = inline_args.pop(0)
                            left = inline_args.pop(0)
                            if argtkn.char == '+':
                                inline_args.append(Node(add_op, left, right))
                            elif argtkn.char == '-':
                                inline_args.append(Node(sub_op, left, right))
                        elif argtkn.char == ',':
                            pass
                        else:
                            raise ValueError(
                                f"Unexpected token {repr(argtkn)}")
                statement = opcode.build(stack, inline_args)
                if isinstance(statement, Statement):
                    while len(stack) > 0:
                        statements.append(Eval(stack.pop()))
                    statements.append(statement)
        else:
            raise ValueError(f"Unexpected token {repr(token)}")

    return SeqState(*statements)


def build_opcodes(opcodes: Iterable[Opcode]) -> Dict[str, Opcode]:
    """Build a dictionary of opcodes.
    """
    opcode_dict = {}
    for opcode in opcodes:
        opcode_dict[opcode.code] = opcode
    return opcode_dict


class DupOp(Opcode):
    def __init__(self) -> None:
        super().__init__('dup', 1)

    def build(self, stack: Deque, inline_args: Optional[list] = None):
        if len(stack) < self.stk_args:
            raise ValueError("Not enough stack args")
        val = stack.pop()
        stack.append(val)
        stack.append(val)


class CallOp(StkOp):
    def __init__(self, **flags) -> None:
        super().__init__('call', 0, 1, has_effect=True, **flags)

    def build(self, stack: deque, inline_args: Optional[list] = None) -> None:
        self.stk_args = inline_args.pop().value
        super().build(stack, inline_args)
        self.stk_args = 0
        return None


class GrabRegOp(Opcode):
    def __init__(self, **flags) -> None:
        super().__init__('grabreg', 0, 1, **flags)

    def build(self, stack: Deque, inline_args: Optional[list] = None):
        return Register(inline_args[0].value)


class PutRegOp(MoveOp):
    def __init__(self) -> None:
        super().__init__('putreg', IntType)

    def build(self, stack: Deque, inline_args: Optional[list] = None) -> Move:
        val = stack.pop()
        dest = Register(inline_args[0].value)
        return Move(dest, val, IntType)


opcode_list = [add_op, sub_op, JmpOp(), PushOp(), PushframeOp(), BrzOp(),
               StkJmp(), ClearOp(), StkOp('ret', 1, has_effect=True), DupOp(),
               StkOp('lequ', 2, forcond=True), StkOp('lshift', 2),
               StkOp('deref', 1), StkOp('equ', 2, forcond=True),
               StkOp('savexp'), StkOp('call', 0, 1, has_effect=True),
               StkOp('or', 2), StkOp('eor', 2), StkOp('and', 2),
               StkOp('nequ', 2, forcond=True), StkOp('gequ', 2, forcond=True),
               StkOp('less', 2, forcond=True),
               StkOp('great', 2, forcond=True),
               StkOp('ugequ', 2, forcond=True),
               StkOp('ulequ', 2, forcond=True),
               StkOp('ugreat', 2, forcond=True),
               StkOp('rshift', 2),
               CallOp(), StkOp('mod', 2), StkOp('div', 2), StkOp('mult', 2),
               StkOp('neg', 1), StkOp('compl', 1),
               StkOp('lognot', 1, forcond=True),
               StkOp('log', 1, forcond=True),
               GrabRegOp(), StkOp('cderef', 2),
               PutRegOp(), MoveOp('store', IntType),
               MoveOp('cstore', CharType), StkOp('drop', 2)]


def codegen(backend: Callable[[Statement], str], seq: SeqState) -> str:
    out = ''
    for statement in seq:
        out += backend(statement)
    return out


def word(i: int) -> int:
    """Convert an integer into an unsigned 16bit value.
    """
    i &= 0xFFFF
    if i < 0:
        i = ~i+1
    return i
