C64-dev/reference/powerc-re/tools/pcopt.py

#!/usr/bin/env python3
"""pcopt - peephole optimizer for Power C object files (.o).

Reads an object produced by the original Power C `cc`, decodes every compiled
function into 6502 instructions, applies rewrites that are checked against a
register/flag liveness analysis, and writes a smaller/faster object with all
offsets, relocations, labels and branches fixed up. Anything it cannot decode
with certainty is left untouched. The original compiler, linker and libraries
are never modified.

usage: pcopt.py in.o out.o [--lib libdir] [--no-rules] [-v]
"""
import sys, os, glob
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from objfmt import Obj, pn

# ------------------------------------------------------------------ 6502 table
_T = """
00 brk imp;01 ora izx;05 ora zp;06 asl zp;08 php imp;09 ora imm;0a asl acc;0d ora abs;0e asl abs;
10 bpl rel;11 ora izy;15 ora zpx;16 asl zpx;18 clc imp;19 ora aby;1d ora abx;1e asl abx;
20 jsr abs;21 and izx;24 bit zp;25 and zp;26 rol zp;28 plp imp;29 and imm;2a rol acc;2c bit abs;2d and abs;2e rol abs;
30 bmi rel;31 and izy;35 and zpx;36 rol zpx;38 sec imp;39 and aby;3d and abx;3e rol abx;
40 rti imp;41 eor izx;45 eor zp;46 lsr zp;48 pha imp;49 eor imm;4a lsr acc;4c jmp abs;4d eor abs;4e lsr abs;
50 bvc rel;51 eor izy;55 eor zpx;56 lsr zpx;58 cli imp;59 eor aby;5d eor abx;5e lsr abx;
60 rts imp;61 adc izx;65 adc zp;66 ror zp;68 pla imp;69 adc imm;6a ror acc;6c jmp ind;6d adc abs;6e ror abs;
70 bvs rel;71 adc izy;75 adc zpx;76 ror zpx;78 sei imp;79 adc aby;7d adc abx;7e ror abx;
81 sta izx;84 sty zp;85 sta zp;86 stx zp;88 dey imp;8a txa imp;8c sty abs;8d sta abs;8e stx abs;
90 bcc rel;91 sta izy;94 sty zpx;95 sta zpx;96 stx zpy;98 tya imp;99 sta aby;9a txs imp;9d sta abx;
a0 ldy imm;a1 lda izx;a2 ldx imm;a4 ldy zp;a5 lda zp;a6 ldx zp;a8 tay imp;a9 lda imm;aa tax imp;ac ldy abs;ad lda abs;ae ldx abs;
b0 bcs rel;b1 lda izy;b4 ldy zpx;b5 lda zpx;b6 ldx zpy;b8 clv imp;b9 lda aby;ba tsx imp;bc ldy abx;bd lda abx;be ldx aby;
c0 cpy imm;c1 cmp izx;c4 cpy zp;c5 cmp zp;c6 dec zp;c8 iny imp;c9 cmp imm;ca dex imp;cc cpy abs;cd cmp abs;ce dec abs;
d0 bne rel;d1 cmp izy;d5 cmp zpx;d6 dec zpx;d8 cld imp;d9 cmp aby;dd cmp abx;de dec abx;
e0 cpx imm;e1 sbc izx;e4 cpx zp;e5 sbc zp;e6 inc zp;e8 inx imp;e9 sbc imm;ea nop imp;ec cpx abs;ed sbc abs;ee inc abs;
f0 beq rel;f1 sbc izy;f5 sbc zpx;f6 inc zpx;f8 sed imp;f9 sbc aby;fd sbc abx;fe inc abx
"""
DEC = {}   # opcode -> (mnemonic, mode)
ENC = {}   # (mnemonic, mode) -> opcode
for _item in _T.replace('\n', '').split(';'):
    _item = _item.strip()
    if not _item:
        continue
    _c, _mn, _md = _item.split()
    DEC[int(_c, 16)] = (_mn, _md)
    ENC[(_mn, _md)] = int(_c, 16)
SIZE = {'imp': 1, 'acc': 1, 'imm': 2, 'zp': 2, 'zpx': 2, 'zpy': 2, 'izx': 2, 'izy': 2, 'rel': 2,
        'abs': 3, 'abx': 3, 'aby': 3, 'ind': 3}
BRANCH = {'bpl', 'bmi', 'bvc', 'bvs', 'bcc', 'bcs', 'bne', 'beq'}
INVERT = {'bpl': 'bmi', 'bmi': 'bpl', 'bvc': 'bvs', 'bvs': 'bvc', 'bcc': 'bcs', 'bcs': 'bcc', 'bne': 'beq', 'beq': 'bne'}
REGS = {'A', 'X', 'Y', 'C', 'N', 'Z', 'V'}
NZ = {'N', 'Z'}


