#!/usr/bin/env python3
# m6502.py - minimal 6502 opcode table, instruction decoder and recursive-descent code tracer.
from dataclasses import dataclass

# addressing modes
IMP, ACC, IMM, ZP, ZPX, ZPY, ABS, ABX, ABY, IND, IZX, IZY, REL = range(13)
MODE_LEN = {IMP:1, ACC:1, IMM:2, ZP:2, ZPX:2, ZPY:2, ABS:3, ABX:3, ABY:3, IND:3, IZX:2, IZY:2, REL:2}

OPCODES = {}
def _def(mnemonic, table):
    for opcode, mode in table.items():
        OPCODES[opcode] = (mnemonic, mode)

_def("adc", {0x69:IMM,0x65:ZP,0x75:ZPX,0x6D:ABS,0x7D:ABX,0x79:ABY,0x61:IZX,0x71:IZY})
_def("and", {0x29:IMM,0x25:ZP,0x35:ZPX,0x2D:ABS,0x3D:ABX,0x39:ABY,0x21:IZX,0x31:IZY})
_def("asl", {0x0A:ACC,0x06:ZP,0x16:ZPX,0x0E:ABS,0x1E:ABX})
_def("bcc", {0x90:REL}); _def("bcs", {0xB0:REL}); _def("beq", {0xF0:REL}); _def("bmi", {0x30:REL})
_def("bne", {0xD0:REL}); _def("bpl", {0x10:REL}); _def("bvc", {0x50:REL}); _def("bvs", {0x70:REL})
_def("bit", {0x24:ZP,0x2C:ABS})
_def("brk", {0x00:IMP})
_def("clc", {0x18:IMP}); _def("cld", {0xD8:IMP}); _def("cli", {0x58:IMP}); _def("clv", {0xB8:IMP})
_def("cmp", {0xC9:IMM,0xC5:ZP,0xD5:ZPX,0xCD:ABS,0xDD:ABX,0xD9:ABY,0xC1:IZX,0xD1:IZY})
_def("cpx", {0xE0:IMM,0xE4:ZP,0xEC:ABS}); _def("cpy", {0xC0:IMM,0xC4:ZP,0xCC:ABS})
_def("dec", {0xC6:ZP,0xD6:ZPX,0xCE:ABS,0xDE:ABX}); _def("dex", {0xCA:IMP}); _def("dey", {0x88:IMP})
_def("eor", {0x49:IMM,0x45:ZP,0x55:ZPX,0x4D:ABS,0x5D:ABX,0x59:ABY,0x41:IZX,0x51:IZY})
_def("inc", {0xE6:ZP,0xF6:ZPX,0xEE:ABS,0xFE:ABX}); _def("inx", {0xE8:IMP}); _def("iny", {0xC8:IMP})
_def("jmp", {0x4C:ABS,0x6C:IND}); _def("jsr", {0x20:ABS})
_def("lda", {0xA9:IMM,0xA5:ZP,0xB5:ZPX,0xAD:ABS,0xBD:ABX,0xB9:ABY,0xA1:IZX,0xB1:IZY})
_def("ldx", {0xA2:IMM,0xA6:ZP,0xB6:ZPY,0xAE:ABS,0xBE:ABY})
_def("ldy", {0xA0:IMM,0xA4:ZP,0xB4:ZPX,0xAC:ABS,0xBC:ABX})
_def("lsr", {0x4A:ACC,0x46:ZP,0x56:ZPX,0x4E:ABS,0x5E:ABX})
_def("nop", {0xEA:IMP})
_def("ora", {0x09:IMM,0x05:ZP,0x15:ZPX,0x0D:ABS,0x1D:ABX,0x19:ABY,0x01:IZX,0x11:IZY})
_def("pha", {0x48:IMP}); _def("php", {0x08:IMP}); _def("pla", {0x68:IMP}); _def("plp", {0x28:IMP})
_def("rol", {0x2A:ACC,0x26:ZP,0x36:ZPX,0x2E:ABS,0x3E:ABX})
_def("ror", {0x6A:ACC,0x66:ZP,0x76:ZPX,0x6E:ABS,0x7E:ABX})
_def("rti", {0x40:IMP}); _def("rts", {0x60:IMP})
_def("sbc", {0xE9:IMM,0xE5:ZP,0xF5:ZPX,0xED:ABS,0xFD:ABX,0xF9:ABY,0xE1:IZX,0xF1:IZY})
_def("sec", {0x38:IMP}); _def("sed", {0xF8:IMP}); _def("sei", {0x78:IMP})
_def("sta", {0x85:ZP,0x95:ZPX,0x8D:ABS,0x9D:ABX,0x99:ABY,0x81:IZX,0x91:IZY})
_def("stx", {0x86:ZP,0x96:ZPY,0x8E:ABS}); _def("sty", {0x84:ZP,0x94:ZPX,0x8C:ABS})
_def("tax", {0xAA:IMP}); _def("tay", {0xA8:IMP}); _def("tsx", {0xBA:IMP})
_def("txa", {0x8A:IMP}); _def("txs", {0x9A:IMP}); _def("tya", {0x98:IMP})

