#!/usr/bin/env python3
"""x86-32 displacement cross-referencer for the SOTS1 exe.

Why this exists
---------------
Ghidra indexes immediate operands but NOT ModRM displacements. On an MSVC C++
binary `lea reg,[reg+disp]` is *the* idiom for taking the address of a member --
it is how every std::vector / std::string / embedded sub-object is passed to a
method or ctor. `find-constant-uses 0x274` therefore cannot see the append site
for a vector member at +0x274. This tool answers the question Ghidra can't:

    "what code takes the address of, or accesses, offset N off some object?"

How it works
------------
x86 is variable-length and not self-synchronising, so a naive byte scan for
"8D 8E <disp32>" produces confident garbage. Instead we run a full instruction
*length* decoder (prefixes / 1-,2-,3-byte opcodes / ModRM / SIB / disp / imm)
seeded from Ghidra's 41k function starts, so every decode begins on a real
instruction boundary. `--brute` runs the naive scan too, purely so the
false-positive rate of the naive method can be *measured* rather than guessed.

Correctness notes that matter (each of these is a silent-garbage source):
  * mod=0,rm=5      -> disp32 with NO base: an absolute global address, not a
                       member offset. Excluded from member queries by default.
  * mod=0,rm=4,
    sib.base=5      -> disp32, no base, index only. Same: excluded.
  * mod=1           -> disp8 SIGN-EXTENDED. -0x08 must not be reported as 0xf8.
  * mod=3           -> register operand, no memory, no displacement at all.
  * 0x67 addr-size  -> 16-bit ModRM, entirely different layout. Skipped, flagged.

Usage
-----
  uv run python3 tools/x86disp.py build                 # decode + cache index
  uv run python3 tools/x86disp.py query 0x274           # who touches +0x274
  uv run python3 tools/x86disp.py query 0x29c --lea     # lea only
  uv run python3 tools/x86disp.py func 0x0080db00       # dump one function
  uv run python3 tools/x86disp.py stats
"""
import bisect
import json
import os
import struct
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.dirname(HERE)
EXE = os.path.join(REPO, "dumps", "sots.exe")
FUNCS = os.path.join(REPO, "dumps", "functions.json")
INDEX = os.path.join(REPO, "dumps", "dispindex.json")

R32 = ["eax", "ecx", "edx", "ebx", "esp", "ebp", "esi", "edi"]

# ---------------------------------------------------------------- opcode maps
# value: (has_modrm, imm_kind)
#   imm kinds: 0 none | 'b' 1 | 'w' 2 | 'd' 4 | 'z' 2-if-66-else-4
#              'p' far ptr 6 (4 with 66) | 'a' moffs (addr-size) | 'e' enter 3
#              'g6' F6 group (b if reg<2) | 'g7' F7 group (z if reg<2)
_ONE = {}


def _fill(rng, modrm, imm):
    for o in rng:
        _ONE[o] = (modrm, imm)


# 00..3F: the eight ALU ops, each  /r /r /r /r  AL,ib  eAX,iz  + 2 seg ops
for _base in (0x00, 0x08, 0x10, 0x18, 0x20, 0x28, 0x30, 0x38):
    _fill(range(_base, _base + 4), True, 0)
    _ONE[_base + 4] = (False, "b")
    _ONE[_base + 5] = (False, "z")
    _ONE[_base + 6] = (False, 0)          # PUSH seg / prefix-or-ascii-adjust
    _ONE[_base + 7] = (False, 0)
