Fix parsing of pseudocode tags

5 files changed, 251 insertions, 53 deletions

diff --git a/idatui/app.py b/idatui/app.py
index 75dc9d1..52f474b 100644
--- a/idatui/app.py
+++ b/idatui/app.py
@@ -51,7 +51,7 @@ from textual.widgets import (
 from textual.widgets.option_list import Option
 
 from . import diag, findings, graph, kittygfx, search
-from .domain import Func, Head, ListingModel, Program, Struct
+from .domain import Func, Head, ListingModel, Program, Struct, split_pseudocode_line
 from .edit_ctl import EditController
 from .errors import IDAConnectionError
 from .highlight import CTextArea, highlight_c
@@ -193,12 +193,6 @@ _S_DECOMP_WAIT = Style(color="#7c8b9e", italic=True)  # 'decompiling' label
 _S_DECOMP_DOTS = Style(color="#626c7a")  # trailing ellipsis
 _S_LINK = Style(bgcolor="#233044")  # split view: rows linked to the other pane's cursor
 
-# Hex-Rays appends a `/*0xEA*/` address marker to each pseudocode line (we fetch
-# with include_addresses so we have a per-line anchor). Matched here to extract
-# the address and strip the marker from the display.
-_ADDR_MARK_RE = re.compile(r"/\*\s*0x([0-9A-Fa-f]+)\s*\*/")
-_ADDR_MARK_STRIP_RE = re.compile(r"\s*/\*\s*0x[0-9A-Fa-f]+\s*\*/")
-
 
 @dataclass
 class BinaryState:
@@ -1786,9 +1780,9 @@ class DecompView(SearchMixin, NavMixin, ColumnCursor, ScrollView, can_focus=True
         line_eas: list[int | None] = []
         clean: list[str] = []
         for ln in text.split("\n"):
-            eas = _ADDR_MARK_RE.findall(ln)
-            line_eas.append(int(eas[-1], 16) if eas else None)
-            clean.append(_ADDR_MARK_STRIP_RE.sub("", ln).rstrip())
+            source, line_ea = split_pseudocode_line(ln)
+            line_eas.append(line_ea)
+            clean.append(source)
         seglists = highlight_c("\n".join(clean))
         self._strips = [Strip(segs) for segs in seglists]
         self._texts = ["".join(seg.text for seg in segs) for segs in seglists]
@@ -8225,7 +8219,7 @@ class IdaTui(App):
         lines = (dec.code or "").splitlines()
         if not (0 <= line_idx < len(lines)):
             return 0
-        clean = _ADDR_MARK_STRIP_RE.sub("", lines[line_idx])
+        clean, _ = split_pseudocode_line(lines[line_idx])
         m = re.search(rf"\b{re.escape(name)}\b", clean)
         return m.start() if m else 0
 
@@ -8255,7 +8249,8 @@ class IdaTui(App):
         best: tuple[int, int] | None = None
         best_key: tuple[int, int] | None = None
         for i, ln in enumerate(lines):
-            m = pat.search(_ADDR_MARK_STRIP_RE.sub("", ln))
+            clean, _ = split_pseudocode_line(ln)
+            m = pat.search(clean)
             if not m:
                 continue
             # rank: nearest to the anchor; tie -> the line at/after the anchor.
@@ -8277,10 +8272,9 @@ class IdaTui(App):
             return -1
         best_idx, best_ea = -1, -1
         for i, line in enumerate(dec.code.splitlines()):
-            for m in re.findall(r"/\*\s*0x([0-9A-Fa-f]+)\s*\*/", line):
-                e = int(m, 16)
-                if best_ea < e <= ea:
-                    best_ea, best_idx = e, i
+            _, e = split_pseudocode_line(line)
+            if e is not None and best_ea < e <= ea:
+                best_ea, best_idx = e, i
         return best_idx
 
     @staticmethod
diff --git a/idatui/domain.py b/idatui/domain.py
index 1e5a863..941cc59 100644
--- a/idatui/domain.py
+++ b/idatui/domain.py
@@ -2626,6 +2626,21 @@ def _parse_xrefs(payload) -> list[Xref]:
     return out
 
 
+_PSEUDOCODE_ANCHOR_RE = re.compile(r"[ \t]*/\*[ \t]*0x([0-9A-Fa-f]+)[ \t]*\*/[ \t]*$")
+
+
+def split_pseudocode_line(line: str) -> tuple[str, int | None]:
+    """Separate the legacy trailing address annotation from a pseudocode line.
+
+    Only the final annotation is metadata. Identical text inside a string or an
+    earlier source comment must survive, both for display and cursor columns.
+    """
+    match = _PSEUDOCODE_ANCHOR_RE.search(line)
+    if match is None:
+        return line.rstrip(), None
+    return line[: match.start()].rstrip(), int(match.group(1), 16)
+
+
 def _parse_decompilation(ea: int, payload) -> Decompilation:
     if not isinstance(payload, dict):
         return Decompilation(ea, None, True, "unexpected payload", False, None)
diff --git a/idatui/remote_tools.py b/idatui/remote_tools.py
index 6681379..7176015 100644
--- a/idatui/remote_tools.py
+++ b/idatui/remote_tools.py
@@ -14,8 +14,8 @@ injected BODY, which the IDA Nexus port deletes):
 
 Keeping the originals rather than paraphrasing them is deliberate: this is the
 most performance-tuned and most behaviour-sensitive code in the project (the
-span walker is a single regex pass because a per-character loop was the most
-expensive thing the listing did, and the cycle only offers stops that change
+span walker skips text runs with a regex because a per-character loop was the
+most expensive thing the listing did, and the cycle only offers stops that change
 what you see). A re-implementation drifts from it silently.
 
 This file is SOURCE SHIPPED AS TEXT to the database process; it is never
@@ -252,7 +252,7 @@ _IDATUI_TAGS = None
 _IDATUI_OPND_TAGS = None
 
 
-_IDATUI_CTL = None  # re: a tag = one of three control chars plus its argument
+_IDATUI_CTL = None  # re: find a control byte; the kernel determines its size
 
 
 _IDATUI_TAGINFO = None
@@ -292,12 +292,11 @@ def _idatui_spans(line):
     if _IDATUI_CTL is None:
         import re as _re
 
-        # One capturing split gives [text, tag, text, tag, ..., text] in a
-        # single C pass. A per-character python loop over the line used to be
-        # the most expensive thing the `heads` tool did, and a line is ~54
-        # characters but only ~13 tags -- everything between two tags is already
-        # exactly one span's worth of text.
-        _IDATUI_CTL = _re.compile("([\\x01\\x02\\x03](?s:.))")
+        # Search for control bytes, not two-byte pairs: semantic tags have
+        # kernel-dependent headers, possibly containing control bytes themselves.
+        # Still skip ordinary text in C rather than walking it character by
+        # character (previously the most expensive part of `heads`).
+        _IDATUI_CTL = _re.compile("[\\x01-\\x04]")
     if _IDATUI_TAGINFO is None:
         _IDATUI_TAGINFO = {
             tag: (_IDATUI_TAGS.get(tag, "text"), _IDATUI_OPND_TAGS.get(tag))
@@ -305,43 +304,41 @@ def _idatui_spans(line):
         }
     taginfo = _IDATUI_TAGINFO
     plain_tag = ("text", None)
-    on, off, esc = "\x01", "\x02", "\x03"
+    on, off, esc, inv = "\x01", "\x02", "\x03", "\x04"
     addr_tag = chr(getattr(ida_lines, "COLOR_ADDR", 0x28))
-    addr_len = int(getattr(ida_lines, "COLOR_ADDR_SIZE", 16))
-    parts = _IDATUI_CTL.split(line)
+    semantic_tag = "\x36"
+    skipcode = ida_lines.tag_skipcode
     spans, stack = [], []  # stack entries: (kind, operand index|None)
     kind, opnd = "text", None  # state the current run of text belongs to
     pend = ""
-    skip = 0  # characters of an address payload still due
-    i, n = 0, len(parts)
+    i, n = 0, len(line)
     while i < n:
-        txt = parts[i]
-        i += 1
-        if skip:
-            if len(txt) <= skip:
-                skip -= len(txt)
-                txt = ""
-            else:
-                txt = txt[skip:]
-                skip = 0
-        if txt:
-            pend += txt
-        if i >= n:
+        match = _IDATUI_CTL.search(line, i)
+        if match is None:
+            pend += line[i:]
             break
-        pair = parts[i]
-        i += 1
-        if skip:  # a tag INSIDE an address payload: 2 chars
-            skip = skip - 2 if skip > 2 else 0
+        at = match.start()
+        pend += line[i:at]
+        size = skipcode(line[at:])
+        if size <= 0:  # malformed code: retain it and always make progress
+            pend += line[at]
+            i = at + 1
             continue
-        ch = pair[0]
+        i = min(at + size, n)
+        ch = line[at]
+        if ch == inv:  # one byte, no following argument to consume
+            continue
+        if at + 1 >= n:
+            break
+        tag = line[at + 1]
         if ch == esc:  # escaped literal: keep the char it guards
-            pend += pair[1]
+            pend += tag
             continue
-        tag = pair[1]
-        if ch == on and tag == addr_tag:
-            # An embedded target address, not display text: 16 hex digits that
-            # must not reach the screen. Deliberately NOT a span boundary.
-            skip = addr_len
+        if tag == semantic_tag or (ch == on and tag == addr_tag):
+            # Structural metadata is invisible and does not change colour or
+            # operand ownership. Older kernels emit headerless semantic groups;
+            # newer ones include a kind and sometimes a TID. Only skipcode knows
+            # the correct length, so never hard-code these payload sizes.
             continue
         if pend:
             spans.append([kind, pend, opnd])
diff --git a/tests/test_pseudocode.py b/tests/test_pseudocode.py
new file mode 100644
index 0000000..b3b9556
--- /dev/null
+++ b/tests/test_pseudocode.py
@@ -0,0 +1,58 @@
+"""IDA-free regressions for pseudocode display text and address annotations."""
+
+from __future__ import annotations
+
+import os
+import sys
+
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from idatui.domain import split_pseudocode_line  # noqa: E402
+
+NEEDS_IDA = False
+PASS = FAIL = 0
+
+
+def check(name, line, expected):
+    global PASS, FAIL
+    actual = split_pseudocode_line(line)
+    if actual == expected:
+        PASS += 1
+        print(f"  ok   {name}")
+    else:
+        FAIL += 1
+        print(f"  FAIL {name}: {actual!r} != {expected!r}")
+
+
+def main():
+    check(
+        "trailing annotation", "  return 42; /*0x401000*/", ("  return 42;", 0x401000)
+    )
+    source = '  printf("before /*0x1234*/ after", 42);'
+    check("marker-shaped string survives", source + " /*0x401000*/", (source, 0x401000))
+    check("string without annotation is not an anchor", source, (source, None))
+    source = "  x = /*0x1234*/ y;"
+    check(
+        "inline source comment survives", source + " /*0x401000*/", (source, 0x401000)
+    )
+    source = "  return x; /*0x1234*/"
+    check(
+        "only the final annotation is removed",
+        source + " /*0x401000*/",
+        (source, 0x401000),
+    )
+    check(
+        "non-address source comment",
+        "  return x; /* explanation */",
+        ("  return x; /* explanation */", None),
+    )
+    check(
+        "annotation whitespace and hex casing", "  x;\t/* 0xAbCd */  ", ("  x;", 0xABCD)
+    )
+    check("address zero is valid", "x; /*0x0*/", ("x;", 0))
+    check("blank line", "   ", ("", None))
+    print(f"\n{PASS} passed, {FAIL} failed")
+    return bool(FAIL)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())
diff --git a/tests/test_remote_tags.py b/tests/test_remote_tags.py
new file mode 100644
index 0000000..2c2b8aa
--- /dev/null
+++ b/tests/test_remote_tags.py
@@ -0,0 +1,134 @@
+"""IDA-free regressions for the source-shipped listing tag parser."""
+
+from __future__ import annotations
+
+import ast
+import re
+from pathlib import Path
+from types import SimpleNamespace
+
+NEEDS_IDA = False
+PASS = FAIL = 0
+ON, OFF, ESC, INV = "\x01", "\x02", "\x03", "\x04"
+ADDR, SEM = "\x28", "\x36"
+INSN, REG, NUM, OPND = "\x05", "\x29", "\x0c", "\x40"
+
+
+def load_parser(semantic_headers=True):
+    # remote_tools is shipped as source and imports IDAPython at module scope.
+    # Compile the actual parser, not a copy, with only its IDA surface faked.
+    path = Path(__file__).resolve().parents[1] / "idatui" / "remote_tools.py"
+    tree = ast.parse(path.read_text())
+    node = next(
+        n
+        for n in tree.body
+        if isinstance(n, ast.FunctionDef) and n.name == "_idatui_spans"
+    )
+    calls = []
+
+    def skipcode(at):
+        calls.append(at)
+        if at[0] == INV:
+            return 1
+        if at.startswith(ON + ADDR):
+            return 18
+        if semantic_headers and at.startswith(ON + SEM):
+            return 19 if at[2] == "\x03" else 3
+        return 2
+
+    ns = {
+        "_re": re,
+        "ida_lines": SimpleNamespace(
+            COLOR_ADDR=0x28, COLOR_ADDR_SIZE=16, tag_skipcode=skipcode
+        ),
+        "_IDATUI_TAGS": {INSN: "insn", REG: "reg", NUM: "num"},
+        "_IDATUI_OPND_TAGS": {OPND: 0},
+        "_IDATUI_CTL": None,
+        "_IDATUI_TAGINFO": None,
+    }
+    exec(compile(ast.Module(body=[node], type_ignores=[]), str(path), "exec"), ns)
+    return ns["_idatui_spans"], calls
+
+
+def tag(t, text):
+    return ON + t + text + OFF + t
+
+
+def check(name, actual, expected):
+    global PASS, FAIL
+    if actual == expected:
+        PASS += 1
+        print(f"  ok   {name}")
+    else:
+        FAIL += 1
+        print(f"  FAIL {name}: {actual!r} != {expected!r}")
+
+
+def main():
+    parse, calls = load_parser()
+    check(
+        "ordinary colours and operand coordinates",
+        parse(
+            tag(INSN, "mov")
+            + "   "
+            + tag(OPND, tag(REG, "eax") + ", " + tag(NUM, "42"))
+        ),
+        (
+            [
+                ["insn", "mov"],
+                ["text", " "],
+                ["reg", "eax"],
+                ["text", ", "],
+                ["num", "42"],
+            ],
+            [[4, 11, 0]],
+        ),
+    )
+    # Semantic metadata can itself contain bytes that look like presentation
+    # controls. The kernel's header length, not a two-byte regex, owns them.
+    name = ON + SEM + "\x01" + "main" + OFF + SEM
+    check(
+        "semantic name preserves enclosing colour",
+        parse(tag(INSN, "call " + name)),
+        ([["insn", "call main"]], []),
+    )
+    typed = ON + SEM + "\x03" + "0123456789ABCDEF" + tag(REG, "MY_TYPE") + OFF + SEM
+    check(
+        "type-id payload is invisible and operand survives",
+        parse(tag(OPND, typed + "+" + tag(NUM, "4"))),
+        ([["reg", "MY_TYPE"], ["text", "+"], ["num", "4"]], [[0, 9, 0]]),
+    )
+    check(
+        "address payload is invisible",
+        parse(tag(REG, "a" + ON + ADDR + "0123456789ABCDEF" + "b")),
+        ([["reg", "ab"]], []),
+    )
+    check(
+        "inverse is a one-byte code, including at end",
+        parse(tag(REG, "a" + INV + "b" + INV)),
+        ([["reg", "ab"]], []),
+    )
+    check(
+        "escaped controls are literal",
+        parse("a" + ESC + ON + "b"),
+        ([["text", "a" + ON + "b"]], []),
+    )
+    check("kernel skipcode is consulted", bool(calls), True)
+    old_parse, _ = load_parser(semantic_headers=False)
+    check(
+        "headerless semantic group preserves enclosing colour",
+        old_parse(tag(INSN, "call " + tag(SEM, "main"))),
+        ([["insn", "call main"]], []),
+    )
+    check(
+        "unknown presentation tags retain their text",
+        parse(tag("\x7f", "unknown")),
+        ([["text", "unknown"]], []),
+    )
+    check("empty line", parse(""), ([], []))
+    print(f"\n{PASS} passed, {FAIL} failed")
+    return bool(FAIL)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())