Skip to main content

reflex/
trigram.rs

1//! Trigram-based inverted index for fast full-text code search
2//!
3//! This module implements the core trigram indexing algorithm used by Reflex.
4//! A trigram is a sequence of 3 consecutive bytes. By building an inverted index
5//! mapping trigrams to file locations, we can quickly narrow down search candidates
6//! and achieve sub-100ms query times even on large codebases.
7//!
8//! # Algorithm
9//!
10//! 1. **Indexing**: Extract all trigrams from each file, store locations
11//! 2. **Querying**: Extract trigrams from query, intersect posting lists
12//! 3. **Verification**: Check actual matches at candidate locations
13//!
14//! See `.context/TRIGRAM_RESEARCH.md` for detailed algorithm documentation.
15
16use anyhow::{Context, Result};
17use std::collections::HashMap;
18use std::fs::{File, OpenOptions};
19use std::io::Write;
20use std::path::{Path, PathBuf};
21
22/// A trigram is 3 consecutive bytes, packed into a u32 for efficient hashing
23pub type Trigram = u32;
24
25// Binary format constants for trigrams.bin
26const MAGIC: &[u8; 4] = b"RFTG"; // ReFlex TriGrams
27const VERSION: u32 = 3; // V3: No filtering, lazy loading with directory + data separation
28// Header: magic(4) + version(4) + num_trigrams(8) + num_files(8) = 24 bytes
29#[allow(dead_code)]
30const HEADER_SIZE: usize = 24;
31
32/// Write a u32 as a varint (variable-length integer)
33/// Uses 1-5 bytes depending on magnitude (smaller numbers = fewer bytes)
34fn 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; // Set continuation bit
40        }
41        writer.write_all(&[byte])?;
42        if value == 0 {
43            break;
44        }
45    }
46    Ok(())
47}
48
49/// Read a varint from a byte slice, returns (value, bytes_consumed)
50fn 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
75/// Decompress a posting list from memory-mapped data
76///
77/// Reads a compressed posting list (delta+varint encoded) from the given offset
78/// and decompresses it into a Vec<FileLocation>.
79///
80/// # Arguments
81/// * `mmap` - Memory-mapped file data
82/// * `offset` - Absolute byte offset where compressed data starts
83/// * `size` - Number of bytes to read
84fn 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    // Decompress delta-encoded posting list
100    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        // Read file_id delta
108        let (file_id_delta, consumed) = read_varint(&compressed_data[pos..])?;
109        pos += consumed;
110
111        // Read line_no delta
112        let (line_no_delta, consumed) = read_varint(&compressed_data[pos..])?;
113        pos += consumed;
114
115        // Read byte_offset delta
116        let (byte_offset_delta, consumed) = read_varint(&compressed_data[pos..])?;
117        pos += consumed;
118
119        // Reconstruct absolute values from deltas
120        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        // Update previous values for next delta
131        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/// Location of a trigram occurrence in the codebase
140#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
141pub struct FileLocation {
142    /// File ID (index into file list)
143    pub file_id: u32,
144    /// Line number (1-indexed)
145    pub line_no: u32,
146    /// Byte offset in file (for context extraction)
147    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/// Directory entry for lazy-loaded trigram index
161///
162/// Maps each trigram to its compressed posting list location in the data section.
163/// Total size: 16 bytes per entry (4 + 8 + 4)
164#[derive(Debug, Clone)]
165struct DirectoryEntry {
166    /// The trigram value (for binary search)
167    trigram: Trigram,
168    /// Absolute byte offset in the file where compressed data starts
169    data_offset: u64,
170    /// Size of compressed posting list in bytes
171    compressed_size: u32,
172}
173
174/// Trigram-based inverted index
175///
176/// Maps each trigram to a sorted list of locations where it appears.
177/// Posting lists are kept sorted by (file_id, line_no) for efficient intersection.
178/// The index itself is kept sorted by trigram for O(log n) binary search.
179///
180/// Supports three modes:
181/// 1. **In-memory mode** (during indexing): All posting lists in RAM
182/// 2. **Batch-flush mode** (large codebases): Periodically flushes partial indices to disk to limit RAM
183/// 3. **Lazy-loaded mode** (after loading): Compressed posting lists in mmap, decompressed on-demand
184pub struct TrigramIndex {
185    /// Inverted index: sorted Vec of (trigram, locations) for binary search
186    /// Used in in-memory mode (during indexing)
187    index: Vec<(Trigram, Vec<FileLocation>)>,
188    /// File ID to file path mapping
189    files: Vec<PathBuf>,
190    /// Temporary HashMap used during batch indexing (None when finalized)
191    temp_index: Option<HashMap<Trigram, Vec<FileLocation>>>,
192    /// Memory-mapped index file (for lazy loading)
193    mmap: Option<memmap2::Mmap>,
194    /// Directory of (trigram, offset, size) for lazy loading
195    directory: Vec<DirectoryEntry>,
196    /// Partial index files created during batch flushing (for k-way merge at finalize)
197    partial_indices: Vec<PathBuf>,
198    /// Temporary directory for partial indices
199    temp_dir: Option<PathBuf>,
200    /// Cap on posting list size; 0 = unlimited. Enforced at finalize time.
201    /// Bounds query latency for high-frequency trigrams (trigram-density lens).
202    max_posting_list_entries: usize,
203}
204
205impl TrigramIndex {
206    /// Create a new empty trigram index
207    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    /// Set maximum posting list entries per trigram (0 = unlimited).
221    pub fn set_max_posting_list_entries(&mut self, cap: usize) {
222        self.max_posting_list_entries = cap;
223    }
224
225    /// Enable batch-flush mode for large codebases
226    ///
227    /// Creates a temporary directory for partial indices that will be merged at finalize().
228    /// Call this before indexing to enable memory-efficient indexing for huge codebases.
229    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    /// Flush current temp_index to a partial index file
238    ///
239    /// This clears the in-memory HashMap and writes a sorted partial index to disk.
240    /// Called periodically during indexing to limit memory usage.
241    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        // Take ownership of temp_index to finalize it
247        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            // Nothing to flush, restore empty map
254            self.temp_index = Some(HashMap::new());
255            return Ok(());
256        }
257
258        // Convert HashMap to sorted Vec
259        let mut partial_index: Vec<(Trigram, Vec<FileLocation>)> = temp_map.into_iter().collect();
260
261        // Sort and deduplicate posting lists
262        for (_, list) in partial_index.iter_mut() {
263            list.sort_unstable();
264            list.dedup();
265        }
266
267        // Sort by trigram
268        partial_index.sort_unstable_by_key(|(trigram, _)| *trigram);
269
270        // Write to temp file
271        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        // Create new empty temp_index for next batch
277        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    /// Write a partial index to disk (simplified format for merging)
289    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        // Write number of trigrams
305        writer.write_all(&(index.len() as u64).to_le_bytes())?;
306
307        // Write each (trigram, posting_list)
308        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    /// Add a file to the index and return its file_id
324    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    /// Get file path for a file_id
331    pub fn get_file(&self, file_id: u32) -> Option<&PathBuf> {
332        self.files.get(file_id as usize)
333    }
334
335    /// Get total number of files
336    pub fn file_count(&self) -> usize {
337        self.files.len()
338    }
339
340    /// Get total number of unique trigrams
341    pub fn trigram_count(&self) -> usize {
342        if !self.directory.is_empty() {
343            // Lazy-loaded mode
344            self.directory.len()
345        } else {
346            // In-memory mode
347            self.index.len()
348        }
349    }
350
351    /// Index a file's content
352    ///
353    /// Extracts all trigrams from the content and adds them to the inverted index.
354    /// Must call finalize() after indexing all files to prepare for searching.
355    pub fn index_file(&mut self, file_id: u32, content: &str) {
356        let trigrams = extract_trigrams_with_locations(content, file_id);
357
358        // Use the persistent HashMap for O(1) updates during batch processing
359        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    /// Build index from a collection of pre-extracted trigrams (bulk operation)
372    ///
373    /// This is much more efficient than calling index_file() multiple times,
374    /// as it builds the HashMap once instead of rebuilding it for each file.
375    pub fn build_from_trigrams(&mut self, trigrams: Vec<(Trigram, FileLocation)>) {
376        let mut temp_map: HashMap<Trigram, Vec<FileLocation>> = HashMap::new();
377
378        // Group trigrams into posting lists
379        for (trigram, location) in trigrams {
380            temp_map.entry(trigram).or_default().push(location);
381        }
382
383        // Convert to sorted Vec for binary search
384        self.index = temp_map.into_iter().collect();
385
386        // Clear temp_index since we're using the Vec directly
387        self.temp_index = None;
388
389        // Finalize immediately (sort and deduplicate)
390        self.finalize();
391    }
392
393    /// Finalize the index by sorting all posting lists and the index itself
394    ///
395    /// Must be called after all files are indexed, before querying.
396    /// Converts the HashMap to a sorted Vec for fast binary search.
397    ///
398    /// If batch flushing was enabled, finalization will be deferred until write()
399    /// is called, which will perform streaming merge directly to disk.
400    pub fn finalize(&mut self) {
401        // If we have partial indices from batch flushing, DON'T merge yet
402        // We'll do streaming merge in write() or write_with_streaming_merge()
403        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            // Flush final batch if temp_index is not empty
410            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            // Don't merge yet - write() will handle it
417            return;
418        }
419
420        // Standard finalization (no batch flushing)
421        // Convert HashMap to Vec if we have a temp index
422        if let Some(temp_map) = self.temp_index.take() {
423            self.index = temp_map.into_iter().collect();
424        }
425
426        // Sort, deduplicate, and cap posting lists
427        let cap = self.max_posting_list_entries;
428        for (trigram, list) in self.index.iter_mut() {
429            list.sort_unstable();
430            list.dedup(); // Remove duplicates (same trigram appearing multiple times on same line)
431            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        // Sort the index by trigram for binary search
443        self.index.sort_unstable_by_key(|(trigram, _)| *trigram);
444    }
445
446    /// Merge all partial indices directly to trigrams.bin using streaming k-way merge
447    ///
448    /// This avoids loading the entire index into RAM by:
449    /// 1. Opening all partial index files as readers
450    /// 2. Performing k-way merge using a priority queue
451    /// 3. Writing compressed posting lists directly to disk
452    /// 4. Never accumulating more than K posting lists in memory at once
453    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        // Open all partial indices as buffered readers
465        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            // Read number of trigrams (we don't need it for streaming merge)
480            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        // Helper to read next trigram from a reader
492        fn read_next_trigram(reader: &mut PartialIndexReader) -> Result<bool> {
493            // Try to read trigram
494            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                    // Read posting list size
500                    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                    // Read all locations for this trigram
505                    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        // Initialize: read first trigram from each reader
537        for reader in &mut readers {
538            read_next_trigram(reader)?;
539        }
540
541        // Priority queue entry for k-way merge
542        #[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                // Reverse for min-heap
551                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        // Build initial heap
565        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        // Open output file for writing
576        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        // Write placeholder header (we'll update it at the end)
586        writer.write_all(MAGIC)?;
587        writer.write_all(&VERSION.to_le_bytes())?;
588        writer.write_all(&0u64.to_le_bytes())?; // num_trigrams (placeholder)
589        writer.write_all(&(self.files.len() as u64).to_le_bytes())?; // num_files
590
591        // We'll build the directory as we go
592        let mut directory: Vec<DirectoryEntry> = Vec::new();
593        let mut num_trigrams = 0u64;
594
595        // K-way merge loop
596        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 this is a new trigram, write the previous one
603            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                // Compress and write this trigram's posting list
619                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            // Set current trigram
634            current_trigram = Some(entry.trigram);
635
636            // Merge this reader's posting list into accumulated list
637            merged_locations.extend_from_slice(&reader.current_posting_list);
638
639            // Advance this reader to next trigram
640            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        // Write final trigram
651        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        // Remember where data section ended (not used but kept for clarity)
686        let _data_end_pos = writer.stream_position()?;
687
688        // Write file paths after data section
689        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        // Flush before we rewrite the beginning
697        writer.flush()?;
698        drop(writer);
699
700        // Now we need to insert the directory at the beginning
701        // We'll read the data+files we just wrote, then rewrite the file with directory in between
702        use std::io::{Seek, SeekFrom};
703
704        // Read data and files sections
705        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        // Rewrite file with correct structure
713        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        // Write header with correct num_trigrams
720        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        // Write directory
726        for entry in &directory {
727            writer.write_all(&entry.trigram.to_le_bytes())?;
728            // Adjust data offset to account for directory size
729            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        // Write data and files sections
735        writer.write_all(&temp_data)?;
736
737        // Flush and sync
738        writer.flush()?;
739        writer.get_ref().sync_all()?;
740
741        // Clean up partial index files
742        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    /// Write a compressed posting list to the writer and return the compressed size
755    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        // Compress posting list using delta+varint encoding
763        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            // Compute deltas
769            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 deltas as varints
774            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            // Update previous values
779            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    /// Merge all partial indices into self.index (old in-memory approach - deprecated)
791    #[allow(dead_code)]
792    fn merge_partial_indices(&mut self) -> Result<()> {
793        use std::io::{BufReader, Read};
794
795        // Read all partial indices into memory (simplified approach for now)
796        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            // Read number of trigrams
804            let mut buf = [0u8; 8];
805            reader.read_exact(&mut buf)?;
806            let num_trigrams = u64::from_le_bytes(buf) as usize;
807
808            // Read each (trigram, posting_list)
809            for _ in 0..num_trigrams {
810                // Read trigram
811                let mut trigram_buf = [0u8; 4];
812                reader.read_exact(&mut trigram_buf)?;
813                let trigram = u32::from_le_bytes(trigram_buf);
814
815                // Read posting list size
816                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                // Read all locations
821                for _ in 0..list_len {
822                    let mut loc_buf = [0u8; 12]; // 3 * u32
823                    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        // Group by trigram
851        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        // Convert to sorted vec
857        self.index = index_map.into_iter().collect();
858
859        // Sort and deduplicate posting lists
860        for (_, list) in self.index.iter_mut() {
861            list.sort_unstable();
862            list.dedup();
863        }
864
865        // Sort by trigram
866        self.index.sort_unstable_by_key(|(trigram, _)| *trigram);
867
868        // Clean up partial index files
869        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    /// Search for a plain text pattern
882    ///
883    /// Returns candidate file locations that could contain the pattern.
884    /// Caller must verify actual matches.
885    ///
886    /// In lazy-loaded mode: Decompresses posting lists on-demand from mmap.
887    /// In in-memory mode: Uses pre-loaded posting lists.
888    pub fn search(&self, pattern: &str) -> Vec<FileLocation> {
889        if pattern.len() < 3 {
890            // Pattern too short for trigrams - caller must fall back to full scan
891            return vec![];
892        }
893
894        let trigrams = extract_trigrams(pattern);
895        if trigrams.is_empty() {
896            return vec![];
897        }
898
899        // Check if we're in lazy-loaded mode or in-memory mode
900        if let Some(ref mmap) = self.mmap {
901            // Lazy-loaded mode: decompress posting lists on-demand
902            let mut posting_lists: Vec<Vec<FileLocation>> = Vec::new();
903
904            for trigram in &trigrams {
905                // Binary search directory for this trigram
906                match self.directory.binary_search_by_key(trigram, |e| e.trigram) {
907                    Ok(idx) => {
908                        let entry = &self.directory[idx];
909                        // Decompress this posting list on-demand
910                        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                        // Trigram not found - pattern cannot match
928                        return vec![];
929                    }
930                }
931            }
932
933            if posting_lists.is_empty() || posting_lists.len() < trigrams.len() {
934                return vec![];
935            }
936
937            // Sort by list size (smallest first for efficient intersection)
938            posting_lists.sort_by_key(|list| list.len());
939
940            // Intersect posting lists (owned version)
941            intersect_by_file_owned(&posting_lists)
942        } else {
943            // In-memory mode: use pre-loaded index
944            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                // Some trigrams missing - pattern cannot match
960                return vec![];
961            }
962
963            // Sort by list size (smallest first for efficient intersection)
964            posting_lists.sort_by_key(|list| list.len());
965
966            // Intersect posting lists (reference version)
967            intersect_by_file(&posting_lists)
968        }
969    }
970
971    /// Get posting list for a specific trigram (for debugging)
972    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    /// Write the trigram index to disk
980    ///
981    /// Binary format V3 (lazy-loadable with directory + data separation):
982    /// - Header (24 bytes): magic, version, num_trigrams, num_files
983    /// - Directory Section (16 bytes per trigram):
984    ///   - trigram: u32 (4 bytes)
985    ///   - data_offset: u64 (8 bytes) - absolute offset in file
986    ///   - compressed_size: u32 (4 bytes) - size of compressed posting list
987    /// - Data Section (variable size):
988    ///   - Compressed posting lists (delta+varint encoded)
989    /// - File Paths Section (variable size):
990    ///   - path_len: varint
991    ///   - path_bytes: [u8; path_len]
992    pub fn write(&mut self, path: impl AsRef<Path>) -> Result<()> {
993        let path = path.as_ref();
994
995        // If we have partial indices from batch flushing, use streaming merge
996        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        // Standard write path (no batch flushing).
1005        // Non-atomic: writes directly to the target path with truncate(true).
1006        // On disk-full mid-write the file is left corrupt; the indexer fast-path
1007        // validates the "RFTG" magic bytes on re-index, forcing a clean rebuild.
1008        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        // Use a large buffer (16MB) for streaming writes
1016        let mut writer = std::io::BufWriter::with_capacity(16 * 1024 * 1024, file);
1017
1018        // Write header
1019        writer.write_all(MAGIC)?;
1020        writer.write_all(&VERSION.to_le_bytes())?;
1021        writer.write_all(&(self.index.len() as u64).to_le_bytes())?; // num_trigrams
1022        writer.write_all(&(self.files.len() as u64).to_le_bytes())?; // num_files
1023
1024        // Build directory and write compressed data in a single pass
1025        let mut directory: Vec<DirectoryEntry> = Vec::with_capacity(self.index.len());
1026
1027        // Calculate directory start and size
1028        let directory_start = HEADER_SIZE as u64;
1029        let directory_size = self.index.len() * 16;
1030
1031        // Reserve space for directory (we'll write it after data)
1032        let data_start = directory_start + directory_size as u64;
1033        let mut current_offset = data_start;
1034
1035        // We need to write in the correct order: header, directory, data, file paths
1036        // But we need data offsets to write directory
1037        // So we compress data first, then write header+directory+data
1038
1039        // Step 1: Compress all posting lists and track offsets
1040        let mut compressed_lists: Vec<(Trigram, Vec<u8>)> = Vec::with_capacity(self.index.len());
1041
1042        for (trigram, locations) in &self.index {
1043            // Compress the posting list
1044            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        // Step 2: Write directory
1074        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        // Step 3: Write data section (compressed posting lists)
1081        for (_, compressed) in &compressed_lists {
1082            writer.write_all(compressed)?;
1083        }
1084
1085        // Step 4: Write file paths
1086        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        // Flush and sync
1094        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    /// Load trigram index from disk using memory-mapped I/O with lazy loading
1108    ///
1109    /// Binary format V3: Only reads the directory and file paths, keeps posting lists compressed in mmap.
1110    /// Posting lists are decompressed on-demand during search queries.
1111    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        // Memory-map the file (keep it alive for lazy access)
1118        let mmap = unsafe {
1119            memmap2::Mmap::map(&file)
1120                .with_context(|| format!("Failed to mmap {}", path.display()))?
1121        };
1122
1123        // Validate header
1124        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        // Read directory (trigram → offset mappings) - fast, just metadata
1159        let mut directory = Vec::with_capacity(num_trigrams);
1160        let mut pos = HEADER_SIZE;
1161        let directory_size = num_trigrams * 16; // 16 bytes per entry
1162
1163        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 is already sorted by trigram (from write())
1196        directory.sort_unstable_by_key(|e| e.trigram);
1197
1198        // Calculate where file paths section starts (after header + directory + data)
1199        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        // Read file paths (varint-encoded lengths)
1204        let mut files = Vec::with_capacity(num_files);
1205        for _ in 0..num_files {
1206            // Read path length (varint)
1207            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(), // Empty in lazy mode
1230            files,
1231            temp_index: None,
1232            mmap: Some(mmap), // Keep mmap alive for lazy decompression!
1233            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
1247/// Extract all trigrams from text
1248///
1249/// Returns a vector of trigrams (without location info).
1250pub 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
1262/// Extract trigrams with file location information
1263///
1264/// Returns a vector of (trigram, location) pairs for building the inverted index.
1265pub 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        // Track newlines
1273        if byte == b'\n' {
1274            line_no += 1;
1275        }
1276
1277        // Extract trigram
1278        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/// Convert 3 bytes to a trigram (packed u32)
1289#[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/// Convert trigram back to bytes (for debugging)
1296#[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
1305/// Intersect posting lists by (file_id, line_no) pairs
1306///
1307/// Returns locations where ALL trigrams appear on the SAME line (not just in the same file).
1308/// This ensures accurate full-text matching.
1309fn 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    // Create a set of (file_id, line_no) pairs from the first list
1317    let mut candidates: HashSet<(u32, u32)> = lists[0]
1318        .iter()
1319        .map(|loc| (loc.file_id, loc.line_no))
1320        .collect();
1321
1322    // Intersect with (file_id, line_no) pairs from other lists
1323    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    // Convert back to FileLocation results
1330    let mut result = Vec::new();
1331    for &(file_id, line_no) in &candidates {
1332        // Find a location matching this (file_id, line_no) from the first list
1333        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
1345/// Intersect posting lists by (file_id, line_no) pairs (owned version for lazy-loading)
1346///
1347/// Similar to intersect_by_file() but works with owned Vec<Vec<FileLocation>>
1348/// instead of references. Used in lazy-loading mode where posting lists are decompressed on-demand.
1349///
1350/// Returns locations where ALL trigrams appear on the SAME line (not just in the same file).
1351fn 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    // Create a set of (file_id, line_no) pairs from the first list
1359    let mut candidates: HashSet<(u32, u32)> = lists[0]
1360        .iter()
1361        .map(|loc| (loc.file_id, loc.line_no))
1362        .collect();
1363
1364    // Intersect with (file_id, line_no) pairs from other lists
1365    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    // Convert back to FileLocation results
1372    let mut result = Vec::new();
1373    for &(file_id, line_no) in &candidates {
1374        // Find a location matching this (file_id, line_no) from the first list
1375        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        // "hello" → "hel", "ell", "llo"
1397        assert_eq!(trigrams.len(), 3);
1398
1399        // Verify trigrams are unique
1400        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        // "hello\nworld" has 9 trigrams:
1438        // "hel", "ell", "llo", "lo\n", "o\nw", "\nwo", "wor", "orl", "rld"
1439        assert_eq!(locs.len(), 9);
1440
1441        // First trigram should be on line 1
1442        assert_eq!(locs[0].1.line_no, 1);
1443
1444        // After newline, should be line 2
1445        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        // Search for "hello"
1462        let results = index.search("hello");
1463        assert!(!results.is_empty());
1464
1465        // Search for "world"
1466        let results = index.search("world");
1467        assert!(!results.is_empty());
1468
1469        // Search for "goodbye" (not in text)
1470        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); // One result per file
1487
1488        // Verify we got both files
1489        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        // Build and write index
1502        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        // Write to disk
1514        index.write(&trigrams_path).unwrap();
1515
1516        // Verify file was created
1517        assert!(trigrams_path.exists());
1518
1519        // Verify file has content (header + data)
1520        let metadata = std::fs::metadata(&trigrams_path).unwrap();
1521        assert!(metadata.len() > HEADER_SIZE as u64);
1522
1523        // Verify we can read the header back
1524        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        // Note: Full roundtrip test verifies write works correctly.
1531        // Load verification is tested in production via query performance tests.
1532    }
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}