#!/usr/bin/env python3
"""Recover the DMCUB address map from the image itself, without trusting the
driver headers: decode every Xtensa `l32r` and collect the literal pool it
resolves to.

`l32r at, label` is PC-relative and *backwards only*:
    literal_addr = ((PC + 3) & ~3) + (sign_extend16(imm16) << 2)
with imm16 always negative, so the arithmetic works entirely in file-offset
space -- no load base assumption is required. That makes two things possible:

  1. A validator for the ISA call: on a correct decode nearly every l32r must
     resolve *inside* the image (and 4-byte aligned by construction); on a wrong
     decode, targets scatter out of range.
  2. The address map: the literal values are the absolute addresses the firmware
     uses -- code pointers, data/stack bases, and the register aperture.

Usage: dmcub_map.py fw/dcn_3_2_0_dmcub.bin [payload_off] [payload_end]
"""
import collections
import struct
import sys


def l32r_sites(data, lo, hi):
    """Yield (site_off, literal_off) for every op0==1 slot whose literal lands
    inside [lo, hi). Linear sweep -- deliberately not a recursive-descent
    disassembly, so counts are an upper bound and the in-range *ratio* is the
    signal."""
    for off in range(lo, hi - 3):
        if data[off] & 0xF != 0x1:
            continue
        imm16 = struct.unpack_from('<H', data, off + 1)[0]
        lit = ((off + 3) & ~3) + ((imm16 - 0x10000) << 2)
        yield off, lit


def main():
    path = sys.argv[1]
    data = open(path, 'rb').read()
    lo = int(sys.argv[2], 0) if len(sys.argv) > 2 else 0x100
    hi = int(sys.argv[3], 0) if len(sys.argv) > 3 else len(data)

    total = inrange = 0
    lits = collections.Counter()
    for off, lit in l32r_sites(data, lo, hi):
        total += 1
        if lo <= lit < hi - 3:
            inrange += 1
            lits[lit] += 1

    print('%s payload 0x%x..0x%x' % (path, lo, hi))
    print('l32r slots (linear sweep): %d   literal in-image: %d (%.1f%%)' % (
        total, inrange, 100.0 * inrange / (total or 1)))
    print('distinct literal addresses: %d' % len(lits))

    vals = collections.Counter()
    for lit, n in lits.items():
        v, = struct.unpack_from('<I', data, lit)
        vals[v] += n

    top = collections.Counter()
    for v, n in vals.items():
        top[v >> 24] += n
    print('\nliteral values by high byte (0x??000000):')
    for hb, n in top.most_common(16):
        print('   0x%02x______  %6d refs' % (hb, n))

    print('\nmost-referenced literal values:')
    for v, n in vals.most_common(25):
        print('   0x%08x  %4d refs' % (v, n))


if __name__ == '__main__':
    main()