1use parking_lot::Mutex;
33use serde::{Deserialize, Serialize};
34use std::cmp::Reverse;
35use std::collections::BinaryHeap;
36use std::num::NonZero;
37use std::path::Path;
38use std::sync::Arc;
39use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
40use tokio::sync::RwLock;
41
42use rayon::prelude::*;
43
44pub struct FileIndex {
49 files: Vec<String>,
51 directories: Vec<String>,
53 last_built: std::time::Instant,
55}
56
57fn build_parallel_walker(
59 search_directory: &Path,
60 exclude: &[String],
61 threads: usize,
62 respect_gitignore: bool,
63 follow_links: bool,
64) -> anyhow::Result<ignore::WalkParallel> {
65 let mut walk_builder = ignore::WalkBuilder::new(search_directory);
66 vtcode_commons::walk::apply_defaults(&mut walk_builder);
67
68 walk_builder.threads(threads);
70 walk_builder.follow_links(follow_links);
71 walk_builder.require_git(false); if !respect_gitignore {
74 walk_builder
75 .git_ignore(false)
76 .git_global(false)
77 .git_exclude(false)
78 .ignore(false)
79 .parents(false);
80 }
81
82 if !exclude.is_empty() {
83 let mut override_builder = ignore::overrides::OverrideBuilder::new(search_directory);
84 for exclude_pattern in exclude {
85 let pattern = format!("!{exclude_pattern}");
86 override_builder.add(&pattern)?;
87 }
88 walk_builder.overrides(override_builder.build()?);
89 }
90
91 Ok(walk_builder.build_parallel())
92}
93
94impl FileIndex {
95 fn build_from_directory(
98 search_directory: &Path,
99 exclude: &[String],
100 respect_gitignore: bool,
101 threads: usize,
102 ) -> anyhow::Result<Self> {
103 let walker =
104 build_parallel_walker(search_directory, exclude, threads, respect_gitignore, true)?;
105
106 let files_arc = Arc::new(Mutex::new(Vec::new()));
108 let dirs_arc = Arc::new(Mutex::new(Vec::new()));
109
110 walker.run(|| {
111 let files_clone = files_arc.clone();
112 let dirs_clone = dirs_arc.clone();
113 let search_dir = search_directory.to_path_buf();
114
115 Box::new(move |result| {
116 let entry = match result {
117 Ok(e) => e,
118 Err(_) => return ignore::WalkState::Continue,
119 };
120
121 if let Some(rel_path) = entry
123 .path()
124 .strip_prefix(&search_dir)
125 .ok()
126 .and_then(|p| p.to_str())
127 && !rel_path.is_empty()
128 {
129 if entry.path().is_dir() {
130 dirs_clone.lock().push(rel_path.to_string());
131 } else {
132 files_clone.lock().push(rel_path.to_string());
133 }
134 }
135
136 ignore::WalkState::Continue
137 })
138 });
139
140 let files = Arc::try_unwrap(files_arc)
141 .map_err(|arc| {
142 anyhow::anyhow!(
143 "failed to unwrap files arc, {} references remain",
144 Arc::strong_count(&arc)
145 )
146 })?
147 .into_inner();
148 let directories = Arc::try_unwrap(dirs_arc)
149 .map_err(|arc| {
150 anyhow::anyhow!(
151 "failed to unwrap dirs arc, {} references remain",
152 Arc::strong_count(&arc)
153 )
154 })?
155 .into_inner();
156
157 Ok(Self {
158 files,
159 directories,
160 last_built: std::time::Instant::now(),
161 })
162 }
163
164 fn query(
167 &self,
168 pattern_text: &str,
169 limit: usize,
170 match_type_filter: Option<MatchType>,
171 ) -> Vec<(u32, String, MatchType)> {
172 let mut heaps = Vec::new();
177
178 if match_type_filter.is_none_or(|t| t == MatchType::File) {
179 heaps.push(score_paths_top_k(
180 &self.files,
181 limit,
182 pattern_text,
183 MatchType::File,
184 ));
185 }
186
187 if match_type_filter.is_none_or(|t| t == MatchType::Directory) {
188 heaps.push(score_paths_top_k(
189 &self.directories,
190 limit,
191 pattern_text,
192 MatchType::Directory,
193 ));
194 }
195
196 merge_top_k(heaps, limit)
197 .into_sorted_vec()
198 .into_iter()
199 .map(|Reverse(item)| item)
200 .collect()
201 }
202}
203
204fn score_paths_top_k(
212 paths: &[String],
213 limit: usize,
214 pattern_text: &str,
215 match_type: MatchType,
216) -> BinaryHeap<Reverse<(u32, String, MatchType)>> {
217 const CHUNK: usize = 1024;
218
219 if paths.len() <= CHUNK {
222 let mut list = BestMatchesList::new(limit, pattern_text);
223 for path in paths {
224 list.record_match(path, match_type);
225 }
226 return list.matches;
227 }
228
229 let heaps: Vec<_> = paths
230 .par_chunks(CHUNK)
231 .map_init(
232 || BestMatchesList::new(limit, pattern_text),
233 |list, chunk| {
234 for path in chunk {
235 list.record_match(path, match_type);
236 }
237 std::mem::take(&mut list.matches)
238 },
239 )
240 .collect();
241
242 merge_top_k(heaps, limit)
243}
244
245fn merge_top_k(
251 heaps: Vec<BinaryHeap<Reverse<(u32, String, MatchType)>>>,
252 limit: usize,
253) -> BinaryHeap<Reverse<(u32, String, MatchType)>> {
254 let mut merged = BinaryHeap::with_capacity(limit);
255 for heap in heaps {
256 for Reverse(item) in heap.into_vec() {
257 push_top_match(&mut merged, limit, item.0, item.1, item.2);
258 }
259 }
260 merged
261}
262
263pub struct FileIndexCache {
265 cache: Arc<RwLock<Option<Arc<FileIndex>>>>,
266 search_directory: std::path::PathBuf,
267 exclude: Vec<String>,
268 respect_gitignore: bool,
269 threads: usize,
270}
271
272impl FileIndexCache {
273 pub fn new(
274 search_directory: std::path::PathBuf,
275 exclude: impl IntoIterator<Item = String>,
276 respect_gitignore: bool,
277 threads: usize,
278 ) -> Self {
279 Self {
280 cache: Arc::new(RwLock::new(None)),
281 search_directory,
282 exclude: exclude.into_iter().collect(),
283 respect_gitignore,
284 threads,
285 }
286 }
287
288 pub async fn get_or_build(&self) -> anyhow::Result<Arc<FileIndex>> {
290 {
292 let guard = self.cache.read().await;
293 if let Some(index) = guard.as_ref() {
294 if index.last_built.elapsed() < std::time::Duration::from_secs(300) {
296 return Ok(Arc::clone(index));
297 }
298 }
299 }
300
301 let index = Arc::new(FileIndex::build_from_directory(
303 &self.search_directory,
304 &self.exclude,
305 self.respect_gitignore,
306 self.threads,
307 )?);
308
309 {
311 let mut guard = self.cache.write().await;
312 *guard = Some(Arc::clone(&index));
313 }
314 Ok(index)
315 }
316
317 pub fn refresh_background(&self) -> Option<Arc<FileIndex>> {
320 let search_directory = self.search_directory.clone();
322 let exclude = self.exclude.clone();
323 let respect_gitignore = self.respect_gitignore;
324 let threads = self.threads;
325 let cache = self.cache.clone();
326
327 tokio::spawn(async move {
328 match FileIndex::build_from_directory(
329 &search_directory,
330 &exclude,
331 respect_gitignore,
332 threads,
333 ) {
334 Ok(new_index) => {
335 let mut guard = cache.write().await;
336 *guard = Some(Arc::new(new_index));
337 }
338 Err(e) => {
339 tracing::error!("failed to rebuild file index: {e}");
340 }
341 }
342 });
343
344 let guard = self.cache.blocking_read();
346 guard.as_ref().map(Arc::clone)
347 }
348
349 pub fn update_file(&self, path: &str, is_added: bool) {
352 let mut guard = self.cache.blocking_write();
353 let Some(existing) = guard.take() else { return };
354
355 let mut index = Arc::try_unwrap(existing).unwrap_or_else(|arc| (*arc).clone());
356 if is_added {
357 if Path::new(path).is_dir() {
358 index.directories.push(path.to_string());
359 } else {
360 index.files.push(path.to_string());
361 }
362 } else {
363 index.files.retain(|p| p != path);
364 index.directories.retain(|p| p != path);
365 }
366 index.last_built = std::time::Instant::now();
367 *guard = Some(Arc::new(index));
368 }
369
370 pub async fn index_age(&self) -> Option<std::time::Duration> {
372 let guard = self.cache.read().await;
373 guard.as_ref().map(|idx| idx.last_built.elapsed())
374 }
375}
376
377impl Clone for FileIndex {
379 fn clone(&self) -> Self {
380 Self {
381 files: self.files.clone(),
382 directories: self.directories.clone(),
383 last_built: self.last_built,
384 }
385 }
386}
387
388#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)]
396#[serde(rename_all = "lowercase")]
397pub enum MatchType {
398 File,
399 Directory,
400}
401
402#[derive(Debug, Clone, Serialize, Deserialize)]
403pub struct FileMatch {
404 pub score: u32,
405 pub path: String,
406 pub match_type: MatchType,
407 #[serde(skip_serializing_if = "Option::is_none")]
408 pub indices: Option<Vec<u32>>,
409}
410
411#[derive(Debug)]
413pub struct FileSearchResults {
414 pub matches: Vec<FileMatch>,
415 pub total_match_count: usize,
416}
417
418pub struct FileSearchConfig {
420 pub pattern_text: String,
421 pub limit: NonZero<usize>,
422 pub search_directory: std::path::PathBuf,
423 pub exclude: Vec<String>,
424 pub threads: NonZero<usize>,
425 pub cancel_flag: Arc<AtomicBool>,
426 pub compute_indices: bool,
427 pub respect_gitignore: bool,
428}
429
430pub use vtcode_commons::paths::file_name_from_path;
431
432struct BestMatchesList {
437 matches: BinaryHeap<Reverse<(u32, String, MatchType)>>,
438 limit: usize,
439 matcher: nucleo_matcher::Matcher,
440 haystack_buf: Vec<char>,
441 pattern: PatternStorage,
443}
444
445enum PatternStorage {
447 Ascii(Vec<u8>),
449 Unicode(Vec<char>),
451}
452
453impl BestMatchesList {
454 fn new(limit: usize, pattern_text: &str) -> Self {
455 let pattern = if pattern_text.is_ascii() {
459 PatternStorage::Ascii(pattern_text.to_ascii_lowercase().into_bytes())
460 } else {
461 PatternStorage::Unicode(pattern_text.to_lowercase().chars().collect())
462 };
463
464 Self {
465 matches: BinaryHeap::new(),
466 limit,
467 matcher: nucleo_matcher::Matcher::new(nucleo_matcher::Config::DEFAULT),
468 haystack_buf: Vec::with_capacity(256),
469 pattern,
470 }
471 }
472
473 fn record_match(&mut self, path: &str, match_type: MatchType) -> bool {
478 let haystack = nucleo_matcher::Utf32Str::new(path, &mut self.haystack_buf);
480 let needle = match &self.pattern {
481 PatternStorage::Ascii(bytes) => nucleo_matcher::Utf32Str::Ascii(bytes),
482 PatternStorage::Unicode(chars) => nucleo_matcher::Utf32Str::Unicode(chars),
483 };
484 let Some(score) = self.matcher.fuzzy_match(haystack, needle) else {
485 return false;
486 };
487
488 push_top_match(
489 &mut self.matches,
490 self.limit,
491 score as u32,
492 path.to_string(),
493 match_type,
494 );
495 true
496 }
497}
498
499fn push_top_match(
500 matches: &mut BinaryHeap<Reverse<(u32, String, MatchType)>>,
501 limit: usize,
502 score: u32,
503 path: String,
504 match_type: MatchType,
505) -> bool {
506 let candidate = (score, path, match_type);
507 if matches.len() < limit {
508 matches.push(Reverse(candidate));
509 return true;
510 }
511
512 let Some(minimum) = matches.peek().map(|entry| &entry.0) else {
513 return false;
514 };
515
516 if &candidate <= minimum {
517 return false;
518 }
519
520 matches.pop();
521 matches.push(Reverse(candidate));
522 true
523}
524
525pub async fn run_with_index(
539 config: FileSearchConfig,
540 index_cache: &FileIndexCache,
541) -> anyhow::Result<FileSearchResults> {
542 let limit = config.limit.get();
543 let cancel_flag = &config.cancel_flag;
544 let compute_indices = config.compute_indices;
545
546 let index = index_cache.get_or_build().await?;
548
549 if cancel_flag.load(Ordering::Relaxed) {
551 return Ok(FileSearchResults {
552 matches: Vec::new(),
553 total_match_count: 0,
554 });
555 }
556
557 let matched_paths = index.query(&config.pattern_text, limit, None);
559 let total_match_count = matched_paths.len();
560
561 let matches = matched_paths
563 .into_iter()
564 .map(|(score, path, match_type)| FileMatch {
565 score,
566 path,
567 match_type,
568 indices: if compute_indices {
569 Some(Vec::new())
570 } else {
571 None
572 },
573 })
574 .collect();
575
576 Ok(FileSearchResults {
577 matches,
578 total_match_count,
579 })
580}
581
582pub fn run(config: FileSearchConfig) -> anyhow::Result<FileSearchResults> {
592 run_with_policy(config, true, false)
593}
594
595pub fn run_bounded_no_follow(config: FileSearchConfig) -> anyhow::Result<FileSearchResults> {
601 run_bounded_no_follow_with_visit(config, |_| {})
602}
603
604fn run_bounded_no_follow_with_visit(
605 config: FileSearchConfig,
606 mut visit: impl FnMut(&Path),
607) -> anyhow::Result<FileSearchResults> {
608 let limit = config.limit.get();
609 let search_directory = &config.search_directory;
610 let mut walk_builder = ignore::WalkBuilder::new(search_directory);
611 vtcode_commons::walk::apply_defaults(&mut walk_builder);
612 walk_builder
613 .follow_links(false)
614 .require_git(false)
615 .sort_by_file_path(|left, right| left.cmp(right));
616
617 if !config.respect_gitignore {
618 walk_builder
619 .git_ignore(false)
620 .git_global(false)
621 .git_exclude(false)
622 .ignore(false)
623 .parents(false);
624 }
625
626 if !config.exclude.is_empty() {
627 let mut override_builder = ignore::overrides::OverrideBuilder::new(search_directory);
628 for exclude_pattern in &config.exclude {
629 override_builder.add(&format!("!{exclude_pattern}"))?;
630 }
631 walk_builder.overrides(override_builder.build()?);
632 }
633
634 let mut matches = BestMatchesList::new(limit, &config.pattern_text);
635 let mut matching_count = 0usize;
636 for result in walk_builder.build() {
637 if config.cancel_flag.load(Ordering::Relaxed) {
638 break;
639 }
640 let entry = match result {
641 Ok(entry) => entry,
642 Err(_) => continue,
643 };
644 visit(entry.path());
645 if !entry
646 .file_type()
647 .is_some_and(|file_type| file_type.is_file())
648 {
649 continue;
650 }
651 let Some(relative_path) = entry
652 .path()
653 .strip_prefix(search_directory)
654 .ok()
655 .and_then(|path| path.to_str())
656 .filter(|path| !path.is_empty())
657 else {
658 continue;
659 };
660 if matches.record_match(relative_path, MatchType::File) {
661 matching_count += 1;
662 if matching_count >= limit {
663 break;
664 }
665 }
666 }
667
668 let matches = matches
669 .matches
670 .into_sorted_vec()
671 .into_iter()
672 .map(|Reverse((score, path, match_type))| FileMatch {
673 score,
674 path,
675 match_type,
676 indices: config.compute_indices.then(Vec::new),
677 })
678 .collect();
679
680 Ok(FileSearchResults {
681 matches,
682 total_match_count: matching_count + usize::from(matching_count >= limit),
685 })
686}
687
688fn run_with_policy(
689 config: FileSearchConfig,
690 follow_links: bool,
691 files_only: bool,
692) -> anyhow::Result<FileSearchResults> {
693 let limit = config.limit.get();
694 let search_directory = &config.search_directory;
695 let exclude = &config.exclude;
696 let threads = config.threads.get();
697 let cancel_flag = &config.cancel_flag;
698 let compute_indices = config.compute_indices;
699 let respect_gitignore = config.respect_gitignore;
700
701 let walker = build_parallel_walker(
702 search_directory,
703 exclude,
704 threads,
705 respect_gitignore,
706 follow_links,
707 )?;
708
709 let best_matchers_per_worker: Vec<Arc<Mutex<BestMatchesList>>> = (0..threads)
712 .map(|_| {
713 Arc::new(Mutex::new(BestMatchesList::new(
714 limit,
715 &config.pattern_text,
716 )))
717 })
718 .collect();
719
720 let total_match_count = Arc::new(AtomicUsize::new(0));
721
722 let worker_counter = AtomicUsize::new(0);
725 let worker_count = best_matchers_per_worker.len();
726 walker.run(|| {
727 let worker_id = worker_counter.fetch_add(1, Ordering::Relaxed) % worker_count;
728 let best_list = best_matchers_per_worker[worker_id].clone();
729 let cancel_flag_clone = cancel_flag.clone();
730 let total_match_count_clone = total_match_count.clone();
731
732 Box::new(move |result| {
733 if cancel_flag_clone.load(Ordering::Relaxed) {
735 return ignore::WalkState::Quit;
736 }
737
738 let entry = match result {
739 Ok(e) => e,
740 Err(_) => return ignore::WalkState::Continue,
741 };
742
743 let relative_path = entry
745 .path()
746 .strip_prefix(search_directory)
747 .ok()
748 .and_then(|p| p.to_str());
749
750 let path_to_match = match relative_path {
751 Some(p) if !p.is_empty() => p,
752 _ => return ignore::WalkState::Continue, };
754
755 let match_type = if entry.path().is_dir() {
756 MatchType::Directory
757 } else {
758 MatchType::File
759 };
760
761 if files_only && match_type == MatchType::Directory {
762 return ignore::WalkState::Continue;
763 }
764
765 {
767 let mut list = best_list.lock();
768 if list.record_match(path_to_match, match_type) {
769 total_match_count_clone.fetch_add(1, Ordering::Relaxed);
770 }
771 }
772
773 ignore::WalkState::Continue
774 })
775 });
776
777 let worker_heaps: Vec<BinaryHeap<Reverse<(u32, String, MatchType)>>> = best_matchers_per_worker
779 .into_iter()
780 .map(|arc| std::mem::take(&mut arc.lock().matches))
781 .collect();
782 let merged_matches = merge_top_k(worker_heaps, limit);
783
784 let matches = merged_matches
786 .into_sorted_vec()
787 .into_iter()
788 .map(|Reverse((score, path, match_type))| FileMatch {
789 score,
790 path,
791 match_type,
792 indices: if compute_indices {
793 Some(Vec::new())
794 } else {
795 None
796 },
797 })
798 .collect();
799
800 Ok(FileSearchResults {
801 matches,
802 total_match_count: total_match_count.load(Ordering::Relaxed),
803 })
804}
805
806#[cfg(test)]
807mod tests {
808 use super::{FileSearchConfig, run_bounded_no_follow, run_bounded_no_follow_with_visit};
809 use std::num::NonZero;
810 use std::sync::Arc;
811 use std::sync::atomic::AtomicBool;
812 use tempfile::TempDir;
813
814 fn bounded_paths(workspace: &std::path::Path) -> Vec<String> {
815 run_bounded_no_follow(FileSearchConfig {
816 pattern_text: "widget".to_string(),
817 limit: NonZero::new(2).expect("non-zero limit"),
818 search_directory: workspace.to_path_buf(),
819 exclude: Vec::new(),
820 threads: NonZero::new(4).expect("non-zero threads"),
821 cancel_flag: Arc::new(AtomicBool::new(false)),
822 compute_indices: false,
823 respect_gitignore: true,
824 })
825 .expect("bounded path search")
826 .matches
827 .into_iter()
828 .map(|candidate| candidate.path)
829 .collect()
830 }
831
832 #[test]
833 fn bounded_path_selection_is_stable_across_repeated_walks() {
834 let workspace = TempDir::new().expect("workspace");
835 for directory in ["z", "a", "m", "b", "y"] {
836 let directory = workspace.path().join(directory);
837 std::fs::create_dir(&directory).expect("fixture directory");
838 std::fs::write(directory.join("widget.rs"), "fn widget() {}\n")
839 .expect("fixture source");
840 }
841
842 let expected = bounded_paths(workspace.path());
843 assert_eq!(expected.len(), 2);
844 for _ in 0..20 {
845 assert_eq!(bounded_paths(workspace.path()), expected);
846 }
847 }
848
849 #[test]
850 fn bounded_path_selection_is_the_sorted_prefix_and_stops_early() {
851 let workspace = TempDir::new().expect("workspace");
852 for directory in ["z", "a", "m", "b", "y"] {
853 let directory = workspace.path().join(directory);
854 std::fs::create_dir(&directory).expect("fixture directory");
855 std::fs::write(directory.join("widget.rs"), "fn widget() {}\n")
856 .expect("fixture source");
857 }
858 let mut visited = Vec::new();
859
860 let results = run_bounded_no_follow_with_visit(
861 FileSearchConfig {
862 pattern_text: "widget".to_string(),
863 limit: NonZero::new(2).expect("non-zero limit"),
864 search_directory: workspace.path().to_path_buf(),
865 exclude: Vec::new(),
866 threads: NonZero::new(4).expect("non-zero threads"),
867 cancel_flag: Arc::new(AtomicBool::new(false)),
868 compute_indices: false,
869 respect_gitignore: true,
870 },
871 |path| visited.push(path.to_path_buf()),
872 )
873 .expect("bounded path search");
874 let mut paths = results
875 .matches
876 .into_iter()
877 .map(|candidate| candidate.path)
878 .collect::<Vec<_>>();
879 paths.sort();
880
881 assert_eq!(paths, vec!["a/widget.rs", "b/widget.rs"]);
882 assert!(
883 visited.len() < 11,
884 "the bounded route must stop before traversing the complete fixture tree"
885 );
886 assert_eq!(results.total_match_count, 3);
887 }
888}