class I:
    """one instruction. tgt = target instruction for branches and local jmp/jsr"""
    __slots__ = ('mn', 'md', 'val', 'ref', 'reloc', 'tgt', 'old', 'new')

    def __init__(s, mn, md, val=0, ref=None, reloc=False, tgt=None, old=None):
        s.mn, s.md, s.val, s.ref, s.reloc, s.tgt, s.old = mn, md, val, ref, reloc, tgt, old
        s.new = None

    @property
    def size(s):
        return SIZE[s.md]

    def __repr__(s):
        return '%s %s %s%s%s' % (s.mn, s.md, hex(s.val), ' ref=%s' % (s.ref,) if s.ref else '', ' R' if s.reloc else '')


def replace(insts, k, n, new):
    """replace insts[k:k+n] by the list new, retargeting anything that pointed at the first old one"""
    first = insts[k]
    if new:
        for x in insts:
            if x.tgt is first:
                x.tgt = new[0]
    else:
        nxt = insts[k + n] if k + n < len(insts) else None
        for x in insts:
            if x.tgt is first:
                x.tgt = nxt
    insts[k:k + n] = new


# ------------------------------------------------------------------ effects
def effects(i):
    """(reads, writes) over A X Y C N Z V for one non-flow-ending instruction"""
    mn, md = i.mn, i.md
    r, w = set(), set()
    if md in ('abx', 'zpx', 'izx'):
        r.add('X')
    if md in ('aby', 'zpy', 'izy'):
        r.add('Y')
    if mn == 'lda': w |= {'A', 'N', 'Z'}
    elif mn == 'ldx': w |= {'X', 'N', 'Z'}
    elif mn == 'ldy': w |= {'Y', 'N', 'Z'}
    elif mn == 'sta': r.add('A')
    elif mn == 'stx': r.add('X')
    elif mn == 'sty': r.add('Y')
    elif mn == 'tax': r.add('A'); w |= {'X', 'N', 'Z'}
    elif mn == 'tay': r.add('A'); w |= {'Y', 'N', 'Z'}
    elif mn == 'txa': r.add('X'); w |= {'A', 'N', 'Z'}
    elif mn == 'tya': r.add('Y'); w |= {'A', 'N', 'Z'}
    elif mn in ('adc', 'sbc'): r |= {'A', 'C'}; w |= {'A', 'C', 'N', 'Z', 'V'}
    elif mn in ('and', 'ora', 'eor'): r.add('A'); w |= {'A', 'N', 'Z'}
    elif mn == 'cmp': r.add('A'); w |= {'C', 'N', 'Z'}
    elif mn == 'cpx': r.add('X'); w |= {'C', 'N', 'Z'}
    elif mn == 'cpy': r.add('Y'); w |= {'C', 'N', 'Z'}
    elif mn == 'bit': r.add('A'); w |= {'N', 'Z', 'V'}
    elif mn in ('asl', 'lsr', 'rol', 'ror'):
        if md == 'acc':
            r.add('A'); w.add('A')
        if mn in ('rol', 'ror'):
            r.add('C')
        w |= {'C', 'N', 'Z'}
    elif mn in ('inc', 'dec'): w |= NZ
    elif mn in ('inx', 'dex'): r.add('X'); w |= {'X', 'N', 'Z'}
    elif mn in ('iny', 'dey'): r.add('Y'); w |= {'Y', 'N', 'Z'}
    elif mn in ('clc', 'sec'): w.add('C')
    elif mn in ('bne', 'beq'): r.add('Z')
    elif mn in ('bpl', 'bmi'): r.add('N')
    elif mn in ('bcc', 'bcs'): r.add('C')
    elif mn in ('bvc', 'bvs'): r.add('V')
    elif mn == 'pha': r.add('A')
    elif mn == 'pla': w |= {'A', 'N', 'Z'}
    elif mn == 'nop': pass
    else:
        raise ValueError('effects: unhandled ' + mn)
    return r, w


TEMP = 0x033C
HELPER_READS = {}   # runtime symbol -> registers it may read on entry
CALL_READS = {'A', 'X', 'Y'}


def helper_reads(code, start, refs_at=None, reloc=None):
    """registers a runtime helper may read (on some path) before writing them.
    Explores all paths through the helper; jsr/jmp to unknown code counts as a read of A,X,Y."""
    result = set()
    for r in REGS:
        seen = set()
        stack = [start]
        live = False
        steps = 0
        while stack and not live:
            p = stack.pop()
            while True:
                steps += 1
                if steps > 3000 or p >= len(code) or p in seen:
                    live = live or (steps > 3000 or p >= len(code))
                    break
                seen.add(p)
                op = code[p]
                if op not in DEC:
                    live = True
                    break
                mn, md = DEC[op]
                n = SIZE[md]
                if mn == 'rts':
                    break
                if mn in ('jmp', 'jsr'):
                    if md == 'ind':
                        live = True
                        break
                    tgt = code[p + 1] | code[p + 2] << 8
                    if refs_at is not None and p in refs_at:
                        if r in CALL_READS:
                            live = True
                        if mn == 'jmp':
                            break
                        p += n
                        continue
                    if reloc is not None and p in reloc:       # a jump/call inside this object
                        if mn == 'jmp':
                            p = tgt
                            continue
                        if mn == 'jsr':                  # nested call: A,X,Y are its arguments, flags never are
                            if r in CALL_READS:
                                live = True
                                break
                            p += n
                            continue
                    live = True            # unresolved target: be conservative
                    break
                if mn in ('rti', 'brk', 'php', 'plp', 'tsx', 'txs', 'sei', 'cli', 'sed', 'cld') or md == 'ind':
                    live = True
                    break
                try:
                    rr, ww = effects(I(mn, md))
                except ValueError:
                    live = True
                    break
                if r in rr:
                    live = True
                    break
                if r in ww:
                    break
                if mn in BRANCH:
                    off = code[p + 1]
                    stack.append(p + 2 + (off - 256 if off > 127 else off))
                p += n
        if live:
            result.add(r)
    return result


