sift_core/index/trigram/
mod.rs

1pub mod builder;
2pub mod file_table;
3pub mod key;
4pub mod storage;
5
6mod planner;
7
8use std::cmp::Reverse;
9use std::collections::BinaryHeap;
10use std::path::{Path, PathBuf};
11
12use crate::index::{CorpusKind, FileId, IndexConfig, PlanMode, QueryPlanOutput};
13use crate::query::QuerySpec;
14
15use self::builder::{IndexTables, build_tables};
16use self::file_table::FileFingerprint;
17use self::planner::{TrigramCandidatePlan, TrigramPlanner};
18pub use key::Trigram;
19
20/// Errors specific to opening or persisting a trigram index.
21#[derive(Debug, thiserror::Error)]
22pub enum TrigramIndexError {
23    #[error("index component missing: {0}")]
24    MissingComponent(PathBuf),
25
26    #[error("IO error: {0}")]
27    Io(#[from] std::io::Error),
28}
29
30#[derive(Debug)]
31pub struct TrigramIndex {
32    root: PathBuf,
33    pub(crate) fingerprints: Vec<FileFingerprint>,
34    trigram_sets: storage::trigram_sets::TrigramSets,
35    lexicon: storage::lexicon::Lexicon,
36    postings: storage::postings::Postings,
37    corpus_kind: CorpusKind,
38}
39
40impl TrigramIndex {
41    /// Write tables to `dir` as persistence files and return an mmap-backed index.
42    pub(crate) fn create_in_dir(
43        tables: &IndexTables,
44        root: &Path,
45        corpus_kind: CorpusKind,
46        dir: &Path,
47    ) -> crate::Result<Self> {
48        std::fs::create_dir_all(dir)?;
49
50        let files_path = dir.join(crate::FILES_BIN);
51        let lexicon_path = dir.join(crate::LEXICON_BIN);
52        let postings_path = dir.join(crate::POSTINGS_BIN);
53        let trigrams_path = dir.join(crate::TRIGRAMS_BIN);
54
55        let ((fr, lr), (pr, tr)) = rayon::join(
56            || {
57                rayon::join(
58                    || file_table::FileTable::create(&files_path, &tables.fingerprints),
59                    || storage::lexicon::Lexicon::create(&lexicon_path, &tables.lexicon),
60                )
61            },
62            || {
63                rayon::join(
64                    || storage::postings::Postings::create(&postings_path, &tables.postings),
65                    || {
66                        storage::trigram_sets::TrigramSets::create(
67                            &trigrams_path,
68                            &tables.file_trigrams,
69                        )
70                    },
71                )
72            },
73        );
74
75        let files = fr.map_err(crate::Error::Io)?;
76        let lexicon = lr.map_err(crate::Error::Io)?;
77        let postings = pr.map_err(crate::Error::Io)?;
78        let trigram_sets = tr.map_err(crate::Error::Io)?;
79
80        let root = root.to_path_buf();
81        let fingerprints = files.to_fingerprints().map_err(crate::Error::Io)?;
82        Self::validate_file_paths(&fingerprints, &files_path)?;
83
84        Self::validate_lexicon_postings(&lexicon, &postings)?;
85
86        Ok(Self {
87            root,
88            fingerprints,
89            trigram_sets,
90            lexicon,
91            postings,
92            corpus_kind,
93        })
94    }
95
96    #[must_use]
97    pub fn file_path(&self, id: FileId) -> Option<&Path> {
98        self.fingerprints.get(id.get()).map(|fp| fp.path.as_path())
99    }
100
101    #[must_use]
102    pub fn file_abs_path(&self, id: FileId) -> Option<PathBuf> {
103        self.fingerprints
104            .get(id.get())
105            .map(|fp| self.root.join(&fp.path))
106    }
107
108    /// Returns an explanation of how a query would be handled.
109    #[must_use]
110    pub fn explain(&self, query: &QuerySpec<'_>) -> QueryPlanOutput {
111        let mode = match TrigramPlanner::build(query) {
112            Some(_) => PlanMode::IndexedCandidates,
113            None => PlanMode::FullScan,
114        };
115        QueryPlanOutput {
116            pattern: query.patterns.to_vec().join("|"),
117            mode,
118        }
119    }
120
121    #[must_use]
122    pub fn root(&self) -> &Path {
123        &self.root
124    }
125
126    #[must_use]
127    pub const fn corpus_kind(&self) -> CorpusKind {
128        self.corpus_kind
129    }
130
131    #[must_use]
132    pub fn candidates(&self, query: &QuerySpec<'_>) -> Vec<crate::Candidate> {
133        TrigramPlanner::build(query).map_or_else(
134            || self.resolve_all_candidates(),
135            |plan| self.resolve_candidates(self.trigram_candidate_ids(&plan)),
136        )
137    }
138
139    #[must_use]
140    pub fn all_files(&self) -> Vec<crate::Candidate> {
141        self.resolve_all_candidates()
142    }
143
144    /// Build a new trigram index from the corpus described in `config`.
145    ///
146    /// # Errors
147    ///
148    /// Returns an error if corpus walking, extraction, or file I/O fails.
149    pub fn build(config: &IndexConfig<'_>, output_dir: &Path) -> crate::Result<Self> {
150        let tables = build_tables(config)?;
151        let root = config.corpus.root.canonicalize()?;
152        Self::create_in_dir(&tables, &root, config.corpus.kind, output_dir)
153    }
154
155    /// Open a previously persisted trigram index from `index_dir`.
156    ///
157    /// # Errors
158    ///
159    /// Returns an error if persistence files are missing or malformed.
160    pub fn open(index_dir: &Path, root: &Path, corpus_kind: CorpusKind) -> crate::Result<Self> {
161        Ok(Self::open_tables(index_dir, root, corpus_kind)?)
162    }
163
164    /// Update the index from the current corpus, reusing per-file trigram sets
165    /// for unchanged files.
166    ///
167    /// Returns `Ok(Some(index))` if a new snapshot was written, or `Ok(None)`
168    /// if no files changed.
169    ///
170    /// # Errors
171    ///
172    /// Returns an error if corpus walking, extraction, or file I/O fails.
173    pub fn update(
174        &self,
175        config: &IndexConfig<'_>,
176        output_dir: &Path,
177    ) -> crate::Result<Option<Self>> {
178        use rayon::prelude::*;
179        use std::collections::HashMap;
180
181        let paths = crate::index::trigram::builder::CorpusWalker::new(config).collect()?;
182        let fingerprints =
183            crate::index::trigram::builder::FingerprintCollector::new(config.corpus.root, &paths)
184                .collect()?;
185
186        if fingerprints == self.fingerprints {
187            return Ok(None);
188        }
189
190        let previous_sets = self.trigram_sets.to_vec().map_err(crate::Error::Io)?;
191
192        let lookup: HashMap<(&Path, i64, u64), &storage::trigram_sets::TrigramSet> = self
193            .fingerprints
194            .iter()
195            .zip(previous_sets.iter())
196            .map(|(fp, set)| ((fp.path.as_path(), fp.mtime_secs, fp.size), set))
197            .collect();
198
199        let file_trigrams: Vec<storage::trigram_sets::TrigramSet> = fingerprints
200            .par_iter()
201            .map(|fp| {
202                if let Some(set) = lookup.get(&(fp.path.as_path(), fp.mtime_secs, fp.size)) {
203                    return Ok((*set).clone());
204                }
205                let abs = config.corpus.root.join(&fp.path);
206                storage::trigram_sets::TrigramSet::from_file(&abs).map_err(crate::Error::Io)
207            })
208            .collect::<crate::Result<_>>()?;
209
210        let (lexicon, postings) =
211            crate::index::trigram::builder::PostingAssembler::new(&file_trigrams).assemble()?;
212
213        let tables = IndexTables {
214            fingerprints,
215            file_trigrams,
216            lexicon,
217            postings,
218        };
219
220        let root = config.corpus.root.canonicalize()?;
221        let index = Self::create_in_dir(&tables, &root, config.corpus.kind, output_dir)?;
222        Ok(Some(index))
223    }
224
225    fn posting_bytes_slice(&self, tri: Trigram) -> &[u8] {
226        let Some(entry) = self.lexicon.get(tri.to_bytes()) else {
227            return &[];
228        };
229        let start = usize::try_from(entry.offset).unwrap_or(usize::MAX);
230        let payload_len = self.postings.payload_len();
231        let end = self.lexicon.posting_byte_end(entry.offset, payload_len);
232        self.postings.slice(start, end.saturating_sub(start))
233    }
234
235    fn candidate_file_ids(&self, arms: &[Vec<u8>]) -> Vec<u32> {
236        if arms.is_empty() {
237            return Vec::new();
238        }
239        if arms.len() == 1 {
240            return self.posting_ids_for_literal(&arms[0]).unwrap_or_default();
241        }
242        let mut id_lists: Vec<Vec<u32>> = Vec::with_capacity(arms.len());
243        for arm in arms {
244            if let Some(ids) = self.posting_ids_for_literal(arm) {
245                id_lists.push(ids);
246            }
247        }
248        PostingOps::merge_sorted_runs(id_lists)
249    }
250
251    fn posting_ids_for_literal(&self, lit: &[u8]) -> Option<Vec<u32>> {
252        if lit.len() < 3 {
253            return None;
254        }
255        let trigrams: Vec<Trigram> = Trigram::windows(lit).collect();
256        if trigrams.is_empty() {
257            return None;
258        }
259        let mut slices: Vec<&[u8]> = Vec::with_capacity(trigrams.len());
260        for tri in &trigrams {
261            let s = self.posting_bytes_slice(*tri);
262            if s.is_empty() {
263                return None;
264            }
265            slices.push(s);
266        }
267        slices.sort_unstable_by_key(|slice| slice.len());
268        let ids = PostingOps::intersect_sorted_slices(&slices);
269        if ids.is_empty() { None } else { Some(ids) }
270    }
271
272    fn trigram_candidate_ids(&self, plan: &TrigramCandidatePlan) -> Vec<FileId> {
273        let raw = self.candidate_file_ids(&plan.arms);
274        raw.into_iter()
275            .filter_map(|id| usize::try_from(id).ok().map(FileId::new))
276            .collect()
277    }
278
279    fn candidate_from_fingerprint(&self, fp: &FileFingerprint) -> crate::Candidate {
280        let rel_path = fp.path.clone();
281        let abs_path = self.root.join(&fp.path);
282        crate::Candidate::new(rel_path, abs_path)
283    }
284
285    fn resolve_candidates(&self, ids: impl IntoIterator<Item = FileId>) -> Vec<crate::Candidate> {
286        ids.into_iter()
287            .filter_map(|id| {
288                self.fingerprints
289                    .get(id.get())
290                    .map(|fp| self.candidate_from_fingerprint(fp))
291            })
292            .collect()
293    }
294
295    fn resolve_all_candidates(&self) -> Vec<crate::Candidate> {
296        self.fingerprints
297            .iter()
298            .map(|fp| self.candidate_from_fingerprint(fp))
299            .collect()
300    }
301
302    fn open_tables(
303        dir: &Path,
304        root: &Path,
305        corpus_kind: CorpusKind,
306    ) -> Result<Self, TrigramIndexError> {
307        let files_path = dir.join(crate::FILES_BIN);
308        let lexicon_path = dir.join(crate::LEXICON_BIN);
309        let postings_path = dir.join(crate::POSTINGS_BIN);
310        let trigrams_path = dir.join(crate::TRIGRAMS_BIN);
311
312        for p in [&files_path, &lexicon_path, &postings_path, &trigrams_path] {
313            if !p.is_file() {
314                return Err(TrigramIndexError::MissingComponent(p.clone()));
315            }
316        }
317
318        let files = file_table::FileTable::open(&files_path).map_err(TrigramIndexError::Io)?;
319        let fingerprints = files.to_fingerprints().map_err(TrigramIndexError::Io)?;
320        Self::validate_file_paths(&fingerprints, &files_path)?;
321
322        let lexicon =
323            storage::lexicon::Lexicon::open(&lexicon_path).map_err(TrigramIndexError::Io)?;
324        let postings =
325            storage::postings::Postings::open(&postings_path).map_err(TrigramIndexError::Io)?;
326
327        Self::validate_lexicon_postings(&lexicon, &postings)?;
328
329        let trigram_sets = storage::trigram_sets::TrigramSets::open(&trigrams_path)
330            .map_err(TrigramIndexError::Io)?;
331
332        Ok(Self {
333            root: root.to_path_buf(),
334            fingerprints,
335            trigram_sets,
336            lexicon,
337            postings,
338            corpus_kind,
339        })
340    }
341
342    fn validate_lexicon_postings(
343        lexicon: &storage::lexicon::Lexicon,
344        postings: &storage::postings::Postings,
345    ) -> Result<(), TrigramIndexError> {
346        let payload_len = postings.payload_len();
347        for entry in lexicon {
348            let start = usize::try_from(entry.offset).map_err(|_| {
349                TrigramIndexError::Io(std::io::Error::new(
350                    std::io::ErrorKind::InvalidData,
351                    format!(
352                        "lexicon entry {:?} offset {} exceeds usize",
353                        entry.trigram, entry.offset,
354                    ),
355                ))
356            })?;
357            let end = lexicon.posting_byte_end(entry.offset, payload_len);
358            if start > end || end > payload_len {
359                return Err(TrigramIndexError::Io(std::io::Error::new(
360                    std::io::ErrorKind::InvalidData,
361                    format!(
362                        "lexicon entry {:?} posting range [{start},{end}) exceeds payload_len {payload_len}",
363                        entry.trigram,
364                    ),
365                )));
366            }
367            let slice = postings.slice(start, end.saturating_sub(start));
368            let decoded_count = storage::postings::Postings::validate_list(slice).map_err(|e| {
369                TrigramIndexError::Io(std::io::Error::new(
370                    std::io::ErrorKind::InvalidData,
371                    format!("posting list for trigram {:?}: {e}", entry.trigram),
372                ))
373            })?;
374            if decoded_count != entry.len as usize {
375                return Err(TrigramIndexError::Io(std::io::Error::new(
376                    std::io::ErrorKind::InvalidData,
377                    format!(
378                        "lexicon entry {:?} claims len {} but posting list has {decoded_count} entries",
379                        entry.trigram, entry.len,
380                    ),
381                )));
382            }
383        }
384        Ok(())
385    }
386
387    fn validate_file_paths(
388        fingerprints: &[FileFingerprint],
389        _meta_path: &Path,
390    ) -> Result<(), TrigramIndexError> {
391        for fp in fingerprints {
392            if fp.path.as_os_str().is_empty()
393                || fp.path.is_absolute()
394                || fp
395                    .path
396                    .components()
397                    .any(|c| matches!(c, std::path::Component::ParentDir))
398            {
399                return Err(TrigramIndexError::Io(std::io::Error::new(
400                    std::io::ErrorKind::InvalidData,
401                    format!("invalid file path in index: {}", fp.path.display()),
402                )));
403            }
404        }
405        Ok(())
406    }
407}
408
409struct PostingOps;
410
411impl PostingOps {
412    fn intersect_sorted_slices(slices: &[&[u8]]) -> Vec<u32> {
413        if slices.is_empty() {
414            return Vec::new();
415        }
416        if slices.len() == 1 {
417            return storage::postings::Postings::decode_sorted(slices[0])
418                .expect("postings validated at open");
419        }
420        let mut ordered: Vec<&[u8]> = slices.to_vec();
421        ordered.sort_unstable_by_key(|slice| slice.len());
422        let mut cur = storage::postings::Postings::decode_sorted(ordered[0])
423            .expect("postings validated at open");
424        for s in &ordered[1..] {
425            cur = storage::postings::Postings::intersect_sorted(&cur, s)
426                .expect("postings validated at open");
427            if cur.is_empty() {
428                break;
429            }
430        }
431        cur
432    }
433
434    fn merge_sorted_runs(lists: Vec<Vec<u32>>) -> Vec<u32> {
435        if lists.is_empty() {
436            return Vec::new();
437        }
438        if lists.len() == 1 {
439            return lists.into_iter().next().unwrap_or_default();
440        }
441
442        let total: usize = lists.iter().map(Vec::len).sum();
443        let mut heap: BinaryHeap<Reverse<(u32, usize)>> = BinaryHeap::with_capacity(lists.len());
444        let mut positions = vec![0usize; lists.len()];
445
446        for (list_idx, list) in lists.iter().enumerate() {
447            if let Some(&first) = list.first() {
448                heap.push(Reverse((first, list_idx)));
449            }
450        }
451
452        let mut out = Vec::with_capacity(total);
453        let mut last = None;
454        while let Some(Reverse((value, list_idx))) = heap.pop() {
455            if last != Some(value) {
456                out.push(value);
457                last = Some(value);
458            }
459
460            positions[list_idx] += 1;
461            if let Some(&next) = lists[list_idx].get(positions[list_idx]) {
462                heap.push(Reverse((next, list_idx)));
463            }
464        }
465        out
466    }
467}
468
469#[cfg(test)]
470mod tests {
471    use super::*;
472    use std::path::PathBuf;
473
474    fn encode(ids: &[u32]) -> Vec<u8> {
475        let mut buf = Vec::new();
476        let mut prev = 0u64;
477        for (i, &value) in ids.iter().enumerate() {
478            let raw = if i == 0 {
479                u64::from(value)
480            } else {
481                u64::from(value) - prev
482            };
483            let mut varint_buf = unsigned_varint::encode::u64_buffer();
484            let encoded = unsigned_varint::encode::u64(raw, &mut varint_buf);
485            buf.extend_from_slice(encoded);
486            prev = u64::from(value);
487        }
488        buf
489    }
490
491    #[test]
492    fn merge_sorted_runs_preserves_order_and_uniqueness() {
493        let merged =
494            PostingOps::merge_sorted_runs(vec![vec![1, 3, 7], vec![1, 2, 7, 9], vec![4, 7, 8]]);
495        assert_eq!(merged, vec![1, 2, 3, 4, 7, 8, 9]);
496    }
497
498    #[test]
499    fn intersect_sorted_posting_byte_slices_handles_smallest_first_order() {
500        let a = encode(&[1, 3, 5, 7, 9]);
501        let b = encode(&[3, 7]);
502        let c = encode(&[0, 3, 4, 7, 8]);
503        let slices = vec![a.as_slice(), b.as_slice(), c.as_slice()];
504        let ids = PostingOps::intersect_sorted_slices(&slices);
505        assert_eq!(ids, vec![3, 7]);
506    }
507
508    #[test]
509    fn merge_sorted_runs_empty_input_returns_empty() {
510        let merged = PostingOps::merge_sorted_runs(vec![]);
511        assert!(merged.is_empty());
512    }
513
514    #[test]
515    fn merge_sorted_runs_single_list_returns_as_is() {
516        let merged = PostingOps::merge_sorted_runs(vec![vec![1, 2, 3]]);
517        assert_eq!(merged, vec![1, 2, 3]);
518    }
519
520    #[test]
521    fn merge_sorted_runs_with_empty_lists_mixed_in() {
522        let merged = PostingOps::merge_sorted_runs(vec![vec![1, 3], vec![], vec![2, 3]]);
523        assert_eq!(merged, vec![1, 2, 3]);
524    }
525
526    #[test]
527    fn intersect_sorted_posting_byte_slices_empty_input_returns_empty() {
528        let ids = PostingOps::intersect_sorted_slices(&[]);
529        assert!(ids.is_empty());
530    }
531
532    #[test]
533    fn intersect_sorted_slices_single_returns_decoded_ids() {
534        let a = encode(&[1, 3, 5]);
535        let ids = PostingOps::intersect_sorted_slices(&[a.as_slice()]);
536        assert_eq!(ids, vec![1, 3, 5]);
537    }
538
539    #[test]
540    #[should_panic(expected = "postings validated at open")]
541    fn intersect_sorted_slices_invalid_varint_panics() {
542        let a = &[0xff];
543        let _ids = PostingOps::intersect_sorted_slices(&[a]);
544    }
545
546    #[test]
547    fn intersect_sorted_slices_no_overlap_returns_empty() {
548        let a = encode(&[1, 2, 3]);
549        let b = encode(&[4, 5, 6]);
550        let ids = PostingOps::intersect_sorted_slices(&[a.as_slice(), b.as_slice()]);
551        assert!(ids.is_empty());
552    }
553
554    #[test]
555    fn open_tables_rejects_count_mismatch() {
556        use crate::index::trigram::storage::format::{
557            FILES_MAGIC, LEXICON_MAGIC, POSTINGS_MAGIC, TRIGRAMS_MAGIC,
558        };
559        use tempfile::TempDir;
560        let tmp = TempDir::new().expect("create temp dir");
561        let dir = tmp.path().join("index");
562        std::fs::create_dir(&dir).expect("create index dir");
563
564        // files.bin: empty
565        let mut files = FILES_MAGIC.to_vec();
566        files.extend_from_slice(&0u32.to_le_bytes());
567        std::fs::write(dir.join("files.bin"), &files).expect("write files");
568
569        // lexicon.bin: one entry, trigram "abc", offset=0, len=3
570        let mut lex = LEXICON_MAGIC.to_vec();
571        lex.extend_from_slice(&1u32.to_le_bytes()); // count=1
572        lex.extend_from_slice(b"abc");
573        lex.extend_from_slice(&0u64.to_le_bytes()); // offset=0
574        lex.extend_from_slice(&3u32.to_le_bytes()); // len=3 — claims 3 entries
575        std::fs::write(dir.join("lexicon.bin"), &lex).expect("write lexicon");
576
577        // postings.bin: only 2 encoded entries (lexicon claims 3)
578        let mut posting_payload = Vec::new();
579        let mut buf = unsigned_varint::encode::u64_buffer();
580        posting_payload.extend_from_slice(unsigned_varint::encode::u64(0, &mut buf));
581        let mut buf2 = unsigned_varint::encode::u64_buffer();
582        posting_payload.extend_from_slice(unsigned_varint::encode::u64(1, &mut buf2));
583        let mut pb = POSTINGS_MAGIC.to_vec();
584        pb.extend_from_slice(&u32::try_from(posting_payload.len()).unwrap().to_le_bytes());
585        pb.extend_from_slice(&posting_payload);
586        std::fs::write(dir.join("postings.bin"), &pb).expect("write postings");
587
588        // trigrams.bin: empty
589        let mut tri = TRIGRAMS_MAGIC.to_vec();
590        tri.extend_from_slice(&0u32.to_le_bytes());
591        std::fs::write(dir.join("trigrams.bin"), &tri).expect("write trigrams");
592
593        let result = TrigramIndex::open_tables(
594            &dir,
595            Path::new("/root"),
596            crate::index::CorpusKind::Directory,
597        );
598        assert!(result.is_err());
599        let err = result.unwrap_err().to_string();
600        assert!(
601            err.contains("claims len") || err.contains("entries"),
602            "expected count mismatch error, got: {err}",
603        );
604    }
605
606    #[test]
607    fn validate_file_paths_accepts_normal_relative_paths() {
608        let fps = vec![
609            FileFingerprint {
610                path: PathBuf::from("a.txt"),
611                mtime_secs: 0,
612                size: 0,
613            },
614            FileFingerprint {
615                path: PathBuf::from("sub/b.txt"),
616                mtime_secs: 0,
617                size: 0,
618            },
619        ];
620        let result = TrigramIndex::validate_file_paths(&fps, Path::new("/meta.json"));
621        assert!(result.is_ok());
622    }
623
624    #[test]
625    fn validate_file_paths_rejects_absolute_paths() {
626        let abs = std::env::current_dir().unwrap().join("a.txt");
627        let fps = vec![FileFingerprint {
628            path: abs,
629            mtime_secs: 0,
630            size: 0,
631        }];
632        let result = TrigramIndex::validate_file_paths(&fps, Path::new("/meta.json"));
633        assert!(result.is_err());
634    }
635
636    #[test]
637    fn validate_file_paths_rejects_empty_paths() {
638        let fps = vec![FileFingerprint {
639            path: PathBuf::from(""),
640            mtime_secs: 0,
641            size: 0,
642        }];
643        let result = TrigramIndex::validate_file_paths(&fps, Path::new("/meta.json"));
644        assert!(result.is_err());
645    }
646
647    #[test]
648    fn validate_file_paths_rejects_parent_dir_paths() {
649        let fps = vec![FileFingerprint {
650            path: PathBuf::from("../escape.txt"),
651            mtime_secs: 0,
652            size: 0,
653        }];
654        let result = TrigramIndex::validate_file_paths(&fps, Path::new("/meta.json"));
655        assert!(result.is_err());
656    }
657}
sift_core/index/trigram/mod.rs

sift_core/index/trigram/
mod.rs