_fill(range(0x40, 0x60), False, 0)        # INC/DEC/PUSH/POP r32
_fill([0x60, 0x61], False, 0)             # PUSHA/POPA
_ONE[0x62] = (True, 0)                    # BOUND
_ONE[0x63] = (True, 0)                    # ARPL
_ONE[0x68] = (False, "z")                 # PUSH iz
_ONE[0x69] = (True, "z")                  # IMUL r,Ev,iz
_ONE[0x6A] = (False, "b")                 # PUSH ib
_ONE[0x6B] = (True, "b")                  # IMUL r,Ev,ib
_fill(range(0x6C, 0x70), False, 0)        # INS/OUTS
_fill(range(0x70, 0x80), False, "b")      # Jcc rel8
_ONE[0x80] = (True, "b")
_ONE[0x81] = (True, "z")
_ONE[0x82] = (True, "b")
_ONE[0x83] = (True, "b")
_fill(range(0x84, 0x90), True, 0)         # TEST/XCHG/MOV/MOV-seg/LEA/POP Ev
_fill(range(0x90, 0x9A), False, 0)        # NOP/XCHG/CWDE/CDQ
_ONE[0x9A] = (False, "p")                 # CALLF
_fill(range(0x9B, 0xA0), False, 0)
_fill(range(0xA0, 0xA4), False, "a")      # MOV moffs
_fill(range(0xA4, 0xA8), False, 0)        # MOVS/CMPS
_ONE[0xA8] = (False, "b")
_ONE[0xA9] = (False, "z")
_fill(range(0xAA, 0xB0), False, 0)        # STOS/LODS/SCAS
_fill(range(0xB0, 0xB8), False, "b")      # MOV r8,ib
_fill(range(0xB8, 0xC0), False, "z")      # MOV r32,iz
_ONE[0xC0] = (True, "b")
_ONE[0xC1] = (True, "b")
_ONE[0xC2] = (False, "w")                 # RET imm16
_ONE[0xC3] = (False, 0)
_ONE[0xC4] = (True, 0)                    # LES
_ONE[0xC5] = (True, 0)                    # LDS
_ONE[0xC6] = (True, "b")                  # MOV Eb,Ib
_ONE[0xC7] = (True, "z")                  # MOV Ev,Iz
_ONE[0xC8] = (False, "e")                 # ENTER iw,ib
_ONE[0xC9] = (False, 0)
_ONE[0xCA] = (False, "w")                 # RETF imm16
_fill([0xCB, 0xCC], False, 0)
_ONE[0xCD] = (False, "b")                 # INT ib
_fill([0xCE, 0xCF], False, 0)
_fill(range(0xD0, 0xD4), True, 0)         # shift group by 1 / by CL
_ONE[0xD4] = (False, "b")                 # AAM
_ONE[0xD5] = (False, "b")                 # AAD
_fill([0xD6, 0xD7], False, 0)
_fill(range(0xD8, 0xE0), True, 0)         # x87 -- always ModRM
_fill(range(0xE0, 0xE4), False, "b")      # LOOP*/JECXZ
_fill(range(0xE4, 0xE8), False, "b")      # IN/OUT ib
_ONE[0xE8] = (False, "z")                 # CALL rel32
_ONE[0xE9] = (False, "z")                 # JMP rel32
_ONE[0xEA] = (False, "p")                 # JMPF
_ONE[0xEB] = (False, "b")                 # JMP rel8
_fill(range(0xEC, 0xF0), False, 0)        # IN/OUT DX
_fill(range(0xF0, 0xF6), False, 0)        # LOCK/INT1/REP*/HLT/CMC
_ONE[0xF6] = (True, "g6")
_ONE[0xF7] = (True, "g7")
_fill(range(0xF8, 0xFE), False, 0)
_ONE[0xFE] = (True, 0)
_ONE[0xFF] = (True, 0)

_TWO = {}


def _fill2(rng, modrm, imm):
    for o in rng:
        _TWO[o] = (modrm, imm)


_fill2(range(0x00, 0x05), True, 0)
_fill2(range(0x05, 0x0D), False, 0)
_TWO[0x0D] = (True, 0)
_TWO[0x0E] = (False, 0)
_TWO[0x0F] = (True, "b")                  # 3DNow!
_fill2(range(0x10, 0x18), True, 0)
_fill2(range(0x18, 0x20), True, 0)        # hint-NOP / prefetch
_fill2(range(0x20, 0x25), True, 0)
_fill2(range(0x28, 0x30), True, 0)
_fill2(range(0x30, 0x38), False, 0)
_fill2(range(0x40, 0x50), True, 0)        # CMOVcc
_fill2(range(0x50, 0x70), True, 0)        # SSE/MMX
_TWO[0x70] = (True, "b")
_fill2(range(0x71, 0x74), True, "b")
_fill2(range(0x74, 0x77), True, 0)
_TWO[0x77] = (False, 0)                   # EMMS
_fill2(range(0x78, 0x80), True, 0)
_fill2(range(0x80, 0x90), False, "z")     # Jcc rel32
_fill2(range(0x90, 0xA0), True, 0)        # SETcc
_fill2([0xA0, 0xA1, 0xA2], False, 0)
_TWO[0xA3] = (True, 0)                    # BT
_TWO[0xA4] = (True, "b")                  # SHLD ib
_TWO[0xA5] = (True, 0)                    # SHLD CL
_fill2([0xA8, 0xA9, 0xAA], False, 0)
_TWO[0xAB] = (True, 0)                    # BTS
_TWO[0xAC] = (True, "b")                  # SHRD ib
_fill2([0xAD, 0xAE, 0xAF], True, 0)
_fill2(range(0xB0, 0xBA), True, 0)
_TWO[0xBA] = (True, "b")                  # group8 BT/BTS/BTR/BTC ib
_fill2(range(0xBB, 0xC0), True, 0)
_fill2([0xC0, 0xC1], True, 0)             # XADD
_TWO[0xC2] = (True, "b")                  # CMPPS
_TWO[0xC3] = (True, 0)                    # MOVNTI
_fill2([0xC4, 0xC5, 0xC6], True, "b")
_TWO[0xC7] = (True, 0)                    # group9 CMPXCHG8B
_fill2(range(0xC8, 0xD0), False, 0)       # BSWAP
_fill2(range(0xD0, 0x100), True, 0)       # MMX/SSE bulk

