1"""Streaming XML helpers for xctrace export output.
2 
3Instruments XML deduplicates repeated values with `id`/`ref` attributes that
4can span the whole document, so we stream rows with iterparse while keeping
5a global id cache for later ref lookups.
6"""
7from __future__ import annotations
8 
9import xml.etree.ElementTree as ET
10from collections.abc import Iterator
11from dataclasses import dataclass
12 
13 
14@dataclass(frozen=True)
15class Column:
16    mnemonic: str          # e.g. "time", "weight", "stack"
17    engineering_type: str  # e.g. "sample-time", "weight", "tagged-backtrace"
18 
19 
20class RowStream:
21    """Iterate <row> elements of a single <table> schema export.
22 
23    Yields `dict[str, Element]` keyed by column mnemonic. Elements inside a
24    yielded row are live ET elements (rooted in the id cache where applicable
25    so ref resolution via `resolve()` remains valid after the row is yielded).
26    """
27 
28    def __init__(self, xml_bytes: bytes):
29        self._xml = xml_bytes
30        self.columns: list[Column] = []
31        self._id_cache: dict[str, ET.Element] = {}
32 
33    def resolve(self, element: ET.Element) -> ET.Element:
34        """If the element is a ref, return the referenced element; else self."""
35        ref = element.get("ref")
36        if ref is None:
37            return element
38        target = self._id_cache.get(ref)
39        if target is None:
40            return element  # unresolved; return the ref element itself
41        return target
42 
43    def __iter__(self) -> Iterator[dict[str, ET.Element]]:
44        # iterparse fires `end` events once an element is fully parsed, so ids
45        # are visible to descendants via the cache. We only need `end` events;
46        # row bodies are reconstructed from the end element itself in _row_dict.
47        #
48        # NOTE: we intentionally don't call `elem.clear()` after yielding a row.
49        # Instruments' XML is a single shared doc where any row can `ref` an
50        # `id` defined earlier (threads, processes, stacks, metadata), and
51        # clearing would break those later lookups. The tradeoff is peak RAM
52        # ≈ document size. That's fine for typical traces up to a few hundred
53        # MB; very large exports may need a smarter pass that first indexes
54        # referenced ids and only retains those.
55        schema_seen = False
56 
57        context = ET.iterparse(_bytes_to_file(self._xml), events=("end",))
58        for _event, elem in context:
59            eid = elem.get("id")
60            if eid is not None:
61                self._id_cache[eid] = elem
62 
63            if elem.tag == "schema" and not schema_seen:
64                self.columns = _parse_columns(elem)
65                schema_seen = True
66                continue
67 
68            if elem.tag == "row":
69                yield _row_dict(elem, self.columns)
70                # Do not clear elem — children referenced via id may still be needed.
71 
72 
73def _parse_columns(schema_el: ET.Element) -> list[Column]:
74    cols: list[Column] = []
75    for col in schema_el.findall("col"):
76        mnemonic = (col.findtext("mnemonic") or "").strip()
77        etype = (col.findtext("engineering-type") or "").strip()
78        if mnemonic:
79            cols.append(Column(mnemonic=mnemonic, engineering_type=etype))
80    return cols
81 
82 
83def _row_dict(row_el: ET.Element, cols: list[Column]) -> dict[str, ET.Element]:
84    # Row children map positionally to columns. <sentinel/> marks a missing
85    # optional value for that column.
86    result: dict[str, ET.Element] = {}
87    children = list(row_el)
88    for idx, child in enumerate(children):
89        if idx >= len(cols):
90            break
91        if child.tag == "sentinel":
92            continue
93        result[cols[idx].mnemonic] = child
94    return result
95 
96 
97def _bytes_to_file(data: bytes):
98    import io
99    return io.BytesIO(data)
100 
101 
102# --- Extraction helpers ---------------------------------------------------
103 
104def int_text(elem: ET.Element | None) -> int | None:
105    if elem is None or elem.text is None:
106        return None
107    try:
108        return int(elem.text)
109    except ValueError:
110        return None
111 
112 
113def str_text(elem: ET.Element | None) -> str | None:
114    if elem is None or elem.text is None:
115        return None
116    return elem.text
117 
118 
119def fmt_attr(elem: ET.Element | None) -> str | None:
120    """Return the human-readable `fmt` attribute if present."""
121    if elem is None:
122        return None
123    return elem.get("fmt")
124 
125 
126def extract_thread(thread_el: ET.Element, stream: RowStream) -> dict:
127    """Parse a <thread> element into name, tid, process dict.
128 
129    Handles ref-style threads by resolving through the stream's id cache.
130    """
131    resolved = stream.resolve(thread_el)
132    name = resolved.get("fmt", "")
133    tid_el = resolved.find("tid")
134    process_el = resolved.find("process")
135    process = extract_process(process_el, stream) if process_el is not None else None
136    return {
137        "name": name,
138        "tid": int_text(tid_el),
139        "process": process,
140        "is_main": name.startswith("Main Thread") if name else False,
141    }
142 
143 
144def extract_process(process_el: ET.Element, stream: RowStream) -> dict:
145    resolved = stream.resolve(process_el)
146    name = resolved.get("fmt", "")
147    pid_el = resolved.find("pid")
148    return {
149        "name": _clean_process_name(name),
150        "pid": int_text(pid_el),
151    }
152 
153 
154def _clean_process_name(fmt: str) -> str:
155    # "NowPlaying Gigs (28401)" -> "NowPlaying Gigs"
156    if " (" in fmt and fmt.endswith(")"):
157        return fmt.rsplit(" (", 1)[0]
158    return fmt
159 
160 
161def extract_backtrace(
162    bt_el: ET.Element, stream: RowStream, max_frames: int = 20
163) -> list[dict]:
164    """Return a list of frame dicts from a <tagged-backtrace> or <backtrace>.
165 
166    Frames are ordered leaf-first (top of stack first), matching Instruments'
167    display order.
168    """
169    resolved = stream.resolve(bt_el)
170    inner = resolved.find("backtrace")
171    if inner is None:
172        inner = resolved
173    frames: list[dict] = []
174    for frame_el in inner.findall("frame"):
175        f = stream.resolve(frame_el)
176        frames.append({
177            "name": f.get("name") or "",
178            "addr": f.get("addr") or "",
179        })
180        if len(frames) >= max_frames:
181            break
182    return frames
183 
184 
185def top_symbol(frames: list[dict]) -> str:
186    """Pick the leaf symbol, falling back to addr if unsymbolicated."""
187    if not frames:
188        return "<empty-stack>"
189    first = frames[0]
190    return first.get("name") or first.get("addr") or "<unknown>"
191 
192 
193def first_present(row: dict, *keys: str) -> ET.Element | None:
194    """Return the first row column whose key exists.
195 
196    `row[key] or row[other_key]` is unsafe here: Element is falsy when it has
197    no children (a common case for leaf <event-time>, <start-time>, etc.), so
198    `or` short-circuits past valid leaf elements. This walks keys explicitly.
199    """
200    for key in keys:
201        el = row.get(key)
202        if el is not None:
203            return el
204    return None
205 
206 
207def in_window(time_ns: int | None, window: tuple[int, int] | None) -> bool:
208    """Return True if time_ns is inside [start, end] (inclusive), or window is None."""
209    if window is None:
210        return True
211    if time_ns is None:
212        return False
213    start, end = window
214    return start <= time_ns <= end
215 
216 
217def event_overlaps_window(
218    start_ns: int, end_ns: int, window: tuple[int, int] | None
219) -> bool:
220    """Return True if [start, end] overlaps [window.start, window.end]."""
221    if window is None:
222        return True
223    w_start, w_end = window
224    return not (end_ns < w_start or start_ns > w_end)
225