1use anyhow::{Context, Result};
17use std::collections::HashMap;
18use std::fs::{File, OpenOptions};
19use std::io::Write;
20use std::path::{Path, PathBuf};
21
22pub type Trigram = u32;
24
25const MAGIC: &[u8; 4] = b"RFTG"; const VERSION: u32 = 3; #[allow(dead_code)]
30const HEADER_SIZE: usize = 24;
31
32fn write_varint(writer: &mut impl Write, mut value: u32) -> std::io::Result<()> {
35 loop {
36 let mut byte = (value & 0x7F) as u8;
37 value >>= 7;
38 if value != 0 {
39 byte |= 0x80; }
41 writer.write_all(&[byte])?;
42 if value == 0 {
43 break;
44 }
45 }
46 Ok(())
47}
48
49fn read_varint(data: &[u8]) -> Result<(u32, usize)> {
51 let mut value: u32 = 0;
52 let mut shift = 0;
53 let mut pos = 0;
54
55 loop {
56 if pos >= data.len() {
57 anyhow::bail!("Truncated varint");
58 }
59 let byte = data[pos];
60 pos += 1;
61
62 value |= ((byte & 0x7F) as u32) << shift;
63 if byte & 0x80 == 0 {
64 break;
65 }
66 shift += 7;
67 if shift >= 32 {
68 anyhow::bail!("Varint too large");
69 }
70 }
71
72 Ok((value, pos))
73}
74
75fn decompress_posting_list(mmap: &[u8], offset: u64, size: u32) -> Result<Vec<FileLocation>> {
85 let start = offset as usize;
86 let end = start + size as usize;
87
88 if end > mmap.len() {
89 anyhow::bail!(
90 "Posting list out of bounds: offset={}, size={}, mmap_len={}",
91 offset,
92 size,
93 mmap.len()
94 );
95 }
96
97 let compressed_data = &mmap[start..end];
98
99 let mut locations = Vec::new();
101 let mut pos = 0;
102 let mut prev_file_id = 0u32;
103 let mut prev_line_no = 0u32;
104 let mut prev_byte_offset = 0u32;
105
106 while pos < compressed_data.len() {
107 let (file_id_delta, consumed) = read_varint(&compressed_data[pos..])?;
109 pos += consumed;
110
111 let (line_no_delta, consumed) = read_varint(&compressed_data[pos..])?;
113 pos += consumed;
114
115 let (byte_offset_delta, consumed) = read_varint(&compressed_data[pos..])?;
117 pos += consumed;
118
119 let file_id = prev_file_id.wrapping_add(file_id_delta);
121 let line_no = prev_line_no.wrapping_add(line_no_delta);
122 let byte_offset = prev_byte_offset.wrapping_add(byte_offset_delta);
123
124 locations.push(FileLocation {
125 file_id,
126 line_no,
127 byte_offset,
128 });
129
130 prev_file_id = file_id;
132 prev_line_no = line_no;
133 prev_byte_offset = byte_offset;
134 }
135
136 Ok(locations)
137}
138
139#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
141pub struct FileLocation {
142 pub file_id: u32,
144 pub line_no: u32,
146 pub byte_offset: u32,
148}
149
150impl FileLocation {
151 pub fn new(file_id: u32, line_no: u32, byte_offset: u32) -> Self {
152 Self {
153 file_id,
154 line_no,
155 byte_offset,
156 }
157 }
158}
159
160#[derive(Debug, Clone)]
165struct DirectoryEntry {
166 trigram: Trigram,
168 data_offset: u64,
170 compressed_size: u32,
172}
173
174pub struct TrigramIndex {
185 index: Vec<(Trigram, Vec<FileLocation>)>,
188 files: Vec<PathBuf>,
190 temp_index: Option<HashMap<Trigram, Vec<FileLocation>>>,
192 mmap: Option<memmap2::Mmap>,
194 directory: Vec<DirectoryEntry>,
196 partial_indices: Vec<PathBuf>,
198 temp_dir: Option<PathBuf>,
200 max_posting_list_entries: usize,
203}
204
205impl TrigramIndex {
206 pub fn new() -> Self {
208 Self {
209 index: Vec::new(),
210 files: Vec::new(),
211 temp_index: Some(HashMap::new()),
212 mmap: None,
213 directory: Vec::new(),
214 partial_indices: Vec::new(),
215 temp_dir: None,
216 max_posting_list_entries: 0,
217 }
218 }
219
220 pub fn set_max_posting_list_entries(&mut self, cap: usize) {
222 self.max_posting_list_entries = cap;
223 }
224
225 pub fn enable_batch_flush(&mut self, temp_dir: PathBuf) -> Result<()> {
230 std::fs::create_dir_all(&temp_dir)
231 .context("Failed to create temp directory for batch flushing")?;
232 self.temp_dir = Some(temp_dir);
233 log::info!("Enabled batch-flush mode for trigram index");
234 Ok(())
235 }
236
237 pub fn flush_batch(&mut self) -> Result<()> {
242 let temp_dir = self.temp_dir.as_ref().ok_or_else(|| {
243 anyhow::anyhow!("Batch flush not enabled - call enable_batch_flush() first")
244 })?;
245
246 let temp_map = self
248 .temp_index
249 .take()
250 .ok_or_else(|| anyhow::anyhow!("No temp index to flush"))?;
251
252 if temp_map.is_empty() {
253 self.temp_index = Some(HashMap::new());
255 return Ok(());
256 }
257
258 let mut partial_index: Vec<(Trigram, Vec<FileLocation>)> = temp_map.into_iter().collect();
260
261 for (_, list) in partial_index.iter_mut() {
263 list.sort_unstable();
264 list.dedup();
265 }
266
267 partial_index.sort_unstable_by_key(|(trigram, _)| *trigram);
269
270 let partial_file = temp_dir.join(format!("partial_{}.bin", self.partial_indices.len()));
272 self.write_partial_index(&partial_file, &partial_index)?;
273
274 self.partial_indices.push(partial_file);
275
276 self.temp_index = Some(HashMap::new());
278
279 log::debug!(
280 "Flushed batch {} with {} trigrams to disk",
281 self.partial_indices.len(),
282 partial_index.len()
283 );
284
285 Ok(())
286 }
287
288 fn write_partial_index(
290 &self,
291 path: &Path,
292 index: &[(Trigram, Vec<FileLocation>)],
293 ) -> Result<()> {
294 use std::io::BufWriter;
295
296 let file = OpenOptions::new()
297 .create(true)
298 .write(true)
299 .truncate(true)
300 .open(path)?;
301
302 let mut writer = BufWriter::with_capacity(16 * 1024 * 1024, file);
303
304 writer.write_all(&(index.len() as u64).to_le_bytes())?;
306
307 for (trigram, locations) in index {
309 writer.write_all(&trigram.to_le_bytes())?;
310 writer.write_all(&(locations.len() as u32).to_le_bytes())?;
311
312 for loc in locations {
313 writer.write_all(&loc.file_id.to_le_bytes())?;
314 writer.write_all(&loc.line_no.to_le_bytes())?;
315 writer.write_all(&loc.byte_offset.to_le_bytes())?;
316 }
317 }
318
319 writer.flush()?;
320 Ok(())
321 }
322
323 pub fn add_file(&mut self, path: PathBuf) -> u32 {
325 let file_id = self.files.len() as u32;
326 self.files.push(path);
327 file_id
328 }
329
330 pub fn get_file(&self, file_id: u32) -> Option<&PathBuf> {
332 self.files.get(file_id as usize)
333 }
334
335 pub fn file_count(&self) -> usize {
337 self.files.len()
338 }
339
340 pub fn trigram_count(&self) -> usize {
342 if !self.directory.is_empty() {
343 self.directory.len()
345 } else {
346 self.index.len()
348 }
349 }
350
351 pub fn index_file(&mut self, file_id: u32, content: &str) {
356 let trigrams = extract_trigrams_with_locations(content, file_id);
357
358 if let Some(ref mut temp_map) = self.temp_index {
360 for (trigram, location) in trigrams {
361 temp_map
362 .entry(trigram)
363 .or_insert_with(Vec::new)
364 .push(location);
365 }
366 } else {
367 panic!("Cannot call index_file() after finalize(). Index is read-only.");
368 }
369 }
370
371 pub fn build_from_trigrams(&mut self, trigrams: Vec<(Trigram, FileLocation)>) {
376 let mut temp_map: HashMap<Trigram, Vec<FileLocation>> = HashMap::new();
377
378 for (trigram, location) in trigrams {
380 temp_map.entry(trigram).or_default().push(location);
381 }
382
383 self.index = temp_map.into_iter().collect();
385
386 self.temp_index = None;
388
389 self.finalize();
391 }
392
393 pub fn finalize(&mut self) {
401 if !self.partial_indices.is_empty() {
404 log::info!(
405 "Deferring finalization - will stream merge {} partial indices during write()",
406 self.partial_indices.len()
407 );
408
409 if let Some(ref temp_map) = self.temp_index
411 && !temp_map.is_empty()
412 {
413 self.flush_batch().expect("Failed to flush final batch");
414 }
415
416 return;
418 }
419
420 if let Some(temp_map) = self.temp_index.take() {
423 self.index = temp_map.into_iter().collect();
424 }
425
426 let cap = self.max_posting_list_entries;
428 for (trigram, list) in self.index.iter_mut() {
429 list.sort_unstable();
430 list.dedup(); if cap > 0 && list.len() > cap {
432 log::warn!(
433 "Trigram 0x{:06X} posting list has {} entries (cap {}); truncating.",
434 trigram,
435 list.len(),
436 cap
437 );
438 list.truncate(cap);
439 }
440 }
441
442 self.index.sort_unstable_by_key(|(trigram, _)| *trigram);
444 }
445
446 fn merge_partial_indices_to_file(&mut self, output_path: &Path) -> Result<()> {
454 use std::cmp::Ordering;
455 use std::collections::BinaryHeap;
456 use std::io::{BufReader, BufWriter, Read};
457
458 log::info!(
459 "Streaming merge of {} partial indices to {:?}",
460 self.partial_indices.len(),
461 output_path
462 );
463
464 struct PartialIndexReader {
466 reader: BufReader<File>,
467 current_trigram: Option<Trigram>,
468 current_posting_list: Vec<FileLocation>,
469 reader_id: usize,
470 }
471
472 let mut readers: Vec<PartialIndexReader> = Vec::new();
473
474 for (idx, partial_path) in self.partial_indices.iter().enumerate() {
475 let file = File::open(partial_path)
476 .with_context(|| format!("Failed to open partial index: {:?}", partial_path))?;
477 let mut reader = BufReader::with_capacity(16 * 1024 * 1024, file);
478
479 let mut buf = [0u8; 8];
481 reader.read_exact(&mut buf)?;
482
483 readers.push(PartialIndexReader {
484 reader,
485 current_trigram: None,
486 current_posting_list: Vec::new(),
487 reader_id: idx,
488 });
489 }
490
491 fn read_next_trigram(reader: &mut PartialIndexReader) -> Result<bool> {
493 let mut trigram_buf = [0u8; 4];
495 match reader.reader.read_exact(&mut trigram_buf) {
496 Ok(_) => {
497 let trigram = u32::from_le_bytes(trigram_buf);
498
499 let mut len_buf = [0u8; 4];
501 reader.reader.read_exact(&mut len_buf)?;
502 let list_len = u32::from_le_bytes(len_buf) as usize;
503
504 let mut locations = Vec::with_capacity(list_len);
506 for _ in 0..list_len {
507 let mut loc_buf = [0u8; 12];
508 reader.reader.read_exact(&mut loc_buf)?;
509
510 let file_id =
511 u32::from_le_bytes([loc_buf[0], loc_buf[1], loc_buf[2], loc_buf[3]]);
512 let line_no =
513 u32::from_le_bytes([loc_buf[4], loc_buf[5], loc_buf[6], loc_buf[7]]);
514 let byte_offset =
515 u32::from_le_bytes([loc_buf[8], loc_buf[9], loc_buf[10], loc_buf[11]]);
516
517 locations.push(FileLocation {
518 file_id,
519 line_no,
520 byte_offset,
521 });
522 }
523
524 reader.current_trigram = Some(trigram);
525 reader.current_posting_list = locations;
526 Ok(true)
527 }
528 Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => {
529 reader.current_trigram = None;
530 Ok(false)
531 }
532 Err(e) => Err(e.into()),
533 }
534 }
535
536 for reader in &mut readers {
538 read_next_trigram(reader)?;
539 }
540
541 #[derive(Eq, PartialEq)]
543 struct HeapEntry {
544 trigram: Trigram,
545 reader_id: usize,
546 }
547
548 impl Ord for HeapEntry {
549 fn cmp(&self, other: &Self) -> Ordering {
550 other
552 .trigram
553 .cmp(&self.trigram)
554 .then_with(|| other.reader_id.cmp(&self.reader_id))
555 }
556 }
557
558 impl PartialOrd for HeapEntry {
559 fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
560 Some(self.cmp(other))
561 }
562 }
563
564 let mut heap: BinaryHeap<HeapEntry> = BinaryHeap::new();
566 for reader in &readers {
567 if let Some(trigram) = reader.current_trigram {
568 heap.push(HeapEntry {
569 trigram,
570 reader_id: reader.reader_id,
571 });
572 }
573 }
574
575 let file = OpenOptions::new()
577 .create(true)
578 .write(true)
579 .truncate(true)
580 .open(output_path)
581 .with_context(|| format!("Failed to create {}", output_path.display()))?;
582
583 let mut writer = BufWriter::with_capacity(16 * 1024 * 1024, file);
584
585 writer.write_all(MAGIC)?;
587 writer.write_all(&VERSION.to_le_bytes())?;
588 writer.write_all(&0u64.to_le_bytes())?; writer.write_all(&(self.files.len() as u64).to_le_bytes())?; let mut directory: Vec<DirectoryEntry> = Vec::new();
593 let mut num_trigrams = 0u64;
594
595 let mut current_trigram: Option<Trigram> = None;
597 let mut merged_locations: Vec<FileLocation> = Vec::new();
598
599 while let Some(entry) = heap.pop() {
600 let reader = &mut readers[entry.reader_id];
601
602 if let Some(trigram) = current_trigram.filter(|&t| t != entry.trigram) {
604 merged_locations.sort_unstable();
605 merged_locations.dedup();
606
607 let cap = self.max_posting_list_entries;
608 if cap > 0 && merged_locations.len() > cap {
609 log::warn!(
610 "Trigram 0x{:06X} posting list has {} entries (cap {}); truncating.",
611 trigram,
612 merged_locations.len(),
613 cap
614 );
615 merged_locations.truncate(cap);
616 }
617
618 let data_offset = writer.stream_position()?;
620 let compressed_size =
621 self.write_compressed_posting_list(&mut writer, &merged_locations)?;
622
623 directory.push(DirectoryEntry {
624 trigram,
625 data_offset,
626 compressed_size,
627 });
628
629 num_trigrams += 1;
630 merged_locations.clear();
631 }
632
633 current_trigram = Some(entry.trigram);
635
636 merged_locations.extend_from_slice(&reader.current_posting_list);
638
639 if read_next_trigram(reader)?
641 && let Some(next_trigram) = reader.current_trigram
642 {
643 heap.push(HeapEntry {
644 trigram: next_trigram,
645 reader_id: entry.reader_id,
646 });
647 }
648 }
649
650 if let Some(trigram) = current_trigram {
652 merged_locations.sort_unstable();
653 merged_locations.dedup();
654
655 let cap = self.max_posting_list_entries;
656 if cap > 0 && merged_locations.len() > cap {
657 log::warn!(
658 "Trigram 0x{:06X} posting list has {} entries (cap {}); truncating.",
659 trigram,
660 merged_locations.len(),
661 cap
662 );
663 merged_locations.truncate(cap);
664 }
665
666 let data_offset = writer.stream_position()?;
667 let compressed_size =
668 self.write_compressed_posting_list(&mut writer, &merged_locations)?;
669
670 directory.push(DirectoryEntry {
671 trigram,
672 data_offset,
673 compressed_size,
674 });
675
676 num_trigrams += 1;
677 }
678
679 log::info!(
680 "Merged {} trigrams from {} partial indices",
681 num_trigrams,
682 self.partial_indices.len()
683 );
684
685 let _data_end_pos = writer.stream_position()?;
687
688 for file_path in &self.files {
690 let path_str = file_path.to_string_lossy();
691 let path_bytes = path_str.as_bytes();
692 write_varint(&mut writer, path_bytes.len() as u32)?;
693 writer.write_all(path_bytes)?;
694 }
695
696 writer.flush()?;
698 drop(writer);
699
700 use std::io::{Seek, SeekFrom};
703
704 let mut temp_data = Vec::new();
706 {
707 let mut file = File::open(output_path)?;
708 file.seek(SeekFrom::Start(HEADER_SIZE as u64))?;
709 file.read_to_end(&mut temp_data)?;
710 }
711
712 let file = OpenOptions::new()
714 .write(true)
715 .truncate(true)
716 .open(output_path)?;
717 let mut writer = BufWriter::with_capacity(16 * 1024 * 1024, file);
718
719 writer.write_all(MAGIC)?;
721 writer.write_all(&VERSION.to_le_bytes())?;
722 writer.write_all(&num_trigrams.to_le_bytes())?;
723 writer.write_all(&(self.files.len() as u64).to_le_bytes())?;
724
725 for entry in &directory {
727 writer.write_all(&entry.trigram.to_le_bytes())?;
728 let adjusted_offset = entry.data_offset + (directory.len() * 16) as u64;
730 writer.write_all(&adjusted_offset.to_le_bytes())?;
731 writer.write_all(&entry.compressed_size.to_le_bytes())?;
732 }
733
734 writer.write_all(&temp_data)?;
736
737 writer.flush()?;
739 writer.get_ref().sync_all()?;
740
741 for partial_path in &self.partial_indices {
743 let _ = std::fs::remove_file(partial_path);
744 }
745 if let Some(ref temp_dir) = self.temp_dir {
746 let _ = std::fs::remove_dir(temp_dir);
747 }
748
749 log::info!("Wrote {} trigrams to {:?}", num_trigrams, output_path);
750
751 Ok(())
752 }
753
754 fn write_compressed_posting_list(
756 &self,
757 writer: &mut impl Write,
758 locations: &[FileLocation],
759 ) -> Result<u32> {
760 let mut compressed = Vec::new();
761
762 let mut prev_file_id = 0u32;
764 let mut prev_line_no = 0u32;
765 let mut prev_byte_offset = 0u32;
766
767 for loc in locations {
768 let file_id_delta = loc.file_id.wrapping_sub(prev_file_id);
770 let line_no_delta = loc.line_no.wrapping_sub(prev_line_no);
771 let byte_offset_delta = loc.byte_offset.wrapping_sub(prev_byte_offset);
772
773 write_varint(&mut compressed, file_id_delta)?;
775 write_varint(&mut compressed, line_no_delta)?;
776 write_varint(&mut compressed, byte_offset_delta)?;
777
778 prev_file_id = loc.file_id;
780 prev_line_no = loc.line_no;
781 prev_byte_offset = loc.byte_offset;
782 }
783
784 let compressed_size = compressed.len() as u32;
785 writer.write_all(&compressed)?;
786
787 Ok(compressed_size)
788 }
789
790 #[allow(dead_code)]
792 fn merge_partial_indices(&mut self) -> Result<()> {
793 use std::io::{BufReader, Read};
794
795 let mut all_entries: Vec<(Trigram, FileLocation)> = Vec::new();
797
798 for partial_path in &self.partial_indices {
799 let file = File::open(partial_path)
800 .with_context(|| format!("Failed to open partial index: {:?}", partial_path))?;
801 let mut reader = BufReader::with_capacity(16 * 1024 * 1024, file);
802
803 let mut buf = [0u8; 8];
805 reader.read_exact(&mut buf)?;
806 let num_trigrams = u64::from_le_bytes(buf) as usize;
807
808 for _ in 0..num_trigrams {
810 let mut trigram_buf = [0u8; 4];
812 reader.read_exact(&mut trigram_buf)?;
813 let trigram = u32::from_le_bytes(trigram_buf);
814
815 let mut len_buf = [0u8; 4];
817 reader.read_exact(&mut len_buf)?;
818 let list_len = u32::from_le_bytes(len_buf) as usize;
819
820 for _ in 0..list_len {
822 let mut loc_buf = [0u8; 12]; reader.read_exact(&mut loc_buf)?;
824
825 let file_id =
826 u32::from_le_bytes([loc_buf[0], loc_buf[1], loc_buf[2], loc_buf[3]]);
827 let line_no =
828 u32::from_le_bytes([loc_buf[4], loc_buf[5], loc_buf[6], loc_buf[7]]);
829 let byte_offset =
830 u32::from_le_bytes([loc_buf[8], loc_buf[9], loc_buf[10], loc_buf[11]]);
831
832 all_entries.push((
833 trigram,
834 FileLocation {
835 file_id,
836 line_no,
837 byte_offset,
838 },
839 ));
840 }
841 }
842 }
843
844 log::info!(
845 "Read {} total trigram entries from {} partial indices",
846 all_entries.len(),
847 self.partial_indices.len()
848 );
849
850 let mut index_map: HashMap<Trigram, Vec<FileLocation>> = HashMap::new();
852 for (trigram, location) in all_entries {
853 index_map.entry(trigram).or_default().push(location);
854 }
855
856 self.index = index_map.into_iter().collect();
858
859 for (_, list) in self.index.iter_mut() {
861 list.sort_unstable();
862 list.dedup();
863 }
864
865 self.index.sort_unstable_by_key(|(trigram, _)| *trigram);
867
868 for partial_path in &self.partial_indices {
870 let _ = std::fs::remove_file(partial_path);
871 }
872 if let Some(ref temp_dir) = self.temp_dir {
873 let _ = std::fs::remove_dir(temp_dir);
874 }
875
876 log::info!("Merged into final index with {} trigrams", self.index.len());
877
878 Ok(())
879 }
880
881 pub fn search(&self, pattern: &str) -> Vec<FileLocation> {
889 if pattern.len() < 3 {
890 return vec![];
892 }
893
894 let trigrams = extract_trigrams(pattern);
895 if trigrams.is_empty() {
896 return vec![];
897 }
898
899 if let Some(ref mmap) = self.mmap {
901 let mut posting_lists: Vec<Vec<FileLocation>> = Vec::new();
903
904 for trigram in &trigrams {
905 match self.directory.binary_search_by_key(trigram, |e| e.trigram) {
907 Ok(idx) => {
908 let entry = &self.directory[idx];
909 match decompress_posting_list(
911 mmap,
912 entry.data_offset,
913 entry.compressed_size,
914 ) {
915 Ok(locations) => posting_lists.push(locations),
916 Err(e) => {
917 log::warn!(
918 "Failed to decompress posting list for trigram {}: {}",
919 trigram,
920 e
921 );
922 return vec![];
923 }
924 }
925 }
926 Err(_) => {
927 return vec![];
929 }
930 }
931 }
932
933 if posting_lists.is_empty() || posting_lists.len() < trigrams.len() {
934 return vec![];
935 }
936
937 posting_lists.sort_by_key(|list| list.len());
939
940 intersect_by_file_owned(&posting_lists)
942 } else {
943 let mut posting_lists: Vec<&Vec<FileLocation>> = trigrams
945 .iter()
946 .filter_map(|t| {
947 self.index
948 .binary_search_by_key(t, |(trigram, _)| *trigram)
949 .ok()
950 .map(|idx| &self.index[idx].1)
951 })
952 .collect();
953
954 if posting_lists.is_empty() {
955 return vec![];
956 }
957
958 if posting_lists.len() < trigrams.len() {
959 return vec![];
961 }
962
963 posting_lists.sort_by_key(|list| list.len());
965
966 intersect_by_file(&posting_lists)
968 }
969 }
970
971 pub fn get_posting_list(&self, trigram: Trigram) -> Option<&Vec<FileLocation>> {
973 self.index
974 .binary_search_by_key(&trigram, |(t, _)| *t)
975 .ok()
976 .map(|idx| &self.index[idx].1)
977 }
978
979 pub fn write(&mut self, path: impl AsRef<Path>) -> Result<()> {
993 let path = path.as_ref();
994
995 if !self.partial_indices.is_empty() {
997 log::info!(
998 "Using streaming merge to write {} partial indices",
999 self.partial_indices.len()
1000 );
1001 return self.merge_partial_indices_to_file(path);
1002 }
1003
1004 let file = OpenOptions::new()
1009 .create(true)
1010 .write(true)
1011 .truncate(true)
1012 .open(path)
1013 .with_context(|| format!("Failed to create {}", path.display()))?;
1014
1015 let mut writer = std::io::BufWriter::with_capacity(16 * 1024 * 1024, file);
1017
1018 writer.write_all(MAGIC)?;
1020 writer.write_all(&VERSION.to_le_bytes())?;
1021 writer.write_all(&(self.index.len() as u64).to_le_bytes())?; writer.write_all(&(self.files.len() as u64).to_le_bytes())?; let mut directory: Vec<DirectoryEntry> = Vec::with_capacity(self.index.len());
1026
1027 let directory_start = HEADER_SIZE as u64;
1029 let directory_size = self.index.len() * 16;
1030
1031 let data_start = directory_start + directory_size as u64;
1033 let mut current_offset = data_start;
1034
1035 let mut compressed_lists: Vec<(Trigram, Vec<u8>)> = Vec::with_capacity(self.index.len());
1041
1042 for (trigram, locations) in &self.index {
1043 let mut compressed = Vec::new();
1045 let mut prev_file_id = 0u32;
1046 let mut prev_line_no = 0u32;
1047 let mut prev_byte_offset = 0u32;
1048
1049 for loc in locations {
1050 let file_id_delta = loc.file_id.wrapping_sub(prev_file_id);
1051 let line_no_delta = loc.line_no.wrapping_sub(prev_line_no);
1052 let byte_offset_delta = loc.byte_offset.wrapping_sub(prev_byte_offset);
1053
1054 write_varint(&mut compressed, file_id_delta)?;
1055 write_varint(&mut compressed, line_no_delta)?;
1056 write_varint(&mut compressed, byte_offset_delta)?;
1057
1058 prev_file_id = loc.file_id;
1059 prev_line_no = loc.line_no;
1060 prev_byte_offset = loc.byte_offset;
1061 }
1062
1063 directory.push(DirectoryEntry {
1064 trigram: *trigram,
1065 data_offset: current_offset,
1066 compressed_size: compressed.len() as u32,
1067 });
1068 current_offset += compressed.len() as u64;
1069
1070 compressed_lists.push((*trigram, compressed));
1071 }
1072
1073 for entry in &directory {
1075 writer.write_all(&entry.trigram.to_le_bytes())?;
1076 writer.write_all(&entry.data_offset.to_le_bytes())?;
1077 writer.write_all(&entry.compressed_size.to_le_bytes())?;
1078 }
1079
1080 for (_, compressed) in &compressed_lists {
1082 writer.write_all(compressed)?;
1083 }
1084
1085 for file_path in &self.files {
1087 let path_str = file_path.to_string_lossy();
1088 let path_bytes = path_str.as_bytes();
1089 write_varint(&mut writer, path_bytes.len() as u32)?;
1090 writer.write_all(path_bytes)?;
1091 }
1092
1093 writer.flush()?;
1095 writer.get_ref().sync_all()?;
1096
1097 log::info!(
1098 "Wrote lazy-loadable trigram index: {} trigrams, {} files to {:?}",
1099 self.index.len(),
1100 self.files.len(),
1101 path
1102 );
1103
1104 Ok(())
1105 }
1106
1107 pub fn load(path: impl AsRef<Path>) -> Result<Self> {
1112 let path = path.as_ref();
1113
1114 let file =
1115 File::open(path).with_context(|| format!("Failed to open {}", path.display()))?;
1116
1117 let mmap = unsafe {
1119 memmap2::Mmap::map(&file)
1120 .with_context(|| format!("Failed to mmap {}", path.display()))?
1121 };
1122
1123 if mmap.len() < HEADER_SIZE {
1125 anyhow::bail!(
1126 "trigrams.bin too small (expected at least {} bytes)",
1127 HEADER_SIZE
1128 );
1129 }
1130
1131 if &mmap[0..4] != MAGIC {
1132 anyhow::bail!("Invalid trigrams.bin (wrong magic bytes)");
1133 }
1134
1135 let version = u32::from_le_bytes([mmap[4], mmap[5], mmap[6], mmap[7]]);
1136 if version != VERSION {
1137 anyhow::bail!(
1138 "Unsupported trigrams.bin version: {} (expected {}). Please re-index with 'reflex index'.",
1139 version,
1140 VERSION
1141 );
1142 }
1143
1144 let num_trigrams = u64::from_le_bytes([
1145 mmap[8], mmap[9], mmap[10], mmap[11], mmap[12], mmap[13], mmap[14], mmap[15],
1146 ]) as usize;
1147
1148 let num_files = u64::from_le_bytes([
1149 mmap[16], mmap[17], mmap[18], mmap[19], mmap[20], mmap[21], mmap[22], mmap[23],
1150 ]) as usize;
1151
1152 log::debug!(
1153 "Loading lazy trigram index: {} trigrams, {} files",
1154 num_trigrams,
1155 num_files
1156 );
1157
1158 let mut directory = Vec::with_capacity(num_trigrams);
1160 let mut pos = HEADER_SIZE;
1161 let directory_size = num_trigrams * 16; for _ in 0..num_trigrams {
1164 if pos + 16 > mmap.len() {
1165 anyhow::bail!("Truncated directory entry at pos={}", pos);
1166 }
1167
1168 let trigram =
1169 u32::from_le_bytes([mmap[pos], mmap[pos + 1], mmap[pos + 2], mmap[pos + 3]]);
1170 pos += 4;
1171
1172 let data_offset = u64::from_le_bytes([
1173 mmap[pos],
1174 mmap[pos + 1],
1175 mmap[pos + 2],
1176 mmap[pos + 3],
1177 mmap[pos + 4],
1178 mmap[pos + 5],
1179 mmap[pos + 6],
1180 mmap[pos + 7],
1181 ]);
1182 pos += 8;
1183
1184 let compressed_size =
1185 u32::from_le_bytes([mmap[pos], mmap[pos + 1], mmap[pos + 2], mmap[pos + 3]]);
1186 pos += 4;
1187
1188 directory.push(DirectoryEntry {
1189 trigram,
1190 data_offset,
1191 compressed_size,
1192 });
1193 }
1194
1195 directory.sort_unstable_by_key(|e| e.trigram);
1197
1198 let data_section_size: u64 = directory.iter().map(|e| e.compressed_size as u64).sum();
1200 let files_section_offset = HEADER_SIZE + directory_size + data_section_size as usize;
1201 pos = files_section_offset;
1202
1203 let mut files = Vec::with_capacity(num_files);
1205 for _ in 0..num_files {
1206 let (path_len, consumed) = read_varint(&mmap[pos..])?;
1208 pos += consumed;
1209 let path_len = path_len as usize;
1210
1211 if pos + path_len > mmap.len() {
1212 anyhow::bail!("Truncated file path at pos={}", pos);
1213 }
1214
1215 let path_bytes = &mmap[pos..pos + path_len];
1216 let path_str = std::str::from_utf8(path_bytes).context("Invalid UTF-8 in file path")?;
1217 files.push(PathBuf::from(path_str));
1218 pos += path_len;
1219 }
1220
1221 log::info!(
1222 "Loaded lazy trigram index: {} trigrams, {} files (directory: {} KB)",
1223 num_trigrams,
1224 num_files,
1225 directory_size / 1024
1226 );
1227
1228 Ok(Self {
1229 index: Vec::new(), files,
1231 temp_index: None,
1232 mmap: Some(mmap), directory,
1234 partial_indices: Vec::new(),
1235 temp_dir: None,
1236 max_posting_list_entries: 0,
1237 })
1238 }
1239}
1240
1241impl Default for TrigramIndex {
1242 fn default() -> Self {
1243 Self::new()
1244 }
1245}
1246
1247pub fn extract_trigrams(text: &str) -> Vec<Trigram> {
1251 let bytes = text.as_bytes();
1252 let mut trigrams = Vec::new();
1253
1254 for i in 0..bytes.len().saturating_sub(2) {
1255 let trigram = bytes_to_trigram(&bytes[i..i + 3]);
1256 trigrams.push(trigram);
1257 }
1258
1259 trigrams
1260}
1261
1262pub fn extract_trigrams_with_locations(text: &str, file_id: u32) -> Vec<(Trigram, FileLocation)> {
1266 let bytes = text.as_bytes();
1267 let mut result = Vec::new();
1268
1269 let mut line_no = 1;
1270
1271 for (i, &byte) in bytes.iter().enumerate() {
1272 if byte == b'\n' {
1274 line_no += 1;
1275 }
1276
1277 if i + 2 < bytes.len() {
1279 let trigram = bytes_to_trigram(&bytes[i..i + 3]);
1280 let location = FileLocation::new(file_id, line_no, i as u32);
1281 result.push((trigram, location));
1282 }
1283 }
1284
1285 result
1286}
1287
1288#[inline]
1290fn bytes_to_trigram(bytes: &[u8]) -> Trigram {
1291 debug_assert_eq!(bytes.len(), 3);
1292 (bytes[0] as u32) << 16 | (bytes[1] as u32) << 8 | (bytes[2] as u32)
1293}
1294
1295#[allow(dead_code)]
1297fn trigram_to_bytes(trigram: Trigram) -> [u8; 3] {
1298 [
1299 ((trigram >> 16) & 0xFF) as u8,
1300 ((trigram >> 8) & 0xFF) as u8,
1301 (trigram & 0xFF) as u8,
1302 ]
1303}
1304
1305fn intersect_by_file(lists: &[&Vec<FileLocation>]) -> Vec<FileLocation> {
1310 if lists.is_empty() {
1311 return vec![];
1312 }
1313
1314 use std::collections::HashSet;
1315
1316 let mut candidates: HashSet<(u32, u32)> = lists[0]
1318 .iter()
1319 .map(|loc| (loc.file_id, loc.line_no))
1320 .collect();
1321
1322 for &list in &lists[1..] {
1324 let list_pairs: HashSet<(u32, u32)> =
1325 list.iter().map(|loc| (loc.file_id, loc.line_no)).collect();
1326 candidates.retain(|pair| list_pairs.contains(pair));
1327 }
1328
1329 let mut result = Vec::new();
1331 for &(file_id, line_no) in &candidates {
1332 if let Some(loc) = lists[0]
1334 .iter()
1335 .find(|loc| loc.file_id == file_id && loc.line_no == line_no)
1336 {
1337 result.push(*loc);
1338 }
1339 }
1340
1341 result.sort_unstable();
1342 result
1343}
1344
1345fn intersect_by_file_owned(lists: &[Vec<FileLocation>]) -> Vec<FileLocation> {
1352 if lists.is_empty() {
1353 return vec![];
1354 }
1355
1356 use std::collections::HashSet;
1357
1358 let mut candidates: HashSet<(u32, u32)> = lists[0]
1360 .iter()
1361 .map(|loc| (loc.file_id, loc.line_no))
1362 .collect();
1363
1364 for list in &lists[1..] {
1366 let list_pairs: HashSet<(u32, u32)> =
1367 list.iter().map(|loc| (loc.file_id, loc.line_no)).collect();
1368 candidates.retain(|pair| list_pairs.contains(pair));
1369 }
1370
1371 let mut result = Vec::new();
1373 for &(file_id, line_no) in &candidates {
1374 if let Some(loc) = lists[0]
1376 .iter()
1377 .find(|loc| loc.file_id == file_id && loc.line_no == line_no)
1378 {
1379 result.push(*loc);
1380 }
1381 }
1382
1383 result.sort_unstable();
1384 result
1385}
1386
1387#[cfg(test)]
1388mod tests {
1389 use super::*;
1390
1391 #[test]
1392 fn test_extract_trigrams() {
1393 let text = "hello";
1394 let trigrams = extract_trigrams(text);
1395
1396 assert_eq!(trigrams.len(), 3);
1398
1399 let expected = vec![
1401 bytes_to_trigram(b"hel"),
1402 bytes_to_trigram(b"ell"),
1403 bytes_to_trigram(b"llo"),
1404 ];
1405 assert_eq!(trigrams, expected);
1406 }
1407
1408 #[test]
1409 fn test_extract_trigrams_short() {
1410 assert_eq!(extract_trigrams("ab").len(), 0);
1411 assert_eq!(extract_trigrams("abc").len(), 1);
1412 }
1413
1414 #[test]
1415 fn test_bytes_to_trigram() {
1416 let trigram1 = bytes_to_trigram(b"abc");
1417 let trigram2 = bytes_to_trigram(b"abc");
1418 let trigram3 = bytes_to_trigram(b"xyz");
1419
1420 assert_eq!(trigram1, trigram2);
1421 assert_ne!(trigram1, trigram3);
1422 }
1423
1424 #[test]
1425 fn test_trigram_roundtrip() {
1426 let original = b"foo";
1427 let trigram = bytes_to_trigram(original);
1428 let recovered = trigram_to_bytes(trigram);
1429 assert_eq!(original, &recovered);
1430 }
1431
1432 #[test]
1433 fn test_extract_with_locations() {
1434 let text = "hello\nworld";
1435 let locs = extract_trigrams_with_locations(text, 0);
1436
1437 assert_eq!(locs.len(), 9);
1440
1441 assert_eq!(locs[0].1.line_no, 1);
1443
1444 let world_start = text.find("world").unwrap();
1446 let world_trigram_idx = locs
1447 .iter()
1448 .position(|(_, loc)| loc.byte_offset as usize == world_start)
1449 .unwrap();
1450 assert_eq!(locs[world_trigram_idx].1.line_no, 2);
1451 }
1452
1453 #[test]
1454 fn test_trigram_index_basic() {
1455 let mut index = TrigramIndex::new();
1456
1457 let file_id = index.add_file(PathBuf::from("test.txt"));
1458 index.index_file(file_id, "hello world");
1459 index.finalize();
1460
1461 let results = index.search("hello");
1463 assert!(!results.is_empty());
1464
1465 let results = index.search("world");
1467 assert!(!results.is_empty());
1468
1469 let results = index.search("goodbye");
1471 assert!(results.is_empty());
1472 }
1473
1474 #[test]
1475 fn test_search_multifile() {
1476 let mut index = TrigramIndex::new();
1477
1478 let file1 = index.add_file(PathBuf::from("file1.txt"));
1479 let file2 = index.add_file(PathBuf::from("file2.txt"));
1480
1481 index.index_file(file1, "extract_symbols is here");
1482 index.index_file(file2, "extract_symbols is also here");
1483 index.finalize();
1484
1485 let results = index.search("extract_symbols");
1486 assert_eq!(results.len(), 2); let file_ids: Vec<u32> = results.iter().map(|loc| loc.file_id).collect();
1490 assert!(file_ids.contains(&file1));
1491 assert!(file_ids.contains(&file2));
1492 }
1493
1494 #[test]
1495 fn test_persistence_write() {
1496 use tempfile::TempDir;
1497
1498 let temp = TempDir::new().unwrap();
1499 let trigrams_path = temp.path().join("trigrams.bin");
1500
1501 let mut index = TrigramIndex::new();
1503 let file1 = index.add_file(PathBuf::from("src/main.rs"));
1504 let file2 = index.add_file(PathBuf::from("src/lib.rs"));
1505
1506 index.index_file(file1, "fn main() { println!(\"hello\"); }");
1507 index.index_file(
1508 file2,
1509 "pub fn hello() -> String { String::from(\"hello\") }",
1510 );
1511 index.finalize();
1512
1513 index.write(&trigrams_path).unwrap();
1515
1516 assert!(trigrams_path.exists());
1518
1519 let metadata = std::fs::metadata(&trigrams_path).unwrap();
1521 assert!(metadata.len() > HEADER_SIZE as u64);
1522
1523 use std::io::Read;
1525 let mut file = File::open(&trigrams_path).unwrap();
1526 let mut magic = [0u8; 4];
1527 file.read_exact(&mut magic).unwrap();
1528 assert_eq!(&magic, MAGIC);
1529
1530 }
1533
1534 #[test]
1535 fn test_posting_list_cap_enforced() {
1536 let cap: usize = 10;
1537 let content = "aaa ".repeat(200);
1538 let mut index = TrigramIndex::new();
1539 index.set_max_posting_list_entries(cap);
1540 let file_id = index.add_file(PathBuf::from("dense.txt"));
1541 index.index_file(file_id, &content);
1542 index.finalize();
1543 let aaa = bytes_to_trigram(b"aaa");
1544 let list = index
1545 .get_posting_list(aaa)
1546 .expect("aaa trigram should exist");
1547 assert!(list.len() <= cap, "cap exceeded: {} > {}", list.len(), cap);
1548 }
1549
1550 #[test]
1551 fn test_posting_list_cap_zero_means_unlimited() {
1552 let repetitions = 50;
1553 let content = "aaa ".repeat(repetitions);
1554 let mut index = TrigramIndex::new();
1555 index.set_max_posting_list_entries(0);
1556 let file_id = index.add_file(PathBuf::from("dense.txt"));
1557 index.index_file(file_id, &content);
1558 index.finalize();
1559 let aaa = bytes_to_trigram(b"aaa");
1560 let list = index
1561 .get_posting_list(aaa)
1562 .expect("aaa trigram should exist");
1563 assert!(
1564 list.len() >= repetitions,
1565 "expected >= {} entries, got {}",
1566 repetitions,
1567 list.len()
1568 );
1569 }
1570}