#!/usr/bin/env python3
# discover.py - find additional code entry points that plain recursive descent cannot see:
#   * interrupt / indirect-jump vectors written with LDA #lo / STA vec, LDA #hi / STA vec+1
#   * address tables copied into the operand of a JSR/JMP (self-modifying dispatch)
#   * address tables inside data runs (3+ consecutive words that point at plausible code)
# Every candidate is only accepted by the caller after a clean transactional trace.
from m6502 import decode, IMM, ABS, ABX, ABY, IND


def instructionStarts(mem, codeSet):
    starts = {}
    nextStart = 0
    for a in sorted(codeSet):
        if a < nextStart:
            continue
        insn = decode(mem, a)
        if insn.isLegal():
            starts[a] = insn
            nextStart = a + insn.length
    return starts


def vectorEntries(mem, starts, vectorAddrs):
    """values stored to vector lo/hi bytes via immediate loads -> set of candidate addresses"""
    lo = {v: set() for v in vectorAddrs}
    hi = {v: set() for v in vectorAddrs}
    order = sorted(starts)
    for i, a in enumerate(order):
        insn = starts[a]
        if insn.mnemonic not in ("sta", "stx", "sty") or insn.mode != ABS:
            continue
        t = insn.target
        reg = {"sta": "lda", "stx": "ldx", "sty": "ldy"}[insn.mnemonic]
        # look back up to 4 instructions for the immediate load of the same register
        val = None
        for j in range(i - 1, max(-1, i - 5), -1):
            p = starts[order[j]]
            if p.mnemonic == reg and p.mode == IMM:
                val = p.operand
                break
            if p.mnemonic in ("jsr", "jmp", "rts", "rti") or (p.mnemonic == reg):
                break
        if val is None:
            continue
        if t in lo:
            lo[t].add(val)
        if (t - 1) in hi:
            hi[t - 1].add(val)
    out = set()
    for v in vectorAddrs:
        for l in lo[v]:
            for h in hi[v]:
                out.add(l | (h << 8))
    return out


def patchedCallTables(mem, starts):
    """Find tables whose entries are copied into the operand of a JSR/JMP.

    Two shapes exist:
      word table   lda tbl,x / sta jsr+1 ; lda tbl+1,x / sta jsr+2   (x = index*2)
      split table  lda lo,x  / sta jsr+1 ; lda hi,x   / sta jsr+2    (x = index)
    Returns a list of (kind, baseLo, baseHi) with kind in {"word", "split"}.
    """
    order = sorted(starts)
    pos = {a: i for i, a in enumerate(order)}
    found = []
    for a, insn in starts.items():
        if insn.mnemonic != "sta" or insn.mode != ABS:
            continue
        target = insn.target                      # operand byte being patched
        for back in (1, 2):
            base = target - back
            if base not in starts:
                continue
            callInsn = starts[base]
            if callInsn.mnemonic not in ("jsr", "jmp") or callInsn.mode != ABS:
                continue
            if back != 1:
                continue                          # anchor the search on the low-byte store
            loSrc = _sourceTable(order, pos, starts, a)
            hiStore = _findStoreTo(order, pos, starts, a, target + 1)
            if loSrc is None or hiStore is None:
                continue
            hiSrc = _sourceTable(order, pos, starts, hiStore)
            if hiSrc is None:
                continue
            found.append(("word" if hiSrc == loSrc + 1 else "split", loSrc, hiSrc))
    return found


def _sourceTable(order, pos, starts, storeAddr, lookBack=4):
    """base address of the lda tbl,x / tbl,y that feeds the store at storeAddr"""
    i = pos[storeAddr]
    for j in range(i - 1, max(-1, i - 1 - lookBack), -1):
        p = starts[order[j]]
        if p.mnemonic == "lda" and p.mode in (ABX, ABY):
            return p.target
        if p.mnemonic in ("jsr", "jmp", "rts", "rti", "lda"):
            return None
    return None


def _findStoreTo(order, pos, starts, fromAddr, wantTarget, lookAhead=6):
    """address of the next sta wantTarget after fromAddr"""
    i = pos[fromAddr]
    for j in range(i + 1, min(len(order), i + 1 + lookAhead)):
        p = starts[order[j]]
        if p.mnemonic == "sta" and p.mode == ABS and p.target == wantTarget:
            return p.addr
    return None


def splitTableEntries(mem, baseLo, baseHi, isPlausible, maxEntries=64):
    """entries of a split low/high byte table, stopping at the first implausible pair"""
    out = []
    for i in range(maxEntries):
        w = mem[baseLo + i] | (mem[baseHi + i] << 8)
        if not isPlausible(w):
            break
        out.append(w)
    return out


def wordsFrom(mem, base, isPlausible, maxEntries=64):
    """read consecutive words from base while they look like code addresses"""
    out = []
    a = base
    gaps = 0
    for _ in range(maxEntries):
        w = mem[a] | (mem[a + 1] << 8)
        if isPlausible(w):
            out.append(w)
            gaps = 0
        elif w in (0xFFFF, 0x0000) and gaps == 0 and out:
            gaps += 1                      # tolerate a single placeholder entry
        else:
            break
        a += 2
    return out


def addressTablesInData(mem, codeSet, ranges, isPlausible, minRun=3):
    """scan data bytes inside the given ranges for runs of plausible code addresses"""
    tables = []
    for lo, hi in ranges:
        a = lo
        while a + 1 < hi:
            if a in codeSet:
                a += 1
                continue
            run = []
            b = a
            while b + 1 < hi and b not in codeSet and (b + 1) not in codeSet:
                w = mem[b] | (mem[b + 1] << 8)
                if isPlausible(w):
                    run.append(w)
                    b += 2
                else:
                    break
            if len(run) >= minRun:
                tables.append((a, run))
                a = b
            else:
                a += 1
    return tables