def load_lib(libdir):
    for f in glob.glob(os.path.join(libdir, '*.obj')):
        try:
            o = Obj.parse(open(f, 'rb').read())
        except Exception:
            continue
        for nm, kind, off in o.defs:
            if nm.startswith(b'C$') and kind == 1 and off < len(o.code):
                HELPER_READS[bytes(nm)] = helper_reads(o.code, off, set(x[2] for x in o.refs), set(o.reloc))


# ------------------------------------------------------------------ function discovery / decode
def find_functions(o):
    code = o.code
    starts = sorted(set(off for nm, k, off in o.defs
                        if off + 1 < len(code) and code[off] == 0x85 and code[off + 1] == 0xfb))
    ret_sites = sorted(off for nm, k, off in o.refs if bytes(nm) == b'C$106' and k == 0)
    funcs = []
    for n, s in enumerate(starts):
        limit = starts[n + 1] if n + 1 < len(starts) else len(code)
        cand = [r for r in ret_sites if s <= r < limit]
        if not cand:
            continue
        end = cand[-1] + 3
        if end > limit:
            continue
        funcs.append((s, end))
    return funcs


def decode_func(o, start, end):
    code = o.code
    refs_at = {}
    for nm, kind, off in o.refs:
        if start <= off < end:
            refs_at.setdefault(off, []).append((bytes(nm), kind))
    reloc_at = set(x for x in o.reloc if start <= x < end)
    insts, pos, at = [], start, {}
    while pos < end:
        op = code[pos]
        if op not in DEC:
            return None
        mn, md = DEC[op]
        n = SIZE[md]
        if pos + n > end:
            return None
        v = 0
        if n == 2: v = code[pos + 1]
        elif n == 3: v = code[pos + 1] | code[pos + 2] << 8
        refs = refs_at.get(pos, [])
        if len(refs) > 1:
            return None
        ref = refs[0] if refs else None
        if md == 'ind' or mn in ('rti', 'brk', 'php', 'plp', 'pha', 'pla', 'txs', 'tsx', 'sei', 'cli', 'sed', 'cld'):
            return None
        i = I(mn, md, v, ref, pos in reloc_at, None, pos)
        if ref and md not in ('abs', 'abx', 'aby', 'imm'):
            return None
        if i.reloc and md not in ('abs', 'abx', 'aby'):
            return None
        at[pos] = i
        insts.append(i)
        pos += n
    if pos != end:
        return None
    for p in refs_at:
        if p not in at: return None
    for p in reloc_at:
        if p not in at: return None
    for i in insts:
        if i.md == 'rel':
            t = i.old + 2 + (i.val - 256 if i.val > 127 else i.val)
            if t not in at: return None
            i.tgt = at[t]
        elif i.reloc and i.mn in ('jmp', 'jsr'):
            if start <= i.val < end:
                if i.val not in at: return None
                i.tgt = at[i.val]
        elif i.reloc and start <= i.val < end:
            return None
    return insts


# ------------------------------------------------------------------ analysis
def targets_of(insts):
    return set(id(i.tgt) for i in insts if i.tgt is not None)


def dead_after(insts, k, regs, tset=None):
    """subset of regs that no path leaving insts[k] can read before it is rewritten (conservative:
    unknown control flow counts as a read)"""
    return dead_from(insts, [k + 1], regs)


def dead_from(insts, starts, regs):
    """like dead_after but from explicit start indices (all of them must be dead)"""
    index = {id(x): n for n, x in enumerate(insts)}
    out = set()
    for r in regs:
        seen = set()
        stack = list(starts)
        live = False
        steps = 0
        while stack and not live:
            j = stack.pop()
            while True:
                steps += 1
                if steps > 4000:
                    live = True
                    break
                if j >= len(insts):
                    live = True
                    break
                if j in seen:
                    break
                seen.add(j)
                i = insts[j]
                mn = i.mn
                if mn == 'rts':
                    live = True
                    break
                if mn == 'jmp':
                    if i.tgt is not None:
                        j = index[id(i.tgt)]
                        continue
                    rd = HELPER_READS.get(i.ref[0]) if (i.ref and i.ref[1] == 0) else None
                    if rd is None:
                        rd = CALL_READS               # unknown callee: A,X,Y are arguments, flags never are
                    if r in rd:
                        live = True
                    break
                if mn == 'jsr':
                    rd = HELPER_READS.get(i.ref[0]) if (i.ref and i.ref[1] == 0) else None
                    if rd is None:
                        rd = CALL_READS
                    if r in rd:
                        live = True
                        break
                    j += 1
                    continue
                rr, ww = effects(i)
                if r in rr:
                    live = True
                    break
                if r in ww:
                    break                      # killed on this path
                if mn in BRANCH:
                    stack.append(index[id(i.tgt)])
                j += 1
        if not live:
            out.add(r)
    return out


