1use arc_swap::ArcSwapOption;
33use parking_lot::Mutex;
34use serde::{Deserialize, Serialize};
35use std::cmp::Reverse;
36use std::collections::BinaryHeap;
37use std::num::NonZero;
38use std::path::Path;
39use std::sync::Arc;
40use std::sync::atomic::{AtomicBool, Ordering};
41use tokio::sync::RwLock;
42
43use rayon::prelude::*;
44use vtcode_commons::StringId;
45
46pub struct FileIndex {
51 files: Vec<StringId>,
52 directories: Vec<StringId>,
53 path_texts_by_id: Arc<Vec<Arc<str>>>,
57 interner: vtcode_commons::StringInterner,
59 last_built: std::time::Instant,
60}
61
62fn build_parallel_walker(
64 search_directory: &Path,
65 exclude: &[String],
66 threads: usize,
67 respect_gitignore: bool,
68 follow_links: bool,
69) -> anyhow::Result<ignore::WalkParallel> {
70 let mut walk_builder = ignore::WalkBuilder::new(search_directory);
71 vtcode_commons::walk::apply_defaults(&mut walk_builder);
72
73 walk_builder.threads(threads);
75 walk_builder.follow_links(follow_links);
76 walk_builder.require_git(false); if !respect_gitignore {
79 walk_builder
80 .git_ignore(false)
81 .git_global(false)
82 .git_exclude(false)
83 .ignore(false)
84 .parents(false);
85 }
86
87 if !exclude.is_empty() {
88 let mut override_builder = ignore::overrides::OverrideBuilder::new(search_directory);
89 for exclude_pattern in exclude {
90 let pattern = format!("!{exclude_pattern}");
91 override_builder.add(&pattern)?;
92 }
93 walk_builder.overrides(override_builder.build()?);
94 }
95
96 Ok(walk_builder.build_parallel())
97}
98
99impl FileIndex {
100 fn build_from_directory(
103 search_directory: &Path,
104 exclude: &[String],
105 respect_gitignore: bool,
106 threads: usize,
107 ) -> anyhow::Result<Self> {
108 let walker = build_parallel_walker(search_directory, exclude, threads, respect_gitignore, true)?;
109
110 let worker_results = Arc::new(Mutex::new(Vec::new()));
112
113 walker.run(|| {
114 let search_dir = search_directory.to_path_buf();
115 let mut state = IndexWorkerGuard {
116 results: Arc::clone(&worker_results),
117 files: Vec::new(),
118 directories: Vec::new(),
119 };
120
121 Box::new(move |result| {
122 let entry = match result {
123 Ok(e) => e,
124 Err(_) => return ignore::WalkState::Continue,
125 };
126
127 if let Some(rel_path) = entry.path().strip_prefix(&search_dir).ok().and_then(|p| p.to_str())
129 && !rel_path.is_empty()
130 {
131 if entry.file_type().is_some_and(|file_type| file_type.is_dir()) {
132 state.directories.push(rel_path.to_string());
133 } else {
134 state.files.push(rel_path.to_string());
135 }
136 }
137
138 ignore::WalkState::Continue
139 })
140 });
141
142 let worker_results = Arc::try_unwrap(worker_results)
143 .map_err(|results| {
144 anyhow::anyhow!("file index worker results still referenced: {}", Arc::strong_count(&results))
145 })?
146 .into_inner();
147 let (files, directories) =
148 worker_results
149 .into_iter()
150 .fold((Vec::new(), Vec::new()), |(mut files, mut directories), result| {
151 files.extend(result.files);
152 directories.extend(result.directories);
153 (files, directories)
154 });
155
156 let mut interner = vtcode_commons::StringInterner::new();
157 let mut path_texts_by_id = Vec::with_capacity(files.len() + directories.len());
158 let interned_files: Vec<StringId> = files
159 .iter()
160 .map(|path| intern_path(path, &mut interner, &mut path_texts_by_id))
161 .collect();
162 let interned_dirs: Vec<StringId> = directories
163 .iter()
164 .map(|path| intern_path(path, &mut interner, &mut path_texts_by_id))
165 .collect();
166
167 Ok(Self {
168 files: interned_files,
169 directories: interned_dirs,
170 path_texts_by_id: Arc::new(path_texts_by_id),
171 interner,
172 last_built: std::time::Instant::now(),
173 })
174 }
175
176 fn query(
179 &self,
180 pattern_text: &str,
181 limit: usize,
182 match_type_filter: Option<MatchType>,
183 ) -> Vec<(u32, StringId, MatchType)> {
184 let mut heaps = Vec::new();
189
190 if match_type_filter.is_none_or(|t| t == MatchType::File) {
191 heaps.push(score_paths_top_k(
192 &self.files,
193 self.path_texts_by_id.as_slice(),
194 limit,
195 pattern_text,
196 MatchType::File,
197 ));
198 }
199
200 if match_type_filter.is_none_or(|t| t == MatchType::Directory) {
201 heaps.push(score_paths_top_k(
202 &self.directories,
203 self.path_texts_by_id.as_slice(),
204 limit,
205 pattern_text,
206 MatchType::Directory,
207 ));
208 }
209
210 merge_top_k(heaps, limit)
211 .into_sorted_vec()
212 .into_iter()
213 .map(|Reverse(item)| item)
214 .collect()
215 }
216}
217
218fn intern_path(
219 path: &str,
220 interner: &mut vtcode_commons::StringInterner,
221 path_texts_by_id: &mut Vec<Arc<str>>,
222) -> StringId {
223 let path_id = interner.intern(path);
224 let path_index = path_id.as_u32() as usize;
225 if path_index == path_texts_by_id.len() {
226 path_texts_by_id.push(Arc::from(path));
227 }
228 path_id
229}
230
231fn score_paths_top_k(
239 paths: &[StringId],
240 path_texts_by_id: &[Arc<str>],
241 limit: usize,
242 pattern_text: &str,
243 match_type: MatchType,
244) -> BinaryHeap<Reverse<(u32, StringId, MatchType)>> {
245 const CHUNK: usize = 1024;
246
247 if paths.len() <= CHUNK {
250 let mut list = BestMatchesList::new(limit, pattern_text);
251 for &path_id in paths {
252 if let Some(path) = path_texts_by_id.get(path_id.as_u32() as usize) {
253 list.record_match(path_id, path, match_type);
254 }
255 }
256 return list.matches;
257 }
258
259 let heaps: Vec<_> = paths
260 .par_chunks(CHUNK)
261 .map_init(
262 || BestMatchesList::new(limit, pattern_text),
263 |list, chunk| {
264 for &path_id in chunk {
265 if let Some(path) = path_texts_by_id.get(path_id.as_u32() as usize) {
266 list.record_match(path_id, path, match_type);
267 }
268 }
269 std::mem::take(&mut list.matches)
270 },
271 )
272 .collect();
273
274 merge_top_k(heaps, limit)
275}
276
277fn merge_top_k(
283 heaps: Vec<BinaryHeap<Reverse<(u32, StringId, MatchType)>>>,
284 limit: usize,
285) -> BinaryHeap<Reverse<(u32, StringId, MatchType)>> {
286 let mut merged = BinaryHeap::with_capacity(limit);
287 for heap in heaps {
288 for Reverse(item) in heap.into_vec() {
289 push_top_match(&mut merged, limit, item.0, item.1, item.2);
290 }
291 }
292 merged
293}
294
295pub struct FileIndexCache {
297 cache: Arc<RwLock<Option<Arc<FileIndex>>>>,
298 snapshot: Arc<ArcSwapOption<FileIndex>>,
302 build_gate: Arc<tokio::sync::Semaphore>,
305 search_directory: std::path::PathBuf,
306 exclude: Vec<String>,
307 respect_gitignore: bool,
308 threads: usize,
309}
310
311impl FileIndexCache {
312 pub fn new(
313 search_directory: std::path::PathBuf,
314 exclude: impl IntoIterator<Item = String>,
315 respect_gitignore: bool,
316 threads: usize,
317 ) -> Self {
318 Self {
319 cache: Arc::new(RwLock::new(None)),
320 snapshot: Arc::new(ArcSwapOption::empty()),
321 build_gate: Arc::new(tokio::sync::Semaphore::new(1)),
322 search_directory,
323 exclude: exclude.into_iter().collect(),
324 respect_gitignore,
325 threads,
326 }
327 }
328
329 pub async fn get_or_build(&self) -> anyhow::Result<Arc<FileIndex>> {
331 {
333 let guard = self.cache.read().await;
334 if let Some(index) = guard.as_ref() {
335 if index.last_built.elapsed() < std::time::Duration::from_secs(300) {
337 return Ok(Arc::clone(index));
338 }
339 }
340 }
341
342 let _build_permit = self.build_gate.acquire().await?;
345 {
346 let guard = self.cache.read().await;
347 if let Some(index) = guard.as_ref()
348 && index.last_built.elapsed() < std::time::Duration::from_secs(300)
349 {
350 return Ok(Arc::clone(index));
351 }
352 }
353
354 let search_directory = self.search_directory.clone();
358 let exclude = self.exclude.clone();
359 let respect_gitignore = self.respect_gitignore;
360 let threads = self.threads;
361 let index = Arc::new(
362 tokio::task::spawn_blocking(move || {
363 FileIndex::build_from_directory(&search_directory, &exclude, respect_gitignore, threads)
364 })
365 .await??,
366 );
367
368 {
370 let mut guard = self.cache.write().await;
371 *guard = Some(Arc::clone(&index));
372 self.snapshot.store(Some(Arc::clone(&index)));
373 }
374 Ok(index)
375 }
376
377 pub fn refresh_background(&self) -> Option<Arc<FileIndex>> {
383 let runtime = match tokio::runtime::Handle::try_current() {
384 Ok(runtime) => runtime,
385 Err(error) => {
386 tracing::debug!(%error, "cannot refresh file index without a Tokio runtime");
387 return self.snapshot.load_full();
388 }
389 };
390
391 let search_directory = self.search_directory.clone();
393 let exclude = self.exclude.clone();
394 let respect_gitignore = self.respect_gitignore;
395 let threads = self.threads;
396 let cache = self.cache.clone();
397 let snapshot = Arc::clone(&self.snapshot);
398 let build_gate = Arc::clone(&self.build_gate);
399
400 runtime.spawn(async move {
401 let _build_permit = match build_gate.acquire_owned().await {
402 Ok(permit) => permit,
403 Err(error) => {
404 tracing::error!(%error, "file index build gate closed");
405 return;
406 }
407 };
408
409 match tokio::task::spawn_blocking(move || {
410 FileIndex::build_from_directory(&search_directory, &exclude, respect_gitignore, threads)
411 })
412 .await
413 {
414 Ok(Ok(new_index)) => {
415 let new_index = Arc::new(new_index);
416 let mut guard = cache.write().await;
417 *guard = Some(Arc::clone(&new_index));
418 snapshot.store(Some(new_index));
419 }
420 Ok(Err(error)) => {
421 tracing::error!(%error, "failed to rebuild file index");
422 }
423 Err(error) => {
424 tracing::error!(%error, "file index rebuild task failed");
425 }
426 }
427 });
428
429 self.snapshot.load_full()
430 }
431
432 pub fn update_file(&self, path: &str, is_added: bool) {
435 let mut guard = self.cache.blocking_write();
436 let Some(existing) = guard.take() else { return };
437
438 let mut index = Arc::try_unwrap(existing).unwrap_or_else(|arc| (*arc).clone());
439 if is_added {
440 let path_id = intern_path(path, &mut index.interner, Arc::make_mut(&mut index.path_texts_by_id));
441 let is_directory = self.search_directory.join(path).is_dir();
442 if is_directory {
443 index.files.retain(|&existing| existing != path_id);
444 if !index.directories.contains(&path_id) {
445 index.directories.push(path_id);
446 }
447 } else {
448 index.directories.retain(|&existing| existing != path_id);
449 if !index.files.contains(&path_id) {
450 index.files.push(path_id);
451 }
452 }
453 } else {
454 let Some(path_id) = index.files.iter().chain(index.directories.iter()).copied().find(|&path_id| {
455 index
456 .path_texts_by_id
457 .get(path_id.as_u32() as usize)
458 .is_some_and(|value| value.as_ref() == path)
459 }) else {
460 let index = Arc::new(index);
461 *guard = Some(Arc::clone(&index));
462 self.snapshot.store(Some(index));
463 return;
464 };
465 index.files.retain(|&existing| existing != path_id);
466 index.directories.retain(|&existing| existing != path_id);
467 }
468 index.last_built = std::time::Instant::now();
469 let index = Arc::new(index);
470 *guard = Some(Arc::clone(&index));
471 self.snapshot.store(Some(index));
472 }
473
474 pub async fn index_age(&self) -> Option<std::time::Duration> {
476 let guard = self.cache.read().await;
477 guard.as_ref().map(|idx| idx.last_built.elapsed())
478 }
479}
480
481impl Clone for FileIndex {
483 fn clone(&self) -> Self {
484 Self {
485 files: self.files.clone(),
486 directories: self.directories.clone(),
487 path_texts_by_id: Arc::clone(&self.path_texts_by_id),
488 interner: self.interner.clone(),
489 last_built: self.last_built,
490 }
491 }
492}
493
494#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
502#[serde(rename_all = "lowercase")]
503pub enum MatchType {
504 File,
505 Directory,
506}
507
508#[derive(Debug, Clone, Serialize, Deserialize)]
509pub struct FileMatch {
510 pub score: u32,
511 pub path: String,
512 pub match_type: MatchType,
513 #[serde(skip_serializing_if = "Option::is_none")]
514 pub indices: Option<Vec<u32>>,
515}
516
517#[derive(Debug)]
519pub struct FileSearchResults {
520 pub matches: Vec<FileMatch>,
521 pub total_match_count: usize,
522}
523
524pub struct FileSearchConfig {
526 pub pattern_text: String,
527 pub limit: NonZero<usize>,
528 pub search_directory: std::path::PathBuf,
529 pub exclude: Vec<String>,
530 pub threads: NonZero<usize>,
531 pub cancel_flag: Arc<AtomicBool>,
532 pub compute_indices: bool,
533 pub respect_gitignore: bool,
534}
535
536pub use vtcode_commons::paths::file_name_from_path;
537
538struct BestMatchesList {
543 matches: BinaryHeap<Reverse<(u32, StringId, MatchType)>>,
544 limit: usize,
545 matcher: nucleo_matcher::Matcher,
546 haystack_buf: Vec<char>,
547 pattern: PatternStorage,
549 owned_matches: BinaryHeap<Reverse<OwnedMatchCandidate>>,
550}
551
552#[derive(Debug, Eq, PartialEq)]
553struct OwnedMatchCandidate {
554 score: u32,
555 path: String,
556 match_type: MatchType,
557}
558
559impl Ord for OwnedMatchCandidate {
560 fn cmp(&self, other: &Self) -> std::cmp::Ordering {
561 self.score
562 .cmp(&other.score)
563 .then_with(|| other.path.cmp(&self.path))
564 .then_with(|| other.match_type.cmp(&self.match_type))
565 }
566}
567
568impl PartialOrd for OwnedMatchCandidate {
569 fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
570 Some(self.cmp(other))
571 }
572}
573
574struct SearchAggregation {
575 matches: BinaryHeap<Reverse<OwnedMatchCandidate>>,
576 total_match_count: usize,
577}
578
579struct IndexWorkerResult {
580 files: Vec<String>,
581 directories: Vec<String>,
582}
583
584struct IndexWorkerGuard {
585 results: Arc<Mutex<Vec<IndexWorkerResult>>>,
586 files: Vec<String>,
587 directories: Vec<String>,
588}
589
590impl Drop for IndexWorkerGuard {
591 fn drop(&mut self) {
592 self.results.lock().push(IndexWorkerResult {
593 files: std::mem::take(&mut self.files),
594 directories: std::mem::take(&mut self.directories),
595 });
596 }
597}
598
599struct SearchWorkerGuard {
600 results: Arc<Mutex<SearchAggregation>>,
601 best_list: BestMatchesList,
602 total_match_count: usize,
603 limit: usize,
604}
605
606impl Drop for SearchWorkerGuard {
607 fn drop(&mut self) {
608 let worker_matches = std::mem::take(&mut self.best_list.owned_matches);
609 let mut results = self.results.lock();
610 results.total_match_count += self.total_match_count;
611 for Reverse(candidate) in worker_matches {
612 push_owned_match(&mut results.matches, self.limit, candidate);
613 }
614 }
615}
616
617enum PatternStorage {
619 Ascii(Vec<u8>),
621 Unicode(Vec<char>),
623}
624
625impl BestMatchesList {
626 fn new(limit: usize, pattern_text: &str) -> Self {
627 let pattern = if pattern_text.is_ascii() {
631 PatternStorage::Ascii(pattern_text.to_ascii_lowercase().into_bytes())
632 } else {
633 PatternStorage::Unicode(pattern_text.to_lowercase().chars().collect())
634 };
635
636 Self {
637 matches: BinaryHeap::new(),
638 limit,
639 matcher: nucleo_matcher::Matcher::new(nucleo_matcher::Config::DEFAULT),
640 haystack_buf: Vec::with_capacity(256),
641 pattern,
642 owned_matches: BinaryHeap::new(),
643 }
644 }
645
646 fn score_path(&mut self, path: &str) -> Option<u32> {
649 let haystack = nucleo_matcher::Utf32Str::new(path, &mut self.haystack_buf);
650 let needle = match &self.pattern {
651 PatternStorage::Ascii(bytes) => nucleo_matcher::Utf32Str::Ascii(bytes),
652 PatternStorage::Unicode(chars) => nucleo_matcher::Utf32Str::Unicode(chars),
653 };
654 self.matcher.fuzzy_match(haystack, needle).map(|score| score as u32)
655 }
656
657 fn record_match(&mut self, path_id: StringId, path: &str, match_type: MatchType) -> bool {
659 let Some(score) = self.score_path(path) else {
660 return false;
661 };
662 push_top_match(&mut self.matches, self.limit, score, path_id, match_type);
663 true
664 }
665
666 fn record_scored_match(&mut self, path_id: StringId, score: u32, match_type: MatchType) {
667 push_top_match(&mut self.matches, self.limit, score, path_id, match_type);
668 }
669
670 fn record_owned_match(&mut self, path: &str, score: u32, match_type: MatchType) {
671 if self.owned_matches.len() == self.limit {
672 let Some(Reverse(worst)) = self.owned_matches.peek() else {
673 return;
674 };
675 if !owned_match_is_better(score, path, match_type, worst) {
676 return;
677 }
678 }
679
680 push_owned_match(
681 &mut self.owned_matches,
682 self.limit,
683 OwnedMatchCandidate { score, path: path.to_owned(), match_type },
684 );
685 }
686}
687
688fn push_owned_match(
689 matches: &mut BinaryHeap<Reverse<OwnedMatchCandidate>>,
690 limit: usize,
691 candidate: OwnedMatchCandidate,
692) {
693 if matches.len() == limit {
694 let Some(Reverse(worst)) = matches.peek() else {
695 return;
696 };
697 if !owned_match_is_better(candidate.score, candidate.path.as_str(), candidate.match_type, worst) {
698 return;
699 }
700 matches.pop();
701 }
702
703 matches.push(Reverse(candidate));
704}
705
706fn owned_match_is_better(score: u32, path: &str, match_type: MatchType, worst: &OwnedMatchCandidate) -> bool {
707 score
708 .cmp(&worst.score)
709 .then_with(|| worst.path.as_str().cmp(path))
710 .then_with(|| worst.match_type.cmp(&match_type))
711 .is_gt()
712}
713
714fn push_top_match(
715 matches: &mut BinaryHeap<Reverse<(u32, StringId, MatchType)>>,
716 limit: usize,
717 score: u32,
718 path: StringId,
719 match_type: MatchType,
720) -> bool {
721 let candidate = (score, path, match_type);
722 if matches.len() < limit {
723 matches.push(Reverse(candidate));
724 return true;
725 }
726
727 let Some(minimum) = matches.peek().map(|entry| &entry.0) else {
728 return false;
729 };
730
731 if &candidate <= minimum {
732 return false;
733 }
734
735 matches.pop();
736 matches.push(Reverse(candidate));
737 true
738}
739
740pub async fn run_with_index(
754 config: FileSearchConfig,
755 index_cache: &FileIndexCache,
756) -> anyhow::Result<FileSearchResults> {
757 let limit = config.limit.get();
758 let cancel_flag = &config.cancel_flag;
759 let compute_indices = config.compute_indices;
760
761 let index = index_cache.get_or_build().await?;
763
764 if cancel_flag.load(Ordering::Relaxed) {
766 return Ok(FileSearchResults { matches: Vec::new(), total_match_count: 0 });
767 }
768
769 let index_for_results = index.clone();
772 let matched_paths = tokio::task::spawn_blocking({
773 let pattern_text = config.pattern_text.clone();
774 move || Ok::<_, anyhow::Error>(index.query(&pattern_text, limit, None))
775 })
776 .await??;
777
778 let total_match_count = matched_paths.len();
779
780 let matches = matched_paths
782 .into_iter()
783 .filter_map(|(score, path_id, match_type)| {
784 let path = index_for_results.path_texts_by_id.get(path_id.as_u32() as usize)?.to_string();
785 Some(FileMatch {
786 score,
787 path,
788 match_type,
789 indices: if compute_indices { Some(Vec::new()) } else { None },
790 })
791 })
792 .collect();
793
794 Ok(FileSearchResults { matches, total_match_count })
795}
796
797pub fn run(config: FileSearchConfig) -> anyhow::Result<FileSearchResults> {
807 run_with_policy(config, true, false)
808}
809
810pub fn run_bounded_no_follow(config: FileSearchConfig) -> anyhow::Result<FileSearchResults> {
816 run_bounded_no_follow_with_visit(config, |_| {})
817}
818
819fn run_bounded_no_follow_with_visit(
820 config: FileSearchConfig,
821 mut visit: impl FnMut(&Path),
822) -> anyhow::Result<FileSearchResults> {
823 let limit = config.limit.get();
824 let search_directory = &config.search_directory;
825 let mut walk_builder = ignore::WalkBuilder::new(search_directory);
826 vtcode_commons::walk::apply_defaults(&mut walk_builder);
827 walk_builder
828 .follow_links(false)
829 .require_git(false)
830 .sort_by_file_path(|left, right| left.cmp(right));
831
832 if !config.respect_gitignore {
833 walk_builder
834 .git_ignore(false)
835 .git_global(false)
836 .git_exclude(false)
837 .ignore(false)
838 .parents(false);
839 }
840
841 if !config.exclude.is_empty() {
842 let mut override_builder = ignore::overrides::OverrideBuilder::new(search_directory);
843 for exclude_pattern in &config.exclude {
844 override_builder.add(&format!("!{exclude_pattern}"))?;
845 }
846 walk_builder.overrides(override_builder.build()?);
847 }
848
849 let interner = Arc::new(Mutex::new(vtcode_commons::StringInterner::new()));
850 let mut matches = BestMatchesList::new(limit, &config.pattern_text);
851 let mut matching_count = 0usize;
852 for result in walk_builder.build() {
853 if config.cancel_flag.load(Ordering::Relaxed) {
854 break;
855 }
856 let entry = match result {
857 Ok(entry) => entry,
858 Err(_) => continue,
859 };
860 visit(entry.path());
861 if !entry.file_type().is_some_and(|file_type| file_type.is_file()) {
862 continue;
863 }
864 let Some(relative_path) = entry
865 .path()
866 .strip_prefix(search_directory)
867 .ok()
868 .and_then(|path| path.to_str())
869 .filter(|path| !path.is_empty())
870 else {
871 continue;
872 };
873 let Some(score) = matches.score_path(relative_path) else {
874 continue;
875 };
876 let path_id = interner.lock().intern(relative_path);
877 matches.record_scored_match(path_id, score, MatchType::File);
878 matching_count += 1;
879 if matching_count >= limit {
880 break;
881 }
882 }
883
884 let interner_guard = interner.lock();
885 let matches = matches
886 .matches
887 .into_sorted_vec()
888 .into_iter()
889 .filter_map(|Reverse((score, path_id, match_type))| {
890 let path = interner_guard.get(path_id)?.to_string();
891 Some(FileMatch {
892 score,
893 path,
894 match_type,
895 indices: config.compute_indices.then(Vec::new),
896 })
897 })
898 .collect();
899
900 Ok(FileSearchResults {
901 matches,
902 total_match_count: matching_count + usize::from(matching_count >= limit),
905 })
906}
907
908fn run_with_policy(
909 config: FileSearchConfig,
910 follow_links: bool,
911 files_only: bool,
912) -> anyhow::Result<FileSearchResults> {
913 let limit = config.limit.get();
914 let search_directory = &config.search_directory;
915 let exclude = &config.exclude;
916 let threads = config.threads.get();
917 let cancel_flag = &config.cancel_flag;
918 let compute_indices = config.compute_indices;
919 let respect_gitignore = config.respect_gitignore;
920
921 let walker = build_parallel_walker(search_directory, exclude, threads, respect_gitignore, follow_links)?;
922
923 let worker_results = Arc::new(Mutex::new(SearchAggregation { matches: BinaryHeap::new(), total_match_count: 0 }));
924 walker.run(|| {
925 let mut state = SearchWorkerGuard {
926 results: Arc::clone(&worker_results),
927 best_list: BestMatchesList::new(limit, &config.pattern_text),
928 total_match_count: 0,
929 limit,
930 };
931
932 Box::new(move |result| {
933 if cancel_flag.load(Ordering::Relaxed) {
935 return ignore::WalkState::Quit;
936 }
937
938 let entry = match result {
939 Ok(e) => e,
940 Err(_) => return ignore::WalkState::Continue,
941 };
942
943 let relative_path = entry.path().strip_prefix(search_directory).ok().and_then(|p| p.to_str());
945
946 let path_to_match = match relative_path {
947 Some(p) if !p.is_empty() => p,
948 _ => return ignore::WalkState::Continue, };
950
951 let Some(file_type) = entry.file_type() else {
952 return ignore::WalkState::Continue;
953 };
954 let match_type = if file_type.is_dir() {
955 MatchType::Directory
956 } else {
957 MatchType::File
958 };
959
960 if files_only && match_type == MatchType::Directory {
961 return ignore::WalkState::Continue;
962 }
963
964 let Some(score) = state.best_list.score_path(path_to_match) else {
966 return ignore::WalkState::Continue;
967 };
968 state.best_list.record_owned_match(path_to_match, score, match_type);
969 state.total_match_count += 1;
970
971 ignore::WalkState::Continue
972 })
973 });
974
975 let results = Arc::try_unwrap(worker_results)
976 .map_err(|results| {
977 anyhow::anyhow!("file search worker results still referenced: {}", Arc::strong_count(&results))
978 })?
979 .into_inner();
980 let total_match_count = results.total_match_count;
981 let mut candidates = results
982 .matches
983 .into_iter()
984 .map(|Reverse(candidate)| candidate)
985 .collect::<Vec<_>>();
986 candidates.sort_by(|left, right| {
987 right
988 .score
989 .cmp(&left.score)
990 .then_with(|| left.path.cmp(&right.path))
991 .then_with(|| left.match_type.cmp(&right.match_type))
992 });
993 candidates.truncate(limit);
994
995 let matches = candidates
997 .into_iter()
998 .map(|candidate| FileMatch {
999 score: candidate.score,
1000 path: candidate.path,
1001 match_type: candidate.match_type,
1002 indices: if compute_indices { Some(Vec::new()) } else { None },
1003 })
1004 .collect();
1005
1006 Ok(FileSearchResults { matches, total_match_count })
1007}
1008
1009#[cfg(test)]
1010mod tests {
1011 use super::{
1012 FileIndexCache, FileSearchConfig, MatchType, run, run_bounded_no_follow, run_bounded_no_follow_with_visit,
1013 run_with_index,
1014 };
1015 use std::num::NonZero;
1016 use std::sync::Arc;
1017 use std::sync::atomic::AtomicBool;
1018 use tempfile::TempDir;
1019
1020 #[tokio::test(flavor = "current_thread")]
1021 async fn concurrent_index_builds_share_async_cache_entry() {
1022 let workspace = TempDir::new().expect("workspace");
1023 std::fs::write(workspace.path().join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1024
1025 let cache = FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), true, 1);
1026 let (first, second) = tokio::join!(cache.get_or_build(), cache.get_or_build());
1027 let first = first.expect("build file index");
1028 let second = second.expect("reuse file index");
1029
1030 assert!(Arc::ptr_eq(&first, &second));
1031 }
1032
1033 #[tokio::test(flavor = "current_thread")]
1034 async fn background_refresh_is_safe_when_called_from_tokio() {
1035 let workspace = TempDir::new().expect("workspace");
1036 std::fs::write(workspace.path().join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1037
1038 let cache = FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), false, 1);
1039 cache.get_or_build().await.expect("initial index");
1040
1041 assert!(cache.refresh_background().is_some());
1042 }
1043
1044 #[test]
1045 fn background_refresh_returns_snapshot_without_runtime() {
1046 let workspace = TempDir::new().expect("workspace");
1047 std::fs::write(workspace.path().join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1048
1049 let runtime = tokio::runtime::Runtime::new().expect("Tokio runtime");
1050 let cache = FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), false, 1);
1051 runtime.block_on(cache.get_or_build()).expect("initial index");
1052 drop(runtime);
1053
1054 assert!(cache.refresh_background().is_some());
1055 }
1056
1057 #[tokio::test(flavor = "current_thread")]
1058 async fn incremental_directory_updates_use_the_cache_root() {
1059 let workspace = TempDir::new().expect("workspace");
1060 std::fs::write(workspace.path().join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1061 std::fs::create_dir(workspace.path().join("new_directory")).expect("fixture directory");
1062
1063 let cache = Arc::new(FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), false, 1));
1064 cache.get_or_build().await.expect("initial index");
1065
1066 tokio::task::spawn_blocking({
1067 let cache = Arc::clone(&cache);
1068 move || cache.update_file("new_directory", true)
1069 })
1070 .await
1071 .expect("incremental directory update task");
1072
1073 let index = cache.get_or_build().await.expect("updated index");
1074 let matches = index.query("new_directory", 16, None);
1075 assert!(matches.iter().any(|(_, _, match_type)| *match_type == MatchType::Directory));
1076 }
1077
1078 fn indexed_search_config(
1079 workspace: &std::path::Path,
1080 pattern: &str,
1081 cancel_flag: Arc<AtomicBool>,
1082 ) -> FileSearchConfig {
1083 FileSearchConfig {
1084 pattern_text: pattern.to_string(),
1085 limit: NonZero::new(16).expect("non-zero limit"),
1086 search_directory: workspace.to_path_buf(),
1087 exclude: Vec::new(),
1088 threads: NonZero::new(1).expect("non-zero threads"),
1089 cancel_flag,
1090 compute_indices: false,
1091 respect_gitignore: false,
1092 }
1093 }
1094
1095 fn result_signature(results: &super::FileSearchResults) -> Vec<(u32, String, MatchType)> {
1096 results
1097 .matches
1098 .iter()
1099 .map(|candidate| (candidate.score, candidate.path.clone(), candidate.match_type))
1100 .collect()
1101 }
1102
1103 #[tokio::test(flavor = "current_thread")]
1104 async fn indexed_search_preserves_scores_order_and_match_types() {
1105 let workspace = TempDir::new().expect("workspace");
1106 std::fs::create_dir_all(workspace.path().join("src/widget_dir")).expect("fixture directory");
1107 std::fs::write(workspace.path().join("src/widget.rs"), "fn widget() {}\n").expect("fixture source");
1108 std::fs::write(workspace.path().join("src/widget_test.rs"), "fn widget_test() {}\n").expect("fixture source");
1109 std::fs::write(workspace.path().join("README.md"), "widget documentation\n").expect("fixture docs");
1110
1111 let cache = FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), false, 1);
1112 let first =
1113 run_with_index(indexed_search_config(workspace.path(), "widget", Arc::new(AtomicBool::new(false))), &cache)
1114 .await
1115 .expect("indexed search");
1116 let second =
1117 run_with_index(indexed_search_config(workspace.path(), "widget", Arc::new(AtomicBool::new(false))), &cache)
1118 .await
1119 .expect("repeat indexed search");
1120
1121 assert_eq!(result_signature(&first), result_signature(&second));
1122 assert!(
1123 first
1124 .matches
1125 .iter()
1126 .any(|candidate| candidate.match_type == MatchType::Directory)
1127 );
1128 assert!(first.matches.iter().any(|candidate| candidate.match_type == MatchType::File));
1129 assert!(first.matches.windows(2).all(|window| window[0].score >= window[1].score));
1130 }
1131
1132 #[tokio::test(flavor = "current_thread")]
1133 async fn indexed_search_honors_cancellation_before_scoring() {
1134 let workspace = TempDir::new().expect("workspace");
1135 std::fs::write(workspace.path().join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1136 let cache = FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), false, 1);
1137 let cancel_flag = Arc::new(AtomicBool::new(true));
1138
1139 let results = run_with_index(indexed_search_config(workspace.path(), "widget", cancel_flag), &cache)
1140 .await
1141 .expect("cancelled indexed search");
1142 assert!(results.matches.is_empty());
1143 assert_eq!(results.total_match_count, 0);
1144 }
1145
1146 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1147 async fn incremental_updates_do_not_mutate_an_old_search_index() {
1148 let workspace = TempDir::new().expect("workspace");
1149 std::fs::write(workspace.path().join("old_widget.rs"), "fn old_widget() {}\n").expect("fixture source");
1150 let cache = Arc::new(FileIndexCache::new(workspace.path().to_path_buf(), Vec::new(), false, 1));
1151 let old_index = cache.get_or_build().await.expect("initial index");
1152
1153 tokio::task::spawn_blocking({
1154 let cache = Arc::clone(&cache);
1155 move || cache.update_file("new_widget.rs", true)
1156 })
1157 .await
1158 .expect("incremental update task");
1159
1160 let old_matches = old_index.query("widget", 16, None);
1161 let new_index = cache.get_or_build().await.expect("updated index");
1162 let new_matches = new_index.query("widget", 16, None);
1163 let old_paths = old_matches
1164 .iter()
1165 .filter_map(|(_, path_id, _)| old_index.path_texts_by_id.get(path_id.as_u32() as usize))
1166 .map(AsRef::as_ref)
1167 .collect::<Vec<&str>>();
1168 let new_paths = new_matches
1169 .iter()
1170 .filter_map(|(_, path_id, _)| new_index.path_texts_by_id.get(path_id.as_u32() as usize))
1171 .map(AsRef::as_ref)
1172 .collect::<Vec<&str>>();
1173
1174 assert!(!old_paths.contains(&"new_widget.rs"));
1175 assert!(new_paths.contains(&"new_widget.rs"));
1176 }
1177
1178 fn uncached_search_config(
1179 workspace: &std::path::Path,
1180 pattern: &str,
1181 limit: usize,
1182 cancel_flag: Arc<AtomicBool>,
1183 ) -> FileSearchConfig {
1184 FileSearchConfig {
1185 pattern_text: pattern.to_string(),
1186 limit: NonZero::new(limit).expect("non-zero limit"),
1187 search_directory: workspace.to_path_buf(),
1188 exclude: Vec::new(),
1189 threads: NonZero::new(4).expect("non-zero threads"),
1190 cancel_flag,
1191 compute_indices: false,
1192 respect_gitignore: false,
1193 }
1194 }
1195
1196 #[test]
1197 fn uncached_parallel_search_aggregates_worker_paths_and_counts() {
1198 let workspace = TempDir::new().expect("workspace");
1199 for index in 0..64 {
1200 std::fs::create_dir(workspace.path().join(format!("widget_dir_{index:03}"))).expect("fixture directory");
1201 std::fs::write(workspace.path().join(format!("widget_file_{index:03}.rs")), "fn widget() {}\n")
1202 .expect("fixture source");
1203 }
1204
1205 let results = run(uncached_search_config(workspace.path(), "widget", 32, Arc::new(AtomicBool::new(false))))
1206 .expect("uncached parallel search");
1207
1208 assert_eq!(results.total_match_count, 128);
1209 assert_eq!(results.matches.len(), 32);
1210 let unique_paths = results
1211 .matches
1212 .iter()
1213 .map(|candidate| candidate.path.as_str())
1214 .collect::<std::collections::HashSet<_>>();
1215 assert_eq!(unique_paths.len(), results.matches.len());
1216 assert!(results.matches.iter().all(|candidate| candidate.path.contains("widget")));
1217 }
1218
1219 #[test]
1220 fn uncached_parallel_search_honors_cancellation() {
1221 let workspace = TempDir::new().expect("workspace");
1222 for index in 0..32 {
1223 std::fs::write(workspace.path().join(format!("widget_{index:03}.rs")), "fn widget() {}\n")
1224 .expect("fixture source");
1225 }
1226 let cancel_flag = Arc::new(AtomicBool::new(true));
1227
1228 let results = run(uncached_search_config(workspace.path(), "widget", 16, cancel_flag))
1229 .expect("cancelled uncached search");
1230
1231 assert!(results.matches.is_empty());
1232 assert_eq!(results.total_match_count, 0);
1233 }
1234
1235 #[test]
1236 fn uncached_equal_score_results_are_lexically_ordered_and_repeatable() {
1237 let workspace = TempDir::new().expect("workspace");
1238 for name in ["a_x.rs", "b_x.rs", "c_x.rs", "d_x.rs", "e_x.rs"] {
1239 std::fs::write(workspace.path().join(name), "fn x() {}\n").expect("fixture source");
1240 }
1241
1242 let search = || {
1243 run(uncached_search_config(workspace.path(), "x", 3, Arc::new(AtomicBool::new(false))))
1244 .expect("uncached equal-score search")
1245 .matches
1246 .into_iter()
1247 .map(|candidate| (candidate.score, candidate.path))
1248 .collect::<Vec<_>>()
1249 };
1250
1251 let expected = search();
1252 for _ in 0..20 {
1253 assert_eq!(search(), expected);
1254 }
1255 assert!(
1256 expected.windows(2).all(|window| {
1257 window[0].0 > window[1].0 || (window[0].0 == window[1].0 && window[0].1 <= window[1].1)
1258 })
1259 );
1260 }
1261
1262 fn bounded_paths(workspace: &std::path::Path) -> Vec<String> {
1263 run_bounded_no_follow(FileSearchConfig {
1264 pattern_text: "widget".to_string(),
1265 limit: NonZero::new(2).expect("non-zero limit"),
1266 search_directory: workspace.to_path_buf(),
1267 exclude: Vec::new(),
1268 threads: NonZero::new(4).expect("non-zero threads"),
1269 cancel_flag: Arc::new(AtomicBool::new(false)),
1270 compute_indices: false,
1271 respect_gitignore: true,
1272 })
1273 .expect("bounded path search")
1274 .matches
1275 .into_iter()
1276 .map(|candidate| candidate.path)
1277 .collect()
1278 }
1279
1280 #[test]
1281 fn bounded_path_selection_is_stable_across_repeated_walks() {
1282 let workspace = TempDir::new().expect("workspace");
1283 for directory in ["z", "a", "m", "b", "y"] {
1284 let directory = workspace.path().join(directory);
1285 std::fs::create_dir(&directory).expect("fixture directory");
1286 std::fs::write(directory.join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1287 }
1288
1289 let expected = bounded_paths(workspace.path());
1290 assert_eq!(expected.len(), 2);
1291 for _ in 0..20 {
1292 assert_eq!(bounded_paths(workspace.path()), expected);
1293 }
1294 }
1295
1296 #[test]
1297 fn bounded_path_selection_is_the_sorted_prefix_and_stops_early() {
1298 let workspace = TempDir::new().expect("workspace");
1299 for directory in ["z", "a", "m", "b", "y"] {
1300 let directory = workspace.path().join(directory);
1301 std::fs::create_dir(&directory).expect("fixture directory");
1302 std::fs::write(directory.join("widget.rs"), "fn widget() {}\n").expect("fixture source");
1303 }
1304 let mut visited = Vec::new();
1305
1306 let results = run_bounded_no_follow_with_visit(
1307 FileSearchConfig {
1308 pattern_text: "widget".to_string(),
1309 limit: NonZero::new(2).expect("non-zero limit"),
1310 search_directory: workspace.path().to_path_buf(),
1311 exclude: Vec::new(),
1312 threads: NonZero::new(4).expect("non-zero threads"),
1313 cancel_flag: Arc::new(AtomicBool::new(false)),
1314 compute_indices: false,
1315 respect_gitignore: true,
1316 },
1317 |path| visited.push(path.to_path_buf()),
1318 )
1319 .expect("bounded path search");
1320 let mut paths = results.matches.into_iter().map(|candidate| candidate.path).collect::<Vec<_>>();
1321 paths.sort();
1322
1323 assert_eq!(paths, vec!["a/widget.rs", "b/widget.rs"]);
1324 assert!(visited.len() < 11, "the bounded route must stop before traversing the complete fixture tree");
1325 assert_eq!(results.total_match_count, 3);
1326 }
1327}