sift_core/index/trigram/
search.rs

1use std::cmp::Reverse;
2use std::collections::BinaryHeap;
3
4use regex_syntax::ast::parse::Parser as AstParser;
5use regex_syntax::hir::literal::{ExtractKind, Extractor};
6use regex_syntax::hir::{self, Hir};
7
8use crate::index::{FileId, PlanMode, QueryPlanOutput};
9use crate::query::QuerySpec;
10
11use super::TrigramIndex;
12use super::key::Trigram;
13use super::storage;
14
15impl TrigramIndex {
16    /// Produce narrowed candidate files for the query.
17    /// Returns `None` if the query can't be narrowed (full scan required).
18    #[must_use]
19    pub fn candidates(&self, query: &QuerySpec<'_>) -> Option<Vec<crate::Candidate>> {
20        let arms = Self::extract_literal_arms(query)?;
21        Some(
22            self.candidate_file_ids(&arms)
23                .into_iter()
24                .filter_map(|id| {
25                    let fid = FileId::new(usize::try_from(id).ok()?);
26                    let fp = self.fingerprints.get(fid.get())?;
27                    Some(crate::Candidate::with_metadata(
28                        fp.path.clone(),
29                        self.root.join(&fp.path),
30                        Some(fp.size),
31                        None,
32                    ))
33                })
34                .collect(),
35        )
36    }
37
38    /// Returns an explanation of how a query would be handled.
39    #[must_use]
40    pub fn explain(&self, query: &QuerySpec<'_>) -> QueryPlanOutput {
41        let mode = match Self::extract_literal_arms(query) {
42            Some(_) => PlanMode::IndexedCandidates,
43            None => PlanMode::FullScan,
44        };
45        QueryPlanOutput {
46            pattern: query.patterns.to_vec().join("|"),
47            mode,
48        }
49    }
50
51    #[must_use]
52    pub(crate) fn all_files(&self) -> Vec<crate::Candidate> {
53        self.fingerprints
54            .iter()
55            .map(|fp| {
56                crate::Candidate::with_metadata(
57                    fp.path.clone(),
58                    self.root.join(&fp.path),
59                    Some(fp.size),
60                    None,
61                )
62            })
63            .collect()
64    }
65
66    fn candidate_file_ids(&self, arms: &[Vec<u8>]) -> Vec<u32> {
67        if arms.is_empty() {
68            return Vec::new();
69        }
70        if arms.len() == 1 {
71            return self.posting_ids_for_literal(&arms[0]).unwrap_or_default();
72        }
73        let mut id_lists: Vec<Vec<u32>> = Vec::with_capacity(arms.len());
74        for arm in arms {
75            if let Some(ids) = self.posting_ids_for_literal(arm) {
76                id_lists.push(ids);
77            }
78        }
79        Self::merge_sorted_runs(id_lists)
80    }
81
82    fn posting_ids_for_literal(&self, lit: &[u8]) -> Option<Vec<u32>> {
83        if lit.len() < 3 {
84            return None;
85        }
86        let trigrams: Vec<Trigram> = Trigram::windows(lit).collect();
87        if trigrams.is_empty() {
88            return None;
89        }
90        let mut slices: Vec<&[u8]> = Vec::with_capacity(trigrams.len());
91        for tri in &trigrams {
92            let s = self.posting_bytes_slice(*tri);
93            if s.is_empty() {
94                return None;
95            }
96            slices.push(s);
97        }
98        slices.sort_unstable_by_key(|slice| slice.len());
99        let ids = Self::intersect_sorted_slices(&slices);
100        if ids.is_empty() { None } else { Some(ids) }
101    }
102
103    fn posting_bytes_slice(&self, tri: Trigram) -> &[u8] {
104        let Some(entry) = self.lexicon.get(tri.to_bytes()) else {
105            return &[];
106        };
107        let start = usize::try_from(entry.offset).unwrap_or(usize::MAX);
108        let payload_len = self.postings.payload_len();
109        let end = self.lexicon.posting_byte_end(entry.offset, payload_len);
110        self.postings.slice(start, end.saturating_sub(start))
111    }
112
113    fn intersect_sorted_slices(slices: &[&[u8]]) -> Vec<u32> {
114        if slices.is_empty() {
115            return Vec::new();
116        }
117        if slices.len() == 1 {
118            return storage::postings::Postings::decode_sorted(slices[0])
119                .expect("postings validated at open");
120        }
121        let mut ordered: Vec<&[u8]> = slices.to_vec();
122        ordered.sort_unstable_by_key(|slice| slice.len());
123        let mut cur = storage::postings::Postings::decode_sorted(ordered[0])
124            .expect("postings validated at open");
125        for s in &ordered[1..] {
126            cur = storage::postings::Postings::intersect_sorted(&cur, s)
127                .expect("postings validated at open");
128            if cur.is_empty() {
129                break;
130            }
131        }
132        cur
133    }
134
135    fn merge_sorted_runs(lists: Vec<Vec<u32>>) -> Vec<u32> {
136        if lists.is_empty() {
137            return Vec::new();
138        }
139        if lists.len() == 1 {
140            return lists.into_iter().next().unwrap_or_default();
141        }
142
143        let total: usize = lists.iter().map(Vec::len).sum();
144        let mut heap: BinaryHeap<Reverse<(u32, usize)>> = BinaryHeap::with_capacity(lists.len());
145        let mut positions = vec![0usize; lists.len()];
146
147        for (list_idx, list) in lists.iter().enumerate() {
148            if let Some(&first) = list.first() {
149                heap.push(Reverse((first, list_idx)));
150            }
151        }
152
153        let mut out = Vec::with_capacity(total);
154        let mut last = None;
155        while let Some(Reverse((value, list_idx))) = heap.pop() {
156            if last != Some(value) {
157                out.push(value);
158                last = Some(value);
159            }
160
161            positions[list_idx] += 1;
162            if let Some(&next) = lists[list_idx].get(positions[list_idx]) {
163                heap.push(Reverse((next, list_idx)));
164            }
165        }
166        out
167    }
168
169    // -----------------------------------------------------------------------
170    // Literal extraction (absorbs TrigramPlanner)
171    // -----------------------------------------------------------------------
172
173    /// Extract literal byte arms from a query spec.
174    /// Returns `None` if no usable literals >= 3 bytes can be extracted.
175    fn extract_literal_arms(spec: &QuerySpec<'_>) -> Option<Vec<Vec<u8>>> {
176        if spec.invert_match() {
177            return None;
178        }
179        let mut literal_arms: Vec<Vec<u8>> = Vec::new();
180        for p in spec.patterns {
181            let arms = if spec.fixed_strings() {
182                Self::fixed_string_literals(p.as_bytes(), spec.case_insensitive())
183            } else {
184                Self::plan_pattern(
185                    p.as_str(),
186                    spec.case_insensitive(),
187                    spec.word_regexp(),
188                    spec.line_regexp(),
189                )?
190            };
191            for lit in arms {
192                if lit.len() < 3 {
193                    return None;
194                }
195                literal_arms.push(lit);
196            }
197        }
198        if literal_arms.is_empty() {
199            return None;
200        }
201        Some(literal_arms)
202    }
203
204    fn plan_pattern(
205        pattern: &str,
206        case_insensitive: bool,
207        word_regexp: bool,
208        line_regexp: bool,
209    ) -> Option<Vec<Vec<u8>>> {
210        let hir = Self::build_configured_hir(pattern, case_insensitive)?;
211        let shaped = Self::shape_hir(hir, word_regexp, line_regexp);
212        let lits = Self::extract_literals(&shaped);
213        if lits.is_empty() { None } else { Some(lits) }
214    }
215
216    fn build_configured_hir(pattern: &str, case_insensitive: bool) -> Option<Hir> {
217        let ast = AstParser::new().parse(pattern).ok()?;
218        let mut builder = regex_syntax::hir::translate::TranslatorBuilder::new();
219        builder.unicode(true);
220        if case_insensitive {
221            builder.case_insensitive(true);
222        }
223        let mut translator = builder.build();
224        let hir = translator.translate(pattern, &ast).ok()?;
225        Some(hir)
226    }
227
228    fn shape_hir(hir: Hir, word_regexp: bool, line_regexp: bool) -> Hir {
229        if line_regexp {
230            Self::wrap_line(hir)
231        } else if word_regexp {
232            Self::wrap_word(hir)
233        } else {
234            hir
235        }
236    }
237
238    fn wrap_word(hir: Hir) -> Hir {
239        Hir::concat(vec![
240            Hir::look(hir::Look::WordStartHalfUnicode),
241            hir,
242            Hir::look(hir::Look::WordEndHalfUnicode),
243        ])
244    }
245
246    fn wrap_line(hir: Hir) -> Hir {
247        Hir::concat(vec![
248            Hir::look(hir::Look::StartLF),
249            hir,
250            Hir::look(hir::Look::EndLF),
251        ])
252    }
253
254    fn extract_literals(hir: &Hir) -> Vec<Vec<u8>> {
255        let extractor_prefix = Extractor::new();
256        let extractor_suffix = {
257            let mut e = Extractor::new();
258            e.kind(ExtractKind::Suffix);
259            e
260        };
261
262        let seq_prefix = extractor_prefix.extract(hir);
263        let seq_suffix = extractor_suffix.extract(hir);
264
265        let lits_prefix = seq_prefix.literals();
266        let lits_suffix = seq_suffix.literals();
267
268        Self::pick_better_lits(lits_prefix, lits_suffix)
269    }
270
271    fn pick_better_lits(
272        lits_a: Option<&[regex_syntax::hir::literal::Literal]>,
273        lits_b: Option<&[regex_syntax::hir::literal::Literal]>,
274    ) -> Vec<Vec<u8>> {
275        fn total_bytes(lits: Option<&[regex_syntax::hir::literal::Literal]>) -> usize {
276            lits.map_or(0, |l| l.iter().map(|lit| lit.as_bytes().len()).sum())
277        }
278
279        let a_count = lits_a.map_or(0, <[regex_syntax::hir::literal::Literal]>::len);
280        let b_count = lits_b.map_or(0, <[regex_syntax::hir::literal::Literal]>::len);
281        let a_has = a_count > 0;
282        let b_has = b_count > 0;
283
284        let lits = match (a_has, b_has) {
285            (true, false) => lits_a,
286            (false, true) => lits_b,
287            (false, false) => return Vec::new(),
288            (true, true) => {
289                let a_total = total_bytes(lits_a);
290                let b_total = total_bytes(lits_b);
291                if a_total >= b_total { lits_a } else { lits_b }
292            }
293        };
294
295        let lits = match lits {
296            Some(l) if !l.is_empty() => l,
297            _ => return Vec::new(),
298        };
299
300        let mut out = Vec::new();
301        for lit in lits {
302            let bytes = lit.as_bytes();
303            if bytes.len() >= 3 {
304                out.push(bytes.to_vec());
305            }
306        }
307        out
308    }
309
310    fn fixed_string_literals(lit: &[u8], case_insensitive: bool) -> Vec<Vec<u8>> {
311        if case_insensitive {
312            vec![lit.to_ascii_lowercase()]
313        } else {
314            vec![lit.to_vec()]
315        }
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322    use crate::query::QueryFlags;
323    use std::path::Path;
324
325    fn encode(ids: &[u32]) -> Vec<u8> {
326        let mut buf = Vec::new();
327        let mut prev = 0u64;
328        for (i, &value) in ids.iter().enumerate() {
329            let raw = if i == 0 {
330                u64::from(value)
331            } else {
332                u64::from(value) - prev
333            };
334            let mut varint_buf = unsigned_varint::encode::u64_buffer();
335            let encoded = unsigned_varint::encode::u64(raw, &mut varint_buf);
336            buf.extend_from_slice(encoded);
337            prev = u64::from(value);
338        }
339        buf
340    }
341
342    fn narrow(
343        patterns: &[String],
344        case_insensitive: bool,
345        word_regexp: bool,
346        line_regexp: bool,
347    ) -> bool {
348        let mut flags = QueryFlags::empty();
349        if case_insensitive {
350            flags |= QueryFlags::CASE_INSENSITIVE;
351        }
352        if word_regexp {
353            flags |= QueryFlags::WORD_REGEXP;
354        }
355        if line_regexp {
356            flags |= QueryFlags::LINE_REGEXP;
357        }
358        let spec = QuerySpec { patterns, flags };
359        TrigramIndex::extract_literal_arms(&spec).is_some()
360    }
361
362    fn full_scan(
363        patterns: &[String],
364        case_insensitive: bool,
365        word_regexp: bool,
366        line_regexp: bool,
367    ) -> bool {
368        let mut flags = QueryFlags::empty();
369        if case_insensitive {
370            flags |= QueryFlags::CASE_INSENSITIVE;
371        }
372        if word_regexp {
373            flags |= QueryFlags::WORD_REGEXP;
374        }
375        if line_regexp {
376            flags |= QueryFlags::LINE_REGEXP;
377        }
378        let spec = QuerySpec { patterns, flags };
379        TrigramIndex::extract_literal_arms(&spec).is_none()
380    }
381
382    #[test]
383    fn merge_sorted_runs_preserves_order_and_uniqueness() {
384        let merged =
385            TrigramIndex::merge_sorted_runs(vec![vec![1, 3, 7], vec![1, 2, 7, 9], vec![4, 7, 8]]);
386        assert_eq!(merged, vec![1, 2, 3, 4, 7, 8, 9]);
387    }
388
389    #[test]
390    fn intersect_sorted_posting_byte_slices_handles_smallest_first_order() {
391        let a = encode(&[1, 3, 5, 7, 9]);
392        let b = encode(&[3, 7]);
393        let c = encode(&[0, 3, 4, 7, 8]);
394        let slices = vec![a.as_slice(), b.as_slice(), c.as_slice()];
395        let ids = TrigramIndex::intersect_sorted_slices(&slices);
396        assert_eq!(ids, vec![3, 7]);
397    }
398
399    #[test]
400    fn merge_sorted_runs_empty_input_returns_empty() {
401        let merged = TrigramIndex::merge_sorted_runs(vec![]);
402        assert!(merged.is_empty());
403    }
404
405    #[test]
406    fn merge_sorted_runs_single_list_returns_as_is() {
407        let merged = TrigramIndex::merge_sorted_runs(vec![vec![1, 2, 3]]);
408        assert_eq!(merged, vec![1, 2, 3]);
409    }
410
411    #[test]
412    fn merge_sorted_runs_with_empty_lists_mixed_in() {
413        let merged = TrigramIndex::merge_sorted_runs(vec![vec![1, 3], vec![], vec![2, 3]]);
414        assert_eq!(merged, vec![1, 2, 3]);
415    }
416
417    #[test]
418    fn intersect_sorted_posting_byte_slices_empty_input_returns_empty() {
419        let ids = TrigramIndex::intersect_sorted_slices(&[]);
420        assert!(ids.is_empty());
421    }
422
423    #[test]
424    fn intersect_sorted_slices_single_returns_decoded_ids() {
425        let a = encode(&[1, 3, 5]);
426        let ids = TrigramIndex::intersect_sorted_slices(&[a.as_slice()]);
427        assert_eq!(ids, vec![1, 3, 5]);
428    }
429
430    #[test]
431    #[should_panic(expected = "postings validated at open")]
432    fn intersect_sorted_slices_invalid_varint_panics() {
433        let a = &[0xff];
434        TrigramIndex::intersect_sorted_slices(&[a]);
435    }
436
437    #[test]
438    fn intersect_sorted_slices_no_overlap_returns_empty() {
439        let a = encode(&[1, 2, 3]);
440        let b = encode(&[4, 5, 6]);
441        let ids = TrigramIndex::intersect_sorted_slices(&[a.as_slice(), b.as_slice()]);
442        assert!(ids.is_empty());
443    }
444
445    #[test]
446    fn literal_narrows() {
447        assert!(narrow(&["beta".to_string()], false, false, false));
448    }
449
450    #[test]
451    fn dot_star_full_scan() {
452        assert!(full_scan(&[".*".to_string()], false, false, false));
453    }
454
455    #[test]
456    fn alternation_narrows() {
457        assert!(narrow(&[r"foo|bar".to_string()], false, false, false));
458    }
459
460    #[test]
461    fn word_literal_narrows() {
462        assert!(narrow(&["beta".to_string()], false, true, false));
463    }
464
465    #[test]
466    fn line_regexp_narrows() {
467        assert!(narrow(&["beta".to_string()], false, false, true));
468    }
469
470    #[test]
471    fn case_insensitive_narrows() {
472        assert!(narrow(&["beta".to_string()], true, false, false));
473    }
474
475    #[test]
476    fn required_literal_inside_regex_narrows() {
477        assert!(narrow(&["[A-Z]+_RESUME".to_string()], false, false, false));
478    }
479
480    #[test]
481    fn unicode_class_full_scan() {
482        assert!(full_scan(&[r"\p{Greek}".to_string()], false, false, false));
483    }
484
485    #[test]
486    fn no_literal_full_scan() {
487        assert!(full_scan(
488            &[r"\w{5}\s+\w{5}".to_string()],
489            false,
490            false,
491            false
492        ));
493    }
494
495    #[test]
496    fn short_literal_full_scan() {
497        assert!(full_scan(&["ab".to_string()], false, false, false));
498    }
499
500    #[test]
501    fn fixed_string_narrows() {
502        let spec = QuerySpec {
503            patterns: &["beta.gamma".to_string()],
504            flags: QueryFlags::FIXED_STRINGS,
505        };
506        assert!(TrigramIndex::extract_literal_arms(&spec).is_some());
507    }
508
509    #[test]
510    fn open_tables_accepts_count_mismatch() {
511        use crate::index::trigram::storage::format::{
512            FILES_MAGIC, LEXICON_MAGIC, POSTINGS_MAGIC, TRIGRAMS_MAGIC,
513        };
514        use tempfile::TempDir;
515        let tmp = TempDir::new().expect("create temp dir");
516        let dir = tmp.path().join("index");
517        std::fs::create_dir(&dir).expect("create index dir");
518
519        let mut files = FILES_MAGIC.to_vec();
520        files.extend_from_slice(&0u32.to_le_bytes());
521        std::fs::write(dir.join("files.bin"), &files).expect("write files");
522
523        let mut lex = LEXICON_MAGIC.to_vec();
524        lex.extend_from_slice(&1u32.to_le_bytes());
525        lex.extend_from_slice(b"abc");
526        lex.extend_from_slice(&0u64.to_le_bytes());
527        lex.extend_from_slice(&3u32.to_le_bytes());
528        std::fs::write(dir.join("lexicon.bin"), &lex).expect("write lexicon");
529
530        let mut posting_payload = Vec::new();
531        let mut buf = unsigned_varint::encode::u64_buffer();
532        posting_payload.extend_from_slice(unsigned_varint::encode::u64(0, &mut buf));
533        let mut buf2 = unsigned_varint::encode::u64_buffer();
534        posting_payload.extend_from_slice(unsigned_varint::encode::u64(1, &mut buf2));
535        let mut pb = POSTINGS_MAGIC.to_vec();
536        pb.extend_from_slice(&u32::try_from(posting_payload.len()).unwrap().to_le_bytes());
537        pb.extend_from_slice(&posting_payload);
538        std::fs::write(dir.join("postings.bin"), &pb).expect("write postings");
539
540        let mut tri = TRIGRAMS_MAGIC.to_vec();
541        tri.extend_from_slice(&0u32.to_le_bytes());
542        std::fs::write(dir.join("trigrams.bin"), &tri).expect("write trigrams");
543
544        // Posting count mismatches are caught at build time.
545        // The open path skips content-level validation for speed.
546        let result = TrigramIndex::open(
547            &dir,
548            Path::new("/root"),
549            crate::index::CorpusKind::Directory,
550        );
551        assert!(result.is_ok());
552    }
553}
sift_core/index/trigram/search.rs

sift_core/index/trigram/
search.rs
