prog8/il65/emit/assignment.py

"""
Programming Language for 6502/6510 microprocessors, codename 'Sick'
This is the code generator for assignment statements.

Written by Irmen de Jong (irmen@razorvine.net) - license: GNU GPL 3.0
"""

from typing import Callable
from ..plyparse import Scope, Assignment, AugAssignment, Register, LiteralValue, SymbolName, VarDef, Dereference
from . import CodeError, preserving_registers, to_hex, Context
from ..datatypes import REGISTER_BYTES, VarType, DataType
from ..compile import Zeropage


def generate_assignment(ctx: Context) -> None:
    assert isinstance(ctx.stmt, Assignment)
    ctx.out("\v\t\t\t; " + ctx.stmt.lineref)
    ctx.out("\v; @todo assignment")
    # @todo assignment


def generate_aug_assignment(ctx: Context) -> None:
    # for instance: value += 3  (value = 0-255 for now)
    # left: Register, SymbolName, or Dereference. right: Expression/LiteralValue
    out = ctx.out
    stmt = ctx.stmt
    assert isinstance(stmt, AugAssignment)
    out("\v\t\t\t; " + stmt.lineref)
    lvalue = stmt.left
    rvalue = stmt.right
    if isinstance(lvalue, Register):
        if isinstance(rvalue, LiteralValue):
            if type(rvalue.value) is int:
                if 0 <= rvalue.value <= 255:
                    if stmt.operator not in ("<<=", ">>=") or rvalue.value != 0:
                        _generate_aug_reg_int(out, lvalue, stmt.operator, rvalue.value, "", ctx.scope)
                else:
                    raise CodeError("aug. assignment value must be 0..255", rvalue)
            else:
                raise CodeError("constant integer literal or variable required for now", rvalue)   # XXX
        elif isinstance(rvalue, SymbolName):
            symdef = ctx.scope.lookup(rvalue.name)
            if isinstance(symdef, VarDef):
                if symdef.vartype == VarType.CONST:
                    if symdef.datatype.isinteger() and 0 <= symdef.value.const_value() <= 255:  # type: ignore
                        _generate_aug_reg_int(out, lvalue, stmt.operator, symdef.value.const_value(), "", ctx.scope)   # type: ignore
                    else:
                        raise CodeError("aug. assignment value must be integer 0..255", rvalue)
                elif symdef.datatype == DataType.BYTE:
                    _generate_aug_reg_int(out, lvalue, stmt.operator, 0, symdef.name, ctx.scope)
                else:
                    raise CodeError("variable must be of type byte for now", rvalue)   # XXX
            else:
                raise CodeError("can only use variable name as symbol for aug assign rvalue", rvalue)
        elif isinstance(rvalue, Register):
            if lvalue.datatype == DataType.BYTE and rvalue.datatype == DataType.WORD:
                raise CodeError("cannot assign a combined 16-bit register to a single 8-bit register", rvalue)
            _generate_aug_reg_reg(out, lvalue, stmt.operator, rvalue, ctx.scope)
        elif isinstance(rvalue, Dereference):
            if isinstance(rvalue.operand, (LiteralValue, SymbolName)):
                if isinstance(rvalue.operand, LiteralValue):
                    what = to_hex(rvalue.operand.value)
                else:
                    symdef = rvalue.my_scope().lookup(rvalue.operand.name)
                    if isinstance(symdef, VarDef) and symdef.vartype == VarType.MEMORY:
                        what = to_hex(symdef.value.value)  # type: ignore
                    else:
                        what = rvalue.operand.name
                out("\vpha\n\vtya\n\vpha")   # save A, Y on stack
                out("\vlda  " + what)
                out("\vsta  il65_lib.SCRATCH_ZPWORD1")
                out("\vlda  {:s}+1".format(what))
                out("\vsta  il65_lib.SCRATCH_ZPWORD1+1")
                out("\vldy  #0")
                out("\vlda  (il65_lib.SCRATCH_ZPWORD1), y")
                a_reg = Register(name="A", sourceref=stmt.sourceref)        # type: ignore
                _generate_aug_reg_reg(out, lvalue, stmt.operator, a_reg, ctx.scope)
                out("\vst{:s}  il65_lib.SCRATCH_ZP1".format(lvalue.name.lower()))
                out("\vpla\n\vtay\n\vpla")  # restore A, Y from stack
                out("\vld{:s}  il65_lib.SCRATCH_ZP1".format(lvalue.name.lower()))
            elif isinstance(rvalue.operand, Register):
                assert rvalue.operand.datatype == DataType.WORD
                if rvalue.datatype != DataType.BYTE:
                    raise CodeError("aug. assignment value must be a byte, 0..255", rvalue)
                reg = rvalue.operand.name
                out("\vst{:s}  il65_lib.SCRATCH_ZPWORD1".format(reg[0].lower()))
                out("\vst{:s}  il65_lib.SCRATCH_ZPWORD1+1".format(reg[1].lower()))
                out("\vpha\n\vtya\n\vpha")   # save A, Y on stack
                out("\vldy  #0")
                out("\vlda  (il65_lib.SCRATCH_ZPWORD1), y")
                a_reg = Register(name="A", sourceref=stmt.sourceref)        # type: ignore
                _generate_aug_reg_reg(out, lvalue, stmt.operator, a_reg, ctx.scope)
                if lvalue.name != 'X':
                    out("\vst{:s}  il65_lib.SCRATCH_ZP1".format(lvalue.name.lower()))
                out("\vpla\n\vtay\n\vpla")  # restore A, Y from stack
                if lvalue.name != 'X':
                    out("\vld{:s}  il65_lib.SCRATCH_ZP1".format(lvalue.name.lower()))
            else:
                raise CodeError("invalid dereference operand type", rvalue)
        else:
            raise CodeError("invalid rvalue type", rvalue)
    else:
        raise CodeError("aug. assignment only implemented for registers for now", stmt.sourceref)  # XXX


