1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
|
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 = {}
|
