1use std::collections::{HashMap, HashSet};
27use std::path::{Path, PathBuf};
28use std::sync::{Arc, Mutex, OnceLock};
29use std::time::{Duration, Instant};
30
31use ignore::WalkBuilder;
32
33use crate::tools::ctx_search::{is_binary_ext, is_generated_file, MAX_FILE_SIZE, MAX_WALK_DEPTH};
34
35const TTL: Duration = Duration::from_secs(15);
38
39const MAX_FILES: usize = 200_000;
41
42const MAX_POSTING_ENTRIES: usize = 12_000_000;
47
48const MAX_TOTAL_ENTRIES: usize = 48_000_000;
52
53const BLOOM_BITS_PER_ITEM: usize = 12;
57const BLOOM_K: usize = 7;
58const BLOOM_MIN_BITS: usize = 64;
60const BLOOM_MAX_BITS: usize = 1 << 20;
61
62fn is_word_byte(b: u8) -> bool {
64 b.is_ascii_alphanumeric() || b == b'_'
65}
66
67fn pack(b0: u8, b1: u8, b2: u8) -> u32 {
68 (u32::from(b0) << 16) | (u32::from(b1) << 8) | u32::from(b2)
69}
70
71enum Narrowing {
80 Postings(HashMap<u32, Vec<u32>>),
81 Blooms(Vec<FileBloom>),
82}
83
84struct FileBloom {
87 bits: Vec<u64>,
89}
90
91#[inline]
94fn mix64(mut x: u64) -> u64 {
95 x = x.wrapping_add(0x9E37_79B9_7F4A_7C15);
96 x = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9);
97 x = (x ^ (x >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB);
98 x ^ (x >> 31)
99}
100
101impl FileBloom {
102 fn with_capacity(distinct_trigrams: usize) -> Self {
103 let target = distinct_trigrams
104 .saturating_mul(BLOOM_BITS_PER_ITEM)
105 .next_power_of_two()
106 .clamp(BLOOM_MIN_BITS, BLOOM_MAX_BITS);
107 FileBloom {
108 bits: vec![0u64; target / 64],
109 }
110 }
111
112 #[inline]
113 fn m_bits(&self) -> usize {
114 self.bits.len() * 64
115 }
116
117 #[inline]
119 fn probes(&self, trigram: u32) -> impl Iterator<Item = usize> + '_ {
120 let m = self.m_bits();
121 let mask = m - 1; let h = mix64(u64::from(trigram));
123 let h1 = (h & 0xFFFF_FFFF) as usize;
124 let h2 = ((h >> 32) as usize) | 1; (0..BLOOM_K).map(move |i| h1.wrapping_add(i.wrapping_mul(h2)) & mask)
126 }
127
128 fn insert(&mut self, trigram: u32) {
129 for p in self.probes(trigram).collect::<Vec<_>>() {
130 self.bits[p / 64] |= 1u64 << (p % 64);
131 }
132 }
133
134 fn maybe_contains(&self, trigram: u32) -> bool {
135 self.probes(trigram)
136 .all(|p| self.bits[p / 64] & (1u64 << (p % 64)) != 0)
137 }
138}
139
140pub struct SearchIndex {
142 files: Vec<PathBuf>,
143 narrowing: Narrowing,
145 respect_gitignore: bool,
146 allow_secret_paths: bool,
147 built_at: Instant,
148}
149
150impl SearchIndex {
151 pub fn build(root: &str, respect_gitignore: bool, allow_secret_paths: bool) -> Option<Self> {
154 let root_path = Path::new(root);
155 if !root_path.exists() {
156 return None;
157 }
158
159 let walker = WalkBuilder::new(root_path)
160 .hidden(true)
161 .max_depth(Some(MAX_WALK_DEPTH))
162 .git_ignore(respect_gitignore)
163 .git_global(respect_gitignore)
164 .git_exclude(respect_gitignore)
165 .build();
166
167 let mut files: Vec<PathBuf> = Vec::new();
168 let mut per_file_trigrams: Vec<Vec<u32>> = Vec::new();
172 let mut total_entries: usize = 0;
173 let mut scratch: HashSet<u32> = HashSet::new();
174
175 for entry in walker.filter_map(std::result::Result::ok) {
176 if entry.file_type().is_none_or(|ft| ft.is_dir()) {
177 continue;
178 }
179 if entry.file_type().is_some_and(|ft| ft.is_symlink()) {
180 continue;
181 }
182 let path = entry.path();
183 if is_binary_ext(path) || is_generated_file(path) {
184 continue;
185 }
186 if !allow_secret_paths && crate::core::io_boundary::is_secret_like(path).is_some() {
187 continue;
188 }
189 if let Ok(meta) = std::fs::metadata(path) {
190 if meta.len() > MAX_FILE_SIZE {
191 continue;
192 }
193 }
194 let Ok(content) = std::fs::read_to_string(path) else {
197 continue;
198 };
199
200 if files.len() >= MAX_FILES {
201 return None; }
203
204 scratch.clear();
205 let bytes = content.as_bytes();
206 if bytes.len() >= 3 {
207 for w in bytes.windows(3) {
208 if is_word_byte(w[0]) && is_word_byte(w[1]) && is_word_byte(w[2]) {
209 scratch.insert(pack(w[0], w[1], w[2]));
210 }
211 }
212 }
213 total_entries += scratch.len();
214 if total_entries > MAX_TOTAL_ENTRIES {
215 return None; }
217 let mut tris: Vec<u32> = scratch.iter().copied().collect();
218 tris.sort_unstable();
219 files.push(path.to_path_buf());
220 per_file_trigrams.push(tris);
221 }
222
223 let narrowing = build_narrowing(&per_file_trigrams, total_entries);
224
225 Some(Self {
226 files,
227 narrowing,
228 respect_gitignore,
229 allow_secret_paths,
230 built_at: Instant::now(),
231 })
232 }
233
234 fn is_fresh(&self) -> bool {
235 self.built_at.elapsed() < TTL
236 }
237
238 fn config_matches(&self, respect_gitignore: bool, allow_secret_paths: bool) -> bool {
239 self.respect_gitignore == respect_gitignore && self.allow_secret_paths == allow_secret_paths
240 }
241
242 pub fn candidate_paths(&self, pattern: &str, ext: Option<&str>) -> CandidateSet {
251 if let Some(ids) = self.literal_candidates(pattern) {
252 let paths = ids
253 .into_iter()
254 .map(|id| self.files[id as usize].clone())
255 .filter(|p| ext_matches(p, ext))
256 .collect();
257 CandidateSet::Narrowed(paths)
258 } else {
259 let paths = self
260 .files
261 .iter()
262 .filter(|p| ext_matches(p, ext))
263 .cloned()
264 .collect();
265 CandidateSet::FullList(paths)
266 }
267 }
268
269 fn literal_candidates(&self, pattern: &str) -> Option<Vec<u32>> {
273 let bytes = pattern.as_bytes();
274 if bytes.len() < 3 || !bytes.iter().all(|&b| is_word_byte(b)) {
275 return None;
276 }
277 let mut tris: Vec<u32> = bytes.windows(3).map(|w| pack(w[0], w[1], w[2])).collect();
279 tris.sort_unstable();
280 tris.dedup();
281
282 match &self.narrowing {
283 Narrowing::Postings(trigrams) => Some(Self::postings_intersect(trigrams, &tris)),
284 Narrowing::Blooms(blooms) => Some(Self::bloom_scan(blooms, &tris)),
285 }
286 }
287
288 fn postings_intersect(trigrams: &HashMap<u32, Vec<u32>>, tris: &[u32]) -> Vec<u32> {
291 let mut lists: Vec<&Vec<u32>> = Vec::with_capacity(tris.len());
292 for &tri in tris {
293 match trigrams.get(&tri) {
294 None => return Vec::new(),
296 Some(list) => lists.push(list),
297 }
298 }
299 lists.sort_by_key(|l| l.len());
300
301 let mut acc: Vec<u32> = lists[0].clone();
302 for list in &lists[1..] {
303 acc = intersect_sorted(&acc, list);
304 if acc.is_empty() {
305 break;
306 }
307 }
308 acc
309 }
310
311 fn bloom_scan(blooms: &[FileBloom], tris: &[u32]) -> Vec<u32> {
315 let mut out = Vec::new();
316 for (fid, bloom) in blooms.iter().enumerate() {
317 if tris.iter().all(|&t| bloom.maybe_contains(t)) {
318 out.push(fid as u32);
319 }
320 }
321 out
322 }
323}
324
325fn build_narrowing(per_file: &[Vec<u32>], total_entries: usize) -> Narrowing {
327 if total_entries <= MAX_POSTING_ENTRIES {
328 let mut trigrams: HashMap<u32, Vec<u32>> = HashMap::new();
329 for (fid, tris) in per_file.iter().enumerate() {
330 for &t in tris {
331 trigrams.entry(t).or_default().push(fid as u32);
333 }
334 }
335 Narrowing::Postings(trigrams)
336 } else {
337 let blooms = per_file
338 .iter()
339 .map(|tris| {
340 let mut b = FileBloom::with_capacity(tris.len());
341 for &t in tris {
342 b.insert(t);
343 }
344 b
345 })
346 .collect();
347 Narrowing::Blooms(blooms)
348 }
349}
350
351pub enum CandidateSet {
353 Narrowed(Vec<PathBuf>),
355 FullList(Vec<PathBuf>),
357}
358
359impl CandidateSet {
360 pub fn into_paths(self) -> Vec<PathBuf> {
361 match self {
362 CandidateSet::Narrowed(p) | CandidateSet::FullList(p) => p,
363 }
364 }
365}
366
367fn ext_matches(path: &Path, ext: Option<&str>) -> bool {
368 match ext {
369 None => true,
370 Some(want) => path.extension().and_then(|e| e.to_str()) == Some(want),
371 }
372}
373
374fn intersect_sorted(a: &[u32], b: &[u32]) -> Vec<u32> {
376 let mut out = Vec::new();
377 let (mut i, mut j) = (0, 0);
378 while i < a.len() && j < b.len() {
379 match a[i].cmp(&b[j]) {
380 std::cmp::Ordering::Less => i += 1,
381 std::cmp::Ordering::Greater => j += 1,
382 std::cmp::Ordering::Equal => {
383 out.push(a[i]);
384 i += 1;
385 j += 1;
386 }
387 }
388 }
389 out
390}
391
392struct CacheEntry {
397 index: Option<Arc<SearchIndex>>,
398 building: bool,
399}
400
401static CACHE: OnceLock<Mutex<HashMap<String, CacheEntry>>> = OnceLock::new();
402
403fn cache() -> &'static Mutex<HashMap<String, CacheEntry>> {
404 CACHE.get_or_init(|| Mutex::new(HashMap::new()))
405}
406
407fn index_disabled() -> bool {
410 std::env::var("LEAN_CTX_DISABLE_SEARCH_INDEX")
411 .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
412}
413
414pub fn get_fresh(
418 root: &str,
419 respect_gitignore: bool,
420 allow_secret_paths: bool,
421) -> Option<Arc<SearchIndex>> {
422 if !respect_gitignore || index_disabled() {
424 return None;
425 }
426
427 let mut needs_build = false;
428 let result = {
429 let mut map = cache()
430 .lock()
431 .unwrap_or_else(std::sync::PoisonError::into_inner);
432 let entry = map.entry(root.to_string()).or_insert(CacheEntry {
433 index: None,
434 building: false,
435 });
436 match &entry.index {
437 Some(idx)
438 if idx.config_matches(respect_gitignore, allow_secret_paths) && idx.is_fresh() =>
439 {
440 Some(Arc::clone(idx))
441 }
442 Some(idx) if idx.config_matches(respect_gitignore, allow_secret_paths) => {
443 needs_build = !entry.building;
445 if needs_build {
446 entry.building = true;
447 }
448 Some(Arc::clone(idx))
449 }
450 _ => {
451 needs_build = !entry.building;
452 if needs_build {
453 entry.building = true;
454 }
455 None
456 }
457 }
458 };
459
460 if needs_build {
461 spawn_build(root.to_string(), respect_gitignore, allow_secret_paths);
462 }
463 result
464}
465
466pub fn ensure_background(root: &str, respect_gitignore: bool, allow_secret_paths: bool) {
469 if !respect_gitignore || index_disabled() {
470 return;
471 }
472 let needs_build = {
473 let mut map = cache()
474 .lock()
475 .unwrap_or_else(std::sync::PoisonError::into_inner);
476 let entry = map.entry(root.to_string()).or_insert(CacheEntry {
477 index: None,
478 building: false,
479 });
480 let fresh = entry.index.as_ref().is_some_and(|idx| {
481 idx.config_matches(respect_gitignore, allow_secret_paths) && idx.is_fresh()
482 });
483 if fresh || entry.building {
484 false
485 } else {
486 entry.building = true;
487 true
488 }
489 };
490 if needs_build {
491 spawn_build(root.to_string(), respect_gitignore, allow_secret_paths);
492 }
493}
494
495pub fn warm_blocking(root: &str, respect_gitignore: bool, allow_secret_paths: bool) -> bool {
499 if !respect_gitignore || index_disabled() {
500 return false;
501 }
502 let Some(idx) = SearchIndex::build(root, respect_gitignore, allow_secret_paths) else {
503 return false;
504 };
505 let mut map = cache()
506 .lock()
507 .unwrap_or_else(std::sync::PoisonError::into_inner);
508 map.insert(
509 root.to_string(),
510 CacheEntry {
511 index: Some(Arc::new(idx)),
512 building: false,
513 },
514 );
515 true
516}
517
518fn spawn_build(root: String, respect_gitignore: bool, allow_secret_paths: bool) {
519 std::thread::spawn(move || {
520 let built = std::panic::catch_unwind(|| {
521 SearchIndex::build(&root, respect_gitignore, allow_secret_paths)
522 })
523 .ok()
524 .flatten();
525
526 let mut map = cache()
527 .lock()
528 .unwrap_or_else(std::sync::PoisonError::into_inner);
529 if let Some(entry) = map.get_mut(&root) {
530 entry.building = false;
531 if let Some(idx) = built {
532 entry.index = Some(Arc::new(idx));
533 }
534 }
535 });
536}
537
538#[cfg(test)]
539mod tests {
540 use super::*;
541
542 fn corpus() -> tempfile::TempDir {
543 let dir = tempfile::tempdir().unwrap();
544 std::fs::write(dir.path().join("a.rs"), "fn handler() {}\nlet x = 1;\n").unwrap();
545 std::fs::write(dir.path().join("b.rs"), "fn other() {}\n// nothing here\n").unwrap();
546 std::fs::write(dir.path().join("c.txt"), "handler appears in text too\n").unwrap();
547 dir
548 }
549
550 #[test]
551 fn narrows_to_files_containing_literal() {
552 let dir = corpus();
553 let idx = SearchIndex::build(dir.path().to_str().unwrap(), true, false).unwrap();
554 let cands = idx.candidate_paths("handler", None);
555 let paths = cands.into_paths();
556 assert!(paths.iter().any(|p| p.ends_with("a.rs")));
558 assert!(paths.iter().any(|p| p.ends_with("c.txt")));
559 assert!(!paths.iter().any(|p| p.ends_with("b.rs")));
560 }
561
562 #[test]
563 fn absent_trigram_yields_empty_candidates() {
564 let dir = corpus();
565 let idx = SearchIndex::build(dir.path().to_str().unwrap(), true, false).unwrap();
566 match idx.candidate_paths("zzzqqq", None) {
567 CandidateSet::Narrowed(p) => assert!(p.is_empty()),
568 CandidateSet::FullList(_) => panic!("pure literal should narrow"),
569 }
570 }
571
572 #[test]
573 fn ext_filter_restricts_candidates() {
574 let dir = corpus();
575 let idx = SearchIndex::build(dir.path().to_str().unwrap(), true, false).unwrap();
576 let paths = idx.candidate_paths("handler", Some("rs")).into_paths();
577 assert!(paths.iter().all(|p| p.extension().unwrap() == "rs"));
578 assert!(paths.iter().any(|p| p.ends_with("a.rs")));
579 }
580
581 #[test]
582 fn regex_query_falls_back_to_full_list() {
583 let dir = corpus();
584 let idx = SearchIndex::build(dir.path().to_str().unwrap(), true, false).unwrap();
585 match idx.candidate_paths("fn .*\\(\\)", None) {
586 CandidateSet::FullList(p) => assert!(!p.is_empty()),
587 CandidateSet::Narrowed(_) => panic!("metachar query must not narrow"),
588 }
589 }
590
591 #[test]
592 fn short_query_falls_back_to_full_list() {
593 let dir = corpus();
594 let idx = SearchIndex::build(dir.path().to_str().unwrap(), true, false).unwrap();
595 assert!(matches!(
596 idx.candidate_paths("fn", None),
597 CandidateSet::FullList(_)
598 ));
599 }
600
601 #[test]
605 fn narrowing_has_identical_recall_to_full_scan() {
606 use regex::Regex;
607 use std::collections::BTreeSet;
608
609 let dir = tempfile::tempdir().unwrap();
610 let samples = [
612 (
613 "auth/login.rs",
614 "fn authenticate(user) {}\nlet token = mint();\n",
615 ),
616 (
617 "auth/session.rs",
618 "struct Session;\n// authenticate again here\n",
619 ),
620 ("db/pool.rs", "fn connect() {}\nlet retries = 3;\n"),
621 (
622 "ui/button.tsx",
623 "export const Button = () => authenticate;\n",
624 ),
625 ("readme.md", "This project uses authenticate flows.\n"),
626 ("unrelated.rs", "fn helper() { let v = 1; }\n"),
627 ];
628 for (rel, content) in samples {
629 let p = dir.path().join(rel);
630 std::fs::create_dir_all(p.parent().unwrap()).unwrap();
631 std::fs::write(p, content).unwrap();
632 }
633 let idx = SearchIndex::build(dir.path().to_str().unwrap(), true, false).unwrap();
634
635 let full_scan = |pat: &str| -> BTreeSet<String> {
636 let re = Regex::new(pat).unwrap();
637 let mut hits = BTreeSet::new();
638 for (rel, content) in samples {
639 for (i, line) in content.lines().enumerate() {
640 if re.is_match(line) {
641 hits.insert(format!("{rel}:{}", i + 1));
642 }
643 }
644 }
645 hits
646 };
647
648 for query in ["authenticate", "Session", "retries", "token"] {
649 let re = Regex::new(query).unwrap();
650 let candidates = idx.candidate_paths(query, None).into_paths();
651 let mut narrowed = BTreeSet::new();
652 for path in &candidates {
653 let content = std::fs::read_to_string(path).unwrap();
654 let rel = path.strip_prefix(dir.path()).unwrap().to_string_lossy();
655 for (i, line) in content.lines().enumerate() {
656 if re.is_match(line) {
657 narrowed.insert(format!("{}:{}", rel.replace('\\', "/"), i + 1));
658 }
659 }
660 }
661 assert_eq!(
662 narrowed,
663 full_scan(query),
664 "recall mismatch for query {query:?}"
665 );
666 }
667 }
668
669 #[test]
670 fn intersect_sorted_basic() {
671 assert_eq!(
672 intersect_sorted(&[1, 2, 3, 5], &[2, 3, 4, 5]),
673 vec![2, 3, 5]
674 );
675 assert_eq!(intersect_sorted(&[1, 2], &[3, 4]), Vec::<u32>::new());
676 }
677
678 fn trigrams_of(s: &str) -> Vec<u32> {
681 let mut set = HashSet::new();
682 let b = s.as_bytes();
683 if b.len() >= 3 {
684 for w in b.windows(3) {
685 if is_word_byte(w[0]) && is_word_byte(w[1]) && is_word_byte(w[2]) {
686 set.insert(pack(w[0], w[1], w[2]));
687 }
688 }
689 }
690 let mut v: Vec<u32> = set.into_iter().collect();
691 v.sort_unstable();
692 v
693 }
694
695 #[test]
696 fn file_bloom_has_no_false_negatives() {
697 let tris = trigrams_of("fn authenticate(user) { let token = mint(); }");
698 let mut bloom = FileBloom::with_capacity(tris.len());
699 for &t in &tris {
700 bloom.insert(t);
701 }
702 assert!(tris.iter().all(|&t| bloom.maybe_contains(t)));
704 }
705
706 #[test]
710 fn bloom_tier_is_superset_of_postings_tier() {
711 let mut seed = 0x1234_5678_9abc_def0u64;
713 let mut rng = || {
714 seed = seed
715 .wrapping_mul(6364136223846793005)
716 .wrapping_add(1442695040888963407);
717 (seed >> 33) as u32
718 };
719 let mut per_file: Vec<Vec<u32>> = Vec::new();
720 for _ in 0..80 {
721 let n = 50 + (rng() % 250) as usize;
722 let mut s = HashSet::new();
723 for _ in 0..n {
724 s.insert(rng() & 0x00FF_FFFF);
725 }
726 let mut v: Vec<u32> = s.into_iter().collect();
727 v.sort_unstable();
728 per_file.push(v);
729 }
730 let total: usize = per_file.iter().map(Vec::len).sum();
731
732 let postings = build_narrowing(&per_file, total); let blooms = build_narrowing(&per_file, MAX_POSTING_ENTRIES + 1); let (Narrowing::Postings(pt), Narrowing::Blooms(bl)) = (&postings, &blooms) else {
735 panic!("unexpected narrowing tiers");
736 };
737
738 for f in &per_file {
741 if f.len() < 3 {
742 continue;
743 }
744 let q = vec![f[0], f[f.len() / 2], f[f.len() - 1]];
745 let exact = SearchIndex::postings_intersect(pt, &q);
746 let bloom: HashSet<u32> = SearchIndex::bloom_scan(bl, &q).into_iter().collect();
747 for id in exact {
748 assert!(
749 bloom.contains(&id),
750 "Bloom tier dropped a true match (false negative) for {q:?}"
751 );
752 }
753 }
754 }
755
756 #[test]
759 fn bloom_tier_end_to_end_recall() {
760 let samples = [
761 (
762 "auth_login.rs",
763 "fn authenticate(user) {}\nlet token = mint();\n",
764 ),
765 (
766 "auth_session.rs",
767 "struct Session;\n// authenticate again here\n",
768 ),
769 ("db_pool.rs", "fn connect() {}\nlet retries = 3;\n"),
770 (
771 "ui_button.tsx",
772 "export const Button = () => authenticate;\n",
773 ),
774 ("readme.md", "This project uses authenticate flows.\n"),
775 ("unrelated.rs", "fn helper() { let v = 1; }\n"),
776 ];
777 let files: Vec<PathBuf> = samples.iter().map(|(rel, _)| PathBuf::from(rel)).collect();
778 let per_file: Vec<Vec<u32>> = samples.iter().map(|(_, c)| trigrams_of(c)).collect();
779
780 let idx = SearchIndex {
781 files,
782 narrowing: build_narrowing(&per_file, MAX_POSTING_ENTRIES + 1),
783 respect_gitignore: true,
784 allow_secret_paths: false,
785 built_at: Instant::now(),
786 };
787 assert!(
788 matches!(idx.narrowing, Narrowing::Blooms(_)),
789 "test must exercise the Bloom tier"
790 );
791
792 for query in ["authenticate", "Session", "retries", "token"] {
793 let cands: HashSet<String> = idx
794 .candidate_paths(query, None)
795 .into_paths()
796 .iter()
797 .map(|p| p.to_string_lossy().to_string())
798 .collect();
799 for (rel, content) in samples {
800 if content.contains(query) {
801 assert!(
802 cands.contains(rel),
803 "Bloom tier dropped real match {rel} for query {query:?}"
804 );
805 }
806 }
807 }
808 }
809}