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

sift_core/index/trigram/
search.rs