# minimal mnemonics -- enough to read a report, not a full disassembler
_ALU = ["add", "or", "adc", "sbb", "and", "sub", "xor", "cmp"]
_G1 = _ALU
_G5 = ["inc", "dec", "call", "callf", "jmp", "jmpf", "push", "?"]
_G3 = ["test", "test", "not", "neg", "mul", "imul", "div", "idiv"]
_SHIFT = ["rol", "ror", "rcl", "rcr", "shl", "shr", "shl", "sar"]


def mnemonic(op2, op, reg):
    """Best-effort mnemonic. op2 True => the opcode was 0F-escaped."""
    if op2:
        if 0x10 <= op <= 0x17 or 0x28 <= op <= 0x2F or 0x51 <= op <= 0x5F:
            return "sse"
        if 0x40 <= op <= 0x4F:
            return "cmov"
        if 0x90 <= op <= 0x9F:
            return "setcc"
        if op in (0xB6, 0xB7):
            return "movzx"
        if op in (0xBE, 0xBF):
            return "movsx"
        if op == 0xAF:
            return "imul"
        if op in (0x6E, 0x6F, 0x7E, 0x7F, 0xD6):
            return "movq/movd"
        return f"0f{op:02x}"
    if op < 0x40 and (op & 7) < 6:
        return _ALU[op >> 3]
    if op in (0x88, 0x89, 0x8A, 0x8B, 0xC6, 0xC7):
        return "mov"
    if op == 0x8D:
        return "lea"
    if op in (0x84, 0x85):
        return "test"
    if op in (0x86, 0x87):
        return "xchg"
    if op in (0x80, 0x81, 0x82, 0x83):
        return _G1[reg]
    if op in (0xC0, 0xC1, 0xD0, 0xD1, 0xD2, 0xD3):
        return _SHIFT[reg]
    if op in (0xF6, 0xF7):
        return _G3[reg]
    if op == 0xFF:
        return _G5[reg]
    if op == 0xFE:
        return ["inc", "dec"][reg] if reg < 2 else "?"
    if op == 0x8F:
        return "pop"
    if op in (0x69, 0x6B):
        return "imul"
    if 0xD8 <= op <= 0xDF:
        return "x87"
    if op == 0x62:
        return "bound"
    return f"{op:02x}"


class Desync(Exception):
    pass