def temp_dead_after(insts, k, lo, hi):
    """is temp-stack memory [lo,hi) provably not read on any path after insts[k]?
    Reaching the function epilogue (jmp C$106) rewrites the temp stack, so that ends a path."""
    return temp_dead_from(insts, [k + 1], lo, hi)


def temp_dead_from(insts, starts, lo, hi):
    index = {id(x): n for n, x in enumerate(insts)}
    seen = set()
    stack = list(starts)
    steps = 0
    while stack:
        j = stack.pop()
        while True:
            steps += 1
            if steps > 4000 or j >= len(insts):
                return False
            if j in seen:
                break
            seen.add(j)
            i = insts[j]
            mn = i.mn
            if mn == 'jmp':
                if i.tgt is not None:
                    j = index[id(i.tgt)]
                    continue
                if i.ref and i.ref[0] == b'C$106' and i.ref[1] == 0:
                    break
                return False
            if mn == 'rts':
                return False
            if mn == 'jsr':
                # a helper/function reads only the top A bytes of the temp stack
                prev = insts[j - 1] if j > 0 else None
                if (prev is not None and prev.mn == 'lda' and prev.md == 'imm' and not prev.ref
                        and TEMP + prev.val <= lo):
                    j += 1
                    continue
                return False
            if i.md in ('abs', 'abx', 'aby') and not i.ref and not i.reloc and lo <= i.val < hi:
                if i.mn in ('sta', 'stx', 'sty') and i.md == 'abs':
                    break                      # overwritten before any read on this path
                return False
            if i.md in ('abx', 'aby') and not i.ref and not i.reloc and i.val < hi and i.val + 255 >= lo:
                return False
            if mn in BRANCH:
                stack.append(index[id(i.tgt)])
            j += 1
    return True


def mem_key(i):
    return (i.md, i.val, i.ref, i.reloc)


# ------------------------------------------------------------------ rules
def rule_forward(insts, stats):
    """X/Y already hold the value of a memory slot: drop the reload / use txa,tya"""
    tset = targets_of(insts)
    xmem, ymem = set(), set()
    for k, i in enumerate(insts):
        if id(i) in tset:
            xmem.clear(); ymem.clear()
        mn, md = i.mn, i.md
        if md in ('abs', 'zp') and not i.ref and not i.reloc:
            key = mem_key(i)
            if mn == 'ldx' and key in xmem or mn == 'ldy' and key in ymem:
                if dead_after(insts, k, NZ, tset) == NZ:
                    replace(insts, k, 1, [])
                    stats['reload'] += 1
                    return True
            if mn == 'lda' and (key in xmem or key in ymem):
                n = I('txa' if key in xmem else 'tya', 'imp', old=i.old)
                replace(insts, k, 1, [n])
                stats['lda->t'] += 1
                return True
        if mn in ('jsr', 'jmp'):
            xmem.clear(); ymem.clear()
            continue
        if mn in ('sta', 'stx', 'sty') and md in ('izx', 'izy', 'abx', 'aby', 'zpx', 'zpy'):
            xmem.clear(); ymem.clear()           # may alias anything
            continue
        r, w = effects(i)
        if 'X' in w: xmem.clear()
        if 'Y' in w: ymem.clear()
        if mn in ('stx', 'sty', 'sta', 'inc', 'dec', 'asl', 'lsr', 'rol', 'ror') and md in ('abs', 'zp'):
            key = mem_key(i)
            xmem.discard(key); ymem.discard(key)
            if not i.ref and not i.reloc:
                if mn == 'stx': xmem.add(key)
                if mn == 'sty': ymem.add(key)
            elif mn in ('stx', 'sty') or i.reloc or i.ref:
                # a relocated store to a variable: the slot is still a fixed address
                if mn == 'stx': xmem.add(key)
                if mn == 'sty': ymem.add(key)
    return False


def rule_jmp_next(insts, stats):
    for k in range(len(insts) - 1):
        i = insts[k]
        if i.mn == 'jmp' and i.tgt is not None and i.tgt is insts[k + 1] and not i.ref:
            replace(insts, k, 1, [])
            stats['jmp-next'] += 1
            return True
    return False


