#!/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()