Skip to main content

thalir_core/
source_location.rs

1use serde::{Deserialize, Serialize};
2use std::path::{Path, PathBuf};
3use std::sync::{Arc, RwLock};
4use thiserror::Error;
5
6pub const MAX_LINES: usize = 1_000_000;
7
8pub const MAX_FILE_SIZE: usize = u32::MAX as usize;
9
10#[derive(Error, Debug)]
11pub enum SourceLocationError {
12    #[error("File exceeds maximum line count: {0} lines")]
13    TooManyLines(usize),
14    #[error("File exceeds maximum size: {0} bytes")]
15    FileTooLarge(usize),
16
17    #[error("Invalid UTF-8 at byte offset {0}")]
18    InvalidUtf8(usize),
19    #[error("Invalid file ID: {0}")]
20    InvalidFileId(u32),
21    #[error("Invalid source span: file_id={0}, start={1}, len={2}")]
22    InvalidSpan(u32, u32, u32),
23    #[error("IO error: {0}")]
24    Io(#[from] std::io::Error),
25}
26
27pub type Result<T> = std::result::Result<T, SourceLocationError>;
28
29#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
30pub struct SourceSpan {
31    pub file_id: u32,
32    pub start: u32,
33    pub len: u32,
34}
35
36pub const INVALID_SPAN: SourceSpan = SourceSpan {
37    file_id: u32::MAX,
38    start: 0,
39    len: 0,
40};
41
42impl SourceSpan {
43    pub fn new(file_id: u32, start: u32, len: u32) -> Self {
44        Self {
45            file_id,
46            start,
47            len,
48        }
49    }
50
51    #[inline]
52    pub fn is_valid(&self) -> bool {
53        self.file_id != u32::MAX && self.len > 0
54    }
55
56    #[inline]
57    pub fn end(&self) -> u32 {
58        self.start.saturating_add(self.len)
59    }
60
61    pub fn contains(&self, other: &SourceSpan) -> bool {
62        self.file_id == other.file_id && other.start >= self.start && other.end() <= self.end()
63    }
64
65    pub fn merge(&self, other: &SourceSpan) -> Option<SourceSpan> {
66        if self.file_id != other.file_id {
67            return None;
68        }
69
70        let start = self.start.min(other.start);
71        let end = self.end().max(other.end());
72
73        Some(SourceSpan::new(
74            self.file_id,
75            start,
76            end.saturating_sub(start),
77        ))
78    }
79
80    pub fn merge_all(spans: &[Option<SourceSpan>]) -> Option<SourceSpan> {
81        const MAX_SPAN_MERGE: usize = 1000;
82
83        let valid_spans: Vec<_> = spans
84            .iter()
85            .filter_map(|s| s.filter(|span| span.is_valid()))
86            .take(MAX_SPAN_MERGE)
87            .collect();
88
89        if valid_spans.is_empty() {
90            return None;
91        }
92
93        let first = valid_spans[0];
94        valid_spans
95            .iter()
96            .skip(1)
97            .try_fold(first, |acc, span| acc.merge(span))
98    }
99}
100
101impl Default for SourceSpan {
102    fn default() -> Self {
103        INVALID_SPAN
104    }
105}
106
107#[derive(Debug, Clone)]
108pub struct SourceFiles {
109    files: Arc<RwLock<Vec<SourceFile>>>,
110}
111
112impl SourceFiles {
113    pub fn new() -> Self {
114        Self {
115            files: Arc::new(RwLock::new(Vec::new())),
116        }
117    }
118
119    pub fn add_file(&self, path: PathBuf, text: String) -> Result<u32> {
120        if text.len() > MAX_FILE_SIZE {
121            return Err(SourceLocationError::FileTooLarge(text.len()));
122        }
123
124        let line_starts = compute_line_starts(&text)?;
125
126        let mut files = self.files.write().unwrap();
127        let file_id = files.len() as u32;
128
129        files.push(SourceFile {
130            path,
131            text: Arc::new(text),
132            line_starts: Arc::new(line_starts),
133        });
134
135        Ok(file_id)
136    }
137
138    pub fn get_file(&self, file_id: u32) -> Option<SourceFile> {
139        let files = self.files.read().unwrap();
140        files.get(file_id as usize).cloned()
141    }
142
143    pub fn to_line_col(&self, span: SourceSpan) -> Option<(u32, u32)> {
144        if !span.is_valid() {
145            return None;
146        }
147
148        let file = self.get_file(span.file_id)?;
149        let line_starts = file.line_starts.as_ref();
150
151        let line_idx = line_starts
152            .partition_point(|&start| start <= span.start as usize)
153            .saturating_sub(1);
154
155        if line_idx >= line_starts.len() {
156            return None;
157        }
158
159        let line_start = line_starts[line_idx];
160        let line = (line_idx + 1) as u32;
161
162        let column = file.text[line_start..span.start as usize].chars().count() as u32 + 1;
163
164        Some((line, column))
165    }
166
167    pub fn snippet(&self, span: SourceSpan, context_lines: usize) -> Option<String> {
168        if !span.is_valid() {
169            return None;
170        }
171
172        let file = self.get_file(span.file_id)?;
173        let (line, _col) = self.to_line_col(span)?;
174
175        let start_line = (line as usize).saturating_sub(context_lines);
176        let end_line = (line as usize).saturating_add(context_lines);
177
178        let mut snippet = String::with_capacity(1024);
179
180        for line_no in start_line..=end_line {
181            if let Some(line_text) = file.get_line(line_no) {
182                snippet.push_str(&format!("{:4} | {}\n", line_no, line_text));
183            }
184        }
185
186        Some(snippet)
187    }
188
189    pub fn relative_path(&self, file_id: u32, workspace_root: &Path) -> Option<PathBuf> {
190        let file = self.get_file(file_id)?;
191
192        file.path
193            .strip_prefix(workspace_root)
194            .ok()
195            .map(|p| p.to_path_buf())
196            .or_else(|| Some(PathBuf::from(format!("external/file_{}", file_id))))
197    }
198
199    pub fn file_count(&self) -> usize {
200        self.files.read().unwrap().len()
201    }
202
203    pub fn get_all_files(&self) -> Vec<SourceFile> {
204        self.files.read().unwrap().clone()
205    }
206}
207
208impl Default for SourceFiles {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214#[derive(Debug, Clone)]
215pub struct SourceFile {
216    pub path: PathBuf,
217    pub text: Arc<String>,
218    pub line_starts: Arc<Vec<usize>>,
219}
220
221impl SourceFile {
222    pub fn get_line(&self, line: usize) -> Option<&str> {
223        if line == 0 {
224            return None;
225        }
226
227        let line_idx = line - 1;
228        if line_idx >= self.line_starts.len() {
229            return None;
230        }
231
232        let start = self.line_starts[line_idx];
233        let end = if line_idx + 1 < self.line_starts.len() {
234            self.line_starts[line_idx + 1]
235        } else {
236            self.text.len()
237        };
238
239        let text_bytes = self.text.as_bytes();
240        let start_boundary = find_char_boundary(text_bytes, start);
241        let end_boundary = find_char_boundary(text_bytes, end);
242
243        std::str::from_utf8(&text_bytes[start_boundary..end_boundary])
244            .ok()
245            .map(|s| s.trim_end_matches(&['\r', '\n'][..]))
246    }
247
248    pub fn line_count(&self) -> usize {
249        self.line_starts.len()
250    }
251}
252
253fn compute_line_starts(text: &str) -> Result<Vec<usize>> {
254    let mut line_starts = Vec::with_capacity(1000);
255    line_starts.push(0);
256
257    let bytes = text.as_bytes();
258    let mut i = 0;
259
260    while i < bytes.len() {
261        if line_starts.len() >= MAX_LINES {
262            return Err(SourceLocationError::TooManyLines(line_starts.len()));
263        }
264
265        match bytes[i] {
266            b'\r' => {
267                if i + 1 < bytes.len() && bytes[i + 1] == b'\n' {
268                    i += 2;
269                } else {
270                    i += 1;
271                }
272                line_starts.push(i);
273            }
274            b'\n' => {
275                i += 1;
276                line_starts.push(i);
277            }
278            _ => {
279                i += 1;
280            }
281        }
282    }
283
284    Ok(line_starts)
285}
286
287fn find_char_boundary(bytes: &[u8], mut index: usize) -> usize {
288    if index >= bytes.len() {
289        return bytes.len();
290    }
291
292    while index > 0 && (bytes[index] & 0b1100_0000) == 0b1000_0000 {
293        index -= 1;
294    }
295
296    index
297}
298
299#[derive(Debug, Clone)]
300pub struct SourceInfo {
301    pub file_id: u32,
302    pub path: PathBuf,
303    pub line: u32,
304    pub column: u32,
305    pub length: u32,
306}
307
308impl SourceFiles {
309    pub fn get_source_info(&self, span: SourceSpan) -> Option<SourceInfo> {
310        let (line, column) = self.to_line_col(span)?;
311        let file = self.get_file(span.file_id)?;
312
313        Some(SourceInfo {
314            file_id: span.file_id,
315            path: file.path.clone(),
316            line,
317            column,
318            length: span.len,
319        })
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326
327    #[test]
328    fn test_source_span_basics() {
329        let span = SourceSpan::new(0, 10, 20);
330        assert!(span.is_valid());
331        assert_eq!(span.end(), 30);
332
333        assert_eq!(INVALID_SPAN.is_valid(), false);
334    }
335
336    #[test]
337    fn test_span_merge() {
338        let span1 = SourceSpan::new(0, 10, 5);
339        let span2 = SourceSpan::new(0, 20, 5);
340
341        let merged = span1.merge(&span2).unwrap();
342        assert_eq!(merged.start, 10);
343        assert_eq!(merged.len, 15);
344
345        let span3 = SourceSpan::new(1, 10, 5);
346        assert!(span1.merge(&span3).is_none());
347    }
348
349    #[test]
350    fn test_span_contains() {
351        let outer = SourceSpan::new(0, 10, 20);
352        let inner = SourceSpan::new(0, 15, 5);
353
354        assert!(outer.contains(&inner));
355        assert!(!inner.contains(&outer));
356    }
357
358    #[test]
359    fn test_line_starts_unix() {
360        let text = "line 1\nline 2\nline 3";
361        let line_starts = compute_line_starts(text).unwrap();
362        assert_eq!(line_starts, vec![0, 7, 14]);
363    }
364
365    #[test]
366    fn test_line_starts_windows() {
367        let text = "line 1\r\nline 2\r\nline 3";
368        let line_starts = compute_line_starts(text).unwrap();
369        assert_eq!(line_starts, vec![0, 8, 16]);
370    }
371
372    #[test]
373    fn test_line_starts_mixed() {
374        let text = "line 1\nline 2\r\nline 3\rline 4";
375        let line_starts = compute_line_starts(text).unwrap();
376        assert_eq!(line_starts, vec![0, 7, 15, 22]);
377    }
378
379    #[test]
380    fn test_to_line_col() {
381        let files = SourceFiles::new();
382        let file_id = files
383            .add_file(
384                PathBuf::from("test.sol"),
385                "line 1\nline 2\nline 3".to_string(),
386            )
387            .unwrap();
388
389        let span = SourceSpan::new(file_id, 7, 6);
390        let (line, col) = files.to_line_col(span).unwrap();
391        assert_eq!(line, 2);
392        assert_eq!(col, 1);
393
394        let span = SourceSpan::new(file_id, 10, 1);
395        let (line, col) = files.to_line_col(span).unwrap();
396        assert_eq!(line, 2);
397        assert_eq!(col, 4);
398    }
399
400    #[test]
401    fn test_to_line_col_zero_offset() {
402        let files = SourceFiles::new();
403        let file_id = files
404            .add_file(PathBuf::from("test.sol"), "hello\nworld".to_string())
405            .unwrap();
406
407        let span = SourceSpan::new(file_id, 0, 5);
408        let (line, col) = files.to_line_col(span).unwrap();
409        assert_eq!(line, 1);
410        assert_eq!(col, 1);
411    }
412
413    #[test]
414    fn test_snippet_extraction() {
415        let files = SourceFiles::new();
416        let file_id = files
417            .add_file(
418                PathBuf::from("test.sol"),
419                "line 1\nline 2\nline 3\nline 4\nline 5".to_string(),
420            )
421            .unwrap();
422
423        let span = SourceSpan::new(file_id, 7, 6);
424        let snippet = files.snippet(span, 1).unwrap();
425
426        assert!(snippet.contains("1 | line 1"));
427        assert!(snippet.contains("2 | line 2"));
428        assert!(snippet.contains("3 | line 3"));
429        assert!(!snippet.contains("4 | line 4"));
430    }
431
432    #[test]
433    fn test_unicode_handling() {
434        let files = SourceFiles::new();
435        let text = "hello ๐Ÿ‘จโ€๐Ÿ‘ฉโ€๐Ÿ‘งโ€๐Ÿ‘ฆ world";
436        let file_id = files
437            .add_file(PathBuf::from("unicode.sol"), text.to_string())
438            .unwrap();
439
440        let span = SourceSpan::new(file_id, 6, 25);
441        let (line, col) = files.to_line_col(span).unwrap();
442        assert_eq!(line, 1);
443        assert!(col > 1);
444    }
445
446    #[test]
447    fn test_resource_limits() {
448        let huge_text = "x\n".repeat(MAX_LINES + 1);
449        let result = compute_line_starts(&huge_text);
450        assert!(matches!(result, Err(SourceLocationError::TooManyLines(_))));
451    }
452
453    #[test]
454    fn test_get_line() {
455        let files = SourceFiles::new();
456        let file_id = files
457            .add_file(
458                PathBuf::from("test.sol"),
459                "line 1\nline 2\nline 3".to_string(),
460            )
461            .unwrap();
462
463        let file = files.get_file(file_id).unwrap();
464
465        assert_eq!(file.get_line(1).unwrap(), "line 1");
466        assert_eq!(file.get_line(2).unwrap(), "line 2");
467        assert_eq!(file.get_line(3).unwrap(), "line 3");
468        assert!(file.get_line(4).is_none());
469        assert!(file.get_line(0).is_none());
470    }
471
472    #[test]
473    fn test_thread_safety() {
474        use std::thread;
475
476        let files = SourceFiles::new();
477        let file_id = files
478            .add_file(
479                PathBuf::from("test.sol"),
480                "line 1\nline 2\nline 3".to_string(),
481            )
482            .unwrap();
483
484        let handles: Vec<_> = (0..10)
485            .map(|i| {
486                let files_clone = files.clone();
487                thread::spawn(move || {
488                    let span = SourceSpan::new(file_id, (i % 3) * 7, 6);
489                    files_clone.to_line_col(span)
490                })
491            })
492            .collect();
493
494        for handle in handles {
495            handle.join().unwrap();
496        }
497    }
498
499    #[test]
500    fn test_relative_path_anonymization() {
501        let files = SourceFiles::new();
502        let file_id = files
503            .add_file(
504                PathBuf::from("/home/user/secret/contract.sol"),
505                "test".to_string(),
506            )
507            .unwrap();
508
509        let workspace = Path::new("/home/user/project");
510        let relative = files.relative_path(file_id, workspace).unwrap();
511
512        assert!(relative.to_string_lossy().contains("external"));
513    }
514}