def rule_zero_pair(insts, stats):
    """ldx #K; ldy #K; stx M; sty N  (X,Y,A dead after)  ->  lda #K; sta M; sta N"""
    tset = targets_of(insts)
    for k in range(len(insts) - 3):
        a, b, c, d = insts[k:k + 4]
        if not (a.mn == 'ldx' and a.md == 'imm' and b.mn == 'ldy' and b.md == 'imm' and not a.ref and not b.ref):
            continue
        if a.val != b.val:
            continue
        if not (c.mn == 'stx' and d.mn == 'sty' and c.md in ('abs', 'zp') and d.md in ('abs', 'zp')):
            continue
        if any(id(x) in tset for x in (b, c, d)):
            continue
        want = {'A', 'X', 'Y', 'N', 'Z'}
        if dead_after(insts, k + 3, want, tset) != want:
            continue
        new = [I('lda', 'imm', a.val, old=a.old),
               I('sta', c.md, c.val, c.ref, c.reloc, old=c.old),
               I('sta', d.md, d.val, d.ref, d.reloc, old=d.old)]
        # c and d may be branch targets only if they were in tset (checked); a was handled by replace()
        replace(insts, k, 4, new)
        stats['zero-pair'] += 1
        return True
    return False


def rule_inc16(insts, stats):
    """clc;txa;adc #1;tax;tya;adc #0;tay -> inx;bne +1;iny   (A,C,N,Z,V dead after)"""
    pat = [('clc', 'imp'), ('txa', 'imp'), ('adc', 'imm'), ('tax', 'imp'), ('tya', 'imp'), ('adc', 'imm'), ('tay', 'imp')]
    tset = targets_of(insts)
    for k in range(len(insts) - 7):
        w = insts[k:k + 7]
        if any((x.mn, x.md) != p for x, p in zip(w, pat)):
            continue
        if w[2].val != 1 or w[5].val != 0 or any(x.ref or x.reloc for x in w):
            continue
        if any(id(x) in tset for x in w[1:]):
            continue
        want = {'A', 'C', 'N', 'Z', 'V'}
        if dead_after(insts, k + 6, want, tset) != want:
            continue
        nxt = insts[k + 7]
        inx = I('inx', 'imp', old=w[0].old)
        iny = I('iny', 'imp')
        bne = I('bne', 'rel', 0, tgt=nxt)
        replace(insts, k, 7, [inx, bne, iny])
        stats['inc16'] += 1
        return True
    return False


def rule_dup_ldy(insts, stats):
    tset = targets_of(insts)
    for k in range(len(insts) - 1):
        a, b = insts[k], insts[k + 1]
        if (a.mn == b.mn and a.mn in ('ldy', 'ldx') and a.md == b.md == 'imm' and a.val == b.val
                and not a.ref and not b.ref and id(b) not in tset):
            replace(insts, k + 1, 1, [])
            stats['dup-ld'] += 1
            return True
    return False


def layout_of(insts, base):
    lay, p = {}, base
    for i in insts:
        lay[id(i)] = p
        p += i.size
    return lay


def rule_branch_over_jmp(insts, base, stats):
    """bXX L; jmp T; L:  ->  b!XX T   when T is in branch range"""
    tset = targets_of(insts)
    lay = layout_of(insts, base)
    for k in range(len(insts) - 2):
        b, j, l = insts[k], insts[k + 1], insts[k + 2]
        if b.mn in BRANCH and b.tgt is l and j.mn == 'jmp' and j.tgt is not None and not j.ref and id(j) not in tset:
            rel = lay[id(j.tgt)] - (lay[id(b)] + 2)
            if -120 <= rel <= 120:
                nb = I(INVERT[b.mn], 'rel', 0, tgt=j.tgt, old=b.old)
                replace(insts, k, 2, [nb])
                stats['bra-over-jmp'] += 1
                return True
    return False


# ------------------------------------------------------------------ inline the compare helpers
CMP = {b'C$0063': ('s', 'lt'), b'C$0067': ('s', 'gt'), b'C$0071': ('s', 'le'), b'C$0075': ('s', 'ge'),
       b'C$0065': ('u', 'lt'), b'C$0069': ('u', 'gt'), b'C$0073': ('u', 'le'), b'C$0077': ('u', 'ge'),
       b'C$0079': ('e', 'eq'), b'C$0082': ('e', 'ne')}
ALLREGS = {'A', 'X', 'Y', 'C', 'N', 'Z', 'V'}


def opd(i):
    """operand descriptor of a load instruction"""
    return (i.md, i.val, i.ref, i.reloc)


def E(mn, d):
    return I(mn, d[0], d[1], d[2], d[3])


def imm_src(i):
    """a source that is stable (not temp-stack memory)"""
    if i.md == 'imm':
        return True
    if i.md in ('abs', 'zp'):
        return i.ref is not None or i.reloc or not (TEMP <= i.val < 0x0400)
    return False


