diff options
| -rw-r--r-- | idatui/app.py | 18 | ||||
| -rw-r--r-- | idatui/codemode_client.py | 2 | ||||
| -rw-r--r-- | idatui/domain.py | 78 | ||||
| -rw-r--r-- | idatui/remote_tools.py | 119 | ||||
| -rw-r--r-- | tests/test_scenarios.py | 32 |
5 files changed, 244 insertions, 5 deletions
diff --git a/idatui/app.py b/idatui/app.py index a806d65..3f67756 100644 --- a/idatui/app.py +++ b/idatui/app.py @@ -1185,11 +1185,21 @@ class ListingView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=Tru model = self.model if model is None: return - # Load just enough to render the viewport around the cursor, so the - # listing appears immediately even on a huge segment; the rest streams - # in via _grow. (load_all here would blank the pane for seconds.) + # One call gets the WHOLE row index -- every row's address, kind and + # size, and the page boundaries -- so the scrollbar is right immediately + # and _grow has nothing left to stream. The rows arrive text-less and + # materialise a page at a time as they are read. + # + # It is an optimisation, not a contract: an older or unhappy backend + # returns nothing usable and we stream exactly as before. height = max(self.size.height, 1) - model.ensure(self.cursor + height + 2 * ListingModel.PAGE) + if model.build_from_index(): + # Render the viewport HERE, on this worker thread. Reading a + # skeleton page fetches it, and doing that lazily from render_line + # would put an RPC on the UI loop for the first paint. + model.window(max(self.cursor - height, 0), height * 3) + else: + model.ensure(self.cursor + height + 2 * ListingModel.PAGE) self.app.call_from_thread(self._on_primed, len(model), model.complete) if not model.complete: self._grow() diff --git a/idatui/codemode_client.py b/idatui/codemode_client.py index 2361e08..61ff6d8 100644 --- a/idatui/codemode_client.py +++ b/idatui/codemode_client.py @@ -1209,7 +1209,7 @@ _OPERATIONS["pc_num_format"] = _remote_op( #: remote_tools.segment_index: the alternative is fetching every row. _OPERATIONS["segment_index"] = _remote_op( 'segment_index(addr=a["addr"], end=a.get("end", ""),' - ' page_rows=int(a.get("page_rows", 500)))') + ' page_rows=int(a.get("page_rows", 500)), detail=bool(a.get("detail", False)))') _HEADS = _remote_op( 'heads(addr=a["addr"], count=int(a.get("count", 200)),' diff --git a/idatui/domain.py b/idatui/domain.py index 60dbbbf..ecc6a80 100644 --- a/idatui/domain.py +++ b/idatui/domain.py @@ -17,9 +17,11 @@ Textual worker threads; the internal prefetch pool is separate and small. from __future__ import annotations +import array import bisect import re import threading +from base64 import b64decode from concurrent.futures import ThreadPoolExecutor from collections.abc import Sequence from dataclasses import dataclass, field, replace @@ -802,6 +804,82 @@ class ListingModel: with self._lock: return self._max_raw + def build_from_index(self) -> bool: + """Populate the whole row index from ONE call instead of streaming it. + + ``segment_index(detail=True)`` walks the segment and returns every row's + address, kind and size as packed arrays, plus the page boundaries a + refetch would use. That is everything this model needs to know how many + rows there are and where each one lives -- all that is missing is the + rendered text, which is exactly what a skeleton page is missing too. + + So the rows land marked ``_SKELETON_GEN`` and the FIRST read of any page + materialises it through the existing ``_ensure_text``/``_ensure_page`` + path, the same one a rename uses. Measured on bash: 594ms and one call, + against 1827ms and 458 for streaming the same thing. + + Returns False if the backend cannot supply it, in which case the caller + should stream as before -- this is an optimisation, not a new contract. + """ + try: + idx = self._prog.client.invoke( + "segment_index", addr=hex(self.seg_start), end=hex(self.seg_end), + page_rows=self.PAGE, detail=True) + except Exception: # noqa: BLE001 -- fall back to streaming + return False + if not isinstance(idx, dict) or idx.get("error") or "eas" not in idx: + return False + try: + eas = array.array("Q"); eas.frombytes(b64decode(idx["eas"])) + kinds = array.array("B"); kinds.frombytes(b64decode(idx["kinds"])) + sizes = array.array("I"); sizes.frombytes(b64decode(idx["sizes"])) + except Exception: # noqa: BLE001 + return False + names = idx.get("kind_names") or [] + anchors = idx.get("anchors") or [] + n = len(eas) + if not (n == len(kinds) == len(sizes)) or not anchors: + return False + + heads: list[Head] = [] + row_at: list[int] = [] + by_ea: dict[int, int] = {} + rows = 0 + ap = heads.append + rap = row_at.append + for i in range(n): + ea = eas[i] + kind = names[kinds[i]] if kinds[i] < len(names) else "unknown" + size = sizes[i] + ap(Head(ea, kind, size, "")) + rap(rows) + # Banner/label rows are display-only; navigation must land on the + # real head at that address. Same rule as the streaming loader. + if kind not in ("sep", "funchdr", "label"): + by_ea.setdefault(ea, rows) + rows += size if (kind == "unknown" and size > 1) else 1 + + with self._lock: + self._heads = heads + self._head_eas = list(eas) + self._head_gen = [self._SKELETON_GEN] * n + self._row_at = row_at + self._by_ea = by_ea + self._rows = rows + # Anchors are [logical_row, ea, head_index] at the exact boundaries + # heads(count=PAGE) pages on, so _ensure_page can refetch one page + # and have it line up head for head. + self._page_head = [a[2] for a in anchors] + self._page_addr = [_as_int(a[1]) for a in anchors] + self._page_digest = [None] * len(anchors) + self._page_rows = [ + (anchors[k + 1][2] if k + 1 < len(anchors) else n) - anchors[k][2] + for k in range(len(anchors))] + self._skeleton = True + self._done = True + self._next = None + return True + def load_next_page(self, text: bool = True) -> int: """Load one more page of heads; returns how many were added. diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py index 614fba4..0e26965 100644 --- a/idatui/remote_tools.py +++ b/idatui/remote_tools.py @@ -510,10 +510,127 @@ def _idatui_func_footer_rows(ea, func): ] +#: Row kinds, as small ints, for the packed detail index. Order is frozen: the +#: client decodes by position. +_IDATUI_KINDS = ("code", "data", "unknown", "sep", "funchdr", "label", "member") +_IDATUI_KIND_ID = {k: i for i, k in enumerate(_IDATUI_KINDS)} + + +def _idatui_segment_detail(addr, end, page_rows): + """Every listing ROW of a segment as packed arrays, with no text. + + ``{eas, kinds, sizes}`` are raw buffers -- uint64, uint8, uint32, one entry + per row in listing order -- so the client can build its whole row index + (addresses, spans, ea->row map) from ONE call instead of 458 pages. + + This re-implements the row sequence that ``_rows_for`` emits rather than + calling it, because building the dicts is most of what a page costs and + skipping them is the entire point. That duplication is the risk, so it is + covered by a test that walks a whole segment and compares this against the + real ``heads()`` output row for row -- if the two ever drift, that fails. + """ + import array + import ida_segment + + start = parse_address(addr) + seg = ida_segment.getseg(start) + if not seg: + return {"addr": str(addr), "error": "no segment", "rows": 0, "anchors": []} + lo, hi = seg.start_ea, seg.end_ea + if end: + try: + hi = min(hi, parse_address(end)) + except Exception: + pass + + K_CODE = _IDATUI_KIND_ID["code"]; K_DATA = _IDATUI_KIND_ID["data"] + K_UNK = _IDATUI_KIND_ID["unknown"]; K_SEP = _IDATUI_KIND_ID["sep"] + K_FUNC = _IDATUI_KIND_ID["funchdr"]; K_LABEL = _IDATUI_KIND_ID["label"] + K_MEMBER = _IDATUI_KIND_ID["member"] + + eas = array.array("Q") + kinds = array.array("B") + sizes = array.array("I") + ea_ap, kind_ap, size_ap = eas.append, kinds.append, sizes.append + + get_flags = ida_bytes.get_flags + get_item_end = ida_bytes.get_item_end + get_item_size = ida_bytes.get_item_size + next_head = ida_bytes.next_head + get_ea_name = ida_name.get_ea_name + get_func = idaapi.get_func + BAD = idaapi.BADADDR + + # Anchors mark where heads(addr=..., count=page_rows) would START each page, + # so a client can refetch exactly one page. heads() stops once it has + # emitted >= count PHYSICAL rows, checked before the next head -- so a + # boundary is the first head at which the running physical count reached the + # limit. Anchoring every N LOGICAL rows instead looks equivalent (the two + # are the same number until a segment contains an undefined run) and then + # silently yields pages that do not line up with a refetch. + anchors = [] + page_phys = 0 # physical rows emitted into the page being filled + rows = 0 # logical rows so far (what the scrollbar counts) + fn = None + ea = ida_bytes.get_item_head(lo) + while ea != BAD and ea < hi: + if not anchors or page_phys >= page_rows: + anchors.append([rows, hex(ea), len(eas)]) + page_phys = 0 + before = len(eas) + f = get_flags(ea) + cls = f & _MS_CLS + if cls != _FF_CODE and cls != _FF_DATA: + nh = next_head(ea, hi) + stop = nh if (nh != BAD and ea < nh <= hi) else hi + run = stop - ea + ea_ap(ea); kind_ap(K_UNK); size_ap(run) + rows += run if run > 1 else 1 + page_phys += len(eas) - before + ea = stop + continue + if fn is None or not (fn.start_ea <= ea < fn.end_ea): + fn = get_func(ea) + at_start = fn is not None and fn.start_ea == ea + if at_start: + for k in (K_SEP, K_SEP, K_FUNC): # blank, banner, `name proc` + ea_ap(ea); kind_ap(k); size_ap(0) + rows += 3 + elif cls == _FF_CODE and get_ea_name(ea): + ea_ap(ea); kind_ap(K_LABEL); size_ap(0) + rows += 1 + ea_ap(ea); kind_ap(K_CODE if cls == _FF_CODE else K_DATA) + size_ap(int(get_item_size(ea))); rows += 1 + if cls == _FF_DATA: + for m in _idatui_struct_member_rows(ea): + ea_ap(int(m["ea"], 16) if isinstance(m["ea"], str) else m["ea"]) + kind_ap(_IDATUI_KIND_ID.get(m.get("kind", "member"), K_MEMBER)) + size_ap(int(m.get("size", 0) or 0)); rows += 1 + item_end = get_item_end(ea) + if fn is not None and item_end >= fn.end_ea: + for k in (K_FUNC, K_SEP): # `name endp`, separator + ea_ap(ea); kind_ap(k); size_ap(0) + rows += 2 + page_phys += len(eas) - before + ea = item_end if item_end > ea else ea + 1 + + # base64, not raw bytes: the client packs answers with json.dumps(default=str), + # which turns a bytes object into its repr -- 4 characters per byte and + # unparseable at the other end. Learned by watching 2.97MB arrive as 11.26MB. + import base64 + b64 = base64.b64encode + return {"addr": hex(lo), "end": hex(hi), "rows": rows, "heads": len(eas), + "anchors": anchors, "kind_names": list(_IDATUI_KINDS), + "eas": b64(eas.tobytes()).decode(), + "kinds": b64(kinds.tobytes()).decode(), + "sizes": b64(sizes.tobytes()).decode()} + + def segment_index( addr: Annotated[str, "Any address in the segment to index"], end: Annotated[str, "Optional exclusive end address; default = segment end"] = "", page_rows: Annotated[int, "Rows between anchors (default 500)"] = 500, + detail: Annotated[bool, "Also return every row's ea/kind/size as packed arrays"] = False, ) -> dict: """How many listing rows a segment has, and where to seek into it. @@ -542,6 +659,8 @@ def segment_index( start = parse_address(addr) except Exception as e: return {"addr": str(addr), "error": str(e), "rows": 0, "anchors": []} + if detail: + return _idatui_segment_detail(addr, end, count) import ida_segment seg = ida_segment.getseg(start) if not seg: diff --git a/tests/test_scenarios.py b/tests/test_scenarios.py index b256b80..0bbcf17 100644 --- a/tests/test_scenarios.py +++ b/tests/test_scenarios.py @@ -459,6 +459,38 @@ async def s_segment_index(c: Ctx): c.check("every anchor points at the row it claims", not bad, f"{len(bad)} wrong, first={bad[:2]}") + # A model built from the index must be INDISTINGUISHABLE from a streamed + # one. That is the invariant the whole optimisation rests on: _prime builds + # from the index now, and every read path -- rendering, goto, xrefs, search, + # rename refresh -- indexes into these arrays assuming they were produced + # the old way. Comparing row counts alone would miss a shifted _row_at or a + # _by_ea that sends a jump to the wrong line. + # + # BOTH models are built here, back to back. Comparing against the app's + # long-lived model instead is wrong by one row and flaky: earlier scenarios + # rename and define things, so that model describes the database as it was + # at boot, not as it is now. + from idatui.domain import ListingModel # noqa: PLC0415 + args = (app.program, model.seg_start, model.seg_end, model.name) + idx_model, streamed = ListingModel(*args), ListingModel(*args) + if not idx_model.build_from_index(): + c.check("build_from_index works", False) + return + while not streamed.complete: + if streamed.load_next_page(text=False) == 0: + break + c.check("an index-built model is complete immediately", idx_model.complete) + c.check("index-built model has the streamed row count", + len(idx_model) == len(streamed), f"{len(idx_model)} vs {len(streamed)}") + for field in ("_row_at", "_head_eas", "_by_ea", "_page_head", + "_page_addr", "_page_rows"): + a, b = getattr(idx_model, field), getattr(streamed, field) + c.check(f"index-built {field} matches streaming", a == b, + f"len {len(a)} vs {len(b)}") + c.check("index-built rows carry the same ea/kind/size", + [(h.ea, h.kind, h.size) for h in idx_model._heads] + == [(h.ea, h.kind, h.size) for h in streamed._heads]) + @scenario("skeleton_pages") async def s_skeleton_pages(c: Ctx): |