# watcom_hi16.py — IDA 9.x plugin
#
# Watcom emits  load4(&a) >> 16  to read a 16-bit struct member at offsetof(a)+2,
# because a dword load + sar is cheaper than movsx. Hex-Rays renders this as
# *(int *)&s->a >> 16, naming the WRONG field. This rewrites those expressions
# to reference the actual member, when a 2-byte member exists at that offset.
#
# Install: drop in ~/.idapro/plugins/  (or %APPDATA%\Hex-Rays\IDA Pro\plugins\)
# Toggle at runtime:  import watcom_hi16; watcom_hi16.ENABLED = False

import ida_hexrays as hr
import ida_typeinf as ti
import ida_idaapi

SHIFT   = 16      # only rewrite >> 16
HALF    = 2       # target member size in bytes
ENABLED = True
VERBOSE = False

_stats = {"rewrites": 0, "funcs": 0}


# ---------------------------------------------------------------- type lookup

def _member_at(stype, off):
    """tinfo_t of the HALF-byte member at byte offset `off`, or None."""
    if stype is None:
        return None
    st = ti.tinfo_t(stype)
    st.remove_ptr_or_array()
    if not st.is_struct():
        return None
    udt = ti.udt_type_data_t()
    if not st.get_udt_details(udt):
        return None
    for i in range(udt.size()):
        m = udt[i]
        if m.offset // 8 == off and m.size // 8 == HALF:
            return ti.tinfo_t(m.type)
    return None


# ---------------------------------------------------------------- collection

class _Collector(hr.ctree_visitor_t):
    """Find rewritable nodes. Collect only — never mutate during traversal."""

    def __init__(self):
        hr.ctree_visitor_t.__init__(self, hr.CV_FAST)
        self.hits = []

    def visit_expr(self, e):
        try:
            if e.op not in (hr.cot_sshr, hr.cot_ushr):
                return 0
            if e.y.op != hr.cot_num or e.y.numval() != SHIFT:
                return 0

            x = e.x
            if x.op != hr.cot_ptr or x.type.get_size() != 4:
                return 0
            x = x.x
            while x.op == hr.cot_cast:
                x = x.x

            # tolerate  ((char *)&fld + N)
            extra = 0
            if x.op == hr.cot_add and x.y.op == hr.cot_num:
                extra = x.y.numval()
                x = x.x
                while x.op == hr.cot_cast:
                    x = x.x

            if x.op == hr.cot_ref:
                x = x.x
            if x.op not in (hr.cot_memptr, hr.cot_memref):
                return 0

            if x.op == hr.cot_memptr:
                stype = x.x.type.get_pointed_object()
            else:
                stype = x.x.type

            newoff = x.m + extra + HALF
            mtype  = _member_at(stype, newoff)
            if mtype is not None:
                self.hits.append((e, x, newoff, mtype))

        except Exception as ex:
            if VERBOSE:
                print("[hi16] skip:", ex)
        return 0      # SWIG director requires an int on every path


# ---------------------------------------------------------------- rewriting

def fix_cfunc(cfunc):
    """Rewrite every hi-half read in `cfunc`. Returns count."""
    c = _Collector()
    c.apply_to(cfunc.body, None)

    n = 0
    for e, mem, newoff, mtype in c.hits:
        try:
            new = hr.cexpr_t()
            new.op   = mem.op
            new.x    = hr.cexpr_t(mem.x)
            new.m    = newoff
            new.type = mtype
            e.replace_by(new)
            n += 1
        except Exception as ex:
            if VERBOSE:
                print("[hi16] rewrite failed:", ex)
    return n


# ---------------------------------------------------------------- hook

class _Hooks(hr.Hexrays_Hooks):
    def maturity(self, cfunc, new_maturity):
        if ENABLED and new_maturity == hr.CMAT_FINAL:
            try:
                n = fix_cfunc(cfunc)
                if n:
                    _stats["rewrites"] += n
                    _stats["funcs"]    += 1
                    if VERBOSE:
                        print("[hi16] %d rewrites" % n)
            except Exception as ex:
                print("[hi16] ERROR:", ex)
        return 0


_hooks = None


def _install():
    global _hooks
    if _hooks is not None:
        _hooks.unhook()
        _hooks = None
    if not hr.init_hexrays_plugin():
        print("[hi16] decompiler unavailable")
        return False
    _hooks = _Hooks()
    _hooks.hook()
    return True


def uninstall():
    global _hooks
    if _hooks is not None:
        _hooks.unhook()
        _hooks = None
        print("[hi16] uninstalled")


def stats():
    print("[hi16] %d rewrites across %d functions"
          % (_stats["rewrites"], _stats["funcs"]))


# ---------------------------------------------------------------- plugin glue

class WatcomHi16(ida_idaapi.plugin_t):
    flags = ida_idaapi.PLUGIN_HIDE
    wanted_name = "Watcom hi-half reader"
    wanted_hotkey = ""
    comment = "Rewrites *(int*)&a >> 16 to the real 16-bit member at +2"
    help = comment

    def init(self):
        if not _install():
            return ida_idaapi.PLUGIN_SKIP
        print("[hi16] installed")
        return ida_idaapi.PLUGIN_KEEP

    def run(self, arg):
        stats()

    def term(self):
        uninstall()


def PLUGIN_ENTRY():
    return WatcomHi16()