#!/usr/bin/env python3
"""Name every DCN register the DMCUB firmware touches, and say which function
touches it, read or write.

How register access actually works in this firmware (this was a wrong turn in the
first pass -- register addresses are *not* l32r literals):

    l32r  a6, =0x60018030        ; base of a 2057-entry register-address table
    l32i  a10, a6, 0xe8          ; a10 = absolute register address (0x5ffc____)
    l32i  a9, a10, 0             ; read  |  s32i a9, a10, 0  ; write

The table holds absolute addresses in the DMCUB's register aperture. The aperture
is recovered by voting `value = APERTURE + 4*dword_offset` against AMD's published
offsets and then *validating semantically*: at 0x5ffcd300 the table's runs decode
as coherent register blocks (DMCUB_REGION3_CW*, DMCUB_INTERRUPT_*, DMCUB_SCRATCH*,
VPG0_* then VPG1_* in instance order), 1597/2153 values resolve to a known
register and 1446 neighbour pairs share a block prefix. Competing apertures score
strictly worse; see --vote.

usage:
  dmcub_regs.py <image> <offset_header> table              dump the register table
  dmcub_regs.py <image> <offset_header> access [--grep RE] register touches by function
  dmcub_regs.py <image> <offset_header> vote               aperture evidence + controls
"""
import argparse
import collections
import importlib.util
import os
import re
import struct

HERE = os.path.dirname(os.path.abspath(__file__))


def _load(name):
    s = importlib.util.spec_from_file_location(name, os.path.join(HERE, name + '.py'))
    m = importlib.util.module_from_spec(s)
    s.loader.exec_module(m)
    return m


dis = _load('dmcub_dis')
regmap = _load('dmcub_regmap')

APERTURE = 0x5ffcd300
TABLE_LO, TABLE_HI = 0x5f800000, 0x60000000