def emit_compare(kind, op, A, B, T, F):
    """inline code: falls into T when the comparison holds, else jumps to F"""
    def jf():
        return I('jmp', 'abs', 0, None, True, F)
    if kind == 'e':
        if op == 'eq':
            j = jf()
            return [E('lda', A[0]), E('cmp', B[0]), I('bne', 'rel', 0, tgt=j), E('lda', A[1]), E('cmp', B[1]),
                    I('beq', 'rel', 0, tgt=T), j]
        return [E('lda', A[0]), E('cmp', B[0]), I('bne', 'rel', 0, tgt=T), E('lda', A[1]), E('cmp', B[1]),
                I('bne', 'rel', 0, tgt=T), jf()]
    swap = op in ('gt', 'le')
    p, q = (B, A) if swap else (A, B)
    out = [E('lda', p[0]), E('cmp', q[0]), E('lda', p[1]), E('sbc', q[1])]
    if kind == 's':
        br = I('bmi' if op in ('lt', 'gt') else 'bpl', 'rel', 0, tgt=T)
        out += [I('bvc', 'rel', 0, tgt=br), I('eor', 'imm', 0x80), br, jf()]
    else:
        br = I('bcc' if op in ('lt', 'gt') else 'bcs', 'rel', 0, tgt=T)
        out += [br, jf()]
    return out


def match_consumer(insts, c):
    """the ldx/ldy/txa/bne.. idioms that test a 0/1 result stored at some base; (base, length) or None"""
    def slot(i, mn):
        return i.mn == mn and i.md == 'abs' and not i.ref and not i.reloc
    n = len(insts)
    if c + 7 < n:
        w = insts[c:c + 7]
        if (slot(w[0], 'ldx') and slot(w[1], 'ldy') and w[1].val == w[0].val + 1 and w[2].mn == 'txa'
                and w[3].mn == 'bne' and w[4].mn == 'tya' and w[5].mn == 'bne' and w[6].mn == 'jmp'
                and w[3].tgt is w[5].tgt and w[6].tgt is not None and not w[6].ref
                and w[3].tgt is insts[c + 7]):
            return w[0].val, 7
    if c + 5 < n:
        w = insts[c:c + 5]
        if (slot(w[0], 'lda') and slot(w[2], 'lda') and w[0].val == w[2].val + 1 and w[1].mn == 'bne'
                and w[3].mn == 'bne' and w[4].mn == 'jmp' and w[1].tgt is w[3].tgt and w[4].tgt is not None
                and not w[4].ref and w[1].tgt is insts[c + 5]):
            return w[2].val, 5
    return None


def push_pair(w, s0):
    """ldx src; ldy src; stx s0; sty s0+1 with stable sources -> (lo, hi) operand descriptors"""
    if not (w[0].mn == 'ldx' and w[1].mn == 'ldy' and w[2].mn == 'stx' and w[3].mn == 'sty'):
        return None
    if not (w[2].md == 'abs' and not w[2].ref and not w[2].reloc and w[2].val == s0 and
            w[3].md == 'abs' and not w[3].ref and not w[3].reloc and w[3].val == s0 + 1):
        return None
    if not (imm_src(w[0]) and imm_src(w[1])):
        return None
    return opd(w[0]), opd(w[1])


def rule_compare(insts, stats):
    tset = targets_of(insts)
    for j, jm in enumerate(insts):
        if jm.mn != 'jsr' or not jm.ref or jm.ref[1] != 0 or jm.ref[0] not in CMP or j == 0:
            continue
        ld = insts[j - 1]
        if not (ld.mn == 'lda' and ld.md == 'imm' and not ld.ref and ld.val >= 4):
            continue
        base = TEMP + ld.val - 4
        m = match_consumer(insts, j + 1)
        if not m or m[0] != base:
            continue
        last = j + m[1]                      # index of the final jmp F
        T = insts[last + 1]
        F = insts[last].tgt
        if id(jm) in tset or any(id(x) in tset for x in insts[j + 2:last + 1]):
            continue
        keep = id(insts[j + 1]) in tset          # the result test is a join point (&&, ||): keep it
        kind, op = CMP[jm.ref[0]]
        start = j - 1                        # the lda #n
        A = [('abs', base, None, False), ('abs', base + 1, None, False)]
        B = [('abs', base + 2, None, False), ('abs', base + 3, None, False)]
        combined = False
        if j - 1 >= 8:
            pre = insts[j - 9:j - 1]
            pa = push_pair(pre[0:4], base)
            pb = push_pair(pre[4:8], base + 2)
            if pa and pb and not any(id(x) in tset for x in pre[1:]):
                A, B = list(pa), list(pb)
                start = j - 9
                combined = True
        new = emit_compare(kind, op, A, B, T, F)
        idxT = last + 1
        idxF = [n for n, x in enumerate(insts) if x is F][0]
        if dead_from(insts, [idxT, idxF], ALLREGS) != ALLREGS:
            continue
        # the 0/1 result is no longer stored in the temp stack: nobody may read it afterwards
        if not temp_dead_from(insts, [idxT, idxF], base, base + 2):
            continue
        replace(insts, start, (j if keep else last) + 1 - start, new)
        stats['cmp-inline' if combined else 'cmp-slots'] += 1
        if keep:
            stats['cmp-keep'] += 1
        return True
    return False


RET_SYM = b'C$RETURN'