def _generate_aug_reg_int(out: Callable, lvalue: Register, operator: str, rvalue: int, rname: str, scope: Scope) -> None:
    if rname:
        right_str = rname
    else:
        # an immediate value is provided in rvalue
        right_str = "#" + to_hex(rvalue)
    if operator == "+=":
        if lvalue.name == "A":
            out("\vclc")
            out("\vadc  " + right_str)
        elif lvalue.name == "X":
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vclc")
                out("\vadc  " + right_str)
                out("\vtax")
        elif lvalue.name == "Y":
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vclc")
                out("\vadc  " + right_str)
                out("\vtay")
        else:
            raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo +=.word
    elif operator == "-=":
        if lvalue.name == "A":
            out("\vsec")
            out("\vsbc  " + right_str)
        elif lvalue.name == "X":
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vsec")
                out("\vsbc  " + right_str)
                out("\vtax")
        elif lvalue.name == "Y":
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vsec")
                out("\vsbc  " + right_str)
                out("\vtay")
        else:
            raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo -=.word
    elif operator == "&=":
        if lvalue.name == "A":
            out("\vand  " + right_str)
        elif lvalue.name == "X":
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vand  " + right_str)
                out("\vtax")
        elif lvalue.name == "Y":
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vand  " + right_str)
                out("\vtay")
        else:
            raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo &=.word
    elif operator == "|=":
        if lvalue.name == "A":
            out("\vora  " + right_str)
        elif lvalue.name == "X":
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vora  " + right_str)
                out("\vtax")
        elif lvalue.name == "Y":
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vora  " + right_str)
                out("\vtay")
        else:
            raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo |=.word
    elif operator == "^=":
        if lvalue.name == "A":
            out("\veor  " + right_str)
        elif lvalue.name == "X":
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\veor  " + right_str)
                out("\vtax")
        elif lvalue.name == "Y":
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\veor  " + right_str)
                out("\vtay")
        else:
            raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo ^=.word
    elif operator == ">>=":
        if rname:
            assert lvalue.name in REGISTER_BYTES, "only single registers for now"  # @todo >>=.word
            if lvalue.name == "A":
                preserve_regs = {'X'}
            elif lvalue.name == "X":
                preserve_regs = {'A'}
                out("\vtxa")
            elif lvalue.name == "Y":
                preserve_regs = {'A', 'X'}
                out("\vtya")
            with preserving_registers(preserve_regs, scope, out):
                out("\vldx  " + rname)
                out("\vjsr  il65_lib.lsr_A_by_X")
                # out("\vbeq  +")
                # out("-\vlsr  a")
                # out("\vdex")
                # out("\vbne  -")
                # put A back into target register
                if lvalue.name == "X":
                    out("\vtax")
                if lvalue.name == "Y":
                    out("\vtay")
        else:
            def shifts_A(times: int) -> None:
                if times >= 8:
                    out("\vlda  #0")
                else:
                    for _ in range(min(8, times)):
                        out("\vlsr  a")
            if lvalue.name == "A":
                shifts_A(rvalue)
            elif lvalue.name == "X":
                with preserving_registers({'A'}, scope, out):
                    out("\vtxa")
                    shifts_A(rvalue)
                    out("\vtax")
            elif lvalue.name == "Y":
                with preserving_registers({'A'}, scope, out):
                    out("\vtya")
                    shifts_A(rvalue)
                    out("\vtay")
            else:
                raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo >>=.word
    elif operator == "<<=":
        if rname:
            assert lvalue.name in REGISTER_BYTES, "only single registers for now"   # @todo <<=.word
            if lvalue.name == "A":
                preserve_regs = {'X'}
            elif lvalue.name == "X":
                preserve_regs = {'A'}
                out("\vtxa")
            elif lvalue.name == "Y":
                preserve_regs = {'A', 'X'}
                out("\vtya")
            with preserving_registers(preserve_regs, scope, out):
                out("\vldx  " + rname)
                out("\vjsr  il65_lib.asl_A_by_X")
                # out("\vbeq  +")
                # out("-\vasl  a")
                # out("\vdex")
                # out("\vbne  -")
                # put A back into target register
                if lvalue.name == "X":
                    out("\vtax")
                elif lvalue.name == "Y":
                    out("\vtay")
        else:
            def shifts_A(times: int) -> None:
                if times >= 8:
                    out("\vlda  #0")
                else:
                    for _ in range(min(8, times)):
                        out("\vasl  a")
            if lvalue.name == "A":
                shifts_A(rvalue)
            elif lvalue.name == "X":
                with preserving_registers({'A'}, scope, out):
                    out("\vtxa")
                    shifts_A(rvalue)
                    out("\vtax")
            elif lvalue.name == "Y":
                with preserving_registers({'A'}, scope, out):
                    out("\vtya")
                    shifts_A(rvalue)
                    out("\vtay")
            else:
                raise CodeError("unsupported register for aug assign", str(lvalue))  # @todo <<=.word


