mirror of
https://github.com/valentderah/totk-pc-recomp.git
synced 2026-11-11 01:41:02 +00:00
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
250 lines
9.9 KiB
Python
250 lines
9.9 KiB
Python
"""Differential test of the lifter on the game's own instructions: lifted C against Unicorn, one instruction at a time.
|
|
|
|
Samples distinct instruction words from the .text of each ExeFS module, one per opcode signature first, lifts each
|
|
with the portable FP/SIMD lowering, runs it from random register state, and compares X0-X30, SP, NZCV, PC (for
|
|
branches), V0-V31 and the scratch buffer (for loads and stores) with Unicorn as a Cortex-A57.
|
|
|
|
Skipped: system instructions and exceptions, whose single step cannot be compared, and exclusives and ARMv8.1
|
|
atomics, which atomic.py covers. Encodings the lifter rejects are counted per module and mnemonic; --coverage lifts
|
|
every distinct word instead of testing.
|
|
|
|
Adapted from LAS Recomp tools/test_lift.py @ da35b83 (MIT, Copyright (c) 2026 Peiikai).
|
|
|
|
Usage: lift.py [--exefs DIR] [--modules main,sdk,subsdk0] [--encodings N] [--cases K] [--seed S] [--jobs J]
|
|
lift.py --coverage
|
|
"""
|
|
import argparse
|
|
import os
|
|
import platform
|
|
import random
|
|
import sys
|
|
import time
|
|
import zlib
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import unicorn as uc
|
|
|
|
from common import BUF, BUF_SIZE, ROOT, Emu, build_harness, diff_state, dis, mnemonic, rand_fpcr, rand_v, rand_x, \
|
|
run_harness
|
|
from lifter import Lifter, Unsupported, bits
|
|
from nso import Nso
|
|
|
|
BUF_MID = BUF + 0x10000
|
|
# Test load addresses, below 2^39 and clear of BUF.
|
|
BASES = {'main': 0x7000000000, 'subsdk0': 0x7008000000, 'sdk': 0x7010000000}
|
|
|
|
|
|
def is_mem(w):
|
|
return bits(w, 25, 1) == 0 and bits(w, 27, 1) == 1
|
|
|
|
|
|
def classify(w):
|
|
"""'skip' for instructions whose single-step behaviour cannot be compared."""
|
|
if (w & 0xFFC00000) == 0xD5000000: # system: hints, barriers, MRS/MSR, cache maintenance
|
|
return 'skip'
|
|
if (w & 0x3F000000) == 0x08000000 and (bits(w, 23, 1) == 0 or bits(w, 21, 1)):
|
|
return 'skip' # exclusives need a paired load; CAS is ARMv8.1; both are covered by atomic.py
|
|
if (w & 0x3F200C00) == 0x38200000: # LSE atomics, ARMv8.1: atomic.py
|
|
return 'skip'
|
|
if (w & 0xFF000000) == 0xD4000000: # SVC, BRK, HLT
|
|
return 'skip'
|
|
if (w & 0xFFFF0000) == 0: # UDF
|
|
return 'skip'
|
|
return 'ok'
|
|
|
|
|
|
def class_keys(words):
|
|
"""Opcode signature per word: register fields (Rd, Rn, Rm) cleared."""
|
|
return words & np.uint32(0xFFE0FC00)
|
|
|
|
|
|
def to_test_body(ins, pc):
|
|
if ins.kind == 'op':
|
|
return ins.code
|
|
nxt = f'0x{pc + 4:x}ull'
|
|
if ins.kind == 'b':
|
|
return f'c->pc = 0x{ins.target:x}ull;'
|
|
if ins.kind == 'bl':
|
|
return f'c->x[30] = {nxt}; c->pc = 0x{ins.target:x}ull;'
|
|
if ins.kind == 'bcond':
|
|
return f'c->pc = ({ins.cond}) ? 0x{ins.target:x}ull : {nxt};'
|
|
if ins.kind == 'br':
|
|
return f'c->pc = {ins.target};'
|
|
if ins.kind == 'blr':
|
|
return f'{{ uint64_t t = {ins.target}; c->x[30] = {nxt}; c->pc = t; }}'
|
|
if ins.kind == 'ret':
|
|
return 'c->pc = c->x[30];'
|
|
raise Unsupported(ins.kind)
|
|
|
|
|
|
def pick_module(name, nso, base, count, r, unsupported):
|
|
"""Distinct words of .text, one per opcode signature first, lifted. Returns [(word, pc, body)]."""
|
|
lo, size = nso.text
|
|
words = np.frombuffer(bytes(nso.image[lo:lo + size]), dtype='<u4')
|
|
uniq, first = np.unique(words, return_index=True)
|
|
keys = class_keys(uniq)
|
|
order = list(range(len(uniq)))
|
|
r.shuffle(order)
|
|
seen, primary, rest = set(), [], []
|
|
for i in order:
|
|
k = int(keys[i])
|
|
(rest if k in seen else primary).append(i)
|
|
seen.add(k)
|
|
lifter = Lifter(base, 'c')
|
|
picks = []
|
|
for i in primary + rest:
|
|
if len(picks) >= count:
|
|
break
|
|
w = int(uniq[i])
|
|
if classify(w) == 'skip':
|
|
continue
|
|
pc = base + lo + int(first[i]) * 4
|
|
try:
|
|
body = to_test_body(lifter.lift(w, pc), pc)
|
|
except Unsupported:
|
|
unsupported[(name, mnemonic(w))] += 1
|
|
continue
|
|
picks.append((w, pc, body))
|
|
print(f'{name}: {len(words)} words, {len(uniq)} distinct, {len(seen)} signatures, {len(picks)} picked')
|
|
return picks
|
|
|
|
|
|
def coverage(name, nso, base):
|
|
"""Lift every distinct word of .text. Returns {kind mnemonic: [distinct words, occurrences]} of words the lifter
|
|
rejects (Unsupported) or turns into an undefined-instruction trap (rt_udf)."""
|
|
lo, size = nso.text
|
|
words = np.frombuffer(bytes(nso.image[lo:lo + size]), dtype='<u4')
|
|
uniq, first, counts = np.unique(words, return_index=True, return_counts=True)
|
|
lifter = Lifter(base, 'c')
|
|
rejected = {}
|
|
for w, i, n in zip(uniq.tolist(), first.tolist(), counts.tolist()):
|
|
try:
|
|
ins = lifter.lift(w, base + lo + i * 4)
|
|
if ins.kind != 'trap' or not ins.code.startswith('rt_udf'):
|
|
continue
|
|
key = f'udf {mnemonic(w)}'
|
|
except Unsupported:
|
|
key = f'unsupported {mnemonic(w)}'
|
|
e = rejected.setdefault(key, [0, 0])
|
|
e[0] += 1
|
|
e[1] += n
|
|
total = sum(e[1] for e in rejected.values())
|
|
print(f'{name}: {len(words)} words, {len(uniq)} distinct, {total} rejected '
|
|
f'({sum(e[0] for e in rejected.values())} distinct)')
|
|
for mn, (d, n) in sorted(rejected.items(), key=lambda kv: -kv[1][1]):
|
|
print(f' {mn:24} {n:8} occurrences {d:6} distinct')
|
|
return rejected
|
|
|
|
|
|
def make_case(r, w):
|
|
if is_mem(w):
|
|
x = [BUF_MID + r.randrange(0, 0x100, 8) for _ in range(31)]
|
|
if (bits(w, 21, 1) and bits(w, 10, 2) == 2) or (w & 0xBFA00000) == 0x0C800000: # register offset
|
|
rm = bits(w, 16, 5)
|
|
if rm != bits(w, 5, 5) and rm < 31:
|
|
x[rm] = r.randrange(0, 0x100)
|
|
else:
|
|
x = [rand_x(r) for _ in range(31)]
|
|
sp = BUF_MID + r.randrange(0, 0x100, 16)
|
|
nzcv = r.getrandbits(4) << 28
|
|
v = [rand_v(r) for _ in range(32)]
|
|
fpcr = rand_fpcr(r)
|
|
return fpcr, x, sp, nzcv, v
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
|
|
ap.add_argument('--exefs', type=Path, default=ROOT / 'build' / 'extracted' / 'v121' / 'exefs')
|
|
ap.add_argument('--modules', default='main,sdk,subsdk0')
|
|
ap.add_argument('--encodings', type=int, default=20000, help='per module')
|
|
ap.add_argument('--cases', type=int, default=3, help='random states per encoding')
|
|
ap.add_argument('--seed', type=int, default=1)
|
|
ap.add_argument('--jobs', type=int, default=os.cpu_count() or 4)
|
|
ap.add_argument('--show', type=int, default=40)
|
|
ap.add_argument('--cflags', default='-march=x86-64-v3' if platform.machine().lower() in ('amd64', 'x86_64') else '')
|
|
ap.add_argument('--coverage', action='store_true',
|
|
help='only lift every distinct word and list the rejected ones (capstone "?" = undecodable data)')
|
|
args = ap.parse_args()
|
|
r = random.Random(args.seed)
|
|
t0 = time.time()
|
|
names = args.modules.split(',')
|
|
if args.coverage:
|
|
for name in names:
|
|
coverage(name, Nso(str(args.exefs / name)), BASES[name])
|
|
print(f'({time.time() - t0:.0f}s)')
|
|
return 0
|
|
|
|
work = ROOT / 'build' / 'test' / 'lift'
|
|
work.mkdir(parents=True, exist_ok=True)
|
|
images, picks, unsupported = [], [], Counter()
|
|
for name in names:
|
|
nso = Nso(str(args.exefs / name))
|
|
base = BASES[name]
|
|
img = work / f'{name}.img'
|
|
img.write_bytes(bytes(nso.image))
|
|
images.append((base, img, nso.image))
|
|
picks += [(name,) + p for p in pick_module(name, nso, base, args.encodings, r, unsupported)]
|
|
print(f'lifted {len(picks)} encodings ({time.time() - t0:.0f}s)')
|
|
|
|
t1 = time.time()
|
|
exe = build_harness(work, [p[3] for p in picks], args.jobs, args.cflags)
|
|
print(f'compiled ({time.time() - t1:.0f}s)')
|
|
|
|
emu = Emu([(base, data) for base, _, data in images])
|
|
mu = emu.mu
|
|
buf0 = bytes(r.getrandbits(8) for _ in range(4096)) * (BUF_SIZE // 4096)
|
|
mu.mem_write(BUF, buf0)
|
|
cases, expect = [], []
|
|
t2 = time.time()
|
|
for k, (name, w, pc, _) in enumerate(picks):
|
|
mem = is_mem(w)
|
|
for _ in range(args.cases):
|
|
fpcr, x, sp, nzcv, v = make_case(r, w)
|
|
try:
|
|
ox, osp, onz, opc, ov = emu.run(pc, fpcr, x, sp, nzcv, v)
|
|
except uc.UcError:
|
|
if mem:
|
|
mu.mem_write(BUF, buf0)
|
|
continue
|
|
mh = 0
|
|
if mem:
|
|
mh = zlib.crc32(bytes(mu.mem_read(BUF, BUF_SIZE)))
|
|
mu.mem_write(BUF, buf0)
|
|
cases.append((k, fpcr, x, sp, nzcv, v, mem))
|
|
expect.append((ox, osp, onz, opc, ov, mh))
|
|
print(f'unicorn: {len(cases)} cases ({time.time() - t2:.0f}s)')
|
|
|
|
t3 = time.time()
|
|
recs, err = run_harness(exe, work, buf0, cases, [(base, img) for base, img, _ in images])
|
|
print(f'harness: {len(recs)} of {len(cases)} cases ({time.time() - t3:.0f}s)')
|
|
if err:
|
|
k = cases[len(recs)][0] if len(recs) < len(cases) else None
|
|
where = f', stopped at {picks[k][1]:08x} {dis(picks[k][1])}' if k is not None else ''
|
|
print(f'harness failed: {err}{where}')
|
|
|
|
fails, nfail = {}, 0
|
|
for case, want, got in zip(cases, expect, recs):
|
|
k = case[0]
|
|
name, w, pc, _ = picks[k]
|
|
d = diff_state(got, want, check_pc=got[3] != 0, check_mem=case[6])
|
|
if d:
|
|
nfail += 1
|
|
fails.setdefault(w, (name, pc, d[:4]))
|
|
|
|
print('unsupported encodings (distinct words) by module and mnemonic:')
|
|
for (name, mn), n in sorted(unsupported.items(), key=lambda kv: (kv[0][0], -kv[1])):
|
|
print(f' {name:8} {mn:10} {n}')
|
|
print(f'{len(picks)} encodings, {len(cases)} cases: {nfail} failing cases, {len(fails)} failing encodings '
|
|
f'({time.time() - t0:.0f}s total)')
|
|
if fails:
|
|
print(' failing by mnemonic:', dict(Counter(mnemonic(w) for w in fails).most_common(40)))
|
|
for w, (name, pc, d) in list(fails.items())[:args.show]:
|
|
print(f' {name:8} {w:08x} {dis(w, pc):40} {"; ".join(d)}')
|
|
return 1 if fails or err else 0
|
|
|
|
|
|
if __name__ == '__main__':
|
|
sys.exit(main())
|