Fix parsing of pseudocode tags

3 files changed, 59 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])