#!/usr/bin/env python3
"""Recover the DMCUB core's MMIO aperture, then name every register the firmware
touches, using AMD's own published offsets (`dcn_3_2_0_offset.h`).

Method (no assumed constant anywhere):
  1. Collect every `l32r` literal reachable from a code-region function.
  2. AMD register headers give a dword offset plus a segment id (BASE_IDX) per
     register. If the DMCUB sees registers through a linear aperture, then
         literal = APERTURE(seg) + 4 * dword_offset
     so each (literal, register) pair votes for one aperture value. A real
     aperture shows up as a spike far above the combinatorial noise floor; if
     there is no spike, the mapping is not linear and we say so.
  3. With the winning aperture, translate every literal back to a register name.

Usage:
  dmcub_regmap.py fw/dcn_3_2_0_dmcub.bin ref/dcn_3_2_0_offset.h [--seg N] [--grep PAT]
"""
import argparse
import collections
import re
import sys
import importlib.util
import os

HERE = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location('dmcub_cfg', os.path.join(HERE, 'dmcub_cfg.py'))
cfg = importlib.util.module_from_spec(spec)
spec.loader.exec_module(cfg)


def parse_offsets(path):
    """name -> (dword_offset, base_idx)"""
    off = {}
    idx = {}
    re_off = re.compile(r'^#define\s+reg(\w+)\s+(0x[0-9a-fA-F]+|\d+)\s*$')
    re_idx = re.compile(r'^#define\s+reg(\w+)_BASE_IDX\s+(\d+)\s*$')
    for line in open(path):
        m = re_idx.match(line)
        if m:
            idx[m.group(1)] = int(m.group(2))
            continue
        m = re_off.match(line)
        if m:
            off[m.group(1)] = int(m.group(2), 0)
    return {n: (o, idx.get(n)) for n, o in off.items()}


def collect_literals(img, code_lo=None):
    """(literal value) -> Counter of referencing function VAs, code region only."""
    all_ents = img.entries()
    if code_lo is None:
        code_lo = img.code_start(all_ents)
    ents = [e for e in all_ents if e >= code_lo]
    refs = collections.defaultdict(collections.Counter)
    for i, s in enumerate(ents):
        e = ents[i + 1] if i + 1 < len(ents) else len(img.payload)
        fva = img.va(s)
        for ins in img.md.disasm(img.payload[s:e], fva):
            if ins.mnemonic != 'l32r':
                continue
            try:
                la = int(ins.op_str.split()[-1], 0)
            except ValueError:
                continue
            v = img.word(la)
            if v is not None:
                refs[v][fva] += 1
    return ents, refs


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('image')
    ap.add_argument('offsets')
    ap.add_argument('--seg', type=int, default=None, help='only this BASE_IDX')
    ap.add_argument('--grep', default=None, help='report only registers matching regex')
    ap.add_argument('--aperture', default=None, help='skip the vote, force this aperture')
    ap.add_argument('--range', default='0x40000000:0x60000000',
                    help='literal window treated as candidate MMIO addresses')
    args = ap.parse_args()

    lo, hi = (int(x, 0) for x in args.range.split(':'))
    regs = parse_offsets(args.offsets)
    img = cfg.Image(args.image)
    ents, refs = collect_literals(img)
    print('%s: %d code functions, %d distinct literals' % (
        os.path.basename(args.image), len(ents), len(refs)))

    cands = [v for v in refs if lo <= v < hi]
    print('literals in candidate MMIO window 0x%x..0x%x: %d distinct, %d refs' % (
        lo, hi, len(cands), sum(sum(refs[v].values()) for v in cands)))

    by_seg = collections.defaultdict(list)
    for name, (off, seg) in regs.items():
        if seg is not None and (args.seg is None or seg == args.seg):
            by_seg[seg].append((name, off))

    if args.aperture:
        best = {args.seg or 0: int(args.aperture, 0)}
    else:
        best = {}
        for seg, lst in sorted(by_seg.items()):
            votes = collections.Counter()
            offs = set(o for _, o in lst)
            for v in cands:
                for o in offs:
                    a = v - 4 * o
                    if a >= 0:
                        votes[a] += 1
            if not votes:
                continue
            top = votes.most_common(4)
            # noise floor estimate: median vote count among top 1000
            floor = sorted(votes.values())[-1000] if len(votes) > 1000 else 1
            print('seg %d: %d registers; top apertures %s (noise floor ~%d)' % (
                seg, len(lst), ' '.join('0x%08x:%d' % (a, n) for a, n in top), floor))
            if top and top[0][1] >= max(4, 3 * floor):
                best[seg] = top[0][0]

    if not best:
        print('\nNO linear aperture found -- do not infer register access from literals.')
        return

    for seg, ap_base in sorted(best.items()):
        print('\n=== seg %d aperture 0x%08x : named register touches' % (seg, ap_base))
        byoff = {}
        for name, off in by_seg[seg]:
            byoff.setdefault(off, []).append(name)
        named = 0
        rows = []
        for v in cands:
            if (v - ap_base) % 4:
                continue
            off = (v - ap_base) // 4
            if off in byoff:
                names = [n for n in byoff[off] if not args.grep or re.search(args.grep, n)]
                if not names:
                    continue
                named += 1
                rows.append((v, '/'.join(sorted(names)), refs[v]))
        rows.sort()
        for v, name, fns in rows:
            print('  0x%08x  %-46s  %d ref(s): %s' % (
                v, name, sum(fns.values()),
                ' '.join('0x%08x' % f for f in sorted(fns))))
        print('  -> %d of %d candidate literals resolve to a seg-%d register name' % (
            named, len(cands), seg))


if __name__ == '__main__':
    main()