def decode(buf, i, end):
    """Decode one instruction at buf[i]. Returns (length, info|None).

    info = dict(mnem, base, index, scale, disp, dispsize, modrm_reg, is_lea)
    for instructions with a memory operand carrying a displacement; None
    otherwise. Raises Desync on an unknown/invalid opcode.
    """
    start = i
    opsize66 = False
    addr67 = False
    while i < end:
        b = buf[i]
        if b == 0x66:
            opsize66 = True
            i += 1
        elif b == 0x67:
            addr67 = True
            i += 1
        elif b in (0xF0, 0xF2, 0xF3, 0x2E, 0x36, 0x3E, 0x26, 0x64, 0x65):
            i += 1
        else:
            break
    else:
        raise Desync("prefix run to end")
    if i - start > 14:
        raise Desync("prefix flood")

    op = buf[i]
    i += 1
    op2 = op3 = False
    if op == 0x0F:
        if i >= end:
            raise Desync("truncated 0f")
        op = buf[i]
        i += 1
        op2 = True
        if op in (0x38, 0x3A):
            three_imm = "b" if op == 0x3A else 0
            if i >= end:
                raise Desync("truncated 0f3x")
            op = buf[i]
            i += 1
            op3 = True
            has_modrm, imm = True, three_imm
        else:
            ent = _TWO.get(op)
            if ent is None:
                raise Desync(f"unknown 0f{op:02x}")
            has_modrm, imm = ent
    else:
        ent = _ONE.get(op)
        if ent is None:
            raise Desync(f"unknown {op:02x}")
        has_modrm, imm = ent

    info = None
    if has_modrm:
        if i >= end:
            raise Desync("truncated modrm")
        modrm = buf[i]
        i += 1
        mod = modrm >> 6
        reg = (modrm >> 3) & 7
        rm = modrm & 7
        if imm == "g6":
            imm = "b" if reg < 2 else 0
        elif imm == "g7":
            imm = "z" if reg < 2 else 0

        if mod != 3:
            if addr67:
                # 16-bit ModRM: different table entirely. Rare in MSVC code and
                # never the member-address idiom -- skip rather than mis-decode.
                raise Desync("16-bit addressing (0x67)")
            base = index = None
            scale = 1
            if rm == 4:
                if i >= end:
                    raise Desync("truncated sib")
                sib = buf[i]
                i += 1
                scale = 1 << (sib >> 6)
                idx = (sib >> 3) & 7
                bse = sib & 7
                index = None if idx == 4 else R32[idx]
                if bse == 5 and mod == 0:
                    base = None          # disp32 absolute + index
                else:
                    base = R32[bse]
            elif rm == 5 and mod == 0:
                base = None              # disp32 absolute
            else:
                base = R32[rm]

            disp = 0
            dispsize = 0
            if mod == 1:
                if i >= end:
                    raise Desync("truncated disp8")
                disp = struct.unpack_from("<b", buf, i)[0]   # SIGN-EXTENDED
                dispsize = 1
                i += 1
            elif mod == 2 or base is None:
                if i + 4 > end:
                    raise Desync("truncated disp32")
                disp = struct.unpack_from("<i", buf, i)[0]
                dispsize = 4
                i += 4
            if dispsize:
                info = {
                    "mnem": mnemonic(op2, op, reg), "base": base,
                    "index": index, "scale": scale, "disp": disp,
                    "dispsize": dispsize, "reg": R32[reg],
                    "lea": (not op2 and op == 0x8D),
                }
    else:
        reg = 0

    n = 0
    if imm == "b":
        n = 1
    elif imm == "w":
        n = 2
    elif imm == "d":
        n = 4
    elif imm == "z":
        n = 2 if opsize66 else 4
    elif imm == "p":
        n = 4 if opsize66 else 6
    elif imm == "a":
        n = 2 if addr67 else 4
    elif imm == "e":
        n = 3
    i += n
    if i > end:
        raise Desync("truncated imm")
    if i == start:
        raise Desync("zero length")
    return i - start, info


# ------------------------------------------------------------------- PE / IO
def load_pe(path):
    """Return (image_base, [(va_start, va_end, bytes, name)]) for exec sections."""
    data = open(path, "rb").read()
    pe = struct.unpack_from("<I", data, 0x3C)[0]
    assert data[pe:pe + 4] == b"PE\0\0", "not a PE"
    nsec = struct.unpack_from("<H", data, pe + 6)[0]
    optsz = struct.unpack_from("<H", data, pe + 20)[0]
    base = struct.unpack_from("<I", data, pe + 24 + 28)[0]
    secs = []
    off = pe + 24 + optsz
    for k in range(nsec):
        s = off + k * 40
        name = data[s:s + 8].rstrip(b"\0").decode("latin1")
        vsize, vaddr, rsize, raddr = struct.unpack_from("<IIII", data, s + 8)
        chars = struct.unpack_from("<I", data, s + 36)[0]
        if not (chars & 0x20000000):        # IMAGE_SCN_MEM_EXECUTE
            continue
        n = min(vsize, rsize) if vsize else rsize
        secs.append((base + vaddr, base + vaddr + n, data[raddr:raddr + n], name))
    return base, secs


def load_funcs():
    with open(FUNCS) as fh:
        raw = json.load(fh)
    fl = sorted((int(a, 16), n, sz) for a, (n, sz) in raw.items())
    return fl