BRANCHES = {"bcc","bcs","beq","bmi","bne","bpl","bvc","bvs"}
FLOW_END = {"rts","rti","brk"}


@dataclass
class InsnT:
    addr: int
    opcode: int
    mnemonic: str
    mode: int
    length: int
    operand: int          # raw operand value (or 0)
    target: int           # effective address for ABS/ZP/REL/etc (None for IMP/ACC/IMM)

    def isLegal(self):
        return self.mnemonic != "???"


def decode(mem, addr):
    opcode = mem[addr]
    if opcode not in OPCODES:
        return InsnT(addr, opcode, "???", IMP, 1, 0, None)
    mnemonic, mode = OPCODES[opcode]
    length = MODE_LEN[mode]
    operand = 0
    target = None
    if length == 2:
        operand = mem[(addr+1) & 0xFFFF]
    elif length == 3:
        operand = mem[(addr+1) & 0xFFFF] | (mem[(addr+2) & 0xFFFF] << 8)
    if mode == REL:
        target = (addr + 2 + ((operand ^ 0x80) - 0x80)) & 0xFFFF
    elif mode in (ZP, ZPX, ZPY, ABS, ABX, ABY, IND, IZX, IZY):
        target = operand
    return InsnT(addr, opcode, mnemonic, mode, length, operand, target)


def formatOperand(insn, labelFor=None):
    m = insn.mode
    def lab(a, width):
        if labelFor:
            l = labelFor(a)
            if l:
                return l
        return f"${a:0{width}X}"
    if m in (IMP,):
        return ""
    if m == ACC:
        return "a"
    if m == IMM:
        return f"#${insn.operand:02X}"
    if m == ZP:
        return lab(insn.operand, 2)
    if m == ZPX:
        return lab(insn.operand, 2) + ",x"
    if m == ZPY:
        return lab(insn.operand, 2) + ",y"
    if m == ABS:
        return lab(insn.operand, 4)
    if m == ABX:
        return lab(insn.operand, 4) + ",x"
    if m == ABY:
        return lab(insn.operand, 4) + ",y"
    if m == IND:
        return "(" + lab(insn.operand, 4) + ")"
    if m == IZX:
        return "(" + lab(insn.operand, 2) + ",x)"
    if m == IZY:
        return "(" + lab(insn.operand, 2) + "),y"
    if m == REL:
        return lab(insn.target, 4)
    return "?"


def trace(mem, entries, isValid=lambda a: True, noReturn=frozenset(), inlineArgs=None, known=frozenset()):
    """Recursive-descent: returns (codeSet, callTargets, jumpTargets, indirectJumps, illegalHits).
    inlineArgs: dict subroutineAddr -> number of inline bytes that follow a JSR to it."""
    code = set()
    calls = {}
    jumps = {}
    indirects = []
    illegal = []
    work = list(entries)
    seen = set(entries)
    while work:
        pc = work.pop()
        while True:
            if pc in code or pc in known or not isValid(pc):
                break
            insn = decode(mem, pc)
            if not insn.isLegal():
                illegal.append(pc)
                break
            for i in range(insn.length):
                code.add((pc + i) & 0xFFFF)
            nxt = (pc + insn.length) & 0xFFFF
            if insn.mnemonic in BRANCHES:
                if insn.target not in seen:
                    seen.add(insn.target); work.append(insn.target)
                pc = nxt
            elif insn.mnemonic == "jsr":
                calls.setdefault(insn.target, []).append(pc)
                if insn.target not in seen:
                    seen.add(insn.target); work.append(insn.target)
                if inlineArgs and insn.target in inlineArgs:
                    nxt = (nxt + inlineArgs[insn.target]) & 0xFFFF
                if insn.target in noReturn:
                    break
                pc = nxt
            elif insn.mnemonic == "jmp":
                if insn.mode == IND:
                    indirects.append(pc)
                    break
                jumps.setdefault(insn.target, []).append(pc)
                if insn.target not in seen:
                    seen.add(insn.target); work.append(insn.target)
                break
            elif insn.mnemonic in FLOW_END:
                break
            else:
                pc = nxt
    return code, calls, jumps, indirects, illegal