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}