# --------------------------------------------------------------------- build
def build():
    _, secs = load_pe(EXE)
    funcs = load_funcs()
    starts = [f[0] for f in funcs]
    sites = []          # [va, disp, base, index, scale, mnem, lea, funcidx, hex]
    covered = 0
    total_code = sum(e - s for s, e, _, _ in secs)
    desyncs = 0
    decoded_ins = 0

    for fi, (fva, fname, fsz) in enumerate(funcs):
        sec = next((s for s in secs if s[0] <= fva < s[1]), None)
        if sec is None or fsz <= 0:
            continue
        sva, eva, buf, _ = sec
        # Sweep to the NEXT function start, not fva+fsz: Ghidra's sizeInBytes
        # understates ~11% of bodies (it clips valid epilogues mid-instruction).
        # Sites past fsz are still attributed to this function but flagged, so
        # a caller can tell body-proper from possible inter-function padding.
        j = bisect.bisect_right(starts, fva)
        limit = min(starts[j] if j < len(starts) else eva, eva)
        body_end = fva + fsz
        i = fva - sva
        end = limit - sva
        while i < end:
            try:
                ln, info = decode(buf, i, end)
            except Desync:
                desyncs += 1
                break
            decoded_ins += 1
            if info and info["base"] is not None and info["disp"] != 0:
                va = sva + i
                sites.append([va, info["disp"], info["base"], info["index"],
                              info["scale"], info["mnem"], info["lea"], fi,
                              buf[i:i + min(ln, 10)].hex(),
                              1 if va < body_end else 0, info["reg"]])
            i += ln
        covered += max(0, i - (fva - sva))

    idx = {}
    for k, s in enumerate(sites):
        idx.setdefault(str(s[1]), []).append(k)
    out = {
        "sites": sites,
        "byDisp": idx,
        "funcs": [[f[0], f[1], f[2]] for f in funcs],
        "stats": {"functions": len(funcs), "instructions": decoded_ins,
                  "sites": len(sites), "desyncs": desyncs,
                  "bytesCovered": covered, "codeBytes": total_code},
    }
    with open(INDEX, "w") as fh:
        json.dump(out, fh)
    st = out["stats"]
    print(f"functions swept : {st['functions']}")
    print(f"instructions    : {st['instructions']}")
    print(f"disp sites      : {st['sites']}")
    print(f"desyncs         : {st['desyncs']} "
          f"({100.0 * st['desyncs'] / st['functions']:.2f}% of functions)")
    print(f"code coverage   : {st['bytesCovered']}/{st['codeBytes']} "
          f"({100.0 * st['bytesCovered'] / st['codeBytes']:.1f}%)")


# --------------------------------------------------------------------- query
def load_index():
    if not os.path.exists(INDEX):
        sys.exit("no index; run: uv run python3 tools/x86disp.py build")
    with open(INDEX) as fh:
        return json.load(fh)


THIS_REGS = ("ecx", "esi", "edi", "ebx")   # typical `this` carriers in MSVC


def fmt(site, funcs):
    va, disp, base, index, scale, mnem, lea, fi, hx, inbody, dst = site
    fva, fname, _ = funcs[fi]
    ea = f"[{base}"
    if index:
        ea += f"+{index}*{scale}"
    ea += f"{'+' if disp >= 0 else '-'}0x{abs(disp):x}]"
    txt = f"{mnem} {dst},{ea}" if lea else f"{mnem} {ea}"
    return (f"  0x{va:08x}  {txt:<32} "
            f"{fname}+0x{va - fva:x}{'' if inbody else ' [past-body]'}  ({hx})")


def query(argv):
    want = int(argv[0], 0)
    only_lea = "--lea" in argv
    only_this = "--this" in argv
    show_all = "--all" in argv
    d = load_index()
    funcs = d["funcs"]
    keys = d["byDisp"].get(str(want), [])
    hits = [d["sites"][k] for k in keys]
    if only_lea:
        hits = [h for h in hits if h[6]]
    if only_this:
        hits = [h for h in hits if h[2] in THIS_REGS]
    hits.sort(key=lambda h: h[0])

    ranked, other = [], []
    for h in hits:
        (ranked if h[2] in THIS_REGS else other).append(h)
    print(f"displacement 0x{want:x} ({want}): {len(hits)} site(s)")
    print(f"\n== base is a likely `this` ({'/'.join(THIS_REGS)}): {len(ranked)} ==")
    byfn = {}
    for h in ranked:
        byfn.setdefault(h[7], []).append(h)
    for fi in sorted(byfn, key=lambda f: funcs[f][0]):
        print(f" {funcs[fi][1]}  @0x{funcs[fi][0]:08x}")
        for h in byfn[fi]:
            print(fmt(h, funcs))
    if other:
        print(f"\n== other base regs (esp/ebp = locals, eax/edx = temps): {len(other)} ==")
        if show_all:
            for h in other:
                print(fmt(h, funcs))
        else:
            cnt = {}
            for h in other:
                cnt[h[2]] = cnt.get(h[2], 0) + 1
            print("  " + ", ".join(f"{k}:{v}" for k, v in sorted(cnt.items()))
                  + "   (--all to list)")


