Adopt idb_events and remote module features from ida-codemode

4 files changed, 117 insertions, 140 deletions

diff --git a/experiments/bench_ops.py b/experiments/bench_ops.py
index 178b4bc..9cc51fa 100644
--- a/experiments/bench_ops.py
+++ b/experiments/bench_ops.py
@@ -27,6 +27,7 @@ ida-codemode is installed **editable** into both venvs, so checking that repo ou
 swaps the backend under the TUI with no reinstall -- which is what makes this A/B
 cheap.
 """
+
 from __future__ import annotations
 
 import argparse
@@ -34,6 +35,7 @@ import os
 import statistics
 import time
 
+from idatui import remote_ops
 from idatui.codemode_client import CodeModeClient
 
 
@@ -59,7 +61,7 @@ def main() -> int:
 
     # Work on the biggest function we can find, so the payload-heavy operations
     # are actually payload-heavy.
-    index = client.invoke("list_funcs", queries=[{"offset": 0, "count": 60}])
+    index = client.call(remote_ops.list_funcs, queries=[{"offset": 0, "count": 60}])
     funcs = (index.get("result") or [{}])[0].get("data") or []
     if not funcs:
         print("VERDICT: FAIL - no functions")
@@ -71,19 +73,30 @@ def main() -> int:
         # Synthetic: isolates the per-operation floor (execute_sync marshalling).
         ("empty round trip", lambda: handle.execute_python("result = 1")),
         # Payload-dominated: what _PACK_EPILOGUE was written for.
-        ("list_funcs 500", lambda: client.invoke(
-            "list_funcs", queries=[{"offset": 0, "count": 500}])),
-        ("heads 200 (listing page)", lambda: client.invoke(
-            "heads", addr=ea, count=200, annotate=True)),
+        (
+            "list_funcs 500",
+            lambda: client.call(
+                remote_ops.list_funcs, queries=[{"offset": 0, "count": 500}]
+            ),
+        ),
+        (
+            "heads 200 (listing page)",
+            lambda: client.call(remote_ops.heads, addr=ea, count=200, annotate=True),
+        ),
         # IDA-work-dominated: Hex-Rays, nothing upstream can move.
-        ("decompile (warm)", lambda: client.invoke("decompile", addr=ea)),
-        ("flowchart (graph)", lambda: client.invoke("flowchart", addr=ea)),
+        ("decompile (warm)", lambda: client.call(remote_ops.decompile, addr=ea)),
+        ("flowchart (graph)", lambda: client.call(remote_ops.flowchart, addr=ea)),
         # Round-trip-dominated: small payload, so only the floor matters.
-        ("xrefs_to", lambda: client.invoke("xref_query", direction="to", addr=ea)),
+        (
+            "xrefs_to",
+            lambda: client.call(remote_ops.xref_query, direction="to", addr=ea),
+        ),
     ]
 
-    print(f"# target={os.path.basename(args.target)} func={ea} reps={args.reps} "
-          f"backend={client.backend}")
+    print(
+        f"# target={os.path.basename(args.target)} func={ea} reps={args.reps} "
+        f"backend={client.backend}"
+    )
     results = {}
     for name, fn in ops:
         try:
diff --git a/experiments/bench_pack_trace.py b/experiments/bench_pack_trace.py
index c6ca980..edbb9f7 100644
--- a/experiments/bench_pack_trace.py
+++ b/experiments/bench_pack_trace.py
@@ -1,59 +1,49 @@
-"""Measure ``_PACK_EPILOGUE`` against the live ida-codemode runtime.
+"""Measure cold installation versus warm calls for typed remote modules.
 
-Our snippets return one pre-serialised JSON STRING instead of a structure, to
-dodge to_jsonable()'s Python-level walk of the result (a 200-row listing page
-is ~10k small objects). ida-codemode 0.3.2 gave that path a C fast path --
-``serialization.dumps_json`` hands the structure straight to ``json.dumps`` and
-only falls back to the walker for values the encoder rejects -- so the packing
-now costs a double encode (escaping the whole payload as a string literal) to
-avoid a walk that may no longer happen.
-
-This script answers whether packing still pays. Its sibling question, the
-``sys.settrace`` strip, is settled: 0.3.2 deleted the trace hook, the workaround
-measured 0.99x, and it has been removed.
+Historical note: this file used to benchmark ``_PACK_EPILOGUE``. Application
+scripts are no longer strings and packing is gone; the relevant design cost is
+now the one-time content-addressed module installation versus steady-state calls.
 
 Usage::
 
-    PYTHONPATH=. ~/ida-venv/bin/python experiments/bench_pack_trace.py [FILE]
+    PYTHONPATH=. python experiments/bench_pack_trace.py [FILE]
 """
