1use std::fs::File;
19use std::io::{self, BufReader, BufWriter, Read, Write};
20use std::path::Path;
21
22use memmap2::Mmap;
23
24use super::index::{SemanticIndex, StringTable};
25use super::types::{Edge, Reference, Scope, Symbol, Token};
26use crate::core::error::{Error, Result};
27
28const MAGIC: [u8; 8] = *b"GRPTRACE";
34
35const VERSION: u32 = 1;
37
38const HEADER_SIZE: usize = 64;
40
41#[derive(Debug, Clone, Copy)]
61#[repr(C)]
62struct Header {
63 magic: [u8; 8],
65 version: u32,
67 _reserved: u32,
69 symbol_count: u32,
71 token_count: u32,
73 reference_count: u32,
75 scope_count: u32,
77 edge_count: u32,
79 file_count: u32,
81 string_size: u32,
83 _padding: [u8; 20],
85}
86
87impl Header {
88 fn new(index: &SemanticIndex) -> Self {
89 Self {
90 magic: MAGIC,
91 version: VERSION,
92 _reserved: 0,
93 symbol_count: index.symbols.len() as u32,
94 token_count: index.tokens.len() as u32,
95 reference_count: index.references.len() as u32,
96 scope_count: index.scopes.len() as u32,
97 edge_count: index.edges.len() as u32,
98 file_count: index.files.len() as u32,
99 string_size: index.strings.byte_size() as u32,
100 _padding: [0; 20],
101 }
102 }
103
104 fn validate(&self) -> Result<()> {
105 if self.magic != MAGIC {
106 return Err(Error::IndexError {
107 message: "Invalid trace index file (bad magic)".into(),
108 });
109 }
110 let version = self.version;
111 if version != VERSION {
112 return Err(Error::IndexError {
113 message: format!(
114 "Unsupported trace index version {} (expected {})",
115 version, VERSION
116 ),
117 });
118 }
119 Ok(())
120 }
121
122 fn as_bytes(&self) -> &[u8] {
123 unsafe {
124 std::slice::from_raw_parts(
125 self as *const Self as *const u8,
126 std::mem::size_of::<Self>(),
127 )
128 }
129 }
130
131 fn from_bytes(bytes: &[u8]) -> Result<Self> {
132 if bytes.len() < std::mem::size_of::<Self>() {
133 return Err(Error::IndexError {
134 message: "Invalid trace index file (header too small)".into(),
135 });
136 }
137
138 let header = unsafe { std::ptr::read_unaligned(bytes.as_ptr() as *const Self) };
140
141 header.validate()?;
142 Ok(header)
143 }
144}
145
146const _: () = {
148 assert!(std::mem::size_of::<Header>() == HEADER_SIZE);
149};
150
151pub fn save_index(index: &SemanticIndex, path: impl AsRef<Path>) -> Result<()> {
157 let path = path.as_ref();
158 let file = File::create(path)?;
159 let mut writer = BufWriter::with_capacity(64 * 1024, file);
160
161 let header = Header::new(index);
163 writer.write_all(header.as_bytes())?;
164
165 write_slice(&mut writer, &index.symbols)?;
167
168 write_slice(&mut writer, &index.tokens)?;
170
171 write_slice(&mut writer, &index.references)?;
173
174 write_slice(&mut writer, &index.scopes)?;
176
177 write_slice(&mut writer, &index.edges)?;
179
180 for path in &index.files {
182 let path_bytes = path.to_string_lossy().as_bytes().to_vec();
183 let len = path_bytes.len() as u32;
184 writer.write_all(&len.to_le_bytes())?;
185 writer.write_all(&path_bytes)?;
186 }
187
188 writer.write_all(index.strings.as_bytes())?;
190
191 writer.flush()?;
192 Ok(())
193}
194
195fn write_slice<T, W: Write>(writer: &mut W, slice: &[T]) -> io::Result<()> {
197 let bytes = unsafe {
198 std::slice::from_raw_parts(
199 slice.as_ptr() as *const u8,
200 slice.len() * std::mem::size_of::<T>(),
201 )
202 };
203 writer.write_all(bytes)
204}
205
206pub fn load_index(path: impl AsRef<Path>) -> Result<SemanticIndex> {
214 let path = path.as_ref();
215 let file = File::open(path)?;
216 let mmap = unsafe { Mmap::map(&file)? };
217
218 let header = Header::from_bytes(&mmap)?;
220 let mut offset = HEADER_SIZE;
221
222 let symbols: Vec<Symbol> = read_vec(&mmap, &mut offset, header.symbol_count as usize)?;
224
225 let tokens: Vec<Token> = read_vec(&mmap, &mut offset, header.token_count as usize)?;
227
228 let references: Vec<Reference> = read_vec(&mmap, &mut offset, header.reference_count as usize)?;
230
231 let scopes: Vec<Scope> = read_vec(&mmap, &mut offset, header.scope_count as usize)?;
233
234 let edges: Vec<Edge> = read_vec(&mmap, &mut offset, header.edge_count as usize)?;
236
237 let mut files = Vec::with_capacity(header.file_count as usize);
239 for _ in 0..header.file_count {
240 if offset + 4 > mmap.len() {
241 return Err(Error::IndexError {
242 message: "Truncated trace index file (files section)".into(),
243 });
244 }
245 let len = u32::from_le_bytes([
246 mmap[offset],
247 mmap[offset + 1],
248 mmap[offset + 2],
249 mmap[offset + 3],
250 ]) as usize;
251 offset += 4;
252
253 if offset + len > mmap.len() {
254 return Err(Error::IndexError {
255 message: "Truncated trace index file (file path)".into(),
256 });
257 }
258 let path_str =
259 std::str::from_utf8(&mmap[offset..offset + len]).map_err(|e| Error::IndexError {
260 message: format!("Invalid UTF-8 in file path: {}", e),
261 })?;
262 files.push(path_str.into());
263 offset += len;
264 }
265
266 let string_bytes = if offset < mmap.len() {
268 mmap[offset..].to_vec()
269 } else {
270 Vec::new()
271 };
272 let strings = StringTable::from_bytes(string_bytes);
273
274 let mut index = SemanticIndex {
276 symbols,
277 tokens,
278 references,
279 scopes,
280 edges,
281 symbol_by_name: Default::default(),
282 token_by_name: Default::default(),
283 incoming_edges: Default::default(),
284 outgoing_edges: Default::default(),
285 refs_to_symbol: Default::default(),
286 files,
287 strings,
288 entry_points: Default::default(),
289 };
290
291 index.rebuild_lookups();
293
294 Ok(index)
295}
296
297fn read_vec<T: Clone>(mmap: &Mmap, offset: &mut usize, count: usize) -> Result<Vec<T>> {
299 let size = count * std::mem::size_of::<T>();
300 if *offset + size > mmap.len() {
301 return Err(Error::IndexError {
302 message: format!(
303 "Truncated trace index file at offset {} (need {} bytes, have {})",
304 offset,
305 size,
306 mmap.len() - *offset
307 ),
308 });
309 }
310
311 let slice = &mmap[*offset..*offset + size];
312 *offset += size;
313
314 let result = unsafe {
316 let ptr = slice.as_ptr() as *const T;
317 std::slice::from_raw_parts(ptr, count).to_vec()
318 };
319
320 Ok(result)
321}
322
323pub fn load_index_streaming(path: impl AsRef<Path>) -> Result<SemanticIndex> {
331 let path = path.as_ref();
332 let file = File::open(path)?;
333 let mut reader = BufReader::with_capacity(64 * 1024, file);
334
335 let mut header_bytes = [0u8; HEADER_SIZE];
337 reader.read_exact(&mut header_bytes)?;
338 let header = Header::from_bytes(&header_bytes)?;
339
340 let symbols: Vec<Symbol> = read_vec_streaming(&mut reader, header.symbol_count as usize)?;
342
343 let tokens: Vec<Token> = read_vec_streaming(&mut reader, header.token_count as usize)?;
345
346 let references: Vec<Reference> =
348 read_vec_streaming(&mut reader, header.reference_count as usize)?;
349
350 let scopes: Vec<Scope> = read_vec_streaming(&mut reader, header.scope_count as usize)?;
352
353 let edges: Vec<Edge> = read_vec_streaming(&mut reader, header.edge_count as usize)?;
355
356 let mut files = Vec::with_capacity(header.file_count as usize);
358 for _ in 0..header.file_count {
359 let mut len_bytes = [0u8; 4];
360 reader.read_exact(&mut len_bytes)?;
361 let len = u32::from_le_bytes(len_bytes) as usize;
362
363 let mut path_bytes = vec![0u8; len];
364 reader.read_exact(&mut path_bytes)?;
365 let path_str = std::str::from_utf8(&path_bytes).map_err(|e| Error::IndexError {
366 message: format!("Invalid UTF-8 in file path: {}", e),
367 })?;
368 files.push(path_str.into());
369 }
370
371 let mut string_bytes = Vec::new();
373 reader.read_to_end(&mut string_bytes)?;
374 let strings = StringTable::from_bytes(string_bytes);
375
376 let mut index = SemanticIndex {
378 symbols,
379 tokens,
380 references,
381 scopes,
382 edges,
383 symbol_by_name: Default::default(),
384 token_by_name: Default::default(),
385 incoming_edges: Default::default(),
386 outgoing_edges: Default::default(),
387 refs_to_symbol: Default::default(),
388 files,
389 strings,
390 entry_points: Default::default(),
391 };
392
393 index.rebuild_lookups();
395
396 Ok(index)
397}
398
399fn read_vec_streaming<T: Clone, R: Read>(reader: &mut R, count: usize) -> io::Result<Vec<T>> {
401 let size = count * std::mem::size_of::<T>();
402 let mut bytes = vec![0u8; size];
403 reader.read_exact(&mut bytes)?;
404
405 let result = unsafe {
407 let ptr = bytes.as_ptr() as *const T;
408 std::slice::from_raw_parts(ptr, count).to_vec()
409 };
410
411 Ok(result)
412}
413
414pub fn trace_index_path(project_root: impl AsRef<Path>) -> std::path::PathBuf {
420 project_root.as_ref().join(".greppy").join("trace.idx")
421}
422
423pub fn trace_index_exists(project_root: impl AsRef<Path>) -> bool {
425 trace_index_path(project_root).exists()
426}
427
428#[cfg(test)]
433mod tests {
434 use super::*;
435 use crate::trace::types::{RefKind, ScopeKind, SymbolFlags, SymbolKind, TokenKind};
436 use tempfile::tempdir;
437
438 fn create_test_index() -> SemanticIndex {
439 let mut index = SemanticIndex::new();
440
441 let file_id = index.add_file("src/main.rs".into());
443
444 let name1 = index.strings.intern("main");
446 let name2 = index.strings.intern("helper");
447
448 index.add_symbol(
449 Symbol::new(
450 0,
451 name1,
452 file_id,
453 SymbolKind::Function,
454 SymbolFlags::IS_ENTRY_POINT,
455 1,
456 10,
457 ),
458 "main",
459 );
460 index.add_symbol(
461 Symbol::new(
462 1,
463 name2,
464 file_id,
465 SymbolKind::Function,
466 SymbolFlags::empty(),
467 12,
468 20,
469 ),
470 "helper",
471 );
472
473 index.add_token(
475 Token::new(0, name1, file_id, 1, 4, TokenKind::Identifier, 0),
476 "main",
477 );
478 index.add_token(
479 Token::new(1, name2, file_id, 5, 4, TokenKind::Call, 0),
480 "helper",
481 );
482
483 index.add_reference(Reference::new(1, 1, RefKind::Call));
485
486 index.add_scope(Scope::file_scope(0, file_id, 25));
488 index.add_scope(Scope::new(1, ScopeKind::Function, file_id, 0, 1, 10, name1));
489
490 index.add_edge(Edge::new(0, 1, 5));
492
493 index
494 }
495
496 #[test]
497 fn test_save_load_roundtrip() {
498 let dir = tempdir().unwrap();
499 let path = dir.path().join("trace.idx");
500
501 let original = create_test_index();
502 save_index(&original, &path).unwrap();
503
504 let loaded = load_index(&path).unwrap();
505
506 assert_eq!(loaded.symbols.len(), original.symbols.len());
508 assert_eq!(loaded.tokens.len(), original.tokens.len());
509 assert_eq!(loaded.references.len(), original.references.len());
510 assert_eq!(loaded.scopes.len(), original.scopes.len());
511 assert_eq!(loaded.edges.len(), original.edges.len());
512 assert_eq!(loaded.files.len(), original.files.len());
513
514 assert!(loaded.symbols_by_name("main").is_some());
516 assert!(loaded.symbols_by_name("helper").is_some());
517 assert_eq!(loaded.entry_points.len(), 1);
518
519 assert_eq!(loaded.callers(1), &[0]);
521 assert_eq!(loaded.callees(0), &[1]);
522 }
523
524 #[test]
525 fn test_streaming_load() {
526 let dir = tempdir().unwrap();
527 let path = dir.path().join("trace.idx");
528
529 let original = create_test_index();
530 save_index(&original, &path).unwrap();
531
532 let loaded = load_index_streaming(&path).unwrap();
533
534 assert_eq!(loaded.symbols.len(), original.symbols.len());
535 assert_eq!(loaded.tokens.len(), original.tokens.len());
536 }
537
538 #[test]
539 fn test_header_validation() {
540 let bad_magic = [0u8; HEADER_SIZE];
542 assert!(Header::from_bytes(&bad_magic).is_err());
543
544 let mut valid = [0u8; HEADER_SIZE];
546 valid[..8].copy_from_slice(&MAGIC);
547 valid[8..12].copy_from_slice(&VERSION.to_le_bytes());
548 assert!(Header::from_bytes(&valid).is_ok());
549 }
550}