def func_dump(argv):
    fva = int(argv[0], 0)
    d = load_index()
    funcs = d["funcs"]
    fi = next((i for i, f in enumerate(funcs) if f[0] == fva), None)
    if fi is None:
        sys.exit(f"no function at 0x{fva:08x}")
    print(f"{funcs[fi][1]} @ 0x{fva:08x} size {funcs[fi][2]}")
    for s in d["sites"]:
        if s[7] == fi:
            print(fmt(s, funcs))


def stats():
    d = load_index()
    st = d["stats"]
    for k, v in st.items():
        print(f"{k:16}: {v}")
    hist = {}
    for s in d["sites"]:
        if s[2] in THIS_REGS:
            hist[s[1]] = hist.get(s[1], 0) + 1
    print(f"\ndistinct displacements off ecx/esi/edi/ebx: {len(hist)}")
    top = sorted(hist.items(), key=lambda kv: -kv[1])[:15]
    print("most common: " + ", ".join(f"0x{k:x}({v})" for k, v in top))


_STORE = {0x89, 0x01, 0x29, 0x39, 0x09, 0x21, 0x31, 0x19, 0x11,
          0x88, 0x00, 0x08, 0x10, 0x18, 0x20, 0x28, 0x30, 0x38}


def _reg_form(raw):
    """Render the mod=3 (register-to-register) forms, which are exactly the
    ones that carry the arithmetic you care about when recovering a struct
    stride: `imul ecx,ecx,0x2c`, `sar edx,3`, `sub ecx,edi`."""
    p = 0
    while p < len(raw) and raw[p] in (0x66, 0x67, 0xF0, 0xF2, 0xF3,
                                      0x2E, 0x36, 0x3E, 0x26, 0x64, 0x65):
        p += 1
    op = raw[p]
    if op == 0x0F or p + 1 >= len(raw):
        return mnemonic(op == 0x0F, raw[p + 1] if op == 0x0F else op, 0)
    m = raw[p + 1]
    mod, reg, rm = m >> 6, (m >> 3) & 7, m & 7
    if mod != 3:
        # memory operand: render it and show which register is the other half,
        # so a load reads `mov ecx,[edi+4]` not a bare `mov`.
        q = p + 2
        ea = "["
        if rm == 4:
            sib = raw[q]; q += 1
            bse, ix = sib & 7, (sib >> 3) & 7
            ea += ("" if (bse == 5 and mod == 0) else R32[bse])
            if ix != 4:
                ea += f"+{R32[ix]}*{1 << (sib >> 6)}"
        elif not (rm == 5 and mod == 0):
            ea += R32[rm]
        if mod == 1 and q < len(raw):
            dv = struct.unpack_from("<b", raw, q)[0]
            ea += f"{'+' if dv >= 0 else '-'}0x{abs(dv):x}"
        elif mod == 2 and q + 4 <= len(raw):
            dv = struct.unpack_from("<i", raw, q)[0]
            ea += f"{'+' if dv >= 0 else '-'}0x{abs(dv):x}"
        ea += "]"
        nm = mnemonic(False, op, reg)
        if op in _STORE:
            return f"{nm} {ea},{R32[reg]}"
        if op in (0x80, 0x81, 0x83, 0xC6, 0xC7, 0xF6, 0xF7, 0xFF, 0xFE,
                  0xC0, 0xC1, 0xD0, 0xD1, 0xD2, 0xD3) or op >= 0xD8:
            return f"{nm} {ea}"
        return f"{nm} {R32[reg]},{ea}"
    r, x = R32[reg], R32[rm]
    if op in (0x89, 0x01, 0x29, 0x39, 0x09, 0x21, 0x31, 0x19, 0x11, 0x85, 0x87):
        a, b = x, r                       # Ev,Gv  -> dst is rm
    else:
        a, b = r, x                       # Gv,Ev  -> dst is reg
    nm = mnemonic(False, op, reg)
    if op in (0xC0, 0xC1):                # shift group, imm8
        return f"{_SHIFT[reg]} {x},{raw[p + 2]}"
    if op in (0xD1, 0xD3):
        return f"{_SHIFT[reg]} {x},{'cl' if op == 0xD3 else '1'}"
    if op == 0x6B:                        # imul r,rm,imm8  <- the stride!
        return f"imul {r},{x},0x{struct.unpack_from('<b', raw, p + 2)[0] & 0xff:x}"
    if op == 0x69:
        return f"imul {r},{x},0x{struct.unpack_from('<I', raw, p + 2)[0]:x}"
    if op in (0x83, 0x81, 0x80):
        v = (raw[p + 2] if op != 0x81
             else struct.unpack_from("<I", raw, p + 2)[0])
        return f"{_G1[reg]} {x},0x{v:x}"
    if op in (0xF7, 0xF6):
        return f"{_G3[reg]} {x}"
    if op == 0xFF:
        return f"{_G5[reg]} {x}"
    return f"{nm} {a},{b}"