+
 from __future__ import annotations
 
 import argparse
-import json
 import os
 import statistics
 import time
 
-from idatui import codemode_client as cc
+from idatui import remote_ops
 from idatui.codemode_client import CodeModeClient
 
 
-def _time(fn, reps: int) -> tuple[float, float]:
-    """Best-of and median wall time in ms; best-of resists co-tenant noise."""
+def timed(function, reps: int = 1) -> tuple[object, float]:
     samples = []
+    result = None
     for _ in range(reps):
         started = time.perf_counter()
-        fn()
+        result = function()
         samples.append((time.perf_counter() - started) * 1000.0)
-    return min(samples), statistics.median(samples)
+    return result, statistics.median(samples)
 
 
 def main() -> int:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("target", nargs="?", default="targets/bash")
-    ap.add_argument("--reps", type=int, default=25)
-    ap.add_argument("--rows", type=int, default=200)
-    args = ap.parse_args()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("target", nargs="?", default="targets/bash")
+    parser.add_argument("--reps", type=int, default=25)
+    parser.add_argument("--rows", type=int, default=200)
+    args = parser.parse_args()
 
-    target = os.path.abspath(args.target)
-    client = CodeModeClient(target)
-    client.connect()
+    client = CodeModeClient(os.path.abspath(args.target)).connect()
 
-    # A real listing page: the flow the workaround was tuned for.
-    # list_funcs answers {"result": [{"data": [...], "total": N}]}.
-    index = client.invoke("list_funcs", queries=[{"offset": 0, "count": 40}])
+    index, operations_cold = timed(
+        lambda: client.call(remote_ops.list_funcs, queries=[{"offset": 0, "count": 40}])
+    )
     funcs = (index.get("result") or [{}])[0].get("data") or []
-    biggest = max(funcs, key=lambda f: f.get("size") or 0, default=None)
+    biggest = max(funcs, key=lambda function: function.get("size") or 0, default=None)
     if not biggest:
         print("VERDICT: FAIL - no functions")
         return 1
@@ -61,34 +51,26 @@ def main() -> int:
     if isinstance(addr, int):
         addr = hex(addr)
 
-    page = lambda: client.invoke(  # noqa: E731
-        "heads", addr=addr, count=args.rows, annotate=True)
-
-    # _script() reads _PACK_EPILOGUE at CALL time, so both variants share one
-    # process -- one lease, one warm database, one fair baseline. _unpack()
-    # passes an unpacked answer through untouched, so plain `result` works.
-    packed_epilogue = cc._PACK_EPILOGUE
-    results = {}
-    for packing in (True, False):
-        cc._PACK_EPILOGUE = packed_epilogue if packing else "\nresult\n"
-        payload = page()
-        rows = len(payload.get("heads", []))
-        for _ in range(3):  # warm caches; the first sample is always an outlier
-            page()
-        results[packing] = (rows, *_time(page, args.reps))
-    cc._PACK_EPILOGUE = packed_epilogue
+    _, operations_warm = timed(
+        lambda: client.call(
+            remote_ops.list_funcs, queries=[{"offset": 0, "count": 40}]
+        ),
+        args.reps,
+    )
+    page = lambda: client.call(  # noqa: E731
+        remote_ops.heads, addr=addr, count=args.rows, annotate=True
+    )
+    payload, tools_cold = timed(page)
+    _, tools_warm = timed(page, args.reps)
 
-    size = len(json.dumps(payload, separators=(",", ":"), default=str)) / 1024
-    print(f"target        {os.path.basename(target)}  backend={client.backend}")
-    print(f"listing page  func {addr}, {results[True][0]} rows, "
-          f"{size:.1f} KiB of JSON, {args.reps} reps")
-    for packing, label in ((True, "packed string (current)"),
-                           (False, "plain structure")):
-        rows, best, med = results[packing]
-        print(f"  {label:<26} best {best:7.2f}ms   median {med:7.2f}ms"
-              f"   rows={rows}")
-    print(f"  packing buys               "
-          f"{results[False][1] / results[True][1]:.2f}x")
+    print(f"target          {os.path.basename(args.target)}  backend={client.backend}")
+    print(f"function        {addr}")
+    print(f"listing rows    {len(payload.get('heads', []))}")
+    print(
+        f"operations.py   cold {operations_cold:8.3f}ms  warm {operations_warm:8.3f}ms"
+    )
+    print(f"remote_tools.py cold {tools_cold:8.3f}ms  warm {tools_warm:8.3f}ms")
+    print(f"tools install overhead {tools_cold / max(tools_warm, 0.001):.2f}x one time")
     client.close()
     return 0
 