def _generate_aug_reg_reg(out: Callable, lvalue: Register, operator: str, rvalue: Register, scope: Scope) -> None:
    if operator == "+=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo +=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\vclc")
            out("\vadc  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vclc")
                out("\vadc  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vclc")
                out("\vadc  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo +=.word
    elif operator == "-=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo -=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\vsec")
            out("\vsbc  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vsec")
                out("\vsbc  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vsec")
                out("\vsbc  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo -=.word
    elif operator == "&=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo &=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\vand  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vand  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vand  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo &=.word
    elif operator == "|=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo |=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\vora  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vora  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vora  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo |=.word
    elif operator == "^=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo ^=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\veor  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\veor  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\veor  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo ^=.word
    elif operator == ">>=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo >>=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\vlsr  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vlsr  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vlsr  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo >>=.word
    elif operator == "<<=":
        if rvalue.name not in REGISTER_BYTES:
            raise CodeError("unsupported rvalue register for aug assign", str(rvalue))  # @todo <<=.word
        if lvalue.name == "A":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            out("\vasl  " + to_hex(Zeropage.SCRATCH_B1))
        elif lvalue.name == "X":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtxa")
                out("\vasl  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtax")
        elif lvalue.name == "Y":
            out("\vst{:s}  {:s}".format(rvalue.name.lower(), to_hex(Zeropage.SCRATCH_B1)))
            with preserving_registers({'A'}, scope, out):
                out("\vtya")
                out("\vasl  " + to_hex(Zeropage.SCRATCH_B1))
                out("\vtay")
        else:
            raise CodeError("unsupported lvalue register for aug assign", str(lvalue))  # @todo <<=.word