1use crate::payload::{DiagnosticSpan, Utf16Position};
2use std::sync::{Arc, OnceLock};
3
4#[derive(Debug, Clone, Copy, PartialEq, Eq)]
5pub struct LineCol {
6 pub line: usize,
7 pub column: usize,
8}
9
10impl LineCol {
11 pub const fn new(line: usize, column: usize) -> Self {
12 Self { line, column }
13 }
14}
15
16#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
17pub enum SourceMapError {
18 #[error("byte offset {offset} is outside source length {source_len}")]
19 OffsetOutOfBounds { offset: usize, source_len: usize },
20 #[error("byte offset {offset} is not a UTF-8 character boundary")]
21 OffsetNotCharBoundary { offset: usize },
22 #[error("range start {start} is after range end {end}")]
23 ReversedRange { start: usize, end: usize },
24}
25
26#[derive(Debug, Clone)]
27pub struct SourceMap {
28 source: Arc<str>,
29 line_starts: Arc<[usize]>,
30 line_metrics: Arc<[OnceLock<LineMetric>]>,
31}
32
33impl SourceMap {
34 pub fn new(source: impl Into<Arc<str>>) -> Self {
35 let source = source.into();
36 let line_starts = line_starts(source.as_ref());
37 let line_metrics = (0..line_starts.len())
38 .map(|_| OnceLock::new())
39 .collect::<Vec<_>>();
40 Self {
41 source,
42 line_starts: Arc::from(line_starts.into_boxed_slice()),
43 line_metrics: Arc::from(line_metrics.into_boxed_slice()),
44 }
45 }
46
47 pub fn source(&self) -> &str {
48 self.source.as_ref()
49 }
50
51 pub fn source_arc(&self) -> Arc<str> {
52 Arc::clone(&self.source)
53 }
54
55 pub fn source_len(&self) -> usize {
56 self.source.len()
57 }
58
59 pub fn line_starts(&self) -> &[usize] {
60 &self.line_starts
61 }
62
63 pub fn line_col(&self, offset: usize) -> Result<LineCol, SourceMapError> {
64 let metrics = self.offset_metrics(offset)?;
65 Ok(LineCol::new(
66 metrics.line_index + 1,
67 metrics.char_column + 1,
68 ))
69 }
70
71 pub fn utf16_position(&self, offset: usize) -> Result<Utf16Position, SourceMapError> {
72 let metrics = self.offset_metrics(offset)?;
73 Ok(Utf16Position {
74 line: metrics.line_index,
75 character: metrics.utf16_column,
76 })
77 }
78
79 pub fn span(&self, start: usize, end: usize) -> Result<DiagnosticSpan, SourceMapError> {
80 if start > end {
81 return Err(SourceMapError::ReversedRange { start, end });
82 }
83
84 let start_lc = self.line_col(start)?;
85 let end_lc = self.line_col(end)?;
86 let lsp_start = self.utf16_position(start)?;
87 let lsp_end = self.utf16_position(end)?;
88
89 Ok(DiagnosticSpan::new(
90 start,
91 end,
92 start_lc.line,
93 start_lc.column,
94 end_lc.line,
95 end_lc.column,
96 lsp_start,
97 lsp_end,
98 ))
99 }
100
101 pub fn whole_source_span(&self) -> Result<DiagnosticSpan, SourceMapError> {
102 self.span(0, self.source.len())
103 }
104
105 pub fn line_bounds(&self, line_index: usize) -> Option<(usize, usize)> {
106 let line = self.line_metric(line_index)?;
107 Some((line.start, line.content_end))
108 }
109
110 pub fn byte_offset_for_utf16_position(&self, position: Utf16Position) -> Option<usize> {
111 let line = self.line_metric(position.line)?;
112 match line.utf16_columns.binary_search(&position.character) {
113 Ok(boundary_index) => Some(line.start + line.byte_boundaries[boundary_index]),
114 Err(boundary_index) if boundary_index >= line.utf16_columns.len() => {
115 Some(line.content_end)
116 }
117 Err(_) => None,
118 }
119 }
120
121 fn validate_offset(&self, offset: usize) -> Result<(), SourceMapError> {
122 if offset > self.source.len() {
123 return Err(SourceMapError::OffsetOutOfBounds {
124 offset,
125 source_len: self.source.len(),
126 });
127 }
128 if !self.source.is_char_boundary(offset) {
129 return Err(SourceMapError::OffsetNotCharBoundary { offset });
130 }
131 Ok(())
132 }
133
134 fn line_index_for_offset(&self, offset: usize) -> usize {
135 match self.line_starts.binary_search(&offset) {
136 Ok(index) => index,
137 Err(0) => 0,
138 Err(index) => index - 1,
139 }
140 }
141
142 fn offset_metrics(&self, offset: usize) -> Result<OffsetMetrics, SourceMapError> {
143 self.validate_offset(offset)?;
144 let line_index = self.line_index_for_offset(offset);
145 let line = self
146 .line_metric(line_index)
147 .expect("validated source offset should map to a cached line");
148 let clamped = offset.clamp(line.start, line.content_end);
149 let relative = clamped - line.start;
150 let boundary_index = line
151 .byte_boundaries
152 .binary_search(&relative)
153 .expect("validated source offset should map to a cached line boundary");
154
155 Ok(OffsetMetrics {
156 line_index,
157 char_column: boundary_index,
158 utf16_column: line.utf16_columns[boundary_index],
159 })
160 }
161
162 fn line_metric(&self, line_index: usize) -> Option<&LineMetric> {
163 let slot = self.line_metrics.get(line_index)?;
164 Some(slot.get_or_init(|| {
165 let start = self.line_starts[line_index];
166 let next_start = self
167 .line_starts
168 .get(line_index + 1)
169 .copied()
170 .unwrap_or(self.source.len());
171 line_metric(self.source.as_ref(), start, next_start)
172 }))
173 }
174
175 #[cfg(test)]
176 fn cached_line_metric_count(&self) -> usize {
177 self.line_metrics
178 .iter()
179 .filter(|metric| metric.get().is_some())
180 .count()
181 }
182
183 #[cfg(test)]
184 fn cached_line_boundary_count(&self, line_index: usize) -> Option<usize> {
185 self.line_metrics
186 .get(line_index)
187 .and_then(OnceLock::get)
188 .map(|line| line.byte_boundaries.len())
189 }
190}
191
192#[derive(Debug, Clone)]
193struct LineMetric {
194 start: usize,
195 content_end: usize,
196 byte_boundaries: Vec<usize>,
197 utf16_columns: Vec<usize>,
198}
199
200#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201struct OffsetMetrics {
202 line_index: usize,
203 char_column: usize,
204 utf16_column: usize,
205}
206
207pub(crate) fn whole_text_span_without_source_copy(text: &str) -> DiagnosticSpan {
208 let mut end_line = 1usize;
209 let mut end_column = 1usize;
210 let mut end_lsp_line = 0usize;
211 let mut end_lsp_character = 0usize;
212 let mut chars = text.chars().peekable();
213
214 while let Some(ch) = chars.next() {
215 match ch {
216 '\r' => {
217 if chars.peek() == Some(&'\n') {
218 chars.next();
219 }
220 end_line += 1;
221 end_column = 1;
222 end_lsp_line += 1;
223 end_lsp_character = 0;
224 }
225 '\n' => {
226 end_line += 1;
227 end_column = 1;
228 end_lsp_line += 1;
229 end_lsp_character = 0;
230 }
231 _ => {
232 end_column += 1;
233 end_lsp_character += ch.len_utf16();
234 }
235 }
236 }
237
238 DiagnosticSpan::new(
239 0,
240 text.len(),
241 1,
242 1,
243 end_line,
244 end_column,
245 Utf16Position {
246 line: 0,
247 character: 0,
248 },
249 Utf16Position {
250 line: end_lsp_line,
251 character: end_lsp_character,
252 },
253 )
254}
255
256fn line_starts(source: &str) -> Vec<usize> {
257 let mut starts = vec![0];
258 let bytes = source.as_bytes();
259 let mut idx = 0usize;
260 while idx < bytes.len() {
261 match bytes[idx] {
262 b'\r' => {
263 idx += 1;
264 if bytes.get(idx) == Some(&b'\n') {
265 idx += 1;
266 }
267 starts.push(idx);
268 }
269 b'\n' => {
270 idx += 1;
271 starts.push(idx);
272 }
273 _ => {
274 idx += 1;
275 }
276 }
277 }
278 starts
279}
280
281fn line_metric(source: &str, start: usize, next_start: usize) -> LineMetric {
282 let content_end = line_content_end(source.as_bytes(), start, next_start);
283 let line = &source[start..content_end];
284 let mut byte_boundaries = Vec::with_capacity(line.chars().count() + 1);
285 let mut utf16_columns = Vec::with_capacity(byte_boundaries.capacity());
286 let mut utf16 = 0usize;
287
288 byte_boundaries.push(0);
289 utf16_columns.push(0);
290
291 for (relative, ch) in line.char_indices() {
292 utf16 += ch.len_utf16();
293 byte_boundaries.push(relative + ch.len_utf8());
294 utf16_columns.push(utf16);
295 }
296
297 LineMetric {
298 start,
299 content_end,
300 byte_boundaries,
301 utf16_columns,
302 }
303}
304
305fn line_content_end(bytes: &[u8], start: usize, next_start: usize) -> usize {
306 let mut end = next_start;
307 if end > start && bytes.get(end - 1) == Some(&b'\n') {
308 end -= 1;
309 }
310 if end > start && bytes.get(end - 1) == Some(&b'\r') {
311 end -= 1;
312 }
313 end
314}
315
316#[cfg(test)]
317mod tests {
318 use super::*;
319
320 #[test]
321 fn maps_ascii_offsets_to_one_based_cli_positions() {
322 let map = SourceMap::new("flowchart TD\nA-->B\n");
323
324 assert_eq!(map.line_col(0).unwrap(), LineCol::new(1, 1));
325 assert_eq!(map.line_col(13).unwrap(), LineCol::new(2, 1));
326 assert_eq!(map.line_col(map.source_len()).unwrap(), LineCol::new(3, 1));
327 }
328
329 #[test]
330 fn maps_utf8_offsets_to_lsp_utf16_positions() {
331 let map = SourceMap::new("flowchart TD\nA[🤓]-->B\n");
332 let emoji_start = map.source().find('🤓').unwrap();
333 let emoji_end = emoji_start + "🤓".len();
334 let after_bracket = emoji_end + 1;
335
336 assert_eq!(
337 map.utf16_position(emoji_start).unwrap(),
338 Utf16Position {
339 line: 1,
340 character: 2
341 }
342 );
343 assert_eq!(
344 map.utf16_position(after_bracket).unwrap(),
345 Utf16Position {
346 line: 1,
347 character: 5
348 }
349 );
350 }
351
352 #[test]
353 fn crlf_line_bounds_and_positions_ignore_carriage_return() {
354 let source = "flowchart TD\r\nA[🤓]-->B\r\n";
355 let map = SourceMap::new(source);
356 let first_cr = source.find('\r').unwrap();
357 let first_lf = source.find('\n').unwrap();
358
359 assert_eq!(map.line_bounds(0), Some((0, first_cr)));
360 assert_eq!(
361 map.utf16_position(first_cr).unwrap(),
362 Utf16Position {
363 line: 0,
364 character: "flowchart TD".len(),
365 }
366 );
367 assert_eq!(
368 map.utf16_position(first_lf).unwrap(),
369 Utf16Position {
370 line: 0,
371 character: "flowchart TD".len(),
372 }
373 );
374 assert_eq!(
375 map.byte_offset_for_utf16_position(Utf16Position {
376 line: 0,
377 character: "flowchart TD".len(),
378 }),
379 Some(first_cr)
380 );
381 }
382
383 #[test]
384 fn bare_cr_line_bounds_and_positions_treat_carriage_return_as_line_ending() {
385 let source = "flowchart TD\rA-->B\rC-->D";
386 let map = SourceMap::new(source);
387 let first_cr = source.find('\r').unwrap();
388 let second_line_start = first_cr + 1;
389 let second_cr = source[second_line_start..].find('\r').unwrap() + second_line_start;
390
391 assert_eq!(map.line_bounds(0), Some((0, first_cr)));
392 assert_eq!(map.line_bounds(1), Some((second_line_start, second_cr)));
393 assert_eq!(
394 map.utf16_position(second_line_start).unwrap(),
395 Utf16Position {
396 line: 1,
397 character: 0,
398 }
399 );
400 assert_eq!(
401 map.byte_offset_for_utf16_position(Utf16Position {
402 line: 2,
403 character: 0,
404 }),
405 Some(second_cr + 1)
406 );
407 }
408
409 #[test]
410 fn utf16_position_past_line_end_clamps_to_content_end() {
411 let source = "flowchart TD\nA[🤓]-->B\n";
412 let map = SourceMap::new(source);
413 let second_line_start = source.find("A[").unwrap();
414 let second_line_end = source[second_line_start..].find('\n').unwrap() + second_line_start;
415
416 assert_eq!(
417 map.byte_offset_for_utf16_position(Utf16Position {
418 line: 1,
419 character: 10_000,
420 }),
421 Some(second_line_end)
422 );
423 }
424
425 #[test]
426 fn dense_span_conversion_uses_cached_line_metrics() {
427 let nodes = (0..512)
428 .map(|index| format!("N{index}[🤓]"))
429 .collect::<Vec<_>>()
430 .join(" ");
431 let source = format!("flowchart TD {nodes}");
432 let map = SourceMap::new(source.clone());
433
434 assert_eq!(map.cached_line_metric_count(), 0);
435 assert_eq!(map.cached_line_boundary_count(0), None);
436
437 for offset in source.match_indices('N').map(|(offset, _)| offset) {
438 let end = source[offset..].find('[').map(|len| offset + len).unwrap();
439 let span = map.span(offset, end).unwrap();
440 assert_eq!(span.lsp_range.start.line, 0);
441 assert!(span.lsp_range.end.character > span.lsp_range.start.character);
442 }
443 assert_eq!(map.cached_line_metric_count(), 1);
444 assert_eq!(
445 map.cached_line_boundary_count(0),
446 Some(source.chars().count() + 1)
447 );
448 }
449
450 #[test]
451 fn rejects_non_char_boundary_offsets() {
452 let map = SourceMap::new("flowchart TD\nA[🤓]\n");
453 let inside_emoji = map.source().find('🤓').unwrap() + 1;
454
455 assert_eq!(
456 map.line_col(inside_emoji).unwrap_err(),
457 SourceMapError::OffsetNotCharBoundary {
458 offset: inside_emoji
459 }
460 );
461 }
462
463 #[test]
464 fn builds_diagnostic_span_with_cli_and_lsp_positions() {
465 let map = SourceMap::new("flowchart TD\nA[🤓]-->B\n");
466 let start = map.source().find('A').unwrap();
467 let end = map.source().find("-->").unwrap();
468 let span = map.span(start, end).unwrap();
469
470 assert_eq!(span.byte_start, start);
471 assert_eq!(span.byte_end, end);
472 assert_eq!(span.line, 2);
473 assert_eq!(span.column, 1);
474 assert_eq!(span.end_line, 2);
475 assert_eq!(span.end_column, 5);
476 assert_eq!(span.lsp_range.start.line, 1);
477 assert_eq!(span.lsp_range.start.character, 0);
478 assert_eq!(span.lsp_range.end.character, 5);
479 }
480
481 #[test]
482 fn whole_text_span_without_source_copy_matches_source_map_span() {
483 for source in [
484 "flowchart TD\nA[🤓]-->B\n",
485 "flowchart TD\r\nA[🤓]-->B",
486 "flowchart TD\r\nA[🤓]-->B\r",
487 "flowchart TD\r\r\nA[🤓]-->B",
488 ] {
489 assert_eq!(
490 whole_text_span_without_source_copy(source),
491 SourceMap::new(source).whole_source_span().unwrap()
492 );
493 }
494 }
495}