def dis(argv):
    """Linear disassembly window: `dis <va> [count]`. Length-accurate; the
    mnemonics are minimal but call/jmp targets are resolved to function names,
    which is what you actually need to follow a `lea this,[obj+N]; call` pair."""
    va = int(argv[0], 0)
    n = int(argv[1]) if len(argv) > 1 else 24
    _, secs = load_pe(EXE)
    d = load_index()
    fmap = {f[0]: f[1] for f in d["funcs"]}
    sec = next((s for s in secs if s[0] <= va < s[1]), None)
    if sec is None:
        sys.exit("address not in an executable section")
    sva, eva, buf, _ = sec
    i = va - sva
    for _ in range(n):
        try:
            ln, info = decode(buf, i, len(buf))
        except Desync as e:
            print(f"0x{sva + i:08x}  <desync: {e}>")
            return
        raw = buf[i:i + ln]
        cur = sva + i
        txt = ""
        # resolve the two rel32 forms and the rel8 jumps by hand
        if raw[0] in (0xE8, 0xE9) and ln == 5:
            tgt = cur + 5 + struct.unpack_from("<i", raw, 1)[0]
            txt = (f"{'call' if raw[0] == 0xE8 else 'jmp'} 0x{tgt:08x}"
                   f"  {fmap.get(tgt, '')}")
        elif raw[0] == 0x0F and 0x80 <= raw[1] <= 0x8F and ln == 6:
            tgt = cur + 6 + struct.unpack_from("<i", raw, 2)[0]
            txt = f"jcc 0x{tgt:08x}"
        elif raw[0] == 0xEB or 0x70 <= raw[0] <= 0x7F:
            tgt = cur + ln + struct.unpack_from("<b", raw, ln - 1)[0]
            txt = f"jmp/jcc 0x{tgt:08x}"
        elif 0x50 <= raw[0] <= 0x57:
            txt = f"push {R32[raw[0] - 0x50]}"
        elif 0x58 <= raw[0] <= 0x5F:
            txt = f"pop {R32[raw[0] - 0x58]}"
        elif raw[0] == 0x6A:
            txt = f"push 0x{raw[1]:x}"
        elif raw[0] == 0x68:
            txt = f"push 0x{struct.unpack_from('<I', raw, 1)[0]:x}"
        elif 0xB8 <= raw[0] <= 0xBF and ln == 5:
            txt = f"mov {R32[raw[0] - 0xB8]},0x{struct.unpack_from('<I', raw, 1)[0]:x}"
        elif raw[0] == 0xC3:
            txt = "ret"
        elif raw[0] == 0xC2:
            txt = f"ret 0x{struct.unpack_from('<H', raw, 1)[0]:x}"
        else:
            txt = _reg_form(raw)
        lbl = fmap.get(cur, "")
        print(f"0x{cur:08x}  {raw.hex():<18} {txt:<44} {lbl}")
        i += ln