diff --git a/experiments/profile_client.py b/experiments/profile_client.py
index 77f16e7..bac130e 100644
--- a/experiments/profile_client.py
+++ b/experiments/profile_client.py
@@ -10,6 +10,7 @@ nothing else here can see it.
 Time spent in `invoke` is the backend + transport; everything below it in the
 `tottime` list is ours and is what this file is for.
 """
+
 from __future__ import annotations
 
 import argparse
@@ -19,6 +20,7 @@ import os
 import pstats
 import time
 
+from idatui import remote_ops
 from idatui.codemode_client import CodeModeClient
 from idatui.domain import Program
 
@@ -28,14 +30,15 @@ def main() -> int:
     ap.add_argument("binary", nargs="?", default="targets/bash")
     ap.add_argument("--pages", type=int, default=60)
     ap.add_argument("--lines", type=int, default=16)
-    ap.add_argument("--text", action="store_true",
-                    help="load full pages instead of skeletons")
+    ap.add_argument(
+        "--text", action="store_true", help="load full pages instead of skeletons"
+    )
     args = ap.parse_args()
 
     client = CodeModeClient(os.path.abspath(args.binary))
     client.connect()
     program = Program(client)
-    regions = client.invoke("file_regions")
+    regions = client.call(remote_ops.file_regions)
     rows = regions.get("regions") or regions.get("result") or []
     text_seg = next((r for r in rows if ".text" in str(r.get("name", ""))), rows[0])
     model = program.listing(int(str(text_seg["start"]), 16))
@@ -57,8 +60,10 @@ def main() -> int:
     pr.disable()
     wall = (time.perf_counter() - started) * 1000
 
-    print(f"# {os.path.basename(args.binary)}  pages={loaded}  "
-          f"text={want_text}  {wall:.0f}ms  ({wall/max(loaded,1):.2f}ms/page)")
+    print(
+        f"# {os.path.basename(args.binary)}  pages={loaded}  "
+        f"text={want_text}  {wall:.0f}ms  ({wall / max(loaded, 1):.2f}ms/page)"
+    )
     buf = io.StringIO()
     pstats.Stats(pr, stream=buf).sort_stats("tottime").print_stats(args.lines)
     print(buf.getvalue())
diff --git a/experiments/profile_remote.py b/experiments/profile_remote.py
index ca673db..97f4fb1 100644
--- a/experiments/profile_remote.py
+++ b/experiments/profile_remote.py
@@ -1,94 +1,71 @@
-"""Profile an operation INSIDE the database process.
+"""Profile persistent ida-tui operations inside the IDA process.
 
-`bench_ops.py` says how long an operation takes; this says where that time
-goes. The snippet ships cProfile into the Code Mode sandbox, runs the real
-remote-library function there in a loop, and returns the stats as text -- so
-the split between IDA's own calls and OUR python in `remote_tools.py` is
-visible, which no client-side timer can see.
+The profiler itself is a typed ``RemoteModule`` function in ``remote_tools.py``;
+this file contains no generated Python source or knowledge of remote module names.
 
-    PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py [BINARY]
-    PYTHONPATH=. ~/ida-venv/bin/python experiments/profile_remote.py --op decompile
+Usage::
 
-Read the `tottime` column: time in that function excluding subcalls. IDA
-builtins (generate_disasm_line, get_flags, next_head...) are the floor; a
-python frame from ida_tui_remote near the top is ours, and ours is fixable.
+    PYTHONPATH=. python experiments/profile_remote.py [BINARY]
+    PYTHONPATH=. python experiments/profile_remote.py --op decompile
 """
+
 from __future__ import annotations
 
 import argparse
 import os
-import sys
-
-from idatui.codemode_client import CodeModeClient, _REMOTE_MODULE, _script
 
