|
| 1 | +//! Node-exact capture extraction and the shared innermost-wins flatten. |
| 2 | +//! |
| 3 | +//! This is the **one resolver both consumers share**: the render producer |
| 4 | +//! (`Registry::highlight` / `UserGrammars::highlight`) and the editor's |
| 5 | +//! semantic-token extractor (`quarto-lsp-core::tokens`) both extract spans |
| 6 | +//! with [`captures_from_tree`] and collapse them with [`flatten_spans`], so |
| 7 | +//! a code cell is coloured identically in the editor and the rendered HTML |
| 8 | +//! by construction. |
| 9 | +//! |
| 10 | +//! It replaces the old `tree-sitter-highlight` `HighlightEvent`-stream walk |
| 11 | +//! (`collect_spans`), which was lossy for same-start nested captures: it |
| 12 | +//! dropped the inner capture's end boundary, stretching it to the outer |
| 13 | +//! capture's end (bd-98k6). `tree_sitter::Query::captures()` returns |
| 14 | +//! node-exact byte ranges, so the boundary survives. |
| 15 | +
|
| 16 | +use tree_sitter::{Node, Query, QueryCursor, StreamingIterator}; |
| 17 | + |
| 18 | +use crate::encoding::HighlightSpan; |
| 19 | + |
| 20 | +/// Run a `QueryCursor` over `root` and return one [`HighlightSpan`] per |
| 21 | +/// capture — node-exact `(start_byte, end_byte, capture_name)`, **unflattened** |
| 22 | +/// (captures from one query over one tree nest; [`flatten_spans`] collapses |
| 23 | +/// them). |
| 24 | +/// |
| 25 | +/// `source` is the byte slice the tree was parsed from, used as the query's |
| 26 | +/// text provider for predicate evaluation. |
| 27 | +pub(crate) fn captures_from_tree( |
| 28 | + query: &Query, |
| 29 | + root: Node<'_>, |
| 30 | + source: &[u8], |
| 31 | +) -> Vec<HighlightSpan> { |
| 32 | + let names = query.capture_names(); |
| 33 | + let mut cursor = QueryCursor::new(); |
| 34 | + let mut spans = Vec::new(); |
| 35 | + let mut it = cursor.captures(query, root, source); |
| 36 | + while let Some((m, capture_index)) = it.next() { |
| 37 | + let cap = m.captures[*capture_index]; |
| 38 | + let name = names.get(cap.index as usize).copied().unwrap_or("unknown"); |
| 39 | + let node = cap.node; |
| 40 | + spans.push(HighlightSpan { |
| 41 | + start: node.start_byte(), |
| 42 | + end: node.end_byte(), |
| 43 | + capture: name.to_string(), |
| 44 | + }); |
| 45 | + } |
| 46 | + spans |
| 47 | +} |
| 48 | + |
| 49 | +/// Collapse nested/overlapping capture spans into a flat, non-overlapping, |
| 50 | +/// start-sorted run where the **innermost (narrowest) capture wins each byte**. |
| 51 | +/// |
| 52 | +/// Captures from one `Query` over one tree are nested-or-disjoint (CST node |
| 53 | +/// ranges never partially overlap), so for any byte the covering captures form |
| 54 | +/// a strict nesting chain and "narrowest" is unambiguous. The wider span is |
| 55 | +/// split around the narrower one. |
| 56 | +/// |
| 57 | +/// **Tie-break.** The one residual ambiguity is two captures at *equal extent* |
| 58 | +/// (identical start *and* end — two patterns matching the same node). This |
| 59 | +/// cannot arise from the structural nesting of a single tree, only from a query |
| 60 | +/// that captures one node twice. It is resolved deterministically: the capture |
| 61 | +/// appearing **later in the tree-sitter capture stream wins** (the stable sort |
| 62 | +/// below preserves capture-stream order for equal-extent spans, and the paint |
| 63 | +/// buffer lets the last-painted span win). No built-in-language corpus golden |
| 64 | +/// contains a genuine equal-extent collision — the |
| 65 | +/// `tests/fixtures/user-grammar-equal-extent` fixture exists to exercise this |
| 66 | +/// path. |
| 67 | +/// |
| 68 | +/// Zero-width spans (`start == end`) are dropped. The output is idempotent: |
| 69 | +/// `flatten_spans(flatten_spans(x)) == flatten_spans(x)`. |
| 70 | +/// |
| 71 | +/// This stays multi-line: a span crossing a newline is emitted as one span. |
| 72 | +/// The editor's LSP conversion splits per-line separately; the render/HTML path |
| 73 | +/// wants multi-line spans, so the split must **not** live here. |
| 74 | +pub fn flatten_spans(mut spans: Vec<HighlightSpan>) -> Vec<HighlightSpan> { |
| 75 | + spans.retain(|s| s.end > s.start); |
| 76 | + if spans.is_empty() { |
| 77 | + return Vec::new(); |
| 78 | + } |
| 79 | + |
| 80 | + let min = spans.iter().map(|s| s.start).min().expect("non-empty"); |
| 81 | + let max = spans.iter().map(|s| s.end).max().expect("non-empty"); |
| 82 | + |
| 83 | + // Stable sort: start ascending, then width descending so that at a given |
| 84 | + // start the wider span is painted first and the narrower overwrites it |
| 85 | + // (innermost-wins). Stability keeps equal-extent captures in capture-stream |
| 86 | + // order, so the last one painted (the later one in the stream) wins the tie. |
| 87 | + spans.sort_by(|a, b| { |
| 88 | + a.start |
| 89 | + .cmp(&b.start) |
| 90 | + .then_with(|| (b.end - b.start).cmp(&(a.end - a.start))) |
| 91 | + }); |
| 92 | + |
| 93 | + // Per-byte owner = index (into the sorted `spans`) of the winning capture. |
| 94 | + let len = max - min; |
| 95 | + let mut owner: Vec<Option<usize>> = vec![None; len]; |
| 96 | + for (idx, s) in spans.iter().enumerate() { |
| 97 | + for byte in &mut owner[(s.start - min)..(s.end - min)] { |
| 98 | + *byte = Some(idx); |
| 99 | + } |
| 100 | + } |
| 101 | + |
| 102 | + // Run-length encode contiguous same-owner runs, skipping uncovered gaps. |
| 103 | + // Keying runs on the owner *index* (not the capture string) keeps the |
| 104 | + // output idempotent: two adjacent flattened spans that happen to share a |
| 105 | + // capture name stay distinct, so re-flattening reproduces them exactly. |
| 106 | + let mut out: Vec<HighlightSpan> = Vec::new(); |
| 107 | + let mut run_owner: Option<usize> = None; |
| 108 | + let mut run_start = 0usize; |
| 109 | + for (offset, &o) in owner.iter().enumerate() { |
| 110 | + if o != run_owner { |
| 111 | + if let Some(idx) = run_owner { |
| 112 | + out.push(HighlightSpan { |
| 113 | + start: run_start + min, |
| 114 | + end: offset + min, |
| 115 | + capture: spans[idx].capture.clone(), |
| 116 | + }); |
| 117 | + } |
| 118 | + run_owner = o; |
| 119 | + run_start = offset; |
| 120 | + } |
| 121 | + } |
| 122 | + if let Some(idx) = run_owner { |
| 123 | + out.push(HighlightSpan { |
| 124 | + start: run_start + min, |
| 125 | + end: len + min, |
| 126 | + capture: spans[idx].capture.clone(), |
| 127 | + }); |
| 128 | + } |
| 129 | + out |
| 130 | +} |
| 131 | + |
| 132 | +#[cfg(test)] |
| 133 | +mod tests { |
| 134 | + use super::*; |
| 135 | + |
| 136 | + fn span(start: usize, end: usize, capture: &str) -> HighlightSpan { |
| 137 | + HighlightSpan { |
| 138 | + start, |
| 139 | + end, |
| 140 | + capture: capture.to_string(), |
| 141 | + } |
| 142 | + } |
| 143 | + |
| 144 | + #[test] |
| 145 | + fn flatten_innermost_wins_over_nested() { |
| 146 | + // Outer `string` [0,14] with an inner `variable` [5,9]: the inner |
| 147 | + // splits the outer into [0,5] and [9,14]. |
| 148 | + let out = flatten_spans(vec![span(0, 14, "string"), span(5, 9, "variable")]); |
| 149 | + assert_eq!( |
| 150 | + out, |
| 151 | + vec![ |
| 152 | + span(0, 5, "string"), |
| 153 | + span(5, 9, "variable"), |
| 154 | + span(9, 14, "string"), |
| 155 | + ] |
| 156 | + ); |
| 157 | + } |
| 158 | + |
| 159 | + #[test] |
| 160 | + fn flatten_drops_fully_covered_wrapper() { |
| 161 | + // `embedded` [3,9] is fully tiled by its children → it wins no byte |
| 162 | + // and disappears (the python f-string `{name}` case). |
| 163 | + let out = flatten_spans(vec![ |
| 164 | + span(0, 12, "string"), |
| 165 | + span(3, 9, "embedded"), |
| 166 | + span(3, 4, "punctuation.special"), |
| 167 | + span(4, 8, "variable"), |
| 168 | + span(8, 9, "punctuation.special"), |
| 169 | + ]); |
| 170 | + assert_eq!( |
| 171 | + out, |
| 172 | + vec![ |
| 173 | + span(0, 3, "string"), |
| 174 | + span(3, 4, "punctuation.special"), |
| 175 | + span(4, 8, "variable"), |
| 176 | + span(8, 9, "punctuation.special"), |
| 177 | + span(9, 12, "string"), |
| 178 | + ] |
| 179 | + ); |
| 180 | + } |
| 181 | + |
| 182 | + #[test] |
| 183 | + fn flatten_keeps_disjoint_spans() { |
| 184 | + let out = flatten_spans(vec![span(0, 3, "keyword"), span(4, 7, "function")]); |
| 185 | + assert_eq!(out, vec![span(0, 3, "keyword"), span(4, 7, "function")]); |
| 186 | + } |
| 187 | + |
| 188 | + #[test] |
| 189 | + fn flatten_sorts_and_drops_zero_width() { |
| 190 | + let out = flatten_spans(vec![span(4, 7, "b"), span(2, 2, "empty"), span(0, 3, "a")]); |
| 191 | + assert_eq!(out, vec![span(0, 3, "a"), span(4, 7, "b")]); |
| 192 | + } |
| 193 | + |
| 194 | + #[test] |
| 195 | + fn flatten_is_idempotent() { |
| 196 | + let input = vec![ |
| 197 | + span(0, 14, "property"), |
| 198 | + span(0, 4, "type"), |
| 199 | + span(5, 6, "operator"), |
| 200 | + span(7, 14, "string"), |
| 201 | + ]; |
| 202 | + let once = flatten_spans(input); |
| 203 | + let twice = flatten_spans(once.clone()); |
| 204 | + assert_eq!(once, twice); |
| 205 | + } |
| 206 | + |
| 207 | + #[test] |
| 208 | + fn flatten_handles_equal_extent_tie() { |
| 209 | + // Two captures at the identical range collapse to exactly one span. |
| 210 | + // Later-in-stream wins (here: `property`, appended second). |
| 211 | + let out = flatten_spans(vec![span(0, 4, "type"), span(0, 4, "property")]); |
| 212 | + assert_eq!(out, vec![span(0, 4, "property")]); |
| 213 | + } |
| 214 | + |
| 215 | + #[test] |
| 216 | + fn flatten_empty() { |
| 217 | + assert!(flatten_spans(vec![]).is_empty()); |
| 218 | + } |
| 219 | +} |
0 commit comments