def rule_return_copy(insts, stats):
    """stx T; sty T+1; lda #2; ldx #2; jsr C$0112 (copy to the return slot)  ->  stx RET; sty RET+1"""
    tset = targets_of(insts)
    for k in range(len(insts) - 4):
        w = insts[k:k + 5]
        if not (w[0].mn == 'stx' and w[1].mn == 'sty' and w[2].mn == 'lda' and w[3].mn == 'ldx' and w[4].mn == 'jsr'):
            continue
        if not (w[4].ref and w[4].ref[0] == b'C$0112' and w[4].ref[1] == 0):
            continue
        if not (w[2].md == 'imm' and w[2].val == 2 and not w[2].ref and w[3].md == 'imm' and w[3].val == 2 and not w[3].ref):
            continue
        if not all(x.md == 'abs' and not x.ref and not x.reloc for x in (w[0], w[1])):
            continue
        if w[0].val != TEMP or w[1].val != TEMP + 1:
            continue
        if any(id(x) in tset for x in w[1:]):
            continue
        if dead_after(insts, k + 4, ALLREGS) != ALLREGS:
            continue
        if not temp_dead_after(insts, k + 4, TEMP, TEMP + 0x40):
            continue
        replace(insts, k, 5, [I('stx', 'abs', 0, (RET_SYM, 0), False), I('sty', 'abs', 0, (RET_SYM, 4), False)])
        stats['ret-copy'] += 1
        return True
    return False


def rule_eq_branch(insts, base, stats):
    """16 bit ==/!= against a constant: shorten
         cx; bne J; cy; beq T; J: jmp F; T:   ->  cx; bne F; cy; bne F; T:
         cx; bne T; cy; bne T; jmp F; T:      ->  cx; bne T; cy; beq F; T:     (F within branch range)"""
    tset = targets_of(insts)
    lay = layout_of(insts, base)
    CMPS = ('cpx', 'cpy', 'cmp')
    for k in range(len(insts) - 5):
        w = insts[k:k + 6]
        if not (w[0].mn in CMPS and w[1].mn == 'bne' and w[2].mn in CMPS and w[3].mn in ('beq', 'bne')
                and w[4].mn == 'jmp' and w[4].tgt is not None and not w[4].ref):
            continue
        F = w[4].tgt
        if id(w[2]) in tset or id(w[4]) in tset and w[1].tgt is not w[4]:
            continue
        rel_ok = lambda src: -120 <= lay[id(F)] - (lay[id(src)] + 2) <= 120
        if w[3].mn == 'beq' and w[1].tgt is w[4] and w[3].tgt is w[5]:
            # the jmp is only reachable through the first bne
            if sum(1 for x in insts if x.tgt is w[4]) != 1 or not (rel_ok(w[1]) and rel_ok(w[3])):
                continue
            nb1 = I('bne', 'rel', 0, tgt=F, old=w[1].old)
            nb3 = I('bne', 'rel', 0, tgt=F, old=w[3].old)
            replace(insts, k + 1, 4, [nb1, w[2], nb3])
            # w[2] kept in place; fix up the replaced slice: [bne, cy, bne]
            stats['eq-branch'] += 1
            return True
        if w[3].mn == 'bne' and w[1].tgt is w[5] and w[3].tgt is w[5] and id(w[4]) not in tset:
            if not rel_ok(w[3]):
                continue
            nb3 = I('beq', 'rel', 0, tgt=F, old=w[3].old)
            replace(insts, k + 3, 2, [nb3])
            stats['ne-branch'] += 1
            return True
    return False


def rule_add_const(insts, stats):
    """clc;txa;adc #K;tax;tya;adc #0;tay  ->  clc;txa;adc #K;tax;bcc +1;iny   (A,C,N,Z,V dead after)"""
    pat = [('clc', 'imp'), ('txa', 'imp'), ('adc', 'imm'), ('tax', 'imp'), ('tya', 'imp'), ('adc', 'imm'), ('tay', 'imp')]
    tset = targets_of(insts)
    for k in range(len(insts) - 7):
        w = insts[k:k + 7]
        if any((x.mn, x.md) != p for x, p in zip(w, pat)):
            continue
        if w[5].val != 0 or w[5].ref or w[2].ref or w[2].reloc or w[2].val == 1:
            continue
        if any(id(x) in tset for x in w[1:]):
            continue
        if dead_after(insts, k + 6, {'A', 'C', 'N', 'Z', 'V'}) != {'A', 'C', 'N', 'Z', 'V'}:
            continue
        nxt = insts[k + 7]
        bcc = I('bcc', 'rel', 0, tgt=nxt)
        iny = I('iny', 'imp')
        replace(insts, k + 4, 3, [bcc, iny])
        stats['add-const'] += 1
        return True
    return False


def optimize_func(insts, base, stats):
    while True:
        if rule_compare(insts, stats): continue
        if rule_return_copy(insts, stats): continue
        if rule_forward(insts, stats): continue
        if rule_jmp_next(insts, stats): continue
        if rule_zero_pair(insts, stats): continue
        if rule_inc16(insts, stats): continue
        if rule_dup_ldy(insts, stats): continue
        if rule_eq_branch(insts, base, stats): continue
        if rule_add_const(insts, stats): continue
        if rule_branch_over_jmp(insts, base, stats): continue
        break