def cohort(argv):
    """Rank functions by how many offsets of a KNOWN class layout they touch.

    This is the answer to the real precision problem. A bare query for 0x274
    returns ~99 sites and the tool cannot know which base register holds a
    ServerPlayer. But a function that touches 0x274 *and* 0x29c *and* 0x244
    *and* 0x254 is not doing that by coincidence -- those are ServerPlayer's
    members. Feed the offsets we have already pinned and the class's own
    methods float to the top.

      uv run python3 tools/x86disp.py cohort 0x274 --known=0x29c,0x244,0x254,...

    WARNING, learned the hard way on the ObservedTech hunt: use this as a
    RANKER, never as a hard filter. It finds fat *class methods* (the
    ServerPlayer serializer scores 50/50) but it actively hides narrow helpers.
    `RecordObservedTech` -- the actual append site -- touches only 0x274 and
    0x278 and nothing else on ServerPlayer, so every --min>=1 setting drops it.
    Read the full `query` output before trusting a cohort shortlist.
    """
    want = int(argv[0], 0)
    known = []
    minhits = 2
    for a in argv[1:]:
        if a.startswith("--known="):
            known = [int(x, 0) for x in a[8:].split(",")]
        elif a.startswith("--min="):
            minhits = int(a[6:])
    if not known:
        sys.exit("need --known=<comma-separated offsets of the same class>")
    d = load_index()
    funcs = d["funcs"]
    kset = set(known)
    touched = {}          # funcidx -> set(offset)
    target = {}           # funcidx -> [sites at `want`]
    for s in d["sites"]:
        if s[2] not in THIS_REGS:
            continue
        if s[1] in kset:
            touched.setdefault(s[7], set()).add(s[1])
        if s[1] == want:
            target.setdefault(s[7], []).append(s)
    rows = []
    for fi, sites in target.items():
        hits = touched.get(fi, set()) - {want}
        if len(hits) >= minhits:
            rows.append((len(hits), fi, sites, hits))
    rows.sort(key=lambda r: -r[0])
    print(f"functions touching 0x{want:x} AND >={minhits} other known offsets "
          f"of this class: {len(rows)} of {len(target)} candidates "
          f"({100.0 * len(rows) / max(1, len(target)):.0f}% kept)")
    for n, fi, sites, hits in rows:
        fva, fname, _ = funcs[fi]
        hx = ",".join(f"0x{h:x}" for h in sorted(hits))
        print(f"\n {fname} @0x{fva:08x}   [{n} co-hits: {hx}]")
        for s in sites:
            print(fmt(s, funcs))


# ------------------------------------------------- naive scan, for FP measure
def brute(argv):
    """Naive 'search the raw bytes' scan -- the thing you'd write without a
    decoder. Only exists so we can put a NUMBER on how wrong it is.

    `--ops 8d` (default) = lea only. `--ops 8d,89,8b,01,03,39,3b` = the wider
    scan you would actually need, since a member is read/written far more often
    than its address is taken. The wider the opcode set and the smaller the
    displacement, the worse the naive method gets -- that is the point.
    """
    want = int(argv[0], 0)
    ops = {0x8D}
    for a in argv[1:]:
        if a.startswith("--ops="):
            ops = {int(x, 16) for x in a[6:].split(",")}
    _, secs = load_pe(EXE)
    d = load_index()
    real = {s[0] for s in d["sites"] if s[1] == want}
    found = []
    d8 = want if -128 <= want <= 127 else None
    pat32 = struct.pack("<i", want)
    for sva, eva, buf, _ in secs:
        for i in range(len(buf) - 6):
            if buf[i] not in ops:
                continue
            m = buf[i + 1]
            mod, rm = m >> 6, m & 7
            if mod == 3 or mod == 0:
                continue
            k = i + 2 + (1 if rm == 4 else 0)
            if mod == 1 and d8 is not None:
                if struct.unpack_from("<b", buf, k)[0] == want:
                    found.append(sva + i)
            elif mod == 2 and buf[k:k + 4] == pat32:
                found.append(sva + i)
    fp = [a for a in found if a not in real]
    tag = ",".join(f"{o:02x}" for o in sorted(ops))
    print(f"naive byte-scan (opcodes {tag}) for disp 0x{want:x}: "
          f"{len(found)} candidate(s)")
    print(f"  real instructions at that address : {len(found) - len(fp)}")
    print(f"  NOT an instruction boundary       : {len(fp)}  "
          f"({100.0 * len(fp) / max(1, len(found)):.1f}% false positive)")
    missed = len(real) - (len(found) - len(fp))
    print(f"  real sites the naive scan MISSED  : {missed} of {len(real)}")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    cmd, rest = sys.argv[1], sys.argv[2:]
    {"build": lambda a: build(), "query": query, "func": func_dump,
     "stats": lambda a: stats(), "brute": brute, "cohort": cohort,
     "dis": dis}[cmd](rest)