diff --git a/idatui/worker.py b/idatui/worker.py index 556e69a..04252cc 100644 --- a/idatui/worker.py +++ b/idatui/worker.py @@ -119,6 +119,36 @@ def recv(sock: socket.socket): # --------------------------------------------------------------------------- # # worker # --------------------------------------------------------------------------- # +def _use_fast_decompile() -> None: + """Point ida-pro-mcp's decompile tools at our per-line loop. + + The shipped ``decompile_function_safe`` allocates three ctree_item_t SWIG + objects per pseudocode line (two of which it never reads) and formats an + item description per line to recover the ``/*0xEA*/`` marker: 121us a line, + which on a warm cfunc is most of what the tool costs. The replacement lives + in server/patch_server.py and is differentially checked against the original + by .auto/check_decomp.py. + + Every consumer binds the name at import time (``from .utils import ...``), + so rebinding it on ``utils`` alone would miss them; rebind on each module + that imported it, and leave anything unexpected exactly as it was. + """ + try: + from ida_pro_mcp.ida_mcp import api_types, utils + fast = api_types._idatui_decompile_function_safe + except Exception: # noqa: BLE001 -- never let this stop the worker booting + return + utils.decompile_function_safe = fast + # Rebind only on modules that are ALREADY imported. Importing one to rebind + # it would be work the worker had not chosen to do, on the boot path. + prefix = "ida_pro_mcp.ida_mcp." + for name, mod in list(sys.modules.items()): + if not name.startswith(prefix) or mod is None: + continue + if getattr(mod, "decompile_function_safe", None) is not None: + mod.decompile_function_safe = fast + + def _ensure_tools_injected() -> None: """Inject idatui's custom tools (heads/read_raw/resolve_names/func_types/...) into the installed ida_pro_mcp, idempotently, so the worker is self-sufficient @@ -190,6 +220,8 @@ def _open_and_register(binpath: str, load_args: str = ""): # importing the package registers all api_*/patched tools against MCP_SERVER from ida_pro_mcp.ida_mcp import MCP_SERVER # noqa: WPS433 + _use_fast_decompile() + import ida_nalt module = os.path.basename(ida_nalt.get_root_filename() or binpath) diff --git a/server/patch_server.py b/server/patch_server.py index 6667e12..64424e4 100644 --- a/server/patch_server.py +++ b/server/patch_server.py @@ -1039,6 +1039,67 @@ def decomp_map( return {"addr": hex(func.start_ea), "lines": lines} +def _idatui_decompile_function_safe(ea, include_addresses=True): + """ida-pro-mcp's ``decompile_function_safe``, with the three costs the same + sweep had in ``decomp_map`` taken out. Byte-identical output -- it is + differentially checked against the original over every function of a real + binary by ``.auto/check_decomp.py``. + + The shipped version costs 121us per pseudocode line, which is more than the + line's share of Hex-Rays itself on a warm cfunc: + + * it allocates THREE ctree_item_t SWIG objects per line, and ``_head`` and + ``_tail`` are never read -- ``get_line_item`` accepts None for both. + * it calls ``dstr()`` per line. That formats a whole 'EA: description' + string at 24us a call, and consecutive lines of a multi-line expression + report the same ctree item, so memoising by ``obj_id`` (unique within a + cfunc) skips most of them. + + 18 991 lines of bash: 2 302ms -> 429ms. + """ + import ida_lines + import ida_hexrays as _hx + from ida_pro_mcp.ida_mcp.utils import compact_whitespace, decompile_checked + from ida_pro_mcp.ida_mcp.sync import IDAError + try: + cfunc = decompile_checked(ea) + item = _hx.ctree_item_t() + get_line_item = cfunc.get_line_item + tag_remove = ida_lines.tag_remove + ea_of_id = {} + lines = [] + for sl in cfunc.get_pseudocode(): + line = sl.line + line_ea = None + if include_addresses and get_line_item(line, 0, False, None, + item, None): + it = item.it + oid = it.obj_id if it is not None else None + if oid is not None and oid in ea_of_id: + line_ea = ea_of_id[oid] + else: + dstr = item.dstr() + if dstr: + ds = dstr.split(": ") + if len(ds) == 2: + try: + line_ea = int(ds[0], 16) + except ValueError: + pass + if oid is not None: + ea_of_id[oid] = line_ea + text = compact_whitespace(tag_remove(line)) + if line_ea is not None: + lines.append(f"{text} /*{line_ea:#x}*/") + else: + lines.append(text) + return "\\n".join(lines), None + except IDAError as e: + return None, str(e) + except Exception as e: + return None, f"Decompilation failed at {hex(ea)}: {e}" + + _idatui_strings_cache = {}