class Regs:
    def __init__(self, image, offsets, aperture=APERTURE):
        self.a = dis.Analysis(image)
        self.aperture = aperture
        self.off2names = collections.defaultdict(list)
        for n, (o, seg) in regmap.parse_offsets(offsets).items():
            self.off2names[o].append((seg, n))
        self.runs = self._runs()

    def _runs(self):
        p = self.a.img.payload
        runs, cur = [], []
        for off in range(0, len(p) - 3, 4):
            v, = struct.unpack_from('<I', p, off)
            if TABLE_LO <= v < TABLE_HI and v % 4 == 0:
                cur.append((self.a.img.va(off), v))
            else:
                if len(cur) >= 8:
                    runs.append(cur)
                cur = []
        if len(cur) >= 8:
            runs.append(cur)
        return runs

    def name(self, addr, prefer_seg=None):
        if addr is None or addr < self.aperture or (addr - self.aperture) % 4:
            return None
        cands = self.off2names.get((addr - self.aperture) // 4)
        if not cands:
            return None
        names = [n for seg, n in cands]
        return '/'.join(sorted(names, key=lambda n: (0 if n.startswith('DMCUB') else 1, len(n))))

    # ---- light abstract interpretation to resolve register accesses ------
    def accesses(self):
        """yield (func_va, insn_va, 'R'|'W', reg_addr, name)"""
        tablebases = set()
        for r in self.runs:
            tablebases.add(r[0][0])
        out = []
        for fva in sorted(self.a.funcs):
            # val[reg] = ('lit', v) | ('regptr', addr)
            val = {}
            for x in self.a.dis(fva):
                parts = [p.strip() for p in x.op_str.split(',')]
                m = x.mnemonic
                if m == 'l32r':
                    v = self.a.lit(x)
                    val[parts[0]] = ('lit', v)
                    continue
                if m in ('mov', 'mov.n') and len(parts) == 2:
                    if parts[1] in val:
                        val[parts[0]] = val[parts[1]]
                    else:
                        val.pop(parts[0], None)
                    continue
                if m in ('l32i', 'l32i.n') and len(parts) == 3:
                    dst, base = parts[0], parts[1]
                    try:
                        off = int(parts[2], 0)
                    except ValueError:
                        off = None
                    b = val.get(base)
                    if b and b[0] == 'lit' and b[1] is not None and off is not None:
                        w = self.a.img.word(b[1] + off)
                        if w is not None and TABLE_LO <= w < TABLE_HI:
                            val[dst] = ('regptr', w)
                            continue
                    if b and b[0] == 'regptr' and off == 0:
                        out.append((fva, x.address, 'R', b[1], self.name(b[1])))
                    val.pop(dst, None)
                    continue
                if m in ('s32i', 's32i.n') and len(parts) == 3:
                    base = parts[1]
                    try:
                        off = int(parts[2], 0)
                    except ValueError:
                        off = None
                    b = val.get(base)
                    if b and b[0] == 'regptr' and off == 0:
                        out.append((fva, x.address, 'W', b[1], self.name(b[1])))
                    continue
                # anything else kills the destination
                if parts and re.fullmatch(r'a\d+', parts[0]):
                    val.pop(parts[0], None)
        return out


def vote(r, offsets):
    """Aperture evidence with controls: score candidates by named-fraction and
    block-prefix coherence of neighbours in the register tables."""
    vals = [v for run in r.runs for _, v in run]
    votes = collections.Counter()
    offs = set(r.off2names)
    for v in vals[:800]:
        for o in offs:
            votes[v - 4 * o] += 1
    cands = [a for a, _ in votes.most_common(200)]

    def prefix(n):
        return n.split('_')[0]

    def score(A):
        named = coh = 0
        for run in r.runs:
            prev = None
            for _, v in run:
                if v < A or (v - A) % 4:
                    prev = None
                    continue
                c = r.off2names.get((v - A) // 4)
                if not c:
                    prev = None
                    continue
                named += 1
                pfx = set(prefix(n) for _, n in c)
                if prev and prev & pfx:
                    coh += 1
                prev = pfx
        return named, coh

    rows = sorted(((score(A), A) for A in cands), reverse=True)
    print('register-address values in tables: %d (%d distinct)' % (len(vals), len(set(vals))))
    for (named, coh), A in rows[:6]:
        print('  aperture 0x%08x  named %d/%d (%.0f%%)  coherent-neighbour pairs %d%s' % (
            A, named, len(vals), 100.0 * named / len(vals), coh,
            '   <= used' if A == r.aperture else ''))


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('image')
    ap.add_argument('offsets')
    ap.add_argument('cmd', choices=['table', 'access', 'vote'])
    ap.add_argument('--grep', default=None)
    ap.add_argument('--aperture', default=hex(APERTURE))
    args = ap.parse_args()

    r = Regs(args.image, args.offsets, int(args.aperture, 0))
    if args.cmd == 'vote':
        vote(r, args.offsets)
        return
    if args.cmd == 'table':
        for run in r.runs:
            print('=== table @VA 0x%08x  n=%d' % (run[0][0], len(run)))
            for va, v in run:
                n = r.name(v)
                if args.grep and not (n and re.search(args.grep, n)):
                    continue
                print('  [0x%03x] 0x%08x  %s' % (va - run[0][0], v, n or '?'))
        return
    acc = r.accesses()
    byreg = collections.defaultdict(list)
    for fva, ia, rw, addr, name in acc:
        byreg[(addr, name)].append((fva, ia, rw))
    print('resolved register accesses: %d (%d distinct registers)' % (len(acc), len(byreg)))
    for (addr, name), uses in sorted(byreg.items(), key=lambda kv: (kv[0][1] or '', kv[0][0])):
        if args.grep and not (name and re.search(args.grep, name)):
            continue
        rd = sum(1 for _, _, rw in uses if rw == 'R')
        wr = len(uses) - rd
        print('  0x%08x %-44s R=%-3d W=%-3d  %s' % (
            addr, name or '?', rd, wr,
            ' '.join(sorted(set('0x%08x' % f for f, _, _ in uses)))))


if __name__ == '__main__':
    main()