-# Runs in the database process. `a` is the bound argument dict.
-PROFILE = '''
-import cProfile, pstats, io, sys
-_m = sys.modules.get(%(mod)r)
-if _m is None:
-    result = {"error": "remote lib not installed yet"}
-else:
-    call = a["call"]
-    reps = int(a["reps"])
-    ns = {"_m": _m, "a": a}
-    src = "for _ in range(%%d):\\n    _m.%%s" %% (reps, call)
-    code = compile(src, "<profile>", "exec")
-    pr = cProfile.Profile()
-    pr.enable()
-    exec(code, ns)
-    pr.disable()
-    buf = io.StringIO()
-    st = pstats.Stats(pr, stream=buf).sort_stats("tottime")
-    st.print_stats(int(a["lines"]))
-    result = {"stats": buf.getvalue(), "total": st.total_tt, "reps": reps}
-''' % {"mod": _REMOTE_MODULE}
+from idatui import remote_ops
+from idatui.codemode_client import CodeModeClient
 
 CALLS = {
-    # One full listing page, exactly as the background grower asks for it.
-    "heads": 'heads(addr=a["addr"], count=500, annotate=True)',
-    "heads_plain": 'heads(addr=a["addr"], count=500, annotate=False)',
-    "heads_skeleton": 'heads(addr=a["addr"], count=500, annotate=True, text=False)',
-    "decompile": 'decompile(a["addr"])',
-    "disasm": 'disasm(a["addr"], 500)',
+    "heads": ("heads", {"count": 500, "annotate": True}),
+    "heads_plain": ("heads", {"count": 500, "annotate": False}),
+    "heads_skeleton": (
+        "heads",
+        {"count": 500, "annotate": True, "text": False},
+    ),
+    "decompile": ("decompile", {}),
 }
 
 
 def main() -> int:
-    ap = argparse.ArgumentParser()
-    ap.add_argument("binary", nargs="?", default="targets/bash")
-    ap.add_argument("--op", default="heads", choices=sorted(CALLS))
-    ap.add_argument("--reps", type=int, default=20)
-    ap.add_argument("--lines", type=int, default=18)
-    ap.add_argument("--addr", default=None, help="default: the .text start")
-    args = ap.parse_args()
-
-    client = CodeModeClient(os.path.abspath(args.binary))
-    client.connect()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("binary", nargs="?", default="targets/bash")
+    parser.add_argument("--op", default="heads", choices=sorted(CALLS))
+    parser.add_argument("--reps", type=int, default=20)
+    parser.add_argument("--addr", default=None, help="default: the .text start")
+    args = parser.parse_args()
 
+    client = CodeModeClient(os.path.abspath(args.binary)).connect()
     addr = args.addr
     if addr is None:
-        regions = client.invoke("file_regions")
+        regions = client.call(remote_ops.file_regions)
         rows = regions.get("regions") or regions.get("result") or []
-        text = next((r for r in rows if ".text" in str(r.get("name", ""))), None)
+        text = next(
+            (row for row in rows if ".text" in str(row.get("name", ""))),
+            None,
+        )
         addr = (text or rows[0])["start"] if rows else "0x0"
-    print(f"# {os.path.basename(args.binary)}  op={args.op}  addr={addr}  reps={args.reps}")
 
-    # Prime: the remote lib installs lazily, and its lru_caches must be warm or
-    # the profile measures cache misses that the real workload never pays.
-    client.invoke("heads", addr=addr, count=500, annotate=True)
+    operation, call_args = CALLS[args.op]
+    call_args = {"addr": addr, **call_args}
+    print(
+        f"# {os.path.basename(args.binary)}  op={args.op}  "
+        f"addr={addr}  reps={args.reps}"
+    )
 
-    # _script binds the args as JSON and adds the pack epilogue, exactly as a
-    # real operation is shipped -- so this measures the same path, not a
-    # special one.
-    out = client._unpack(client.execute_python(_script(
-        {"call": CALLS[args.op], "addr": addr,
-         "reps": args.reps, "lines": args.lines}, PROFILE), timeout=600))
-    if "error" in out:
-        print("FAILED:", out["error"])
-        return 1
+    # Install the persistent tool module and warm its caches before profiling.
+    client.call(remote_ops.heads, addr=addr, count=500, annotate=True)
+    out = client.call(
+        remote_ops.profile_remote,
+        operation=operation,
+        args=call_args,
+        reps=args.reps,
+    )
     per = out["total"] / out["reps"] * 1000
-    print(f"# {out['total']*1000:.0f}ms total, {per:.1f}ms per call\n")
+    print(f"# {out['total'] * 1000:.0f}ms total, {per:.1f}ms per call\n")
     print(out["stats"])
+    client.close()
     return 0