# ------------------------------------------------------------------ object rewrite
def map_addr(newpos, v):
    """new value of a code-relative address operand; tolerates small negative/positive addends (array-1)"""
    if v in newpos:
        return newpos[v]
    for d in (1, 2, 3, 4):                     # address = target - d  (e.g. a[0]-1 indexed with x)
        if (v + d) & 0xFFFF in newpos:
            return (newpos[(v + d) & 0xFFFF] - d) & 0xFFFF
    for d in (1, 2, 3, 4):                     # address = target + d
        if v - d in newpos:
            return newpos[v - d] + d
    raise ValueError('cannot map address %x' % v)


def optimize_obj(o, rules=True, verbose=False):
    stats = {k: 0 for k in ('reload', 'lda->t', 'jmp-next', 'zero-pair', 'inc16', 'dup-ld', 'bra-over-jmp',
                            'cmp-inline', 'cmp-slots', 'cmp-keep', 'ret-copy', 'eq-branch', 'ne-branch', 'add-const')}
    code = o.code
    plan_in = []
    for s, e in find_functions(o):
        insts = decode_func(o, s, e)
        if insts is None:
            if verbose: print('  cannot decode function at', hex(s))
            continue
        plan_in.append((s, e, insts))
    if not plan_in:
        return o, stats
    for s, e, insts in plan_in:
        if rules:
            optimize_func(insts, s, stats)
    # lay out the new code
    newcode = bytearray()
    newpos = {}
    plan = []
    pos = 0
    for s, e, insts in plan_in:
        for b in range(pos, s):
            newpos[b] = len(newcode) + (b - pos)
        newcode += code[pos:s]
        base = len(newcode)
        p = base
        for i in insts:
            i.new = p
            p += i.size
        for i in insts:
            if i.old is not None:
                newpos.setdefault(i.old, i.new)
        plan.append((s, e, insts, base))
        newcode += bytes(p - base)
        pos = e
    for b in range(pos, len(code)):
        newpos[b] = len(newcode) + (b - pos)
    newcode += code[pos:]
    newpos[len(code)] = len(newcode)
    # old instruction starts whose instruction vanished map to the next surviving one
    for s, e, insts, base in plan:
        pass
    inside = lambda off: any(s <= off < e for s, e, _, _ in plan)
    new_reloc, new_refs = [], []
    for s, e, insts, base in plan:
        for i in insts:
            b = bytearray([ENC[(i.mn, i.md)]])
            if i.md == 'rel':
                rel = i.tgt.new - (i.new + 2)
                if not -128 <= rel <= 127:
                    raise ValueError('branch out of range')
                b.append(rel & 255)
            elif i.size == 2:
                b.append(i.val & 255)
            elif i.size == 3:
                v = i.val
                if i.reloc:
                    v = i.tgt.new if i.tgt is not None else map_addr(newpos, v)
                b += bytes((v & 255, v >> 8 & 255))
            newcode[i.new:i.new + len(b)] = b
            if i.reloc:
                new_reloc.append(i.new)
            if i.ref:
                new_refs.append([i.ref[0], i.ref[1], i.new])
    # relocated words that live in data (pointers into the code)
    for off in o.reloc:
        if inside(off):
            continue
        n_off = newpos[off]
        old_v = code[off + 1] | code[off + 2] << 8
        nv = newpos[old_v]
        newcode[n_off + 1] = nv & 255
        newcode[n_off + 2] = nv >> 8
        new_reloc.append(n_off)
    for nm, kind, off in o.refs:
        if not inside(off):
            new_refs.append([nm, kind, newpos[off]])
    n = Obj()
    n.code = newcode
    n.reloc = sorted(new_reloc)
    n.defs = [[nm, k, newpos[off]] for nm, k, off in o.defs]
    n.refs = sorted(new_refs, key=lambda r: r[2])
    n.res = o.res
    n.trailing = o.trailing
    return n, stats


def main():
    args = sys.argv[1:]
    verbose = '-v' in args
    args = [a for a in args if a != '-v']
    rules = '--no-rules' not in args
    args = [a for a in args if a != '--no-rules']
    if '--lib' in args:
        k = args.index('--lib')
        load_lib(args[k + 1])
        del args[k:k + 2]
    src, dst = args[0], args[1]
    o = Obj.parse(open(src, 'rb').read())
    try:
        n, st = optimize_obj(o, rules, verbose)
    except Exception as ex:
        print('%s: UNCHANGED (%s)' % (os.path.basename(src), ex))
        open(dst, 'wb').write(open(src, 'rb').read())
        return
    open(dst, 'wb').write(n.build())
    tot = {k: v for k, v in st.items() if v}
    print('%s: code %d -> %d bytes (%d saved) %s' % (os.path.basename(src), len(o.code), len(n.code),
                                                    len(o.code) - len(n.code), tot))


if __name__ == '__main__':
    main()